diff --git a/aibridge/intercept/chatcompletions/base.go b/aibridge/intercept/chatcompletions/base.go index b8f896b92e..2c42c9c9df 100644 --- a/aibridge/intercept/chatcompletions/base.go +++ b/aibridge/intercept/chatcompletions/base.go @@ -27,6 +27,7 @@ import ( "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/recorder" "github.com/coder/coder/v2/aibridge/tracing" + "github.com/coder/coder/v2/aibridge/utils" "github.com/coder/quartz" ) @@ -188,7 +189,7 @@ func (i *interceptionBase) unmarshalArgs(in string) (args recorder.ToolArgs) { } // writeUpstreamError marshals and writes a given error. -func (i *interceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *responseError) { +func (i *interceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *ResponseError) { if oaiErr == nil { return } @@ -234,10 +235,10 @@ func (i *interceptionBase) markKeyOnError(ctx context.Context, key *keypool.Key, ) } -// processKeyPoolError translates a keypool exhaustion error +// 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 { +func ProcessKeyPoolError(err error) *ResponseError { var transient *keypool.TransientKeyPoolError switch { case errors.As(err, &transient): @@ -292,7 +293,7 @@ func calculateActualInputTokenUsage(in openai.CompletionUsage) int64 { in.PromptTokensDetails.CachedTokens /* The aggregated number of text input tokens that has been cached from previous requests. */ } -func getErrorResponse(err error) *responseError { +func getErrorResponse(err error) *ResponseError { var apiErr *openai.Error if !errors.As(err, &apiErr) { return nil @@ -300,16 +301,16 @@ func getErrorResponse(err error) *responseError { return newErrorResponse(apiErr.Message, apiErr.Type, apiErr.Code, apiErr.StatusCode, keypool.ParseRetryAfter(apiErr.Response)) } -var _ error = &responseError{} +var _ error = &ResponseError{} -type responseError struct { +type ResponseError struct { ErrorObject *shared.ErrorObject `json:"error"` StatusCode int `json:"-"` RetryAfter time.Duration `json:"-"` } -func newErrorResponse(msg, errType, code string, status int, retryAfter time.Duration) *responseError { - return &responseError{ +func newErrorResponse(msg, errType, code string, status int, retryAfter time.Duration) *ResponseError { + return &ResponseError{ ErrorObject: &shared.ErrorObject{ Code: code, Message: msg, @@ -320,9 +321,19 @@ func newErrorResponse(msg, errType, code string, status int, retryAfter time.Dur } } -func (a *responseError) Error() string { - if a.ErrorObject == nil { +func (e *ResponseError) Error() string { + if e.ErrorObject == nil { return "" } - return a.ErrorObject.Message + return e.ErrorObject.Message +} + +// ToResponse marshals e into an *http.Response shaped for the +// OpenAI API. +func (e *ResponseError) ToResponse() *http.Response { + body, err := json.Marshal(e) + if err != nil { + body = []byte(`{"error":{"type":"error","message":"error marshaling upstream error","code":"server_error"}}`) + } + return utils.NewJSONErrorResponse(e.StatusCode, e.RetryAfter, body) } diff --git a/aibridge/intercept/chatcompletions/base_test.go b/aibridge/intercept/chatcompletions/base_test.go index 4f7ffc7bcd..85b3cf5c23 100644 --- a/aibridge/intercept/chatcompletions/base_test.go +++ b/aibridge/intercept/chatcompletions/base_test.go @@ -127,7 +127,7 @@ func TestProcessKeyPoolError(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() - got := processKeyPoolError(tc.err) + got := ProcessKeyPoolError(tc.err) if tc.expectedNil { require.Nil(t, got) return @@ -207,7 +207,7 @@ func TestWriteUpstreamError(t *testing.T) { tests := []struct { name string - respErr *responseError + respErr *ResponseError expectStatus int // Empty string means the header should be absent. expectRetryAfter string diff --git a/aibridge/intercept/chatcompletions/blocking.go b/aibridge/intercept/chatcompletions/blocking.go index 23d46c75cf..aef420ee32 100644 --- a/aibridge/intercept/chatcompletions/blocking.go +++ b/aibridge/intercept/chatcompletions/blocking.go @@ -225,7 +225,7 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req // The failover loop may return a keypool exhaustion // error. Check before the SDK-error path. - if keyErr := processKeyPoolError(err); keyErr != nil { + if keyErr := ProcessKeyPoolError(err); keyErr != nil { i.writeUpstreamError(w, keyErr) return xerrors.Errorf("key pool exhausted: %w", err) } diff --git a/aibridge/intercept/chatcompletions/streaming.go b/aibridge/intercept/chatcompletions/streaming.go index b84898c264..38aa2d7154 100644 --- a/aibridge/intercept/chatcompletions/streaming.go +++ b/aibridge/intercept/chatcompletions/streaming.go @@ -144,7 +144,7 @@ func (i *StreamingInterception) ProcessRequest(w http.ResponseWriter, r *http.Re var currentKey *keypool.Key if walker != nil { key, err := walker.Next() - if respErr := processKeyPoolError(err); respErr != nil { + 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 @@ -474,7 +474,7 @@ func (i *StreamingInterception) newStream(ctx context.Context, svc openai.ChatCo // 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 { +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)) diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 14c276370e..8e32dd3a32 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -433,7 +433,7 @@ func filterBedrockBetaFlags(headers http.Header, model string) { } // writeUpstreamError marshals and writes a given error. -func (i *interceptionBase) writeUpstreamError(w http.ResponseWriter, antErr *responseError) { +func (i *interceptionBase) writeUpstreamError(w http.ResponseWriter, antErr *ResponseError) { if antErr == nil { return } @@ -537,10 +537,10 @@ func (i *interceptionBase) markKeyOnError(ctx context.Context, key *keypool.Key, ) } -// processKeyPoolError translates a keypool exhaustion error +// ProcessKeyPoolError translates a keypool exhaustion error // into a developer-facing responseError shaped for the Anthropic // API. Returns nil if err is not an exhaustion error. -func processKeyPoolError(err error) *responseError { +func ProcessKeyPoolError(err error) *ResponseError { var transient *keypool.TransientKeyPoolError switch { case errors.As(err, &transient): @@ -562,7 +562,7 @@ func processKeyPoolError(err error) *responseError { } } -func getErrorResponse(err error) *responseError { +func getErrorResponse(err error) *ResponseError { var apierr *anthropic.Error if !errors.As(err, &apierr) { return nil @@ -583,17 +583,17 @@ func getErrorResponse(err error) *responseError { return newErrorResponse(msg, errType, apierr.StatusCode, keypool.ParseRetryAfter(apierr.Response)) } -var _ error = &responseError{} +var _ error = &ResponseError{} -type responseError struct { +type ResponseError struct { *anthropic.ErrorResponse StatusCode int `json:"-"` RetryAfter time.Duration `json:"-"` } -func newErrorResponse(msg, errType string, status int, retryAfter time.Duration) *responseError { - return &responseError{ +func newErrorResponse(msg, errType string, status int, retryAfter time.Duration) *ResponseError { + return &ResponseError{ ErrorResponse: &shared.ErrorResponse{ Error: shared.ErrorObjectUnion{ Message: msg, @@ -606,9 +606,19 @@ func newErrorResponse(msg, errType string, status int, retryAfter time.Duration) } } -func (a *responseError) Error() string { - if a.ErrorResponse == nil { +func (e *ResponseError) Error() string { + if e.ErrorResponse == nil { return "" } - return a.ErrorResponse.Error.Message + return e.ErrorResponse.Error.Message +} + +// ToResponse marshals e into an *http.Response shaped for the +// Anthropic API. +func (e *ResponseError) ToResponse() *http.Response { + body, err := json.Marshal(e) + if err != nil { + body = []byte(`{"type":"error","error":{"type":"error","message":"error marshaling upstream error"}}`) + } + return utils.NewJSONErrorResponse(e.StatusCode, e.RetryAfter, body) } diff --git a/aibridge/intercept/messages/base_test.go b/aibridge/intercept/messages/base_test.go index a6accceeaa..52cdc2ab7b 100644 --- a/aibridge/intercept/messages/base_test.go +++ b/aibridge/intercept/messages/base_test.go @@ -1039,7 +1039,7 @@ func TestProcessKeyPoolError(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() - got := processKeyPoolError(tc.err) + got := ProcessKeyPoolError(tc.err) if tc.expectedNil { require.Nil(t, got) return @@ -1119,7 +1119,7 @@ func TestWriteUpstreamError(t *testing.T) { tests := []struct { name string - respErr *responseError + respErr *ResponseError expectStatus int // Empty string means the header should be absent. expectRetryAfter string diff --git a/aibridge/intercept/messages/blocking.go b/aibridge/intercept/messages/blocking.go index 7956acf088..dba348bfc8 100644 --- a/aibridge/intercept/messages/blocking.go +++ b/aibridge/intercept/messages/blocking.go @@ -114,7 +114,7 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req // The failover loop may return a keypool exhaustion // error. Check before the SDK-error path. - if keyErr := processKeyPoolError(err); keyErr != nil { + if keyErr := ProcessKeyPoolError(err); keyErr != nil { i.writeUpstreamError(w, keyErr) return xerrors.Errorf("key pool exhausted: %w", err) } diff --git a/aibridge/intercept/messages/streaming.go b/aibridge/intercept/messages/streaming.go index 6143da2710..b095e8e1b8 100644 --- a/aibridge/intercept/messages/streaming.go +++ b/aibridge/intercept/messages/streaming.go @@ -175,7 +175,7 @@ newStream: var currentKey *keypool.Key if walker != nil { key, err := walker.Next() - if respErr := processKeyPoolError(err); respErr != nil { + 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 @@ -599,10 +599,10 @@ newStream: } // mapStreamError converts a mid-stream upstream error or -// processing error into a relayable responseError. Returns nil +// 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 { +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)) diff --git a/aibridge/intercept/responses/base.go b/aibridge/intercept/responses/base.go index 059162c278..d35f95e461 100644 --- a/aibridge/intercept/responses/base.go +++ b/aibridge/intercept/responses/base.go @@ -35,6 +35,7 @@ import ( "github.com/coder/coder/v2/aibridge/mcp" "github.com/coder/coder/v2/aibridge/recorder" "github.com/coder/coder/v2/aibridge/tracing" + "github.com/coder/coder/v2/aibridge/utils" "github.com/coder/quartz" ) @@ -142,7 +143,7 @@ func (i *responsesInterceptionBase) validateRequest(ctx context.Context, w http. } // writeUpstreamError marshals and writes a given error. -func (i *responsesInterceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *responseError) { +func (i *responsesInterceptionBase) writeUpstreamError(w http.ResponseWriter, oaiErr *ResponseError) { if oaiErr == nil { return } @@ -188,10 +189,10 @@ func (i *responsesInterceptionBase) markKeyOnError(ctx context.Context, key *key ) } -// processKeyPoolError translates a keypool exhaustion error -// into a developer-facing responseError shaped for the OpenAI +// 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 { +func ProcessKeyPoolError(err error) *ResponseError { var transient *keypool.TransientKeyPoolError switch { case errors.As(err, &transient): @@ -215,8 +216,8 @@ func processKeyPoolError(err error) *responseError { } } -func newErrorResponse(msg, errType, code string, status int, retryAfter time.Duration) *responseError { - return &responseError{ +func newErrorResponse(msg, errType, code string, status int, retryAfter time.Duration) *ResponseError { + return &ResponseError{ ErrorObject: &shared.ErrorObject{ Code: code, Message: msg, @@ -227,19 +228,29 @@ func newErrorResponse(msg, errType, code string, status int, retryAfter time.Dur } } -var _ error = &responseError{} +var _ error = &ResponseError{} -type responseError struct { +type ResponseError struct { ErrorObject *shared.ErrorObject `json:"error"` StatusCode int `json:"-"` RetryAfter time.Duration `json:"-"` } -func (a *responseError) Error() string { - if a.ErrorObject == nil { +func (e *ResponseError) Error() string { + if e.ErrorObject == nil { return "" } - return a.ErrorObject.Message + return e.ErrorObject.Message +} + +// ToResponse marshals e into an *http.Response shaped for the +// OpenAI API. +func (e *ResponseError) ToResponse() *http.Response { + body, err := json.Marshal(e) + if err != nil { + body = []byte(`{"error":{"type":"error","message":"error marshaling upstream error","code":"server_error"}}`) + } + return utils.NewJSONErrorResponse(e.StatusCode, e.RetryAfter, body) } // sendCustomErr sends custom responses.Error error to the client diff --git a/aibridge/intercept/responses/base_test.go b/aibridge/intercept/responses/base_test.go index 4651b48f2f..7e965c7257 100644 --- a/aibridge/intercept/responses/base_test.go +++ b/aibridge/intercept/responses/base_test.go @@ -432,7 +432,7 @@ func TestProcessKeyPoolError(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { t.Parallel() - got := processKeyPoolError(tc.err) + got := ProcessKeyPoolError(tc.err) if tc.expectedNil { require.Nil(t, got) return @@ -512,7 +512,7 @@ func TestWriteUpstreamError(t *testing.T) { tests := []struct { name string - respErr *responseError + respErr *ResponseError expectStatus int // Empty string means the header should be absent. expectRetryAfter string diff --git a/aibridge/intercept/responses/blocking.go b/aibridge/intercept/responses/blocking.go index cec59307ea..d2dd7a3de8 100644 --- a/aibridge/intercept/responses/blocking.go +++ b/aibridge/intercept/responses/blocking.go @@ -103,7 +103,7 @@ func (i *BlockingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r * // The failover loop may return a keypool exhaustion // error. Render it here. if upstreamErr != nil { - if keyErr := processKeyPoolError(upstreamErr); keyErr != nil { + if keyErr := ProcessKeyPoolError(upstreamErr); keyErr != nil { i.writeUpstreamError(w, keyErr) return xerrors.Errorf("key pool exhausted: %w", upstreamErr) } diff --git a/aibridge/intercept/responses/streaming.go b/aibridge/intercept/responses/streaming.go index 6730dcb04d..6d1e70c9ca 100644 --- a/aibridge/intercept/responses/streaming.go +++ b/aibridge/intercept/responses/streaming.go @@ -135,7 +135,7 @@ func (i *StreamingResponsesInterceptor) ProcessRequest(w http.ResponseWriter, r var currentKey *keypool.Key if walker != nil { key, err := walker.Next() - if respErr := processKeyPoolError(err); respErr != nil { + 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 diff --git a/aibridge/internal/testutil/mockprovider.go b/aibridge/internal/testutil/mockprovider.go index 9bdcb10694..0c56cf2c9e 100644 --- a/aibridge/internal/testutil/mockprovider.go +++ b/aibridge/internal/testutil/mockprovider.go @@ -6,8 +6,10 @@ import ( "go.opentelemetry.io/otel/trace" + "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/intercept" + "github.com/coder/coder/v2/aibridge/keypool" ) type MockProvider struct { @@ -31,6 +33,10 @@ func (m *MockProvider) InjectAuthHeader(h *http.Header) { m.InjectAuthHeaderFunc(h) } } + +func (*MockProvider) KeyFailoverConfig(_ slog.Logger) keypool.KeyFailoverConfig { + return keypool.KeyFailoverConfig{} +} func (*MockProvider) CircuitBreakerConfig() *config.CircuitBreaker { return nil } func (*MockProvider) APIDumpDir() string { return "" } func (m *MockProvider) CreateInterceptor(w http.ResponseWriter, r *http.Request, tracer trace.Tracer) (intercept.Interceptor, error) { diff --git a/aibridge/keypool/failover.go b/aibridge/keypool/failover.go new file mode 100644 index 0000000000..f650cab258 --- /dev/null +++ b/aibridge/keypool/failover.go @@ -0,0 +1,120 @@ +package keypool + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + + "github.com/coder/coder/v2/aibridge/utils" +) + +// KeyFailoverConfig is the per-provider configuration consumed by +// NewKeyFailoverTransport. +type KeyFailoverConfig struct { + // Pool is the key pool to walk. Nil disables key failover. + Pool *Pool + + // IsBYOK returns true when the request already carries + // user-supplied auth. BYOK requests skip key failover. + IsBYOK func(*http.Request) bool + + // InjectAuthKey writes the key value into the outbound headers + // in the format the provider expects. + InjectAuthKey func(*http.Header, string) + + // MarkKey marks the key based on the upstream response. + // Returns true when the response is a key-specific error, + // causing the walker to advance and retry with the next key. + MarkKey func(ctx context.Context, key *Key, resp *http.Response) bool + + // BuildExhaustedResponse returns the response sent to the + // client when the walker has no more keys to try. + BuildExhaustedResponse func(err error) *http.Response +} + +// keyFailoverTransport retries inner across the key pool on +// key-specific failures. +type keyFailoverTransport struct { + inner http.RoundTripper + config KeyFailoverConfig +} + +// NewKeyFailoverTransport returns an http.RoundTripper backed by +// keyFailoverTransport. If config.Pool is nil, inner is returned +// unchanged. +func NewKeyFailoverTransport(inner http.RoundTripper, config KeyFailoverConfig) http.RoundTripper { + if config.Pool == nil { + return inner + } + return &keyFailoverTransport{ + inner: inner, + config: config, + } +} + +// RoundTrip is invoked by the proxy once per outer client request, +// after Rewrite has applied proxy headers. +// +// For centralized requests it walks the key pool, retrying on +// key-specific failures until one key succeeds or the pool is +// exhausted. BYOK requests skip the failover loop. +func (t *keyFailoverTransport) RoundTrip(req *http.Request) (*http.Response, error) { + if t.config.IsBYOK(req) { + return t.inner.RoundTrip(req) + } + + // Buffer once so retries can replay the body. + body, err := bufferBody(req) + if err != nil { + return nil, err + } + + // Fresh walker per request, independent of other inflight requests. + walker := t.config.Pool.Walker() + for { + key, err := walker.Next() + if err != nil { + resp := t.config.BuildExhaustedResponse(err) + if resp == nil { + // Fallback if BuildExhaustedResponse returns nil. + body := []byte(fmt.Sprintf(`{"error":"key pool exhausted: %s"}`, err)) + resp = utils.NewJSONErrorResponse(http.StatusBadGateway, 0, body) + } + return resp, nil + } + + // Clone per attempt so the original request isn't mutated. + outReq := req.Clone(req.Context()) + if body != nil { + outReq.Body = io.NopCloser(bytes.NewReader(body)) + } + t.config.InjectAuthKey(&outReq.Header, key.Value()) + + resp, rtErr := t.inner.RoundTrip(outReq) + if rtErr != nil { + // Transport-level error, not a key issue. + return resp, rtErr + } + // MarkKey returns true on key-specific failures (e.g. 401/403/429). + if t.config.MarkKey(req.Context(), key, resp) { + // Drain and retry with the next key. + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + continue + } + // Success or non-key error, forward as-is. + return resp, nil + } +} + +// bufferBody reads the request body fully so it can be replayed +// across key-failover retries. Returns nil for a nil body. +func bufferBody(req *http.Request) ([]byte, error) { + if req.Body == nil { + return nil, nil + } + defer req.Body.Close() + return io.ReadAll(req.Body) +} diff --git a/aibridge/keypool/failover_test.go b/aibridge/keypool/failover_test.go new file mode 100644 index 0000000000..c8fdc81c29 --- /dev/null +++ b/aibridge/keypool/failover_test.go @@ -0,0 +1,69 @@ +package keypool_test + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/aibridge/keypool" + "github.com/coder/quartz" +) + +// errFakeRoundTripperCalled is returned by fakeRoundTripper if it +// ever gets invoked. The constructor identity tests should never +// trigger a RoundTrip call. +var errFakeRoundTripperCalled = xerrors.New("fakeRoundTripper should not be invoked") + +// fakeRoundTripper is a no-op http.RoundTripper used to check +// constructor identity in tests. +type fakeRoundTripper struct{} + +func (*fakeRoundTripper) RoundTrip(*http.Request) (*http.Response, error) { + return nil, errFakeRoundTripperCalled +} + +func TestNewKeyFailoverTransport(t *testing.T) { + t.Parallel() + + pool, err := keypool.New([]string{"k0"}, quartz.NewMock(t)) + require.NoError(t, err) + + tests := []struct { + name string + // Constructor input. + config keypool.KeyFailoverConfig + // Whether the constructor returns inner unchanged. + expectSame bool + }{ + { + // Pool is nil: failover is disabled, inner is returned unchanged. + name: "pool_nil_returns_inner", + config: keypool.KeyFailoverConfig{}, + expectSame: true, + }, + { + // Pool is set: inner is wrapped in a key-failover transport. + name: "pool_set_returns_wrapper", + config: keypool.KeyFailoverConfig{Pool: pool}, + expectSame: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + inner := &fakeRoundTripper{} + got := keypool.NewKeyFailoverTransport(inner, tc.config) + + if tc.expectSame { + assert.Same(t, inner, got) + } else { + assert.NotSame(t, inner, got) + } + }) + } +} diff --git a/aibridge/passthrough.go b/aibridge/passthrough.go index 3fcb1b34dd..0dc6beb480 100644 --- a/aibridge/passthrough.go +++ b/aibridge/passthrough.go @@ -13,6 +13,7 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge/intercept/apidump" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/metrics" "github.com/coder/coder/v2/aibridge/provider" "github.com/coder/coder/v2/aibridge/tracing" @@ -41,13 +42,17 @@ func newPassthroughRouter(prov provider.Provider, logger slog.Logger, m *metrics ExpectContinueTimeout: 1 * time.Second, } - // Build a reverse proxy to the upstream, reused across all requests for this provider. - // All request modifications happen in Rewrite. + // Build the passthrough proxy, reused across all requests for this provider. + // Rewrite sets proxy headers. For centralized requests, KeyFailoverTransport + // handles auth and failover. BYOK requests pass through. proxy := &httputil.ReverseProxy{ Rewrite: func(pr *httputil.ProxyRequest) { - rewritePassthroughRequest(pr, provBaseURL, prov) + rewritePassthroughRequest(pr, provBaseURL) }, - Transport: apidump.NewPassthroughMiddleware(t, prov.APIDumpDir(), prov.Name(), logger, quartz.NewReal()), + Transport: keypool.NewKeyFailoverTransport( + apidump.NewPassthroughMiddleware(t, prov.APIDumpDir(), prov.Name(), logger, quartz.NewReal()), + prov.KeyFailoverConfig(logger), + ), ErrorHandler: func(rw http.ResponseWriter, req *http.Request, e error) { logger.Warn(req.Context(), "reverse proxy error", slog.Error(e), slog.F("path", req.URL.Path)) http.Error(rw, "upstream proxy error", http.StatusBadGateway) @@ -67,8 +72,8 @@ func newPassthroughRouter(prov provider.Provider, logger slog.Logger, m *metrics } // rewritePassthroughRequest configures the outbound request for the upstream and -// applies proxy headers and provider auth. -func rewritePassthroughRequest(pr *httputil.ProxyRequest, provBaseURL *url.URL, prov provider.Provider) { +// applies proxy headers. +func rewritePassthroughRequest(pr *httputil.ProxyRequest, provBaseURL *url.URL) { pr.SetURL(provBaseURL) // Rewrite sets "X-Forwarded-For" to just last hop (clients IP address). @@ -87,9 +92,6 @@ func rewritePassthroughRequest(pr *httputil.ProxyRequest, provBaseURL *url.URL, if _, ok := pr.Out.Header["User-Agent"]; !ok { pr.Out.Header.Set("User-Agent", "aibridge") // TODO: use build tag. } - - // Inject provider auth. - prov.InjectAuthHeader(&pr.Out.Header) } // newInvalidBaseURLHandler returns a handler that always returns 502 diff --git a/aibridge/passthrough_test.go b/aibridge/passthrough_test.go index 33a8c62b6b..6679feb932 100644 --- a/aibridge/passthrough_test.go +++ b/aibridge/passthrough_test.go @@ -2,20 +2,27 @@ package aibridge //nolint:testpackage // tests unexported newPassthroughRouter import ( "crypto/tls" + "io" "maps" "net" "net/http" "net/http/httptest" "net/http/httputil" "net/url" + "strings" "sync/atomic" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "go.opentelemetry.io/otel" "cdr.dev/slog/v3/sloggers/slogtest" + "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/provider" + "github.com/coder/quartz" ) var testTracer = otel.Tracer("bridge_test") @@ -171,25 +178,6 @@ func TestRewritePassthroughRequest(t *testing.T) { "User-Agent": {"custom-agent/1.0"}, }, }, - { - name: "injects_auth_header", - reqPath: "http://client-host/chat", - reqRemoteAddr: "1.1.1.1:1111", - provider: &testutil.MockProvider{ - URL: "https://upstream-host/base", - InjectAuthHeaderFunc: func(h *http.Header) { - h.Set("Authorization", "Bearer test-token") - }, - }, - expectURL: "https://upstream-host/base/chat", - expectHeaders: http.Header{ - "X-Forwarded-Host": {"client-host"}, - "X-Forwarded-Proto": {"http"}, - "X-Forwarded-For": {"1.1.1.1"}, - "User-Agent": {"aibridge"}, - "Authorization": {"Bearer test-token"}, - }, - }, { name: "appends_remote_addr_to_existing_forwarded_for_chain", reqPath: "http://client-host/chat", @@ -260,7 +248,7 @@ func TestRewritePassthroughRequest(t *testing.T) { Out: r.Clone(r.Context()), } - rewritePassthroughRequest(pr, provBaseURL, tc.provider) + rewritePassthroughRequest(pr, provBaseURL) assert.Equal(t, tc.expectURL, pr.Out.URL.String()) assert.Equal(t, "", pr.Out.Host) @@ -301,3 +289,285 @@ func TestPassthroughRouterReusesProxyInstance(t *testing.T) { assert.EqualValues(t, 1, newConnections.Load()) } + +// TestPassthrough_KeyFailover exercises the KeyFailoverTransport +// end-to-end through the passthrough proxy, parameterised over +// providers (anthropic, openai). Each scenario asserts the upstream +// request count, the response status and Retry-After, and the final +// pool state. +func TestPassthrough_KeyFailover(t *testing.T) { + t.Parallel() + + type upstreamResponse struct { + statusCode int + body string + headers map[string]string + } + + const ( + rateLimitBody = `{"error":"rate"}` + authErrorBody = `{"error":"unauthorized"}` + serverErrorBody = `{"error":"server"}` + successBody = `{"data":[]}` + ) + + // providers parameterises the table over the two providers + // that support key failover. Each entry encapsulates the + // provider-specific bits the test needs: how the mock upstream + // extracts the key from the request, how a BYOK request sets + // it, and how the provider is constructed for a given pool. + providers := []struct { + name string + extractKey func(*http.Request) string + setBYOK func(*http.Request, string) + newProvider func(baseURL string, pool *keypool.Pool) provider.Provider + }{ + { + name: "anthropic", + extractKey: func(r *http.Request) string { + return r.Header.Get("X-Api-Key") + }, + setBYOK: func(r *http.Request, key string) { + r.Header.Set("X-Api-Key", key) + }, + newProvider: func(baseURL string, pool *keypool.Pool) provider.Provider { + return provider.NewAnthropic(config.Anthropic{ + BaseURL: baseURL, + KeyPool: pool, + }, nil) + }, + }, + { + name: "openai", + extractKey: func(r *http.Request) string { + return strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + }, + setBYOK: func(r *http.Request, key string) { + r.Header.Set("Authorization", "Bearer "+key) + }, + newProvider: func(baseURL string, pool *keypool.Pool) provider.Provider { + cfg := config.OpenAI{BaseURL: baseURL} + if pool != nil { + cfg.KeyPool = pool + } + return provider.NewOpenAI(cfg) + }, + }, + } + + 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 API key value. + 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 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 user-supplied key returning 429. + // Then: 1 request, 429 forwarded as-is, 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"}, + body: rateLimitBody, + }, + }, + expectedRequestCount: 1, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "5", + }, + } + + for _, prov := range providers { + for _, tc := range tests { + t.Run(prov.name+"/"+tc.name, func(t *testing.T) { + t.Parallel() + + // Mock upstream: counts requests and returns + // scripted responses keyed by API key. 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[prov.extractKey(r)] + 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) + + 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) + } + + p := prov.newProvider(upstream.URL, pool) + // IgnoreErrors: MarkKey logs at ERROR level when a + // key is marked permanent (401/403); slogtest would + // otherwise fail those scenarios. + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + handler := newPassthroughRouter(p, logger, nil, testTracer) + + req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) + if tc.byokKey != "" { + prov.setBYOK(req, tc.byokKey) + } + w := httptest.NewRecorder() + handler.ServeHTTP(w, req) + + 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") + } + }) + } + } +} diff --git a/aibridge/provider/anthropic.go b/aibridge/provider/anthropic.go index 48f376fbce..64128cee5f 100644 --- a/aibridge/provider/anthropic.go +++ b/aibridge/provider/anthropic.go @@ -1,6 +1,7 @@ package provider import ( + "context" "fmt" "io" "net/http" @@ -11,6 +12,7 @@ import ( "go.opentelemetry.io/otel/trace" "golang.org/x/xerrors" + "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge/circuitbreaker" "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/intercept" @@ -219,6 +221,25 @@ func (p *Anthropic) InjectAuthHeader(headers *http.Header) { } } +func (p *Anthropic) KeyFailoverConfig(logger slog.Logger) keypool.KeyFailoverConfig { + name := p.Name() + return keypool.KeyFailoverConfig{ + Pool: p.cfg.KeyPool, + IsBYOK: func(r *http.Request) bool { + return r.Header.Get("X-Api-Key") != "" || r.Header.Get("Authorization") != "" + }, + InjectAuthKey: func(h *http.Header, key string) { + h.Set("X-Api-Key", key) + }, + MarkKey: func(ctx context.Context, key *keypool.Key, resp *http.Response) bool { + return keypool.MarkKeyOnStatus(ctx, key, resp, logger, name) + }, + BuildExhaustedResponse: func(err error) *http.Response { + return messages.ProcessKeyPoolError(err).ToResponse() + }, + } +} + func (p *Anthropic) CircuitBreakerConfig() *config.CircuitBreaker { return p.cfg.CircuitBreaker } diff --git a/aibridge/provider/copilot.go b/aibridge/provider/copilot.go index 9557b05d44..b68513ecec 100644 --- a/aibridge/provider/copilot.go +++ b/aibridge/provider/copilot.go @@ -12,10 +12,12 @@ import ( "go.opentelemetry.io/otel/trace" "golang.org/x/xerrors" + "cdr.dev/slog/v3" "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/intercept/responses" + "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/tracing" "github.com/coder/coder/v2/aibridge/utils" ) @@ -111,6 +113,12 @@ func (*Copilot) AuthHeader() string { // The original Authorization header flows through untouched from the client. func (*Copilot) InjectAuthHeader(_ *http.Header) {} +// KeyFailoverConfig returns a config with a nil Pool, which makes +// the KeyFailoverTransport short-circuit. Copilot is always BYOK. +func (*Copilot) KeyFailoverConfig(_ slog.Logger) keypool.KeyFailoverConfig { + return keypool.KeyFailoverConfig{} +} + func (p *Copilot) CircuitBreakerConfig() *config.CircuitBreaker { return p.circuitBreaker } diff --git a/aibridge/provider/openai.go b/aibridge/provider/openai.go index c670626b11..894dda7194 100644 --- a/aibridge/provider/openai.go +++ b/aibridge/provider/openai.go @@ -1,6 +1,7 @@ package provider import ( + "context" "encoding/json" "fmt" "io" @@ -12,6 +13,7 @@ import ( "go.opentelemetry.io/otel/trace" "golang.org/x/xerrors" + "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/intercept" "github.com/coder/coder/v2/aibridge/intercept/chatcompletions" @@ -218,6 +220,25 @@ func (p *OpenAI) InjectAuthHeader(headers *http.Header) { } } +func (p *OpenAI) KeyFailoverConfig(logger slog.Logger) keypool.KeyFailoverConfig { + name := p.Name() + return keypool.KeyFailoverConfig{ + Pool: p.cfg.KeyPool, + IsBYOK: func(r *http.Request) bool { + return r.Header.Get("Authorization") != "" + }, + InjectAuthKey: func(h *http.Header, key string) { + h.Set("Authorization", "Bearer "+key) + }, + MarkKey: func(ctx context.Context, key *keypool.Key, resp *http.Response) bool { + return keypool.MarkKeyOnStatus(ctx, key, resp, logger, name) + }, + BuildExhaustedResponse: func(err error) *http.Response { + return chatcompletions.ProcessKeyPoolError(err).ToResponse() + }, + } +} + func (p *OpenAI) CircuitBreakerConfig() *config.CircuitBreaker { return p.circuitBreaker } diff --git a/aibridge/provider/provider.go b/aibridge/provider/provider.go index cd09b6fc31..587dfd85ce 100644 --- a/aibridge/provider/provider.go +++ b/aibridge/provider/provider.go @@ -6,8 +6,10 @@ import ( "go.opentelemetry.io/otel/trace" "golang.org/x/xerrors" + "cdr.dev/slog/v3" "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/intercept" + "github.com/coder/coder/v2/aibridge/keypool" ) var ErrUnknownRoute = xerrors.New("unknown route") @@ -76,7 +78,12 @@ type Provider interface { // token in. AuthHeader() string // InjectAuthHeader allows [Provider]s to set its authentication header. + // TODO(ssncferreira): remove. Auth is now applied per-attempt by + // KeyFailoverTransport (see [Provider.KeyFailoverConfig]). InjectAuthHeader(*http.Header) + // KeyFailoverConfig returns the per-provider configuration for + // automatic key failover on passthrough routes. + KeyFailoverConfig(logger slog.Logger) keypool.KeyFailoverConfig // CircuitBreakerConfig returns the circuit breaker configuration for the provider. CircuitBreakerConfig() *config.CircuitBreaker diff --git a/aibridge/utils/http.go b/aibridge/utils/http.go new file mode 100644 index 0000000000..e41feb1f39 --- /dev/null +++ b/aibridge/utils/http.go @@ -0,0 +1,34 @@ +package utils + +import ( + "bytes" + "fmt" + "io" + "math" + "net/http" + "strconv" + "time" +) + +// NewJSONErrorResponse builds an *http.Response with a JSON body +// and optional Retry-After header. Used to synthesize bridge-side +// error responses (e.g. key-pool exhaustion, marshaling +// fallbacks). Retry-After is set to whole seconds (rounded up) +// when retryAfter is positive, and omitted otherwise. +func NewJSONErrorResponse(status int, retryAfter time.Duration, body []byte) *http.Response { + h := http.Header{} + h.Set("Content-Type", "application/json") + if retryAfter > 0 { + h.Set("Retry-After", strconv.Itoa(int(math.Ceil(retryAfter.Seconds())))) + } + return &http.Response{ + Status: fmt.Sprintf("%d %s", status, http.StatusText(status)), + StatusCode: status, + Proto: "HTTP/1.1", + ProtoMajor: 1, + ProtoMinor: 1, + Header: h, + Body: io.NopCloser(bytes.NewReader(body)), + ContentLength: int64(len(body)), + } +} diff --git a/aibridge/utils/http_test.go b/aibridge/utils/http_test.go new file mode 100644 index 0000000000..337c42f12f --- /dev/null +++ b/aibridge/utils/http_test.go @@ -0,0 +1,91 @@ +package utils_test + +import ( + "io" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/aibridge/utils" +) + +func TestNewJSONErrorResponse(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + status int + retryAfter time.Duration + body []byte + // Empty string means the header should be absent. + expectRetryAfter string + }{ + { + // Permanent exhaustion: 502 with no Retry-After. + name: "permanent_no_retry_after", + status: http.StatusBadGateway, + retryAfter: 0, + body: []byte(`{"error":"permanent"}`), + expectRetryAfter: "", + }, + { + // Transient exhaustion with zero retryAfter: no Retry-After. + name: "transient_no_retry_after", + status: http.StatusTooManyRequests, + retryAfter: 0, + body: []byte(`{"error":"rate"}`), + expectRetryAfter: "", + }, + { + // Transient exhaustion: 429 with Retry-After in seconds. + name: "transient_with_retry_after", + status: http.StatusTooManyRequests, + retryAfter: 60 * time.Second, + body: []byte(`{"error":"rate"}`), + expectRetryAfter: "60", + }, + { + // Transient exhaustion with negative retryAfter: Retry-After header omitted. + name: "transient_negative_retry_after", + status: http.StatusTooManyRequests, + retryAfter: -1 * time.Second, + body: []byte(`{"error":"rate"}`), + expectRetryAfter: "", + }, + { + // Transient exhaustion with 500ms retryAfter rounds up to Retry-After: 1. + name: "transient_under_one_second_rounds_up", + status: http.StatusTooManyRequests, + retryAfter: 500 * time.Millisecond, + body: []byte(`{"error":"rate"}`), + expectRetryAfter: "1", + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + resp := utils.NewJSONErrorResponse(tc.status, tc.retryAfter, tc.body) + require.NotNil(t, resp) + + assert.Equal(t, tc.status, resp.StatusCode) + assert.Equal(t, "application/json", resp.Header.Get("Content-Type")) + assert.Equal(t, int64(len(tc.body)), resp.ContentLength) + + if tc.expectRetryAfter == "" { + assert.Empty(t, resp.Header.Get("Retry-After")) + } else { + assert.Equal(t, tc.expectRetryAfter, resp.Header.Get("Retry-After")) + } + + body, err := io.ReadAll(resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + assert.Equal(t, tc.body, body) + }) + } +}