From 1dedb2097dcf50845f5169f7ff25425a3857f187 Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Sat, 11 Jul 2026 20:41:32 +0800 Subject: [PATCH] fix(grok): persist quota exhaustion as rate limit --- backend/internal/repository/account_repo.go | 32 +++ .../account_repo_integration_test.go | 19 ++ .../internal/service/account_test_service.go | 11 +- .../service/account_test_service_grok_test.go | 83 ++++++ backend/internal/service/grok_media.go | 7 +- .../internal/service/grok_quota_service.go | 7 + .../service/grok_quota_service_test.go | 20 ++ .../service/openai_gateway_cc_pipeline.go | 7 +- .../openai_gateway_chat_completions_raw.go | 3 +- .../internal/service/openai_gateway_grok.go | 188 ++++++++++++- .../service/openai_gateway_grok_test.go | 264 +++++++++++++++++- .../service/openai_gateway_messages.go | 14 +- .../internal/service/openai_ws_http_bridge.go | 7 + .../__tests__/AccountStatusIndicator.spec.ts | 33 +++ 14 files changed, 649 insertions(+), 46 deletions(-) diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index c3c2f6708a..261be7c14b 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -1286,6 +1286,38 @@ func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetA return nil } +// SetRateLimitedIfLater atomically extends an account-level rate limit. Grok +// requests may finish concurrently, so an older response must not overwrite a +// later reset boundary observed by another request or instance. +func (r *accountRepository) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error { + now := time.Now() + updated, err := r.client.Account.Update(). + Where( + dbaccount.IDEQ(id), + dbaccount.Or( + dbaccount.RateLimitResetAtIsNil(), + dbaccount.RateLimitResetAtLT(resetAt), + ), + ). + SetRateLimitedAt(now). + SetRateLimitResetAt(resetAt). + Save(ctx) + if err != nil { + return err + } + if updated == 0 { + // This instance may not have observed the later value written elsewhere. + // Refresh its local scheduler snapshot even though no outbox event is needed. + r.syncSchedulerAccountSnapshot(ctx, id) + return nil + } + if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil { + logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue extended rate limit failed: account=%d err=%v", id, err) + } + r.syncSchedulerAccountSnapshot(ctx, id) + return nil +} + func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error { if scope == "" { return nil diff --git a/backend/internal/repository/account_repo_integration_test.go b/backend/internal/repository/account_repo_integration_test.go index de7fd3a9f3..f82ad4ca28 100644 --- a/backend/internal/repository/account_repo_integration_test.go +++ b/backend/internal/repository/account_repo_integration_test.go @@ -703,6 +703,25 @@ func (s *AccountRepoSuite) TestSetRateLimited() { s.Require().WithinDuration(resetAt, *got.RateLimitResetAt, time.Second) } +func (s *AccountRepoSuite) TestSetRateLimitedIfLaterDoesNotShortenReset() { + account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-rl-monotonic"}) + later := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second) + earlier := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second) + cacheRecorder := &schedulerCacheRecorder{} + s.repo.schedulerCache = cacheRecorder + + s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, later)) + s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, earlier)) + + got, err := s.repo.GetByID(s.ctx, account.ID) + s.Require().NoError(err) + s.Require().NotNil(got.RateLimitResetAt) + s.Require().WithinDuration(later, *got.RateLimitResetAt, time.Second) + s.Require().Len(cacheRecorder.setAccounts, 2) + s.Require().NotNil(cacheRecorder.setAccounts[1].RateLimitResetAt) + s.Require().WithinDuration(later, *cacheRecorder.setAccounts[1].RateLimitResetAt, time.Second) +} + func (s *AccountRepoSuite) TestClearRateLimit() { account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-clear"}) until := time.Now().Add(1 * time.Hour) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 15e8097b79..e947e2ec8e 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -723,10 +723,19 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * } defer func() { _ = resp.Body.Close() }() - if snapshot := xai.ParseQuotaHeaders(resp.Header, resp.StatusCode); snapshot != nil && s.accountRepo != nil { + now := time.Now() + snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now) + if snapshot != nil && s.accountRepo != nil { + resetAt, limited := grokRateLimitResetAt(snapshot, now) + if limited { + normalizeGrokExhaustedWindowResets(snapshot, resetAt, now) + } _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ grokQuotaSnapshotExtraKey: snapshot, }) + if limited { + persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) + } } if resp.StatusCode != http.StatusOK { diff --git a/backend/internal/service/account_test_service_grok_test.go b/backend/internal/service/account_test_service_grok_test.go index 0223070774..d0c812a003 100644 --- a/backend/internal/service/account_test_service_grok_test.go +++ b/backend/internal/service/account_test_service_grok_test.go @@ -3,6 +3,7 @@ package service import ( + "context" "io" "net/http" "net/http/httptest" @@ -15,6 +16,18 @@ import ( "github.com/tidwall/gjson" ) +type grokAccountTestRateLimitRepo struct { + *mockAccountRepoForGemini + rateLimitedCalls int + resetAt time.Time +} + +func (r *grokAccountTestRateLimitRepo) SetRateLimited(_ context.Context, _ int64, resetAt time.Time) error { + r.rateLimitedCalls++ + r.resetAt = resetAt + return nil +} + func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testing.T) { gin.SetMode(gin.TestMode) @@ -66,3 +79,73 @@ func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testin require.Contains(t, rec.Body.String(), `"model":"grok-4.3"`) require.Contains(t, rec.Body.String(), `"type":"test_complete"`) } + +func TestAccountTestService_Grok429PersistsRateLimitReset(t *testing.T) { + gin.SetMode(gin.TestMode) + + account := &Account{ + ID: 14, + Name: "grok-oauth-limited", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "grok-access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} + repo := &grokAccountTestRateLimitRepo{mockAccountRepoForGemini: baseRepo} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"45"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + }} + svc := &AccountTestService{ + accountRepo: repo, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + httpUpstream: upstream, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/14/test", nil) + + err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault) + + require.Error(t, err) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, time.Now().Add(45*time.Second), repo.resetAt, time.Second) +} + +func TestAccountTestService_Grok429WithoutQuotaHeadersUsesFallback(t *testing.T) { + gin.SetMode(gin.TestMode) + account := &Account{ + ID: 15, Name: "grok-oauth-limited-no-headers", Platform: PlatformGrok, + Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "grok-access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} + repo := &grokAccountTestRateLimitRepo{mockAccountRepoForGemini: baseRepo} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"quota exhausted"}}`)), + }} + svc := &AccountTestService{ + accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, nil), httpUpstream: upstream, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/15/test", nil) + before := time.Now() + + err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault) + + require.Error(t, err) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.resetAt, time.Second) +} diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 0f0e3366ee..01101ef6b5 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -357,11 +357,10 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( requestIDHeader := firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")) requestModel := requestInfo.Model if resp.StatusCode >= 400 { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel) } - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) if err != nil { return nil, err @@ -564,6 +563,9 @@ func (s *OpenAIGatewayService) handleGrokMediaErrorResponse( requestedModel string, ) (*OpenAIForwardResult, error) { body := s.readUpstreamErrorBody(resp) + // Reconcile readiness before configurable passthrough branches can return; + // otherwise a Grok 429 can remain schedulable. + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(body))) if upstreamMsg == "" { upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode) @@ -609,7 +611,6 @@ func (s *OpenAIGatewayService) handleGrokMediaErrorResponse( return nil, fmt.Errorf("upstream error: %d (not in custom error codes) message=%s", resp.StatusCode, upstreamMsg) } - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) kind := "http_error" if s.shouldFailoverUpstreamError(resp.StatusCode) { kind = "failover" diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 08eca2e7e2..17e91dac6f 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -91,9 +91,16 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr defer func() { _ = resp.Body.Close() }() snapshot := xai.ObserveQuotaHeaders(resp.Header, resp.StatusCode, "active_probe") + resetAt, limited := grokRateLimitResetAt(snapshot, time.Now()) + if limited { + normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now()) + } _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ grokQuotaSnapshotExtraKey: snapshot, }) + if limited { + persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) + } result := &GrokQuotaProbeResult{ Source: "active_probe", diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 5400507827..d306969efa 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -19,6 +19,10 @@ import ( type grokQuotaAccountRepo struct { *mockAccountRepoForPlatform updates map[int64]map[string]any + updateCalls int + rateLimitedCalls int + lastRateLimitedID int64 + lastRateLimitResetAt time.Time tempUnschedCalls int lastTempUnschedID int64 lastTempUnschedUntil time.Time @@ -26,6 +30,7 @@ type grokQuotaAccountRepo struct { } func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error { + r.updateCalls++ if r.updates == nil { r.updates = make(map[int64]map[string]any) } @@ -33,6 +38,17 @@ func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates return nil } +func (r *grokQuotaAccountRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error { + r.rateLimitedCalls++ + r.lastRateLimitedID = id + r.lastRateLimitResetAt = resetAt + return nil +} + +func (r *grokQuotaAccountRepo) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error { + return r.SetRateLimited(ctx, id, resetAt) +} + func (r *grokQuotaAccountRepo) SetTempUnschedulable(_ context.Context, id int64, until time.Time, reason string) error { r.tempUnschedCalls++ r.lastTempUnschedID = id @@ -286,6 +302,10 @@ func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { require.NotNil(t, result.Snapshot) require.NotNil(t, result.Snapshot.RetryAfterSeconds) require.Equal(t, 45, *result.Snapshot.RetryAfterSeconds) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Equal(t, account.ID, repo.lastRateLimitedID) + require.WithinDuration(t, time.Now().Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) } func TestGrokQuotaServiceResetQuotaUnsupported(t *testing.T) { diff --git a/backend/internal/service/openai_gateway_cc_pipeline.go b/backend/internal/service/openai_gateway_cc_pipeline.go index ea0148fb2a..690b9991bb 100644 --- a/backend/internal/service/openai_gateway_cc_pipeline.go +++ b/backend/internal/service/openai_gateway_cc_pipeline.go @@ -88,6 +88,9 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( upstreamMsg string, upstreamModel string, ) *UpstreamFailoverError { + if account != nil && account.Platform == PlatformGrok { + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } if !s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody) { return nil } @@ -109,7 +112,9 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( Message: upstreamMsg, Detail: upstreamDetail, }) - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) + if account.Platform != PlatformGrok { + s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) + } return &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 56130670df..554c1b4769 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -174,7 +174,6 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( if resp.StatusCode >= 400 { respBody, upstreamMsg := s.readOpenAIUpstreamError(resp) if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, AccountID: account.ID, @@ -201,7 +200,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( } if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) } // 8. Forward response diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 0484d373c1..1d0d69ffd4 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "io" + "log/slog" "net/http" "strings" "time" @@ -22,6 +23,7 @@ const ( grokComposerImageBridgeMaxOutputTokens = 512 grokUpstreamUserAgent = "sub2api-grok/1.0" grokCLIVersion = "0.2.93" + grokRateLimitFallbackCooldown = 2 * time.Minute ) func (s *OpenAIGatewayService) forwardGrokResponses( @@ -79,7 +81,6 @@ func (s *OpenAIGatewayService) forwardGrokResponses( if resp.StatusCode >= 400 { respBody := s.readUpstreamErrorBody(resp) resp.Body = io.NopCloser(bytes.NewReader(respBody)) - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody)) if upstreamMsg == "" { upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode) @@ -104,7 +105,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( return s.handleErrorResponse(ctx, resp, c, account, patchedBody, upstreamModel) } - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) var usage *OpenAIUsage var firstTokenMs *int @@ -485,7 +486,6 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( if resp.StatusCode >= 400 { respBody := s.readUpstreamErrorBody(resp) - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody)) if upstreamMsg == "" { upstreamMsg = fmt.Sprintf("xAI image bridge upstream returned status %d", resp.StatusCode) @@ -510,7 +510,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg) } - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil) if err != nil { return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err) @@ -664,33 +664,191 @@ func applyGrokCLIHeaders(headers http.Header) { headers.Set("X-Grok-Client-Version", grokCLIVersion) } -func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, accountID int64, snapshot *xai.QuotaSnapshot) { - if s == nil || s.accountRepo == nil || accountID <= 0 || snapshot == nil { +func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, account *Account, snapshot *xai.QuotaSnapshot) { + if s == nil || account == nil || account.ID <= 0 || snapshot == nil { return } - if s.codexSnapshotThrottle != nil && !s.codexSnapshotThrottle.Allow(accountID, time.Now()) { + accountID := account.ID + now := time.Now() + resetAt, hasActiveLimit := grokRateLimitResetAt(snapshot, now) + if hasActiveLimit { + normalizeGrokExhaustedWindowResets(snapshot, resetAt, now) + } + critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit + if s.codexSnapshotThrottle != nil { + allowed := s.codexSnapshotThrottle.Allow(accountID, now) + if !critical && !allowed { + return + } + } + + stateCtx := ctx + if hasActiveLimit { + var cancel context.CancelFunc + stateCtx, cancel = openAIAccountStateContext(ctx) + defer cancel() + } + if s.accountRepo != nil { + _ = s.accountRepo.UpdateExtra(stateCtx, accountID, map[string]any{ + grokQuotaSnapshotExtraKey: snapshot, + }) + } + // Error responses are reconciled by handleGrokAccountUpstreamError, which + // also installs the immediate in-memory scheduling block. Successful + // responses can still consume the last available request/token, so persist + // that exhausted window here as a real rate limit rather than relying only + // on the passive snapshot scheduler check. + if hasActiveLimit { + s.rateLimitGrok(stateCtx, account, resetAt) + } +} + +func parseGrokQuotaSnapshot(headers http.Header, statusCode int, now time.Time) *xai.QuotaSnapshot { + snapshot := xai.ParseQuotaHeaders(headers, statusCode) + if snapshot == nil && statusCode == http.StatusTooManyRequests { + return &xai.QuotaSnapshot{ + StatusCode: statusCode, + UpdatedAt: now.UTC().Format(time.RFC3339), + } + } + return snapshot +} + +func normalizeGrokExhaustedWindowResets(snapshot *xai.QuotaSnapshot, resetAt, now time.Time) { + if snapshot == nil || !resetAt.After(now) { return } - _ = s.accountRepo.UpdateExtra(ctx, accountID, map[string]any{ - grokQuotaSnapshotExtraKey: snapshot, - }) + for _, window := range []*xai.QuotaWindow{snapshot.Requests, snapshot.Tokens} { + if window == nil || window.Remaining == nil || *window.Remaining > 0 { + continue + } + candidate := time.Time{} + if window.ResetUnix != nil && *window.ResetUnix > 0 { + candidate = time.Unix(*window.ResetUnix, 0) + } else if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(window.ResetAt)); err == nil { + candidate = parsed + } + if !candidate.After(now) { + candidate = resetAt + } + resetUnix := candidate.Unix() + window.ResetUnix = &resetUnix + window.ResetAt = candidate.UTC().Format(time.RFC3339) + } +} + +func grokRateLimitResetAt(snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) { + if snapshot == nil { + return time.Time{}, false + } + + // Retry-After is xAI's explicit retry boundary. Use the observation time so + // a persisted snapshot does not start a fresh cooldown every time it is read. + retryAfterExpired := false + var resetAt time.Time + if snapshot.RetryAfterSeconds != nil && *snapshot.RetryAfterSeconds > 0 { + observedAt := now + if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(snapshot.UpdatedAt)); err == nil { + observedAt = parsed + } + retryAfterResetAt := observedAt.Add(time.Duration(*snapshot.RetryAfterSeconds) * time.Second) + if retryAfterResetAt.After(now) { + resetAt = retryAfterResetAt + } else { + retryAfterExpired = true + } + } + + exhausted := false + for _, window := range []*xai.QuotaWindow{snapshot.Requests, snapshot.Tokens} { + if window == nil || window.Remaining == nil || *window.Remaining > 0 { + continue + } + exhausted = true + candidate := time.Time{} + if window.ResetUnix != nil && *window.ResetUnix > 0 { + candidate = time.Unix(*window.ResetUnix, 0) + } else if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(window.ResetAt)); err == nil { + candidate = parsed + } + if candidate.After(now) && candidate.After(resetAt) { + resetAt = candidate + } + } + if !resetAt.IsZero() { + return resetAt, true + } + // An observed Retry-After is an absolute boundary once combined with the + // snapshot timestamp. Do not turn an expired persisted snapshot into a new + // rolling fallback cooldown, but still allow a later explicit window reset. + if retryAfterExpired { + return time.Time{}, false + } + if exhausted || snapshot.StatusCode == http.StatusTooManyRequests { + return now.Add(grokRateLimitFallbackCooldown), true + } + return time.Time{}, false +} + +func normalizeGrokRateLimitResetAt(account *Account, resetAt, now time.Time) time.Time { + if !resetAt.After(now) { + resetAt = now.Add(grokRateLimitFallbackCooldown) + } + if account != nil && account.RateLimitResetAt != nil && account.RateLimitResetAt.After(resetAt) { + resetAt = *account.RateLimitResetAt + } + return resetAt +} + +type grokRateLimitExtendingRepository interface { + SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error +} + +func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account *Account, resetAt time.Time) { + if repo == nil || account == nil || account.ID <= 0 { + return + } + resetAt = normalizeGrokRateLimitResetAt(account, resetAt, time.Now()) + stateCtx, cancel := openAIAccountStateContext(ctx) + defer cancel() + var err error + if extendingRepo, ok := repo.(grokRateLimitExtendingRepository); ok { + err = extendingRepo.SetRateLimitedIfLater(stateCtx, account.ID, resetAt) + } else { + err = repo.SetRateLimited(stateCtx, account.ID, resetAt) + } + if err != nil { + slog.Warn("persist_grok_rate_limit_failed", "account_id", account.ID, "reset_at", resetAt.UTC(), "error", err) + } +} + +func (s *OpenAIGatewayService) rateLimitGrok(ctx context.Context, account *Account, resetAt time.Time) { + if s == nil || account == nil { + return + } + resetAt = normalizeGrokRateLimitResetAt(account, resetAt, time.Now()) + + runtimeUntil := resetAt + if account.TempUnschedulableUntil != nil && account.TempUnschedulableUntil.After(runtimeUntil) { + runtimeUntil = *account.TempUnschedulableUntil + } + s.BlockAccountScheduling(account, runtimeUntil, "429") + persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) } func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) { if s == nil || account == nil { return } + now := time.Now() + s.updateGrokUsageSnapshot(ctx, account, parseGrokQuotaSnapshot(headers, statusCode, now)) switch statusCode { case http.StatusUnauthorized: s.tempUnscheduleGrok(ctx, account, 10*time.Minute, "grok oauth token unauthorized") case http.StatusForbidden: s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok entitlement or subscription tier denied") case http.StatusTooManyRequests: - cooldown := 2 * time.Minute - if snapshot := xai.ParseQuotaHeaders(headers, statusCode); snapshot != nil && snapshot.RetryAfterSeconds != nil && *snapshot.RetryAfterSeconds > 0 { - cooldown = time.Duration(*snapshot.RetryAfterSeconds) * time.Second - } - s.tempUnscheduleGrok(ctx, account, cooldown, "grok rate limited") + // updateGrokUsageSnapshot installs both runtime and durable rate-limit state. default: if statusCode >= 500 { s.tempUnscheduleGrok(ctx, account, 2*time.Minute, "grok upstream temporary error") diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index aa01e43461..43bdc8b33e 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -569,7 +569,7 @@ func TestBindGrokMediaVideoRequestAccountUsesRequestIDStickyHash(t *testing.T) { require.Equal(t, int64(63), accountID) } -func TestForwardGrokMediaErrorHonorsCustomErrorCodes(t *testing.T) { +func TestForwardGrokMedia429ReconcilesRateLimitBeforeCustomErrorBypass(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") gin.SetMode(gin.TestMode) @@ -589,18 +589,20 @@ func TestForwardGrokMediaErrorHonorsCustomErrorCodes(t *testing.T) { "api_key": "api-key", "base_url": "https://xai.test/v1", "custom_error_codes_enabled": true, - "custom_error_codes": []any{float64(http.StatusTooManyRequests)}, + "custom_error_codes": []any{float64(http.StatusBadRequest)}, }, } upstream := &httpUpstreamRecorder{resp: &http.Response{ - StatusCode: http.StatusBadRequest, + StatusCode: http.StatusTooManyRequests, Header: http.Header{ "Content-Type": []string{"application/json"}, "Xai-Request-Id": []string{"xai-error-req"}, + "Retry-After": []string{"45"}, }, Body: io.NopCloser(strings.NewReader(`{"error":{"message":"do not expose this upstream detail"}}`)), }} - svc := &OpenAIGatewayService{httpUpstream: upstream} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{httpUpstream: upstream, accountRepo: repo} result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json") require.Error(t, err) @@ -608,6 +610,9 @@ func TestForwardGrokMediaErrorHonorsCustomErrorCodes(t *testing.T) { require.Equal(t, http.StatusInternalServerError, recorder.Code) require.Contains(t, recorder.Body.String(), "Upstream gateway error") require.NotContains(t, recorder.Body.String(), "do not expose") + require.Equal(t, 1, repo.rateLimitedCalls) + require.Zero(t, repo.tempUnschedCalls) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) } func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *testing.T) { @@ -1137,7 +1142,7 @@ func grokMessagesSSECompletedResponse(responseID string, cachedTokens int) *http } } -func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing.T) { +func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *testing.T) { tests := []struct { name string status int @@ -1161,12 +1166,11 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing wantMaxCooldown: 30*time.Minute + time.Second, }, { - name: "rate limited retry after", - status: http.StatusTooManyRequests, - headers: http.Header{"Retry-After": []string{"45"}}, - wantReason: "grok rate limited", - wantMinCooldown: 44 * time.Second, - wantMaxCooldown: 46 * time.Second, + name: "upstream temporary error", + status: http.StatusInternalServerError, + wantReason: "grok upstream temporary error", + wantMinCooldown: 2*time.Minute - time.Second, + wantMaxCooldown: 2*time.Minute + time.Second, }, } @@ -1181,6 +1185,7 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) require.Equal(t, 1, repo.tempUnschedCalls) + require.Zero(t, repo.rateLimitedCalls) require.Equal(t, account.ID, repo.lastTempUnschedID) require.Equal(t, tt.wantReason, repo.lastTempUnschedReason) require.True(t, repo.lastTempUnschedUntil.After(before.Add(tt.wantMinCooldown))) @@ -1189,10 +1194,83 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing } } -func TestHandleGrokAccountUpstreamErrorDoesNotShortenExistingPause(t *testing.T) { +func TestHandleGrokAccountUpstreamError429SetsRateLimitedFromRetryAfter(t *testing.T) { + account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + before := time.Now() + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil) + + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Equal(t, account.ID, repo.lastRateLimitedID) + require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) +} + +func TestHandleGrokAccountUpstreamError429UsesLatestExhaustedWindowReset(t *testing.T) { + now := time.Now() + requestReset := now.Add(10 * time.Minute).Truncate(time.Second) + tokenReset := now.Add(20 * time.Minute).Truncate(time.Second) + headers := http.Header{ + "X-Ratelimit-Limit-Requests": []string{"10"}, + "X-Ratelimit-Remaining-Requests": []string{"0"}, + "X-Ratelimit-Reset-Requests": []string{fmt.Sprintf("%d", requestReset.Unix())}, + "X-Ratelimit-Limit-Tokens": []string{"1000"}, + "X-Ratelimit-Remaining-Tokens": []string{"0"}, + "X-Ratelimit-Reset-Tokens": []string{fmt.Sprintf("%d", tokenReset.Unix())}, + } + account := &Account{ID: 62, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, headers, nil) + + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, tokenReset, repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) +} + +func TestHandleGrokAccountUpstreamError429UsesFallbackReset(t *testing.T) { + account := &Account{ID: 63, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + before := time.Now() + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, nil, nil) + + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) +} + +func TestGrokRateLimitResetAtUsesFutureWindowAfterRetryAfterExpires(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + observedAt := now.Add(-2 * time.Minute) + windowReset := now.Add(15 * time.Minute) + retryAfter := 30 + snapshot := &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + UpdatedAt: observedAt.Format(time.RFC3339), + RetryAfterSeconds: &retryAfter, + Requests: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(10), + Remaining: grokInt64PtrForTest(0), + ResetUnix: grokInt64PtrForTest(windowReset.Unix()), + }, + } + + resetAt, limited := grokRateLimitResetAt(snapshot, now) + + require.True(t, limited) + require.WithinDuration(t, windowReset, resetAt, time.Second) +} + +func TestHandleGrokAccountUpstreamError429DoesNotShortenExistingPause(t *testing.T) { existingUntil := time.Now().Add(15 * time.Minute) account := &Account{ - ID: 62, + ID: 64, Platform: PlatformGrok, Type: AccountTypeOAuth, TempUnschedulableUntil: &existingUntil, @@ -1203,11 +1281,167 @@ func TestHandleGrokAccountUpstreamErrorDoesNotShortenExistingPause(t *testing.T) svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil) - require.Equal(t, 1, repo.tempUnschedCalls) - require.WithinDuration(t, existingUntil, repo.lastTempUnschedUntil, time.Second) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, time.Now().Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) value, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID) require.True(t, ok) runtimeUntil, ok := value.(time.Time) require.True(t, ok) require.WithinDuration(t, existingUntil, runtimeUntil, time.Second) } + +func TestUpdateGrokUsageSnapshotExhaustedSuccessBypassesThrottleAndSetsRateLimited(t *testing.T) { + account := &Account{ID: 65, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{ + accountRepo: repo, + codexSnapshotThrottle: newAccountWriteThrottle(time.Hour), + } + now := time.Now() + + // Consume the normal snapshot write allowance first. + svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{ + StatusCode: http.StatusOK, + Requests: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(10), + Remaining: grokInt64PtrForTest(9), + }, + UpdatedAt: now.UTC().Format(time.RFC3339), + }) + resetAt := now.Add(30 * time.Minute).Truncate(time.Second) + svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{ + StatusCode: http.StatusOK, + Requests: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(10), + Remaining: grokInt64PtrForTest(0), + ResetUnix: grokInt64PtrForTest(resetAt.Unix()), + ResetAt: resetAt.UTC().Format(time.RFC3339), + }, + UpdatedAt: now.UTC().Format(time.RFC3339), + }) + + require.Equal(t, 2, repo.updateCalls) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Equal(t, account.ID, repo.lastRateLimitedID) + require.WithinDuration(t, resetAt, repo.lastRateLimitResetAt, time.Second) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestUpdateGrokUsageSnapshotAvailableSuccessDoesNotSetRateLimited(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 66, Platform: PlatformGrok, Type: AccountTypeOAuth} + + svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{ + StatusCode: http.StatusOK, + Requests: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(10), + Remaining: grokInt64PtrForTest(1), + }, + UpdatedAt: time.Now().UTC().Format(time.RFC3339), + }) + + require.Equal(t, 1, repo.updateCalls) + require.Zero(t, repo.rateLimitedCalls) +} + +func TestUpdateGrokUsageSnapshotExhaustedSuccessWithoutResetUsesFallback(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 67, Platform: PlatformGrok, Type: AccountTypeOAuth} + before := time.Now() + + svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{ + StatusCode: http.StatusOK, + Tokens: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(2_000_000), + Remaining: grokInt64PtrForTest(0), + }, + UpdatedAt: before.UTC().Format(time.RFC3339), + }) + + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.lastRateLimitResetAt, time.Second) + stored, ok := repo.updates[account.ID][grokQuotaSnapshotExtraKey].(*xai.QuotaSnapshot) + require.True(t, ok) + require.NotNil(t, stored.Tokens.ResetUnix) + paused, _ := shouldAutoPauseGrokQuotaWindow("tokens", stored.Tokens, before.Add(time.Second)) + require.True(t, paused) + paused, _ = shouldAutoPauseGrokQuotaWindow("tokens", stored.Tokens, repo.lastRateLimitResetAt.Add(time.Second)) + require.False(t, paused) +} + +func TestOpenAIWSHTTPBridgeGrok429PersistsRateLimit(t *testing.T) { + repo := &grokQuotaAccountRepo{} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"45"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + }} + svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: upstream} + account := &Account{ID: 68, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1} + before := time.Now() + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), nil, account, "token", + []byte(`{"type":"response.create","model":"grok-4.3","input":"hi"}`), + 64, "grok-4.3", "", "", "", "cache-id", 1, + func([]byte) error { return nil }, + ) + + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestOpenAIWSHTTPBridgeGrokExhaustedSuccessPersistsRateLimit(t *testing.T) { + repo := &grokQuotaAccountRepo{} + resetAt := time.Now().Add(20 * time.Minute).UTC().Truncate(time.Second) + resp := grokMessagesSSECompletedResponse("resp_ws_limited", 0) + resp.Header.Set("X-Ratelimit-Limit-Requests", "10") + resp.Header.Set("X-Ratelimit-Remaining-Requests", "0") + resp.Header.Set("X-Ratelimit-Reset-Requests", fmt.Sprintf("%d", resetAt.Unix())) + upstream := &httpUpstreamRecorder{resp: resp} + svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: upstream} + account := &Account{ID: 69, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1} + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), nil, account, "token", + []byte(`{"type":"response.create","model":"grok-4.3","input":"hi"}`), + 64, "grok-4.3", "", "", "", "cache-id", 1, + func([]byte) error { return nil }, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, resetAt, repo.lastRateLimitResetAt, time.Second) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestFailoverOpenAIUpstreamHTTPErrorUsesOnlyGrokRateLimitPolicy(t *testing.T) { + gin.SetMode(gin.TestMode) + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 70, Platform: PlatformGrok, Type: AccountTypeOAuth} + resp := &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"45"}}, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + failoverErr := svc.failoverOpenAIUpstreamHTTPError( + context.Background(), c, account, resp, + []byte(`{"error":{"message":"rate limited"}}`), "rate limited", "grok-4.3", + ) + + require.NotNil(t, failoverErr) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Zero(t, repo.tempUnschedCalls) +} diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 3615fe4b17..23b0aa54d6 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -321,11 +321,6 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // 8. Handle error response with failover if resp.StatusCode >= 400 { respBody, upstreamMsg := s.readOpenAIUpstreamError(resp) - if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) - } - if previousResponseID != "" && (isOpenAICompatPreviousResponseNotFound(resp.StatusCode, upstreamMsg, respBody) || isOpenAICompatPreviousResponseUnsupported(resp.StatusCode, upstreamMsg, respBody)) { if isOpenAICompatPreviousResponseUnsupported(resp.StatusCode, upstreamMsg, respBody) { s.disableOpenAICompatSessionContinuation(ctx, c, account, promptCacheKey) @@ -345,6 +340,9 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // Non-failover error: return Anthropic-formatted error to client return s.handleAnthropicErrorResponse(resp, c, account, billingModel) } + if account.Platform == PlatformGrok && account.Type == AccountTypeOAuth && !account.IsShadow() { + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + } if account.Type == AccountTypeOAuth && promptCacheKey != "" { if turnState := strings.TrimSpace(resp.Header.Get("x-codex-turn-state")); turnState != "" { @@ -392,10 +390,8 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // Extract and save Codex usage snapshot from response headers (for OAuth accounts). // 排除 spark 影子:其 codex_* 仅由 QueryUsage(/wham/usage bengalfox)更新(外审第7轮 P1)。 - if handleErr == nil && account.Type == AccountTypeOAuth && !account.IsShadow() { - if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) - } else if snapshot := ParseCodexRateLimitHeaders(resp.Header); snapshot != nil { + if handleErr == nil && account.Type == AccountTypeOAuth && !account.IsShadow() && account.Platform != PlatformGrok { + if snapshot := ParseCodexRateLimitHeaders(resp.Header); snapshot != nil { s.updateCodexUsageSnapshot(ctx, account.ID, snapshot) } } diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index 94bb09e358..0afb9181a1 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -12,6 +12,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" ) @@ -221,6 +222,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( if resp.StatusCode >= 400 { respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes)) + if account.Platform == PlatformGrok { + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) if upstreamMsg == "" { upstreamMsg = http.StatusText(resp.StatusCode) @@ -228,6 +232,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( _ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg)) return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg) } + if account.Platform == PlatformGrok { + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + } responseID := "" usage := OpenAIUsage{} diff --git a/frontend/src/components/account/__tests__/AccountStatusIndicator.spec.ts b/frontend/src/components/account/__tests__/AccountStatusIndicator.spec.ts index f758e6b0f6..545c2abac4 100644 --- a/frontend/src/components/account/__tests__/AccountStatusIndicator.spec.ts +++ b/frontend/src/components/account/__tests__/AccountStatusIndicator.spec.ts @@ -13,6 +13,14 @@ vi.mock('vue-i18n', async () => { } }) +vi.mock('@/utils/format', async () => { + const actual = await vi.importActual('@/utils/format') + return { + ...actual, + formatCountdown: () => '1h' + } +}) + function makeAccount(overrides: Partial): Account { return { id: 1, @@ -43,6 +51,31 @@ function makeAccount(overrides: Partial): Account { } describe('AccountStatusIndicator', () => { + it('Grok 账号额度限流时显示自动恢复时间而非临时不可调度', () => { + const wrapper = mount(AccountStatusIndicator, { + props: { + account: makeAccount({ + id: 5, + name: 'grok-free-1', + platform: 'grok', + rate_limited_at: '2026-07-11T12:00:00Z', + rate_limit_reset_at: '2099-07-11T13:00:00Z', + temp_unschedulable_until: '2099-07-11T12:30:00Z', + temp_unschedulable_reason: 'legacy grok rate limited' + }) + }, + global: { + stubs: { + Icon: true + } + } + }) + + expect(wrapper.find('.badge-warning').text()).toBe('admin.accounts.status.rateLimited') + expect(wrapper.text()).toContain('admin.accounts.status.rateLimitedAutoResume') + expect(wrapper.text()).not.toContain('admin.accounts.status.tempUnschedulable') + }) + it('模型限流 + overages 启用 + 无 AICredits key → 显示 ⚡ (credits_active)', () => { const wrapper = mount(AccountStatusIndicator, { props: {