From 6a589075145e45e05d300a601a2c24bfd8f539ae Mon Sep 17 00:00:00 2001 From: HSJ-BanFan Date: Wed, 22 Jul 2026 12:51:54 +0800 Subject: [PATCH] fix: allow account failover on exhausted paid quota --- .../application/gateway/failure_test.go | 5 ++ .../internal/application/gateway/selector.go | 23 +++---- .../internal/application/gateway/service.go | 23 +++++++ .../application/gateway/service_test.go | 60 +++++++++++++++---- 4 files changed, 89 insertions(+), 22 deletions(-) diff --git a/backend/internal/application/gateway/failure_test.go b/backend/internal/application/gateway/failure_test.go index 959d3f402..4d3e1570f 100644 --- a/backend/internal/application/gateway/failure_test.go +++ b/backend/internal/application/gateway/failure_test.go @@ -58,6 +58,11 @@ func TestRetryableResponseHonorsUpstreamRetryVeto(t *testing.T) { if isRetryableResponse(response) { t.Fatal("x-should-retry:false 必须禁止换账号重试") } + response.StatusCode = http.StatusPaymentRequired + if !isRetryableResponse(response) { + t.Fatal("账号级 402 必须忽略上游的原账号重试 veto，允许跨账号故障转移") + } + response.StatusCode = http.StatusInternalServerError response.Header.Set("X-Should-Retry", "true") if !isRetryableResponse(response) { t.Fatal("x-should-retry:true 不应覆盖现有状态码重试策略") diff --git a/backend/internal/application/gateway/selector.go b/backend/internal/application/gateway/selector.go index c973d0de4..f49a24d87 100644 --- a/backend/internal/application/gateway/selector.go +++ b/backend/internal/application/gateway/selector.go @@ -32,6 +32,7 @@ const concurrencySnapshotTTL = 25 * time.Millisecond const maxConcurrencySnapshots = 256 const modelAccessDeniedCooldown = 5 * time.Minute +const unknownPaidQuotaProbeInterval = 24 * time.Hour type candidateSnapshot struct { values []account.RoutingCandidate @@ -540,20 +541,22 @@ func (s *Selector) MarkModelAccessDenied(ctx context.Context, credential account s.invalidateCandidates(credential.Provider) } -// MarkPaidQuotaExhausted 使用已知真实账期将付费账号移出号池，到期后才允许 Billing 探测。 +// MarkPaidQuotaExhausted 将明确的付费额度耗尽移出号池。 +// 已知真实账期时等到账期结束；快照不完整时保守等待后再做 Billing 探测。 func (s *Selector) MarkPaidQuotaExhausted(ctx context.Context, credential account.Credential, billing *account.Billing) bool { - if billing == nil || !billing.IsPaid() { - return false + now := time.Now().UTC() + nextProbeAt := now.Add(unknownPaidQuotaProbeInterval) + if billing != nil && billing.IsPaid() { + if periodEnd, ok := billing.PeriodEnd(); ok && periodEnd.After(now) { + nextProbeAt = periodEnd + } } - periodEnd, ok := billing.PeriodEnd() - if !ok { + if err := s.accounts.SaveQuotaRecovery(ctx, account.QuotaRecovery{ + AccountID: credential.ID, Kind: account.QuotaRecoveryKindPaid, Status: account.QuotaRecoveryStatusExhausted, + ExhaustedAt: &now, NextProbeAt: &nextProbeAt, LastConfirmedAt: &now, UpdatedAt: now, + }); err != nil { return false } - now := time.Now().UTC() - _ = s.accounts.SaveQuotaRecovery(ctx, account.QuotaRecovery{ - AccountID: credential.ID, Kind: account.QuotaRecoveryKindPaid, Status: account.QuotaRecoveryStatusExhausted, - ExhaustedAt: &now, NextProbeAt: &periodEnd, LastConfirmedAt: &now, UpdatedAt: now, - }) _ = s.sticky.DeleteByAccount(ctx, credential.ID) s.invalidateCandidates(credential.Provider) return true diff --git a/backend/internal/application/gateway/service.go b/backend/internal/application/gateway/service.go index 5cf7bb4dc..878ed3eb5 100644 --- a/backend/internal/application/gateway/service.go +++ b/backend/internal/application/gateway/service.go @@ -45,6 +45,8 @@ const finalizationTimeout = 5 * time.Second const textBillingReservationTTL = 2 * time.Hour const mediaBillingReservationTTL = 24 * time.Hour const modelCatalogRefreshTimeout = 30 * time.Second +const unknownQuotaExhaustedCooldown = 24 * time.Hour +const maxQuotaFailoverAttempts = 10 var freeQuotaUsagePattern = regexp.MustCompile(`(?i)tokens\s*\(actual/limit\)\s*:\s*([0-9]+)\s*/\s*([0-9]+)`) @@ -725,6 +727,16 @@ attemptLoop: failureHandled = true } else if lastFailure.QuotaExhausted { failureHandled = s.selector.MarkPaidQuotaExhausted(ctx, credential, lease.Billing) + if !failureHandled { + // spending-limit is definitive for this account even when the cached + // Billing snapshot has no trustworthy period end. Keep it out of the + // pool long enough to avoid immediately poisoning other sticky sessions. + if retryAfter < unknownQuotaExhaustedCooldown { + retryAfter = unknownQuotaExhaustedCooldown + } + s.selector.MarkFailure(ctx, credential, response.StatusCode, retryAfter) + failureHandled = true + } } if s.providers.SupportsCredentialRefresh(credential.Provider) && lastFailure.PermanentAccountDenial { if credential.Provider == accountdomain.ProviderBuild { @@ -753,6 +765,12 @@ attemptLoop: break } } + // Definite account quota exhaustion should not consume the small generic + // retry budget while healthy accounts remain in a large pool. Stored + // responses stay pinned and therefore never receive this extension. + if ownership == nil && lastFailure.AccountScoped && lastFailure.QuotaExhausted && attempts < maxQuotaFailoverAttempts { + attempts++ + } continue } if response.StatusCode >= 200 && response.StatusCode < 300 { @@ -1146,6 +1164,11 @@ func isRetryableResponse(response *provider.Response) bool { if response == nil || !isRetryable(response.StatusCode) { return false } + // X-Should-Retry describes retrying the same upstream account. A 402 is + // account-scoped, so the gateway must still fail over to another account. + if response.StatusCode == http.StatusPaymentRequired { + return true + } return !strings.EqualFold(strings.TrimSpace(response.Header.Get("X-Should-Retry")), "false") } diff --git a/backend/internal/application/gateway/service_test.go b/backend/internal/application/gateway/service_test.go index c9d4bc2a5..8661ecabe 100644 --- a/backend/internal/application/gateway/service_test.go +++ b/backend/internal/application/gateway/service_test.go @@ -102,10 +102,18 @@ func TestGatewayFailsOverBeforeReturningBody(t *testing.T) { if err != nil { t.Fatal(err) } + third, _, err := accountRepo.UpsertByIdentity(ctx, account.Credential{Provider: account.ProviderBuild, Name: "third", SourceKey: "third", EncryptedAccessToken: "three", ExpiresAt: time.Now().Add(time.Hour), Enabled: true, AuthStatus: account.AuthStatusActive, Priority: 50, MaxConcurrent: 1}) + if err != nil { + t.Fatal(err) + } + fourth, _, err := accountRepo.UpsertByIdentity(ctx, account.Credential{Provider: account.ProviderBuild, Name: "fourth", SourceKey: "fourth", EncryptedAccessToken: "four", ExpiresAt: time.Now().Add(time.Hour), Enabled: true, AuthStatus: account.AuthStatusActive, Priority: 25, MaxConcurrent: 1}) + if err != nil { + t.Fatal(err) + } if err := modelRepo.UpsertDiscovered(ctx, account.ProviderBuild, []string{"grok-test"}); err != nil { t.Fatal(err) } - for _, accountID := range []uint64{first.ID, second.ID} { + for _, accountID := range []uint64{first.ID, second.ID, third.ID, fourth.ID} { if err := modelRepo.ReplaceAccountCapabilities(ctx, accountID, []string{"grok-test"}, time.Now().UTC()); err != nil { t.Fatal(err) } @@ -115,7 +123,12 @@ func TestGatewayFailsOverBeforeReturningBody(t *testing.T) { t.Fatal(err) } - adapter := &failoverAdapter{firstID: first.ID} + adapter := &failoverAdapter{ + failedIDs: map[uint64]bool{first.ID: true, second.ID: true, third.ID: true}, + firstStatus: http.StatusPaymentRequired, + firstBody: `{"code":"personal-team-blocked:spending-limit","error":"You have run out of credits or need a Grok subscription."}`, + firstHeader: http.Header{"X-Should-Retry": {"false"}}, + } registry := provider.NewRegistry(adapter) cipher := testCipher(t) sticky := memory.NewStickyStore() @@ -137,7 +150,7 @@ func TestGatewayFailsOverBeforeReturningBody(t *testing.T) { if string(body) != "ok" { t.Fatalf("body = %q", body) } - if len(adapter.attempts) != 2 || adapter.attempts[0] != first.ID || adapter.attempts[1] != second.ID { + if len(adapter.attempts) != 4 || adapter.attempts[0] != first.ID || adapter.attempts[1] != second.ID || adapter.attempts[2] != third.ID || adapter.attempts[3] != fourth.ID { t.Fatalf("attempts = %#v", adapter.attempts) } identity := resolveBuildSessionIdentity(clientKey.ID, account.ProviderBuild, "grok-test", "", "claude-session", nil) @@ -151,15 +164,22 @@ func TestGatewayFailsOverBeforeReturningBody(t *testing.T) { if adapter.lastGrokTurnIndex != "3" { t.Fatalf("Grok turn index = %q, want 3", adapter.lastGrokTurnIndex) } - if boundID, ok, err := sticky.Get(ctx, stickySessionKey(identity.affinityKey), time.Now().UTC()); err != nil || !ok || boundID != second.ID { - t.Fatalf("failover sticky binding = %d, %v, err = %v; want account %d", boundID, ok, err, second.ID) + if boundID, ok, err := sticky.Get(ctx, stickySessionKey(identity.affinityKey), time.Now().UTC()); err != nil || !ok || boundID != fourth.ID { + t.Fatalf("failover sticky binding = %d, %v, err = %v; want account %d", boundID, ok, err, fourth.ID) } - observedAccount, err := accountRepo.Get(ctx, second.ID) + recovery, err := accountRepo.GetQuotaRecovery(ctx, first.ID) + if err != nil { + t.Fatal(err) + } + if recovery.Kind != account.QuotaRecoveryKindPaid || recovery.Status != account.QuotaRecoveryStatusExhausted || recovery.NextProbeAt == nil || time.Until(*recovery.NextProbeAt) < 23*time.Hour { + t.Fatalf("spending-limit account did not enter persistent quota recovery: %#v", recovery) + } + observedAccount, err := accountRepo.Get(ctx, fourth.ID) if err != nil || observedAccount.ObservedModel != "grok-test-build-free" { t.Fatalf("observed account = %#v, err = %v", observedAccount, err) } logs, total, err := auditRepo.List(ctx, 0, 10) - if err != nil || total != 1 || logs[0].AccountID == nil || *logs[0].AccountID != second.ID || logs[0].ClientKeyName != "test-key" || logs[0].ModelPublicID != "grok-test" || logs[0].ModelUpstreamModel != "Build/grok-test" || logs[0].AccountName != "second" || logs[0].CachedInputTokens != 80 || logs[0].AttemptCount != 0 { + if err != nil || total != 1 || logs[0].AccountID == nil || *logs[0].AccountID != fourth.ID || logs[0].ClientKeyName != "test-key" || logs[0].ModelPublicID != "grok-test" || logs[0].ModelUpstreamModel != "Build/grok-test" || logs[0].AccountName != "fourth" || logs[0].CachedInputTokens != 80 || logs[0].AttemptCount != 0 { t.Fatalf("audit = %#v, %d, %v", logs, total, err) } detail, err := auditRepo.Get(ctx, logs[0].ID) @@ -167,7 +187,7 @@ func TestGatewayFailsOverBeforeReturningBody(t *testing.T) { t.Fatalf("audit detail = %#v, err = %v", detail, err) } ownership, err := responseRepo.Get(ctx, "resp-test", clientKey.ID, time.Now().UTC()) - if err != nil || ownership.AccountID != second.ID || ownership.PromptCacheKey != expectedCacheKey || ownership.ReasoningReplayKey != identity.replayKey { + if err != nil || ownership.AccountID != fourth.ID || ownership.PromptCacheKey != expectedCacheKey || ownership.ReasoningReplayKey != identity.replayKey { t.Fatalf("ownership = %#v, err = %v", ownership, err) } @@ -197,7 +217,7 @@ func TestGatewayFailsOverBeforeReturningBody(t *testing.T) { _, _ = io.ReadAll(continued.Body) continued.Finalize(Usage{}, "resp-next", "") _ = continued.Body.Close() - if len(adapter.attempts) != 1 || adapter.attempts[0] != second.ID { + if len(adapter.attempts) != 1 || adapter.attempts[0] != fourth.ID { t.Fatalf("continued attempts = %#v", adapter.attempts) } if adapter.lastPromptCacheKey != expectedCacheKey || adapter.lastReasoningReplayKey != identity.replayKey { @@ -216,7 +236,7 @@ func TestGatewayFailsOverBeforeReturningBody(t *testing.T) { _, _ = io.ReadAll(resource.Body) resource.Finalize(Usage{}, "", "") _ = resource.Body.Close() - if adapter.lastPath != "/responses/resp-test?include=reasoning.encrypted_content" || adapter.lastMethod != http.MethodGet || len(adapter.attempts) != 1 || adapter.attempts[0] != second.ID { + if adapter.lastPath != "/responses/resp-test?include=reasoning.encrypted_content" || adapter.lastMethod != http.MethodGet || len(adapter.attempts) != 1 || adapter.attempts[0] != fourth.ID { t.Fatalf("resource request = %s %s, attempts = %#v", adapter.lastMethod, adapter.lastPath, adapter.attempts) } @@ -1387,6 +1407,7 @@ func runQuotaRefreshWorkers(t *testing.T, service *accountapp.Service) { type failoverAdapter struct { mu sync.Mutex firstID uint64 + failedIDs map[uint64]bool attempts []uint64 lastMethod string lastPath string @@ -1394,6 +1415,9 @@ type failoverAdapter struct { lastReasoningReplayKey string lastGrokTurnIndex string resourceStatus int + firstStatus int + firstBody string + firstHeader http.Header } type ssoUnauthorizedAdapter struct { @@ -1751,14 +1775,26 @@ func (a *failoverAdapter) ForwardResponse(_ context.Context, request provider.Re a.lastReasoningReplayKey = request.ReasoningReplayKey a.lastGrokTurnIndex = request.GrokTurnIndex resourceStatus := a.resourceStatus + firstStatus := a.firstStatus + firstBody := a.firstBody + firstHeader := a.firstHeader.Clone() + failed := request.Credential.ID == a.firstID || a.failedIDs[request.Credential.ID] a.mu.Unlock() status, body := http.StatusOK, "ok" + header := make(http.Header) if request.Method != http.MethodPost && resourceStatus != 0 { status, body = resourceStatus, "missing" - } else if request.Credential.ID == a.firstID { + } else if failed { status, body = http.StatusTooManyRequests, "limited" + if firstStatus != 0 { + status = firstStatus + } + if firstBody != "" { + body = firstBody + } + header = firstHeader } - return &provider.Response{StatusCode: status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil + return &provider.Response{StatusCode: status, Header: header, Body: io.NopCloser(strings.NewReader(body))}, nil } func (a *failoverAdapter) setResourceStatus(status int) {