mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-22 06:40:21 +08:00
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:
@@ -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])
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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: '探测',
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user