Merge pull request #4188 from superman2003/fix/grok-free-quota-429-20260713

feat(grok): improve free quota probing and usage display
This commit is contained in:
Wesley Liddick
2026-07-14 10:14:41 +08:00
committed by GitHub
27 changed files with 2140 additions and 104 deletions
+2 -2
View File
@@ -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)
@@ -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
@@ -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])
}
+372
View File
@@ -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
}
+127
View File
@@ -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
}
@@ -61,6 +61,7 @@ var schedulerNeutralExtraKeyPrefixes = []string{
var schedulerNeutralExtraKeys = map[string]struct{}{
"codex_usage_updated_at": {},
"grok_billing_snapshot": {},
"session_window_utilization": {},
}
@@ -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},
}))
}
@@ -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
@@ -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)
+106 -12
View File
@@ -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) {
+122 -24
View File
@@ -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
@@ -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()
+242 -9
View File
@@ -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
@@ -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()
@@ -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)
@@ -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"))
@@ -431,7 +431,7 @@ func resolveGrokWSUpstreamModel(account *Account, body []byte, originalModel str
}
}
if upstreamModel == "" {
upstreamModel = "grok-4.3"
upstreamModel = grokDefaultResponsesModel
}
return upstreamModel
}
@@ -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)
+2 -1
View File
@@ -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
+10 -9
View File
@@ -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 {
@@ -382,7 +382,15 @@
</div>
</div>
<UsageProgressBar
v-if="grokRequestQuotaBar"
v-if="grokWeeklyBillingBar"
label="7d"
:utilization="grokWeeklyBillingBar.utilization"
:resets-at="grokWeeklyBillingBar.resetsAt"
:show-now-when-idle="true"
color="indigo"
/>
<UsageProgressBar
v-if="!grokWeeklyBillingBar && !grokIsFree && grokRequestQuotaBar"
:label="t('admin.accounts.usageWindow.grokRequests')"
:utilization="grokRequestQuotaBar.utilization"
:resets-at="grokRequestQuotaBar.resetsAt"
@@ -390,13 +398,20 @@
color="indigo"
/>
<UsageProgressBar
v-if="grokTokenQuotaBar"
v-if="!grokWeeklyBillingBar && !grokIsFree && grokTokenQuotaBar"
:label="t('admin.accounts.usageWindow.grokTokens')"
:utilization="grokTokenQuotaBar.utilization"
:resets-at="grokTokenQuotaBar.resetsAt"
:remaining-capacity="true"
color="emerald"
/>
<UsageProgressBar
v-if="grokFreeTokenBar"
label="2M"
:utilization="grokFreeTokenBar.utilization"
:show-now-when-idle="true"
color="emerald"
/>
<div v-if="grokRetryAfterLabel" class="text-[10px] text-amber-600 dark:text-amber-400">
{{ t('admin.accounts.usageWindow.grokRetryAfter', { time: grokRetryAfterLabel }) }}
</div>
@@ -409,7 +424,7 @@
<div v-if="grokQuotaStatusLine" class="text-[10px] text-gray-500 dark:text-gray-400">
{{ grokQuotaStatusLine }}
</div>
<GrokQuotaProbeCell :account="account" />
<GrokQuotaProbeCell :account="account" @probed="handleGrokProbed" />
</div>
<div v-else class="text-xs text-gray-400">-</div>
</template>
@@ -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<number, { data: AccountUsageInfo; ts: number }>()
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 {
@@ -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<string | null> = []
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 {
@@ -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: '<div class="usage-bar">{{ label }}|{{ utilization }}|{{ resetsAt }}|{{ remainingCapacity }}</div>'
},
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: '<div class="usage-bar">{{ label }}|{{ utilization }}</div>'
},
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: '<div class="usage-bar">{{ label }}|{{ utilization }}</div>'
},
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: '<div class="usage-bar">{{ label }}|{{ utilization }}</div>'
},
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: '<div class="usage-bar">{{ label }}|{{ utilization }}</div>'
},
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: '<div class="usage-bar">{{ label }}|{{ utilization }}</div>'
},
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: '<div class="usage-bar">{{ label }}|{{ utilization }}</div>'
},
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: '<div class="usage-bar">{{ label }}|{{ utilization }}|{{ resetsAt }}</div>'
},
AccountQuotaInfo: true,
GrokQuotaProbeCell: {
emits: ['probed'],
template: `<button class="probe" @click="$emit('probed', {
source: 'hybrid_probe',
billing: { period_type: 'weekly', usage_percent: 42, period_end: '2026-07-17T00:00:00Z' },
snapshot: {
headers_observed: true,
updated_at: '2026-07-13T00:00:00Z',
entitlement_status: 'ACTIVE',
requests: { limit: 100, remaining: 20 }
},
local_usage_7d: { requests: 4, tokens: 1000000, cost: 1, standard_cost: 1, user_cost: 0.5 },
local_usage_monthly: { requests: 7, tokens: 1500000, cost: 2, standard_cost: 2, user_cost: 1 },
status_code: 200,
headers_observed: true,
reset_supported: false,
fetched_at: 1
})">probe</button>`
}
}
}
})
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: {
@@ -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<string, unknown>) =>
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"'
})
})
})
@@ -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',
@@ -317,6 +317,7 @@ export default {
claude: 'Claude',
grokRequests: '请求',
grokTokens: 'Token',
grokWeeklyUsage: '周额度已用 {percent}%',
grokUnknown: 'Grok 配额需等待首次上游响应返回 xAI rate-limit 头后显示。',
grokRetryAfter: '{time} 后重试',
grokProbe: '探测',
+37 -4
View File
@@ -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