fix(grok): retry transient billing probes
This commit is contained in:
superman2003 committed
1 parent
b8c640d218
commit
c99e479bcc
2 files changed
+228
-24
No files matched your search
@@ -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 {
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in new issue
Block a user