mirror of
https://github.com/coder/coder.git
synced 2026-09-23 14:03:57 +08:00
Susana Ferreira
01ec5e4577
## Description This PR adds Prometheus metrics for aibridge's API-key failover, giving visibility into key pool health and failover behavior per provider. The following metrics are introduced: - **`key_pool_state`** (gauge): number of keys currently in each state (`valid`, `temporary`, `permanent`) per provider, sampled at scrape time. - **`key_pool_state_transitions_total`** (counter): key state transitions during failover, labeled by `reason` (`rate_limited`, `unauthorized`, `forbidden`). - **`key_pool_exhaustions_total`** (counter): times a pool ran out of usable keys, labeled by `outcome` (`rate_limited`, `auth_failed`). - **`key_pool_failover_attempts`** (histogram): keys attempted before success or exhaustion (per interception for bridged requests, per request for passthrough). ## Changes - Moves `MarkKeyOnStatus` and key-pool error handling onto `*keypool.Pool`. - Attaches metrics to each provider's key pool at install time, on construction and on provider reload. - Adds a scrape-time state collector and a `KeyPools()` accessor on the bridge pool to feed it. - Tracks per-request key attempts in the bridged and passthrough failover paths. - Adds test coverage for the new metrics across the keypool unit tests, the bridged intercept failover tests, and the passthrough failover test. Closes https://github.com/coder/internal/issues/1447 Closes https://linear.app/codercom/issue/AIGOV-198/aibridge-key-failover-observability > [!NOTE] > Initially generated by Claude Opus 4.7, modified and reviewed by @ssncferreira
63 lines
1.6 KiB
Go
63 lines
1.6 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 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.
|
|
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 {
|
|
case http.StatusTooManyRequests:
|
|
cooldown := ParseRetryAfter(resp)
|
|
if cooldown <= 0 {
|
|
cooldown = defaultCooldown
|
|
}
|
|
if key.MarkTemporary(cooldown) {
|
|
if p.metrics != nil {
|
|
p.metrics.KeyPoolStateTransitions.WithLabelValues(p.providerName, reasonRateLimited).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
|
|
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),
|
|
slog.F("status", statusCode))
|
|
return false
|
|
}
|
|
}
|