Files
coder/aibridge/keypool/keymark.go
T
Susana Ferreira dba45cede7 fix: remove 403 from key failover and cooldown on 401 (#27419)
## Problem

When a key returned 401 or 403, the pool marked it permanently
unavailable for the lifetime of that in-memory pool. This is bad UX: a
transient auth failure or a briefly-misconfigured key could take a key
out of rotation until the operator either restarted Coder or
reconfigured the key (even re-saving the same working value).

## Changes

- **403 removed from key failover**: it's a per-request authorization
failure, not a key-level problem, so it's surfaced to the caller as-is
without marking the key or failing over.
- **401 now applies a temporary cooldown** (like 429) so the key
recovers on its own instead of staying blocked.
- When every key is in an auth-failure cooldown, the pool reports a
`502` with no `Retry-After`, but the keys still recover automatically
once the cooldown elapses.

Closes
https://linear.app/codercom/issue/AIGOV-421/ai-gateway-a-quarantined-centralized-key-never-recovers-without-a
Closes
https://linear.app/codercom/issue/AIGOV-533/403s-misclassifying-keys-as-permanently-down-in-ai-gateway

> [!NOTE]
> Initially generated by Claude Opus 4.7, modified and reviewed by
@ssncferreira
2026-07-27 12:06:01 +01:00

55 lines
1.5 KiB
Go

package keypool
import (
"context"
"net/http"
"cdr.dev/slog/v3"
)
// 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,
resp *http.Response,
logger slog.Logger,
) bool {
if resp == nil {
return false
}
statusCode := resp.StatusCode
switch statusCode {
// 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.applyCooldown(cooldown, reason) {
if p.metrics != nil {
p.metrics.KeyPoolStateTransitions.WithLabelValues(p.providerName, string(reason)).Inc()
}
logger.Info(ctx, "key marked temporary",
slog.F("provider", p.providerName),
slog.F("api_key_hint", key.Hint()),
slog.F("status", statusCode),
slog.F("cooldown", cooldown))
}
return true
default:
logger.Debug(ctx, "status is not a key failover trigger",
slog.F("provider", p.providerName),
slog.F("status", statusCode))
return false
}
}