mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add automatic key failover for AI Bridge Anthropic (#24836)
## Description Adds automatic key failover for centralized Anthropic provider. When a key pool is configured, each upstream call walks the pool and tries keys in order until one succeeds or the pool is exhausted. Keys are marked **temporary** on 429 (with cooldown from `Retry-After`) and **permanent** on 401/403. Errors that aren't key-specific don't trigger failover. Each agentic-loop iteration gets its own fresh walker, so a tool-call continuation can fail over independently of the initial request. BYOK is unchanged: BYOK requests run as a single attempt with no failover. ## Changes - `config.Anthropic` carries a `KeyPool`. `Key` remains for BYOK X-Api-Key set per interception. - Blocking interceptor: walks the pool, marks keys on key-specific failures, returns on first success or non-failover error. - Streaming interceptor: per-iteration walker. Pre-stream failures fail over to the next key; mid-stream errors are relayed as SSE events. - New `keypool` error types: `TransientExhaustionError` (carries soonest cooldown) and `ErrPermanentExhaustion`. Replace the prior `ErrAllKeysExhausted`. - Error responses now consistently include the outer `"type": "error"` field. ## Related Issues Related to: https://github.com/coder/internal/issues/1446 Related to: https://linear.app/codercom/issue/AIGOV-197/aibridge-automatic-key-failover-for-bridged-and-passthrough-routes ## Follow-up PRs - Bedrock multi-key support. - Refactor provider vs interceptor config separation. - Record the actually-used key in the interception credential hint after failover. > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
This commit is contained in:
@@ -0,0 +1,128 @@
|
||||
package integrationtest //nolint:testpackage // tests unexported internals
|
||||
|
||||
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/config"
|
||||
"github.com/coder/coder/v2/aibridge/fixtures"
|
||||
"github.com/coder/coder/v2/aibridge/keypool"
|
||||
"github.com/coder/coder/v2/aibridge/provider"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// TestAnthropic_KeyFailover verifies that a pool's key state
|
||||
// persists across distinct client requests: a key marked
|
||||
// temporary on request 1 is still skipped on request 2 without
|
||||
// a wasted upstream attempt.
|
||||
func TestAnthropic_KeyFailover(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fix := fixtures.Parse(t, fixtures.AntSimple)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
streaming bool
|
||||
successBody []byte
|
||||
successCType string
|
||||
}{
|
||||
{
|
||||
name: "blocking",
|
||||
streaming: false,
|
||||
successBody: fix.NonStreaming(),
|
||||
successCType: "application/json",
|
||||
},
|
||||
{
|
||||
name: "streaming",
|
||||
streaming: true,
|
||||
successBody: fix.Streaming(),
|
||||
successCType: "text/event-stream",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
pool, err := keypool.New([]string{"k0", "k1"}, quartz.NewMock(t))
|
||||
require.NoError(t, err)
|
||||
|
||||
var requestCount atomic.Int32
|
||||
var seenKeysMu sync.Mutex
|
||||
var seenKeys []string
|
||||
|
||||
// Mock upstream: k0 always returns 429, k1 returns
|
||||
// the per-test success body.
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requestCount.Add(1)
|
||||
key := 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)
|
||||
|
||||
bridgeServer := newBridgeTestServer(t.Context(), t, upstream.URL,
|
||||
withCustomProvider(provider.NewAnthropic(config.Anthropic{
|
||||
BaseURL: upstream.URL,
|
||||
KeyPool: pool,
|
||||
}, nil)),
|
||||
)
|
||||
|
||||
requestBody, err := sjson.SetBytes(fix.Request(), "stream", tc.streaming)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Request 1: walker starts at k0, fails over to k1
|
||||
// after 429.
|
||||
resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, requestBody)
|
||||
require.NoError(t, err)
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
require.NoError(t, resp.Body.Close())
|
||||
require.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
// Request 2: walker skips the now-temporary k0 and
|
||||
// goes straight to k1 (1 upstream call, not 2).
|
||||
resp, err = bridgeServer.makeRequest(t, http.MethodPost, pathAnthropicMessages, requestBody)
|
||||
require.NoError(t, err)
|
||||
_, _ = io.Copy(io.Discard, resp.Body)
|
||||
require.NoError(t, resp.Body.Close())
|
||||
require.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
|
||||
seenKeysMu.Lock()
|
||||
defer seenKeysMu.Unlock()
|
||||
// Request 1: 2 calls (k0 then k1). Request 2: 1 call (k1).
|
||||
assert.Equal(t, int32(3), requestCount.Load(), "upstream request count")
|
||||
assert.Equal(t, []string{"k0", "k1", "k1"}, seenKeys, "seen keys")
|
||||
|
||||
// Pool state persists: k0 temporary, k1 valid.
|
||||
assert.Equal(t, []keypool.KeyState{
|
||||
keypool.KeyStateTemporary,
|
||||
keypool.KeyStateValid,
|
||||
}, pool.PoolState(), "key states")
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user