From c99e479bcc14de1e154c901b1ad75d71c439a1f2 Mon Sep 17 00:00:00 2001 From: superman2003 <2112076433zcr@gmail.com> Date: Sun, 19 Jul 2026 12:17:42 +0800 Subject: [PATCH] fix(grok): retry transient billing probes --- .../internal/service/grok_quota_service.go | 80 +++++--- .../service/grok_quota_service_test.go | 172 ++++++++++++++++++ 2 files changed, 228 insertions(+), 24 deletions(-) diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index cf28f6c6e..2c703f9f1 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -23,6 +23,8 @@ const ( grokQuotaProbeInput = "hi" grokQuotaDefaultModel = grokDefaultResponsesModel grokBillingExtraKey = "grok_billing_snapshot" + grokBillingMaxAttempts = 2 + grokBillingRetryDelay = 100 * time.Millisecond ) type GrokQuotaProbeResult struct { @@ -356,33 +358,63 @@ func (s *GrokQuotaService) fetchBilling( if err != nil { return nil, 0, infraerrors.Newf(http.StatusBadRequest, "GROK_QUOTA_BASE_URL_INVALID", "invalid Grok base_url: %v", err) } - req, err := http.NewRequestWithContext(ctx, http.MethodGet, billingURL, nil) - if err != nil { - return nil, 0, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build billing request: %v", err) - } - xai.ApplyCLIBillingHeaders(req, token) - // billing 探测与真实转发保持同一套账号级请求头覆写。 - account.ApplyHeaderOverrides(req.Header) - resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 2)) - if err != nil { - return nil, 0, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", err) - } - defer func() { _ = resp.Body.Close() }() + for attempt := 0; attempt < grokBillingMaxAttempts; attempt++ { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, billingURL, nil) + if err != nil { + return nil, 0, infraerrors.Newf(http.StatusInternalServerError, "GROK_QUOTA_PROBE_REQUEST_BUILD_FAILED", "failed to build billing request: %v", err) + } + xai.ApplyCLIBillingHeaders(req, token) + // billing 探测与真实转发保持同一套账号级请求头覆写。 + account.ApplyHeaderOverrides(req.Header) + resp, requestErr := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 2)) - bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) - if resp.StatusCode == http.StatusTooManyRequests { - return nil, resp.StatusCode, nil + statusCode := 0 + var bodyBytes []byte + if requestErr == nil { + statusCode = resp.StatusCode + bodyBytes, _ = io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + _ = resp.Body.Close() + } + + shouldRetry := requestErr != nil || isRetryableGrokBillingStatus(statusCode) + if shouldRetry && attempt+1 < grokBillingMaxAttempts { + timer := time.NewTimer(grokBillingRetryDelay) + select { + case <-timer.C: + case <-ctx.Done(): + timer.Stop() + return nil, statusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", ctx.Err()) + } + continue + } + + if requestErr != nil { + return nil, 0, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed: %v", requestErr) + } + if statusCode == http.StatusTooManyRequests { + return nil, statusCode, nil + } + if statusCode >= 400 { + bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240) + slog.Warn("grok_quota_billing_failed", "account_id", account.ID, "weekly", weekly, "status", statusCode, "body", bodyText) + return nil, statusCode, infraerrors.Newf(mapUpstreamStatus(statusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing returned %d: %s", statusCode, bodyText) + } + payload, err := xai.ParseBillingPayload(bodyBytes) + if err != nil { + return nil, statusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_BILLING_PARSE_ERROR", "failed to parse billing body: %v", err) + } + return xai.BuildBillingSummary(payload.Config), statusCode, nil } - if resp.StatusCode >= 400 { - bodyText := truncate(strings.TrimSpace(string(bodyBytes)), 240) - slog.Warn("grok_quota_billing_failed", "account_id", account.ID, "weekly", weekly, "status", resp.StatusCode, "body", bodyText) - return nil, resp.StatusCode, infraerrors.Newf(mapUpstreamStatus(resp.StatusCode), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing returned %d: %s", resp.StatusCode, bodyText) + return nil, 0, infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_REQUEST_FAILED", "billing request failed") +} + +func isRetryableGrokBillingStatus(statusCode int) bool { + switch statusCode { + case http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout: + return true + default: + return false } - payload, err := xai.ParseBillingPayload(bodyBytes) - if err != nil { - return nil, resp.StatusCode, infraerrors.Newf(http.StatusBadGateway, "GROK_QUOTA_BILLING_PARSE_ERROR", "failed to parse billing body: %v", err) - } - return xai.BuildBillingSummary(payload.Config), resp.StatusCode, nil } func mergeGrokBillingProbeErrors(weeklyStatus, monthlyStatus int, weeklyErr, monthlyErr error) error { diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index ae3dacc4a..9ffb70f00 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -5,6 +5,7 @@ package service import ( "bytes" "context" + "errors" "io" "log/slog" "net/http" @@ -127,6 +128,44 @@ type grokHybridUpstream struct { billingHeaders http.Header } +type grokQuotaUpstreamStep struct { + status int + body string + err error +} + +type grokQuotaSequenceUpstream struct { + httpUpstreamRecorder + mu sync.Mutex + steps []grokQuotaUpstreamStep + requests []*http.Request +} + +func (u *grokQuotaSequenceUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + u.mu.Lock() + defer u.mu.Unlock() + u.requests = append(u.requests, req) + index := len(u.requests) - 1 + if index >= len(u.steps) { + return nil, errors.New("unexpected upstream request") + } + step := u.steps[index] + if step.err != nil { + return nil, step.err + } + return &http.Response{ + StatusCode: step.status, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(step.body)), + }, nil +} + +func (u *grokQuotaSequenceUpstream) snapshotRequests() []*http.Request { + u.mu.Lock() + defer u.mu.Unlock() + return append([]*http.Request(nil), u.requests...) +} + func (u *grokHybridUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { var body []byte if req != nil && req.Body != nil { @@ -228,6 +267,139 @@ func healthyGrokQuotaOAuthAccount(id int64) *Account { } } +func TestGrokQuotaServiceFetchBillingRetries502ThenSucceeds(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(401) + upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{ + {status: http.StatusBadGateway, body: `The origin web server returned an invalid or incomplete response to Cloudflare.`}, + {status: http.StatusOK, body: `{"config":{"currentPeriod":{"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"},"creditUsagePercent":12}}`}, + }} + svc := &GrokQuotaService{httpUpstream: upstream} + + summary, status, err := svc.fetchBilling(context.Background(), account, "access-token", "", true) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, status) + require.NotNil(t, summary) + require.NotNil(t, summary.UsagePercent) + require.Equal(t, 12.0, *summary.UsagePercent) + requests := upstream.snapshotRequests() + require.Len(t, requests, 2) + require.Equal(t, http.MethodGet, requests[0].Method) + require.Equal(t, "/v1/billing", requests[0].URL.Path) + require.Equal(t, "format=credits", requests[0].URL.RawQuery) +} + +func TestGrokQuotaServiceFetchBillingRetriesTransportErrorThenSucceeds(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(402) + upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{ + {err: errors.New("temporary transport failure")}, + {status: http.StatusOK, body: `{"config":{"currentPeriod":{"type":"WEEKLY"},"creditUsagePercent":8}}`}, + }} + svc := &GrokQuotaService{httpUpstream: upstream} + + summary, status, err := svc.fetchBilling(context.Background(), account, "access-token", "", true) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, status) + require.NotNil(t, summary) + require.Len(t, upstream.snapshotRequests(), 2) +} + +func TestGrokQuotaServiceFetchBillingStopsAfterSingleTransientRetry(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(403) + upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{ + {status: http.StatusBadGateway, body: `cloudflare failure`}, + {status: http.StatusBadGateway, body: `cloudflare failure`}, + {status: http.StatusOK, body: `{"config":{"currentPeriod":{"type":"WEEKLY"}}}`}, + }} + svc := &GrokQuotaService{httpUpstream: upstream} + + summary, status, err := svc.fetchBilling(context.Background(), account, "access-token", "", true) + + require.Error(t, err) + require.Nil(t, summary) + require.Equal(t, http.StatusBadGateway, status) + require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err)) + require.Contains(t, infraerrors.Message(err), "billing returned 502: cloudflare failure") + require.Len(t, upstream.snapshotRequests(), 2) +} + +func TestGrokQuotaServiceFetchBillingDoesNotRetryNonTransientStatuses(t *testing.T) { + tests := []struct { + name string + status int + wantErr bool + }{ + {name: "unauthorized", status: http.StatusUnauthorized, wantErr: true}, + {name: "forbidden", status: http.StatusForbidden, wantErr: true}, + {name: "rate limited", status: http.StatusTooManyRequests, wantErr: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(404) + upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{ + {status: tt.status, body: `{"error":{"message":"rejected"}}`}, + {status: http.StatusOK, body: `{"config":{"currentPeriod":{"type":"WEEKLY"}}}`}, + }} + svc := &GrokQuotaService{httpUpstream: upstream} + + summary, status, err := svc.fetchBilling(context.Background(), account, "access-token", "", true) + + if tt.wantErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + require.Nil(t, summary) + require.Equal(t, tt.status, status) + require.Len(t, upstream.snapshotRequests(), 1) + }) + } +} + +func TestIsRetryableGrokBillingStatus(t *testing.T) { + tests := []struct { + status int + want bool + }{ + {status: http.StatusBadGateway, want: true}, + {status: http.StatusServiceUnavailable, want: true}, + {status: http.StatusGatewayTimeout, want: true}, + {status: http.StatusUnauthorized, want: false}, + {status: http.StatusForbidden, want: false}, + {status: http.StatusTooManyRequests, want: false}, + } + + for _, tt := range tests { + t.Run(strconv.Itoa(tt.status), func(t *testing.T) { + require.Equal(t, tt.want, isRetryableGrokBillingStatus(tt.status)) + }) + } +} + +func TestGrokQuotaServiceProbeUsageDoesNotRetryResponsesPost(t *testing.T) { + account := healthyGrokQuotaOAuthAccount(405) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokQuotaSequenceUpstream{steps: []grokQuotaUpstreamStep{ + {status: http.StatusBadGateway, body: `cloudflare failure`}, + {status: http.StatusOK, body: `{"id":"unexpected_retry"}`}, + }} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream, nil) + + result, err := svc.ProbeUsage(context.Background(), account.ID) + + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err)) + requests := upstream.snapshotRequests() + require.Len(t, requests, 1) + require.Equal(t, http.MethodPost, requests[0].Method) + require.Equal(t, "/v1/responses", requests[0].URL.Path) +} + func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { t.Parallel()