mirror of
https://github.com/coder/coder.git
synced 2026-10-31 11:02:50 +08:00
## 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
55 lines
1.5 KiB
Go
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
|
|
}
|
|
}
|