Thanks to visit codestin.com
Credit goes to github.com

Skip to content

fix: remove 403 from key failover and cooldown on 401 #27419

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ssncferreira merged 4 commits into from
Jul 27, 2026
Merged

fix: remove 403 from key failover and cooldown on 401 #27419

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5079d73
fix: remove 403 from key failover and cooldown on 401
ssncferreira Jul 22, 2026
52165ff
chore: improve response message
ssncferreira Jul 23, 2026
d2d219a
docs: update key failover documentation
ssncferreira Jul 23, 2026
74f9900
chore: address comments
ssncferreira Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 7 additions & 7 deletions aibridge/intercept/chatcompletions/base_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
55 changes: 29 additions & 26 deletions aibridge/intercept/keyfailover_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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{
Expand All @@ -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 {
Expand All @@ -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. Contact your Administrator",
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},

@pawbana pawbana Jul 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On StatusForbidden we don't cycle to next key?
I may lack knowledge what provider use StatusForbidden and StatusUnauthorized for but I would assume next key would be tried if first failed with either of those statuses.

@ssncferreira ssncferreira Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is what we agreed to: completely remove 403 Forbidden as a key failover status. A 403 is a per-request authorization decision (the key is valid but not permitted for this action/resource), not a key-health problem, so this status code is no longer handled, and we return the upstream error to the client. This is exactly the issue customers are seeing.

},
{
// A 500 is not a key-specific failure, so it does not fail over.
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down
4 changes: 2 additions & 2 deletions aibridge/intercept/messages/base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -673,12 +673,12 @@ func ResponseErrorFromKeyPool(keyPoolErr *keypool.Error) *ResponseError {
return nil
}
switch keyPoolErr.Kind {
case keypool.ErrorKindPermanent:
case keypool.ErrorKindPermanent, keypool.ErrorKindUnauthorized:

@ssncferreira ssncferreira Jul 23, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When all the keys have failed with 401, we return a 502 Bad Gateway without a Retry-After. If every key is failing auth at the same time, it's almost always a real problem with the keys (revoked, rotated, or misconfigured), so a retry won't help and a Retry-After would just send clients back on a timer when nothing has changed.

We return 502 rather than the 401 because the auth failure is between the gateway and the upstream, on the centralized keys we manage on the server and the client never sees. Passing the 401 back would read as "your request was unauthorized" when the client's request was fine, and there's nothing on their end to fix. Resolving it is on the administrator, who rotates or reconfigures the keys.

A 401 used to take a key out of rotation until the server restarted, so a key that was only temporarily rejected would be blocked permanently. Now each key is retried once its cooldown expires and recovers on its own, transparent to the user and with no restart needed.

return newResponseError(
keyPoolErr.Error(),
string(constant.ValueOf[constant.APIError]()),
http.StatusBadGateway,
keyPoolErr.RetryAfter,
0,
)
case keypool.ErrorKindRateLimited:
return newResponseError(
Expand Down
21 changes: 14 additions & 7 deletions aibridge/intercept/messages/base_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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.
Expand Down
4 changes: 2 additions & 2 deletions aibridge/intercept/openai_errors.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
7 changes: 7 additions & 0 deletions aibridge/intercept/openai_errors_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
14 changes: 7 additions & 7 deletions aibridge/intercept/responses/base_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion aibridge/interception_error.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 7 additions & 1 deletion aibridge/interception_error_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand All @@ -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",
Expand Down
5 changes: 3 additions & 2 deletions aibridge/internal/integrationtest/interception_error_internal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
2 changes: 1 addition & 1 deletion aibridge/keypool/failover.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
42 changes: 17 additions & 25 deletions aibridge/keypool/keymark.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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, string(reason)).Inc()
}
logger.Info(ctx, "key marked temporary",
slog.F("provider", p.providerName),
Expand All @@ -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),
Expand Down
Loading
Loading