From b6dacb4a3c51fc5747d06c1bd02cfd5b214306fa Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Thu, 7 May 2026 15:35:46 +0100 Subject: [PATCH] feat: add automatic key failover for AI Bridge OpenAI (#24847) ## Description Adds automatic key failover for centralized OpenAI provider, covering both chat completions and responses APIs. Same shape as the Anthropic PR: each upstream call walks the configured key pool, keys are marked **temporary** on 429 (with cooldown from `Retry-After`) and **permanent** on 401/403. Each agentic-loop iteration gets its own fresh walker so a tool-call continuation can fail over independently of the initial request. BYOK is unchanged: BYOK requests run as a single attempt with no failover. ## Changes - `config.OpenAI` carries a `KeyPool`. `Key` remains for BYOK Authorization Bearer set per interception. - Chat completions blocking interceptor: walks the pool via `newChatCompletionWithKeyFailover`, marks keys on key-specific failures, returns on first success or non-failover error. - Chat completions streaming interceptor: per-iteration walker. Pre-stream failures fail over to the next key; mid-stream errors are relayed as SSE events. - Responses blocking interceptor: extracts `newResponseWithKeyFailover` parallel to chatcompletions. - Responses streaming interceptor: per-iteration walker, retains the existing buffer-then-forward design. ## Related Issues Related to: https://github.com/coder/internal/issues/1446 Related to: https://linear.app/codercom/issue/AIGOV-197/aibridge-automatic-key-failover-for-bridged-and-passthrough-routes ## Follow-up PRs - Bedrock multi-key support. - Refactor provider vs interceptor config separation. - Record the actually-used key in the interception credential hint after failover. > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira --- aibridge/config/config.go | 11 + .../responses/blocking/http_error.txtar | 12 +- .../responses/streaming/http_error.txtar | 12 +- aibridge/intercept/chatcompletions/base.go | 89 ++- .../intercept/chatcompletions/base_test.go | 200 +++++++ .../intercept/chatcompletions/blocking.go | 56 +- .../chatcompletions/blocking_test.go | 502 +++++++++++++++++ .../intercept/chatcompletions/streaming.go | 158 ++++-- .../chatcompletions/streaming_test.go | 505 +++++++++++++++++- aibridge/intercept/openai_errors.go | 14 + aibridge/intercept/responses/base.go | 120 ++++- aibridge/intercept/responses/base_test.go | 198 +++++++ aibridge/intercept/responses/blocking.go | 58 +- aibridge/intercept/responses/blocking_test.go | 492 +++++++++++++++++ aibridge/intercept/responses/streaming.go | 76 ++- .../intercept/responses/streaming_test.go | 499 +++++++++++++++++ .../integrationtest/keypool_failover_test.go | 132 +++++ .../internal/integrationtest/trace_test.go | 11 +- aibridge/provider/openai.go | 58 +- enterprise/cli/aibridged.go | 11 +- 20 files changed, 3112 insertions(+), 102 deletions(-) create mode 100644 aibridge/intercept/chatcompletions/blocking_test.go create mode 100644 aibridge/intercept/openai_errors.go create mode 100644 aibridge/intercept/responses/blocking_test.go create mode 100644 aibridge/intercept/responses/streaming_test.go diff --git a/aibridge/config/config.go b/aibridge/config/config.go index 676e891c2b..5805741f60 100644 --- a/aibridge/config/config.go +++ b/aibridge/config/config.go @@ -50,11 +50,22 @@ type AWSBedrock struct { BaseURL string } +// OpenAI carries configuration for an OpenAI provider. +// +// Authentication is mutually exclusive across these two fields, +// set per interception in the provider's CreateInterceptor: +// - KeyPool: centralized requests with automatic key failover. +// - Key: BYOK with Authorization Bearer (single attempt, no +// failover). +// +// TODO(ssncferreira): consolidate the authentication fields per +// https://github.com/coder/aibridge/issues/266. type OpenAI struct { // Name is the provider instance name. If empty, defaults to "openai". Name string BaseURL string Key string + KeyPool *keypool.Pool APIDumpDir string CircuitBreaker *CircuitBreaker SendActorHeaders bool diff --git a/aibridge/fixtures/openai/responses/blocking/http_error.txtar b/aibridge/fixtures/openai/responses/blocking/http_error.txtar index 24986a2cea..42183ac8ae 100644 --- a/aibridge/fixtures/openai/responses/blocking/http_error.txtar +++ b/aibridge/fixtures/openai/responses/blocking/http_error.txtar @@ -6,16 +6,16 @@ } -- non-streaming -- -HTTP/2.0 401 Unauthorized -Content-Length: 234 +HTTP/2.0 400 Bad Request +Content-Length: 281 Content-Type: application/json { "error": { - "message": "Incorrect API key provided: sk-***. You can find your API key at https://platform.openai.com/account/api-keys.", - "type": "authentication_error", - "param": null, - "code": "invalid_api_key" + "message": "Input tokens exceed the configured limit of 272000 tokens. Your messages resulted in 3148588 tokens. Please reduce the length of the messages.", + "type": "invalid_request_error", + "param": "messages", + "code": "context_length_exceeded" } } diff --git a/aibridge/fixtures/openai/responses/streaming/http_error.txtar b/aibridge/fixtures/openai/responses/streaming/http_error.txtar index 9c7827fff8..77ecfe255c 100644 --- a/aibridge/fixtures/openai/responses/streaming/http_error.txtar +++ b/aibridge/fixtures/openai/responses/streaming/http_error.txtar @@ -6,16 +6,16 @@ } -- streaming -- -HTTP/2.0 429 Too Many Requests -Content-Length: 176 +HTTP/2.0 400 Bad Request +Content-Length: 281 Content-Type: application/json { "error": { - "message": "Rate limit exceeded. Please try again in 20 seconds.", - "type": "rate_limit_error", - "param": null, - "code": "rate_limit_exceeded" + "message": "Input tokens exceed the configured limit of 272000 tokens. Your messages resulted in 3148588 tokens. Please reduce the length of the messages.", + "type": "invalid_request_error", + "param": "messages", + "code": "context_length_exceeded" } } diff --git a/aibridge/intercept/chatcompletions/base.go b/aibridge/intercept/chatcompletions/base.go index aa84e7dead..b8f896b92e 100644 --- a/aibridge/intercept/chatcompletions/base.go +++ b/aibridge/intercept/chatcompletions/base.go @@ -5,8 +5,11 @@ import ( "encoding/json" "errors" "fmt" + "math" "net/http" + "strconv" "strings" + "time" "github.com/google/uuid" "github.com/openai/openai-go/v3" @@ -20,6 +23,7 @@ import ( aibcontext "github.com/coder/coder/v2/aibridge/context" "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/intercept/apidump" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/recorder" "github.com/coder/coder/v2/aibridge/tracing" @@ -44,8 +48,19 @@ type interceptionBase struct { credential intercept.CredentialInfo } +// newCompletionsService builds the SDK service used for upstream +// calls. BYOK auth is set here. Centralized auth is set +// per-attempt by the failover loop. func (i *interceptionBase) newCompletionsService() openai.ChatCompletionService { - opts := []option.RequestOption{option.WithAPIKey(i.cfg.Key), option.WithBaseURL(i.cfg.BaseURL)} + // TODO(ssncferreira): validate auth is configured per + // https://github.com/coder/aibridge/issues/266. + + var opts []option.RequestOption + // BYOK auth. + if i.cfg.KeyPool == nil { + opts = append(opts, option.WithAPIKey(i.cfg.Key)) + } + opts = append(opts, option.WithBaseURL(i.cfg.BaseURL)) // Add extra headers if configured. // Some providers require additional headers that are not added by the SDK. @@ -179,6 +194,10 @@ func (i *interceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *res } w.Header().Set("Content-Type", "application/json") + // Set Retry-After when a cooldown is configured. + if oaiErr.RetryAfter > 0 { + w.Header().Set("Retry-After", strconv.Itoa(int(math.Ceil(oaiErr.RetryAfter.Seconds())))) + } w.WriteHeader(oaiErr.StatusCode) out, err := json.Marshal(oaiErr) @@ -190,13 +209,58 @@ func (i *interceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *res "type": "error", "message":"error marshaling upstream error", "code": "server_error" - }, + } }`)) } else { _, _ = w.Write(out) } } +// For centralized requests, markKeyOnError extracts an OpenAI +// SDK error from err and marks the key based on its status +// code. Returns true if the status was a key-specific failover +// trigger so callers can retry with the next key. +func (i *interceptionBase) markKeyOnError(ctx context.Context, key *keypool.Key, err error) bool { + if i.cfg.KeyPool == nil { + return false + } + var apiErr *openai.Error + if !errors.As(err, &apiErr) { + return false + } + return keypool.MarkKeyOnStatus( + ctx, key, apiErr.Response, + i.logger, i.providerName, + ) +} + +// processKeyPoolError translates a keypool exhaustion error +// into a developer-facing responseError shaped for the OpenAI +// API. Returns nil if err is not an exhaustion error. +func processKeyPoolError(err error) *responseError { + var transient *keypool.TransientKeyPoolError + switch { + case errors.As(err, &transient): + return newErrorResponse( + "all configured keys are rate-limited", + intercept.OpenAIErrTypeRateLimit, + intercept.OpenAIErrCodeRateLimit, + http.StatusTooManyRequests, + transient.RetryAfter, + ) + case errors.Is(err, keypool.ErrPermanentKeyPool): + return newErrorResponse( + "all configured keys failed authentication", + intercept.OpenAIErrTypeAPI, + intercept.OpenAIErrCodeServer, + http.StatusBadGateway, + 0, + ) + default: + return nil + } +} + func (i *interceptionBase) hasInjectableTools() bool { return i.mcpProxy != nil && len(i.mcpProxy.ListTools()) > 0 } @@ -233,15 +297,7 @@ func getErrorResponse(err error) *responseError { if !errors.As(err, &apiErr) { return nil } - - return &responseError{ - ErrorObject: &shared.ErrorObject{ - Code: apiErr.Code, - Message: apiErr.Message, - Type: apiErr.Type, - }, - StatusCode: apiErr.StatusCode, - } + return newErrorResponse(apiErr.Message, apiErr.Type, apiErr.Code, apiErr.StatusCode, keypool.ParseRetryAfter(apiErr.Response)) } var _ error = &responseError{} @@ -249,15 +305,18 @@ var _ error = &responseError{} type responseError struct { ErrorObject *shared.ErrorObject `json:"error"` StatusCode int `json:"-"` + RetryAfter time.Duration `json:"-"` } -func newErrorResponse(msg error) *responseError { +func newErrorResponse(msg, errType, code string, status int, retryAfter time.Duration) *responseError { return &responseError{ ErrorObject: &shared.ErrorObject{ - Code: "error", - Message: msg.Error(), - Type: "error", + Code: code, + Message: msg, + Type: errType, }, + StatusCode: status, + RetryAfter: retryAfter, } } diff --git a/aibridge/intercept/chatcompletions/base_test.go b/aibridge/intercept/chatcompletions/base_test.go index 67104b9085..4f7ffc7bcd 100644 --- a/aibridge/intercept/chatcompletions/base_test.go +++ b/aibridge/intercept/chatcompletions/base_test.go @@ -1,12 +1,22 @@ package chatcompletions //nolint:testpackage // tests unexported internals import ( + "context" + "net/http" + "net/http/httptest" "testing" + "time" "github.com/openai/openai-go/v3" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/utils" + "github.com/coder/quartz" ) func TestScanForCorrelatingToolCallID(t *testing.T) { @@ -75,3 +85,193 @@ func TestScanForCorrelatingToolCallID(t *testing.T) { }) } } + +func TestProcessKeyPoolError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + expectedNil bool + expectedStatus int + expectedRetryAfter time.Duration + }{ + { + // Transient with valid keys present: 429, no Retry-After. + name: "transient_zero_retry_after", + err: &keypool.TransientKeyPoolError{}, + expectedStatus: http.StatusTooManyRequests, + expectedRetryAfter: 0, + }, + { + // Transient with cooldown: 429, Retry-After set. + name: "transient_with_retry_after", + err: &keypool.TransientKeyPoolError{RetryAfter: 5 * time.Second}, + expectedStatus: http.StatusTooManyRequests, + expectedRetryAfter: 5 * time.Second, + }, + { + // Permanent: 502 api_error. + name: "permanent_returns_502", + err: keypool.ErrPermanentKeyPool, + expectedStatus: http.StatusBadGateway, + }, + { + // Anything else: not a pool-exhaustion error. + name: "non_pool_exhaustion_error_returns_nil", + err: xerrors.New("some other error"), + expectedNil: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got := processKeyPoolError(tc.err) + if tc.expectedNil { + require.Nil(t, got) + return + } + require.NotNil(t, got) + assert.Equal(t, tc.expectedStatus, got.StatusCode) + assert.Equal(t, tc.expectedRetryAfter, got.RetryAfter) + }) + } +} + +func TestMarkKeyOnError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + expectedReturn bool + expectedState keypool.KeyState + }{ + { + // Not an *openai.Error: no status code to act on. + name: "non_api_error_returns_false", + err: xerrors.New("network failure"), + expectedReturn: false, + expectedState: keypool.KeyStateValid, + }, + { + // Rate-limited: temporary cooldown. + name: "429_marks_temporary", + err: &openai.Error{StatusCode: http.StatusTooManyRequests, Response: &http.Response{StatusCode: http.StatusTooManyRequests}}, + expectedReturn: true, + expectedState: keypool.KeyStateTemporary, + }, + { + // Auth failure: mark permanent. + name: "401_marks_permanent", + err: &openai.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}}, + expectedReturn: true, + expectedState: keypool.KeyStatePermanent, + }, + { + // Auth forbidden: mark permanent. + name: "403_marks_permanent", + err: &openai.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}}, + expectedReturn: true, + expectedState: keypool.KeyStatePermanent, + }, + { + // Server errors are not key-specific. + name: "500_does_not_mark", + err: &openai.Error{StatusCode: http.StatusInternalServerError, Response: &http.Response{StatusCode: http.StatusInternalServerError}}, + expectedReturn: false, + expectedState: keypool.KeyStateValid, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + pool, err := keypool.New([]string{"key-0"}, quartz.NewMock(t)) + require.NoError(t, err) + key, err := pool.Walker().Next() + require.NoError(t, err) + + base := &interceptionBase{cfg: config.OpenAI{KeyPool: pool}, logger: slog.Make()} + + got := base.markKeyOnError(context.Background(), key, tc.err) + assert.Equal(t, tc.expectedReturn, got) + assert.Equal(t, tc.expectedState, key.State()) + }) + } +} + +func TestWriteUpstreamError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + respErr *responseError + expectStatus int + // Empty string means the header should be absent. + expectRetryAfter string + // Substring expected in the marshaled body. Empty means no body check. + expectBodyContains string + }{ + { + // Standard error: status, code, and JSON body written. + name: "writes_status_and_body", + respErr: newErrorResponse("upstream failed", "api_error", "server_error", http.StatusBadGateway, 0), + expectStatus: http.StatusBadGateway, + expectBodyContains: `"upstream failed"`, + }, + { + // OpenAI envelope: the code field round-trips into the body. + name: "writes_code_field", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 0), + expectStatus: http.StatusTooManyRequests, + expectBodyContains: `"rate_limit_exceeded"`, + }, + { + // Whole-second retryAfter: emitted as integer seconds. + name: "retry_after_in_seconds", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 60*time.Second), + expectStatus: http.StatusTooManyRequests, + expectRetryAfter: "60", + }, + { + // 500ms rounds up to Retry-After: 1. + name: "retry_after_500ms_rounds_up_to_one", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 500*time.Millisecond), + expectStatus: http.StatusTooManyRequests, + expectRetryAfter: "1", + }, + { + // 200ms rounds up to Retry-After: 1. + name: "retry_after_200ms_rounds_up_to_one", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 200*time.Millisecond), + expectStatus: http.StatusTooManyRequests, + expectRetryAfter: "1", + }, + { + // Negative retryAfter: header omitted. + name: "negative_retry_after_omits_header", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, -1*time.Second), + expectStatus: http.StatusTooManyRequests, + expectRetryAfter: "", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + base := &interceptionBase{logger: slog.Make()} + + w := httptest.NewRecorder() + base.writeUpstreamError(w, tc.respErr) + + assert.Equal(t, tc.expectStatus, w.Code, "status code") + assert.Equal(t, "application/json", w.Header().Get("Content-Type"), "Content-Type header") + assert.Equal(t, tc.expectRetryAfter, w.Header().Get("Retry-After"), "Retry-After header") + if tc.expectBodyContains != "" { + assert.Contains(t, w.Body.String(), tc.expectBodyContains, "response body") + } + }) + } +} diff --git a/aibridge/intercept/chatcompletions/blocking.go b/aibridge/intercept/chatcompletions/blocking.go index 59c8bbb731..23d46c75cf 100644 --- a/aibridge/intercept/chatcompletions/blocking.go +++ b/aibridge/intercept/chatcompletions/blocking.go @@ -223,6 +223,13 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req return xerrors.Errorf("upstream connection closed: %w", err) } + // The failover loop may return a keypool exhaustion + // error. Check before the SDK-error path. + if keyErr := processKeyPoolError(err); keyErr != nil { + i.writeUpstreamError(w, keyErr) + return xerrors.Errorf("key pool exhausted: %w", err) + } + if apiErr := getErrorResponse(err); apiErr != nil { i.writeUpstreamError(w, apiErr) return xerrors.Errorf("openai API error: %w", err) @@ -258,9 +265,54 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req return nil } -func (i *BlockingInterception) newChatCompletion(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (_ *openai.ChatCompletion, outErr error) { - ctx, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) +// newChatCompletion routes between BYOK (single attempt) and +// centralized failover. +func (i *BlockingInterception) newChatCompletion(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (*openai.ChatCompletion, error) { + // BYOK: single attempt, no failover. + if i.cfg.KeyPool == nil { + return i.newChatCompletionWithKey(ctx, svc, opts) + } + return i.newChatCompletionWithKeyFailover(ctx, svc, opts) +} + +// newChatCompletionWithKey performs a single upstream call. +func (i *BlockingInterception) newChatCompletionWithKey(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (_ *openai.ChatCompletion, outErr error) { + _, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) return svc.New(ctx, i.req.ChatCompletionNewParams, opts...) } + +// newChatCompletionWithKeyFailover walks the centralized key +// pool, trying each key until one succeeds or the pool is +// exhausted. Keys are marked temporary on 429 and permanent on +// 401/403. Errors that aren't key-specific don't trigger +// failover and are returned to the caller. +func (i *BlockingInterception) newChatCompletionWithKeyFailover(ctx context.Context, svc openai.ChatCompletionService, opts []option.RequestOption) (*openai.ChatCompletion, error) { + // TODO(ssncferreira): update the interception's credential + // hint with the actually-used key (the successful key on + // success, the last tried key on failure) in the upstack PR. + walker := i.cfg.KeyPool.Walker() + for { + key, err := walker.Next() + if err != nil { + return nil, err + } + + requestOpts := append([]option.RequestOption{}, opts...) + requestOpts = append(requestOpts, + option.WithAPIKey(key.Value()), + // Disable SDK retries because the failover loop + // handles retries via key rotation. + option.WithMaxRetries(0), + ) + completion, err := i.newChatCompletionWithKey(ctx, svc, requestOpts) + // Key-specific failure: try the next key. + if i.markKeyOnError(ctx, key, err) { + continue + } + // Either success (completion, nil) or a non-key error + // (nil, err): nothing to retry, return as-is. + return completion, err + } +} diff --git a/aibridge/intercept/chatcompletions/blocking_test.go b/aibridge/intercept/chatcompletions/blocking_test.go new file mode 100644 index 0000000000..2a6fdb92c7 --- /dev/null +++ b/aibridge/intercept/chatcompletions/blocking_test.go @@ -0,0 +1,502 @@ +package chatcompletions //nolint:testpackage // tests unexported internals + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + + "github.com/google/uuid" + mcplib "github.com/mark3labs/mcp-go/mcp" + "github.com/openai/openai-go/v3" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/intercept" + "github.com/coder/coder/v2/aibridge/internal/testutil" + "github.com/coder/coder/v2/aibridge/keypool" + "github.com/coder/coder/v2/aibridge/mcp" + "github.com/coder/coder/v2/aibridge/utils" + "github.com/coder/quartz" +) + +// OpenAI-shaped response bodies. +const ( + successBody = `{"id":"chatcmpl-01","object":"chat.completion","created":1234567890,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"Hello!"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}` + toolUseBody = `{"id":"chatcmpl-01","object":"chat.completion","created":1234567890,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":null,"tool_calls":[{"id":"call_01","type":"function","function":{"name":"test_tool","arguments":"{}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}` + textCompleteBody = `{"id":"chatcmpl-02","object":"chat.completion","created":1234567890,"model":"gpt-4","choices":[{"index":0,"message":{"role":"assistant","content":"done"},"finish_reason":"stop"}],"usage":{"prompt_tokens":15,"completion_tokens":3,"total_tokens":18}}` + rateLimitBody = `{"error":{"message":"Rate limit exceeded","type":"rate_limit_error","code":"rate_limit_exceeded"}}` + authErrorBody = `{"error":{"message":"Invalid API key","type":"invalid_request_error","code":"invalid_api_key"}}` + serverErrorBody = `{"error":{"message":"Internal server error","type":"server_error","code":"internal_error"}}` +) + +type upstreamResponse struct { + statusCode int + body string + headers map[string]string +} + +// newRequestParams builds a minimal chat-completions request +// for tests. +func newRequestParams(stream bool) *ChatCompletionNewParamsWrapper { + return &ChatCompletionNewParamsWrapper{ + ChatCompletionNewParams: openai.ChatCompletionNewParams{ + Model: "gpt-4", + Messages: []openai.ChatCompletionMessageParamUnion{ + openai.UserMessage("hi"), + }, + }, + Stream: stream, + } +} + +func TestBlockingInterception_KeyFailover(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + // Centralized pool keys. Empty when byokKey is set. + keys []string + // BYOK key. Empty when keys is set. + byokKey string + // Scripted upstream responses keyed by bearer token. + responses map[string]upstreamResponse + expectedRequestCount int32 + expectedStatusCode int + expectedRetryAfter string + // Expected key states after the request, by index in keys. + expectedKeyStates []keypool.KeyState + }{ + { + // Given: 1 valid key returning 200. + // Then: 1 request, 200 response, key remains valid. + name: "single_valid_key", + keys: []string{"k0"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusOK, body: successBody}, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid}, + }, + { + // Given: 2 keys; key-0 returns 429, key-1 returns 200. + // Then: 2 requests, 200 response, key-0 temporary, key-1 valid. + name: "failover_after_429", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + "k1": {statusCode: http.StatusOK, body: successBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 returns 401, key-1 returns 200. + // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + name: "failover_after_401", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + "k1": {statusCode: http.StatusOK, body: successBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 returns 403, key-1 returns 200. + // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + name: "failover_after_403", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusForbidden, body: authErrorBody}, + "k1": {statusCode: http.StatusOK, body: successBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStateValid, + }, + }, + { + // Given: 3 keys; all return 429 with cooldowns 5s, 3s, 10s. + // Then: 3 requests, 429 response with smallest Retry-After, + // all keys temporary. + name: "all_keys_rate_limited", + keys: []string{"k0", "k1", "k2"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + "k1": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "3"}, + body: rateLimitBody, + }, + "k2": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "10"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 3, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "3", + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + }, + }, + { + // Given: 2 keys; both return 401. + // Then: 2 requests, 502 api_error response, both keys permanent. + name: "all_keys_unauthorized", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + "k1": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusBadGateway, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStatePermanent, + }, + }, + { + // Given: 2 keys; key-0 returns 500. + // Then: 1 request, 500 response, both keys remain valid. + name: "server_error_no_failover", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusInternalServerError, body: serverErrorBody}, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusInternalServerError, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, + { + // Given: BYOK with a single key returning 429. + // Then: 1 request, 429 response, no failover, upstream + // Retry-After propagated to the client. + name: "byok_no_failover", + byokKey: "user-byok", + responses: map[string]upstreamResponse{ + "user-byok": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{ + "Retry-After": "5", + // BYOK doesn't set MaxRetries(0); + // suppress SDK retries to test a + // single attempt. + "x-should-retry": "false", + }, + body: rateLimitBody, + }, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "5", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Mock upstream: counts requests and returns + // scripted responses keyed by bearer token. An + // unmapped key falls through to 500 so misconfigured + // cases surface via the status assertion. + var requestCount atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + resp, ok := tc.responses[utils.ExtractBearerToken(r.Header.Get("Authorization"))] + if !ok { + resp = upstreamResponse{statusCode: http.StatusInternalServerError} + } + w.Header().Set("Content-Type", "application/json") + for hk, hv := range resp.headers { + w.Header().Set(hk, hv) + } + w.WriteHeader(resp.statusCode) + _, _ = w.Write([]byte(resp.body)) + })) + t.Cleanup(upstream.Close) + + cfg := config.OpenAI{BaseURL: upstream.URL + "/"} + var pool *keypool.Pool + if len(tc.keys) > 0 { + var err error + pool, err = keypool.New(tc.keys, quartz.NewMock(t)) + require.NoError(t, err) + cfg.KeyPool = pool + } else if tc.byokKey != "" { + cfg.Key = tc.byokKey + } + + interceptor := NewBlockingInterceptor( + uuid.New(), + newRequestParams(false), + config.ProviderOpenAI, + cfg, + http.Header{}, + "Authorization", + otel.Tracer("blocking_test"), + intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""), + ) + interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + w := httptest.NewRecorder() + err := interceptor.ProcessRequest(w, req) + if tc.expectedStatusCode == http.StatusOK { + require.NoError(t, err) + } else { + require.Error(t, err) + } + + assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count") + assert.Equal(t, tc.expectedStatusCode, w.Code, "response status code") + assert.Equal(t, tc.expectedRetryAfter, w.Header().Get("Retry-After"), "Retry-After header") + if pool != nil { + assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") + } + }) + } +} + +// TestBlockingInterception_AgenticLoopFailover covers the +// scenarios that span an agentic-loop continuation: the initial +// client request and the subsequent tool-call continuation can +// each fail over independently. Each iteration gets its own +// walker. +func TestBlockingInterception_AgenticLoopFailover(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + // Scripted upstream responses consumed in order of + // upstream request. + responses []upstreamResponse + expectedRequestCount int32 + expectedSeenKeys []string + expectedStatusCode int + expectedKeyStates []keypool.KeyState + }{ + { + // Given: 2 keys; both upstream calls succeed on key-0. + // Then: 2 requests, 200 response, both keys remain valid. + name: "happy_path", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, body: toolUseBody}, + {statusCode: http.StatusOK, body: textCompleteBody}, + }, + expectedRequestCount: 2, + expectedSeenKeys: []string{"k0", "k0"}, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 succeeds initially, then 429s + // during the agentic continuation, key-1 succeeds. + // Then: 3 requests, 200 response, key-0 temporary, + // key-1 valid. + name: "agentic_failover_to_k1", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, body: toolUseBody}, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + {statusCode: http.StatusOK, body: textCompleteBody}, + }, + expectedRequestCount: 3, + expectedSeenKeys: []string{"k0", "k0", "k1"}, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 succeeds initially, then both + // keys 429 during the agentic continuation. + // Then: 3 requests, 429 response with smallest + // Retry-After, both keys temporary. + name: "agentic_all_keys_fail", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, body: toolUseBody}, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "3"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 3, + expectedSeenKeys: []string{"k0", "k0", "k1"}, + expectedStatusCode: http.StatusTooManyRequests, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var requestCount atomic.Int32 + var seenKeysMu sync.Mutex + var seenKeys []string + + // Mock upstream: returns scripted responses in order, + // records each request's bearer token for assertions. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + idx := int(requestCount.Add(1)) - 1 + seenKeysMu.Lock() + seenKeys = append(seenKeys, utils.ExtractBearerToken(r.Header.Get("Authorization"))) + seenKeysMu.Unlock() + _, _ = io.Copy(io.Discard, r.Body) + + if idx >= len(tc.responses) { + w.WriteHeader(http.StatusInternalServerError) + return + } + resp := tc.responses[idx] + w.Header().Set("Content-Type", "application/json") + for hk, hv := range resp.headers { + w.Header().Set(hk, hv) + } + w.WriteHeader(resp.statusCode) + _, _ = w.Write([]byte(resp.body)) + })) + t.Cleanup(upstream.Close) + + pool, err := keypool.New([]string{"k0", "k1"}, quartz.NewMock(t)) + require.NoError(t, err) + + cfg := config.OpenAI{ + BaseURL: upstream.URL + "/", + KeyPool: pool, + } + + interceptor := NewBlockingInterceptor( + uuid.New(), + newRequestParams(false), + config.ProviderOpenAI, + cfg, + http.Header{}, + "Authorization", + otel.Tracer("blocking_test"), + intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""), + ) + + // Mock proxy with a tool the upstream's tool_use + // response will reference. + proxy := &mockServerProxier{ + tools: []*mcp.Tool{ + { + Client: stubToolCaller{}, + ID: "test_tool", + Name: "test_tool", + ServerName: "coder", + Logger: slog.Make(), + }, + }, + } + interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, proxy) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + w := httptest.NewRecorder() + err = interceptor.ProcessRequest(w, req) + if tc.expectedStatusCode == http.StatusOK { + require.NoError(t, err) + } else { + require.Error(t, err) + } + + assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count") + assert.Equal(t, tc.expectedStatusCode, w.Code, "response status code") + + seenKeysMu.Lock() + defer seenKeysMu.Unlock() + assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys") + assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") + }) + } +} + +// mockServerProxier is a test implementation of mcp.ServerProxier. +type mockServerProxier struct { + tools []*mcp.Tool +} + +func (*mockServerProxier) Init(context.Context) error { + return nil +} + +func (*mockServerProxier) Shutdown(context.Context) error { + return nil +} + +func (m *mockServerProxier) ListTools() []*mcp.Tool { + return m.tools +} + +func (m *mockServerProxier) GetTool(id string) *mcp.Tool { + for _, t := range m.tools { + if t.ID == id { + return t + } + } + return nil +} + +func (*mockServerProxier) CallTool(context.Context, string, any) (*mcplib.CallToolResult, error) { + return nil, nil //nolint:nilnil // mock: no-op implementation +} + +// stubToolCaller is a minimal mcp.ToolCaller that returns a fixed +// text result, so the agentic continuation can proceed. +type stubToolCaller struct{} + +func (stubToolCaller) CallTool(_ context.Context, _ mcplib.CallToolRequest) (*mcplib.CallToolResult, error) { + return mcplib.NewToolResultText("tool result"), nil +} diff --git a/aibridge/intercept/chatcompletions/streaming.go b/aibridge/intercept/chatcompletions/streaming.go index 8dac47dddf..b84898c264 100644 --- a/aibridge/intercept/chatcompletions/streaming.go +++ b/aibridge/intercept/chatcompletions/streaming.go @@ -24,6 +24,7 @@ import ( aibcontext "github.com/coder/coder/v2/aibridge/context" "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/intercept/eventstream" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/recorder" "github.com/coder/coder/v2/aibridge/tracing" @@ -126,9 +127,49 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re lastErr error interceptionErr error ) + for { // TODO add outer loop span (https://github.com/coder/aibridge/issues/67) + + // Per-iteration walker. An iteration is either an agentic + // continuation (sending a tool result back in a new + // stream) or a failover retry (previous key marked, try + // the next one). + var walker *keypool.Walker + if i.cfg.KeyPool != nil { + walker = i.cfg.KeyPool.Walker() + } + var opts []option.RequestOption + var currentKey *keypool.Key + if walker != nil { + key, err := walker.Next() + if respErr := processKeyPoolError(err); respErr != nil { + // Pool exhausted in this iteration. Relay the + // error to the client: as an SSE event if events + // have already been sent, or by direct write + // otherwise. + interceptionErr = respErr + if events.IsStreaming() { + payload, mErr := i.marshalErr(respErr) + if mErr != nil { + logger.Warn(ctx, "failed to marshal exhaustion error", slog.Error(mErr)) + } else if sErr := events.Send(streamCtx, payload); sErr != nil { + logger.Warn(ctx, "failed to relay exhaustion error", slog.Error(sErr)) + } + } else { + i.writeUpstreamError(w, respErr) + } + break + } + currentKey = key + opts = append(opts, + option.WithAPIKey(key.Value()), + // Disable SDK retries because the failover + // loop handles retries via key rotation. + option.WithMaxRetries(0), + ) + } // TODO(ssncferreira): inject actor headers directly in the client-header // middleware instead of using SDK options. @@ -151,7 +192,17 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re var toolCall *openai.FinishedChatCompletionToolCall + // iterationStarted is per-iteration (reset on every + // loop): true once the upstream call has produced any + // events for this iteration. While false, a key-specific + // failure can still fail over to the next key. Distinct + // from events.IsStreaming(), which is stream-wide and + // stays true once iteration 1 has sent any event + // downstream. + var iterationStarted bool + for stream.Next() { + iterationStarted = true chunk := stream.Current() canRelay := processor.process(chunk) @@ -231,40 +282,52 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re }) } - if !events.IsStreaming() { - // response/downstream Stream has not started yet; write error response and exit. - i.writeUpstreamError(w, getErrorResponse(stream.Err())) - return stream.Err() - } - - // Check if the stream encountered any errors. - if streamErr := stream.Err(); streamErr != nil { - if eventstream.IsUnrecoverableError(streamErr) { - logger.Debug(ctx, "stream terminated", slog.Error(streamErr)) - // We can't reflect an error back if there's a connection error or the request context was canceled. - } else if oaiErr := getErrorResponse(streamErr); oaiErr != nil { - logger.Warn(ctx, "openai stream error", slog.Error(streamErr)) - interceptionErr = oaiErr - } else { - logger.Warn(ctx, "unknown stream error", slog.Error(streamErr)) - // Unfortunately, the OpenAI SDK does not support parsing errors received in the stream - // into known types (i.e. [shared.OverloadedError]). - // See https://github.com/openai/openai-go/blob/v2.7.0/packages/ssestream/ssestream.go#L171 - // All it does is wrap the payload in an error - which is all we can return, currently. - interceptionErr = newErrorResponse(xerrors.Errorf("unknown stream error: %w", streamErr)) + if iterationStarted { + // Mid-stream error or logical error: events have + // already streamed for this iteration, so the + // error is relayed as an SSE event. + streamErr := stream.Err() + if respErr := i.mapStreamError(ctx, logger, streamErr, lastErr); respErr != nil { + interceptionErr = respErr + payload, err := i.marshalErr(respErr) + if err != nil { + logger.Warn(ctx, "failed to marshal error", slog.Error(err), slog.F("error_payload", fmt.Sprintf("%+v", respErr))) + } else if err := events.Send(streamCtx, payload); err != nil { + logger.Warn(ctx, "failed to relay error", slog.Error(err), slog.F("payload", payload)) + } + } else if streamErr != nil { + // Unrecoverable (e.g., broken pipe, context + // canceled): can't relay to the client, but record + // the error so it isn't silently swallowed. + interceptionErr = streamErr } - } else if lastErr != nil { - // Otherwise check if any logical errors occurred during processing. - logger.Warn(ctx, "stream processing failed", slog.Error(lastErr)) - interceptionErr = newErrorResponse(xerrors.Errorf("processing error: %w", lastErr)) - } - - if interceptionErr != nil { - payload, err := i.marshalErr(interceptionErr) - if err != nil { - logger.Warn(ctx, "failed to marshal error", slog.Error(err), slog.F("error_payload", fmt.Sprintf("%+v", interceptionErr))) - } else if err := events.Send(streamCtx, payload); err != nil { - logger.Warn(ctx, "failed to relay error", slog.Error(err), slog.F("payload", payload)) + } else { + // Pre-stream failure of this iteration. For + // centralized requests, mark the key and retry with + // the next one. + if currentKey != nil && i.markKeyOnError(ctx, currentKey, stream.Err()) { + continue + } + // Non-key error: relay it. Use mapStreamError so that + // unknown upstream errors (TCP reset, DNS failure, TLS + // error, deadline exceeded) are wrapped in a generic + // response instead of producing a silent HTTP 200. + respErr := i.mapStreamError(ctx, logger, stream.Err(), lastErr) + if respErr != nil { + interceptionErr = respErr + if events.IsStreaming() { + // Prior iterations have streamed, so the SSE + // connection is open: inject as an SSE event. + payload, mErr := i.marshalErr(respErr) + if mErr != nil { + logger.Warn(ctx, "failed to marshal error", slog.Error(mErr)) + } else if sErr := events.Send(streamCtx, payload); sErr != nil { + logger.Warn(ctx, "failed to relay error", slog.Error(sErr)) + } + } else { + // No events streamed yet, write the response directly. + i.writeUpstreamError(w, respErr) + } } } @@ -407,6 +470,35 @@ func (i *StreamingInterception) newStream(ctx context.Context, svc openai.ChatCo return svc.NewStreaming(ctx, openai.ChatCompletionNewParams{}, opts...) } +// mapStreamError converts a mid-stream upstream error or +// processing error into a relayable responseError. Returns nil +// when the error is unrecoverable, in which case nothing can be +// relayed back. +func (*StreamingInterception) mapStreamError(ctx context.Context, logger slog.Logger, streamErr, lastErr error) *responseError { + if streamErr != nil { + if eventstream.IsUnrecoverableError(streamErr) { + logger.Debug(ctx, "stream terminated", slog.Error(streamErr)) + // We can't reflect an error back if there's a connection error or the request context was canceled. + return nil + } + if oaiErr := getErrorResponse(streamErr); oaiErr != nil { + logger.Warn(ctx, "openai stream error", slog.Error(streamErr)) + return oaiErr + } + logger.Warn(ctx, "unknown stream error", slog.Error(streamErr)) + // Unfortunately, the OpenAI SDK does not support parsing errors received in the stream + // into known types (i.e. [shared.OverloadedError]). + // See https://github.com/openai/openai-go/blob/v2.7.0/packages/ssestream/ssestream.go#L171 + // All it does is wrap the payload in an error - which is all we can return, currently. + return newErrorResponse(fmt.Sprintf("unknown stream error: %s", streamErr), intercept.OpenAIErrTypeError, intercept.OpenAIErrTypeError, http.StatusBadGateway, 0) + } + if lastErr != nil { + logger.Warn(ctx, "stream processing failed", slog.Error(lastErr)) + return newErrorResponse(fmt.Sprintf("processing error: %s", lastErr), intercept.OpenAIErrTypeError, intercept.OpenAIErrTypeError, http.StatusBadGateway, 0) + } + return nil +} + type streamProcessor struct { ctx context.Context logger slog.Logger diff --git a/aibridge/intercept/chatcompletions/streaming_test.go b/aibridge/intercept/chatcompletions/streaming_test.go index 640ad197c5..162624375b 100644 --- a/aibridge/intercept/chatcompletions/streaming_test.go +++ b/aibridge/intercept/chatcompletions/streaming_test.go @@ -1,9 +1,12 @@ -package chatcompletions_test +package chatcompletions //nolint:testpackage // tests unexported internals import ( + "io" "net/http" "net/http/httptest" - "strconv" + "strings" + "sync" + "sync/atomic" "testing" "github.com/google/uuid" @@ -16,8 +19,11 @@ import ( "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/intercept" - "github.com/coder/coder/v2/aibridge/intercept/chatcompletions" "github.com/coder/coder/v2/aibridge/internal/testutil" + "github.com/coder/coder/v2/aibridge/keypool" + "github.com/coder/coder/v2/aibridge/mcp" + "github.com/coder/coder/v2/aibridge/utils" + "github.com/coder/quartz" ) // Test that when the upstream provider returns an error before streaming starts, @@ -36,21 +42,21 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) { name: "bad request error", statusCode: http.StatusBadRequest, responseBody: `{"error":{"message":"Invalid request","type":"invalid_request_error","code":"invalid_request"}}`, - expectedErrStr: strconv.Itoa(http.StatusBadRequest), + expectedErrStr: "Invalid request", expectedBody: "invalid_request", }, { name: "rate limit error", statusCode: http.StatusTooManyRequests, responseBody: `{"error":{"message":"Rate limit exceeded","type":"rate_limit_error","code":"rate_limit_exceeded"}}`, - expectedErrStr: strconv.Itoa(http.StatusTooManyRequests), + expectedErrStr: "Rate limit exceeded", expectedBody: "rate_limit", }, { name: "internal server error", statusCode: http.StatusInternalServerError, responseBody: `{"error":{"message":"Internal server error","type":"server_error","code":"internal_error"}}`, - expectedErrStr: strconv.Itoa(http.StatusInternalServerError), + expectedErrStr: "Internal server error", expectedBody: "server_error", }, } @@ -74,7 +80,7 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) { Key: "test-key", } - req := &chatcompletions.ChatCompletionNewParamsWrapper{ + req := &ChatCompletionNewParamsWrapper{ ChatCompletionNewParams: openai.ChatCompletionNewParams{ Model: "gpt-4", Messages: []openai.ChatCompletionMessageParamUnion{ @@ -89,7 +95,7 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) { httpReq := httptest.NewRequest(http.MethodPost, "/chat/completions", nil) tracer := otel.Tracer("test") - interceptor := chatcompletions.NewStreamingInterceptor(uuid.New(), req, config.ProviderOpenAI, cfg, httpReq.Header, "Authorization", tracer, intercept.CredentialInfo{}) + interceptor := NewStreamingInterceptor(uuid.New(), req, config.ProviderOpenAI, cfg, httpReq.Header, "Authorization", tracer, intercept.CredentialInfo{}) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug) interceptor.Setup(logger, &testutil.MockRecorder{}, nil) @@ -110,3 +116,486 @@ func TestStreamingInterception_RelaysUpstreamErrorToClient(t *testing.T) { }) } } + +// OpenAI-shaped SSE body for a successful streaming response. +const streamingSuccessBody = `data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]} + +data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]} + +data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}} + +data: [DONE] + +` + +func TestStreamingInterception_KeyFailover(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + // Centralized pool keys. Empty when byokKey is set. + keys []string + // BYOK key. Empty when keys is set. + byokKey string + // Scripted upstream responses keyed by bearer token. + responses map[string]upstreamResponse + expectedRequestCount int32 + expectedStatusCode int + expectedRetryAfter string + // Expected key states after the request, by index in keys. + expectedKeyStates []keypool.KeyState + }{ + { + // Given: 1 valid key returning a successful stream. + // Then: 1 request, 200 response, key remains valid. + name: "single_valid_key", + keys: []string{"k0"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusOK, + headers: map[string]string{"Content-Type": "text/event-stream"}, + body: streamingSuccessBody, + }, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid}, + }, + { + // Given: 2 keys; key-0 returns 429 pre-stream, key-1 + // streams successfully. + // Then: 2 requests, 200 response, key-0 temporary, key-1 valid. + name: "failover_after_429", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + "k1": { + statusCode: http.StatusOK, + headers: map[string]string{"Content-Type": "text/event-stream"}, + body: streamingSuccessBody, + }, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 returns 401 pre-stream, key-1 + // streams successfully. + // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + name: "failover_after_401", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + "k1": { + statusCode: http.StatusOK, + headers: map[string]string{"Content-Type": "text/event-stream"}, + body: streamingSuccessBody, + }, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 returns 403 pre-stream, key-1 streams. + // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + name: "failover_after_403", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusForbidden, body: authErrorBody}, + "k1": { + statusCode: http.StatusOK, + headers: map[string]string{"Content-Type": "text/event-stream"}, + body: streamingSuccessBody, + }, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStateValid, + }, + }, + { + // Given: 3 keys; all return 429 pre-stream with + // cooldowns 5s, 3s, 10s. + // Then: 3 requests, 429 response with smallest + // Retry-After, all keys temporary. + name: "all_keys_rate_limited", + keys: []string{"k0", "k1", "k2"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + "k1": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "3"}, + body: rateLimitBody, + }, + "k2": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "10"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 3, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "3", + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + }, + }, + { + // Given: 2 keys; both return 401 pre-stream. + // Then: 2 requests, 502 api_error response, both keys permanent. + name: "all_keys_unauthorized", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + "k1": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusBadGateway, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStatePermanent, + }, + }, + { + // Given: 2 keys; key-0 returns 500 pre-stream. + // Then: 1 request, 500 response, both keys remain valid. + name: "server_error_no_failover", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusInternalServerError, body: serverErrorBody}, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusInternalServerError, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, + { + // Given: BYOK with a single key returning 429. + // Then: 1 request, 429 response, no failover, upstream + // Retry-After propagated to the client. + name: "byok_no_failover", + byokKey: "user-byok", + responses: map[string]upstreamResponse{ + "user-byok": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{ + "Retry-After": "5", + // BYOK doesn't set MaxRetries(0); + // suppress SDK retries to test a + // single attempt. + "x-should-retry": "false", + }, + body: rateLimitBody, + }, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "5", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Mock upstream: counts requests and returns + // scripted responses keyed by bearer token. An + // unmapped key falls through to 500 so misconfigured + // cases surface via the status assertion. + var requestCount atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + resp, ok := tc.responses[utils.ExtractBearerToken(r.Header.Get("Authorization"))] + if !ok { + resp = upstreamResponse{statusCode: http.StatusInternalServerError} + } + for hk, hv := range resp.headers { + w.Header().Set(hk, hv) + } + w.WriteHeader(resp.statusCode) + _, _ = w.Write([]byte(resp.body)) + })) + t.Cleanup(upstream.Close) + + cfg := config.OpenAI{BaseURL: upstream.URL + "/"} + var pool *keypool.Pool + if len(tc.keys) > 0 { + var err error + pool, err = keypool.New(tc.keys, quartz.NewMock(t)) + require.NoError(t, err) + cfg.KeyPool = pool + } else if tc.byokKey != "" { + cfg.Key = tc.byokKey + } + + interceptor := NewStreamingInterceptor( + uuid.New(), + newRequestParams(true), + config.ProviderOpenAI, + cfg, + http.Header{}, + "Authorization", + otel.Tracer("streaming_test"), + intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""), + ) + interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + w := httptest.NewRecorder() + err := interceptor.ProcessRequest(w, req) + if tc.expectedStatusCode == http.StatusOK { + require.NoError(t, err) + } else { + require.Error(t, err) + } + + assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count") + assert.Equal(t, tc.expectedStatusCode, w.Code, "response status code") + assert.Equal(t, tc.expectedRetryAfter, w.Header().Get("Retry-After"), "Retry-After header") + if pool != nil { + assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") + } + }) + } +} + +// SSE bodies covering an agentic-continuation flow. +const ( + // First response: a tool_calls delta referencing the + // injected "test_tool". Triggers the agentic continuation + // loop. + toolUseStreamBody = `data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_01","type":"function","function":{"name":"test_tool","arguments":""}}]},"finish_reason":null}]} + +data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]},"finish_reason":null}]} + +data: {"id":"chatcmpl-01","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}} + +data: [DONE] + +` + + // Second response (after the tool result is sent back): + // a plain text completion that ends the loop. + textStreamBody = `data: {"id":"chatcmpl-02","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","content":"done"},"finish_reason":null}]} + +data: {"id":"chatcmpl-02","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":15,"completion_tokens":3,"total_tokens":18}} + +data: [DONE] + +` +) + +// TestStreamingInterception_AgenticLoopFailover covers the +// scenarios that span an agentic-loop continuation: the initial +// client request and the subsequent tool-call continuation can +// each fail over independently. Each iteration gets its own +// walker. +func TestStreamingInterception_AgenticLoopFailover(t *testing.T) { + t.Parallel() + + sseHeaders := map[string]string{"Content-Type": "text/event-stream"} + + tests := []struct { + name string + // Scripted upstream responses consumed in order of + // upstream request. + responses []upstreamResponse + expectedRequestCount int32 + expectedSeenKeys []string + // Substring expected in the response body. Either a + // success marker (e.g. "done") or an error marker + // (e.g. "rate_limit_error"). + expectedBodyContains string + // True when the error must be relayed as an SSE event. + expectErrorAsSSEEvent bool + // True when ProcessRequest is expected to return an + // error (e.g. all keys exhausted). + expectedErr bool + expectedKeyStates []keypool.KeyState + }{ + { + // Given: 2 keys; both upstream calls succeed on key-0. + // Then: 2 requests, success body, both keys remain valid. + name: "happy_path", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody}, + {statusCode: http.StatusOK, headers: sseHeaders, body: textStreamBody}, + }, + expectedRequestCount: 2, + expectedSeenKeys: []string{"k0", "k0"}, + expectedBodyContains: "done", + expectErrorAsSSEEvent: false, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 succeeds initially, then 429s + // during the agentic continuation, key-1 succeeds. + // Then: 3 requests, success body, key-0 temporary, + // key-1 valid. + name: "agentic_failover_to_k1", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody}, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + {statusCode: http.StatusOK, headers: sseHeaders, body: textStreamBody}, + }, + expectedRequestCount: 3, + expectedSeenKeys: []string{"k0", "k0", "k1"}, + expectedBodyContains: "done", + expectErrorAsSSEEvent: false, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 succeeds initially, then both + // keys 429 during the agentic continuation. + // Then: 3 requests, error injected as SSE event, both + // keys temporary. + name: "agentic_all_keys_fail", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody}, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "3"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 3, + expectedSeenKeys: []string{"k0", "k0", "k1"}, + expectedBodyContains: "all configured keys are rate-limited", + expectErrorAsSSEEvent: true, + expectedErr: true, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var requestCount atomic.Int32 + var seenKeysMu sync.Mutex + var seenKeys []string + + // Mock upstream: returns scripted responses in order, + // records each request's bearer token for assertions. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + idx := int(requestCount.Add(1)) - 1 + seenKeysMu.Lock() + seenKeys = append(seenKeys, utils.ExtractBearerToken(r.Header.Get("Authorization"))) + seenKeysMu.Unlock() + _, _ = io.Copy(io.Discard, r.Body) + + if idx >= len(tc.responses) { + w.WriteHeader(http.StatusInternalServerError) + return + } + resp := tc.responses[idx] + for hk, hv := range resp.headers { + w.Header().Set(hk, hv) + } + w.WriteHeader(resp.statusCode) + _, _ = w.Write([]byte(resp.body)) + })) + t.Cleanup(upstream.Close) + + pool, err := keypool.New([]string{"k0", "k1"}, quartz.NewMock(t)) + require.NoError(t, err) + + cfg := config.OpenAI{ + BaseURL: upstream.URL + "/", + KeyPool: pool, + } + + interceptor := NewStreamingInterceptor( + uuid.New(), + newRequestParams(true), + config.ProviderOpenAI, + cfg, + http.Header{}, + "Authorization", + otel.Tracer("streaming_test"), + intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""), + ) + + // Mock proxy with a tool the upstream's tool_calls + // chunks will reference. The stub caller returns a + // fixed text result. + proxy := &mockServerProxier{ + tools: []*mcp.Tool{ + { + Client: stubToolCaller{}, + ID: "test_tool", + Name: "test_tool", + ServerName: "coder", + Logger: slog.Make(), + }, + }, + } + interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, proxy) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + w := httptest.NewRecorder() + err = interceptor.ProcessRequest(w, req) + if tc.expectedErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + + assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count") + body := w.Body.String() + assert.Contains(t, body, tc.expectedBodyContains, "response body") + if tc.expectErrorAsSSEEvent { + // SSE was opened before the failure, so the body + // must start with stream chunks, not a direct + // HTTP error body. + assert.True(t, strings.HasPrefix(body, "data: "), "body must start with SSE chunks") + } + + seenKeysMu.Lock() + defer seenKeysMu.Unlock() + assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys") + assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") + }) + } +} diff --git a/aibridge/intercept/openai_errors.go b/aibridge/intercept/openai_errors.go new file mode 100644 index 0000000000..92e13fd02f --- /dev/null +++ b/aibridge/intercept/openai_errors.go @@ -0,0 +1,14 @@ +package intercept + +// OpenAI error type and code constants used by the chatcompletions +// and responses interceptors. The OpenAI Go SDK does not expose +// these as typed constants, so we define our own. +// See https://platform.openai.com/docs/guides/error-codes. +const ( + OpenAIErrTypeError = "error" + OpenAIErrTypeAPI = "api_error" + OpenAIErrTypeRateLimit = "rate_limit_error" + + OpenAIErrCodeServer = "server_error" + OpenAIErrCodeRateLimit = "rate_limit_exceeded" +) diff --git a/aibridge/intercept/responses/base.go b/aibridge/intercept/responses/base.go index 9affc7d3ea..059162c278 100644 --- a/aibridge/intercept/responses/base.go +++ b/aibridge/intercept/responses/base.go @@ -4,7 +4,10 @@ import ( "bytes" "context" "encoding/json" + "errors" + "fmt" "io" + "math" "net/http" "strconv" "strings" @@ -13,8 +16,10 @@ import ( "time" "github.com/google/uuid" + "github.com/openai/openai-go/v3" "github.com/openai/openai-go/v3/option" "github.com/openai/openai-go/v3/responses" + "github.com/openai/openai-go/v3/shared" "github.com/openai/openai-go/v3/shared/constant" "github.com/tidwall/gjson" "go.opentelemetry.io/otel/attribute" @@ -26,6 +31,7 @@ import ( aibcontext "github.com/coder/coder/v2/aibridge/context" "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/intercept/apidump" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/recorder" "github.com/coder/coder/v2/aibridge/tracing" @@ -53,8 +59,19 @@ type responsesInterceptionBase struct { credential intercept.CredentialInfo } +// newResponsesService builds the SDK service used for upstream +// calls. BYOK auth is set here. Centralized auth is set +// per-attempt by the failover loop. func (i *responsesInterceptionBase) newResponsesService() responses.ResponseService { - opts := []option.RequestOption{option.WithBaseURL(i.cfg.BaseURL), option.WithAPIKey(i.cfg.Key)} + // TODO(ssncferreira): validate auth is configured per + // https://github.com/coder/aibridge/issues/266. + + var opts []option.RequestOption + // BYOK auth. + if i.cfg.KeyPool == nil { + opts = append(opts, option.WithAPIKey(i.cfg.Key)) + } + opts = append(opts, option.WithBaseURL(i.cfg.BaseURL)) // Add extra headers if configured. // Some providers require additional headers that are not added by the SDK. @@ -124,6 +141,107 @@ func (i *responsesInterceptionBase) validateRequest(ctx context.Context, w http. return nil } +// writeUpstreamError marshals and writes a given error. +func (i *responsesInterceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *responseError) { + if oaiErr == nil { + return + } + + w.Header().Set("Content-Type", "application/json") + // Set Retry-After when a cooldown is configured. + if oaiErr.RetryAfter > 0 { + w.Header().Set("Retry-After", strconv.Itoa(int(math.Ceil(oaiErr.RetryAfter.Seconds())))) + } + w.WriteHeader(oaiErr.StatusCode) + + out, err := json.Marshal(oaiErr) + if err != nil { + i.logger.Warn(context.Background(), "failed to marshal upstream error", slog.Error(err), slog.F("error_payload", fmt.Sprintf("%+v", oaiErr))) + // Response has to match expected format. + _, _ = w.Write([]byte(`{ + "error": { + "type": "error", + "message":"error marshaling upstream error", + "code": "server_error" + } +}`)) + } else { + _, _ = w.Write(out) + } +} + +// For centralized requests, markKeyOnError extracts an OpenAI +// SDK error from err and marks the key based on its status +// code. Returns true if the status was a key-specific failover +// trigger so callers can retry with the next key. +func (i *responsesInterceptionBase) markKeyOnError(ctx context.Context, key *keypool.Key, err error) bool { + if i.cfg.KeyPool == nil { + return false + } + var apiErr *openai.Error + if !errors.As(err, &apiErr) { + return false + } + return keypool.MarkKeyOnStatus( + ctx, key, apiErr.Response, + i.logger, i.providerName, + ) +} + +// processKeyPoolError translates a keypool exhaustion error +// into a developer-facing responseError shaped for the OpenAI +// API. Returns nil if err is not an exhaustion error. +func processKeyPoolError(err error) *responseError { + var transient *keypool.TransientKeyPoolError + switch { + case errors.As(err, &transient): + return newErrorResponse( + "all configured keys are rate-limited", + intercept.OpenAIErrTypeRateLimit, + intercept.OpenAIErrCodeRateLimit, + http.StatusTooManyRequests, + transient.RetryAfter, + ) + case errors.Is(err, keypool.ErrPermanentKeyPool): + return newErrorResponse( + "all configured keys failed authentication", + intercept.OpenAIErrTypeAPI, + intercept.OpenAIErrCodeServer, + http.StatusBadGateway, + 0, + ) + default: + return nil + } +} + +func newErrorResponse(msg, errType, code string, status int, retryAfter time.Duration) *responseError { + return &responseError{ + ErrorObject: &shared.ErrorObject{ + Code: code, + Message: msg, + Type: errType, + }, + StatusCode: status, + RetryAfter: retryAfter, + } +} + +var _ error = &responseError{} + +type responseError struct { + ErrorObject *shared.ErrorObject `json:"error"` + StatusCode int `json:"-"` + RetryAfter time.Duration `json:"-"` +} + +func (a *responseError) Error() string { + if a.ErrorObject == nil { + return "" + } + return a.ErrorObject.Message +} + // sendCustomErr sends custom responses.Error error to the client // it should only be called before any data is sent back to the client func (i *responsesInterceptionBase) sendCustomErr(ctx context.Context, w http.ResponseWriter, code int, err error) { diff --git a/aibridge/intercept/responses/base_test.go b/aibridge/intercept/responses/base_test.go index bf1fa198c8..4651b48f2f 100644 --- a/aibridge/intercept/responses/base_test.go +++ b/aibridge/intercept/responses/base_test.go @@ -1,17 +1,25 @@ package responses //nolint:testpackage // tests unexported internals import ( + "context" "net/http" + "net/http/httptest" "testing" "time" "github.com/google/uuid" + "github.com/openai/openai-go/v3" oairesponses "github.com/openai/openai-go/v3/responses" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/xerrors" "cdr.dev/slog/v3" + "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/internal/testutil" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/recorder" + "github.com/coder/quartz" ) func TestRecordPrompt(t *testing.T) { @@ -382,3 +390,193 @@ func TestResponseCopierDoesntSendIfNoResponseReceived(t *testing.T) { require.True(t, mrw.writeCalled) require.True(t, mrw.writeHeaderCalled) } + +func TestProcessKeyPoolError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + expectedNil bool + expectedStatus int + expectedRetryAfter time.Duration + }{ + { + // Transient with valid keys present: 429, no Retry-After. + name: "transient_zero_retry_after", + err: &keypool.TransientKeyPoolError{}, + expectedStatus: http.StatusTooManyRequests, + expectedRetryAfter: 0, + }, + { + // Transient with cooldown: 429, Retry-After set. + name: "transient_with_retry_after", + err: &keypool.TransientKeyPoolError{RetryAfter: 5 * time.Second}, + expectedStatus: http.StatusTooManyRequests, + expectedRetryAfter: 5 * time.Second, + }, + { + // Permanent: 502 api_error. + name: "permanent_returns_502", + err: keypool.ErrPermanentKeyPool, + expectedStatus: http.StatusBadGateway, + }, + { + // Anything else: not a pool-exhaustion error. + name: "non_pool_exhaustion_error_returns_nil", + err: xerrors.New("some other error"), + expectedNil: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got := processKeyPoolError(tc.err) + if tc.expectedNil { + require.Nil(t, got) + return + } + require.NotNil(t, got) + assert.Equal(t, tc.expectedStatus, got.StatusCode) + assert.Equal(t, tc.expectedRetryAfter, got.RetryAfter) + }) + } +} + +func TestMarkKeyOnError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + err error + expectedReturn bool + expectedState keypool.KeyState + }{ + { + // Not an *openai.Error: no status code to act on. + name: "non_api_error_returns_false", + err: xerrors.New("network failure"), + expectedReturn: false, + expectedState: keypool.KeyStateValid, + }, + { + // Rate-limited: temporary cooldown. + name: "429_marks_temporary", + err: &openai.Error{StatusCode: http.StatusTooManyRequests, Response: &http.Response{StatusCode: http.StatusTooManyRequests}}, + expectedReturn: true, + expectedState: keypool.KeyStateTemporary, + }, + { + // Auth failure: mark permanent. + name: "401_marks_permanent", + err: &openai.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}}, + expectedReturn: true, + expectedState: keypool.KeyStatePermanent, + }, + { + // Auth forbidden: mark permanent. + name: "403_marks_permanent", + err: &openai.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}}, + expectedReturn: true, + expectedState: keypool.KeyStatePermanent, + }, + { + // Server errors are not key-specific. + name: "500_does_not_mark", + err: &openai.Error{StatusCode: http.StatusInternalServerError, Response: &http.Response{StatusCode: http.StatusInternalServerError}}, + expectedReturn: false, + expectedState: keypool.KeyStateValid, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + pool, err := keypool.New([]string{"key-0"}, quartz.NewMock(t)) + require.NoError(t, err) + key, err := pool.Walker().Next() + require.NoError(t, err) + + base := &responsesInterceptionBase{cfg: config.OpenAI{KeyPool: pool}, logger: slog.Make()} + + got := base.markKeyOnError(context.Background(), key, tc.err) + assert.Equal(t, tc.expectedReturn, got) + assert.Equal(t, tc.expectedState, key.State()) + }) + } +} + +func TestWriteUpstreamError(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + respErr *responseError + expectStatus int + // Empty string means the header should be absent. + expectRetryAfter string + // Substring expected in the marshaled body. Empty means no body check. + expectBodyContains string + }{ + { + // Standard error: status, code, and JSON body written. + name: "writes_status_and_body", + respErr: newErrorResponse("upstream failed", "api_error", "server_error", http.StatusBadGateway, 0), + expectStatus: http.StatusBadGateway, + expectBodyContains: `"upstream failed"`, + }, + { + // OpenAI envelope: the code field round-trips into the body. + name: "writes_code_field", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 0), + expectStatus: http.StatusTooManyRequests, + expectBodyContains: `"rate_limit_exceeded"`, + }, + { + // Whole-second retryAfter: emitted as integer seconds. + name: "retry_after_in_seconds", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 60*time.Second), + expectStatus: http.StatusTooManyRequests, + expectRetryAfter: "60", + }, + { + // 500ms rounds up to Retry-After: 1. + name: "retry_after_500ms_rounds_up_to_one", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 500*time.Millisecond), + expectStatus: http.StatusTooManyRequests, + expectRetryAfter: "1", + }, + { + // 200ms rounds up to Retry-After: 1. + name: "retry_after_200ms_rounds_up_to_one", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, 200*time.Millisecond), + expectStatus: http.StatusTooManyRequests, + expectRetryAfter: "1", + }, + { + // Negative retryAfter: header omitted. + name: "negative_retry_after_omits_header", + respErr: newErrorResponse("rate limited", "rate_limit_error", "rate_limit_exceeded", http.StatusTooManyRequests, -1*time.Second), + expectStatus: http.StatusTooManyRequests, + expectRetryAfter: "", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + base := &responsesInterceptionBase{logger: slog.Make()} + + w := httptest.NewRecorder() + base.writeUpstreamError(w, tc.respErr) + + assert.Equal(t, tc.expectStatus, w.Code, "status code") + assert.Equal(t, "application/json", w.Header().Get("Content-Type"), "Content-Type header") + assert.Equal(t, tc.expectRetryAfter, w.Header().Get("Retry-After"), "Retry-After header") + if tc.expectBodyContains != "" { + assert.Contains(t, w.Body.String(), tc.expectBodyContains, "response body") + } + }) + } +} diff --git a/aibridge/intercept/responses/blocking.go b/aibridge/intercept/responses/blocking.go index ce98219fc3..cec59307ea 100644 --- a/aibridge/intercept/responses/blocking.go +++ b/aibridge/intercept/responses/blocking.go @@ -100,6 +100,15 @@ func (i *BlockingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r * response, upstreamErr = i.newResponse(ctx, srv, opts) + // The failover loop may return a keypool exhaustion + // error. Render it here. + if upstreamErr != nil { + if keyErr := processKeyPoolError(upstreamErr); keyErr != nil { + i.writeUpstreamError(w, keyErr) + return xerrors.Errorf("key pool exhausted: %w", upstreamErr) + } + } + if upstreamErr != nil || response == nil { break } @@ -135,10 +144,55 @@ func (i *BlockingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r * return errors.Join(upstreamErr, err) } -func (i *BlockingResponsesInterceptor) newResponse(ctx context.Context, srv responses.ResponseService, opts []option.RequestOption) (_ *responses.Response, outErr error) { - ctx, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) +// newResponse routes between BYOK (single attempt) and +// centralized failover. +func (i *BlockingResponsesInterceptor) newResponse(ctx context.Context, srv responses.ResponseService, opts []option.RequestOption) (*responses.Response, error) { + // BYOK: single attempt, no failover. + if i.cfg.KeyPool == nil { + return i.newResponseWithKey(ctx, srv, opts) + } + return i.newResponseWithKeyFailover(ctx, srv, opts) +} + +// newResponseWithKey performs a single upstream call. +func (i *BlockingResponsesInterceptor) newResponseWithKey(ctx context.Context, srv responses.ResponseService, opts []option.RequestOption) (_ *responses.Response, outErr error) { + _, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) // The body is overridden by option.WithRequestBody(reqPayload) in requestOptions return srv.New(ctx, responses.ResponseNewParams{}, opts...) } + +// newResponseWithKeyFailover walks the centralized key pool, +// trying each key until one succeeds or the pool is exhausted. +// Keys are marked temporary on 429 and permanent on 401/403. +// Errors that aren't key-specific don't trigger failover and +// are returned to the caller. +func (i *BlockingResponsesInterceptor) newResponseWithKeyFailover(ctx context.Context, srv responses.ResponseService, opts []option.RequestOption) (*responses.Response, error) { + // TODO(ssncferreira): update the interception's credential + // hint with the actually-used key (the successful key on + // success, the last tried key on failure) in the upstack PR. + walker := i.cfg.KeyPool.Walker() + for { + key, err := walker.Next() + if err != nil { + return nil, err + } + + requestOpts := append([]option.RequestOption{}, opts...) + requestOpts = append(requestOpts, + option.WithAPIKey(key.Value()), + // Disable SDK retries because the failover loop + // handles retries via key rotation. + option.WithMaxRetries(0), + ) + response, err := i.newResponseWithKey(ctx, srv, requestOpts) + // Key-specific failure: try the next key. + if i.markKeyOnError(ctx, key, err) { + continue + } + // Either success (response, nil) or a non-key error + // (nil, err): nothing to retry, return as-is. + return response, err + } +} diff --git a/aibridge/intercept/responses/blocking_test.go b/aibridge/intercept/responses/blocking_test.go new file mode 100644 index 0000000000..8e1e67e0b6 --- /dev/null +++ b/aibridge/intercept/responses/blocking_test.go @@ -0,0 +1,492 @@ +package responses //nolint:testpackage // tests unexported internals + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + + "github.com/google/uuid" + mcplib "github.com/mark3labs/mcp-go/mcp" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/intercept" + "github.com/coder/coder/v2/aibridge/internal/testutil" + "github.com/coder/coder/v2/aibridge/keypool" + "github.com/coder/coder/v2/aibridge/mcp" + "github.com/coder/coder/v2/aibridge/utils" + "github.com/coder/quartz" +) + +// OpenAI Responses API request and response bodies. +const ( + requestBody = `{"input":"hi","model":"gpt-4o-mini"}` + successBody = `{"id":"resp_01","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_01","role":"assistant","content":[{"type":"output_text","text":"Hello!"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}` + toolUseBody = `{"id":"resp_01","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"function_call","id":"fc_01","call_id":"call_01","name":"test_tool","arguments":"{}","status":"completed"}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}}` + textCompleteBody = `{"id":"resp_02","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_02","role":"assistant","content":[{"type":"output_text","text":"done"}]}],"usage":{"input_tokens":15,"output_tokens":3,"total_tokens":18}}` + rateLimitBody = `{"error":{"message":"Rate limit exceeded","type":"rate_limit_error","code":"rate_limit_exceeded"}}` + authErrorBody = `{"error":{"message":"Invalid API key","type":"invalid_request_error","code":"invalid_api_key"}}` + serverErrorBody = `{"error":{"message":"Internal server error","type":"server_error","code":"internal_error"}}` +) + +type upstreamResponse struct { + statusCode int + body string + headers map[string]string +} + +func TestBlockingResponsesInterceptor_KeyFailover(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + // Centralized pool keys. Empty when byokKey is set. + keys []string + // BYOK key. Empty when keys is set. + byokKey string + // Scripted upstream responses keyed by bearer token. + responses map[string]upstreamResponse + expectedRequestCount int32 + expectedStatusCode int + expectedRetryAfter string + // Expected key states after the request, by index in keys. + expectedKeyStates []keypool.KeyState + }{ + { + // Given: 1 valid key returning 200. + // Then: 1 request, 200 response, key remains valid. + name: "single_valid_key", + keys: []string{"k0"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusOK, body: successBody}, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid}, + }, + { + // Given: 2 keys; key-0 returns 429, key-1 returns 200. + // Then: 2 requests, 200 response, key-0 temporary, key-1 valid. + name: "failover_after_429", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + "k1": {statusCode: http.StatusOK, body: successBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 returns 401, key-1 returns 200. + // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + name: "failover_after_401", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + "k1": {statusCode: http.StatusOK, body: successBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 returns 403, key-1 returns 200. + // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + name: "failover_after_403", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusForbidden, body: authErrorBody}, + "k1": {statusCode: http.StatusOK, body: successBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStateValid, + }, + }, + { + // Given: 3 keys; all return 429 with cooldowns 5s, 3s, 10s. + // Then: 3 requests, 429 response with smallest Retry-After, + // all keys temporary. + name: "all_keys_rate_limited", + keys: []string{"k0", "k1", "k2"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + "k1": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "3"}, + body: rateLimitBody, + }, + "k2": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "10"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 3, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "3", + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + }, + }, + { + // Given: 2 keys; both return 401. + // Then: 2 requests, 502 api_error response, both keys permanent. + name: "all_keys_unauthorized", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + "k1": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusBadGateway, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStatePermanent, + }, + }, + { + // Given: 2 keys; key-0 returns 500. + // Then: 1 request, 500 response, both keys remain valid. + name: "server_error_no_failover", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusInternalServerError, body: serverErrorBody}, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusInternalServerError, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, + { + // Given: BYOK with a single key returning 429. + // Then: 1 request, 429 response, no failover. + name: "byok_no_failover", + byokKey: "user-byok", + responses: map[string]upstreamResponse{ + "user-byok": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{ + "Retry-After": "5", + // BYOK doesn't set MaxRetries(0); + // suppress SDK retries to test a + // single attempt. + "x-should-retry": "false", + }, + body: rateLimitBody, + }, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusTooManyRequests, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Mock upstream: counts requests and returns + // scripted responses keyed by bearer token. An + // unmapped key falls through to 500 so misconfigured + // cases surface via the status assertion. + var requestCount atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + resp, ok := tc.responses[utils.ExtractBearerToken(r.Header.Get("Authorization"))] + if !ok { + resp = upstreamResponse{statusCode: http.StatusInternalServerError} + } + w.Header().Set("Content-Type", "application/json") + for hk, hv := range resp.headers { + w.Header().Set(hk, hv) + } + w.WriteHeader(resp.statusCode) + _, _ = w.Write([]byte(resp.body)) + })) + t.Cleanup(upstream.Close) + + cfg := config.OpenAI{BaseURL: upstream.URL + "/"} + var pool *keypool.Pool + if len(tc.keys) > 0 { + var err error + pool, err = keypool.New(tc.keys, quartz.NewMock(t)) + require.NoError(t, err) + cfg.KeyPool = pool + } else if tc.byokKey != "" { + cfg.Key = tc.byokKey + } + + payload, err := NewRequestPayload([]byte(requestBody)) + require.NoError(t, err) + + interceptor := NewBlockingInterceptor( + uuid.New(), + payload, + config.ProviderOpenAI, + cfg, + http.Header{}, + "Authorization", + otel.Tracer("blocking_test"), + intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""), + ) + interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + w := httptest.NewRecorder() + err = interceptor.ProcessRequest(w, req) + if tc.expectedStatusCode == http.StatusOK { + require.NoError(t, err) + } else { + require.Error(t, err) + } + + assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count") + assert.Equal(t, tc.expectedStatusCode, w.Code, "response status code") + assert.Equal(t, tc.expectedRetryAfter, w.Header().Get("Retry-After"), "Retry-After header") + if pool != nil { + assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") + } + }) + } +} + +// TestBlockingResponsesInterceptor_AgenticLoopFailover covers +// the scenarios that span an agentic-loop continuation: the +// initial client request and the subsequent tool-call +// continuation can each fail over independently. Each iteration +// gets its own walker. +func TestBlockingResponsesInterceptor_AgenticLoopFailover(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + // Scripted upstream responses consumed in order of + // upstream request. + responses []upstreamResponse + expectedRequestCount int32 + expectedSeenKeys []string + expectedStatusCode int + expectedKeyStates []keypool.KeyState + }{ + { + // Given: 2 keys; both upstream calls succeed on key-0. + // Then: 2 requests, 200 response, both keys remain valid. + name: "happy_path", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, body: toolUseBody}, + {statusCode: http.StatusOK, body: textCompleteBody}, + }, + expectedRequestCount: 2, + expectedSeenKeys: []string{"k0", "k0"}, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 succeeds initially, then 429s + // during the agentic continuation, key-1 succeeds. + // Then: 3 requests, 200 response, key-0 temporary, + // key-1 valid. + name: "agentic_failover_to_k1", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, body: toolUseBody}, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + {statusCode: http.StatusOK, body: textCompleteBody}, + }, + expectedRequestCount: 3, + expectedSeenKeys: []string{"k0", "k0", "k1"}, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 succeeds initially, then both + // keys 429 during the agentic continuation. + // Then: 3 requests, 429 response with smallest + // Retry-After, both keys temporary. + name: "agentic_all_keys_fail", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, body: toolUseBody}, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "3"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 3, + expectedSeenKeys: []string{"k0", "k0", "k1"}, + expectedStatusCode: http.StatusTooManyRequests, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var requestCount atomic.Int32 + var seenKeysMu sync.Mutex + var seenKeys []string + + // Mock upstream: returns scripted responses in order, + // records each request's bearer token for assertions. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + idx := int(requestCount.Add(1)) - 1 + seenKeysMu.Lock() + seenKeys = append(seenKeys, utils.ExtractBearerToken(r.Header.Get("Authorization"))) + seenKeysMu.Unlock() + _, _ = io.Copy(io.Discard, r.Body) + + if idx >= len(tc.responses) { + w.WriteHeader(http.StatusInternalServerError) + return + } + resp := tc.responses[idx] + w.Header().Set("Content-Type", "application/json") + for hk, hv := range resp.headers { + w.Header().Set(hk, hv) + } + w.WriteHeader(resp.statusCode) + _, _ = w.Write([]byte(resp.body)) + })) + t.Cleanup(upstream.Close) + + pool, err := keypool.New([]string{"k0", "k1"}, quartz.NewMock(t)) + require.NoError(t, err) + + cfg := config.OpenAI{ + BaseURL: upstream.URL + "/", + KeyPool: pool, + } + + payload, err := NewRequestPayload([]byte(requestBody)) + require.NoError(t, err) + + interceptor := NewBlockingInterceptor( + uuid.New(), + payload, + config.ProviderOpenAI, + cfg, + http.Header{}, + "Authorization", + otel.Tracer("blocking_test"), + intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""), + ) + + // Mock proxy with a tool the upstream's function_call + // response will reference. + proxy := &mockServerProxier{ + tools: []*mcp.Tool{ + { + Client: stubToolCaller{}, + ID: "test_tool", + Name: "test_tool", + ServerName: "coder", + Logger: slog.Make(), + }, + }, + } + interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, proxy) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + w := httptest.NewRecorder() + err = interceptor.ProcessRequest(w, req) + if tc.expectedStatusCode == http.StatusOK { + require.NoError(t, err) + } else { + require.Error(t, err) + } + + assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count") + assert.Equal(t, tc.expectedStatusCode, w.Code, "response status code") + + seenKeysMu.Lock() + defer seenKeysMu.Unlock() + assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys") + assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") + }) + } +} + +// mockServerProxier is a test implementation of mcp.ServerProxier. +type mockServerProxier struct { + tools []*mcp.Tool +} + +func (*mockServerProxier) Init(context.Context) error { + return nil +} + +func (*mockServerProxier) Shutdown(context.Context) error { + return nil +} + +func (m *mockServerProxier) ListTools() []*mcp.Tool { + return m.tools +} + +func (m *mockServerProxier) GetTool(id string) *mcp.Tool { + for _, t := range m.tools { + if t.ID == id { + return t + } + } + return nil +} + +func (*mockServerProxier) CallTool(context.Context, string, any) (*mcplib.CallToolResult, error) { + return nil, nil //nolint:nilnil // mock: no-op implementation +} + +// stubToolCaller is a minimal mcp.ToolCaller that returns a fixed +// text result, so the agentic continuation can proceed. +type stubToolCaller struct{} + +func (stubToolCaller) CallTool(_ context.Context, _ mcplib.CallToolRequest) (*mcplib.CallToolResult, error) { + return mcplib.NewToolResultText("tool result"), nil +} diff --git a/aibridge/intercept/responses/streaming.go b/aibridge/intercept/responses/streaming.go index 15847fb4d6..6730dcb04d 100644 --- a/aibridge/intercept/responses/streaming.go +++ b/aibridge/intercept/responses/streaming.go @@ -20,6 +20,7 @@ import ( aibcontext "github.com/coder/coder/v2/aibridge/context" "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/intercept/eventstream" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/recorder" "github.com/coder/coder/v2/aibridge/tracing" @@ -108,36 +109,87 @@ func (i *StreamingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r for shouldLoop { shouldLoop = false - respCopy = responseCopier{} - opts := i.requestOptions(&respCopy) + // Per-iteration walker. An iteration is either an agentic + // continuation (sending a tool result back in a new + // stream) or a failover retry (previous key marked, try + // the next one). + var walker *keypool.Walker + if i.cfg.KeyPool != nil { + walker = i.cfg.KeyPool.Walker() + } - // TODO(ssncferreira): inject actor headers directly in the client-header - // middleware instead of using SDK options. - if actor := aibcontext.ActorFromContext(r.Context()); actor != nil && i.cfg.SendActorHeaders { - opts = append(opts, intercept.ActorHeadersAsOpenAIOpts(actor)...) + // Failover sub-loop: try keys until a stream starts + // successfully or we hit a non-recoverable error. + var stream *ssestream.Stream[responses.ResponseStreamEventUnion] + var startErr error + for { + respCopy = responseCopier{} + opts := i.requestOptions(&respCopy) + + // TODO(ssncferreira): inject actor headers directly in the client-header + // middleware instead of using SDK options. + if actor := aibcontext.ActorFromContext(r.Context()); actor != nil && i.cfg.SendActorHeaders { + opts = append(opts, intercept.ActorHeadersAsOpenAIOpts(actor)...) + } + + var currentKey *keypool.Key + if walker != nil { + key, err := walker.Next() + if respErr := processKeyPoolError(err); respErr != nil { + // Pool exhausted: write the error directly. In + // agentic mode the inner loop buffers events + // instead of streaming them downstream, so the + // SSE connection has not been opened yet. + i.writeUpstreamError(w, respErr) + return xerrors.Errorf("key pool exhausted: %w", err) + } + currentKey = key + opts = append(opts, + option.WithAPIKey(key.Value()), + // Disable SDK retries because the failover + // loop handles retries via key rotation. + option.WithMaxRetries(0), + ) + } + + stream = i.newStream(ctx, srv, opts) + if upstreamErr := stream.Err(); upstreamErr != nil { + // Pre-stream failure of this attempt. For + // centralized requests, mark the key and + // retry with the next one. + if currentKey != nil && i.markKeyOnError(ctx, currentKey, upstreamErr) { + stream.Close() + continue + } + // Non-key error: stop trying and let the + // existing handling below report it. + startErr = upstreamErr + break + } + // Stream started successfully: commit to this key. + break } - stream := i.newStream(ctx, srv, opts) // func scope to defer steam.Close() err := func() error { defer stream.Close() - if upstreamErr := stream.Err(); upstreamErr != nil { + if startErr != nil { // events stream should never be initialized if events.IsStreaming() { i.logger.Warn(ctx, "event stream was initialized when no response was received from upstream") - return upstreamErr + return startErr } // no response received from upstream (eg. client/connection error), return custom error if !respCopy.responseReceived.Load() { - i.sendCustomErr(ctx, w, http.StatusInternalServerError, upstreamErr) - return upstreamErr + i.sendCustomErr(ctx, w, http.StatusInternalServerError, startErr) + return startErr } // forward received response as-is err := respCopy.forwardResp(w) - return errors.Join(upstreamErr, err) + return errors.Join(startErr, err) } for stream.Next() { diff --git a/aibridge/intercept/responses/streaming_test.go b/aibridge/intercept/responses/streaming_test.go new file mode 100644 index 0000000000..88e3b85495 --- /dev/null +++ b/aibridge/intercept/responses/streaming_test.go @@ -0,0 +1,499 @@ +package responses //nolint:testpackage // tests unexported internals + +import ( + "io" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/intercept" + "github.com/coder/coder/v2/aibridge/internal/testutil" + "github.com/coder/coder/v2/aibridge/keypool" + "github.com/coder/coder/v2/aibridge/mcp" + "github.com/coder/coder/v2/aibridge/utils" + "github.com/coder/quartz" +) + +// Streaming request body for the OpenAI Responses API. +const streamingRequestBody = `{"input":"hi","model":"gpt-4o-mini","stream":true}` + +// OpenAI Responses API SSE body for a successful streaming response. +const streamingSuccessBody = `event: response.created +data: {"type":"response.created","response":{"id":"resp_01","object":"response","status":"in_progress"},"sequence_number":0} + +event: response.completed +data: {"type":"response.completed","response":{"id":"resp_01","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_01","role":"assistant","content":[{"type":"output_text","text":"Hello!"}]}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}},"sequence_number":1} + +` + +func TestStreamingResponsesInterceptor_KeyFailover(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + // Centralized pool keys. Empty when byokKey is set. + keys []string + // BYOK key. Empty when keys is set. + byokKey string + // Scripted upstream responses keyed by bearer token. + responses map[string]upstreamResponse + expectedRequestCount int32 + expectedStatusCode int + expectedRetryAfter string + // Expected key states after the request, by index in keys. + expectedKeyStates []keypool.KeyState + }{ + { + // Given: 1 valid key returning a successful stream. + // Then: 1 request, 200 response, key remains valid. + name: "single_valid_key", + keys: []string{"k0"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusOK, + headers: map[string]string{"Content-Type": "text/event-stream"}, + body: streamingSuccessBody, + }, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid}, + }, + { + // Given: 2 keys; key-0 returns 429 pre-stream, key-1 + // streams successfully. + // Then: 2 requests, 200 response, key-0 temporary, key-1 valid. + name: "failover_after_429", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + "k1": { + statusCode: http.StatusOK, + headers: map[string]string{"Content-Type": "text/event-stream"}, + body: streamingSuccessBody, + }, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 returns 401 pre-stream, key-1 + // streams successfully. + // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + name: "failover_after_401", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + "k1": { + statusCode: http.StatusOK, + headers: map[string]string{"Content-Type": "text/event-stream"}, + body: streamingSuccessBody, + }, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 returns 403 pre-stream, key-1 streams. + // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + name: "failover_after_403", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusForbidden, body: authErrorBody}, + "k1": { + statusCode: http.StatusOK, + headers: map[string]string{"Content-Type": "text/event-stream"}, + body: streamingSuccessBody, + }, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusOK, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStateValid, + }, + }, + { + // Given: 3 keys; all return 429 pre-stream with + // cooldowns 5s, 3s, 10s. + // Then: 3 requests, 429 response with smallest + // Retry-After, all keys temporary. + name: "all_keys_rate_limited", + keys: []string{"k0", "k1", "k2"}, + responses: map[string]upstreamResponse{ + "k0": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + "k1": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "3"}, + body: rateLimitBody, + }, + "k2": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "10"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 3, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "3", + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + }, + }, + { + // Given: 2 keys; both return 401 pre-stream. + // Then: 2 requests, 502 api_error response, both keys permanent. + name: "all_keys_unauthorized", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + "k1": {statusCode: http.StatusUnauthorized, body: authErrorBody}, + }, + expectedRequestCount: 2, + expectedStatusCode: http.StatusBadGateway, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStatePermanent, + keypool.KeyStatePermanent, + }, + }, + { + // Given: 2 keys; key-0 returns 500 pre-stream. + // Then: 1 request, 500 response, both keys remain valid. + name: "server_error_no_failover", + keys: []string{"k0", "k1"}, + responses: map[string]upstreamResponse{ + "k0": {statusCode: http.StatusInternalServerError, body: serverErrorBody}, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusInternalServerError, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, + { + // Given: BYOK with a single key returning 429. + // Then: 1 request, 429 response, no failover. + name: "byok_no_failover", + byokKey: "user-byok", + responses: map[string]upstreamResponse{ + "user-byok": { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{ + "Retry-After": "5", + // BYOK doesn't set MaxRetries(0); + // suppress SDK retries to test a + // single attempt. + "x-should-retry": "false", + }, + body: rateLimitBody, + }, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusTooManyRequests, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Mock upstream: counts requests and returns + // scripted responses keyed by bearer token. An + // unmapped key falls through to 500 so misconfigured + // cases surface via the status assertion. + var requestCount atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + _, _ = io.Copy(io.Discard, r.Body) + resp, ok := tc.responses[utils.ExtractBearerToken(r.Header.Get("Authorization"))] + if !ok { + resp = upstreamResponse{statusCode: http.StatusInternalServerError} + } + for hk, hv := range resp.headers { + w.Header().Set(hk, hv) + } + w.WriteHeader(resp.statusCode) + _, _ = w.Write([]byte(resp.body)) + })) + t.Cleanup(upstream.Close) + + cfg := config.OpenAI{BaseURL: upstream.URL + "/"} + var pool *keypool.Pool + if len(tc.keys) > 0 { + var err error + pool, err = keypool.New(tc.keys, quartz.NewMock(t)) + require.NoError(t, err) + cfg.KeyPool = pool + } else if tc.byokKey != "" { + cfg.Key = tc.byokKey + } + + payload, err := NewRequestPayload([]byte(streamingRequestBody)) + require.NoError(t, err) + + interceptor := NewStreamingInterceptor( + uuid.New(), + payload, + config.ProviderOpenAI, + cfg, + http.Header{}, + "Authorization", + otel.Tracer("streaming_test"), + intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""), + ) + interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, nil) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + w := httptest.NewRecorder() + err = interceptor.ProcessRequest(w, req) + if tc.expectedStatusCode == http.StatusOK { + require.NoError(t, err) + } else { + require.Error(t, err) + } + + assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count") + assert.Equal(t, tc.expectedStatusCode, w.Code, "response status code") + assert.Equal(t, tc.expectedRetryAfter, w.Header().Get("Retry-After"), "Retry-After header") + if pool != nil { + assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") + } + }) + } +} + +// SSE bodies covering an agentic-continuation flow. +const ( + // First response: a function_call output referencing the + // injected "test_tool". Triggers the agentic continuation + // loop. + toolUseStreamBody = `event: response.created +data: {"type":"response.created","response":{"id":"resp_01","object":"response","status":"in_progress"},"sequence_number":0} + +event: response.completed +data: {"type":"response.completed","response":{"id":"resp_01","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"function_call","id":"fc_01","call_id":"call_01","name":"test_tool","arguments":"{}","status":"completed"}],"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15}},"sequence_number":1} + +` + + // Second response (after the tool result is sent back): + // a plain text message that ends the loop. + textStreamBody = `event: response.created +data: {"type":"response.created","response":{"id":"resp_02","object":"response","status":"in_progress"},"sequence_number":0} + +event: response.completed +data: {"type":"response.completed","response":{"id":"resp_02","object":"response","status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_02","role":"assistant","content":[{"type":"output_text","text":"done"}]}],"usage":{"input_tokens":15,"output_tokens":3,"total_tokens":18}},"sequence_number":1} + +` +) + +// TestStreamingResponsesInterceptor_AgenticLoopFailover covers +// the scenarios that span an agentic-loop continuation: the +// initial client request and the subsequent tool-call +// continuation can each fail over independently. Each iteration +// gets its own walker. +func TestStreamingResponsesInterceptor_AgenticLoopFailover(t *testing.T) { + t.Parallel() + + sseHeaders := map[string]string{"Content-Type": "text/event-stream"} + + tests := []struct { + name string + // Scripted upstream responses consumed in order of + // upstream request. + responses []upstreamResponse + expectedRequestCount int32 + expectedSeenKeys []string + // Substring expected in the response body. Either a + // success marker (e.g. "done") or an error marker + // (e.g. "rate_limit_error"). + expectedBodyContains string + // True when ProcessRequest is expected to return an + // error (e.g. all keys exhausted). + expectedErr bool + expectedKeyStates []keypool.KeyState + }{ + { + // Given: 2 keys; both upstream calls succeed on key-0. + // Then: 2 requests, success body, both keys remain valid. + name: "happy_path", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody}, + {statusCode: http.StatusOK, headers: sseHeaders, body: textStreamBody}, + }, + expectedRequestCount: 2, + expectedSeenKeys: []string{"k0", "k0"}, + expectedBodyContains: "done", + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 succeeds initially, then 429s + // during the agentic continuation, key-1 succeeds. + // Then: 3 requests, success body, key-0 temporary, + // key-1 valid. + name: "agentic_failover_to_k1", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody}, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + {statusCode: http.StatusOK, headers: sseHeaders, body: textStreamBody}, + }, + expectedRequestCount: 3, + expectedSeenKeys: []string{"k0", "k0", "k1"}, + expectedBodyContains: "done", + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, + }, + { + // Given: 2 keys; key-0 succeeds initially, then both + // keys 429 during the agentic continuation. + // Then: 3 requests, error injected as SSE event, both + // keys temporary. + name: "agentic_all_keys_fail", + responses: []upstreamResponse{ + {statusCode: http.StatusOK, headers: sseHeaders, body: toolUseStreamBody}, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "5"}, + body: rateLimitBody, + }, + { + statusCode: http.StatusTooManyRequests, + headers: map[string]string{"Retry-After": "3"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 3, + expectedSeenKeys: []string{"k0", "k0", "k1"}, + expectedBodyContains: "all configured keys are rate-limited", + expectedErr: true, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + var requestCount atomic.Int32 + var seenKeysMu sync.Mutex + var seenKeys []string + + // Mock upstream: returns scripted responses in order, + // records each request's bearer token for assertions. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + idx := int(requestCount.Add(1)) - 1 + seenKeysMu.Lock() + seenKeys = append(seenKeys, utils.ExtractBearerToken(r.Header.Get("Authorization"))) + seenKeysMu.Unlock() + _, _ = io.Copy(io.Discard, r.Body) + + if idx >= len(tc.responses) { + w.WriteHeader(http.StatusInternalServerError) + return + } + resp := tc.responses[idx] + for hk, hv := range resp.headers { + w.Header().Set(hk, hv) + } + w.WriteHeader(resp.statusCode) + _, _ = w.Write([]byte(resp.body)) + })) + t.Cleanup(upstream.Close) + + pool, err := keypool.New([]string{"k0", "k1"}, quartz.NewMock(t)) + require.NoError(t, err) + + cfg := config.OpenAI{ + BaseURL: upstream.URL + "/", + KeyPool: pool, + } + + payload, err := NewRequestPayload([]byte(streamingRequestBody)) + require.NoError(t, err) + + interceptor := NewStreamingInterceptor( + uuid.New(), + payload, + config.ProviderOpenAI, + cfg, + http.Header{}, + "Authorization", + otel.Tracer("streaming_test"), + intercept.NewCredentialInfo(intercept.CredentialKindCentralized, ""), + ) + + // Mock proxy with a tool the upstream's function_call + // response will reference. The stub caller returns a + // fixed text result. + proxy := &mockServerProxier{ + tools: []*mcp.Tool{ + { + Client: stubToolCaller{}, + ID: "test_tool", + Name: "test_tool", + ServerName: "coder", + Logger: slog.Make(), + }, + }, + } + interceptor.Setup(slog.Make(), &testutil.MockRecorder{}, proxy) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + w := httptest.NewRecorder() + err = interceptor.ProcessRequest(w, req) + if tc.expectedErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + + assert.Equal(t, tc.expectedRequestCount, requestCount.Load(), "upstream request count") + body := w.Body.String() + assert.Contains(t, body, tc.expectedBodyContains, "response body") + + seenKeysMu.Lock() + defer seenKeysMu.Unlock() + assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys") + assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") + }) + } +} diff --git a/aibridge/internal/integrationtest/keypool_failover_test.go b/aibridge/internal/integrationtest/keypool_failover_test.go index bab2552a28..a96b870cd6 100644 --- a/aibridge/internal/integrationtest/keypool_failover_test.go +++ b/aibridge/internal/integrationtest/keypool_failover_test.go @@ -17,9 +17,141 @@ import ( "github.com/coder/coder/v2/aibridge/fixtures" "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/provider" + "github.com/coder/coder/v2/aibridge/utils" "github.com/coder/quartz" ) +// TestOpenAI_KeyFailover verifies that a pool's key state +// persists across distinct client requests for both OpenAI APIs +// (chat completions and responses), in both blocking and +// streaming modes. A key marked temporary on request 1 is +// skipped on request 2 without a wasted upstream attempt. +func TestOpenAI_KeyFailover(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + fixture []byte + path string + streaming bool + successCType string + }{ + { + name: "chatcompletions_blocking", + fixture: fixtures.OaiChatSimple, + path: pathOpenAIChatCompletions, + streaming: false, + successCType: "application/json", + }, + { + name: "chatcompletions_streaming", + fixture: fixtures.OaiChatSimple, + path: pathOpenAIChatCompletions, + streaming: true, + successCType: "text/event-stream", + }, + { + name: "responses_blocking", + fixture: fixtures.OaiResponsesBlockingSimple, + path: pathOpenAIResponses, + streaming: false, + successCType: "application/json", + }, + { + name: "responses_streaming", + fixture: fixtures.OaiResponsesStreamingSimple, + path: pathOpenAIResponses, + streaming: true, + successCType: "text/event-stream", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + fix := fixtures.Parse(t, tc.fixture) + var successBody []byte + if tc.streaming { + successBody = fix.Streaming() + } else { + successBody = fix.NonStreaming() + } + + pool, err := keypool.New([]string{"k0", "k1"}, quartz.NewMock(t)) + require.NoError(t, err) + + var requestCount atomic.Int32 + var seenKeysMu sync.Mutex + var seenKeys []string + + // Mock upstream: k0 always returns 429, k1 returns + // the per-test success body. + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requestCount.Add(1) + key := utils.ExtractBearerToken(r.Header.Get("Authorization")) + seenKeysMu.Lock() + seenKeys = append(seenKeys, key) + seenKeysMu.Unlock() + _, _ = io.Copy(io.Discard, r.Body) + + switch key { + case "k0": + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "60") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = fmt.Fprint(w, `{"error":{"type":"rate_limit_error","message":"rate limited","code":"rate_limit_exceeded"}}`) + case "k1": + w.Header().Set("Content-Type", tc.successCType) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(successBody) + default: + w.WriteHeader(http.StatusInternalServerError) + } + })) + t.Cleanup(upstream.Close) + + bridgeServer := newBridgeTestServer(t.Context(), t, upstream.URL, + withCustomProvider(provider.NewOpenAI(config.OpenAI{ + BaseURL: upstream.URL, + KeyPool: pool, + })), + ) + + requestBody, err := sjson.SetBytes(fix.Request(), "stream", tc.streaming) + require.NoError(t, err) + + // Request 1: walker starts at k0, fails over to k1 + // after 429. + resp, err := bridgeServer.makeRequest(t, http.MethodPost, tc.path, requestBody) + require.NoError(t, err) + _, _ = io.Copy(io.Discard, resp.Body) + require.NoError(t, resp.Body.Close()) + require.Equal(t, http.StatusOK, resp.StatusCode) + + // Request 2: walker skips the now-temporary k0 and + // goes straight to k1 (1 upstream call, not 2). + resp, err = bridgeServer.makeRequest(t, http.MethodPost, tc.path, requestBody) + require.NoError(t, err) + _, _ = io.Copy(io.Discard, resp.Body) + require.NoError(t, resp.Body.Close()) + require.Equal(t, http.StatusOK, resp.StatusCode) + + seenKeysMu.Lock() + defer seenKeysMu.Unlock() + // Request 1: 2 calls (k0 then k1). Request 2: 1 call (k1). + assert.Equal(t, int32(3), requestCount.Load(), "upstream request count") + assert.Equal(t, []string{"k0", "k1", "k1"}, seenKeys, "seen keys") + + // Pool state persists: k0 temporary, k1 valid. + assert.Equal(t, []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, pool.PoolState(), "key states") + }) + } +} + // TestAnthropic_KeyFailover verifies that a pool's key state // persists across distinct client requests: a key marked // temporary on request 1 is still skipped on request 2 without diff --git a/aibridge/internal/integrationtest/trace_test.go b/aibridge/internal/integrationtest/trace_test.go index f3e835ca8a..d70b1ffccf 100644 --- a/aibridge/internal/integrationtest/trace_test.go +++ b/aibridge/internal/integrationtest/trace_test.go @@ -667,13 +667,12 @@ func TestTraceOpenAIErr(t *testing.T) { }, }, { - name: "trace_openai_responses_streaming_http_error", - fixture: fixtures.OaiResponsesStreamingHTTPErr, - streaming: true, - allowOverflow: true, // 429 error causes retries + name: "trace_openai_responses_streaming_http_error", + fixture: fixtures.OaiResponsesStreamingHTTPErr, + streaming: true, path: pathOpenAIResponses, - expectCode: http.StatusTooManyRequests, + expectCode: http.StatusBadRequest, expect: []expectTrace{ {"Intercept", 1, codes.Error}, {"Intercept.CreateInterceptor", 1, codes.Unset}, @@ -689,7 +688,7 @@ func TestTraceOpenAIErr(t *testing.T) { streaming: false, path: pathOpenAIResponses, - expectCode: http.StatusUnauthorized, + expectCode: http.StatusBadRequest, expect: []expectTrace{ {"Intercept", 1, codes.Error}, {"Intercept.CreateInterceptor", 1, codes.Unset}, diff --git a/aibridge/provider/openai.go b/aibridge/provider/openai.go index 80d3eb5eef..c670626b11 100644 --- a/aibridge/provider/openai.go +++ b/aibridge/provider/openai.go @@ -16,8 +16,10 @@ import ( "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/intercept/chatcompletions" "github.com/coder/coder/v2/aibridge/intercept/responses" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/tracing" "github.com/coder/coder/v2/aibridge/utils" + "github.com/coder/quartz" ) const ( @@ -44,6 +46,25 @@ func NewOpenAI(cfg config.OpenAI) *OpenAI { if cfg.BaseURL == "" { cfg.BaseURL = "https://api.openai.com/v1/" } + // Resolve centralized key configuration into KeyPool. + // Precedence: + // 1. cfg.KeyPool (explicit, highest priority). + // 2. cfg.Key (legacy single key). + // After this block cfg.Key is empty so it can only carry a + // BYOK Authorization Bearer set per interception in + // CreateInterceptor. + // TODO(ssncferreira): simplify auth field resolution per + // https://github.com/coder/aibridge/issues/266. + if cfg.KeyPool == nil && cfg.Key != "" { + // keypool.New only fails on empty or duplicate keys, + // neither possible with a single non-empty key. + pool, err := keypool.New([]string{cfg.Key}, quartz.NewReal()) + if err != nil { + panic(fmt.Sprintf("openai provider: build single-key pool: %s", err)) + } + cfg.KeyPool = pool + } + cfg.Key = "" if cfg.CircuitBreaker != nil { cfg.CircuitBreaker.OpenErrorResponse = openAIOpenErrorResponse } @@ -100,20 +121,34 @@ func (p *OpenAI) CreateInterceptor(_ http.ResponseWriter, r *http.Request, trace var interceptor intercept.Interceptor cfg := p.cfg - // At this point the request contains only LLM provider headers. Any - // Coder-specific authentication has already been stripped. + // At this point the request contains only LLM provider headers. + // Any Coder-specific authentication has already been stripped. // // In centralized mode Authorization is absent, so cfg keeps the - // centralized key unchanged. + // KeyPool from provider construction and the failover loop walks + // it. // - // In BYOK mode the user's credential is in Authorization. Replace - // the centralized key with it so it is forwarded upstream. + // In BYOK mode the user's credential is in Authorization, + // populate cfg.Key and clear cfg.KeyPool so failover is disabled. + // + // TODO(ssncferreira): consolidate auth field handling per + // https://github.com/coder/aibridge/issues/266. credKind := intercept.CredentialKindCentralized + var credSecret string if token := utils.ExtractBearerToken(r.Header.Get("Authorization")); token != "" { cfg.Key = token + cfg.KeyPool = nil credKind = intercept.CredentialKindBYOK + credSecret = token + } else if cfg.KeyPool != nil { + // Centralized: use the first key as a placeholder hint. + // TODO(ssncferreira): record the actually-used key in + // the interception record to reflect failover. + if k, err := cfg.KeyPool.Walker().Next(); err == nil { + credSecret = k.Value() + } } - cred := intercept.NewCredentialInfo(credKind, cfg.Key) + cred := intercept.NewCredentialInfo(credKind, credSecret) path := strings.TrimPrefix(r.URL.Path, p.RoutePrefix()) switch path { @@ -171,7 +206,16 @@ func (p *OpenAI) InjectAuthHeader(headers *http.Header) { return } - headers.Set(p.AuthHeader(), "Bearer "+p.cfg.Key) + // Centralized: pull a single key from the pool. No failover + // or exhaustion handling here. + // TODO(ssncferreira): replace with RoundTripper-based auth + // in the upstack passthrough PR. + if p.cfg.KeyPool == nil { + return + } + if key, err := p.cfg.KeyPool.Walker().Next(); err == nil { + headers.Set(p.AuthHeader(), "Bearer "+key.Value()) + } } func (p *OpenAI) CircuitBreakerConfig() *config.CircuitBreaker { diff --git a/enterprise/cli/aibridged.go b/enterprise/cli/aibridged.go index 87f32f5d54..1dd6a2a05d 100644 --- a/enterprise/cli/aibridged.go +++ b/enterprise/cli/aibridged.go @@ -122,15 +122,18 @@ func buildProviders(cfg codersdk.AIBridgeConfig) ([]aibridge.Provider, error) { } switch p.Type { case aibridge.ProviderOpenAI: - // TODO(ssncferreira): pass a keypool.Pool instead. - var key string + var pool *keypool.Pool if len(p.Keys) > 0 { - key = p.Keys[0] + var err error + pool, err = keypool.New(p.Keys, quartz.NewReal()) + if err != nil { + return nil, xerrors.Errorf("create openai key pool for provider %q: %w", name, err) + } } providers = append(providers, aibridge.NewOpenAIProvider(aibridge.OpenAIConfig{ Name: name, BaseURL: p.BaseURL, - Key: key, + KeyPool: pool, APIDumpDir: p.DumpDir, CircuitBreaker: cbConfig, SendActorHeaders: cfg.SendActorHeaders.Value(),