diff --git a/internal/llm/chatgpt_test.go b/internal/llm/chatgpt_test.go index 6d8b03464..16276b778 100644 --- a/internal/llm/chatgpt_test.go +++ b/internal/llm/chatgpt_test.go @@ -14,6 +14,7 @@ import ( "net/http/httptest" "net/url" "strings" + "sync/atomic" "testing" "github.com/alibaba/open-code-review/internal/chatgptauth" @@ -159,3 +160,55 @@ func TestChatGPTResolveEndpoint(t *testing.T) { t.Fatal("missing creds accepted") } } + +// Plan mode hands newOpenAISDK its own retry count and redirect policy; the +// other tests here swap c.sdk out, so only this one sees what the constructor +// actually built. +func TestNewOpenAIResponsesClientPlanModeSDKSettings(t *testing.T) { + for _, tc := range []struct { + name string + plan bool + wantAttempts int32 + wantRedirects int32 + }{ + {name: "api key", plan: false, wantAttempts: 6, wantRedirects: 1}, + {name: "chatgpt plan", plan: true, wantAttempts: 1, wantRedirects: 0}, + } { + t.Run(tc.name, func(t *testing.T) { + var redirect atomic.Bool + var attempts, followed atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/elsewhere" { + followed.Add(1) + w.WriteHeader(http.StatusTeapot) + return + } + if redirect.Load() { + http.Redirect(w, r, "/elsewhere", http.StatusFound) + return + } + attempts.Add(1) + w.Header().Set("Retry-After-Ms", "0") + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(server.Close) + + c := NewOpenAIResponsesClient(ClientConfig{URL: server.URL, APIKey: "k", Model: "m", ChatGPTPlan: tc.plan}) + params := responses.ResponseNewParams{ + Model: "m", + Input: responses.ResponseNewParamsInputUnion{OfString: openai.String("hi")}, + } + + _, _ = c.sdk.Responses.New(context.Background(), params) + if got := attempts.Load(); got != tc.wantAttempts { + t.Errorf("attempts on HTTP 500 = %d, want %d", got, tc.wantAttempts) + } + + redirect.Store(true) + _, _ = c.sdk.Responses.New(context.Background(), params) + if got := followed.Load(); got != tc.wantRedirects { + t.Errorf("redirects followed = %d, want %d", got, tc.wantRedirects) + } + }) + } +} diff --git a/internal/llm/client.go b/internal/llm/client.go index e0b7aa983..3f7311ed1 100644 --- a/internal/llm/client.go +++ b/internal/llm/client.go @@ -18,6 +18,7 @@ import ( "io" "net/http" "os" + "slices" "strings" "sync" "time" @@ -592,12 +593,7 @@ type OpenAIClient struct { // ExtraHeaders are applied per request (not baked into the SDK client) so // SessionKeyTemplateVar can expand to the session key each request carries. func NewOpenAIClient(cfg ClientConfig) *OpenAIClient { - if cfg.Timeout <= 0 { - cfg.Timeout = 5 * time.Minute - } - if cfg.SessionKey == "" { - cfg.SessionKey = NewSessionKey() - } + applyClientDefaults(&cfg) baseURL := strings.TrimRight(cfg.URL, "/") if !strings.HasSuffix(baseURL, "/chat/completions") { cfg.URL = baseURL + "/chat/completions" @@ -605,31 +601,74 @@ func NewOpenAIClient(cfg ClientConfig) *OpenAIClient { sdkBaseURL := strings.TrimSuffix(strings.TrimRight(cfg.URL, "/"), "/chat/completions") + return &OpenAIClient{ + cfg: cfg, + sdk: newOpenAISDK(cfg, sdkBaseURL, 5, httpClientWithHeaderTimeout(cfg.Timeout)), + } +} + +func applyClientDefaults(cfg *ClientConfig) { + if cfg.Timeout <= 0 { + cfg.Timeout = 5 * time.Minute + } + if cfg.SessionKey == "" { + cfg.SessionKey = NewSessionKey() + } +} + +// Callers normalize cfg.URL themselves: the two OpenAI clients keep different +// trailing-slash contracts on it. +func newOpenAISDK(cfg ClientConfig, sdkBaseURL string, maxRetries int, httpClient *http.Client) openai.Client { opts := []openaiopt.RequestOption{ openaiopt.WithAPIKey(cfg.APIKey), openaiopt.WithBaseURL(sdkBaseURL), - openaiopt.WithMaxRetries(5), + openaiopt.WithMaxRetries(maxRetries), openaiopt.WithHeader("User-Agent", userAgent("")), openaiopt.WithRequestTimeout(cfg.Timeout), - openaiopt.WithHTTPClient(httpClientWithHeaderTimeout(cfg.Timeout)), + openaiopt.WithHTTPClient(httpClient), } + opts = append(opts, openaiopt.WithMiddleware(sdkMiddlewares(cfg)...)) + return openai.NewClient(opts...) +} + +// Raw must register before the retry observer: the SDK wraps middlewares +// last-in-innermost, and raw's full-body read plus disk write would otherwise +// inflate the observer's DurationToHeadersMS. +func sdkMiddlewares(cfg ClientConfig) []retryObserver { + var mws []retryObserver if mw := retryCodesMiddleware(cfg.RetryCodes); mw != nil { - opts = append(opts, openaiopt.WithMiddleware(mw)) + mws = append(mws, mw) } - // Raw must register before the retry observer: the SDK wraps middlewares - // last-in-innermost, and raw's full-body read plus disk write would - // otherwise inflate the observer's DurationToHeadersMS. if cfg.rawHolder != nil { - opts = append(opts, openaiopt.WithMiddleware(newRawMiddleware(cfg.rawHolder))) + mws = append(mws, newRawMiddleware(cfg.rawHolder)) } if cfg.retryCollector != nil { - opts = append(opts, openaiopt.WithMiddleware(newRetryObserver(cfg.retryCollector))) + mws = append(mws, newRetryObserver(cfg.retryCollector)) } + return mws +} - return &OpenAIClient{ - cfg: cfg, - sdk: openai.NewClient(opts...), +func (cfg ClientConfig) requestSessionKey(ctx context.Context) string { + if k := SessionKeyFromContext(ctx); k != "" { + return k } + return cfg.SessionKey +} + +func openAIRequestOptions(ctx context.Context, cfg ClientConfig, skipBodyKeys ...string) []openaiopt.RequestOption { + sessionKey := cfg.requestSessionKey(ctx) + + var opts []openaiopt.RequestOption + for k, v := range expandSessionKeyInHeaders(cfg.ExtraHeaders, sessionKey) { + opts = append(opts, openaiopt.WithHeader(k, v)) + } + for k, v := range expandSessionKeyInBody(cfg.ExtraBody, sessionKey) { + if slices.Contains(skipBodyKeys, k) { + continue + } + opts = append(opts, openaiopt.WithJSONSet(k, v)) + } + return opts } // ChatRequest represents the payload for a chat completion call. @@ -645,22 +684,12 @@ type ChatRequest struct { // CompletionsWithCtx sends a chat completion request with context support for cancellation and timeout. // -// The deferred finalizeRequest is the client boundary for the retry report: it is +// The deferred finalizeOnExit is the client boundary for the retry report: it is // the only place that knows the logical request is over, and it covers every exit // path including the streaming branch, the EOF recovery and a panic. Results are // named so the defer can read the error actually returned. func (c *OpenAIClient) CompletionsWithCtx(ctx context.Context, req ChatRequest) (resp *ChatResponse, err error) { - defer func() { - // A panic still has to finalize, or the entry stays unfinalized and Freeze - // drops the whole run's report. The panic value itself is re-raised - // unchanged so agent.go's per-file recovery behaves exactly as before. - if r := recover(); r != nil { - finalizeRequest(ctx, c.cfg.retryCollector, errRequestPanicked) - panic(r) - } - err = describeTimeout(ctx, err) - finalizeRequest(ctx, c.cfg.retryCollector, err) - }() + defer finalizeOnExit(ctx, c.cfg.retryCollector, &err) model := req.Model if model == "" { @@ -669,29 +698,15 @@ func (c *OpenAIClient) CompletionsWithCtx(ctx context.Context, req ChatRequest) params := c.buildOpenAIParams(model, req) - sessionKey := c.cfg.SessionKey - if k := SessionKeyFromContext(ctx); k != "" { - sessionKey = k - } - - var opts []openaiopt.RequestOption - for k, v := range expandSessionKeyInHeaders(c.cfg.ExtraHeaders, sessionKey) { - opts = append(opts, openaiopt.WithHeader(k, v)) - } - for k, v := range expandSessionKeyInBody(c.cfg.ExtraBody, sessionKey) { - // Skip the "stream" key here. The streaming decision below uses a - // dedicated boolean check, and when streaming is enabled the SDK's - // NewStreaming method sets stream=true on the wire itself. When - // streaming is NOT enabled, leaving the key in the body would make - // the API answer with text/event-stream and the non-streaming path - // fails to decode (see issue #647). "stream_options" is owned by the - // streaming branch below for the same reason: providers reject it - // unless stream is true. - if k == "stream" || k == "stream_options" { - continue - } - opts = append(opts, openaiopt.WithJSONSet(k, v)) - } + // Skip the "stream" key here. The streaming decision below uses a + // dedicated boolean check, and when streaming is enabled the SDK's + // NewStreaming method sets stream=true on the wire itself. When + // streaming is NOT enabled, leaving the key in the body would make + // the API answer with text/event-stream and the non-streaming path + // fails to decode (see issue #647). "stream_options" is owned by the + // streaming branch below for the same reason: providers reject it + // unless stream is true. + opts := openAIRequestOptions(ctx, c.cfg, "stream", "stream_options") if stream, ok := c.cfg.ExtraBody["stream"].(bool); ok && stream { if streamOptions, ok := c.cfg.ExtraBody["stream_options"]; !ok { // OpenAI-compatible servers omit token usage from streams unless @@ -1159,12 +1174,7 @@ type AnthropicClient struct { // available to ExtraHeaders/ExtraBody via SessionKeyTemplateVar, applied per // request so it can expand to the session key each request carries. func NewAnthropicClient(cfg ClientConfig) *AnthropicClient { - if cfg.Timeout <= 0 { - cfg.Timeout = 5 * time.Minute - } - if cfg.SessionKey == "" { - cfg.SessionKey = NewSessionKey() - } + applyClientDefaults(&cfg) if !strings.HasSuffix(cfg.URL, "/v1/messages") && !strings.HasSuffix(cfg.URL, "/v1/messages/") { baseURL := strings.TrimRight(cfg.URL, "/") if !strings.HasSuffix(baseURL, "/v1/messages") { @@ -1204,16 +1214,7 @@ func NewAnthropicClient(cfg ClientConfig) *AnthropicClient { ) } - if mw := retryCodesMiddleware(cfg.RetryCodes); mw != nil { - opts = append(opts, option.WithMiddleware(mw)) - } - // Raw before the retry observer; see NewOpenAIClient for why order matters. - if cfg.rawHolder != nil { - opts = append(opts, option.WithMiddleware(newRawMiddleware(cfg.rawHolder))) - } - if cfg.retryCollector != nil { - opts = append(opts, option.WithMiddleware(newRetryObserver(cfg.retryCollector))) - } + opts = append(opts, option.WithMiddleware(sdkMiddlewares(cfg)...)) return &AnthropicClient{ cfg: cfg, @@ -1234,12 +1235,7 @@ func NewAnthropicClient(cfg ClientConfig) *AnthropicClient { // AWS_BEARER_TOKEN_BEDROCK if set. Region comes from AWS_REGION or the active // profile. func NewAnthropicBedrockClient(cfg ClientConfig) *AnthropicClient { - if cfg.Timeout <= 0 { - cfg.Timeout = 5 * time.Minute - } - if cfg.SessionKey == "" { - cfg.SessionKey = NewSessionKey() - } + applyClientDefaults(&cfg) // cfg.URL is deliberately unused: bedrock.WithConfig is appended last and // installs its own base URL from the resolved region, so anything set here @@ -1263,16 +1259,7 @@ func NewAnthropicBedrockClient(cfg ClientConfig) *AnthropicClient { } // ExtraHeaders are applied per request in CompletionsWithCtx, where the // session key template can expand — same as the plain Anthropic client. - if mw := retryCodesMiddleware(cfg.RetryCodes); mw != nil { - opts = append(opts, option.WithMiddleware(mw)) - } - // Raw before the retry observer; see NewOpenAIClient for why order matters. - if cfg.rawHolder != nil { - opts = append(opts, option.WithMiddleware(newRawMiddleware(cfg.rawHolder))) - } - if cfg.retryCollector != nil { - opts = append(opts, option.WithMiddleware(newRetryObserver(cfg.retryCollector))) - } + opts = append(opts, option.WithMiddleware(sdkMiddlewares(cfg)...)) // Load the AWS config here rather than calling bedrock.WithLoadDefaultConfig, // which panics on failure. @@ -1460,14 +1447,7 @@ func anthropicThinkingBudgetTokens(v any) (int64, bool) { // named. A parameter-building failure returns before any HTTP attempt, so // Finalize finds no entry and the request stays out of the report entirely. func (c *AnthropicClient) CompletionsWithCtx(ctx context.Context, req ChatRequest) (resp *ChatResponse, err error) { - defer func() { - if r := recover(); r != nil { - finalizeRequest(ctx, c.cfg.retryCollector, errRequestPanicked) - panic(r) - } - err = describeTimeout(ctx, err) - finalizeRequest(ctx, c.cfg.retryCollector, err) - }() + defer finalizeOnExit(ctx, c.cfg.retryCollector, &err) if c.initErr != nil { return nil, c.initErr @@ -1483,10 +1463,7 @@ func (c *AnthropicClient) CompletionsWithCtx(ctx context.Context, req ChatReques return nil, err } - sessionKey := c.cfg.SessionKey - if k := SessionKeyFromContext(ctx); k != "" { - sessionKey = k - } + sessionKey := c.cfg.requestSessionKey(ctx) var opts []option.RequestOption for k, v := range expandSessionKeyInHeaders(c.cfg.ExtraHeaders, sessionKey) { diff --git a/internal/llm/responses_client.go b/internal/llm/responses_client.go index 4cd8d41bf..316266513 100644 --- a/internal/llm/responses_client.go +++ b/internal/llm/responses_client.go @@ -11,7 +11,6 @@ import ( "net/http" "sort" "strings" - "time" "github.com/alibaba/open-code-review/internal/chatgptauth" openai "github.com/openai/openai-go/v3" @@ -40,12 +39,7 @@ type OpenAIResponsesClient struct { // ExtraHeaders are applied per request (not baked into the SDK client) // so SessionKeyTemplateVar can expand to the session key each request carries. func NewOpenAIResponsesClient(cfg ClientConfig) *OpenAIResponsesClient { - if cfg.Timeout <= 0 { - cfg.Timeout = 5 * time.Minute - } - if cfg.SessionKey == "" { - cfg.SessionKey = NewSessionKey() - } + applyClientDefaults(&cfg) ensureResponsesEndpoint(&cfg) sdkBaseURL := strings.TrimSuffix(strings.TrimRight(cfg.URL, "/"), "/responses") @@ -55,28 +49,9 @@ func NewOpenAIResponsesClient(cfg ClientConfig) *OpenAIResponsesClient { retries = 0 httpClient.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse } } - opts := []openaiopt.RequestOption{ - openaiopt.WithAPIKey(cfg.APIKey), - openaiopt.WithBaseURL(sdkBaseURL), - openaiopt.WithMaxRetries(retries), - openaiopt.WithHeader("User-Agent", userAgent("")), - openaiopt.WithRequestTimeout(cfg.Timeout), - openaiopt.WithHTTPClient(httpClient), - } - if mw := retryCodesMiddleware(cfg.RetryCodes); mw != nil { - opts = append(opts, openaiopt.WithMiddleware(mw)) - } - // Raw before the retry observer; see NewOpenAIClient for why order matters. - if cfg.rawHolder != nil { - opts = append(opts, openaiopt.WithMiddleware(newRawMiddleware(cfg.rawHolder))) - } - if cfg.retryCollector != nil { - opts = append(opts, openaiopt.WithMiddleware(newRetryObserver(cfg.retryCollector))) - } - client := &OpenAIResponsesClient{ cfg: cfg, - sdk: openai.NewClient(opts...), + sdk: newOpenAISDK(cfg, sdkBaseURL, retries, httpClient), } if cfg.ChatGPTPlan { client.chatGPTToken = func(ctx context.Context) (string, error) { @@ -112,18 +87,11 @@ func ensureResponsesEndpoint(cfg *ClientConfig) { // CompletionsWithCtx sends a Responses API request and maps the result back to // the shared ChatResponse shape. // -// The deferred finalizeRequest is this client's boundary for the retry report; +// The deferred finalizeOnExit is this client's boundary for the retry report; // see the OpenAI Chat Completions counterpart for why it is deferred and why the // results are named. func (c *OpenAIResponsesClient) CompletionsWithCtx(ctx context.Context, req ChatRequest) (resp *ChatResponse, err error) { - defer func() { - if r := recover(); r != nil { - finalizeRequest(ctx, c.cfg.retryCollector, errRequestPanicked) - panic(r) - } - err = describeTimeout(ctx, err) - finalizeRequest(ctx, c.cfg.retryCollector, err) - }() + defer finalizeOnExit(ctx, c.cfg.retryCollector, &err) var planToken string if c.cfg.ChatGPTPlan { @@ -150,27 +118,14 @@ func (c *OpenAIResponsesClient) CompletionsWithCtx(ctx context.Context, req Chat params := c.buildResponsesParams(model, req) - sessionKey := c.cfg.SessionKey - if k := SessionKeyFromContext(ctx); k != "" { - sessionKey = k - } - var opts []openaiopt.RequestOption if c.cfg.ChatGPTPlan { opts = append(opts, openaiopt.WithAPIKey(planToken)) } - for k, v := range expandSessionKeyInHeaders(c.cfg.ExtraHeaders, sessionKey) { - opts = append(opts, openaiopt.WithHeader(k, v)) - } - for k, v := range expandSessionKeyInBody(c.cfg.ExtraBody, sessionKey) { - // Streaming is selected only by the resolved provider. Forwarding an - // extra_body.stream setting could conflict with that selection or make an - // ordinary Responses.New call receive SSE, so it is always dropped here. - if k == "stream" { - continue - } - opts = append(opts, openaiopt.WithJSONSet(k, v)) - } + // Streaming is selected only by the resolved provider. Forwarding an + // extra_body.stream setting could conflict with that selection or make an + // ordinary Responses.New call receive SSE, so it is always dropped here. + opts = append(opts, openAIRequestOptions(ctx, c.cfg, "stream")...) var sdkResp *responses.Response if c.cfg.RequiresStreaming { diff --git a/internal/llm/retry_boundary.go b/internal/llm/retry_boundary.go index 0dc832dcb..e4226fb5a 100644 --- a/internal/llm/retry_boundary.go +++ b/internal/llm/retry_boundary.go @@ -140,3 +140,15 @@ func finalizeRequest(ctx context.Context, collector *RetryCollector, reqErr erro } collector.Finalize(meta, reqErr, errors.Is(ctx.Err(), context.Canceled)) } + +// finalizeOnExit must be deferred directly, not wrapped in a closure, or its +// recover never sees the panic. The panic is re-raised after finalizing; see +// errRequestPanicked for why it has to finalize at all. +func finalizeOnExit(ctx context.Context, collector *RetryCollector, err *error) { + if r := recover(); r != nil { + finalizeRequest(ctx, collector, errRequestPanicked) + panic(r) + } + *err = describeTimeout(ctx, *err) + finalizeRequest(ctx, collector, *err) +} diff --git a/internal/llm/retry_boundary_test.go b/internal/llm/retry_boundary_test.go index 67d65de02..d4f6c10e5 100644 --- a/internal/llm/retry_boundary_test.go +++ b/internal/llm/retry_boundary_test.go @@ -713,6 +713,59 @@ func TestFinalizeRequestWithPanicSentinel(t *testing.T) { } } +func TestFinalizeOnExitFinalizesAndRepanics(t *testing.T) { + c := NewRetryCollector() + m := testMeta() + ctx := metaCtx(m) + c.RecordAttempt(m, AttemptRecord{StatusCode: http.StatusOK}, time.Time{}, time.Time{}) + + var recovered any + func() { + defer func() { recovered = recover() }() + var err error + defer finalizeOnExit(ctx, c, &err) + panic("boom") + }() + + if recovered != "boom" { + t.Errorf("recovered %v, want the original panic value", recovered) + } + if req := freezeOne(t, c); req.Outcome != OutcomeFailed { + t.Errorf("outcome = %s, want failed", req.Outcome) + } +} + +// &err is taken at defer time, so this checks the returned error is what gets +// finalized. A clean success is omitted from the report entirely. +func TestFinalizeOnExitReadsReturnedError(t *testing.T) { + finalize := func(retErr error) *RetryReport { + c := NewRetryCollector() + m := testMeta() + ctx := metaCtx(m) + c.RecordAttempt(m, AttemptRecord{StatusCode: http.StatusOK}, time.Time{}, time.Time{}) + + call := func() (err error) { + defer finalizeOnExit(ctx, c, &err) + return retErr + } + _ = call() + + rep, err := c.Freeze("test-run-id") + if err != nil { + t.Fatalf("Freeze: %v", err) + } + return rep + } + + if rep := finalize(nil); rep != nil { + t.Errorf("successful request produced a report: %+v", rep) + } + rep := finalize(errors.New("boom")) + if rep == nil || len(rep.Requests) != 1 || rep.Requests[0].Outcome != OutcomeFailed { + t.Errorf("failed request report = %+v, want one failed request", rep) + } +} + // A nil collector and a context without identity are both no-ops, so no call // site in the three clients needs to guard either. func TestBoundaryHelpersAreInertWithoutCollectorOrMeta(t *testing.T) {