opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(aibridge): apply key pool failover follow-ups (#26130)

Browse Source 
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
This commit is contained in:
Susana Ferreira
2026-06-19 09:01:39 +01:00
committed by GitHub
parent 19aa9f5616
commit fec21e1a28
10 changed files with 290 additions and 257 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aibridge/intercept/keyfailover_test.go
+2 -11
View File
@@ -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")
aibridge/intercept/messages/base.go
+3
View File
@@ -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(
aibridge/intercept/messages/base_internal_test.go
+8
View File
@@ -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)
aibridge/intercept/messages/blocking.go
+15 -8
View File
@@ -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
aibridge/intercept/messages/streaming.go
+8 -7
View File
@@ -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...)
}
aibridge/intercept/openai_errors.go
+3
View File
@@ -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(
aibridge/intercept/openai_errors_test.go
+8
View File
@@ -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)
aibridge/internal/integrationtest/keypool_failover_internal_test.go
+164 -112
View File
@@ -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")
})
}
}
aibridge/internal/testutil/mockupstream.go
+11
View File
@@ -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
aibridge/passthrough_internal_test.go
+68 -119
View File
@@ -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())
}
})
}