mirror of
https://github.com/coder/coder.git
synced 2026-09-22 05:05:20 +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
134 lines
4.0 KiB
Go
134 lines
4.0 KiB
Go
package aibridge
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
"testing"
|
|
"unicode/utf8"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"golang.org/x/xerrors"
|
|
|
|
"github.com/coder/coder/v2/aibridge/circuitbreaker"
|
|
"github.com/coder/coder/v2/aibridge/keypool"
|
|
"github.com/coder/coder/v2/aibridge/recorder"
|
|
)
|
|
|
|
// stubCategorizer is a test errorCategorizer standing in for a provider.
|
|
type stubCategorizer struct {
|
|
result *recorder.ErrorType
|
|
}
|
|
|
|
func (s stubCategorizer) CategorizeError(error) *recorder.ErrorType {
|
|
return s.result
|
|
}
|
|
|
|
func ptr(t recorder.ErrorType) *recorder.ErrorType { return &t }
|
|
|
|
func TestCategorizeInterceptionError(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
cases := []struct {
|
|
name string
|
|
cat stubCategorizer
|
|
err error
|
|
wantType recorder.ErrorType
|
|
wantMsg string
|
|
}{
|
|
{
|
|
name: "nil success",
|
|
err: nil,
|
|
wantType: "",
|
|
wantMsg: "",
|
|
},
|
|
{
|
|
name: "circuit open maps to server error",
|
|
err: circuitbreaker.ErrCircuitOpen,
|
|
wantType: recorder.ErrorTypeServerError,
|
|
wantMsg: circuitbreaker.ErrCircuitOpen.Error(),
|
|
},
|
|
{
|
|
name: "context deadline is timeout",
|
|
err: context.DeadlineExceeded,
|
|
wantType: recorder.ErrorTypeTimeout,
|
|
wantMsg: context.DeadlineExceeded.Error(),
|
|
},
|
|
{
|
|
name: "keypool permanent is unauthorized",
|
|
err: &keypool.Error{Kind: keypool.ErrorKindPermanent},
|
|
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},
|
|
wantType: recorder.ErrorTypeRateLimited,
|
|
wantMsg: (&keypool.Error{Kind: keypool.ErrorKindRateLimited}).Error(),
|
|
},
|
|
{
|
|
name: "keypool unrecognized kind is unknown",
|
|
err: &keypool.Error{Kind: keypool.ErrorKind(-1)},
|
|
wantType: recorder.ErrorTypeUnknown,
|
|
wantMsg: (&keypool.Error{Kind: keypool.ErrorKind(-1)}).Error(),
|
|
},
|
|
{
|
|
name: "context canceled is unknown",
|
|
err: context.Canceled,
|
|
wantType: recorder.ErrorTypeUnknown,
|
|
wantMsg: context.Canceled.Error(),
|
|
},
|
|
{
|
|
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 are permanently unavailable",
|
|
},
|
|
{
|
|
name: "delegated to provider",
|
|
cat: stubCategorizer{result: ptr(recorder.ErrorTypeOverloaded)},
|
|
err: xerrors.New("provider error"),
|
|
wantType: recorder.ErrorTypeOverloaded,
|
|
wantMsg: "provider error",
|
|
},
|
|
{
|
|
name: "provider does not recognize the error",
|
|
err: xerrors.New("mystery"),
|
|
wantType: recorder.ErrorTypeUnknown,
|
|
wantMsg: "mystery",
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
gotType, gotMsg := categorizeInterceptionError(tc.cat, tc.err)
|
|
assert.Equal(t, tc.wantType, gotType)
|
|
assert.Equal(t, tc.wantMsg, gotMsg)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCategorizeInterceptionErrorTruncatesMessage(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
// ASCII: truncated exactly at the byte cap.
|
|
ascii := strings.Repeat("a", maxRecordedErrorMessageBytes*2)
|
|
_, gotMsg := categorizeInterceptionError(stubCategorizer{}, xerrors.New(ascii))
|
|
assert.Len(t, gotMsg, maxRecordedErrorMessageBytes)
|
|
|
|
// Multi-byte: the '€' rune (3 bytes) split at the cap is dropped, leaving
|
|
// valid UTF-8 just below the cap rather than an invalid trailing fragment.
|
|
multibyte := strings.Repeat("€", maxRecordedErrorMessageBytes)
|
|
_, gotMsg = categorizeInterceptionError(stubCategorizer{}, xerrors.New(multibyte))
|
|
assert.True(t, utf8.ValidString(gotMsg), "truncated message must stay valid UTF-8")
|
|
assert.Less(t, len(gotMsg), maxRecordedErrorMessageBytes)
|
|
assert.Positive(t, len(gotMsg))
|
|
}
|