diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index 8bcce231c7..9e968853f4 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -256,8 +256,7 @@ func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Acc break } // OAuth 账号在 401 错误时临时不可调度(给 token 刷新窗口);非 OAuth 账号保持原有 SetError 行为。 - // Antigravity 除外:其 401 由 applyErrorPolicy 的 temp_unschedulable_rules 自行控制。 - if authAccount.Type == AccountTypeOAuth && authAccount.Platform != PlatformAntigravity { + if authAccount.Type == AccountTypeOAuth { // 1. 失效缓存 if s.tokenCacheInvalidator != nil { if err := s.tokenCacheInvalidator.InvalidateToken(ctx, authAccount); err != nil { @@ -299,7 +298,7 @@ func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Acc } shouldDisable = true } else { - // 非 OAuth / Antigravity OAuth:保持 SetError 行为 + // 非 OAuth:保持 SetError 行为 msg := "Authentication failed (401): invalid or expired credentials" if upstreamMsg != "" { msg = "Authentication failed (401): " + upstreamMsg diff --git a/backend/internal/service/ratelimit_service_401_test.go b/backend/internal/service/ratelimit_service_401_test.go index d06c2cb7a9..09afb9314b 100644 --- a/backend/internal/service/ratelimit_service_401_test.go +++ b/backend/internal/service/ratelimit_service_401_test.go @@ -110,9 +110,7 @@ func TestRateLimitService_HandleUpstreamError_OAuth401SetsTempUnschedulable(t *t require.Len(t, invalidator.accounts, 1) }) - t.Run("antigravity_401_uses_SetError", func(t *testing.T) { - // Antigravity 401 由 applyErrorPolicy 的 temp_unschedulable_rules 控制, - // HandleUpstreamError 中走 SetError 路径。 + t.Run("antigravity_401_sets_temp_unschedulable", func(t *testing.T) { repo := &rateLimitAccountRepoStub{} invalidator := &tokenCacheInvalidatorRecorder{} service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) @@ -121,14 +119,22 @@ func TestRateLimitService_HandleUpstreamError_OAuth401SetsTempUnschedulable(t *t ID: 100, Platform: PlatformAntigravity, Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{ + "access_token": "expired-at", + "refresh_token": "rt-100", + }, } shouldDisable := service.HandleUpstreamError(context.Background(), account, 401, http.Header{}, []byte("unauthorized")) require.True(t, shouldDisable) - require.Equal(t, 1, repo.setErrorCalls) - require.Equal(t, 0, repo.tempCalls) - require.Empty(t, invalidator.accounts) + require.Equal(t, 0, repo.setErrorCalls, "Antigravity OAuth 401 must keep status=active so refresh worker can recover it") + require.Equal(t, 1, repo.tempCalls) + require.Equal(t, int64(100), repo.lastTempID) + require.Contains(t, repo.lastTempReason, "invalid or expired credentials") + require.Len(t, invalidator.accounts, 1) + require.Equal(t, int64(100), invalidator.accounts[0].ID) }) } @@ -290,4 +296,27 @@ func TestRateLimitService_HandleUpstreamError_OAuth401NoRefreshTokenSetsError(t require.Equal(t, 1, repo.setErrorCalls) require.Equal(t, 0, repo.tempCalls) }) + + t.Run("antigravity_no_refresh_token_sets_error", func(t *testing.T) { + repo := &rateLimitAccountRepoStub{} + invalidator := &tokenCacheInvalidatorRecorder{} + service := NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + service.SetTokenCacheInvalidator(invalidator) + account := &Account{ + ID: 2883, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": "expired-at", + }, + } + + shouldDisable := service.HandleUpstreamError(context.Background(), account, 401, http.Header{}, []byte("unauthorized")) + + require.True(t, shouldDisable) + require.Equal(t, 1, repo.setErrorCalls, "Antigravity OAuth without refresh_token cannot self-recover") + require.Equal(t, 0, repo.tempCalls) + require.Contains(t, repo.lastErrorMsg, "refresh_token missing") + require.Len(t, invalidator.accounts, 1) + }) } diff --git a/backend/internal/service/token_refresh_service_candidates_test.go b/backend/internal/service/token_refresh_service_candidates_test.go index ec827e4b65..5e2c8cc4aa 100644 --- a/backend/internal/service/token_refresh_service_candidates_test.go +++ b/backend/internal/service/token_refresh_service_candidates_test.go @@ -18,6 +18,7 @@ type tokenRefreshCandidateRepo struct { updatedCredentialIDs []int64 setErrorCalls int setTempUnschedCalls int + clearTempCalls int lastTempUnschedReason string listActiveCalls int } @@ -63,6 +64,11 @@ func (r *tokenRefreshCandidateRepo) SetTempUnschedulable(_ context.Context, _ in return nil } +func (r *tokenRefreshCandidateRepo) ClearTempUnschedulable(context.Context, int64) error { + r.clearTempCalls++ + return nil +} + func isOAuthRefreshPlatform(platform string) bool { switch platform { case PlatformAnthropic, PlatformOpenAI, PlatformGemini, PlatformAntigravity: @@ -128,6 +134,16 @@ func TestTokenRefreshService_ProcessRefreshUsesOAuthRefreshCandidates(t *testing Status: StatusActive, Credentials: map[string]any{"refresh_token": "refresh-token"}, }, + { + ID: 6, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Credentials: map[string]any{"refresh_token": "refresh-token"}, + Extra: map[string]any{"privacy_mode": AntigravityPrivacySet}, + TempUnschedulableUntil: &future, + TempUnschedulableReason: "OAuth 401: unauthorized", + }, }, } svc := &TokenRefreshService{ @@ -140,7 +156,8 @@ func TestTokenRefreshService_ProcessRefreshUsesOAuthRefreshCandidates(t *testing svc.processRefresh() require.Zero(t, repo.listActiveCalls, "TokenRefreshService should not use the broad active-account query") - require.Equal(t, []int64{1}, repo.updatedCredentialIDs) + require.Equal(t, []int64{1, 6}, repo.updatedCredentialIDs) + require.Equal(t, 1, repo.clearTempCalls, "successful refresh should clear the OAuth 401 temp-unschedulable state") } func TestTokenRefreshService_RefreshFailureDoesNotCallPrivacy(t *testing.T) {