diff --git a/backend/internal/handler/openai_gateway_count_tokens.go b/backend/internal/handler/openai_gateway_count_tokens.go index a9c2710b9..bfbd2d72a 100644 --- a/backend/internal/handler/openai_gateway_count_tokens.go +++ b/backend/internal/handler/openai_gateway_count_tokens.go @@ -13,6 +13,48 @@ import ( "go.uber.org/zap" ) +// GrokCountTokens handles Anthropic-compatible count_tokens requests locally. +// The route middleware already authenticates the API key and resolves the +// group; this handler intentionally does not select an account or check billing. +func (h *OpenAIGatewayHandler) GrokCountTokens(c *gin.Context) { + body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg) + if err != nil { + if maxErr, ok := extractMaxBytesError(err); ok { + h.anthropicErrorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) + return + } + h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body") + return + } + if len(body) == 0 { + h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty") + return + } + + bodyRef := service.NewRequestBodyRef(body) + parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic) + if err != nil { + logRequestBodyParseFailure(requestLogger(c, "handler.openai_gateway.grok_count_tokens"), body, err) + h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } + if parsedReq.Model == "" { + h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required") + return + } + + estimated, err := service.EstimateGrokCountTokens(parsedReq.Body.Bytes()) + if err != nil { + requestLogger(c, "handler.openai_gateway.grok_count_tokens").Warn("grok_count_tokens.local_estimate_failed", zap.Error(err)) + h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } + + setOpsRequestContext(c, parsedReq.Model, false) + setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(false, false))) + c.JSON(http.StatusOK, gin.H{"input_tokens": estimated}) +} + // CountTokens handles Anthropic-compatible POST /v1/messages/count_tokens for OpenAI groups. // It validates billing and routes to an OpenAI token-count bridge without taking concurrency slots // or recording usage. diff --git a/backend/internal/pkg/xai/quota.go b/backend/internal/pkg/xai/quota.go index 1387c5bda..fe269d1d8 100644 --- a/backend/internal/pkg/xai/quota.go +++ b/backend/internal/pkg/xai/quota.go @@ -7,6 +7,18 @@ import ( "time" ) +const GrokFreeRolling24hTokenLimit int64 = 1_000_000 + +var grokFreeRolling24hTokenLimits = map[int64]struct{}{ + GrokFreeRolling24hTokenLimit: {}, + 2_000_000: {}, // Legacy Free limit observed before July 2026. +} + +func IsGrokFreeRolling24hTokenLimit(limit int64) bool { + _, ok := grokFreeRolling24hTokenLimits[limit] + return ok +} + type QuotaWindow struct { Limit *int64 `json:"limit,omitempty"` Remaining *int64 `json:"remaining,omitempty"` diff --git a/backend/internal/pkg/xai/quota_test.go b/backend/internal/pkg/xai/quota_test.go index 593edc8eb..983a87e84 100644 --- a/backend/internal/pkg/xai/quota_test.go +++ b/backend/internal/pkg/xai/quota_test.go @@ -61,3 +61,11 @@ func TestObserveQuotaHeadersRecordsNoHeaderProbe(t *testing.T) { require.Nil(t, snapshot.Requests) require.Nil(t, snapshot.Tokens) } + +func TestIsGrokFreeRolling24hTokenLimit(t *testing.T) { + t.Parallel() + + require.True(t, IsGrokFreeRolling24hTokenLimit(GrokFreeRolling24hTokenLimit)) + require.True(t, IsGrokFreeRolling24hTokenLimit(2_000_000), "legacy snapshots remain classifiable") + require.False(t, IsGrokFreeRolling24hTokenLimit(3_000_000)) +} diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index 07b8f1ebd..8e14f1d5a 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -43,6 +43,16 @@ func RegisterGatewayRoutes( isOpenAIGatewayPlatform := func(c *gin.Context) bool { return getGroupPlatform(c) == service.PlatformOpenAI } + countTokensHandler := func(c *gin.Context) { + switch getGroupPlatform(c) { + case service.PlatformOpenAI: + h.OpenAIGateway.CountTokens(c) + case service.PlatformGrok: + h.OpenAIGateway.GrokCountTokens(c) + default: + h.Gateway.CountTokens(c) + } + } modelsHandler := func(c *gin.Context) { if isOpenAIGatewayPlatform(c) && c.Query("client_version") != "" { h.OpenAIGateway.CodexModels(c) @@ -139,26 +149,9 @@ func RegisterGatewayRoutes( } h.Gateway.Messages(c) }) - // /v1/messages/count_tokens: OpenAI uses Anthropic-compat bridge; other - // OpenAI-compatible platforms keep the prior unsupported response. - gateway.POST("/messages/count_tokens", func(c *gin.Context) { - if isOpenAIGatewayPlatform(c) { - h.OpenAIGateway.CountTokens(c) - return - } - if isOpenAIResponsesCompatibleGatewayPlatform(c) { - service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusNotFound, gin.H{ - "type": "error", - "error": gin.H{ - "type": "not_found_error", - "message": "Token counting is not supported for this platform", - }, - }) - return - } - h.Gateway.CountTokens(c) - }) + // /v1/messages/count_tokens: OpenAI bridges upstream, Grok estimates + // locally, and Anthropic-compatible platforms retain their existing path. + gateway.POST("/messages/count_tokens", countTokensHandler) // Codex CLI / Codex app refresh their model picker from the provider's // /models endpoint with a client_version query and expect the ChatGPT // Codex manifest format; other clients keep the OpenAI-style list. @@ -256,6 +249,7 @@ func RegisterGatewayRoutes( h.OpenAIGateway.ResponsesWebSocket(c) }) r.GET("/models", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, modelsHandler) + r.POST("/messages/count_tokens", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, countTokensHandler) codexDirect := r.Group("/backend-api/codex") codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic) { diff --git a/backend/internal/server/routes/gateway_test.go b/backend/internal/server/routes/gateway_test.go index 80e319db7..43a43ff78 100644 --- a/backend/internal/server/routes/gateway_test.go +++ b/backend/internal/server/routes/gateway_test.go @@ -1,6 +1,7 @@ package routes import ( + "encoding/json" "net/http" "net/http/httptest" "strings" @@ -15,6 +16,10 @@ import ( ) func newGatewayRoutesTestRouter(platform ...string) *gin.Engine { + return newGatewayRoutesTestRouterWithConfig(&config.Config{}, platform...) +} + +func newGatewayRoutesTestRouterWithConfig(cfg *config.Config, platform ...string) *gin.Engine { gin.SetMode(gin.TestMode) router := gin.New() @@ -42,7 +47,7 @@ func newGatewayRoutesTestRouter(platform ...string) *gin.Engine { nil, nil, nil, - &config.Config{}, + cfg, ) return router @@ -224,13 +229,22 @@ func TestGatewayRoutesGrokAllowsCLICompatibilityEntrypoints(t *testing.T) { require.NotContains(t, w.Body.String(), "not supported for Grok groups") } - req := httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", strings.NewReader(`{"model":"grok","messages":[{"role":"user","content":"hi"}]}`)) - req.Header.Set("Content-Type", "application/json") - w := httptest.NewRecorder() + countTokensRouter := newGatewayRoutesTestRouterWithConfig(&config.Config{ + Gateway: config.GatewayConfig{MaxBodySize: 1024 * 1024}, + }, service.PlatformGrok) + for _, path := range []string{"/v1/messages/count_tokens", "/messages/count_tokens"} { + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok","messages":[{"role":"user","content":"hi"}]}`)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() - router.ServeHTTP(w, req) - require.Equal(t, http.StatusNotFound, w.Code) - require.Contains(t, w.Body.String(), "Token counting is not supported for this platform") + countTokensRouter.ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code, "path=%s", path) + var response struct { + InputTokens int `json:"input_tokens"` + } + require.NoError(t, json.Unmarshal(w.Body.Bytes(), &response), "path=%s", path) + require.Positive(t, response.InputTokens, "path=%s", path) + } for _, path := range []string{ "/v1/responses", diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index 12dc25a67..e9502ec4c 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -207,6 +207,7 @@ type UsageInfo struct { GrokLastQuotaProbeAt string `json:"grok_last_quota_probe_at,omitempty"` GrokLastHeadersSeenAt string `json:"grok_last_headers_seen_at,omitempty"` GrokLastStatusCode int `json:"grok_last_status_code,omitempty"` + GrokFreeTokenLimit int64 `json:"grok_free_token_limit,omitempty"` GrokLocalUsage *WindowStats `json:"grok_local_usage,omitempty"` GrokLocalUsage24h *WindowStats `json:"grok_local_usage_24h,omitempty"` GrokLocalUsage7d *WindowStats `json:"grok_local_usage_7d,omitempty"` diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index 013cb0404..023c6acc1 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -576,6 +576,14 @@ var proxyQualityTargets = []proxyQualityTarget{ http.StatusOK: {}, }, }, + { + Target: "grok", + URL: "https://api.x.ai/v1/models", + Method: http.MethodGet, + AllowedStatuses: map[int]struct{}{ + http.StatusUnauthorized: {}, + }, + }, } const ( diff --git a/backend/internal/service/admin_service_proxy_quality_test.go b/backend/internal/service/admin_service_proxy_quality_test.go index 023b448fb..2347c2deb 100644 --- a/backend/internal/service/admin_service_proxy_quality_test.go +++ b/backend/internal/service/admin_service_proxy_quality_test.go @@ -93,3 +93,44 @@ func TestRunProxyQualityTarget_AllowedStatusPassForUnauthorized(t *testing.T) { require.Equal(t, http.StatusUnauthorized, item.HTTPStatus) require.Contains(t, item.Message, "目标可达") } + +func TestProxyQualityTargets_IncludesGrok(t *testing.T) { + var grokTarget *proxyQualityTarget + for i := range proxyQualityTargets { + if proxyQualityTargets[i].Target == "grok" { + grokTarget = &proxyQualityTargets[i] + break + } + } + + require.NotNil(t, grokTarget) + require.Equal(t, "https://api.x.ai/v1/models", grokTarget.URL) + require.Equal(t, http.MethodGet, grokTarget.Method) + require.Contains(t, grokTarget.AllowedStatuses, http.StatusUnauthorized) +} + +func TestRunProxyQualityTarget_GrokUnauthorizedPasses(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + t.Errorf("unexpected method: %s", r.Method) + } + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`{"error":"unauthorized"}`)) + })) + defer server.Close() + + target := proxyQualityTarget{ + Target: "grok", + URL: server.URL, + Method: http.MethodGet, + AllowedStatuses: map[int]struct{}{ + http.StatusUnauthorized: {}, + }, + } + + item := runProxyQualityTarget(context.Background(), server.Client(), target) + require.Equal(t, "grok", item.Target) + require.Equal(t, "pass", item.Status) + require.Equal(t, http.StatusUnauthorized, item.HTTPStatus) + require.Contains(t, item.Message, "目标可达") +} diff --git a/backend/internal/service/grok_quota_fetcher.go b/backend/internal/service/grok_quota_fetcher.go index 1f2d6374a..545b1c9e9 100644 --- a/backend/internal/service/grok_quota_fetcher.go +++ b/backend/internal/service/grok_quota_fetcher.go @@ -21,8 +21,9 @@ func NewGrokQuotaFetcher() *GrokQuotaFetcher { func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { now := time.Now() usage := &UsageInfo{ - Source: "passive", - UpdatedAt: &now, + Source: "passive", + UpdatedAt: &now, + GrokFreeTokenLimit: xai.GrokFreeRolling24hTokenLimit, } if account == nil { usage.ErrorCode = "quota_unknown" diff --git a/backend/internal/service/grok_quota_fetcher_test.go b/backend/internal/service/grok_quota_fetcher_test.go index 6bea2814f..a54ecc6c8 100644 --- a/backend/internal/service/grok_quota_fetcher_test.go +++ b/backend/internal/service/grok_quota_fetcher_test.go @@ -19,6 +19,7 @@ func TestGrokQuotaFetcherBuildUsageInfoUnknownUntilFirstSnapshot(t *testing.T) { usage := NewGrokQuotaFetcher().BuildUsageInfo(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth}) require.Equal(t, "passive", usage.Source) + require.Equal(t, xai.GrokFreeRolling24hTokenLimit, usage.GrokFreeTokenLimit) require.Equal(t, "quota_unknown", usage.ErrorCode) require.Contains(t, usage.Error, "unknown until billing is probed") } 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() diff --git a/backend/internal/service/openai_gateway_count_tokens.go b/backend/internal/service/openai_gateway_count_tokens.go index 46b5a3843..eccd1c414 100644 --- a/backend/internal/service/openai_gateway_count_tokens.go +++ b/backend/internal/service/openai_gateway_count_tokens.go @@ -39,6 +39,39 @@ type openAIInputTokensCountPrepared struct { UpstreamModel string } +// EstimateGrokCountTokens estimates an Anthropic-compatible count_tokens request +// locally. Grok does not expose a compatible token-counting endpoint, so this +// path deliberately avoids account selection, credentials, and upstream calls. +func EstimateGrokCountTokens(body []byte) (int, error) { + var anthropicReq apicompat.AnthropicRequest + if err := json.Unmarshal(body, &anthropicReq); err != nil { + return 0, fmt.Errorf("parse anthropic count_tokens request: %w", err) + } + if strings.TrimSpace(anthropicReq.Model) == "" { + return 0, fmt.Errorf("parse anthropic count_tokens request: model is required") + } + + responsesReq, err := apicompat.AnthropicToResponses(&anthropicReq) + if err != nil { + return 0, fmt.Errorf("convert anthropic request to responses: %w", err) + } + + estimated, err := estimateOpenAIInputTokens(openAIInputTokensCountRequest{ + Model: anthropicReq.Model, + Instructions: responsesReq.Instructions, + Input: responsesReq.Input, + Tools: responsesReq.Tools, + ToolChoice: responsesReq.ToolChoice, + }) + if err != nil { + return 0, fmt.Errorf("estimate grok input tokens: %w", err) + } + if estimated < openAIInputTokensFallbackMinimum { + estimated = openAIInputTokensFallbackMinimum + } + return estimated, nil +} + // ForwardCountTokensAsAnthropic bridges Anthropic /v1/messages/count_tokens to // OpenAI POST /v1/responses/input_tokens and returns Anthropic-compatible output. func (s *OpenAIGatewayService) ForwardCountTokensAsAnthropic( diff --git a/backend/internal/service/openai_gateway_count_tokens_test.go b/backend/internal/service/openai_gateway_count_tokens_test.go index 3ab0178a0..9719c9b3f 100644 --- a/backend/internal/service/openai_gateway_count_tokens_test.go +++ b/backend/internal/service/openai_gateway_count_tokens_test.go @@ -236,6 +236,51 @@ func TestEstimateOpenAIInputTokens_RequestSamples(t *testing.T) { } } +func TestEstimateGrokCountTokens_AnthropicRequests(t *testing.T) { + cases := []struct { + name string + body string + }{ + { + name: "simple message", + body: `{"model":"grok-4","messages":[{"role":"user","content":"hello world"}]}`, + }, + { + name: "system blocks and tools", + body: `{ + "model":"grok-4", + "system":[{"type":"text","text":"You are helpful."}], + "messages":[{"role":"user","content":[{"type":"text","text":"look up the weather"}]}], + "tools":[{"name":"lookup_weather","description":"Look up weather","input_schema":{"type":"object","properties":{"city":{"type":"string"}}}}], + "tool_choice":{"type":"auto"} + }`, + }, + { + name: "empty conversation uses positive minimum", + body: `{"model":"grok-4","messages":[]}`, + }, + } + + for _, tt := range cases { + t.Run(tt.name, func(t *testing.T) { + got, err := EstimateGrokCountTokens([]byte(tt.body)) + require.NoError(t, err) + require.Positive(t, got) + }) + } +} + +func TestEstimateGrokCountTokens_RejectsInvalidRequests(t *testing.T) { + for _, body := range []string{ + `{`, + `{"messages":[{"role":"user","content":"hello"}]}`, + `{"model":"grok-4","messages":[{"role":"user","content":{"unexpected":true}}]}`, + } { + _, err := EstimateGrokCountTokens([]byte(body)) + require.Error(t, err, "body=%s", body) + } +} + func TestOpenAIInputTokensEncodingForModel(t *testing.T) { cases := []struct { model string diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go index 9fb916cac..de52115f1 100644 --- a/backend/internal/service/openai_gateway_grok_cache.go +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -6,6 +6,7 @@ import ( "net/http" "strings" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -17,7 +18,6 @@ const ( grokFreeCacheNativeToolsJSON = `[{"type":"web_search"},{"type":"x_search"}]` grokFreeCacheDisabledToolChoice = "none" grokClientToolCacheOptInExtraKey = "grok_client_tool_cache_enabled" - grokFreeRolling24hTokenLimit = int64(2_000_000) ) // resolveGrokCacheIdentity derives one stable, tenant-isolated routing identity @@ -295,7 +295,7 @@ func isKnownGrokFreeAccount(account *Account) bool { } } if snapshot.Tokens != nil && snapshot.Tokens.Limit != nil && - *snapshot.Tokens.Limit == grokFreeRolling24hTokenLimit { + xai.IsGrokFreeRolling24hTokenLimit(*snapshot.Tokens.Limit) { inferredFreeSignal = true } } diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go index 90155aadc..034ba4843 100644 --- a/backend/internal/service/openai_gateway_grok_cache_test.go +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -8,6 +8,7 @@ import ( "sync" "testing" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" @@ -677,7 +678,19 @@ func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing. a := healthyGrokOAuthGatewayTestAccount(9112, "access-token") a.Extra = map[string]any{grokQuotaSnapshotExtraKey: map[string]any{ "headers_observed": true, - "tokens": map[string]any{"limit": grokFreeRolling24hTokenLimit}, + "tokens": map[string]any{"limit": xai.GrokFreeRolling24hTokenLimit}, + }} + return a + }(), + wantMix: true, + }, + { + name: "legacy free rolling token quota", + account: func() *Account { + a := healthyGrokOAuthGatewayTestAccount(9113, "access-token") + a.Extra = map[string]any{grokQuotaSnapshotExtraKey: map[string]any{ + "headers_observed": true, + "tokens": map[string]any{"limit": int64(2_000_000)}, }} return a }(), @@ -699,7 +712,7 @@ func TestGrokFreeMessagesFunctionToolCacheRouteRequiresKnownFreeTier(t *testing. grokBillingExtraKey: map[string]any{"plan": "SuperGrok", "status_code": http.StatusOK}, grokQuotaSnapshotExtraKey: map[string]any{ "headers_observed": true, - "tokens": map[string]any{"limit": grokFreeRolling24hTokenLimit}, + "tokens": map[string]any{"limit": int64(2_000_000)}, }, } return a diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 9a3d6a86e..9a3ac13f0 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -408,7 +408,7 @@ () const USAGE_CACHE_TTL = 5 * 60 * 1000 // 5 minutes -// xAI Free billing exposes a window without usage_percent, so estimate it from local tokens. -const GROK_FREE_TOKEN_LIMIT = 2_000_000 const props = withDefaults( defineProps<{ @@ -1112,8 +1110,10 @@ const grokLocalUsage = computed(() => { }) const grokFreeTokenBar = computed(() => { if (!grokIsFree.value || !grokFreeQuotaUsage.value) return null + const limit = usageInfo.value?.grok_free_token_limit + if (typeof limit !== 'number' || limit <= 0) return null const used = Math.max(0, grokFreeQuotaUsage.value.tokens || 0) - return { utilization: Math.min(100, (used / GROK_FREE_TOKEN_LIMIT) * 100) } + return { utilization: Math.min(100, (used / limit) * 100), limit } }) const grokQuotaUnknown = computed(() => { if (props.account.platform !== 'grok') return false diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index 5f61494c5..203b15018 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -705,11 +705,12 @@ describe('AccountUsageCell', () => { it.each([ { tokens: 0, expected: 0, compact: '0' }, - { tokens: 1_000_000, expected: 50, compact: '1.0M' }, - { tokens: 2_000_000, expected: 100, compact: '2.0M' }, - { tokens: 2_200_000, expected: 100, compact: '2.2M' } - ])('Grok Free derives its 2M quota from local tokens: $tokens -> $expected%', async ({ tokens, expected, compact }) => { + { tokens: 500_000, expected: 50, compact: '500.0K' }, + { tokens: 1_000_000, expected: 100, compact: '1.0M' }, + { tokens: 1_100_000, expected: 100, compact: '1.1M' } + ])('Grok Free derives its 1M quota from local tokens: $tokens -> $expected%', async ({ tokens, expected, compact }) => { getUsage.mockResolvedValue({ + grok_free_token_limit: 1_000_000, grok_billing: { period_type: 'weekly', usage_percent: null, @@ -723,7 +724,7 @@ describe('AccountUsageCell', () => { user_cost: 0 }, grok_request_quota: { limit: 100, remaining: 100 }, - grok_token_quota: { limit: 2_000_000, remaining: 2_000_000 } + grok_token_quota: { limit: 1_000_000, remaining: 1_000_000 } }) const wrapper = mount(AccountUsageCell, { @@ -753,6 +754,7 @@ describe('AccountUsageCell', () => { it('Grok Free uses rolling 24h usage instead of today-only usage', async () => { getUsage.mockResolvedValue({ + grok_free_token_limit: 1_000_000, grok_billing: { period_type: 'weekly', usage_percent: null, plan: '' }, grok_local_usage: { requests: 2, @@ -762,7 +764,7 @@ describe('AccountUsageCell', () => { }, grok_local_usage_24h: { requests: 12, - tokens: 1_500_000, + tokens: 750_000, cost: 0, standard_cost: 0 } @@ -793,7 +795,7 @@ describe('AccountUsageCell', () => { await flushPromises() expect(wrapper.text()).toContain('24h|75|admin.accounts.usageWindow.grokFreeQuota24hHint') - expect(wrapper.text()).toContain('1.5M') + expect(wrapper.text()).toContain('750.0K') expect(wrapper.text()).not.toContain('7d|') expect(wrapper.text()).not.toContain('200.0K') expect(wrapper.text()).not.toContain('250.0K') @@ -801,6 +803,7 @@ describe('AccountUsageCell', () => { it('Grok Free does not substitute today stats when rolling 24h usage is unavailable', async () => { getUsage.mockResolvedValue({ + grok_free_token_limit: 1_000_000, grok_billing: { period_type: 'weekly', usage_percent: null, plan: '' }, grok_local_usage: { requests: 1, @@ -921,8 +924,9 @@ describe('AccountUsageCell', () => { expect(wrapper.text()).not.toContain('2M|') }) - it('Grok credential Free tier keeps the 2M fallback when billing is unavailable', async () => { + it('Grok credential Free tier keeps the 1M fallback when billing is unavailable', async () => { getUsage.mockResolvedValue({ + grok_free_token_limit: 1_000_000, subscription_tier: 'FREE', grok_local_usage_24h: { requests: 3, @@ -950,7 +954,7 @@ describe('AccountUsageCell', () => { await flushPromises() - expect(wrapper.text()).toContain('24h|50') + expect(wrapper.text()).toContain('24h|100') }) it('Grok paid manual probes keep the weekly/local summary when 24h usage is returned', async () => { @@ -1165,6 +1169,7 @@ describe('AccountUsageCell', () => { it('Grok Free manual probes merge rolling 24h usage', async () => { getUsage.mockResolvedValue({ + grok_free_token_limit: 1_000_000, subscription_tier: 'FREE', grok_quota_snapshot_state: 'no_headers' }) @@ -1185,7 +1190,7 @@ describe('AccountUsageCell', () => { template: `