mirror of
https://github.com/coder/coder.git
synced 2026-10-31 11:02:50 +08:00
Susana Ferreira
dba45cede7
## 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
71 lines
1.9 KiB
Go
71 lines
1.9 KiB
Go
package intercept_test
|
|
|
|
import (
|
|
"net/http"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/coder/coder/v2/aibridge/intercept"
|
|
"github.com/coder/coder/v2/aibridge/keypool"
|
|
)
|
|
|
|
func TestResponseErrorFromKeyPool(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
tests := []struct {
|
|
name string
|
|
keyPoolErr *keypool.Error
|
|
expectedStatus int
|
|
expectedRetryAfter time.Duration
|
|
}{
|
|
{
|
|
name: "nil_returns_nil",
|
|
keyPoolErr: nil,
|
|
},
|
|
{
|
|
// Rate-limited with no cooldown: 429, no Retry-After.
|
|
name: "rate_limited_zero_retry_after",
|
|
keyPoolErr: &keypool.Error{Kind: keypool.ErrorKindRateLimited},
|
|
expectedStatus: http.StatusTooManyRequests,
|
|
expectedRetryAfter: 0,
|
|
},
|
|
{
|
|
// Rate-limited with cooldown: 429, Retry-After set.
|
|
name: "rate_limited_with_retry_after",
|
|
keyPoolErr: &keypool.Error{Kind: keypool.ErrorKindRateLimited, RetryAfter: 5 * time.Second},
|
|
expectedStatus: http.StatusTooManyRequests,
|
|
expectedRetryAfter: 5 * time.Second,
|
|
},
|
|
{
|
|
// Permanent: 502 api_error.
|
|
name: "permanent_returns_502",
|
|
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 {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
got := intercept.ResponseErrorFromKeyPool(tc.keyPoolErr)
|
|
if tc.keyPoolErr == nil {
|
|
assert.Nil(t, got)
|
|
return
|
|
}
|
|
require.NotNil(t, got)
|
|
assert.Equal(t, tc.expectedStatus, got.StatusCode)
|
|
assert.Equal(t, tc.expectedRetryAfter, got.RetryAfter)
|
|
})
|
|
}
|
|
}
|