fix: fall back on model support errors during auth rotation

This commit is contained in:
Tam Nhu Tran
2026-03-18 12:43:45 -04:00
parent db63f9b5d6
commit 5135c22cd6
3 changed files with 287 additions and 43 deletions
@@ -108,6 +108,53 @@ func (e *credentialRetryLimitExecutor) Calls() int {
return e.calls
}
type authFallbackExecutor struct {
id string
mu sync.Mutex
executeCalls []string
executeErrors map[string]error
}
func (e *authFallbackExecutor) Identifier() string {
return e.id
}
func (e *authFallbackExecutor) Execute(_ context.Context, auth *Auth, _ cliproxyexecutor.Request, _ cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
e.mu.Lock()
e.executeCalls = append(e.executeCalls, auth.ID)
err := e.executeErrors[auth.ID]
e.mu.Unlock()
if err != nil {
return cliproxyexecutor.Response{}, err
}
return cliproxyexecutor.Response{Payload: []byte(auth.ID)}, nil
}
func (e *authFallbackExecutor) ExecuteStream(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (*cliproxyexecutor.StreamResult, error) {
return nil, &Error{HTTPStatus: 500, Message: "not implemented"}
}
func (e *authFallbackExecutor) Refresh(_ context.Context, auth *Auth) (*Auth, error) {
return auth, nil
}
func (e *authFallbackExecutor) CountTokens(context.Context, *Auth, cliproxyexecutor.Request, cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
return cliproxyexecutor.Response{}, &Error{HTTPStatus: 500, Message: "not implemented"}
}
func (e *authFallbackExecutor) HttpRequest(context.Context, *Auth, *http.Request) (*http.Response, error) {
return nil, nil
}
func (e *authFallbackExecutor) ExecuteCalls() []string {
e.mu.Lock()
defer e.mu.Unlock()
out := make([]string, len(e.executeCalls))
copy(out, e.executeCalls)
return out
}
func newCredentialRetryLimitTestManager(t *testing.T, maxRetryCredentials int) (*Manager, *credentialRetryLimitExecutor) {
t.Helper()
@@ -191,6 +238,76 @@ func TestManager_MaxRetryCredentials_LimitsCrossCredentialRetries(t *testing.T)
}
}
func TestManager_ModelSupportBadRequest_FallsBackAndSuspendsAuth(t *testing.T) {
m := NewManager(nil, nil, nil)
executor := &authFallbackExecutor{
id: "claude",
executeErrors: map[string]error{
"aa-bad-auth": &Error{
HTTPStatus: http.StatusBadRequest,
Message: "invalid_request_error: The requested model is not supported.",
},
},
}
m.RegisterExecutor(executor)
model := "claude-opus-4-6"
badAuth := &Auth{ID: "aa-bad-auth", Provider: "claude"}
goodAuth := &Auth{ID: "bb-good-auth", Provider: "claude"}
reg := registry.GetGlobalRegistry()
reg.RegisterClient(badAuth.ID, "claude", []*registry.ModelInfo{{ID: model}})
reg.RegisterClient(goodAuth.ID, "claude", []*registry.ModelInfo{{ID: model}})
t.Cleanup(func() {
reg.UnregisterClient(badAuth.ID)
reg.UnregisterClient(goodAuth.ID)
})
if _, errRegister := m.Register(context.Background(), badAuth); errRegister != nil {
t.Fatalf("register bad auth: %v", errRegister)
}
if _, errRegister := m.Register(context.Background(), goodAuth); errRegister != nil {
t.Fatalf("register good auth: %v", errRegister)
}
request := cliproxyexecutor.Request{Model: model}
for i := 0; i < 2; i++ {
resp, errExecute := m.Execute(context.Background(), []string{"claude"}, request, cliproxyexecutor.Options{})
if errExecute != nil {
t.Fatalf("execute %d error = %v, want success", i, errExecute)
}
if string(resp.Payload) != goodAuth.ID {
t.Fatalf("execute %d payload = %q, want %q", i, string(resp.Payload), goodAuth.ID)
}
}
got := executor.ExecuteCalls()
want := []string{badAuth.ID, goodAuth.ID, goodAuth.ID}
if len(got) != len(want) {
t.Fatalf("execute calls = %v, want %v", got, want)
}
for i := range want {
if got[i] != want[i] {
t.Fatalf("execute call %d auth = %q, want %q", i, got[i], want[i])
}
}
updatedBad, ok := m.GetByID(badAuth.ID)
if !ok || updatedBad == nil {
t.Fatalf("expected bad auth to remain registered")
}
state := updatedBad.ModelStates[model]
if state == nil {
t.Fatalf("expected model state for %q", model)
}
if !state.Unavailable {
t.Fatalf("expected bad auth model state to be unavailable")
}
if state.NextRetryAfter.IsZero() {
t.Fatalf("expected bad auth model state cooldown to be set")
}
}
func TestManager_MarkResult_RespectsAuthDisableCoolingOverride(t *testing.T) {
prev := quotaCooldownDisabled.Load()
quotaCooldownDisabled.Store(false)