From fec21e1a28423c6a29c378207281c1dfcfa57500 Mon Sep 17 00:00:00 2001 From: Susana Ferreira Date: Fri, 19 Jun 2026 09:01:39 +0100 Subject: [PATCH] refactor(aibridge): apply key pool failover follow-ups (#26130) Applies follow-ups from the key pool failover work: - Add a test verifying key pool state is shared across bridged and passthrough routes. - Refactor the key failover and passthrough tests to use the shared `MockUpstream` helper. - Simplify how the request body option is passed through the Anthropic messages interceptor. - Make `ResponseErrorFromKeyPool` nil-safe and cover it with a test. Closes: https://linear.app/codercom/issue/AIGOV-398/small-follow-up-cleanups-for-key-failover > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira --- aibridge/intercept/keyfailover_test.go | 13 +- aibridge/intercept/messages/base.go | 3 + .../intercept/messages/base_internal_test.go | 8 + aibridge/intercept/messages/blocking.go | 23 +- aibridge/intercept/messages/streaming.go | 15 +- aibridge/intercept/openai_errors.go | 3 + aibridge/intercept/openai_errors_test.go | 8 + .../keypool_failover_internal_test.go | 276 +++++++++++------- aibridge/internal/testutil/mockupstream.go | 11 + aibridge/passthrough_internal_test.go | 187 +++++------- 10 files changed, 290 insertions(+), 257 deletions(-) diff --git a/aibridge/intercept/keyfailover_test.go b/aibridge/intercept/keyfailover_test.go index a20b14d343..700e3429e7 100644 --- a/aibridge/intercept/keyfailover_test.go +++ b/aibridge/intercept/keyfailover_test.go @@ -62,15 +62,6 @@ type interceptorCase struct { newInterceptor func(t *testing.T, streaming bool, upstreamURL string, reqBody []byte, pool *keypool.Pool, byokKey string) intercept.Interceptor } -// keyFromHeader reads the API key an upstream request carried in the named auth -// header. -func keyFromHeader(name string, h http.Header) string { - if name == "Authorization" { - return utils.ExtractBearerToken(h.Get(name)) - } - return h.Get(name) -} - // interceptorCases is the set of interceptors the failover tests run against, // one entry per supported API. var interceptorCases = []interceptorCase{ @@ -371,7 +362,7 @@ func TestInterception_KeyFailover(t *testing.T) { var seenKeys []string for _, r := range upstream.ReceivedRequests() { - seenKeys = append(seenKeys, keyFromHeader(ic.authHeader, r.Header)) + seenKeys = append(seenKeys, testutil.KeyFromHeader(ic.authHeader, r.Header)) } assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys") @@ -546,7 +537,7 @@ func TestInterception_AgenticLoopFailover(t *testing.T) { var seenKeys []string for _, r := range upstream.ReceivedRequests() { - seenKeys = append(seenKeys, keyFromHeader(ic.authHeader, r.Header)) + seenKeys = append(seenKeys, testutil.KeyFromHeader(ic.authHeader, r.Header)) } assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys") diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index b7b414dc8d..dff58c1cad 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -578,6 +578,9 @@ func (i *interceptionBase) markKeyOnError(ctx context.Context, key *keypool.Key, // ResponseErrorFromKeyPool translates a *keypool.Error into // a developer-facing ResponseError shaped for the Anthropic API. func ResponseErrorFromKeyPool(keyPoolErr *keypool.Error) *ResponseError { + if keyPoolErr == nil { + return nil + } switch keyPoolErr.Kind { case keypool.ErrorKindPermanent: return newResponseError( diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index 41133836fa..5a4c5e6f30 100644 --- a/aibridge/intercept/messages/base_internal_test.go +++ b/aibridge/intercept/messages/base_internal_test.go @@ -1041,6 +1041,10 @@ func TestResponseErrorFromKeyPool(t *testing.T) { expectedStatus int expectedRetryAfter time.Duration }{ + { + name: "nil_returns_nil", + keyPoolErr: nil, + }, { // Rate-limited with no cooldown: 429, no Retry-After. name: "rate_limited_zero_retry_after", @@ -1067,6 +1071,10 @@ func TestResponseErrorFromKeyPool(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() got := ResponseErrorFromKeyPool(tc.keyPoolErr) + if tc.keyPoolErr == nil { + assert.Nil(t, got) + return + } require.NotNil(t, got) assert.Equal(t, tc.expectedStatus, got.StatusCode) assert.Equal(t, tc.expectedRetryAfter, got.RetryAfter) diff --git a/aibridge/intercept/messages/blocking.go b/aibridge/intercept/messages/blocking.go index 8ee118b5da..e6b5a8a596 100644 --- a/aibridge/intercept/messages/blocking.go +++ b/aibridge/intercept/messages/blocking.go @@ -112,8 +112,14 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req for { // TODO add outer loop span (https://github.com/coder/aibridge/issues/67) + + // Rebuilt per iteration: i.reqPayload mutates when an agentic + // continuation appends tool results, so withBody must reflect + // the latest payload on every upstream call. + callOpts := []option.RequestOption{i.withBody()} + var keyAttempts int - resp, keyAttempts, err = i.newMessage(ctx, svc) + resp, keyAttempts, err = i.newMessage(ctx, svc, callOpts) totalKeyAttempts += keyAttempts if err != nil { if eventstream.IsConnError(err) { @@ -353,20 +359,19 @@ func (i *BlockingInterception) ProcessRequest(w http.ResponseWriter, r *http.Req // number of key attempts made for this call, and any error. A centralized key // pool fails over across keys, while BYOK and Bedrock authenticate with a // single, fixed credential baked into svc, so they make one attempt. -func (i *BlockingInterception) newMessage(ctx context.Context, svc anthropic.MessageService) (*anthropic.Message, int, error) { +func (i *BlockingInterception) newMessage(ctx context.Context, svc anthropic.MessageService, opts []option.RequestOption) (*anthropic.Message, int, error) { if cp, ok := intercept.AsCentralizedPool(i.cred); ok { - return i.newMessageWithKeyFailover(ctx, svc, cp) + return i.newMessageWithKeyFailover(ctx, svc, cp, opts) } - msg, err := i.newMessageWithKey(intercept.WithCredentialInfo(ctx, i.cred), svc) + msg, err := i.newMessageWithKey(intercept.WithCredentialInfo(ctx, i.cred), svc, opts...) return msg, 0, err } // newMessageWithKey performs a single upstream call. -func (i *BlockingInterception) newMessageWithKey(ctx context.Context, svc anthropic.MessageService, extraOpts ...option.RequestOption) (_ *anthropic.Message, outErr error) { +func (i *BlockingInterception) newMessageWithKey(ctx context.Context, svc anthropic.MessageService, opts ...option.RequestOption) (_ *anthropic.Message, outErr error) { _, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer tracing.EndSpanErr(span, &outErr) - opts := append([]option.RequestOption{i.withBody()}, extraOpts...) return svc.New(ctx, anthropic.MessageNewParams{}, opts...) } @@ -375,7 +380,7 @@ func (i *BlockingInterception) newMessageWithKey(ctx context.Context, svc anthro // 429 and permanent on 401/403. Errors that aren't key-specific don't trigger // failover and are returned to the caller. It returns the upstream message, // the number of key attempts made for this call, and any error. -func (i *BlockingInterception) newMessageWithKeyFailover(ctx context.Context, svc anthropic.MessageService, cp *intercept.CentralizedPool) (*anthropic.Message, int, error) { +func (i *BlockingInterception) newMessageWithKeyFailover(ctx context.Context, svc anthropic.MessageService, cp *intercept.CentralizedPool, opts []option.RequestOption) (*anthropic.Message, int, error) { walker := cp.Pool.Walker() for { key, keyPoolErr := cp.NextKey(walker) @@ -385,12 +390,14 @@ func (i *BlockingInterception) newMessageWithKeyFailover(ctx context.Context, sv ctx = intercept.WithCredentialInfo(ctx, i.cred) i.logger.Debug(ctx, "using centralized api key") - msg, err := i.newMessageWithKey(ctx, svc, + 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), ) + msg, err := i.newMessageWithKey(ctx, svc, requestOpts...) // Key-specific failure: try the next key. if i.markKeyOnError(ctx, key, err) { continue diff --git a/aibridge/intercept/messages/streaming.go b/aibridge/intercept/messages/streaming.go index 22d0dff720..442c4291b0 100644 --- a/aibridge/intercept/messages/streaming.go +++ b/aibridge/intercept/messages/streaming.go @@ -168,11 +168,13 @@ newStream: break } - // Per-iteration: a pool credential advances its failover walker. An - // iteration is either an agentic continuation or a failover retry after - // the previous key was marked. BYOK and Bedrock have no pool and run as - // a single attempt. - var streamOpts []option.RequestOption + // 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). A pool-less credential (BYOK, or pool-less + // centralized such as Bedrock) has no walker and runs as a + // single attempt. + streamOpts := []option.RequestOption{i.withBody()} var currentPoolKey *keypool.Key if cp, isPool := intercept.AsCentralizedPool(i.cred); isPool { walker := cp.Pool.Walker() @@ -686,10 +688,9 @@ func (*StreamingInterception) encodeForStream(payload []byte, typ string) []byte } // newStream traces svc.NewStreaming() call. -func (i *StreamingInterception) newStream(ctx context.Context, svc anthropic.MessageService, extraOpts ...option.RequestOption) *ssestream.Stream[anthropic.MessageStreamEventUnion] { +func (i *StreamingInterception) newStream(ctx context.Context, svc anthropic.MessageService, opts ...option.RequestOption) *ssestream.Stream[anthropic.MessageStreamEventUnion] { _, span := i.tracer.Start(ctx, "Intercept.ProcessRequest.Upstream", trace.WithAttributes(tracing.InterceptionAttributesFromContext(ctx)...)) defer span.End() - opts := append([]option.RequestOption{i.withBody()}, extraOpts...) return svc.NewStreaming(ctx, anthropic.MessageNewParams{}, opts...) } diff --git a/aibridge/intercept/openai_errors.go b/aibridge/intercept/openai_errors.go index faf2e19e3e..80266b1fad 100644 --- a/aibridge/intercept/openai_errors.go +++ b/aibridge/intercept/openai_errors.go @@ -73,6 +73,9 @@ func (e *ResponseError) ToResponse() *http.Response { // ResponseErrorFromKeyPool translates a *keypool.Error into // a developer-facing ResponseError shaped for the OpenAI API. func ResponseErrorFromKeyPool(keyPoolErr *keypool.Error) *ResponseError { + if keyPoolErr == nil { + return nil + } switch keyPoolErr.Kind { case keypool.ErrorKindPermanent: return NewResponseError( diff --git a/aibridge/intercept/openai_errors_test.go b/aibridge/intercept/openai_errors_test.go index 9b49c1e43a..92953c64ef 100644 --- a/aibridge/intercept/openai_errors_test.go +++ b/aibridge/intercept/openai_errors_test.go @@ -21,6 +21,10 @@ func TestResponseErrorFromKeyPool(t *testing.T) { expectedStatus int expectedRetryAfter time.Duration }{ + { + name: "nil_returns_nil", + keyPoolErr: nil, + }, { // Rate-limited with no cooldown: 429, no Retry-After. name: "rate_limited_zero_retry_after", @@ -47,6 +51,10 @@ func TestResponseErrorFromKeyPool(t *testing.T) { t.Run(tc.name, func(t *testing.T) { t.Parallel() got := intercept.ResponseErrorFromKeyPool(tc.keyPoolErr) + if tc.keyPoolErr == nil { + assert.Nil(t, got) + return + } require.NotNil(t, got) assert.Equal(t, tc.expectedStatus, got.StatusCode) assert.Equal(t, tc.expectedRetryAfter, got.RetryAfter) diff --git a/aibridge/internal/integrationtest/keypool_failover_internal_test.go b/aibridge/internal/integrationtest/keypool_failover_internal_test.go index f186fafa36..af92056c87 100644 --- a/aibridge/internal/integrationtest/keypool_failover_internal_test.go +++ b/aibridge/internal/integrationtest/keypool_failover_internal_test.go @@ -1,23 +1,20 @@ package integrationtest import ( - "fmt" "io" "net/http" - "net/http/httptest" - "sync" - "sync/atomic" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/tidwall/sjson" + "github.com/coder/coder/v2/aibridge" "github.com/coder/coder/v2/aibridge/config" "github.com/coder/coder/v2/aibridge/fixtures" + "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/coder/v2/aibridge/utils" "github.com/coder/quartz" ) @@ -30,39 +27,34 @@ func TestOpenAI_KeyFailover(t *testing.T) { t.Parallel() tests := []struct { - name string - fixture []byte - path string - streaming bool - successCType string + name string + fixture []byte + path string + streaming bool }{ { - name: "chatcompletions_blocking", - fixture: fixtures.OaiChatSimple, - path: pathOpenAIChatCompletions, - streaming: false, - successCType: "application/json", + name: "chatcompletions_blocking", + fixture: fixtures.OaiChatSimple, + path: pathOpenAIChatCompletions, + streaming: false, }, { - name: "chatcompletions_streaming", - fixture: fixtures.OaiChatSimple, - path: pathOpenAIChatCompletions, - streaming: true, - successCType: "text/event-stream", + name: "chatcompletions_streaming", + fixture: fixtures.OaiChatSimple, + path: pathOpenAIChatCompletions, + streaming: true, }, { - name: "responses_blocking", - fixture: fixtures.OaiResponsesBlockingSimple, - path: pathOpenAIResponses, - streaming: false, - successCType: "application/json", + name: "responses_blocking", + fixture: fixtures.OaiResponsesBlockingSimple, + path: pathOpenAIResponses, + streaming: false, }, { - name: "responses_streaming", - fixture: fixtures.OaiResponsesStreamingSimple, - path: pathOpenAIResponses, - streaming: true, - successCType: "text/event-stream", + name: "responses_streaming", + fixture: fixtures.OaiResponsesStreamingSimple, + path: pathOpenAIResponses, + streaming: true, }, } @@ -71,45 +63,18 @@ func TestOpenAI_KeyFailover(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(config.ProviderOpenAI, []string{"k0", "k1"}, quartz.NewMock(t), nil) 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) + // Sequential upstream responses: request 1 fails over + // from k0 to k1 (calls 1-2), and request 2 goes straight + // to k1 (call 3). + upstream := testutil.NewMockUpstream(t.Context(), t, + testutil.NewErrorResponse(http.StatusTooManyRequests, "60"), + testutil.NewFixtureResponse(fix), + testutil.NewFixtureResponse(fix), + ) bridgeServer := newBridgeTestServer(t.Context(), t, upstream.URL, withCustomProvider(provider.NewOpenAI(config.OpenAI{ @@ -137,10 +102,11 @@ func TestOpenAI_KeyFailover(t *testing.T) { require.NoError(t, resp.Body.Close()) require.Equal(t, http.StatusOK, resp.StatusCode) - seenKeysMu.Lock() - defer seenKeysMu.Unlock() + var seenKeys []string + for _, r := range upstream.ReceivedRequests() { + seenKeys = append(seenKeys, testutil.KeyFromHeader("Authorization", r.Header)) + } // 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. @@ -162,22 +128,16 @@ func TestAnthropic_KeyFailover(t *testing.T) { fix := fixtures.Parse(t, fixtures.AntSimple) tests := []struct { - name string - streaming bool - successBody []byte - successCType string + name string + streaming bool }{ { - name: "blocking", - streaming: false, - successBody: fix.NonStreaming(), - successCType: "application/json", + name: "blocking", + streaming: false, }, { - name: "streaming", - streaming: true, - successBody: fix.Streaming(), - successCType: "text/event-stream", + name: "streaming", + streaming: true, }, } @@ -188,35 +148,14 @@ func TestAnthropic_KeyFailover(t *testing.T) { pool, err := keypool.New(config.ProviderAnthropic, []string{"k0", "k1"}, quartz.NewMock(t), nil) 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 := r.Header.Get("X-Api-Key") - 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, `{"type":"error","error":{"type":"rate_limit_error","message":"rate limited"}}`) - case "k1": - w.Header().Set("Content-Type", tc.successCType) - w.WriteHeader(http.StatusOK) - _, _ = w.Write(tc.successBody) - default: - w.WriteHeader(http.StatusInternalServerError) - } - })) - t.Cleanup(upstream.Close) + // Sequential upstream responses: request 1 fails over + // from k0 to k1 (calls 1-2), and request 2 goes straight + // to k1 (call 3). + upstream := testutil.NewMockUpstream(t.Context(), t, + testutil.NewErrorResponse(http.StatusTooManyRequests, "60"), + testutil.NewFixtureResponse(fix), + testutil.NewFixtureResponse(fix), + ) bridgeServer := newBridgeTestServer(t.Context(), t, upstream.URL, withCustomProvider(provider.NewAnthropic(config.Anthropic{ @@ -244,10 +183,11 @@ func TestAnthropic_KeyFailover(t *testing.T) { require.NoError(t, resp.Body.Close()) require.Equal(t, http.StatusOK, resp.StatusCode) - seenKeysMu.Lock() - defer seenKeysMu.Unlock() + var seenKeys []string + for _, r := range upstream.ReceivedRequests() { + seenKeys = append(seenKeys, testutil.KeyFromHeader("X-Api-Key", r.Header)) + } // 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. @@ -258,3 +198,115 @@ func TestAnthropic_KeyFailover(t *testing.T) { }) } } + +// TestKeyPool_StateSharing verifies that a key marked unavailable +// by a bridged route is observed in the same state by every other +// route that shares the provider's pool, including other bridged +// routes and passthrough routes. Both paths walk the same +// *keypool.Pool, so state set in one must be visible to all. +func TestKeyPool_StateSharing(t *testing.T) { + t.Parallel() + + // Parse fixtures once so table rows can reference them. + fixAnt := fixtures.Parse(t, fixtures.AntSimple) + fixOaiChat := fixtures.Parse(t, fixtures.OaiChatSimple) + fixOaiResp := fixtures.Parse(t, fixtures.OaiResponsesBlockingSimple) + + type requestStep struct { + method string + path string + body []byte // nil for GET /models passthrough route. + } + + tests := []struct { + name string + providerName string + newProvider func(baseURL string, pool *keypool.Pool) aibridge.Provider + upstreamResponses []testutil.UpstreamResponse + requests []requestStep + expectedSeenKeys []string + }{ + { + // Bridged route fails over k0->k1 (calls 1-2), then + // the passthrough route hits k1 directly (call 3). + name: "anthropic", + providerName: config.ProviderAnthropic, + newProvider: func(baseURL string, pool *keypool.Pool) aibridge.Provider { + return provider.NewAnthropic(config.Anthropic{BaseURL: baseURL, KeyPool: pool}, nil) + }, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusTooManyRequests, "60"), + testutil.NewFixtureResponse(fixAnt), + {Blocking: []byte("{}")}, + }, + requests: []requestStep{ + {method: http.MethodPost, path: pathAnthropicMessages, body: fixAnt.Request()}, + {method: http.MethodGet, path: "/anthropic/v1/models"}, + }, + expectedSeenKeys: []string{"k0", "k1", "k1"}, + }, + { + // Bridged chat completions route fails over k0->k1 + // (calls 1-2), bridged responses route hits k1 + // directly (call 3), then the passthrough route hits + // k1 directly (call 4). + name: "openai", + providerName: config.ProviderOpenAI, + newProvider: func(baseURL string, pool *keypool.Pool) aibridge.Provider { + return provider.NewOpenAI(config.OpenAI{BaseURL: baseURL, KeyPool: pool}) + }, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusTooManyRequests, "60"), + testutil.NewFixtureResponse(fixOaiChat), + testutil.NewFixtureResponse(fixOaiResp), + {Blocking: []byte("{}")}, + }, + requests: []requestStep{ + {method: http.MethodPost, path: pathOpenAIChatCompletions, body: fixOaiChat.Request()}, + {method: http.MethodPost, path: pathOpenAIResponses, body: fixOaiResp.Request()}, + {method: http.MethodGet, path: "/openai/v1/models"}, + }, + expectedSeenKeys: []string{"k0", "k1", "k1", "k1"}, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + pool, err := keypool.New(tc.providerName, []string{"k0", "k1"}, quartz.NewMock(t), nil) + require.NoError(t, err) + + upstream := testutil.NewMockUpstream(t.Context(), t, tc.upstreamResponses...) + + prov := tc.newProvider(upstream.URL, pool) + bridgeServer := newBridgeTestServer(t.Context(), t, upstream.URL, + withCustomProvider(prov), + ) + + // Every request returns 200 to the client: the first + // fails over from k0 (429) to k1 (200) and subsequent + // requests skip the now-temporary k0 and hit k1 + // directly. + for _, req := range tc.requests { + resp, err := bridgeServer.makeRequest(t, req.method, req.path, req.body) + require.NoError(t, err) + _, _ = io.Copy(io.Discard, resp.Body) + require.NoError(t, resp.Body.Close()) + require.Equal(t, http.StatusOK, resp.StatusCode) + } + + var seenKeys []string + for _, r := range upstream.ReceivedRequests() { + seenKeys = append(seenKeys, testutil.KeyFromHeader(prov.AuthHeader(), r.Header)) + } + assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys") + + // Pool state persists across bridged and passthrough routes. + assert.Equal(t, []keypool.KeyState{ + keypool.KeyStateTemporary, + keypool.KeyStateValid, + }, pool.PoolState(), "key states") + }) + } +} diff --git a/aibridge/internal/testutil/mockupstream.go b/aibridge/internal/testutil/mockupstream.go index 242bfde345..c73de50df1 100644 --- a/aibridge/internal/testutil/mockupstream.go +++ b/aibridge/internal/testutil/mockupstream.go @@ -23,6 +23,7 @@ import ( "github.com/coder/coder/v2/aibridge/fixtures" "github.com/coder/coder/v2/aibridge/intercept/eventstream" + "github.com/coder/coder/v2/aibridge/utils" ) // UpstreamResponse defines a single response that MockUpstream will replay @@ -83,6 +84,16 @@ func NewErrorResponse(status int, retryAfter string) UpstreamResponse { return UpstreamResponse{Streaming: rawBytes, Blocking: rawBytes} } +// KeyFromHeader reads the API key an upstream request carried in the named +// auth header. Authorization headers are unwrapped from their "Bearer " +// prefix, and other headers are returned verbatim. +func KeyFromHeader(name string, h http.Header) string { + if name == "Authorization" { + return utils.ExtractBearerToken(h.Get(name)) + } + return h.Get(name) +} + // ReceivedRequest captures the details of a single request handled by MockUpstream. type ReceivedRequest struct { Method string diff --git a/aibridge/passthrough_internal_test.go b/aibridge/passthrough_internal_test.go index c095281bf3..5a0aeb3cb6 100644 --- a/aibridge/passthrough_internal_test.go +++ b/aibridge/passthrough_internal_test.go @@ -2,14 +2,12 @@ package aibridge import ( "crypto/tls" - "io" "maps" "net" "net/http" "net/http/httptest" "net/http/httputil" "net/url" - "strings" "sync/atomic" "testing" @@ -295,42 +293,25 @@ func TestPassthroughRouterReusesProxyInstance(t *testing.T) { // 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. +// providers (anthropic, openai, copilot). Each scenario asserts the +// response status and Retry-After, the keys the upstream actually +// saw, 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 providers exposed // to the failover transport. 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. + // provider-specific bits the test needs: how a BYOK request + // sets its auth header and how the provider is constructed for + // a given pool. providers := []struct { name string byokOnly bool - 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) }, @@ -343,9 +324,6 @@ func TestPassthrough_KeyFailover(t *testing.T) { }, { 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) }, @@ -362,9 +340,6 @@ func TestPassthrough_KeyFailover(t *testing.T) { { name: "copilot", byokOnly: true, - 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) }, @@ -380,11 +355,14 @@ func TestPassthrough_KeyFailover(t *testing.T) { 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 + // Sequential upstream responses replayed by MockUpstream + // in call order. MockUpstream's strict mode asserts the + // upstream call count matches len(upstreamResponses). + upstreamResponses []testutil.UpstreamResponse + // Expected keys the upstream actually saw, in call order. + expectedSeenKeys []string + expectedStatusCode int + expectedRetryAfter string // Expected key states after the request, by index in keys. expectedKeyStates []keypool.KeyState // Expected key_pool_state_transitions_total counts by reason. @@ -397,28 +375,24 @@ func TestPassthrough_KeyFailover(t *testing.T) { // 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}, + upstreamResponses: []testutil.UpstreamResponse{ + {Blocking: []byte("{}")}, }, - expectedRequestCount: 1, - expectedStatusCode: http.StatusOK, - expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid}, + expectedSeenKeys: []string{"k0"}, + 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}, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusTooManyRequests, "5"), + {Blocking: []byte("{}")}, }, - expectedRequestCount: 2, - expectedStatusCode: http.StatusOK, + expectedSeenKeys: []string{"k0", "k1"}, + expectedStatusCode: http.StatusOK, expectedKeyStates: []keypool.KeyState{ keypool.KeyStateTemporary, keypool.KeyStateValid, @@ -430,12 +404,12 @@ func TestPassthrough_KeyFailover(t *testing.T) { // 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}, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusUnauthorized, ""), + {Blocking: []byte("{}")}, }, - expectedRequestCount: 2, - expectedStatusCode: http.StatusOK, + expectedSeenKeys: []string{"k0", "k1"}, + expectedStatusCode: http.StatusOK, expectedKeyStates: []keypool.KeyState{ keypool.KeyStatePermanent, keypool.KeyStateValid, @@ -447,12 +421,12 @@ func TestPassthrough_KeyFailover(t *testing.T) { // 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}, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusForbidden, ""), + {Blocking: []byte("{}")}, }, - expectedRequestCount: 2, - expectedStatusCode: http.StatusOK, + expectedSeenKeys: []string{"k0", "k1"}, + expectedStatusCode: http.StatusOK, expectedKeyStates: []keypool.KeyState{ keypool.KeyStatePermanent, keypool.KeyStateValid, @@ -465,26 +439,14 @@ func TestPassthrough_KeyFailover(t *testing.T) { // 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, - }, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusTooManyRequests, "5"), + testutil.NewErrorResponse(http.StatusTooManyRequests, "3"), + testutil.NewErrorResponse(http.StatusTooManyRequests, "10"), }, - expectedRequestCount: 3, - expectedStatusCode: http.StatusTooManyRequests, - expectedRetryAfter: "3", + expectedSeenKeys: []string{"k0", "k1", "k2"}, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "3", expectedKeyStates: []keypool.KeyState{ keypool.KeyStateTemporary, keypool.KeyStateTemporary, @@ -498,12 +460,12 @@ func TestPassthrough_KeyFailover(t *testing.T) { // 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}, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusUnauthorized, ""), + testutil.NewErrorResponse(http.StatusUnauthorized, ""), }, - expectedRequestCount: 2, - expectedStatusCode: http.StatusBadGateway, + expectedSeenKeys: []string{"k0", "k1"}, + expectedStatusCode: http.StatusBadGateway, expectedKeyStates: []keypool.KeyState{ keypool.KeyStatePermanent, keypool.KeyStatePermanent, @@ -516,11 +478,11 @@ func TestPassthrough_KeyFailover(t *testing.T) { // 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}, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusInternalServerError, ""), }, - expectedRequestCount: 1, - expectedStatusCode: http.StatusInternalServerError, + expectedSeenKeys: []string{"k0"}, + expectedStatusCode: http.StatusInternalServerError, expectedKeyStates: []keypool.KeyState{ keypool.KeyStateValid, keypool.KeyStateValid, @@ -531,16 +493,12 @@ func TestPassthrough_KeyFailover(t *testing.T) { // 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, - }, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusTooManyRequests, "5"), }, - expectedRequestCount: 1, - expectedStatusCode: http.StatusTooManyRequests, - expectedRetryAfter: "5", + expectedSeenKeys: []string{"user-byok"}, + expectedStatusCode: http.StatusTooManyRequests, + expectedRetryAfter: "5", }, } @@ -554,26 +512,11 @@ func TestPassthrough_KeyFailover(t *testing.T) { 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) + // MockUpstream replays the scripted responses in + // call order. Strict mode fails the test if the + // upstream sees a different number of requests + // than tc.upstreamResponses describes. + upstream := testutil.NewMockUpstream(t.Context(), t, tc.upstreamResponses...) reg := prometheus.NewRegistry() m := NewMetrics(reg) @@ -599,9 +542,15 @@ func TestPassthrough_KeyFailover(t *testing.T) { 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") + + var seenKeys []string + for _, r := range upstream.ReceivedRequests() { + seenKeys = append(seenKeys, testutil.KeyFromHeader(p.AuthHeader(), r.Header)) + } + assert.Equal(t, tc.expectedSeenKeys, seenKeys, "seen keys") + if pool != nil { assert.Equal(t, tc.expectedKeyStates, pool.PoolState(), "key states") @@ -627,7 +576,7 @@ func TestPassthrough_KeyFailover(t *testing.T) { hist := promhelp.HistogramValue(t, reg, "key_pool_failover_attempts", prometheus.Labels{"provider": "test"}) require.NotNil(t, hist) assert.Equal(t, uint64(1), hist.GetSampleCount()) - assert.Equal(t, float64(tc.expectedRequestCount), hist.GetSampleSum()) + assert.Equal(t, float64(len(tc.upstreamResponses)), hist.GetSampleSum()) } }) }