diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index e4ef733b32..ba5b78a72f 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -190,9 +190,10 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream) antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository) grokQuotaFetcher := service.NewGrokQuotaFetcher() + grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, usageLogRepository) openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory) usageCache := service.NewUsageCache() - accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService) + accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService) accountTestService := service.NewAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService) crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig) accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator) @@ -207,7 +208,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService, openAIQuotaService) geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService) antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService) - grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream) grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService) proxyHandler := admin.NewProxyHandler(adminService) adminRedeemHandler := admin.NewRedeemHandler(adminService, redeemService) diff --git a/backend/internal/handler/admin/grok_oauth_handler.go b/backend/internal/handler/admin/grok_oauth_handler.go index dafa3076b8..b914a16be7 100644 --- a/backend/internal/handler/admin/grok_oauth_handler.go +++ b/backend/internal/handler/admin/grok_oauth_handler.go @@ -215,7 +215,7 @@ func (h *GrokOAuthHandler) QueryQuota(c *gin.Context) { response.BadRequest(c, "grok quota service is not enabled") return } - result, err := h.quotaService.ProbeUsage(c.Request.Context(), accountID) + result, err := h.quotaService.QueryQuota(c.Request.Context(), accountID) if err != nil { response.ErrorFrom(c, err) return diff --git a/backend/internal/handler/admin/grok_oauth_handler_test.go b/backend/internal/handler/admin/grok_oauth_handler_test.go index 0b0f0d1aba..c1dd527d77 100644 --- a/backend/internal/handler/admin/grok_oauth_handler_test.go +++ b/backend/internal/handler/admin/grok_oauth_handler_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "strings" + "sync" "testing" "time" @@ -41,17 +42,35 @@ func (r *grokQuotaHandlerAccountRepo) UpdateExtra(_ context.Context, id int64, u } type grokQuotaHandlerUpstream struct { - resp *http.Response - lastReq *http.Request - lastBody []byte + mu sync.Mutex + requests []*http.Request + bodies [][]byte } func (u *grokQuotaHandlerUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { - u.lastReq = req + var body []byte if req.Body != nil { - u.lastBody, _ = io.ReadAll(req.Body) + body, _ = io.ReadAll(req.Body) } - return u.resp, nil + u.mu.Lock() + u.requests = append(u.requests, req) + u.bodies = append(u.bodies, body) + u.mu.Unlock() + if req.URL.Path == "/v1/responses" { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "X-Ratelimit-Limit-Requests": []string{"10"}, + "X-Ratelimit-Remaining-Requests": []string{"8"}, + }, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)), + }, nil + } + payload := `{"config":{"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}` + if req.URL.RawQuery == "format=credits" { + payload = `{"config":{"currentPeriod":{"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"}}}` + } + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(payload))}, nil } func (u *grokQuotaHandlerUpstream) DoWithTLS( @@ -77,14 +96,7 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) { "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), }, }} - upstream := &grokQuotaHandlerUpstream{resp: &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{ - "X-Ratelimit-Limit-Requests": []string{"10"}, - "X-Ratelimit-Remaining-Requests": []string{"8"}, - }, - Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)), - }} + upstream := &grokQuotaHandlerUpstream{} quotaService := service.NewGrokQuotaService(repo, nil, service.NewGrokTokenProvider(repo, nil), upstream) handler := NewGrokOAuthHandler(nil, nil, quotaService) @@ -95,12 +107,23 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) { router.ServeHTTP(rec, req) require.Equal(t, http.StatusOK, rec.Code) - require.Contains(t, rec.Body.String(), `"source":"active_probe"`) + require.Contains(t, rec.Body.String(), `"source":"hybrid_probe"`) + require.Contains(t, rec.Body.String(), `"billing":`) + require.Contains(t, rec.Body.String(), `"snapshot":`) require.Contains(t, rec.Body.String(), `"headers_observed":true`) require.NotContains(t, rec.Body.String(), "access-token") - require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) - require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) - require.Contains(t, string(upstream.lastBody), `"store":false`) + upstream.mu.Lock() + requests := append([]*http.Request(nil), upstream.requests...) + bodies := append([][]byte(nil), upstream.bodies...) + upstream.mu.Unlock() + require.Len(t, requests, 3) + for i, upstreamReq := range requests { + require.Equal(t, "Bearer access-token", upstreamReq.Header.Get("Authorization")) + if upstreamReq.URL.String() == xai.DefaultCLIBaseURL+"/responses" { + require.Contains(t, string(bodies[i]), `"model":"grok-4.5"`) + require.Contains(t, string(bodies[i]), `"store":false`) + } + } require.NotNil(t, repo.updates[42]) } diff --git a/backend/internal/pkg/xai/billing.go b/backend/internal/pkg/xai/billing.go new file mode 100644 index 0000000000..15b9c7e50e --- /dev/null +++ b/backend/internal/pkg/xai/billing.go @@ -0,0 +1,372 @@ +package xai + +import ( + "encoding/json" + "fmt" + "math" + "net/http" + "strconv" + "strings" + "time" +) + +const ( + // CLI client identity required by cli-chat-proxy billing endpoints. + CLITokenAuthHeader = "x-xai-token-auth" + CLITokenAuthValue = "xai-grok-cli" + CLIClientVersionHeader = "x-grok-client-version" + // Keep in sync with https://x.ai/cli/stable. + CLIClientVersion = "0.2.93" + CLIUserAgent = "grok-pager/" + CLIClientVersion + " grok-shell/" + CLIClientVersion + " (macos; aarch64)" + + BillingWeeklyPath = "/billing?format=credits" + BillingMonthlyPath = "/billing" + + SuperGrokLimitCents = 15_000 // $150.00 + SuperGrokHeavyLimitCents = 150_000 // $1,500.00 +) + +// BillingPeriod describes the current weekly/monthly window. +type BillingPeriod struct { + Type string `json:"type,omitempty"` + Start string `json:"start,omitempty"` + End string `json:"end,omitempty"` +} + +// BillingProductUsage is per-product usage inside the weekly credits window. +type BillingProductUsage struct { + Product string `json:"product,omitempty"` + UsagePercent *float64 `json:"usagePercent,omitempty"` +} + +// BillingConfig is the nested config object from /v1/billing responses. +type BillingConfig struct { + CurrentPeriod *BillingPeriod `json:"currentPeriod,omitempty"` + CreditUsagePercent *float64 `json:"creditUsagePercent,omitempty"` + ProductUsage []BillingProductUsage `json:"productUsage,omitempty"` + MonthlyLimit json.RawMessage `json:"monthlyLimit,omitempty"` + Used json.RawMessage `json:"used,omitempty"` + BillingPeriodStart string `json:"billingPeriodStart,omitempty"` + BillingPeriodEnd string `json:"billingPeriodEnd,omitempty"` +} + +// BillingPayload is the top-level body from /v1/billing. +type BillingPayload struct { + Config *BillingConfig `json:"config,omitempty"` +} + +// BillingProductSummary is a normalized product usage row for UI. +type BillingProductSummary struct { + Product string `json:"product"` + UsagePercent *float64 `json:"usage_percent,omitempty"` +} + +// BillingSummary is the merged weekly + monthly billing view. +type BillingSummary struct { + PeriodType string `json:"period_type,omitempty"` // weekly | monthly | unknown + UsagePercent *float64 `json:"usage_percent,omitempty"` + PeriodStart string `json:"period_start,omitempty"` + PeriodEnd string `json:"period_end,omitempty"` + ProductUsage []BillingProductSummary `json:"product_usage,omitempty"` + MonthlyLimitCents *float64 `json:"monthly_limit_cents,omitempty"` + UsedCents *float64 `json:"used_cents,omitempty"` + IncludedUsedCents *float64 `json:"included_used_cents,omitempty"` + BillingPeriodStart string `json:"billing_period_start,omitempty"` + BillingPeriodEnd string `json:"billing_period_end,omitempty"` + UsedPercent *float64 `json:"used_percent,omitempty"` + Plan string `json:"plan,omitempty"` // SuperGrok | SuperGrok Heavy | "" + StatusCode int `json:"status_code,omitempty"` + Source string `json:"source,omitempty"` + FetchedAt string `json:"fetched_at,omitempty"` + UpdatedAt string `json:"updated_at,omitempty"` + WeeklyUpdatedAt string `json:"weekly_updated_at,omitempty"` + MonthlyUpdatedAt string `json:"monthly_updated_at,omitempty"` + Partial bool `json:"partial,omitempty"` + FailedWindows []string `json:"failed_windows,omitempty"` +} + +// BuildBillingURL builds weekly or monthly billing URL against the CLI chat proxy. +func BuildBillingURL(formatCredits bool) string { + base := strings.TrimRight(DefaultCLIBaseURL, "/") + if formatCredits { + return base + BillingWeeklyPath + } + return base + BillingMonthlyPath +} + +// ApplyCLIBillingHeaders sets Authorization + CLI identity headers for billing GETs. +func ApplyCLIBillingHeaders(req *http.Request, accessToken string) { + if req == nil { + return + } + token := strings.TrimSpace(accessToken) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + req.Header.Set("Accept", "application/json") + req.Header.Set("Content-Type", "application/json") + req.Header.Set(CLITokenAuthHeader, CLITokenAuthValue) + req.Header.Set(CLIClientVersionHeader, CLIClientVersion) + req.Header.Set("User-Agent", CLIUserAgent) +} + +// ParseBillingPayload unmarshals a billing API response body. +func ParseBillingPayload(body []byte) (*BillingPayload, error) { + if len(body) == 0 { + return nil, fmt.Errorf("empty billing body") + } + var payload BillingPayload + if err := json.Unmarshal(body, &payload); err != nil { + return nil, err + } + return &payload, nil +} + +// BuildBillingSummary normalizes a billing config into a UI-friendly summary. +func BuildBillingSummary(config *BillingConfig) *BillingSummary { + if config == nil { + return nil + } + summary := &BillingSummary{} + period := config.CurrentPeriod + periodType := resolvePeriodType(period) + creditUsage := cloneFloat(config.CreditUsagePercent) + + periodStart := "" + periodEnd := "" + if period != nil { + periodStart = strings.TrimSpace(period.Start) + periodEnd = strings.TrimSpace(period.End) + } + if periodStart == "" { + periodStart = strings.TrimSpace(config.BillingPeriodStart) + } + if periodEnd == "" { + periodEnd = strings.TrimSpace(config.BillingPeriodEnd) + } + + products := make([]BillingProductSummary, 0, len(config.ProductUsage)) + for _, item := range config.ProductUsage { + product := strings.TrimSpace(item.Product) + if product == "" { + continue + } + products = append(products, BillingProductSummary{ + Product: product, + UsagePercent: cloneFloat(item.UsagePercent), + }) + } + + monthlyLimit := parseCentValue(config.MonthlyLimit) + used := parseCentValue(config.Used) + billingStart := strings.TrimSpace(config.BillingPeriodStart) + billingEnd := strings.TrimSpace(config.BillingPeriodEnd) + + var includedUsed *float64 + if used != nil { + if monthlyLimit != nil && *monthlyLimit > 0 { + v := math.Min(*used, *monthlyLimit) + includedUsed = &v + } else { + includedUsed = cloneFloat(used) + } + } + + var usedPercent *float64 + if monthlyLimit != nil && *monthlyLimit > 0 && includedUsed != nil { + v := (*includedUsed / *monthlyLimit) * 100 + usedPercent = &v + } + + hasWeekly := creditUsage != nil || periodType == "weekly" || len(products) > 0 + hasMonthly := monthlyLimit != nil || used != nil || (!hasWeekly && billingEnd != "") + if !hasWeekly && !hasMonthly { + return nil + } + + if hasWeekly { + if periodType == "unknown" { + periodType = "weekly" + } + summary.PeriodType = periodType + summary.UsagePercent = creditUsage + summary.PeriodStart = periodStart + summary.PeriodEnd = periodEnd + } else { + // Monthly-only: do not put monthly % into UsagePercent (weekly bar field). + // Frontend weekly bar only renders when PeriodType == weekly. + summary.PeriodType = "monthly" + summary.PeriodStart = billingStart + summary.PeriodEnd = billingEnd + } + summary.ProductUsage = products + summary.MonthlyLimitCents = monthlyLimit + summary.UsedCents = used + summary.IncludedUsedCents = includedUsed + if hasMonthly { + summary.BillingPeriodStart = billingStart + summary.BillingPeriodEnd = billingEnd + } + summary.UsedPercent = usedPercent + summary.Plan = resolvePlan(monthlyLimit) + return summary +} + +// MergeBillingProbeResult updates successful billing domains while retaining +// the previous value for any domain that could not be refreshed. +func MergeBillingProbeResult(previous, weekly, monthly *BillingSummary, weeklyOK, monthlyOK bool) *BillingSummary { + var out BillingSummary + if previous != nil { + out = *previous + previousUpdatedAt := previous.UpdatedAt + if previousUpdatedAt == "" { + previousUpdatedAt = previous.FetchedAt + } + if out.WeeklyUpdatedAt == "" && (out.UsagePercent != nil || len(out.ProductUsage) > 0) { + out.WeeklyUpdatedAt = previousUpdatedAt + } + if out.MonthlyUpdatedAt == "" && (out.MonthlyLimitCents != nil || out.UsedPercent != nil) { + out.MonthlyUpdatedAt = previousUpdatedAt + } + } + now := time.Now().UTC().Format(time.RFC3339) + + if weeklyOK && weekly != nil { + out.PeriodType = weekly.PeriodType + out.UsagePercent = weekly.UsagePercent + out.PeriodStart = weekly.PeriodStart + out.PeriodEnd = weekly.PeriodEnd + out.ProductUsage = weekly.ProductUsage + out.WeeklyUpdatedAt = now + } + if monthlyOK && monthly != nil { + if out.PeriodType == "" { + out.PeriodType = "monthly" + } + out.MonthlyLimitCents = monthly.MonthlyLimitCents + out.UsedCents = monthly.UsedCents + out.IncludedUsedCents = monthly.IncludedUsedCents + out.BillingPeriodStart = monthly.BillingPeriodStart + out.BillingPeriodEnd = monthly.BillingPeriodEnd + out.UsedPercent = monthly.UsedPercent + out.Plan = monthly.Plan + out.MonthlyUpdatedAt = now + } + + out.Partial = !weeklyOK || !monthlyOK + out.FailedWindows = nil + if !weeklyOK { + out.FailedWindows = append(out.FailedWindows, "weekly") + } + if !monthlyOK { + out.FailedWindows = append(out.FailedWindows, "monthly") + } + if !weeklyOK && !monthlyOK && previous == nil { + return nil + } + return &out +} + +// StampBillingSummary sets fetch metadata. +func StampBillingSummary(summary *BillingSummary, statusCode int, source string) *BillingSummary { + if summary == nil { + return nil + } + now := time.Now().UTC().Format(time.RFC3339) + summary.StatusCode = statusCode + summary.Source = source + summary.FetchedAt = now + summary.UpdatedAt = now + return summary +} + +func resolvePeriodType(period *BillingPeriod) string { + if period == nil { + return "unknown" + } + raw := strings.ToLower(strings.TrimSpace(period.Type)) + if strings.Contains(raw, "weekly") { + return "weekly" + } + if strings.Contains(raw, "monthly") { + return "monthly" + } + return "unknown" +} + +func resolvePlan(monthlyLimitCents *float64) string { + if monthlyLimitCents == nil { + return "" + } + // Allow small float noise. + limit := math.Round(*monthlyLimitCents) + switch limit { + case SuperGrokLimitCents: + return "SuperGrok" + case SuperGrokHeavyLimitCents: + return "SuperGrok Heavy" + default: + return "" + } +} + +func parseCentValue(raw json.RawMessage) *float64 { + if len(raw) == 0 || string(raw) == "null" { + return nil + } + // Object form: {"val": 123} + var obj struct { + Val any `json:"val"` + } + if err := json.Unmarshal(raw, &obj); err == nil && obj.Val != nil { + return anyToFloat(obj.Val) + } + // Bare number / string + var n any + if err := json.Unmarshal(raw, &n); err != nil { + return nil + } + return anyToFloat(n) +} + +func anyToFloat(v any) *float64 { + switch n := v.(type) { + case float64: + return &n + case float32: + f := float64(n) + return &f + case int: + f := float64(n) + return &f + case int64: + f := float64(n) + return &f + case json.Number: + f, err := n.Float64() + if err != nil { + return nil + } + return &f + case string: + s := strings.TrimSpace(n) + if s == "" { + return nil + } + f, err := strconv.ParseFloat(s, 64) + if err != nil { + return nil + } + return &f + default: + return nil + } +} + +func cloneFloat(v *float64) *float64 { + if v == nil { + return nil + } + f := *v + return &f +} diff --git a/backend/internal/pkg/xai/billing_test.go b/backend/internal/pkg/xai/billing_test.go new file mode 100644 index 0000000000..1d863f6a39 --- /dev/null +++ b/backend/internal/pkg/xai/billing_test.go @@ -0,0 +1,127 @@ +package xai + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/stretchr/testify/require" +) + +func TestBuildBillingURL(t *testing.T) { + t.Parallel() + require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing?format=credits", BuildBillingURL(true)) + require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing", BuildBillingURL(false)) +} + +func TestApplyCLIBillingHeaders(t *testing.T) { + t.Parallel() + req, err := http.NewRequest(http.MethodGet, BuildBillingURL(true), nil) + require.NoError(t, err) + + ApplyCLIBillingHeaders(req, " token ") + + require.Equal(t, "Bearer token", req.Header.Get("Authorization")) + require.Equal(t, CLITokenAuthValue, req.Header.Get(CLITokenAuthHeader)) + require.Equal(t, CLIClientVersion, req.Header.Get(CLIClientVersionHeader)) + require.Equal(t, "grok-pager/"+CLIClientVersion+" grok-shell/"+CLIClientVersion+" (macos; aarch64)", req.UserAgent()) +} + +func TestBuildBillingSummaryWeeklyAndMonthly(t *testing.T) { + t.Parallel() + + weeklyBody := []byte(`{ + "config": { + "currentPeriod": {"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"}, + "creditUsagePercent": 2.0, + "productUsage": [{"product":"Api","usagePercent":2.0}] + } + }`) + monthlyBody := []byte(`{ + "config": { + "monthlyLimit": {"val": 15000}, + "used": {"val": 78}, + "billingPeriodStart": "2026-07-01T00:00:00Z", + "billingPeriodEnd": "2026-08-01T00:00:00Z" + } + }`) + + weeklyPayload, err := ParseBillingPayload(weeklyBody) + require.NoError(t, err) + monthlyPayload, err := ParseBillingPayload(monthlyBody) + require.NoError(t, err) + + weekly := BuildBillingSummary(weeklyPayload.Config) + monthly := BuildBillingSummary(monthlyPayload.Config) + require.NotNil(t, weekly) + require.NotNil(t, monthly) + require.Equal(t, "weekly", weekly.PeriodType) + require.InDelta(t, 2.0, *weekly.UsagePercent, 1e-9) + require.Equal(t, "Api", weekly.ProductUsage[0].Product) + require.Equal(t, "SuperGrok", monthly.Plan) + require.InDelta(t, 15000, *monthly.MonthlyLimitCents, 1e-9) + require.InDelta(t, 78, *monthly.UsedCents, 1e-9) + require.InDelta(t, 0.52, *monthly.UsedPercent, 1e-2) + + merged := MergeBillingProbeResult(nil, weekly, monthly, true, true) + require.Equal(t, "weekly", merged.PeriodType) + require.InDelta(t, 2.0, *merged.UsagePercent, 1e-9) + require.Equal(t, "SuperGrok", merged.Plan) + require.InDelta(t, 15000, *merged.MonthlyLimitCents, 1e-9) + require.Equal(t, "2026-08-01T00:00:00Z", merged.BillingPeriodEnd) +} + +func TestParseCentValueBareNumber(t *testing.T) { + t.Parallel() + raw, _ := json.Marshal(15000) + v := parseCentValue(raw) + require.NotNil(t, v) + require.InDelta(t, 15000, *v, 1e-9) +} + +func TestBuildBillingSummaryMonthlyOnlyKeepsWeeklyUsageEmpty(t *testing.T) { + t.Parallel() + payload, err := ParseBillingPayload([]byte(`{"config":{"monthlyLimit":{"val":15000},"used":{"val":7500},"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}`)) + require.NoError(t, err) + + summary := BuildBillingSummary(payload.Config) + require.NotNil(t, summary) + require.Equal(t, "monthly", summary.PeriodType) + require.Nil(t, summary.UsagePercent) + require.InDelta(t, 50, *summary.UsedPercent, 1e-9) +} + +func TestMergeBillingProbeResultRetainsFailedWindow(t *testing.T) { + t.Parallel() + previous := &BillingSummary{ + PeriodType: "weekly", + UsagePercent: floatPointer(100), + PeriodEnd: "2026-07-16T00:00:00Z", + MonthlyLimitCents: floatPointer(15000), + UsedPercent: floatPointer(20), + BillingPeriodEnd: "2026-08-01T00:00:00Z", + WeeklyUpdatedAt: "2026-07-10T00:00:00Z", + MonthlyUpdatedAt: "2026-07-10T00:00:00Z", + FailedWindows: []string{"monthly"}, + } + monthly := &BillingSummary{ + PeriodType: "monthly", + MonthlyLimitCents: floatPointer(15000), + UsedPercent: floatPointer(30), + BillingPeriodEnd: "2026-08-01T00:00:00Z", + } + + merged := MergeBillingProbeResult(previous, nil, monthly, false, true) + require.Equal(t, "weekly", merged.PeriodType) + require.InDelta(t, 100, *merged.UsagePercent, 1e-9) + require.Equal(t, previous.WeeklyUpdatedAt, merged.WeeklyUpdatedAt) + require.InDelta(t, 30, *merged.UsedPercent, 1e-9) + require.NotEqual(t, previous.MonthlyUpdatedAt, merged.MonthlyUpdatedAt) + require.True(t, merged.Partial) + require.Equal(t, []string{"weekly"}, merged.FailedWindows) + require.Equal(t, []string{"monthly"}, previous.FailedWindows) +} + +func floatPointer(value float64) *float64 { + return &value +} diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 261be7c14b..6d26d474c6 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -61,6 +61,7 @@ var schedulerNeutralExtraKeyPrefixes = []string{ var schedulerNeutralExtraKeys = map[string]struct{}{ "codex_usage_updated_at": {}, + "grok_billing_snapshot": {}, "session_window_utilization": {}, } diff --git a/backend/internal/repository/account_repo_grok_billing_test.go b/backend/internal/repository/account_repo_grok_billing_test.go new file mode 100644 index 0000000000..fb41ae5ffa --- /dev/null +++ b/backend/internal/repository/account_repo_grok_billing_test.go @@ -0,0 +1,16 @@ +package repository + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestGrokBillingSnapshotIsSchedulerNeutral(t *testing.T) { + t.Parallel() + + require.True(t, isSchedulerNeutralExtraKey("grok_billing_snapshot")) + require.False(t, shouldEnqueueSchedulerOutboxForExtraUpdates(map[string]any{ + "grok_billing_snapshot": map[string]any{"usage_percent": 50}, + })) +} diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 31549a6b17..222b3f8a4d 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -669,7 +669,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * testModelID := strings.TrimSpace(modelID) if testModelID == "" { - testModelID = "grok-4.3" + testModelID = grokDefaultResponsesModel } if mapped := strings.TrimSpace(account.GetMappedModel(testModelID)); mapped != "" { testModelID = mapped diff --git a/backend/internal/service/account_test_service_grok_test.go b/backend/internal/service/account_test_service_grok_test.go index 497224b713..4b0890ff44 100644 --- a/backend/internal/service/account_test_service_grok_test.go +++ b/backend/internal/service/account_test_service_grok_test.go @@ -80,6 +80,47 @@ func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testin require.Contains(t, rec.Body.String(), `"type":"test_complete"`) } +func TestAccountTestService_TestAccountConnection_GrokDefaultsEmptyModelTo45(t *testing.T) { + gin.SetMode(gin.TestMode) + + account := &Account{ + ID: 16, + Name: "grok-oauth-default-model", + 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), + }, + } + repo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" + + "data: {\"type\":\"response.completed\"}\n\n", + )), + }} + 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/16/test", nil) + + err := svc.TestAccountConnection(c, account.ID, "", "", AccountTestModeDefault) + + require.NoError(t, err) + require.Equal(t, grokDefaultResponsesModel, gjson.GetBytes(upstream.lastBody, "model").String()) + require.Contains(t, recorder.Body.String(), `"model":"grok-4.5"`) +} + func TestAccountTestService_Grok429PersistsRateLimitReset(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index 281122d4f0..9966c66515 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -111,6 +111,7 @@ const ( apiQueryMaxJitter = 800 * time.Millisecond // 用量查询最大随机延迟 windowStatsCacheTTL = 1 * time.Minute openAIProbeCacheTTL = 10 * time.Minute + grokProbeRetryTTL = 1 * time.Minute openAICodexProbeVersion = "0.144.1" ) @@ -122,6 +123,7 @@ type UsageCache struct { apiFlight singleflight.Group // 防止同一账号的并发请求击穿缓存（Anthropic） antigravityFlight singleflight.Group // 防止同一 Antigravity 账号的并发请求击穿缓存 openAIProbeCache sync.Map // accountID -> time.Time + grokProbeCache sync.Map // accountID -> last billing probe attempt } // NewUsageCache 创建 UsageCache 实例 @@ -196,15 +198,18 @@ type UsageInfo struct { AntigravityQuota map[string]*AntigravityModelQuota `json:"antigravity_quota,omitempty"` // Grok / xAI 被动额度快照 - GrokRequestQuota *xai.QuotaWindow `json:"grok_request_quota,omitempty"` - GrokTokenQuota *xai.QuotaWindow `json:"grok_token_quota,omitempty"` - GrokRetryAfterSeconds *int `json:"grok_retry_after_seconds,omitempty"` - GrokEntitlementStatus string `json:"grok_entitlement_status,omitempty"` - GrokQuotaSnapshotState string `json:"grok_quota_snapshot_state,omitempty"` - 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"` - GrokLocalUsage *WindowStats `json:"grok_local_usage,omitempty"` + GrokRequestQuota *xai.QuotaWindow `json:"grok_request_quota,omitempty"` + GrokTokenQuota *xai.QuotaWindow `json:"grok_token_quota,omitempty"` + GrokRetryAfterSeconds *int `json:"grok_retry_after_seconds,omitempty"` + GrokEntitlementStatus string `json:"grok_entitlement_status,omitempty"` + GrokQuotaSnapshotState string `json:"grok_quota_snapshot_state,omitempty"` + 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"` + GrokLocalUsage *WindowStats `json:"grok_local_usage,omitempty"` + GrokLocalUsage7d *WindowStats `json:"grok_local_usage_7d,omitempty"` + GrokLocalUsageMonthly *WindowStats `json:"grok_local_usage_monthly,omitempty"` + GrokBilling *xai.BillingSummary `json:"grok_billing,omitempty"` // Antigravity 账号级信息 SubscriptionTier string `json:"subscription_tier,omitempty"` // 归一化订阅等级: FREE/PRO/ULTRA/UNKNOWN @@ -287,6 +292,7 @@ type AccountUsageService struct { geminiQuotaService *GeminiQuotaService antigravityQuotaFetcher *AntigravityQuotaFetcher grokQuotaFetcher *GrokQuotaFetcher + grokQuotaService *GrokQuotaService openAIQuotaService *OpenAIQuotaService cache *UsageCache identityCache IdentityCache @@ -301,6 +307,7 @@ func NewAccountUsageService( geminiQuotaService *GeminiQuotaService, antigravityQuotaFetcher *AntigravityQuotaFetcher, grokQuotaFetcher *GrokQuotaFetcher, + grokQuotaService *GrokQuotaService, openAIQuotaService *OpenAIQuotaService, cache *UsageCache, identityCache IdentityCache, @@ -313,6 +320,7 @@ func NewAccountUsageService( geminiQuotaService: geminiQuotaService, antigravityQuotaFetcher: antigravityQuotaFetcher, grokQuotaFetcher: grokQuotaFetcher, + grokQuotaService: grokQuotaService, openAIQuotaService: openAIQuotaService, cache: cache, identityCache: identityCache, @@ -358,8 +366,8 @@ func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64, for } if account.Platform == PlatformGrok { - usage, err := s.getGrokUsage(ctx, account) - if err == nil { + usage, err := s.getGrokUsage(ctx, account, forceProbe) + if err == nil && usage != nil && usage.Error == "" { s.tryClearRecoverableAccountError(ctx, account) } return usage, err @@ -930,11 +938,19 @@ func (s *AccountUsageService) getAntigravityUsage(ctx context.Context, account * return usage, nil } -func (s *AccountUsageService) getGrokUsage(ctx context.Context, account *Account) (*UsageInfo, error) { +func (s *AccountUsageService) getGrokUsage(ctx context.Context, account *Account, force bool) (*UsageInfo, error) { if s.grokQuotaFetcher == nil { now := time.Now() return &UsageInfo{UpdatedAt: &now}, nil } + if account != nil && account.IsGrokOAuth() && s.grokQuotaService != nil && (force || grokBillingSnapshotNeedsRefresh(account, time.Now())) && s.shouldProbeGrokBilling(account.ID, time.Now(), force) { + result, err := s.grokQuotaService.ProbeBilling(ctx, account.ID) + if err == nil && result != nil && result.Billing != nil { + mergeAccountExtra(account, map[string]any{grokBillingExtraKey: result.Billing}) + } else if err != nil && force { + return nil, err + } + } usage := s.grokQuotaFetcher.BuildUsageInfo(account) if usage.GrokQuotaSnapshotState == "" { if usage.ErrorCode == "quota_unknown" { @@ -948,12 +964,90 @@ func (s *AccountUsageService) getGrokUsage(ctx context.Context, account *Account if stats, err := s.usageLogRepo.GetAccountTodayStats(ctx, account.ID); err == nil && stats != nil { usage.GrokLocalUsage = windowStatsFromAccountStats(stats) } + usage.GrokLocalUsage7d, usage.GrokLocalUsageMonthly = grokLocalUsageForBilling(ctx, s.usageLogRepo, account.ID, usage.GrokBilling, time.Now().UTC()) } enrichUsageWithAccountError(usage, account) return usage, nil } +func grokLocalUsageForBilling( + ctx context.Context, + repo UsageLogRepository, + accountID int64, + billing *xai.BillingSummary, + now time.Time, +) (*WindowStats, *WindowStats) { + var weekly *WindowStats + var monthly *WindowStats + if repo == nil || accountID <= 0 { + return weekly, monthly + } + if start, ok := currentGrokBillingWindow(billing, true, now); ok { + if stats, err := repo.GetAccountWindowStats(ctx, accountID, start); err == nil { + weekly = windowStatsFromAccountStats(stats) + } else { + slog.Warn("grok_window_usage_query_failed", "account_id", accountID, "window_start", start, "error", err) + } + } + if start, ok := currentGrokBillingWindow(billing, false, now); ok { + if stats, err := repo.GetAccountWindowStats(ctx, accountID, start); err == nil { + monthly = windowStatsFromAccountStats(stats) + } else { + slog.Warn("grok_monthly_usage_query_failed", "account_id", accountID, "window_start", start, "error", err) + } + } + return weekly, monthly +} + +func currentGrokBillingWindow(billing *xai.BillingSummary, weekly bool, now time.Time) (time.Time, bool) { + if billing == nil { + return time.Time{}, false + } + startRaw, endRaw := billing.BillingPeriodStart, billing.BillingPeriodEnd + if weekly { + if billing.PeriodType != "weekly" { + return time.Time{}, false + } + startRaw, endRaw = billing.PeriodStart, billing.PeriodEnd + } + start, startErr := parseTime(strings.TrimSpace(startRaw)) + end, endErr := parseTime(strings.TrimSpace(endRaw)) + if startErr != nil || endErr != nil || now.Before(start) || !now.Before(end) { + return time.Time{}, false + } + return start, true +} + +func grokBillingSnapshotNeedsRefresh(account *Account, now time.Time) bool { + if account == nil { + return false + } + billing, err := grokBillingSnapshotFromExtra(account.Extra) + if err != nil || billing == nil || billing.Partial || len(billing.FailedWindows) > 0 { + return true + } + stamp := strings.TrimSpace(billing.UpdatedAt) + if stamp == "" { + stamp = strings.TrimSpace(billing.FetchedAt) + } + updatedAt, err := parseTime(stamp) + return err != nil || now.Sub(updatedAt) >= openAIProbeCacheTTL +} + +func (s *AccountUsageService) shouldProbeGrokBilling(accountID int64, now time.Time, force bool) bool { + if force || s == nil || s.cache == nil || accountID <= 0 { + return true + } + if cached, ok := s.cache.grokProbeCache.Load(accountID); ok { + if ts, ok := cached.(time.Time); ok && now.Sub(ts) < grokProbeRetryTTL { + return false + } + } + s.cache.grokProbeCache.Store(accountID, now) + return true +} + // recalcAntigravityRemainingSeconds 重新计算 Antigravity UsageInfo 中各窗口的 RemainingSeconds // 用于从缓存取出时更新倒计时，避免返回过时的剩余秒数 func recalcAntigravityRemainingSeconds(info *UsageInfo) { diff --git a/backend/internal/service/grok_quota_fetcher.go b/backend/internal/service/grok_quota_fetcher.go index 0939b78e20..f220fe33b9 100644 --- a/backend/internal/service/grok_quota_fetcher.go +++ b/backend/internal/service/grok_quota_fetcher.go @@ -3,6 +3,8 @@ package service import ( "encoding/json" "fmt" + "net/http" + "strings" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" @@ -24,54 +26,150 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { } if account == nil { usage.ErrorCode = "quota_unknown" - usage.Error = "Grok quota is unknown until the first upstream response includes xAI rate-limit headers" + usage.Error = "Grok quota is unknown until billing is probed or an upstream response includes xAI rate-limit headers" return usage } + billing, _ := grokBillingSnapshotFromExtra(account.Extra) snapshot, err := grokQuotaSnapshotFromExtra(account.Extra) + if billing != nil { + usage.GrokBilling = billing + if billing.Plan != "" { + usage.SubscriptionTier = billing.Plan + usage.SubscriptionTierRaw = billing.Plan + } + if parsedAt, parseErr := time.Parse(time.RFC3339, billing.UpdatedAt); parseErr == nil { + usage.UpdatedAt = &parsedAt + } + if billing.FetchedAt != "" { + usage.GrokLastQuotaProbeAt = billing.FetchedAt + } + usage.GrokQuotaSnapshotState = "billing_observed" + usage.GrokLastStatusCode = billing.StatusCode + switch billing.StatusCode { + case 401: + usage.NeedsReauth = true + usage.ErrorCode = "unauthenticated" + case 403: + usage.IsForbidden = true + usage.ForbiddenType = "forbidden" + usage.ErrorCode = "forbidden" + case 429: + usage.ErrorCode = "rate_limited" + } + } + if err != nil || snapshot == nil { - usage.ErrorCode = "quota_unknown" - usage.Error = "Grok quota is unknown until the first upstream response includes xAI rate-limit headers" + applyGrokCredentialUsageFallback(usage, account) + if billing == nil { + usage.ErrorCode = "quota_unknown" + usage.Error = "Grok quota is unknown until billing is probed or an upstream response includes xAI rate-limit headers" + } return usage } - if parsedAt, err := time.Parse(time.RFC3339, snapshot.UpdatedAt); err == nil { - usage.UpdatedAt = &parsedAt + if parsedAt, parseErr := time.Parse(time.RFC3339, snapshot.UpdatedAt); parseErr == nil { + if billing == nil || usage.UpdatedAt == nil || parsedAt.After(*usage.UpdatedAt) { + usage.UpdatedAt = &parsedAt + } } usage.GrokRequestQuota = snapshot.Requests usage.GrokTokenQuota = snapshot.Tokens usage.GrokRetryAfterSeconds = snapshot.RetryAfterSeconds - usage.SubscriptionTier = snapshot.SubscriptionTier - usage.SubscriptionTierRaw = snapshot.SubscriptionTier - usage.GrokEntitlementStatus = snapshot.EntitlementStatus - usage.GrokLastQuotaProbeAt = snapshot.LastProbeAt + if usage.SubscriptionTier == "" { + usage.SubscriptionTier = snapshot.SubscriptionTier + usage.SubscriptionTierRaw = snapshot.SubscriptionTier + } + if usage.GrokEntitlementStatus == "" { + usage.GrokEntitlementStatus = snapshot.EntitlementStatus + } + if usage.GrokLastQuotaProbeAt == "" { + usage.GrokLastQuotaProbeAt = snapshot.LastProbeAt + } usage.GrokLastHeadersSeenAt = snapshot.LastHeadersSeenAt - usage.GrokLastStatusCode = snapshot.StatusCode + if snapshot.StatusCode >= http.StatusBadRequest || usage.GrokLastStatusCode == 0 { + usage.GrokLastStatusCode = snapshot.StatusCode + } if snapshot.HasObservedHeaders() { - usage.GrokQuotaSnapshotState = "observed" - } else { + if usage.GrokQuotaSnapshotState == "" { + usage.GrokQuotaSnapshotState = "observed" + } + } else if billing == nil { usage.GrokQuotaSnapshotState = "no_headers" usage.ErrorCode = "quota_unknown" usage.Error = "No xAI quota headers observed on the latest Grok probe" } - switch snapshot.StatusCode { - case 401: - usage.NeedsReauth = true - usage.ErrorCode = "unauthenticated" - case 403: - usage.IsForbidden = true - usage.ForbiddenType = "forbidden" - usage.ErrorCode = "forbidden" - if usage.GrokEntitlementStatus == "" { - usage.GrokEntitlementStatus = "forbidden" + if usage.ErrorCode == "" { + switch snapshot.StatusCode { + case 401: + usage.NeedsReauth = true + usage.ErrorCode = "unauthenticated" + case 403: + usage.IsForbidden = true + usage.ForbiddenType = "forbidden" + usage.ErrorCode = "forbidden" + if usage.GrokEntitlementStatus == "" { + usage.GrokEntitlementStatus = "forbidden" + } + case 429: + usage.ErrorCode = "rate_limited" } - case 429: - usage.ErrorCode = "rate_limited" } + applyGrokCredentialUsageFallback(usage, account) return usage } +func applyGrokCredentialUsageFallback(usage *UsageInfo, account *Account) { + if usage == nil || account == nil { + return + } + if usage.SubscriptionTier == "" { + tier := strings.TrimSpace(account.GetCredential("subscription_tier")) + usage.SubscriptionTier = tier + usage.SubscriptionTierRaw = tier + } + if usage.GrokEntitlementStatus == "" { + usage.GrokEntitlementStatus = strings.TrimSpace(account.GetCredential("entitlement_status")) + } +} + +func grokBillingSnapshotFromExtra(extra map[string]any) (*xai.BillingSummary, error) { + if extra == nil { + return nil, nil + } + raw, ok := extra[grokBillingExtraKey] + if !ok || raw == nil { + return nil, nil + } + switch snapshot := raw.(type) { + case *xai.BillingSummary: + return snapshot, nil + case xai.BillingSummary: + return &snapshot, nil + case map[string]any: + data, err := json.Marshal(snapshot) + if err != nil { + return nil, err + } + var out xai.BillingSummary + if err := json.Unmarshal(data, &out); err != nil { + return nil, err + } + return &out, nil + default: + data, err := json.Marshal(raw) + if err != nil { + return nil, fmt.Errorf("marshal grok billing snapshot: %w", err) + } + var out xai.BillingSummary + if err := json.Unmarshal(data, &out); err != nil { + return nil, err + } + return &out, nil + } +} + func grokQuotaSnapshotFromExtra(extra map[string]any) (*xai.QuotaSnapshot, error) { if extra == nil { return nil, nil diff --git a/backend/internal/service/grok_quota_fetcher_test.go b/backend/internal/service/grok_quota_fetcher_test.go index d2d9c14993..1de9b51c9e 100644 --- a/backend/internal/service/grok_quota_fetcher_test.go +++ b/backend/internal/service/grok_quota_fetcher_test.go @@ -20,7 +20,34 @@ func TestGrokQuotaFetcherBuildUsageInfoUnknownUntilFirstSnapshot(t *testing.T) { usage := NewGrokQuotaFetcher().BuildUsageInfo(&Account{Platform: PlatformGrok, Type: AccountTypeOAuth}) require.Equal(t, "passive", usage.Source) require.Equal(t, "quota_unknown", usage.ErrorCode) - require.Contains(t, usage.Error, "unknown until the first upstream response") + require.Contains(t, usage.Error, "unknown until billing is probed") +} + +func TestGrokQuotaFetcherUsesCredentialTierWhenBillingHasNoPlan(t *testing.T) { + t.Parallel() + + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "subscription_tier": " FREE ", + "entitlement_status": " active ", + }, + Extra: map[string]any{ + grokBillingExtraKey: &xai.BillingSummary{ + PeriodType: "weekly", + StatusCode: http.StatusOK, + UpdatedAt: "2030-01-01T00:00:00Z", + }, + }, + } + + usage := NewGrokQuotaFetcher().BuildUsageInfo(account) + + require.NotNil(t, usage.GrokBilling) + require.Equal(t, "FREE", usage.SubscriptionTier) + require.Equal(t, "FREE", usage.SubscriptionTierRaw) + require.Equal(t, "active", usage.GrokEntitlementStatus) } func TestGrokQuotaFetcherBuildUsageInfoFromSnapshot(t *testing.T) { @@ -68,6 +95,32 @@ func TestGrokQuotaFetcherBuildUsageInfoFromSnapshot(t *testing.T) { require.True(t, usage.UpdatedAt.Equal(time.Date(2030, 1, 1, 0, 0, 0, 0, time.UTC))) } +func TestGrokQuotaFetcherSnapshotErrorOverridesSuccessfulBillingStatus(t *testing.T) { + t.Parallel() + + updatedAt := "2030-01-01T00:00:00Z" + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Extra: map[string]any{ + grokBillingExtraKey: &xai.BillingSummary{ + PeriodType: "weekly", + StatusCode: http.StatusOK, + UpdatedAt: updatedAt, + }, + grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + UpdatedAt: updatedAt, + }, + }, + } + + usage := NewGrokQuotaFetcher().BuildUsageInfo(account) + + require.Equal(t, "rate_limited", usage.ErrorCode) + require.Equal(t, http.StatusTooManyRequests, usage.GrokLastStatusCode) +} + func TestGrokQuotaFetcherBuildUsageInfoFromNoHeadersProbe(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 17e91dac6f..aa838cb491 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -7,27 +7,36 @@ import ( "io" "log/slog" "net/http" + "strconv" "strings" + "sync" "time" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "golang.org/x/sync/singleflight" ) const ( grokQuotaUpstreamTimeout = 20 * time.Second grokQuotaProbeInput = "." - grokQuotaDefaultModel = "grok-4.3" + grokQuotaDefaultModel = grokDefaultResponsesModel + grokBillingExtraKey = "grok_billing_snapshot" ) type GrokQuotaProbeResult struct { - Source string `json:"source"` - Model string `json:"model"` - Snapshot *xai.QuotaSnapshot `json:"snapshot,omitempty"` - StatusCode int `json:"status_code,omitempty"` - HeadersObserved bool `json:"headers_observed"` - ResetSupported bool `json:"reset_supported"` - FetchedAt int64 `json:"fetched_at"` + Source string `json:"source"` + Model string `json:"model,omitempty"` + Billing *xai.BillingSummary `json:"billing,omitempty"` + Snapshot *xai.QuotaSnapshot `json:"snapshot,omitempty"` + LocalUsage7d *WindowStats `json:"local_usage_7d,omitempty"` + LocalUsageMonthly *WindowStats `json:"local_usage_monthly,omitempty"` + StatusCode int `json:"status_code,omitempty"` + HeadersObserved bool `json:"headers_observed"` + ResetSupported bool `json:"reset_supported"` + FetchedAt int64 `json:"fetched_at"` + Persisted bool `json:"persisted"` + ProbeError string `json:"probe_error,omitempty"` } type GrokQuotaResetResult struct { @@ -41,6 +50,8 @@ type GrokQuotaService struct { proxyRepo ProxyRepository tokenProvider *GrokTokenProvider httpUpstream HTTPUpstream + usageLogRepo UsageLogRepository + probeFlight singleflight.Group } func NewGrokQuotaService( @@ -48,16 +59,70 @@ func NewGrokQuotaService( proxyRepo ProxyRepository, tokenProvider *GrokTokenProvider, httpUpstream HTTPUpstream, + usageLogRepos ...UsageLogRepository, ) *GrokQuotaService { + var usageLogRepo UsageLogRepository + if len(usageLogRepos) > 0 { + usageLogRepo = usageLogRepos[0] + } return &GrokQuotaService{ accountRepo: accountRepo, proxyRepo: proxyRepo, tokenProvider: tokenProvider, httpUpstream: httpUpstream, + usageLogRepo: usageLogRepo, } } +// QueryQuota combines xAI billing data with an active quota-header probe for +// Free accounts, whose billing response does not include usage_percent. +func (s *GrokQuotaService) QueryQuota(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { + billingResult, billingErr := s.ProbeBilling(ctx, accountID) + if billingErr == nil && billingResult != nil && grokBillingHasAuthoritativeQuota(billingResult.Billing) { + return billingResult, nil + } + + probeResult, probeErr := s.ProbeUsage(ctx, accountID) + if probeErr != nil { + if billingResult != nil && billingResult.Billing != nil { + billingResult.ProbeError = probeErr.Error() + return billingResult, nil + } + return nil, probeErr + } + if probeResult == nil { + if billingErr != nil { + return nil, billingErr + } + return nil, infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_EMPTY", "Grok quota probe returned no result") + } + if billingResult != nil { + probeResult.Source = "hybrid_probe" + probeResult.Billing = billingResult.Billing + probeResult.LocalUsage7d = billingResult.LocalUsage7d + probeResult.LocalUsageMonthly = billingResult.LocalUsageMonthly + probeResult.Persisted = probeResult.Persisted || billingResult.Persisted + } + return probeResult, nil +} + +func grokBillingHasAuthoritativeQuota(billing *xai.BillingSummary) bool { + if billing == nil { + return false + } + return billing.UsagePercent != nil || + billing.UsedPercent != nil || + (billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0) || + strings.TrimSpace(billing.Plan) != "" +} + func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { + return s.runProbeFlight(ctx, "active:"+strconv.FormatInt(accountID, 10), func(sharedCtx context.Context) (*GrokQuotaProbeResult, error) { + return s.probeUsage(sharedCtx, accountID) + }) +} + +func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { account, token, proxyURL, err := s.prepareProbe(ctx, accountID) if err != nil { return nil, err @@ -95,7 +160,7 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr if limited { normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now()) } - _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ + persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ grokQuotaSnapshotExtraKey: snapshot, }) if limited { @@ -110,6 +175,7 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr HeadersObserved: snapshot.HeadersObserved, ResetSupported: false, FetchedAt: time.Now().Unix(), + Persisted: persistErr == nil, } if resp.StatusCode == http.StatusTooManyRequests { return result, nil @@ -123,6 +189,173 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr return result, nil } +// ProbeBilling only calls the xAI billing endpoints. Account usage refreshes +// use this method so opening the account list never consumes model quota. +func (s *GrokQuotaService) ProbeBilling(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { + return s.runProbeFlight(ctx, "billing:"+strconv.FormatInt(accountID, 10), func(sharedCtx context.Context) (*GrokQuotaProbeResult, error) { + return s.probeBilling(sharedCtx, accountID) + }) +} + +func (s *GrokQuotaService) probeBilling(ctx context.Context, accountID int64) (*GrokQuotaProbeResult, error) { + account, token, proxyURL, err := s.prepareProbe(ctx, accountID) + if err != nil { + return nil, err + } + + probeCtx, cancel := context.WithTimeout(ctx, grokQuotaUpstreamTimeout) + defer cancel() + type billingResult struct { + summary *xai.BillingSummary + status int + err error + } + var weekly, monthly billingResult + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + weekly.summary, weekly.status, weekly.err = s.fetchBilling(probeCtx, account, token, proxyURL, true) + }() + go func() { + defer wg.Done() + monthly.summary, monthly.status, monthly.err = s.fetchBilling(probeCtx, account, token, proxyURL, false) + }() + wg.Wait() + + weeklyOK := weekly.summary != nil + monthlyOK := monthly.summary != nil + if !weeklyOK && !monthlyOK { + return nil, mergeGrokBillingProbeErrors(weekly.status, monthly.status, weekly.err, monthly.err) + } + statusCode := preferSuccessfulBillingStatus(weekly.status, monthly.status, weeklyOK, monthlyOK) + previous, _ := grokBillingSnapshotFromExtra(account.Extra) + billing := xai.MergeBillingProbeResult(previous, weekly.summary, monthly.summary, weeklyOK, monthlyOK) + billing = xai.StampBillingSummary(billing, statusCode, "billing_probe") + persistErr := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ + grokBillingExtraKey: billing, + }) + if persistErr != nil { + slog.Warn("grok_billing_persist_failed", "account_id", account.ID, "error", persistErr) + } + localUsage7d, localUsageMonthly := grokLocalUsageForBilling(ctx, s.usageLogRepo, account.ID, billing, time.Now().UTC()) + return &GrokQuotaProbeResult{ + Source: "billing_probe", + Billing: billing, + LocalUsage7d: localUsage7d, + LocalUsageMonthly: localUsageMonthly, + StatusCode: statusCode, + FetchedAt: time.Now().Unix(), + Persisted: persistErr == nil, + }, nil +} + +func (s *GrokQuotaService) runProbeFlight( + ctx context.Context, + key string, + probe func(context.Context) (*GrokQuotaProbeResult, error), +) (*GrokQuotaProbeResult, error) { + if s == nil { + return nil, infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_NOT_CONFIGURED", "grok quota service is not configured") + } + resultCh := s.probeFlight.DoChan(key, func() (any, error) { + sharedCtx, cancel := context.WithTimeout(context.Background(), grokQuotaUpstreamTimeout+5*time.Second) + defer cancel() + return probe(sharedCtx) + }) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case flightResult := <-resultCh: + if flightResult.Err != nil { + return nil, flightResult.Err + } + result, ok := flightResult.Val.(*GrokQuotaProbeResult) + if !ok || result == nil { + return nil, infraerrors.New(http.StatusInternalServerError, "GROK_QUOTA_PROBE_RESULT_INVALID", "invalid Grok quota probe result") + } + cloned := *result + return &cloned, nil + } +} + +func (s *GrokQuotaService) fetchBilling( + ctx context.Context, + account *Account, + token string, + proxyURL string, + weekly bool, +) (*xai.BillingSummary, int, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, xai.BuildBillingURL(weekly), 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) + 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() }() + + bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if resp.StatusCode == http.StatusTooManyRequests { + return nil, resp.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) + } + 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 { + weeklyKey := grokBillingProbeErrorKey(weeklyStatus, weeklyErr) + monthlyKey := grokBillingProbeErrorKey(monthlyStatus, monthlyErr) + if weeklyKey == monthlyKey { + switch { + case weeklyErr != nil: + return weeklyErr + case monthlyErr != nil: + return monthlyErr + case weeklyStatus == http.StatusTooManyRequests: + return infraerrors.New(http.StatusTooManyRequests, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "billing rate limited") + case weeklyStatus != 0 && weeklyStatus != http.StatusOK: + return infraerrors.New(mapUpstreamStatus(weeklyStatus), "GROK_QUOTA_PROBE_UPSTREAM_ERROR", "xAI billing endpoints returned the same upstream error") + default: + return infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_BILLING_EMPTY", "xAI billing endpoints returned no quota data") + } + } + slog.Warn("grok_quota_probe_parts_failed", "weekly_status", weeklyStatus, "weekly_error", weeklyErr, "monthly_status", monthlyStatus, "monthly_error", monthlyErr) + return infraerrors.New(http.StatusBadGateway, "GROK_QUOTA_PROBE_PARTS_FAILED", "weekly and monthly billing probes failed differently").WithMetadata(map[string]string{ + "weekly_status": strconv.Itoa(weeklyStatus), "monthly_status": strconv.Itoa(monthlyStatus), + }) +} + +func grokBillingProbeErrorKey(status int, err error) string { + if err != nil { + return strconv.Itoa(status) + ":" + strconv.Itoa(infraerrors.Code(err)) + ":" + infraerrors.Reason(err) + } + return strconv.Itoa(status) + ":empty" +} + +func preferSuccessfulBillingStatus(weeklyStatus, monthlyStatus int, weeklyOK, monthlyOK bool) int { + if weeklyOK && weeklyStatus >= 200 && weeklyStatus < 300 { + return weeklyStatus + } + if monthlyOK && monthlyStatus >= 200 && monthlyStatus < 300 { + return monthlyStatus + } + if weeklyStatus != 0 { + return weeklyStatus + } + return monthlyStatus +} + func (s *GrokQuotaService) ResetQuota(ctx context.Context, accountID int64) (*GrokQuotaResetResult, error) { if _, err := s.loadGrokOAuthAccount(ctx, accountID); err != nil { return nil, err diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index 2248674899..d49497e69b 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -6,11 +6,14 @@ import ( "context" "io" "net/http" + "strconv" "strings" + "sync" "testing" "time" infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" @@ -63,6 +66,107 @@ type grokQuotaProxyRepo struct { calls int } +type grokQuotaUsageLogRepo struct { + UsageLogRepository + stats *usagestats.AccountStats + err error + calls int +} + +func (r *grokQuotaUsageLogRepo) GetAccountWindowStats(context.Context, int64, time.Time) (*usagestats.AccountStats, error) { + r.calls++ + return r.stats, r.err +} + +type grokHybridUpstream struct { + httpUpstreamRecorder + mu sync.Mutex + requests []*http.Request + bodies [][]byte + weeklyUsagePercent *float64 + monthlyLimitCents *float64 + activeStatus int + activeHeaders http.Header + billingStarted chan struct{} + billingRelease <-chan struct{} + billingStartOnce sync.Once + billingStatus int + billingHeaders http.Header +} + +func (u *grokHybridUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + var body []byte + if req != nil && req.Body != nil { + body, _ = io.ReadAll(req.Body) + } + u.mu.Lock() + u.requests = append(u.requests, req) + u.bodies = append(u.bodies, body) + u.mu.Unlock() + + if req.URL.Path == "/v1/responses" { + status := u.activeStatus + if status == 0 { + status = http.StatusOK + } + headers := u.activeHeaders + if headers == nil { + headers = http.Header{ + "X-Ratelimit-Limit-Tokens": []string{"2000000"}, + "X-Ratelimit-Remaining-Tokens": []string{"1500000"}, + } + } + return &http.Response{StatusCode: status, Header: headers, Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`))}, nil + } + if u.billingStarted != nil { + u.billingStartOnce.Do(func() { close(u.billingStarted) }) + } + if u.billingRelease != nil { + select { + case <-u.billingRelease: + case <-req.Context().Done(): + return nil, req.Context().Err() + } + } + if u.billingStatus != 0 && u.billingStatus != http.StatusOK { + return &http.Response{ + StatusCode: u.billingStatus, + Header: u.billingHeaders, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"billing limited"}}`)), + }, nil + } + + if req.URL.RawQuery == "format=credits" { + usage := "" + if u.weeklyUsagePercent != nil { + usage = `,"creditUsagePercent":` + strconv.FormatFloat(*u.weeklyUsagePercent, 'f', -1, 64) + } + payload := `{"config":{"currentPeriod":{"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"}` + usage + `}}` + return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(payload))}, nil + } + monthlyLimit := "" + if u.monthlyLimitCents != nil { + monthlyLimit = `,"monthlyLimit":{"val":` + strconv.FormatFloat(*u.monthlyLimitCents, 'f', -1, 64) + `}` + } + monthlyPayload := `{"config":{"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"` + monthlyLimit + `}}` + return &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(monthlyPayload)), + }, nil +} + +func (u *grokHybridUpstream) snapshot() ([]*http.Request, [][]byte) { + u.mu.Lock() + defer u.mu.Unlock() + requests := append([]*http.Request(nil), u.requests...) + bodies := make([][]byte, len(u.bodies)) + for i := range u.bodies { + bodies[i] = append([]byte(nil), u.bodies[i]...) + } + return requests, bodies +} + func (r *grokQuotaProxyRepo) GetByID(_ context.Context, id int64) (*Proxy, error) { r.calls++ return r.proxies[id], nil @@ -102,7 +206,7 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { result, err := svc.ProbeUsage(context.Background(), 42) require.NoError(t, err) require.Equal(t, http.StatusOK, result.StatusCode) - require.Equal(t, "grok-4.3", result.Model) + require.Equal(t, "grok-4.5", result.Model) require.True(t, result.HeadersObserved) require.NotNil(t, result.Snapshot) require.True(t, result.Snapshot.HeadersObserved) @@ -115,7 +219,7 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) - require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) require.Contains(t, string(upstream.lastBody), `"max_output_tokens":1`) require.Contains(t, string(upstream.lastBody), `"store":false`) require.NotNil(t, repo.updates[42][grokQuotaSnapshotExtraKey]) @@ -152,8 +256,8 @@ func TestGrokQuotaServiceProbeUsageIgnoresAccountGrokMapping(t *testing.T) { result, err := svc.ProbeUsage(context.Background(), 47) require.NoError(t, err) - require.Equal(t, "grok-4.3", result.Model) - require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "grok-4.5", result.Model) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) require.NotContains(t, string(upstream.lastBody), "grok-composer") } @@ -185,7 +289,7 @@ func TestGrokQuotaServiceProbeUsageReportsProbeModelOnUpstreamError(t *testing.T _, err := svc.ProbeUsage(context.Background(), 48) require.Error(t, err) require.Equal(t, "GROK_QUOTA_PROBE_UPSTREAM_ERROR", infraerrors.Reason(err)) - require.Contains(t, infraerrors.Message(err), `probe model "grok-4.3"`) + require.Contains(t, infraerrors.Message(err), `probe model "grok-4.5"`) } func TestGrokQuotaServiceProbeUsageLoadsProxyWhenAccountEdgeMissing(t *testing.T) { @@ -308,6 +412,299 @@ func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { require.Zero(t, repo.tempUnschedCalls) } +func TestGrokQuotaServiceQueryQuotaFreeFallsBackToGrok45(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 51, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + result, err := svc.QueryQuota(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, "hybrid_probe", result.Source) + require.Equal(t, "grok-4.5", result.Model) + require.NotNil(t, result.Billing) + require.Nil(t, result.Billing.UsagePercent) + require.NotNil(t, result.Snapshot) + require.NotNil(t, result.Snapshot.Tokens) + require.EqualValues(t, 2_000_000, *result.Snapshot.Tokens.Limit) + require.True(t, result.HeadersObserved) + + requests, bodies := upstream.snapshot() + require.Len(t, requests, 3) + responseCalls := 0 + for i, req := range requests { + if req.URL.Path != "/v1/responses" { + continue + } + responseCalls++ + require.Equal(t, http.MethodPost, req.Method) + require.Equal(t, "grok-4.5", gjson.GetBytes(bodies[i], "model").String()) + require.EqualValues(t, 1, gjson.GetBytes(bodies[i], "max_output_tokens").Int()) + } + require.Equal(t, 1, responseCalls) +} + +func TestGrokQuotaServiceQueryQuotaPaidBillingSkipsActiveProbe(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 52, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + usagePercent := 25.0 + upstream := &grokHybridUpstream{weeklyUsagePercent: &usagePercent} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + result, err := svc.QueryQuota(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, "billing_probe", result.Source) + require.NotNil(t, result.Billing) + require.InDelta(t, usagePercent, *result.Billing.UsagePercent, 1e-9) + require.Nil(t, result.Snapshot) + require.Empty(t, result.Model) + + requests, _ := upstream.snapshot() + require.Len(t, requests, 2) + for _, req := range requests { + require.Equal(t, "/v1/billing", req.URL.Path) + } +} + +func TestGrokQuotaServiceQueryQuotaCustomPaidMonthlyLimitSkipsActiveProbe(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 57, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + monthlyLimit := 25_000.0 + upstream := &grokHybridUpstream{monthlyLimitCents: &monthlyLimit} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + result, err := svc.QueryQuota(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, "billing_probe", result.Source) + require.NotNil(t, result.Billing) + require.InDelta(t, monthlyLimit, *result.Billing.MonthlyLimitCents, 1e-9) + require.Nil(t, result.Snapshot) + + requests, _ := upstream.snapshot() + require.Len(t, requests, 2) + for _, req := range requests { + require.Equal(t, "/v1/billing", req.URL.Path) + } +} + +func TestGrokLocalUsageForBillingOnlyReturnsAvailableWindows(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 7, 13, 12, 0, 0, 0, time.UTC) + billing := &xai.BillingSummary{ + PeriodType: "weekly", + PeriodStart: now.Add(-4 * 24 * time.Hour).Format(time.RFC3339), + PeriodEnd: now.Add(3 * 24 * time.Hour).Format(time.RFC3339), + } + + t.Run("valid weekly window", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{stats: &usagestats.AccountStats{Tokens: 1_500_000}} + weekly, monthly := grokLocalUsageForBilling(context.Background(), repo, 57, billing, now) + require.NotNil(t, weekly) + require.EqualValues(t, 1_500_000, weekly.Tokens) + require.Nil(t, monthly) + require.Equal(t, 1, repo.calls) + }) + + t.Run("query failure", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{err: context.DeadlineExceeded} + weekly, monthly := grokLocalUsageForBilling(context.Background(), repo, 57, billing, now) + require.Nil(t, weekly) + require.Nil(t, monthly) + require.Equal(t, 1, repo.calls) + }) + + t.Run("missing billing window", func(t *testing.T) { + repo := &grokQuotaUsageLogRepo{} + weekly, monthly := grokLocalUsageForBilling(context.Background(), repo, 57, nil, now) + require.Nil(t, weekly) + require.Nil(t, monthly) + require.Zero(t, repo.calls) + }) +} + +func TestAccountUsageServiceGrokRefreshUsesBillingOnly(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 54, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{} + quotaService := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + usageService := &AccountUsageService{ + grokQuotaFetcher: NewGrokQuotaFetcher(), + grokQuotaService: quotaService, + cache: NewUsageCache(), + } + + usage, err := usageService.getGrokUsage(context.Background(), account, false) + require.NoError(t, err) + require.NotNil(t, usage.GrokBilling) + require.Nil(t, usage.GrokBilling.UsagePercent) + + requests, _ := upstream.snapshot() + require.Len(t, requests, 2) + for _, req := range requests { + require.Equal(t, http.MethodGet, req.Method) + require.Equal(t, "/v1/billing", req.URL.Path) + } +} + +func TestGrokQuotaServiceProbeFlightsDeduplicateBillingAndSeparateActive(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 55, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + billingStarted := make(chan struct{}) + billingRelease := make(chan struct{}) + upstream := &grokHybridUpstream{billingStarted: billingStarted, billingRelease: billingRelease} + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + type probeOutcome struct { + result *GrokQuotaProbeResult + err error + } + billingOutcomes := make(chan probeOutcome, 2) + go func() { + result, err := svc.ProbeBilling(context.Background(), account.ID) + billingOutcomes <- probeOutcome{result: result, err: err} + }() + <-billingStarted + secondStarted := make(chan struct{}) + go func() { + close(secondStarted) + result, err := svc.ProbeBilling(context.Background(), account.ID) + billingOutcomes <- probeOutcome{result: result, err: err} + }() + <-secondStarted + time.Sleep(25 * time.Millisecond) + + activeResult, err := svc.ProbeUsage(context.Background(), account.ID) + require.NoError(t, err) + require.NotNil(t, activeResult.Snapshot) + close(billingRelease) + for range 2 { + outcome := <-billingOutcomes + require.NoError(t, outcome.err) + require.NotNil(t, outcome.result.Billing) + } + + requests, _ := upstream.snapshot() + billingCalls := 0 + activeCalls := 0 + for _, req := range requests { + switch req.URL.Path { + case "/v1/billing": + billingCalls++ + case "/v1/responses": + activeCalls++ + } + } + require.Equal(t, 2, billingCalls) + require.Equal(t, 1, activeCalls) +} + +func TestGrokQuotaServiceBilling429DoesNotPauseModelScheduling(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 56, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{ + billingStatus: http.StatusTooManyRequests, + billingHeaders: http.Header{"Retry-After": []string{"45"}}, + } + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + result, err := svc.ProbeBilling(context.Background(), account.ID) + + require.Error(t, err) + require.Nil(t, result) + require.Zero(t, repo.rateLimitedCalls) +} + +func TestGrokQuotaServiceQueryQuotaFree429PersistsLimitAndKeepsBilling(t *testing.T) { + t.Parallel() + + account := &Account{ + ID: 53, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &grokHybridUpstream{ + activeStatus: http.StatusTooManyRequests, + activeHeaders: http.Header{"Retry-After": []string{"45"}}, + } + svc := NewGrokQuotaService(repo, nil, NewGrokTokenProvider(repo, nil), upstream) + + result, err := svc.QueryQuota(context.Background(), account.ID) + require.NoError(t, err) + require.Equal(t, http.StatusTooManyRequests, result.StatusCode) + require.NotNil(t, result.Billing) + require.NotNil(t, result.Snapshot) + 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) +} + func TestGrokQuotaServiceResetQuotaUnsupported(t *testing.T) { t.Parallel() diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 379b586136..1906b83780 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -23,6 +23,7 @@ const ( grokComposerImageBridgeMaxOutputTokens = 512 grokUpstreamUserAgent = "sub2api-grok/1.0" grokCLIVersion = "0.2.93" + grokDefaultResponsesModel = "grok-4.5" grokRateLimitFallbackCooldown = 2 * time.Minute ) @@ -41,7 +42,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( upstreamModel := account.GetMappedModel(originalModel) if strings.TrimSpace(upstreamModel) == "" { - upstreamModel = "grok-4.3" + upstreamModel = grokDefaultResponsesModel } cacheIdentity := resolveGrokCacheIdentity(c, body, "", upstreamModel) patchedBody, err := patchGrokResponsesBody(body, upstreamModel) diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index a3edfacf9d..b45e0bccc0 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -853,12 +853,12 @@ func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *tes require.Equal(t, http.StatusOK, recorder.Code) } -func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T) { +func TestForwardGrokResponsesStreamingDefaultsEmptyModelTo45AndSnapshots(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) - body := []byte(`{"model":"grok","input":"hi","stream":true,"reasoning_effort":"high"}`) + body := []byte(`{"input":"hi","stream":true,"reasoning_effort":"high"}`) c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) c.Request.Header.Set("Content-Type", "application/json") c.Request.Header.Set("OpenAI-Beta", "responses=experimental") @@ -905,7 +905,7 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T) accountRepo: repo, } - result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", true, time.Now()) + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "", true, time.Now()) require.NoError(t, err) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index 0afb9181a1..a5aba67c6d 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -431,7 +431,7 @@ func resolveGrokWSUpstreamModel(account *Account, body []byte, originalModel str } } if upstreamModel == "" { - upstreamModel = "grok-4.3" + upstreamModel = grokDefaultResponsesModel } return upstreamModel } diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index 0105c7a331..d2046b9006 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -178,6 +178,51 @@ func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) { require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) } +func TestProxyOpenAIWSHTTPBridgeTurnForGrokDefaultsEmptyModelTo45(t *testing.T) { + gin.SetMode(gin.TestMode) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"resp_grok_default","model":"grok-4.5"}}`, + "", + `data: {"type":"response.completed","response":{"id":"resp_grok_default","model":"grok-4.5","usage":{"input_tokens":1,"output_tokens":1}}}`, + "", + }, "\n"))), + }} + svc := &OpenAIGatewayService{ + cfg: &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}, + httpUpstream: upstream, + } + account := &Account{ + ID: 72, + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{"base_url": xai.DefaultCLIBaseURL}, + } + payload := []byte(`{"type":"response.create","generate":true,"stream":true,"input":"hi"}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil) + var events [][]byte + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), c, account, "access-token", payload, len(payload), + "", "", "", "", "", 1, + func(message []byte) error { + events = append(events, append([]byte(nil), message...)) + return nil + }, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, grokDefaultResponsesModel, gjson.GetBytes(upstream.lastBody, "model").String()) + require.Len(t, events, 2) +} + func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) { gin.SetMode(gin.TestMode) diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 7258ff05a3..d5d9124ec2 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -140,8 +140,9 @@ func ProvideGrokQuotaService( proxyRepo ProxyRepository, tokenProvider *GrokTokenProvider, httpUpstream HTTPUpstream, + usageLogRepo UsageLogRepository, ) *GrokQuotaService { - return NewGrokQuotaService(accountRepo, proxyRepo, tokenProvider, httpUpstream) + return NewGrokQuotaService(accountRepo, proxyRepo, tokenProvider, httpUpstream, usageLogRepo) } // ProvideGeminiTokenProvider creates GeminiTokenProvider with OAuthRefreshAPI injection diff --git a/frontend/src/api/admin/grok.ts b/frontend/src/api/admin/grok.ts index c0055d4dcc..e50fd222ed 100644 --- a/frontend/src/api/admin/grok.ts +++ b/frontend/src/api/admin/grok.ts @@ -4,6 +4,9 @@ */ import { apiClient } from '../client' +import type { GrokBillingSummary, GrokQuotaWindow, WindowStats } from '@/types' + +export type { GrokBillingSummary, GrokQuotaWindow } from '@/types' export interface GrokAuthUrlResponse { auth_url: string @@ -39,13 +42,6 @@ export interface GrokTokenInfo { [key: string]: unknown } -export interface GrokQuotaWindow { - limit?: number | null - remaining?: number | null - reset_unix?: number | null - reset_at?: string | null -} - export interface GrokQuotaSnapshot { requests?: GrokQuotaWindow | null tokens?: GrokQuotaWindow | null @@ -62,13 +58,18 @@ export interface GrokQuotaSnapshot { } export interface GrokQuotaProbeResult { - source: 'active_probe' - model: string + source: 'active_probe' | 'billing_probe' | 'hybrid_probe' + model?: string + billing?: GrokBillingSummary | null snapshot?: GrokQuotaSnapshot | null + local_usage_7d?: WindowStats | null + local_usage_monthly?: WindowStats | null status_code?: number headers_observed: boolean reset_supported: boolean fetched_at: number + persisted?: boolean + probe_error?: string } export interface GrokQuotaResetResult { diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 5d7af74fc6..d751bf0429 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -382,7 +382,15 @@ + +
{{ t('admin.accounts.usageWindow.grokRetryAfter', { time: grokRetryAfterLabel }) }}
@@ -409,7 +424,7 @@
{{ grokQuotaStatusLine }}
- +
-
@@ -602,6 +617,7 @@ import { ref, computed, onMounted, onBeforeUnmount, onUnmounted, watch } from 'vue' import { useI18n } from 'vue-i18n' import { adminAPI } from '@/api/admin' +import type { GrokQuotaProbeResult } from '@/api/admin/grok' import type { Account, AccountUsageInfo, GeminiCredentials, WindowStats } from '@/types' import { buildOpenAIUsageRefreshKey } from '@/utils/accountUsageRefresh' import { enqueueUsageRequest } from '@/utils/usageLoadQueue' @@ -614,6 +630,8 @@ import GrokQuotaProbeCell from './GrokQuotaProbeCell.vue' // Module-level cache shared across all AccountUsageCell instances const _usageCache = new Map() 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<{ @@ -1047,9 +1065,62 @@ const makeGrokQuotaBar = (quota?: { limit?: number | null; remaining?: number | const grokRequestQuotaBar = computed(() => makeGrokQuotaBar(usageInfo.value?.grok_request_quota)) const grokTokenQuotaBar = computed(() => makeGrokQuotaBar(usageInfo.value?.grok_token_quota)) +const grokLocalUsage = computed(() => + props.todayStats || + usageInfo.value?.grok_local_usage || + usageInfo.value?.grok_local_usage_7d || + usageInfo.value?.grok_local_usage_monthly || + null +) +const grokFreeQuotaUsage = computed(() => + usageInfo.value?.grok_local_usage_7d || + props.todayStats || + usageInfo.value?.grok_local_usage || + null +) +const grokBilling = computed(() => usageInfo.value?.grok_billing || null) +const grokWeeklyBillingBar = computed((): GrokQuotaBarInfo | null => { + const billing = grokBilling.value + if (billing?.period_type?.toLowerCase() !== 'weekly' || billing.usage_percent == null) { + return null + } + return { + utilization: Math.min(100, Math.max(0, billing.usage_percent)), + resetsAt: billing.period_end || null + } +}) +const grokPlanLabelIsFree = (value: string) => value.includes('free') || value.includes('basic') +const grokPlanLabelIsPaid = (value: string) => { + return value !== '' && !grokPlanLabelIsFree(value) && !value.includes('unknown') +} +const grokIsFree = computed(() => { + if (props.account.platform !== 'grok' || props.account.type !== 'oauth') return false + const billing = grokBilling.value + if ( + billing?.usage_percent != null || + billing?.used_percent != null || + (billing?.monthly_limit_cents != null && billing.monthly_limit_cents > 0) + ) return false + + const plan = (billing?.plan || '').trim().toLowerCase() + const tier = (usageInfo.value?.subscription_tier || '').trim().toLowerCase() + const entitlement = (usageInfo.value?.grok_entitlement_status || '').toLowerCase() + if (grokPlanLabelIsPaid(plan) || grokPlanLabelIsPaid(tier)) return false + if ( + grokPlanLabelIsFree(plan) || + grokPlanLabelIsFree(tier) || + grokPlanLabelIsFree(entitlement) + ) return true + return billing != null +}) +const grokFreeTokenBar = computed(() => { + if (!grokIsFree.value || !grokFreeQuotaUsage.value) return null + const used = Math.max(0, grokFreeQuotaUsage.value.tokens || 0) + return { utilization: Math.min(100, (used / GROK_FREE_TOKEN_LIMIT) * 100) } +}) const grokQuotaUnknown = computed(() => { if (props.account.platform !== 'grok') return false - if (grokRequestQuotaBar.value || grokTokenQuotaBar.value) return false + if (grokBilling.value || grokFreeTokenBar.value || grokRequestQuotaBar.value || grokTokenQuotaBar.value) return false return usageInfo.value?.grok_quota_snapshot_state !== 'observed' }) const grokQuotaUnknownLabel = computed(() => { @@ -1080,7 +1151,6 @@ const grokQuotaStatusLine = computed(() => { } return parts.length > 0 ? parts.join(' | ') : null }) -const grokLocalUsage = computed(() => usageInfo.value?.grok_local_usage || props.todayStats || null) const grokEntitlementLabel = computed(() => { const status = (usageInfo.value?.grok_entitlement_status || '').trim() return status || null @@ -1283,6 +1353,34 @@ const loadActiveUsage = async () => { } } +const handleGrokProbed = (result: GrokQuotaProbeResult) => { + const current = usageInfo.value + if (!current) return + const snapshot = result.snapshot + const merged: AccountUsageInfo = { + ...current, + grok_billing: result.billing ?? current.grok_billing, + grok_local_usage_7d: result.local_usage_7d ?? current.grok_local_usage_7d, + grok_local_usage_monthly: result.local_usage_monthly ?? current.grok_local_usage_monthly, + grok_request_quota: snapshot?.requests ?? current.grok_request_quota, + grok_token_quota: snapshot?.tokens ?? current.grok_token_quota, + grok_retry_after_seconds: snapshot?.retry_after_seconds ?? current.grok_retry_after_seconds, + grok_entitlement_status: snapshot?.entitlement_status || current.grok_entitlement_status, + grok_quota_snapshot_state: result.billing + ? 'billing_observed' + : snapshot?.headers_observed + ? 'observed' + : current.grok_quota_snapshot_state, + grok_last_quota_probe_at: result.billing?.fetched_at ?? snapshot?.last_probe_at ?? current.grok_last_quota_probe_at, + grok_last_headers_seen_at: snapshot?.last_headers_seen_at ?? current.grok_last_headers_seen_at, + grok_last_status_code: result.status_code ?? snapshot?.status_code ?? current.grok_last_status_code, + error: result.billing || snapshot ? undefined : current.error, + error_code: result.billing || snapshot ? undefined : current.error_code + } + usageInfo.value = merged + _usageCache.set(props.account.id, { data: merged, ts: Date.now() }) +} + // ===== API Key quota progress bars ===== interface QuotaBarInfo { diff --git a/frontend/src/components/account/GrokQuotaProbeCell.vue b/frontend/src/components/account/GrokQuotaProbeCell.vue index 183ab3e7ba..fa6fe29bec 100644 --- a/frontend/src/components/account/GrokQuotaProbeCell.vue +++ b/frontend/src/components/account/GrokQuotaProbeCell.vue @@ -55,6 +55,8 @@ const props = defineProps<{ account: Account }>() +const emit = defineEmits<{ probed: [result: GrokQuotaProbeResult] }>() + const { t } = useI18n() const visible = computed(() => props.account.platform === 'grok' && props.account.type === 'oauth') @@ -92,18 +94,27 @@ const retryAfterLabel = computed(() => { const summary = computed(() => { const snapshot = data.value?.snapshot if (!data.value) return '' - if (!snapshot) return t('admin.accounts.usageWindow.grokNoHeaders') - const parts = [ - formatWindow(t('admin.accounts.usageWindow.grokRequests'), snapshot.requests), - formatWindow(t('admin.accounts.usageWindow.grokTokens'), snapshot.tokens) - ].filter(Boolean) + const billing = data.value.billing + const parts: Array = [] + if (billing?.period_type?.toLowerCase() === 'weekly' && billing.usage_percent != null) { + parts.push(t('admin.accounts.usageWindow.grokWeeklyUsage', { + percent: Math.round(Math.min(100, Math.max(0, billing.usage_percent))) + })) + } + if (snapshot) { + parts.push( + formatWindow(t('admin.accounts.usageWindow.grokRequests'), snapshot.requests), + formatWindow(t('admin.accounts.usageWindow.grokTokens'), snapshot.tokens) + ) + } if (retryAfterLabel.value) { parts.push(t('admin.accounts.usageWindow.grokRetryAfter', { time: retryAfterLabel.value })) } - if (snapshot.entitlement_status) { + if (snapshot?.entitlement_status) { parts.push(snapshot.entitlement_status) } - return parts.length > 0 ? parts.join(' | ') : t('admin.accounts.usageWindow.grokNoHeaders') + const visibleParts = parts.filter((part): part is string => Boolean(part)) + return visibleParts.length > 0 ? visibleParts.join(' | ') : t('admin.accounts.usageWindow.grokNoHeaders') }) const truncatedError = computed(() => { @@ -117,6 +128,8 @@ const handleProbe = async () => { error.value = null try { data.value = await adminAPI.grok.queryQuota(props.account.id) + error.value = data.value.probe_error || null + emit('probed', data.value) } catch (e) { error.value = extractErrorMessage(e) } finally { diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index 2abf6513ca..988c4fc7db 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -660,6 +660,339 @@ describe('AccountUsageCell', () => { expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokTokens|25|true') }) + it('Grok OAuth uses the official weekly billing percentage when available', async () => { + getUsage.mockResolvedValue({ + grok_billing: { + period_type: 'weekly', + usage_percent: 37, + period_end: '2026-07-16T03:25:00Z', + plan: 'SuperGrok' + }, + grok_local_usage: { + requests: 5, + tokens: 2_200_000, + cost: 4.42, + standard_cost: 4.42, + user_cost: 0.44 + }, + grok_request_quota: { limit: 100, remaining: 100 }, + grok_token_quota: { limit: 2_000_000, remaining: 2_000_000 } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4201, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization', 'resetsAt', 'remainingCapacity'], + template: '
{{ label }}|{{ utilization }}|{{ resetsAt }}|{{ remainingCapacity }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('7d|37|2026-07-16T03:25:00Z') + expect(wrapper.text()).not.toContain('admin.accounts.usageWindow.grokRequests|') + expect(wrapper.text()).not.toContain('admin.accounts.usageWindow.grokTokens|') + expect(wrapper.text()).not.toContain('2M|') + }) + + 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 }) => { + getUsage.mockResolvedValue({ + grok_billing: { + period_type: 'weekly', + usage_percent: null, + plan: '' + }, + grok_local_usage: { + requests: 5, + tokens, + cost: 0, + standard_cost: 0, + user_cost: 0 + }, + grok_request_quota: { limit: 100, remaining: 100 }, + grok_token_quota: { limit: 2_000_000, remaining: 2_000_000 } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4300 + expected, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain(`2M|${expected}`) + expect(wrapper.findAll('span').filter((node) => node.text() === compact)).toHaveLength(1) + expect(wrapper.findAll('.usage-bar')).toHaveLength(1) + expect(wrapper.text()).not.toContain('admin.accounts.usageWindow.grokRequests|') + expect(wrapper.text()).not.toContain('admin.accounts.usageWindow.grokTokens|') + }) + + it('Grok Free uses the weekly billing window instead of today-only usage', async () => { + getUsage.mockResolvedValue({ + grok_billing: { period_type: 'weekly', usage_percent: null, plan: '' }, + grok_local_usage: { + requests: 2, + tokens: 200_000, + cost: 0, + standard_cost: 0 + }, + grok_local_usage_7d: { + requests: 12, + tokens: 1_500_000, + cost: 0, + standard_cost: 0 + } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4398, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('2M|75') + expect(wrapper.text()).toContain('200.0K') + }) + + it('Grok Free falls back to refreshed today stats when weekly usage is unavailable', async () => { + getUsage.mockResolvedValue({ + grok_billing: { period_type: 'weekly', usage_percent: null, plan: '' }, + grok_local_usage: { + requests: 1, + tokens: 250_000, + cost: 0, + standard_cost: 0, + user_cost: 0 + } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4399, platform: 'grok', type: 'oauth', extra: {} }), + todayStats: { + requests: 4, + tokens: 1_000_000, + cost: 0, + standard_cost: 0, + user_cost: 0 + } + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('2M|50') + expect(wrapper.text()).toContain('1.0M') + expect(wrapper.text()).not.toContain('250K') + }) + + it('Grok paid plans are not mistaken for Free when weekly usage is temporarily missing', async () => { + getUsage.mockResolvedValue({ + grok_billing: { + period_type: 'weekly', + usage_percent: null, + plan: 'SuperGrok Heavy' + }, + grok_entitlement_status: 'free', + grok_local_usage: { + requests: 2, + tokens: 2_000_000, + cost: 1, + standard_cost: 1 + }, + grok_token_quota: { limit: 1_000, remaining: 250 } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4401, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokTokens|25') + expect(wrapper.text()).not.toContain('2M|') + }) + + it('Grok custom paid monthly limits override stale Free entitlement', async () => { + getUsage.mockResolvedValue({ + grok_billing: { + period_type: 'weekly', + usage_percent: null, + monthly_limit_cents: 25_000, + plan: '' + }, + grok_entitlement_status: 'free', + grok_local_usage: { + requests: 2, + tokens: 2_000_000, + cost: 1, + standard_cost: 1 + }, + grok_token_quota: { limit: 1_000, remaining: 250 } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4402, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokTokens|25') + expect(wrapper.text()).not.toContain('2M|') + }) + + it('Grok credential Free tier keeps the 2M fallback when billing is unavailable', async () => { + getUsage.mockResolvedValue({ + subscription_tier: 'FREE', + grok_local_usage: { + requests: 3, + tokens: 1_000_000, + cost: 0, + standard_cost: 0 + } + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4403, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization'], + template: '
{{ label }}|{{ utilization }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('2M|50') + }) + + it('Grok manual probes merge billing, quota headers, and local usage', async () => { + getUsage.mockResolvedValue({ + grok_quota_snapshot_state: 'no_headers', + error: 'stale error', + error_code: 'quota_unknown' + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ id: 4501, platform: 'grok', type: 'oauth', extra: {} }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization', 'resetsAt'], + template: '
{{ label }}|{{ utilization }}|{{ resetsAt }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: { + emits: ['probed'], + template: `` + } + } + } + }) + + await flushPromises() + await wrapper.get('.probe').trigger('click') + + expect(wrapper.text()).toContain('7d|42|2026-07-17T00:00:00Z') + expect(wrapper.text()).toContain('1.0M') + expect(wrapper.text()).toContain('ACTIVE') + expect(wrapper.text()).not.toContain('stale error') + }) + it('Key 账号在 today stats loading 时显示骨架屏', async () => { const wrapper = mount(AccountUsageCell, { props: { diff --git a/frontend/src/components/account/__tests__/GrokQuotaProbeCell.spec.ts b/frontend/src/components/account/__tests__/GrokQuotaProbeCell.spec.ts new file mode 100644 index 0000000000..9b431b339e --- /dev/null +++ b/frontend/src/components/account/__tests__/GrokQuotaProbeCell.spec.ts @@ -0,0 +1,54 @@ +import { flushPromises, mount } from '@vue/test-utils' +import { beforeEach, describe, expect, it, vi } from 'vitest' +import GrokQuotaProbeCell from '../GrokQuotaProbeCell.vue' +import type { Account } from '@/types' + +const { queryQuota } = vi.hoisted(() => ({ + queryQuota: vi.fn() +})) + +vi.mock('@/api/admin', () => ({ + adminAPI: { + grok: { queryQuota } + } +})) + +vi.mock('vue-i18n', () => ({ + useI18n: () => ({ + t: (key: string, params?: Record) => + params?.percent == null ? key : `${key}:${params.percent}` + }) +})) + +const account = { + id: 99, + platform: 'grok', + type: 'oauth' +} as Account + +describe('GrokQuotaProbeCell', () => { + beforeEach(() => { + queryQuota.mockReset() + }) + + it('keeps billing data while exposing a failed Free quota fallback', async () => { + queryQuota.mockResolvedValue({ + source: 'hybrid_probe', + billing: { period_type: 'weekly', usage_percent: null }, + headers_observed: false, + reset_supported: false, + fetched_at: 1, + probe_error: 'upstream returned 402 for probe model "grok-4.5"' + }) + const wrapper = mount(GrokQuotaProbeCell, { props: { account } }) + + await wrapper.get('button').trigger('click') + await flushPromises() + + expect(wrapper.text()).toContain('upstream returned 402 for probe model "grok-4.5"') + expect(wrapper.emitted('probed')?.[0]?.[0]).toMatchObject({ + billing: { period_type: 'weekly', usage_percent: null }, + probe_error: 'upstream returned 402 for probe model "grok-4.5"' + }) + }) +}) diff --git a/frontend/src/i18n/locales/en/admin/accounts.ts b/frontend/src/i18n/locales/en/admin/accounts.ts index f75c4a8060..c37c969d48 100644 --- a/frontend/src/i18n/locales/en/admin/accounts.ts +++ b/frontend/src/i18n/locales/en/admin/accounts.ts @@ -1205,6 +1205,7 @@ export default { claude: 'Claude', grokRequests: 'Req', grokTokens: 'Tok', + grokWeeklyUsage: 'Weekly {percent}%', grokUnknown: 'Grok quota is unknown until the first upstream response includes xAI rate-limit headers.', grokRetryAfter: 'Retry after {time}', grokProbe: 'Probe', diff --git a/frontend/src/i18n/locales/zh/admin/accounts.ts b/frontend/src/i18n/locales/zh/admin/accounts.ts index 4887a1c045..1e5e6a98df 100644 --- a/frontend/src/i18n/locales/zh/admin/accounts.ts +++ b/frontend/src/i18n/locales/zh/admin/accounts.ts @@ -317,6 +317,7 @@ export default { claude: 'Claude', grokRequests: '请求', grokTokens: 'Token', + grokWeeklyUsage: '周额度已用 {percent}%', grokUnknown: 'Grok 配额需等待首次上游响应返回 xAI rate-limit 头后显示。', grokRetryAfter: '{time} 后重试', grokProbe: '探测', diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index bd8137829e..7cddbace60 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -1020,10 +1020,38 @@ export interface AntigravityModelQuota { } export interface GrokQuotaWindow { - limit?: number - remaining?: number - reset_unix?: number - reset_at?: string + limit?: number | null + remaining?: number | null + reset_unix?: number | null + reset_at?: string | null +} + +export interface GrokBillingProductUsage { + product: string + usage_percent?: number | null +} + +export interface GrokBillingSummary { + period_type?: string + usage_percent?: number | null + period_start?: string + period_end?: string + product_usage?: GrokBillingProductUsage[] + monthly_limit_cents?: number | null + used_cents?: number | null + included_used_cents?: number | null + billing_period_start?: string + billing_period_end?: string + used_percent?: number | null + plan?: string + status_code?: number + source?: string + fetched_at?: string + updated_at?: string + weekly_updated_at?: string + monthly_updated_at?: string + partial?: boolean + failed_windows?: string[] } export interface AccountUsageInfo { @@ -1049,6 +1077,11 @@ export interface AccountUsageInfo { grok_last_headers_seen_at?: string grok_last_status_code?: number grok_local_usage?: WindowStats | null + grok_local_usage_7d?: WindowStats | null + grok_local_usage_monthly?: WindowStats | null + grok_billing?: GrokBillingSummary | null + subscription_tier?: string + subscription_tier_raw?: string ai_credits?: Array<{ credit_type?: string amount?: number