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
2026-09-24 15:04:27 +08:00 · 2026-05-07 14:57:44 +01:00
parent 273e828442
commit f1155ac4d7
17 changed files with 2313 additions and 121 deletions
@@ -15,8 +15,10 @@ import (
 	"github.com/coder/coder/v2/aibridge/config"
 	"github.com/coder/coder/v2/aibridge/intercept"
 	"github.com/coder/coder/v2/aibridge/intercept/messages"
+	"github.com/coder/coder/v2/aibridge/keypool"
 	"github.com/coder/coder/v2/aibridge/tracing"
 	"github.com/coder/coder/v2/aibridge/utils"
+	"github.com/coder/quartz"
 )

 // anthropicForwardHeaders lists headers from incoming requests that should be
@@ -55,6 +57,24 @@ func NewAnthropic(cfg config.Anthropic, bedrockCfg *config.AWSBedrock) *Anthropi
 	if cfg.BaseURL == "" {
 		cfg.BaseURL = "https://api.anthropic.com/"
 	}
+	// Resolve centralized key configuration into KeyPool.
+	// Precedence:
+	//   1. cfg.KeyPool (explicit, highest priority).
+	//   2. cfg.Key (legacy single key).
+	// After this block cfg.Key is empty so it can only carry a
+	// BYOK X-Api-Key set per interception in CreateInterceptor.
+	// TODO(ssncferreira): simplify auth field resolution per
+	// https://github.com/coder/aibridge/issues/266.
+	if cfg.KeyPool == nil && cfg.Key != "" {
+		// keypool.New only fails on empty or duplicate keys,
+		// neither possible with a single non-empty key.
+		pool, err := keypool.New([]string{cfg.Key}, quartz.NewReal())
+		if err != nil {
+			panic(fmt.Sprintf("anthropic provider: build single-key pool: %s", err))
+		}
+		cfg.KeyPool = pool
+	}
+	cfg.Key = ""
 	if cfg.CircuitBreaker != nil {
 		cfg.CircuitBreaker.IsFailure = anthropicIsFailure
 		cfg.CircuitBreaker.OpenErrorResponse = anthropicOpenErrorResponse
@@ -119,29 +139,41 @@ func (p *Anthropic) CreateInterceptor(_ http.ResponseWriter, r *http.Request, tr
 	// Any Coder-specific authentication has already been stripped.
 	//
 	// In centralized mode neither Authorization nor X-Api-Key is
-	// present, so cfg keeps the centralized key unchanged.
+	// present, so cfg keeps the KeyPool from provider construction
+	// and the failover loop walks it.
 	//
-	// In BYOK mode the user's LLM credentials survive intact.
-	// If X-Api-Key is present the user has a personal API key;
-	// overwrite the centralized key with it. If Authorization is
-	// present the user authenticated directly with provider;
-	// set BYOKBearerToken and clear the centralized key.
-	// When both are present, X-Api-Key takes priority to match
-	// claude-code behavior.
+	// In BYOK mode the user's LLM credentials survive intact and
+	// failover is disabled by clearing cfg.KeyPool. If X-Api-Key is
+	// present the user has a personal API key, populate cfg.Key.
+	// If Authorization is present the user authenticated directly
+	// with the provider, populate cfg.BYOKBearerToken. When both
+	// are present, X-Api-Key takes priority to match claude-code
+	// behavior.
+	//
+	// TODO(ssncferreira): consolidate auth field handling per
+	// https://github.com/coder/aibridge/issues/266.
 	credKind := intercept.CredentialKindCentralized
-	credSecret := cfg.Key
+	var credSecret string
 	authHeaderName := p.AuthHeader()
 	if apiKey := r.Header.Get("X-Api-Key"); apiKey != "" {
 		cfg.Key = apiKey
+		cfg.KeyPool = nil
 		authHeaderName = "X-Api-Key"
 		credKind = intercept.CredentialKindBYOK
 		credSecret = apiKey
 	} else if token := utils.ExtractBearerToken(r.Header.Get("Authorization")); token != "" {
 		cfg.BYOKBearerToken = token
-		cfg.Key = ""
+		cfg.KeyPool = nil
 		authHeaderName = "Authorization"
 		credKind = intercept.CredentialKindBYOK
 		credSecret = token
+	} else if cfg.KeyPool != nil {
+		// Centralized: use the first key as a placeholder hint.
+		// TODO(ssncferreira): record the actually-used key in
+		// the interception record to reflect failover.
+		if k, err := cfg.KeyPool.Walker().Next(); err == nil {
+			credSecret = k.Value()
+		}
 	}

 	cred := intercept.NewCredentialInfo(credKind, credSecret)
@@ -175,7 +207,16 @@ func (p *Anthropic) InjectAuthHeader(headers *http.Header) {
 		return
 	}

-	headers.Set(p.AuthHeader(), p.cfg.Key)
+	// Centralized: pull a single key from the pool. No failover
+	// or exhaustion handling here.
+	// TODO(ssncferreira): replace with RoundTripper-based auth
+	// in the upstack passthrough PR.
+	if p.cfg.KeyPool == nil {
+		return
+	}
+	if key, err := p.cfg.KeyPool.Walker().Next(); err == nil {
+		headers.Set(p.AuthHeader(), key.Value())
+	}
 }

 func (p *Anthropic) CircuitBreakerConfig() *config.CircuitBreaker {
@@ -13,6 +13,8 @@ import (
 	"github.com/coder/coder/v2/aibridge/config"
 	"github.com/coder/coder/v2/aibridge/intercept"
 	"github.com/coder/coder/v2/aibridge/internal/testutil"
+	"github.com/coder/coder/v2/aibridge/keypool"
+	"github.com/coder/quartz"
 )

 func TestAnthropic_TypeAndName(t *testing.T) {
@@ -49,6 +51,70 @@ func TestAnthropic_TypeAndName(t *testing.T) {
 	}
 }

+func TestNewAnthropic_KeyResolution(t *testing.T) {
+	t.Parallel()
+
+	pool, err := keypool.New([]string{"pool-key-0", "pool-key-1"}, quartz.NewMock(t))
+	require.NoError(t, err)
+
+	tests := []struct {
+		name         string
+		cfg          config.Anthropic
+		expectedKeys []string
+	}{
+		{
+			// Legacy single-key path: NewAnthropic builds a
+			// pool containing just that key.
+			name:         "key_creates_keypool",
+			cfg:          config.Anthropic{Key: "legacy-key"},
+			expectedKeys: []string{"legacy-key"},
+		},
+		{
+			// Caller supplies the pool directly.
+			name:         "keypool_passed_directly",
+			cfg:          config.Anthropic{KeyPool: pool},
+			expectedKeys: []string{"pool-key-0", "pool-key-1"},
+		},
+		{
+			// Both set: KeyPool wins, Key is ignored.
+			name:         "keypool_takes_precedence_over_key",
+			cfg:          config.Anthropic{Key: "legacy-key", KeyPool: pool},
+			expectedKeys: []string{"pool-key-0", "pool-key-1"},
+		},
+		{
+			// Neither set: no centralized auth available. BYOK
+			// auth is set per-request in CreateInterceptor.
+			name:         "neither_set_no_centralized_auth",
+			cfg:          config.Anthropic{},
+			expectedKeys: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+			p := NewAnthropic(tc.cfg, nil)
+
+			if tc.expectedKeys == nil {
+				assert.Nil(t, p.cfg.KeyPool, "expected no KeyPool")
+				return
+			}
+
+			require.NotNil(t, p.cfg.KeyPool)
+			walker := p.cfg.KeyPool.Walker()
+			var got []string
+			for {
+				key, err := walker.Next()
+				if err != nil {
+					break
+				}
+				got = append(got, key.Value())
+			}
+			assert.Equal(t, tc.expectedKeys, got)
+		})
+	}
+}
+
 func TestAnthropic_CreateInterceptor(t *testing.T) {
 	t.Parallel()