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:
Susana Ferreira
2026-05-07 14:57:44 +01:00
committed by GitHub
parent 273e828442
commit f1155ac4d7
17 changed files with 2313 additions and 121 deletions
@@ -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")
})
}
}