Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 35 additions & 2 deletions sdk/cliproxy/auth/conductor.go
Original file line number Diff line number Diff line change
Expand Up @@ -3734,7 +3734,18 @@ func (m *Manager) MarkResult(ctx context.Context, result Result) {
state.NextRetryAfter = nextTransientErrorRetryAfter(now)
}
default:
state.NextRetryAfter = time.Time{}
if isOAuthInvalidGrantResultError(result.Error) {
if disableCooling {
state.NextRetryAfter = time.Time{}
} else {
next := now.Add(30 * time.Minute)
state.NextRetryAfter = next
suspendReason = "unauthorized"
shouldSuspendModel = true
Comment on lines +3742 to +3746

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Cool down the whole invalid_grant credential

When a revoked Antigravity refresh token returns invalid_grant for one model, this only records the cooldown on that model state and suspends that single auth/model pair. isAuthBlockedForModel treats the same auth as available for a different requested model when that model has no state, so the banned credential can still be selected for every other model and recreate the retry/failover loop per model. Since invalid_grant is credential-level, block the auth across models rather than only result.Model.

Useful? React with 👍 / 👎.

}
} else {
state.NextRetryAfter = time.Time{}
}
}
}

Expand Down Expand Up @@ -3983,6 +3994,20 @@ func isUnauthorizedError(err error) bool {
return strings.Contains(raw, "status 401") || strings.Contains(raw, "401 unauthorized")
}

// isOAuthInvalidGrantMessage checks whether an error message indicates an OAuth
// invalid_grant failure (revoked or expired refresh token). These are credential-level
// errors that should be cooled down, not treated as invalid request payloads.
func isOAuthInvalidGrantMessage(msg string) bool {
return strings.Contains(msg, "invalid_grant")
}
Comment on lines +4002 to +4004

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

We should perform a case-insensitive check when looking for the "invalid_grant" substring. Upstream OAuth providers or intermediate API wrappers might return error codes or messages in uppercase (e.g., "INVALID_GRANT") or mixed case. Converting the message to lowercase first ensures robust detection.

Suggested change
func isOAuthInvalidGrantMessage(msg string) bool {
return strings.Contains(msg, "invalid_grant")
}
func isOAuthInvalidGrantMessage(msg string) bool {
return strings.Contains(strings.ToLower(msg), "invalid_grant")
}


func isOAuthInvalidGrantResultError(err *Error) bool {
if err == nil {
return false
}
return isOAuthInvalidGrantMessage(err.Message) || isOAuthInvalidGrantMessage(err.Code)
}
Comment on lines +4006 to +4011

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Instead of manually checking both err.Message and err.Code separately, we can leverage the err.Error() method. This is cleaner, more maintainable, and consistent with how we check errors in other parts of the codebase (such as in applyAuthFailureState).

Suggested change
func isOAuthInvalidGrantResultError(err *Error) bool {
if err == nil {
return false
}
return isOAuthInvalidGrantMessage(err.Message) || isOAuthInvalidGrantMessage(err.Code)
}
func isOAuthInvalidGrantResultError(err *Error) bool {
if err == nil {
return false
}
return isOAuthInvalidGrantMessage(err.Error())
}


func hasUnauthorizedAuthFailure(auth *Auth) bool {
if auth == nil || auth.LastError == nil {
return false
Expand Down Expand Up @@ -4152,6 +4177,9 @@ func isRequestInvalidError(err error) bool {
switch status {
case http.StatusBadRequest:
msg := err.Error()
if isOAuthInvalidGrantMessage(msg) {
return false
}
return strings.Contains(msg, "invalid_request_error") ||
strings.Contains(msg, "INVALID_ARGUMENT") ||
strings.Contains(msg, "FAILED_PRECONDITION")
Expand Down Expand Up @@ -4241,7 +4269,12 @@ func applyAuthFailureState(auth *Auth, resultErr *Error, retryAfter *time.Durati
auth.NextRetryAfter = nextTransientErrorRetryAfter(now)
}
default:
if auth.StatusMessage == "" {
if resultErr != nil && isOAuthInvalidGrantMessage(resultErr.Error()) {
auth.StatusMessage = "unauthorized"
if !disableCooling {
auth.NextRetryAfter = now.Add(30 * time.Minute)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clear stale retry time when cooling is disabled

When disable_cooling is true and an auth-level invalid_grant is recorded after the auth already has a retry deadline, this branch leaves auth.NextRetryAfter unchanged because it only assigns in the !disableCooling path. The other auth-level cooldown cases above explicitly zero this field when cooling is disabled, otherwise selector checks can keep blocking the auth on a stale deadline despite cooling being disabled.

Useful? React with 👍 / 👎.

}
Comment on lines +4274 to +4280

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

When disableCooling is true, auth.NextRetryAfter is not explicitly cleared, which means any pre-existing cooldown timestamp might be preserved. To be consistent with other status code cases (like 401, 402, 403, 404, etc.) and to ensure cooling is correctly disabled, we should explicitly set auth.NextRetryAfter = time.Time{} when disableCooling is true.

Suggested change
if resultErr != nil && isOAuthInvalidGrantMessage(resultErr.Error()) {
auth.StatusMessage = "unauthorized"
if !disableCooling {
auth.NextRetryAfter = now.Add(30 * time.Minute)
}
if resultErr != nil && isOAuthInvalidGrantMessage(resultErr.Error()) {
auth.StatusMessage = "unauthorized"
if disableCooling {
auth.NextRetryAfter = time.Time{}
} else {
auth.NextRetryAfter = now.Add(30 * time.Minute)
}
}

} else if auth.StatusMessage == "" {
auth.StatusMessage = "request failed"
}
}
Expand Down
81 changes: 81 additions & 0 deletions sdk/cliproxy/auth/cooldown_backoff_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -152,6 +152,87 @@ func TestApplyAuthFailureStateQuotaBackoffOncePerWindow(t *testing.T) {
}
}

func TestMarkResultInvalidGrantAppliesCooldown(t *testing.T) {
manager := NewManager(nil, nil, nil)
auth := &Auth{
ID: "auth-invalid-grant",
Provider: "antigravity",
Metadata: map[string]any{"type": "antigravity"},
}
if _, errRegister := manager.Register(WithSkipPersist(context.Background()), auth); errRegister != nil {
t.Fatalf("Register returned error: %v", errRegister)
}

manager.MarkResult(context.Background(), Result{
AuthID: auth.ID,
Provider: "antigravity",
Model: "claude-sonnet-4-20250514",
Success: false,
Error: &Error{
Code: "invalid_request_error",
Message: `{"error":"invalid_grant","error_description":"Bad Request"}`,
HTTPStatus: http.StatusBadRequest,
},
})

updated, ok := manager.GetByID(auth.ID)
if !ok || updated == nil {
t.Fatal("expected auth to exist")
}
state := updated.ModelStates["claude-sonnet-4-20250514"]
if state == nil {
t.Fatal("expected model state after failure")
}
if state.NextRetryAfter.IsZero() {
t.Fatal("expected cooldown (non-zero NextRetryAfter) for invalid_grant, got zero")
}
if !state.NextRetryAfter.After(time.Now()) {
t.Fatalf("expected cooldown in the future, got %v", state.NextRetryAfter)
}
}

func TestApplyAuthFailureStateInvalidGrantAppliesCooldown(t *testing.T) {
now := time.Now()
invalidGrantErr := &Error{
Code: "invalid_request_error",
Message: `{"error":"invalid_grant","error_description":"Bad Request"}`,
HTTPStatus: http.StatusBadRequest,
}
auth := &Auth{ID: "auth-grant-cooldown"}

applyAuthFailureState(auth, invalidGrantErr, nil, now, false)

if auth.StatusMessage != "unauthorized" {
t.Fatalf("StatusMessage = %q, want %q", auth.StatusMessage, "unauthorized")
}
if auth.NextRetryAfter.IsZero() {
t.Fatal("expected NextRetryAfter to be set for invalid_grant")
}
if !auth.NextRetryAfter.Equal(now.Add(30 * time.Minute)) {
t.Fatalf("NextRetryAfter = %v, want %v", auth.NextRetryAfter, now.Add(30*time.Minute))
}
}

func TestIsRequestInvalidErrorExcludesInvalidGrant(t *testing.T) {
grantErr := &Error{
Code: "invalid_request_error",
Message: `{"error":"invalid_grant","error_description":"Bad Request"}`,
HTTPStatus: http.StatusBadRequest,
}
if isRequestInvalidError(grantErr) {
t.Fatal("isRequestInvalidError should return false for invalid_grant wrapped in invalid_request_error")
}

normalBadReq := &Error{
Code: "invalid_request_error",
Message: "invalid parameter: temperature must be between 0 and 2",
HTTPStatus: http.StatusBadRequest,
}
if !isRequestInvalidError(normalBadReq) {
t.Fatal("isRequestInvalidError should return true for genuine invalid_request_error")
}
}

func TestJitteredCooldownWaitBounds(t *testing.T) {
cases := []struct {
wait time.Duration
Expand Down
Loading