From 5079d73c32691cf4e80970837c526d7a1f2c96fb Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Wed, 22 Jul 2026 14:36:07 +0000 Subject: [PATCH 1/4] fix: remove 403 from key failover and cooldown on 401 --- .../chatcompletions/base_internal_test.go | 14 +-- aibridge/intercept/keyfailover_test.go | 55 ++++++------ aibridge/intercept/messages/base.go | 4 +- .../intercept/messages/base_internal_test.go | 21 +++-- aibridge/intercept/openai_errors.go | 4 +- aibridge/intercept/openai_errors_test.go | 7 ++ .../intercept/responses/base_internal_test.go | 14 +-- aibridge/interception_error.go | 2 +- aibridge/interception_error_internal_test.go | 8 +- .../interception_error_internal_test.go | 5 +- aibridge/keypool/failover.go | 2 +- aibridge/keypool/keymark.go | 42 ++++----- aibridge/keypool/keymark_test.go | 32 +++---- aibridge/keypool/keypool.go | 87 ++++++++++++++----- aibridge/keypool/keypool_test.go | 52 ++++++++++- aibridge/metrics/metrics.go | 4 +- aibridge/passthrough_internal_test.go | 53 ++++++----- 17 files changed, 253 insertions(+), 153 deletions(-) diff --git a/aibridge/intercept/chatcompletions/base_internal_test.go b/aibridge/intercept/chatcompletions/base_internal_test.go index 55baa15a2ac..8d9151fce6e 100644 --- a/aibridge/intercept/chatcompletions/base_internal_test.go +++ b/aibridge/intercept/chatcompletions/base_internal_test.go @@ -236,18 +236,18 @@ func TestMarkKeyOnError(t *testing.T) { expectedState: keypool.KeyStateTemporary, }, { - // Auth failure: mark permanent. - name: "401_marks_permanent", + // Auth failure: temporary cooldown so the key recovers. + name: "401_marks_temporary", err: &openai.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}}, expectedReturn: true, - expectedState: keypool.KeyStatePermanent, + expectedState: keypool.KeyStateTemporary, }, { - // Auth forbidden: mark permanent. - name: "403_marks_permanent", + // Forbidden is per-request, not key-specific. + name: "403_does_not_mark", err: &openai.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}}, - expectedReturn: true, - expectedState: keypool.KeyStatePermanent, + expectedReturn: false, + expectedState: keypool.KeyStateValid, }, { // Server errors are not key-specific. diff --git a/aibridge/intercept/keyfailover_test.go b/aibridge/intercept/keyfailover_test.go index 700e3429e7f..e955d9263df 100644 --- a/aibridge/intercept/keyfailover_test.go +++ b/aibridge/intercept/keyfailover_test.go @@ -178,9 +178,8 @@ var interceptorCases = []interceptorCase{ } // TestInterception_KeyFailover verifies that, within a single interception, the -// centralized key pool fails over across keys (temporary on 429, permanent on -// 401/403) and reports exhaustion, for every interceptor in both blocking and -// streaming mode. +// centralized key pool fails over across keys and reports exhaustion, for every +// interceptor in both blocking and streaming mode. func TestInterception_KeyFailover(t *testing.T) { t.Parallel() @@ -230,33 +229,21 @@ func TestInterception_KeyFailover(t *testing.T) { expectedTransitions: map[string]int{"rate_limited": 1}, }, { - // A 401 marks the key permanent and fails over to the next one. + // A 401 marks the key temporary and fails over to the next one. name: "failover_after_401", keys: []string{k0, k1}, responses: func(s testutil.UpstreamResponse) []testutil.UpstreamResponse { return []testutil.UpstreamResponse{errResp(http.StatusUnauthorized, ""), s} }, expectedStatus: http.StatusOK, - expectedKeyStates: []keypool.KeyState{keypool.KeyStatePermanent, keypool.KeyStateValid}, + expectedKeyStates: []keypool.KeyState{keypool.KeyStateTemporary, keypool.KeyStateValid}, expectedSeenKeys: []string{k0, k1}, expectedTransitions: map[string]int{"unauthorized": 1}, }, - { - // A 403 marks the key permanent and fails over to the next one. - name: "failover_after_403", - keys: []string{k0, k1}, - responses: func(s testutil.UpstreamResponse) []testutil.UpstreamResponse { - return []testutil.UpstreamResponse{errResp(http.StatusForbidden, ""), s} - }, - expectedStatus: http.StatusOK, - expectedKeyStates: []keypool.KeyState{keypool.KeyStatePermanent, keypool.KeyStateValid}, - expectedSeenKeys: []string{k0, k1}, - expectedTransitions: map[string]int{"forbidden": 1}, - }, { // Every key is rate-limited, so the pool is exhausted and the // smallest remaining cooldown is reported. - name: "all_keys_rate_limited", + name: "all_keys_temporary_blocked", keys: []string{k0, k1, k2}, responses: func(testutil.UpstreamResponse) []testutil.UpstreamResponse { return []testutil.UpstreamResponse{ @@ -278,7 +265,9 @@ func TestInterception_KeyFailover(t *testing.T) { expectedExhaustions: map[string]int{"rate_limited": 1}, }, { - // Every key is unauthorized, so the pool is permanently exhausted. + // Every key is unauthorized. Each key cools down and recovers on + // its own, but while all keys are down the exhaustion surfaces as + // an auth failure (502) with no Retry-After. name: "all_keys_unauthorized", keys: []string{k0, k1}, responses: func(testutil.UpstreamResponse) []testutil.UpstreamResponse { @@ -287,11 +276,25 @@ func TestInterception_KeyFailover(t *testing.T) { errResp(http.StatusUnauthorized, ""), } }, - expectedStatus: http.StatusBadGateway, - expectedKeyStates: []keypool.KeyState{keypool.KeyStatePermanent, keypool.KeyStatePermanent}, - expectedSeenKeys: []string{k0, k1}, - expectedTransitions: map[string]int{"unauthorized": 2}, - expectedExhaustions: map[string]int{"auth_failed": 1}, + expectedStatus: http.StatusBadGateway, + expectedRetryAfter: "", + expectedBodyContains: "all configured keys failed authentication", + expectedKeyStates: []keypool.KeyState{keypool.KeyStateTemporary, keypool.KeyStateTemporary}, + expectedSeenKeys: []string{k0, k1}, + expectedTransitions: map[string]int{"unauthorized": 2}, + expectedExhaustions: map[string]int{"auth_failed": 1}, + }, + { + // A 403 is a per-request authorization failure, so it is surfaced + // to the caller without marking the key or failing over. + name: "forbidden_no_failover", + keys: []string{k0, k1}, + responses: func(testutil.UpstreamResponse) []testutil.UpstreamResponse { + return []testutil.UpstreamResponse{errResp(http.StatusForbidden, "")} + }, + expectedStatus: http.StatusForbidden, + expectedKeyStates: []keypool.KeyState{keypool.KeyStateValid, keypool.KeyStateValid}, + expectedSeenKeys: []string{k0}, }, { // A 500 is not a key-specific failure, so it does not fail over. @@ -390,7 +393,7 @@ func TestInterception_KeyFailover(t *testing.T) { gathered, err := reg.Gather() require.NoError(t, err) // One transition per marked key, by reason. - for _, reason := range []string{"rate_limited", "unauthorized", "forbidden"} { + for _, reason := range []string{"rate_limited", "unauthorized"} { if want := tc.expectedTransitions[reason]; want > 0 { assert.True(t, codertestutil.PromCounterHasValue(t, gathered, float64(want), "key_pool_state_transitions_total", ic.provider, reason)) } else { @@ -559,7 +562,7 @@ func TestInterception_AgenticLoopFailover(t *testing.T) { gathered, err := reg.Gather() require.NoError(t, err) // One transition per marked key, by reason. - for _, reason := range []string{"rate_limited", "unauthorized", "forbidden"} { + for _, reason := range []string{"rate_limited", "unauthorized"} { if want := tc.expectedTransitions[reason]; want > 0 { assert.True(t, codertestutil.PromCounterHasValue(t, gathered, float64(want), "key_pool_state_transitions_total", ic.provider, reason)) } else { diff --git a/aibridge/intercept/messages/base.go b/aibridge/intercept/messages/base.go index 1a9ea0dfb85..35e416c63fd 100644 --- a/aibridge/intercept/messages/base.go +++ b/aibridge/intercept/messages/base.go @@ -673,12 +673,12 @@ func ResponseErrorFromKeyPool(keyPoolErr *keypool.Error) *ResponseError { return nil } switch keyPoolErr.Kind { - case keypool.ErrorKindPermanent: + case keypool.ErrorKindPermanent, keypool.ErrorKindUnauthorized: return newResponseError( keyPoolErr.Error(), string(constant.ValueOf[constant.APIError]()), http.StatusBadGateway, - keyPoolErr.RetryAfter, + 0, ) case keypool.ErrorKindRateLimited: return newResponseError( diff --git a/aibridge/intercept/messages/base_internal_test.go b/aibridge/intercept/messages/base_internal_test.go index 939c2fd93fc..4f1aa416200 100644 --- a/aibridge/intercept/messages/base_internal_test.go +++ b/aibridge/intercept/messages/base_internal_test.go @@ -979,6 +979,13 @@ func TestResponseErrorFromKeyPool(t *testing.T) { keyPoolErr: &keypool.Error{Kind: keypool.ErrorKindPermanent}, expectedStatus: http.StatusBadGateway, }, + { + // Auth-failure exhaustion: 502, no Retry-After. + name: "unauthorized_returns_502_without_retry_after", + keyPoolErr: &keypool.Error{Kind: keypool.ErrorKindUnauthorized, RetryAfter: 60 * time.Second}, + expectedStatus: http.StatusBadGateway, + expectedRetryAfter: 0, + }, } for _, tc := range tests { @@ -1020,18 +1027,18 @@ func TestMarkKeyOnError(t *testing.T) { expectedState: keypool.KeyStateTemporary, }, { - // Auth failure: mark permanent. - name: "401_marks_permanent", + // Auth failure: temporary cooldown so the key recovers. + name: "401_marks_temporary", err: &anthropic.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}}, expectedReturn: true, - expectedState: keypool.KeyStatePermanent, + expectedState: keypool.KeyStateTemporary, }, { - // Auth forbidden: mark permanent. - name: "403_marks_permanent", + // Forbidden is per-request, not key-specific. + name: "403_does_not_mark", err: &anthropic.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}}, - expectedReturn: true, - expectedState: keypool.KeyStatePermanent, + expectedReturn: false, + expectedState: keypool.KeyStateValid, }, { // Server errors are not key-specific. diff --git a/aibridge/intercept/openai_errors.go b/aibridge/intercept/openai_errors.go index 80266b1fad2..24b9829dbfb 100644 --- a/aibridge/intercept/openai_errors.go +++ b/aibridge/intercept/openai_errors.go @@ -77,13 +77,13 @@ func ResponseErrorFromKeyPool(keyPoolErr *keypool.Error) *ResponseError { return nil } switch keyPoolErr.Kind { - case keypool.ErrorKindPermanent: + case keypool.ErrorKindPermanent, keypool.ErrorKindUnauthorized: return NewResponseError( keyPoolErr.Error(), OpenAIErrTypeAPI, OpenAIErrCodeServer, http.StatusBadGateway, - keyPoolErr.RetryAfter, + 0, ) case keypool.ErrorKindRateLimited: return NewResponseError( diff --git a/aibridge/intercept/openai_errors_test.go b/aibridge/intercept/openai_errors_test.go index 92953c64eff..24caf453046 100644 --- a/aibridge/intercept/openai_errors_test.go +++ b/aibridge/intercept/openai_errors_test.go @@ -45,6 +45,13 @@ func TestResponseErrorFromKeyPool(t *testing.T) { keyPoolErr: &keypool.Error{Kind: keypool.ErrorKindPermanent}, expectedStatus: http.StatusBadGateway, }, + { + // Auth-failure exhaustion: 502, no Retry-After. + name: "unauthorized_returns_502_without_retry_after", + keyPoolErr: &keypool.Error{Kind: keypool.ErrorKindUnauthorized, RetryAfter: 60 * time.Second}, + expectedStatus: http.StatusBadGateway, + expectedRetryAfter: 0, + }, } for _, tc := range tests { diff --git a/aibridge/intercept/responses/base_internal_test.go b/aibridge/intercept/responses/base_internal_test.go index 69e245f619a..60253bc0d1a 100644 --- a/aibridge/intercept/responses/base_internal_test.go +++ b/aibridge/intercept/responses/base_internal_test.go @@ -526,18 +526,18 @@ func TestMarkKeyOnError(t *testing.T) { expectedState: keypool.KeyStateTemporary, }, { - // Auth failure: mark permanent. - name: "401_marks_permanent", + // Auth failure: temporary cooldown so the key recovers. + name: "401_marks_temporary", err: &openai.Error{StatusCode: http.StatusUnauthorized, Response: &http.Response{StatusCode: http.StatusUnauthorized}}, expectedReturn: true, - expectedState: keypool.KeyStatePermanent, + expectedState: keypool.KeyStateTemporary, }, { - // Auth forbidden: mark permanent. - name: "403_marks_permanent", + // Forbidden is per-request, not key-specific. + name: "403_does_not_mark", err: &openai.Error{StatusCode: http.StatusForbidden, Response: &http.Response{StatusCode: http.StatusForbidden}}, - expectedReturn: true, - expectedState: keypool.KeyStatePermanent, + expectedReturn: false, + expectedState: keypool.KeyStateValid, }, { // Server errors are not key-specific. diff --git a/aibridge/interception_error.go b/aibridge/interception_error.go index e01a71b85d6..58615c37d19 100644 --- a/aibridge/interception_error.go +++ b/aibridge/interception_error.go @@ -62,7 +62,7 @@ func categorizeInterceptionError(c errorCategorizer, err error) (recorder.ErrorT switch keyPoolErr.Kind { case keypool.ErrorKindRateLimited: return recorder.ErrorTypeRateLimited, msg - case keypool.ErrorKindPermanent: + case keypool.ErrorKindPermanent, keypool.ErrorKindUnauthorized: return recorder.ErrorTypeUnauthorized, msg default: return recorder.ErrorTypeUnknown, msg diff --git a/aibridge/interception_error_internal_test.go b/aibridge/interception_error_internal_test.go index d5f08a2121f..5425bef10ee 100644 --- a/aibridge/interception_error_internal_test.go +++ b/aibridge/interception_error_internal_test.go @@ -59,6 +59,12 @@ func TestCategorizeInterceptionError(t *testing.T) { wantType: recorder.ErrorTypeUnauthorized, wantMsg: (&keypool.Error{Kind: keypool.ErrorKindPermanent}).Error(), }, + { + name: "keypool unauthorized is unauthorized", + err: &keypool.Error{Kind: keypool.ErrorKindUnauthorized}, + wantType: recorder.ErrorTypeUnauthorized, + wantMsg: (&keypool.Error{Kind: keypool.ErrorKindUnauthorized}).Error(), + }, { name: "keypool rate limited is rate limited", err: &keypool.Error{Kind: keypool.ErrorKindRateLimited}, @@ -81,7 +87,7 @@ func TestCategorizeInterceptionError(t *testing.T) { name: "wrapped keypool error is unwrapped", err: xerrors.Errorf("key pool exhausted: %w", &keypool.Error{Kind: keypool.ErrorKindPermanent}), wantType: recorder.ErrorTypeUnauthorized, - wantMsg: "key pool exhausted: all configured keys failed authentication", + wantMsg: "key pool exhausted: all configured keys are permanently unavailable", }, { name: "delegated to provider", diff --git a/aibridge/internal/integrationtest/interception_error_internal_test.go b/aibridge/internal/integrationtest/interception_error_internal_test.go index 1f54fb2a938..1975c64833d 100644 --- a/aibridge/internal/integrationtest/interception_error_internal_test.go +++ b/aibridge/internal/integrationtest/interception_error_internal_test.go @@ -18,8 +18,9 @@ import ( // records a categorized upstream error on the ended record. // // The default test provider is centralized (backed by a single-key pool), so a -// 401 exhausts the pool. Both blocking and streaming interceptors preserve the -// *keypool.Error so the cause is categorized as "unauthorized". +// 401 marks the key temporary and exhausts the pool. Both blocking and +// streaming interceptors preserve the *keypool.Error so the cause is +// categorized as "unauthorized", the auth-failure exhaustion outcome. func TestInterceptionUpstreamErrorRecorded(t *testing.T) { t.Parallel() diff --git a/aibridge/keypool/failover.go b/aibridge/keypool/failover.go index 1060c17d8e5..3797427d44d 100644 --- a/aibridge/keypool/failover.go +++ b/aibridge/keypool/failover.go @@ -94,7 +94,7 @@ func (t *keyFailoverTransport) RoundTrip(req *http.Request) (*http.Response, err // Transport-level error, not a key issue. return resp, rtErr } - // MarkKeyOnStatus returns true on key-specific failures (e.g. 401/403/429). + // MarkKeyOnStatus returns true on key-specific failures (e.g. 401/429). if t.config.Pool.MarkKeyOnStatus(req.Context(), key, resp, t.config.Logger) { // Drain and retry with the next key. _, _ = io.Copy(io.Discard, resp.Body) diff --git a/aibridge/keypool/keymark.go b/aibridge/keypool/keymark.go index bb15850b474..b5dea3b220e 100644 --- a/aibridge/keypool/keymark.go +++ b/aibridge/keypool/keymark.go @@ -7,10 +7,10 @@ import ( "cdr.dev/slog/v3" ) -// MarkKeyOnStatus marks key based on a key-specific HTTP -// status code from resp (429 for temporary, 401 or 403 for -// permanent). Returns true if the status was a key-specific -// failover trigger so callers can retry with the next key. +// MarkKeyOnStatus marks key based on a key-specific HTTP status +// code from resp (429 or 401 for temporary). Returns true if the +// status was a key-specific failover trigger so callers can retry +// with the next key. func (p *Pool) MarkKeyOnStatus( ctx context.Context, key *Key, @@ -22,14 +22,21 @@ func (p *Pool) MarkKeyOnStatus( } statusCode := resp.StatusCode switch statusCode { - case http.StatusTooManyRequests: - cooldown := ParseRetryAfter(resp) - if cooldown <= 0 { - cooldown = defaultCooldown + // A 429 rate-limits the key for the provider-supplied cooldown. A 401 + // means the key was rejected, so it cools down for the default period + // and recovers on its own. + case http.StatusTooManyRequests, http.StatusUnauthorized: + cooldown := defaultCooldown + reason := cooldownUnauthorized + if statusCode == http.StatusTooManyRequests { + reason = cooldownRateLimited + if retryAfter := ParseRetryAfter(resp); retryAfter > 0 { + cooldown = retryAfter + } } - if key.MarkTemporary(cooldown) { + if key.applyCooldown(cooldown, reason) { if p.metrics != nil { - p.metrics.KeyPoolStateTransitions.WithLabelValues(p.providerName, reasonRateLimited).Inc() + p.metrics.KeyPoolStateTransitions.WithLabelValues(p.providerName, reason.metricLabel()).Inc() } logger.Info(ctx, "key marked temporary", slog.F("provider", p.providerName), @@ -38,21 +45,6 @@ func (p *Pool) MarkKeyOnStatus( slog.F("cooldown", cooldown)) } return true - case http.StatusUnauthorized, http.StatusForbidden: - if key.MarkPermanent() { - if p.metrics != nil { - reason := reasonUnauthorized - if statusCode == http.StatusForbidden { - reason = reasonForbidden - } - p.metrics.KeyPoolStateTransitions.WithLabelValues(p.providerName, reason).Inc() - } - logger.Warn(ctx, "key marked permanent", - slog.F("provider", p.providerName), - slog.F("api_key_hint", key.Hint()), - slog.F("status", statusCode)) - } - return true default: logger.Debug(ctx, "status is not a key failover trigger", slog.F("provider", p.providerName), diff --git a/aibridge/keypool/keymark_test.go b/aibridge/keypool/keymark_test.go index c90d5912c05..186fb452876 100644 --- a/aibridge/keypool/keymark_test.go +++ b/aibridge/keypool/keymark_test.go @@ -61,18 +61,12 @@ func TestMarkKeyOnStatus(t *testing.T) { expectedReason: "rate_limited", }, { - name: "401_marks_permanent", - statusCode: http.StatusUnauthorized, - expectedReturn: true, - expectedState: keypool.KeyStatePermanent, - expectedReason: "unauthorized", - }, - { - name: "403_marks_permanent", - statusCode: http.StatusForbidden, - expectedReturn: true, - expectedState: keypool.KeyStatePermanent, - expectedReason: "forbidden", + name: "401_marks_temporary", + statusCode: http.StatusUnauthorized, + expectedReturn: true, + expectedState: keypool.KeyStateTemporary, + expectedCooldown: 60 * time.Second, + expectedReason: "unauthorized", }, { name: "200_does_not_mark", @@ -80,6 +74,14 @@ func TestMarkKeyOnStatus(t *testing.T) { expectedReturn: false, expectedState: keypool.KeyStateValid, }, + { + // 403 is a per-request authorization failure, so the key + // is not marked. + name: "403_does_not_mark", + statusCode: http.StatusForbidden, + expectedReturn: false, + expectedState: keypool.KeyStateValid, + }, { name: "500_does_not_mark", statusCode: http.StatusInternalServerError, @@ -121,9 +123,7 @@ func TestMarkKeyOnStatus(t *testing.T) { context.Background(), key, resp, - // 401 and 403 cases legitimately log at error - // level when marking a key permanent. - slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + slogtest.Make(t, nil), ) assert.Equal(t, tc.expectedReturn, got) @@ -133,7 +133,7 @@ func TestMarkKeyOnStatus(t *testing.T) { require.NoError(t, err) // A state transition records one event under its reason, // and other reasons record none. - for _, reason := range []string{"rate_limited", "unauthorized", "forbidden"} { + for _, reason := range []string{"rate_limited", "unauthorized"} { if reason == tc.expectedReason { assert.True(t, codertestutil.PromCounterHasValue(t, gathered, 1, "key_pool_state_transitions_total", providerName, reason)) } else { diff --git a/aibridge/keypool/keypool.go b/aibridge/keypool/keypool.go index 4ca4f76ba40..4a76a029f6d 100644 --- a/aibridge/keypool/keypool.go +++ b/aibridge/keypool/keypool.go @@ -32,6 +32,9 @@ const ( // ErrorKindPermanent means every key is permanently marked // and no key can satisfy the request. ErrorKindPermanent + // ErrorKindUnauthorized means every unavailable key is in a + // cooldown triggered by an authentication failure. + ErrorKindUnauthorized ) // Error is returned when no key is available for the @@ -45,6 +48,8 @@ type Error struct { func (e *Error) Error() string { switch e.Kind { case ErrorKindPermanent: + return "all configured keys are permanently unavailable" + case ErrorKindUnauthorized: return "all configured keys failed authentication" case ErrorKindRateLimited: return fmt.Sprintf("all configured keys are rate-limited (retry after %s)", e.RetryAfter) @@ -71,12 +76,33 @@ const ( // with a zero or negative cooldown duration. const defaultCooldown = 60 * time.Second +// cooldownReason records why a key entered its current cooldown. It is +// meaningful only while the cooldown is active. +type cooldownReason int + +const ( + // cooldownRateLimited means the current cooldown was triggered by a + // rate-limit response (HTTP 429). + cooldownRateLimited cooldownReason = iota + // cooldownUnauthorized means the current cooldown was triggered by an + // authentication failure (HTTP 401). + cooldownUnauthorized +) + +// metricLabel returns the Prometheus label value for the reason, used on +// key-pool state-transition metrics. +func (r cooldownReason) metricLabel() string { + if r == cooldownUnauthorized { + return reasonUnauthorized + } + return reasonRateLimited +} + // Metric label values for the key pool failover metrics. const ( // Reasons for a key_pool_state_transitions_total event. reasonRateLimited = "rate_limited" reasonUnauthorized = "unauthorized" - reasonForbidden = "forbidden" // Outcomes for a key_pool_exhaustions_total event. outcomeRateLimited = "rate_limited" @@ -88,6 +114,9 @@ type Key struct { value string permanent bool cooldownUntil time.Time + // reason records why the current cooldown was applied. It is only + // meaningful while cooldownUntil is active. + reason cooldownReason mu sync.RWMutex clock quartz.Clock @@ -172,26 +201,31 @@ func (k *Key) State() KeyState { return KeyStateValid } -// stateAndCooldown returns the key's state and remaining -// cooldown as a single atomic snapshot. -func (k *Key) stateAndCooldown() (KeyState, time.Duration) { +// stateAndCooldown returns the key's state, remaining cooldown, and the +// reason for the current cooldown as a single atomic snapshot. +func (k *Key) stateAndCooldown() (KeyState, time.Duration, cooldownReason) { k.mu.RLock() defer k.mu.RUnlock() if k.permanent { - return KeyStatePermanent, 0 + return KeyStatePermanent, 0, k.reason } now := k.clock.Now() if now.Before(k.cooldownUntil) { - return KeyStateTemporary, k.cooldownUntil.Sub(now) + return KeyStateTemporary, k.cooldownUntil.Sub(now), k.reason } - return KeyStateValid, 0 + return KeyStateValid, 0, k.reason } -// MarkTemporary marks the key as temporarily unavailable with -// the specified cooldown duration. Returns true if this call -// transitions the key to temporary. +// MarkTemporary marks the key unavailable for the given cooldown. Returns +// true on the valid -> temporary transition. func (k *Key) MarkTemporary(cooldown time.Duration) bool { + return k.applyCooldown(cooldown, cooldownRateLimited) +} + +// applyCooldown marks the key unavailable for the given cooldown, recording +// reason as the cause. Returns true on the valid -> temporary transition. +func (k *Key) applyCooldown(cooldown time.Duration, reason cooldownReason) bool { k.mu.Lock() defer k.mu.Unlock() @@ -215,6 +249,7 @@ func (k *Key) MarkTemporary(cooldown time.Duration) bool { } k.cooldownUntil = newDeadline + k.reason = reason return !inCooldown } @@ -233,46 +268,54 @@ func (k *Key) MarkPermanent() bool { return true } -// keyPoolError returns an Error summarizing why no -// key is currently available. When at least one key is -// temporary, the smallest remaining cooldown is used as the -// retry-after. +// keyPoolError returns an Error summarizing why no key is currently +// available. When at least one key is temporary, the smallest remaining +// cooldown is used as the retry-after. A rate limit anywhere in the pool +// takes precedence, so the exhaustion is classified as unauthorized only +// when every cooldown was triggered by an auth failure. func (p *Pool) keyPoolError() *Error { var retryAfter time.Duration var hasCooldown bool + var isRateLimited bool for i := range p.keys { - state, cooldown := p.keys[i].stateAndCooldown() + state, cooldown, reason := p.keys[i].stateAndCooldown() switch state { // Recoverable now: a key's cooldown expired between the walker's // check and this scan. Return Retry-After: 0 to indicate that // an immediate retry will succeed. case KeyStateValid: return &Error{Kind: ErrorKindRateLimited} - // Recoverable later: track soonest remaining cooldown. + // Recoverable later: track soonest remaining cooldown and reason. case KeyStateTemporary: if !hasCooldown || cooldown < retryAfter { retryAfter = cooldown - hasCooldown = true + } + hasCooldown = true + if reason == cooldownRateLimited { + isRateLimited = true } // Permanent: keep walking to confirm error type. default: } } if hasCooldown { - return &Error{Kind: ErrorKindRateLimited, RetryAfter: retryAfter} + kind := ErrorKindUnauthorized + if isRateLimited { + kind = ErrorKindRateLimited + } + return &Error{Kind: kind, RetryAfter: retryAfter} } return &Error{Kind: ErrorKindPermanent} } -// recordExhaustion increments the exhaustion counter for the outcome -// implied by err.Kind: a rate-limited pool can recover, a permanent -// one cannot. +// recordExhaustion increments the exhaustion counter, labeling the outcome +// as an auth failure for permanent or unauthorized errors, else a rate limit. func (p *Pool) recordExhaustion(err *Error) { if p.metrics == nil { return } outcome := outcomeRateLimited - if err.Kind == ErrorKindPermanent { + if err.Kind == ErrorKindPermanent || err.Kind == ErrorKindUnauthorized { outcome = outcomeAuthFailed } p.metrics.KeyPoolExhaustions.WithLabelValues(p.providerName, outcome).Inc() diff --git a/aibridge/keypool/keypool_test.go b/aibridge/keypool/keypool_test.go index 9880c59e08a..6664cdf97bb 100644 --- a/aibridge/keypool/keypool_test.go +++ b/aibridge/keypool/keypool_test.go @@ -1,6 +1,8 @@ package keypool_test import ( + "context" + "net/http" "sync" "testing" "time" @@ -9,6 +11,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/aibridge/keypool" "github.com/coder/coder/v2/aibridge/metrics" codertestutil "github.com/coder/coder/v2/testutil" @@ -282,6 +285,17 @@ func TestMarkPermanent(t *testing.T) { } } +// markNextByStatus walks to the next available key and cools it down via +// the given HTTP status (401 or 429), so the resulting cooldown carries the +// same reason the failover path records at runtime. +func markNextByStatus(t *testing.T, pool *keypool.Pool, walker *keypool.Walker, status int) { + t.Helper() + key, keyPoolErr := walker.Next() + require.Nil(t, keyPoolErr) + resp := &http.Response{StatusCode: status, Header: make(http.Header)} + pool.MarkKeyOnStatus(context.Background(), key, resp, slogtest.Make(t, nil)) +} + func TestWalkerNext(t *testing.T) { t.Parallel() @@ -499,6 +513,35 @@ func TestWalkerNext(t *testing.T) { expectedValid: []string{}, expectedErr: &keypool.Error{Kind: keypool.ErrorKindRateLimited, RetryAfter: 60 * time.Second}, }, + { + // Given: key-0: temporary (401), key-1: temporary (401). + // Then: key-0: temporary, key-1: temporary. + // Every cooldown is an auth failure, so exhaustion is unauthorized. + name: "all_unauthorized_exhausted", + keys: []string{"key-0", "key-1"}, + setup: func(t *testing.T, pool *keypool.Pool) { + walker := pool.Walker() + markNextByStatus(t, pool, walker, http.StatusUnauthorized) + markNextByStatus(t, pool, walker, http.StatusUnauthorized) + }, + expectedValid: []string{}, + expectedErr: &keypool.Error{Kind: keypool.ErrorKindUnauthorized, RetryAfter: 60 * time.Second}, + }, + { + // Given: key-0: temporary (401), key-1: temporary (429). + // Then: key-0: temporary, key-1: temporary. + // A rate limit anywhere in the pool wins, so exhaustion is + // rate-limited despite the auth failure. + name: "mixed_unauthorized_and_rate_limited_exhausted", + keys: []string{"key-0", "key-1"}, + setup: func(t *testing.T, pool *keypool.Pool) { + walker := pool.Walker() + markNextByStatus(t, pool, walker, http.StatusUnauthorized) + markNextByStatus(t, pool, walker, http.StatusTooManyRequests) + }, + expectedValid: []string{}, + expectedErr: &keypool.Error{Kind: keypool.ErrorKindRateLimited, RetryAfter: 60 * time.Second}, + }, } const providerName = "test-provider" @@ -534,10 +577,11 @@ func TestWalkerNext(t *testing.T) { // exhaustion. assert.Equal(t, len(tc.expectedValid), walker.Attempts()) - // Exhaustion records one event whose outcome reflects the - // error kind: rate-limited keys can recover, permanent cannot. + // Exhaustion records one event whose outcome reflects the error + // kind: a rate limit can recover with a retry, an auth failure or + // permanent marking cannot. wantOutcome := "rate_limited" - if tc.expectedErr.Kind == keypool.ErrorKindPermanent { + if tc.expectedErr.Kind == keypool.ErrorKindPermanent || tc.expectedErr.Kind == keypool.ErrorKindUnauthorized { wantOutcome = "auth_failed" } gathered, err := reg.Gather() @@ -656,7 +700,7 @@ func TestWalkerIndependence(t *testing.T) { require.Nil(t, keyPoolErr) assert.Equal(t, "key-1", key.Value()) - // Simulate 401: mark key-1 permanent. + // Mark key-1 permanent. key.MarkPermanent() // Third attempt: walker advances to key-2. diff --git a/aibridge/metrics/metrics.go b/aibridge/metrics/metrics.go index 3b95c56a78c..554d6b8a788 100644 --- a/aibridge/metrics/metrics.go +++ b/aibridge/metrics/metrics.go @@ -138,12 +138,12 @@ func NewMetrics(reg prometheus.Registerer) *Metrics { // Key pool failover metrics. - // Pessimistic cardinality: 2 providers, 3 reasons = up to 6. + // Pessimistic cardinality: 2 providers, 2 reasons = up to 4. KeyPoolStateTransitions: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{ Subsystem: "key_pool", Name: "state_transitions_total", Help: "The number of API key state transitions during failover " + - "(reason: rate_limited, unauthorized, forbidden).", + "(reason: rate_limited, unauthorized).", }, []string{"provider", "reason"}), // Pessimistic cardinality: 2 providers, 2 outcomes = up to 4. KeyPoolExhaustions: promauto.With(reg).NewCounterVec(prometheus.CounterOpts{ diff --git a/aibridge/passthrough_internal_test.go b/aibridge/passthrough_internal_test.go index 76ca6c17317..7bd519f040b 100644 --- a/aibridge/passthrough_internal_test.go +++ b/aibridge/passthrough_internal_test.go @@ -404,7 +404,7 @@ func TestPassthrough_KeyFailover(t *testing.T) { }, { // Given: 2 keys; key-0 returns 401, key-1 returns 200. - // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. + // Then: 2 requests, 200 response, key-0 temporary, key-1 valid. name: "failover_after_401", keys: []string{"k0", "k1"}, upstreamResponses: []testutil.UpstreamResponse{ @@ -414,33 +414,16 @@ func TestPassthrough_KeyFailover(t *testing.T) { expectedSeenKeys: []string{"k0", "k1"}, expectedStatusCode: http.StatusOK, expectedKeyStates: []keypool.KeyState{ - keypool.KeyStatePermanent, + keypool.KeyStateTemporary, keypool.KeyStateValid, }, expectedTransitions: map[string]int{"unauthorized": 1}, }, - { - // Given: 2 keys; key-0 returns 403, key-1 returns 200. - // Then: 2 requests, 200 response, key-0 permanent, key-1 valid. - name: "failover_after_403", - keys: []string{"k0", "k1"}, - upstreamResponses: []testutil.UpstreamResponse{ - testutil.NewErrorResponse(http.StatusForbidden, ""), - {Blocking: []byte("{}")}, - }, - expectedSeenKeys: []string{"k0", "k1"}, - expectedStatusCode: http.StatusOK, - expectedKeyStates: []keypool.KeyState{ - keypool.KeyStatePermanent, - keypool.KeyStateValid, - }, - expectedTransitions: map[string]int{"forbidden": 1}, - }, { // Given: 3 keys; all return 429 with cooldowns 5s, 3s, 10s. // Then: 3 requests, 429 response with smallest Retry-After, // all keys temporary. - name: "all_keys_rate_limited", + name: "all_keys_temporary_blocked", keys: []string{"k0", "k1", "k2"}, upstreamResponses: []testutil.UpstreamResponse{ testutil.NewErrorResponse(http.StatusTooManyRequests, "5"), @@ -460,7 +443,8 @@ func TestPassthrough_KeyFailover(t *testing.T) { }, { // Given: 2 keys; both return 401. - // Then: 2 requests, 502 response, both keys permanent. + // Then: 2 requests, 502 auth-failure response with no Retry-After, + // both keys temporary and recovering after the cooldown. name: "all_keys_unauthorized", keys: []string{"k0", "k1"}, upstreamResponses: []testutil.UpstreamResponse{ @@ -469,13 +453,29 @@ func TestPassthrough_KeyFailover(t *testing.T) { }, expectedSeenKeys: []string{"k0", "k1"}, expectedStatusCode: http.StatusBadGateway, + expectedRetryAfter: "", expectedKeyStates: []keypool.KeyState{ - keypool.KeyStatePermanent, - keypool.KeyStatePermanent, + keypool.KeyStateTemporary, + keypool.KeyStateTemporary, }, expectedTransitions: map[string]int{"unauthorized": 2}, expectedExhaustions: map[string]int{"auth_failed": 1}, }, + { + // Given: 2 keys; key-0 returns 403. + // Then: 1 request, 403 surfaced as-is, both keys valid. + name: "forbidden_no_failover", + keys: []string{"k0", "k1"}, + upstreamResponses: []testutil.UpstreamResponse{ + testutil.NewErrorResponse(http.StatusForbidden, ""), + }, + expectedSeenKeys: []string{"k0"}, + expectedStatusCode: http.StatusForbidden, + expectedKeyStates: []keypool.KeyState{ + keypool.KeyStateValid, + keypool.KeyStateValid, + }, + }, { // Given: 2 keys; key-0 returns 500. // Then: 1 request, 500 response, both keys remain valid. @@ -532,10 +532,7 @@ func TestPassthrough_KeyFailover(t *testing.T) { } p := prov.newProvider(upstream.URL, pool) - // IgnoreErrors: MarkKey logs at ERROR level when a - // key is marked permanent (401/403); slogtest would - // otherwise fail those scenarios. - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + logger := slogtest.Make(t, nil) handler := newPassthroughRouter(p, logger, nil, testTracer) req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) @@ -560,7 +557,7 @@ func TestPassthrough_KeyFailover(t *testing.T) { gathered, err := reg.Gather() require.NoError(t, err) // One transition per marked key, by reason. - for _, reason := range []string{"rate_limited", "unauthorized", "forbidden"} { + for _, reason := range []string{"rate_limited", "unauthorized"} { if want := tc.expectedTransitions[reason]; want > 0 { assert.True(t, codertestutil.PromCounterHasValue(t, gathered, float64(want), "key_pool_state_transitions_total", "test", reason)) } else { From 52165ff6f9f695b03368b8067e98083df2b99fd8 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 23 Jul 2026 13:03:40 +0000 Subject: [PATCH 2/4] chore: improve response message --- aibridge/intercept/keyfailover_test.go | 2 +- aibridge/keypool/keypool.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/aibridge/intercept/keyfailover_test.go b/aibridge/intercept/keyfailover_test.go index e955d9263df..2e196de449d 100644 --- a/aibridge/intercept/keyfailover_test.go +++ b/aibridge/intercept/keyfailover_test.go @@ -278,7 +278,7 @@ func TestInterception_KeyFailover(t *testing.T) { }, expectedStatus: http.StatusBadGateway, expectedRetryAfter: "", - expectedBodyContains: "all configured keys failed authentication", + expectedBodyContains: "all configured keys failed authentication. Contact your Administrator", expectedKeyStates: []keypool.KeyState{keypool.KeyStateTemporary, keypool.KeyStateTemporary}, expectedSeenKeys: []string{k0, k1}, expectedTransitions: map[string]int{"unauthorized": 2}, diff --git a/aibridge/keypool/keypool.go b/aibridge/keypool/keypool.go index 4a76a029f6d..b7cfa0bfd26 100644 --- a/aibridge/keypool/keypool.go +++ b/aibridge/keypool/keypool.go @@ -50,7 +50,7 @@ func (e *Error) Error() string { case ErrorKindPermanent: return "all configured keys are permanently unavailable" case ErrorKindUnauthorized: - return "all configured keys failed authentication" + return "all configured keys failed authentication. Contact your Administrator" case ErrorKindRateLimited: return fmt.Sprintf("all configured keys are rate-limited (retry after %s)", e.RetryAfter) default: From d2d219a4f545ef9de5cda952f54e36a835decdee Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Thu, 23 Jul 2026 13:03:55 +0000 Subject: [PATCH 3/4] docs: update key failover documentation --- docs/ai-coder/ai-gateway/providers.md | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/docs/ai-coder/ai-gateway/providers.md b/docs/ai-coder/ai-gateway/providers.md index 01ace458246..889689fd3c5 100644 --- a/docs/ai-coder/ai-gateway/providers.md +++ b/docs/ai-coder/ai-gateway/providers.md @@ -322,18 +322,18 @@ a maximum of **5 keys**. ### Failover behavior Every request starts with the first key in the list. If a key is rate-limited -or returns an authentication error, AI Gateway automatically retries the request -with the next available key. - -> [!WARNING] -> A key that fails with an authentication error (`401 Unauthorized` or -> `403 Forbidden`) is permanently disabled and will not be used again until the -> server is restarted or the provider configuration is reloaded. +(`429 Too Many Requests`) or fails authentication (`401 Unauthorized`), AI +Gateway puts that key on a temporary cooldown and retries the request with the +next available key. Keys recover automatically when the cooldown elapses, so +failover stays transparent to end users. Any other response, including a +`403 Forbidden`, is returned to the caller unchanged. If all keys in the pool are exhausted, AI Gateway returns: -- `429 Too Many Requests` when at least one key is rate-limited, with a `Retry-After` header set to the shortest cooldown across all keys. -- `502 Bad Gateway` when every key has failed permanently. +- `429 Too Many Requests` when at least one key is rate-limited, with a `Retry-After` +header set to the shortest cooldown across all keys. +- `502 Bad Gateway` when every key is in an authentication-failure cooldown. +The keys still recover automatically once their cooldowns elapse, so no `Retry-After` is sent. ## Bring Your Own Key From 74f99001286d43319db9d57c4c6485e5e0637184 Mon Sep 17 00:00:00 2001 From: Susana Cardoso Ferreira Date: Mon, 27 Jul 2026 09:04:31 +0000 Subject: [PATCH 4/4] chore: address comments --- aibridge/keypool/keymark.go | 2 +- aibridge/keypool/keypool.go | 21 +++------------------ 2 files changed, 4 insertions(+), 19 deletions(-) diff --git a/aibridge/keypool/keymark.go b/aibridge/keypool/keymark.go index b5dea3b220e..1d867ba3eb7 100644 --- a/aibridge/keypool/keymark.go +++ b/aibridge/keypool/keymark.go @@ -36,7 +36,7 @@ func (p *Pool) MarkKeyOnStatus( } if key.applyCooldown(cooldown, reason) { if p.metrics != nil { - p.metrics.KeyPoolStateTransitions.WithLabelValues(p.providerName, reason.metricLabel()).Inc() + p.metrics.KeyPoolStateTransitions.WithLabelValues(p.providerName, string(reason)).Inc() } logger.Info(ctx, "key marked temporary", slog.F("provider", p.providerName), diff --git a/aibridge/keypool/keypool.go b/aibridge/keypool/keypool.go index b7cfa0bfd26..0809ffd540b 100644 --- a/aibridge/keypool/keypool.go +++ b/aibridge/keypool/keypool.go @@ -78,33 +78,18 @@ const defaultCooldown = 60 * time.Second // cooldownReason records why a key entered its current cooldown. It is // meaningful only while the cooldown is active. -type cooldownReason int +type cooldownReason string const ( // cooldownRateLimited means the current cooldown was triggered by a // rate-limit response (HTTP 429). - cooldownRateLimited cooldownReason = iota + cooldownRateLimited cooldownReason = "rate_limited" // cooldownUnauthorized means the current cooldown was triggered by an // authentication failure (HTTP 401). - cooldownUnauthorized + cooldownUnauthorized cooldownReason = "unauthorized" ) -// metricLabel returns the Prometheus label value for the reason, used on -// key-pool state-transition metrics. -func (r cooldownReason) metricLabel() string { - if r == cooldownUnauthorized { - return reasonUnauthorized - } - return reasonRateLimited -} - -// Metric label values for the key pool failover metrics. const ( - // Reasons for a key_pool_state_transitions_total event. - reasonRateLimited = "rate_limited" - reasonUnauthorized = "unauthorized" - - // Outcomes for a key_pool_exhaustions_total event. outcomeRateLimited = "rate_limited" outcomeAuthFailed = "auth_failed" )