diff --git a/backend/internal/pkg/xai/quota.go b/backend/internal/pkg/xai/quota.go index 11133a05b3..9825f66343 100644 --- a/backend/internal/pkg/xai/quota.go +++ b/backend/internal/pkg/xai/quota.go @@ -43,6 +43,12 @@ type QuotaSnapshot struct { LastProbeAt string `json:"last_probe_at,omitempty"` LastHeadersSeenAt string `json:"last_headers_seen_at,omitempty"` UpdatedAt string `json:"updated_at"` + // Model is the upstream id that produced these rate-limit headers. + Model string `json:"model,omitempty"` + // PlanFrom45Responses is inferred from a grok-4.5 Responses window + // (8300/53M = Heavy). Carried across later non-4.5 overwrites. + PlanFrom45Responses string `json:"plan_from_45_responses,omitempty"` + PlanFrom45ResponsesAt string `json:"plan_from_45_responses_at,omitempty"` } func (s *QuotaSnapshot) HasObservedHeaders() bool { diff --git a/backend/internal/pkg/xai/subscription_tier.go b/backend/internal/pkg/xai/subscription_tier.go index db3e87dee1..44de238e6c 100644 --- a/backend/internal/pkg/xai/subscription_tier.go +++ b/backend/internal/pkg/xai/subscription_tier.go @@ -4,6 +4,17 @@ import ( "encoding/json" "strconv" "strings" + "time" +) + +// GrokQuotaSignalMaxAge bounds how long a grok-4.5 Responses window can +// influence SuperGrok vs Heavy inference. +const GrokQuotaSignalMaxAge = 24 * time.Hour + +const ( + grok45ResponsesModel = "grok-4.5" + grokHeavyQuotaRequestLimit int64 = 8_300 + grokHeavyQuotaTokenLimit int64 = 53_000_000 ) // MapJWTSubscriptionTier maps prod_auth.SubscriptionTier numeric JWT claims @@ -44,8 +55,10 @@ func NormalizeSubscriptionTier(raw string) string { return "supergrok" case "supergrok_lite", "supergroklite": return "supergrok_lite" - case "supergrok_heavy", "supergrokheavy", "supergrokpro": + case "supergrok_heavy", "supergrokheavy": return "supergrok_heavy" + case "supergrok_pro", "supergrokpro": + return "supergrok_pro" case "supergrok_plus", "supergrokplus": return "supergrok_plus" case "x_basic", "xbasic", "basic": @@ -95,3 +108,172 @@ func SubscriptionTierFromJWT(jwt string) string { return "" } } + +// CanonicalGrokPlan resolves SuperGrok vs Heavy when the provider label is +// ambiguous (SuperGrokPro). JWT numeric claims are applied by the caller first. +// Monthly $150/$1500 limits still win when present. +// Rate-limit windows are only used when they came from grok-4.5 Responses. +func CanonicalGrokPlan(monthlyLimitCents *float64, subscriptionTier string, quota *QuotaSnapshot) string { + if plan := resolvePlan(monthlyLimitCents); plan != "" { + return NormalizeSubscriptionTier(plan) + } + + normalized := NormalizeSubscriptionTier(subscriptionTier) + switch normalized { + case "free", "x_basic": + return "free" + case "supergrok_heavy": + return "supergrok_heavy" + case "supergrok_lite": + return "supergrok_lite" + case "supergrok_plus": + return "supergrok_plus" + } + + if isAmbiguousGrokPaidPlan(normalized) { + if hint := Grok45ResponsesPlanHint(quota, time.Time{}); hint != "" { + return hint + } + return "supergrok" + } + return "" +} + +func isAmbiguousGrokPaidPlan(normalized string) bool { + switch normalized { + case "supergrok", "supergrok_pro", "paid", "pro": + return true + default: + return false + } +} + +// IsGrok45ResponsesQuotaModel reports whether model is the grok-4.5 Responses +// id (or a dated grok-4.5-* variant). Empty and other families are false. +func IsGrok45ResponsesQuotaModel(model string) bool { + m := strings.ToLower(strings.TrimSpace(StripGrokProviderPrefix(model))) + return m == grok45ResponsesModel || strings.HasPrefix(m, grok45ResponsesModel+"-") +} + +// Grok45ResponsesPlanHint returns SuperGrok / Heavy inferred from a grok-4.5 +// Responses window. Other models' limits are ignored. +func Grok45ResponsesPlanHint(quota *QuotaSnapshot, now time.Time) string { + if quota == nil { + return "" + } + if plan := NormalizeSubscriptionTier(quota.PlanFrom45Responses); plan == "supergrok" || plan == "supergrok_heavy" { + if isQuotaTimestampFresh(quota.PlanFrom45ResponsesAt, now) { + return plan + } + } + if !IsGrok45ResponsesQuotaModel(quota.Model) || !IsQuotaSnapshotFresh(quota, now) { + return "" + } + if quotaLooksLikeGrokHeavy(quota) { + return "supergrok_heavy" + } + return "" +} + +// ApplyGrok45ResponsesPlanSignal records a grok-4.5 Heavy/SuperGrok hint, or +// copies the previous 4.5 hint when this observation is a different model. +func (s *QuotaSnapshot) ApplyGrok45ResponsesPlanSignal(prev *QuotaSnapshot) { + if s == nil { + return + } + observedAt := firstNonEmptyQuotaTime(s.LastHeadersSeenAt, s.UpdatedAt) + if IsGrok45ResponsesQuotaModel(s.Model) && quotaHasLimitWindow(s) { + if quotaLooksLikeGrokHeavy(s) { + s.PlanFrom45Responses = "supergrok_heavy" + s.PlanFrom45ResponsesAt = observedAt + return + } + s.PlanFrom45Responses = "supergrok" + s.PlanFrom45ResponsesAt = observedAt + return + } + if prev != nil && strings.TrimSpace(prev.PlanFrom45Responses) != "" { + s.PlanFrom45Responses = prev.PlanFrom45Responses + s.PlanFrom45ResponsesAt = prev.PlanFrom45ResponsesAt + } +} + +// QuotaSnapshotObservedAt prefers LastHeadersSeenAt over UpdatedAt so a later +// snapshot rewrite cannot refresh a stale Heavy window. +func QuotaSnapshotObservedAt(snapshot *QuotaSnapshot) (time.Time, bool) { + if snapshot == nil { + return time.Time{}, false + } + return parseQuotaTimestamp(firstNonEmptyQuotaTime(snapshot.LastHeadersSeenAt, snapshot.UpdatedAt)) +} + +// IsQuotaSnapshotFresh reports whether a quota signal is recent enough to +// distinguish SuperGrok from Heavy. +func IsQuotaSnapshotFresh(snapshot *QuotaSnapshot, now time.Time) bool { + observedAt, ok := QuotaSnapshotObservedAt(snapshot) + if !ok { + return false + } + return isTimeFresh(observedAt, now) +} + +func isQuotaTimestampFresh(raw string, now time.Time) bool { + parsed, ok := parseQuotaTimestamp(raw) + if !ok { + return false + } + return isTimeFresh(parsed, now) +} + +func parseQuotaTimestamp(raw string) (time.Time, bool) { + raw = strings.TrimSpace(raw) + if raw == "" { + return time.Time{}, false + } + parsed, err := time.Parse(time.RFC3339, raw) + if err != nil { + return time.Time{}, false + } + return parsed, true +} + +func isTimeFresh(observedAt, now time.Time) bool { + if now.IsZero() { + now = time.Now() + } + age := now.Sub(observedAt) + return age <= GrokQuotaSignalMaxAge && age >= -5*time.Minute +} + +func quotaHasLimitWindow(quota *QuotaSnapshot) bool { + if quota == nil { + return false + } + if quota.Requests != nil && quota.Requests.Limit != nil { + return true + } + return quota.Tokens != nil && quota.Tokens.Limit != nil +} + +func quotaLooksLikeGrokHeavy(quota *QuotaSnapshot) bool { + if quota == nil { + return false + } + var requestLimit, tokenLimit int64 + if quota.Requests != nil && quota.Requests.Limit != nil { + requestLimit = *quota.Requests.Limit + } + if quota.Tokens != nil && quota.Tokens.Limit != nil { + tokenLimit = *quota.Tokens.Limit + } + return requestLimit >= grokHeavyQuotaRequestLimit || tokenLimit >= grokHeavyQuotaTokenLimit +} + +func firstNonEmptyQuotaTime(values ...string) string { + for _, value := range values { + if strings.TrimSpace(value) != "" { + return strings.TrimSpace(value) + } + } + return "" +} diff --git a/backend/internal/pkg/xai/subscription_tier_test.go b/backend/internal/pkg/xai/subscription_tier_test.go index f24a6db217..ef9236d732 100644 --- a/backend/internal/pkg/xai/subscription_tier_test.go +++ b/backend/internal/pkg/xai/subscription_tier_test.go @@ -6,6 +6,7 @@ import ( "encoding/base64" "encoding/json" "testing" + "time" "github.com/stretchr/testify/require" ) @@ -31,7 +32,7 @@ func TestNormalizeSubscriptionTierAliases(t *testing.T) { require.Equal(t, "free", NormalizeSubscriptionTier(" FREE ")) require.Equal(t, "supergrok", NormalizeSubscriptionTier("SuperGrok")) require.Equal(t, "supergrok_heavy", NormalizeSubscriptionTier("SuperGrok Heavy")) - require.Equal(t, "supergrok_heavy", NormalizeSubscriptionTier("SuperGrokPro")) + require.Equal(t, "supergrok_pro", NormalizeSubscriptionTier("SuperGrokPro")) require.Equal(t, "supergrok_lite", NormalizeSubscriptionTier("SuperGrok Lite")) require.Equal(t, "supergrok_lite", NormalizeSubscriptionTier("SuperGrokLite")) require.Equal(t, "x_basic", NormalizeSubscriptionTier("X Basic")) @@ -52,6 +53,89 @@ func TestSubscriptionTierFromJWTUsesNumericClaim(t *testing.T) { require.Empty(t, SubscriptionTierFromJWT("not-a-jwt")) } +func TestCanonicalGrokPlanUsesOnlyGrok45ResponsesWindow(t *testing.T) { + t.Parallel() + + zero := float64(0) + heavyReq, heavyTok := int64(8300), int64(53_000_000) + superReq, superTok := int64(900), int64(15_000_000) + fresh := time.Now().UTC().Format(time.RFC3339) + stale := time.Now().Add(-GrokQuotaSignalMaxAge - time.Hour).UTC().Format(time.RFC3339) + + require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", nil)) + require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrok", nil)) + require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrok Heavy", nil)) + require.Empty(t, CanonicalGrokPlan(&zero, "", nil)) + require.Equal(t, "free", CanonicalGrokPlan(&zero, "free", nil)) + + require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{ + Model: "grok-4.5", + Requests: &QuotaWindow{Limit: &heavyReq}, + Tokens: &QuotaWindow{Limit: &heavyTok}, + LastHeadersSeenAt: fresh, + })) + require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{ + Model: "grok-4.6", + Requests: &QuotaWindow{Limit: &heavyReq}, + Tokens: &QuotaWindow{Limit: &heavyTok}, + LastHeadersSeenAt: fresh, + })) + require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{ + Requests: &QuotaWindow{Limit: &heavyReq}, + Tokens: &QuotaWindow{Limit: &heavyTok}, + LastHeadersSeenAt: fresh, + })) + require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{ + Model: "grok-4.5", + Requests: &QuotaWindow{Limit: &superReq}, + Tokens: &QuotaWindow{Limit: &superTok}, + LastHeadersSeenAt: fresh, + })) + require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{ + Model: "grok-4.5", + Requests: &QuotaWindow{Limit: &heavyReq}, + LastHeadersSeenAt: stale, + })) + require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{ + Model: "grok-4.6", + Requests: &QuotaWindow{Limit: &superReq}, + PlanFrom45Responses: "supergrok_heavy", + PlanFrom45ResponsesAt: fresh, + })) + require.Equal(t, "free", CanonicalGrokPlan(&zero, "free", &QuotaSnapshot{ + Model: "grok-4.5", + Requests: &QuotaWindow{Limit: &heavyReq}, + LastHeadersSeenAt: fresh, + })) + heavyCents := float64(SuperGrokHeavyLimitCents) + require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&heavyCents, "SuperGrokPro", nil)) +} + +func TestApplyGrok45ResponsesPlanSignalCarriesHint(t *testing.T) { + t.Parallel() + + heavyReq := int64(8300) + fresh := time.Now().UTC().Format(time.RFC3339) + prev := &QuotaSnapshot{ + Model: "grok-4.5", + Requests: &QuotaWindow{Limit: &heavyReq}, + LastHeadersSeenAt: fresh, + } + prev.ApplyGrok45ResponsesPlanSignal(nil) + require.Equal(t, "supergrok_heavy", prev.PlanFrom45Responses) + + next := &QuotaSnapshot{ + Model: "grok-4.6", + Requests: &QuotaWindow{Limit: int64Ptr(100)}, + LastHeadersSeenAt: fresh, + } + next.ApplyGrok45ResponsesPlanSignal(prev) + require.Equal(t, "supergrok_heavy", next.PlanFrom45Responses) + require.Equal(t, prev.PlanFrom45ResponsesAt, next.PlanFrom45ResponsesAt) +} + +func int64Ptr(v int64) *int64 { return &v } + func jwtWithClaims(t *testing.T, claims map[string]any) string { t.Helper() payload, err := json.Marshal(claims) diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index 5a62598db2..9e1d29c2f9 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -969,6 +969,7 @@ func (s *AccountTestService) observeGrokTestResponse(ctx context.Context, accoun resp.Body = io.NopCloser(bytes.NewReader(responseBody)) } snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now) + stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx)) if snapshot != nil && s.accountRepo != nil { resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now) if limited { @@ -1073,7 +1074,7 @@ func (s *AccountTestService) testGrokResponsesConnection(c *gin.Context, ctx con } defer func() { _ = resp.Body.Close() }() - s.observeGrokTestResponse(ctx, account, resp) + s.observeGrokTestResponse(withGrokTeamRateLimitModel(ctx, testModelID), account, resp) if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) @@ -1421,7 +1422,7 @@ User query: return s.sendErrorAndEnd(c, fmt.Sprintf("standalone web_search probe failed: %s", err.Error())) } defer func() { _ = resp.Body.Close() }() - s.observeGrokTestResponse(ctx, account, resp) + s.observeGrokTestResponse(withGrokTeamRateLimitModel(ctx, grokDefaultResponsesModel), account, resp) body, _ := io.ReadAll(resp.Body) if resp.StatusCode != http.StatusOK { diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 953d3db32a..2efb2332f6 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -695,7 +695,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel) } - s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) + s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, requestModel), account, resp.Header, resp.StatusCode) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) if err != nil { return nil, err @@ -857,7 +857,7 @@ func (s *OpenAIGatewayService) forwardGrokMediaVideoContent( return s.handleGrokMediaErrorResponse(ctx, contentResp, c, account, contentRequestID, "") } - s.updateGrokUsageFromResponse(ctx, account, contentResp.Header, contentResp.StatusCode) + s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, ""), account, contentResp.Header, contentResp.StatusCode) if err := writeGrokMediaContentResponse(c, contentResp); err != nil { return nil, err } diff --git a/backend/internal/service/grok_model_quota_block.go b/backend/internal/service/grok_model_quota_block.go index 0bc9cd40b8..e6c7841da4 100644 --- a/backend/internal/service/grok_model_quota_block.go +++ b/backend/internal/service/grok_model_quota_block.go @@ -47,6 +47,32 @@ func markGrokModelQuotaBlock(accountID int64, model string, until time.Time) { if max := now.Add(grokModelQuotaBlockMaxTTL); until.After(max) { until = max } + storeGrokModelQuotaBlock(accountID, model, until, now) +} + +const ( + grokModelTransientBlockMinTTL = 500 * time.Millisecond + grokModelTransientBlockMaxTTL = 5 * time.Minute +) + +// markGrokModelTransientBlock soft-blocks a single model for a short capacity +// burst without the free-usage 20m floor (and without unscheduling the account). +func markGrokModelTransientBlock(accountID int64, model string, until time.Time) { + model = strings.TrimSpace(model) + if accountID <= 0 || model == "" || until.IsZero() { + return + } + now := time.Now() + if !until.After(now.Add(grokModelTransientBlockMinTTL)) { + until = now.Add(grokModelTransientBlockMinTTL) + } + if max := now.Add(grokModelTransientBlockMaxTTL); until.After(max) { + until = max + } + storeGrokModelQuotaBlock(accountID, model, until, now) +} + +func storeGrokModelQuotaBlock(accountID int64, model string, until, now time.Time) { key := grokModelQuotaBlockKey(accountID, model) globalGrokModelQuotaBlocks.mu.Lock() defer globalGrokModelQuotaBlocks.mu.Unlock() diff --git a/backend/internal/service/grok_quota_fetcher.go b/backend/internal/service/grok_quota_fetcher.go index caf1afa47c..f41fa0bd2f 100644 --- a/backend/internal/service/grok_quota_fetcher.go +++ b/backend/internal/service/grok_quota_fetcher.go @@ -71,7 +71,7 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { } if err != nil || snapshot == nil { - applyGrokCredentialUsageFallback(usage, account) + applyGrokCredentialUsageFallback(usage, account, billing, nil) 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" @@ -139,7 +139,7 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { usage.ErrorCode = "spending_limit" } } - applyGrokCredentialUsageFallback(usage, account) + applyGrokCredentialUsageFallback(usage, account, billing, snapshot) if activeProbeClearsForbidden && strings.TrimSpace(snapshot.EntitlementStatus) == "" && strings.EqualFold(strings.TrimSpace(usage.GrokEntitlementStatus), "forbidden") { usage.GrokEntitlementStatus = "" @@ -170,31 +170,47 @@ func firstGrokObservationTime(values ...string) (time.Time, bool) { return time.Time{}, false } -func applyGrokCredentialUsageFallback(usage *UsageInfo, account *Account) { +func applyGrokCredentialUsageFallback(usage *UsageInfo, account *Account, billing *xai.BillingSummary, snapshot *xai.QuotaSnapshot) { if usage == nil || account == nil { return } - if usage.SubscriptionTier == "" { - tier := strings.TrimSpace(account.GetCredential("subscription_tier")) - usage.SubscriptionTier = tier - usage.SubscriptionTierRaw = tier - } - applyGrokLiveJWTSubscriptionTier(usage, account) if usage.GrokEntitlementStatus == "" { usage.GrokEntitlementStatus = strings.TrimSpace(account.GetCredential("entitlement_status")) } + applyGrokResolvedSubscriptionTier(usage, account, billing, snapshot) } -func applyGrokLiveJWTSubscriptionTier(usage *UsageInfo, account *Account) { +func applyGrokResolvedSubscriptionTier(usage *UsageInfo, account *Account, billing *xai.BillingSummary, snapshot *xai.QuotaSnapshot) { if usage == nil || account == nil { return } - tier := xai.SubscriptionTierFromJWT(account.GetCredential("access_token")) - if tier == "" { + if jwtTier := xai.SubscriptionTierFromJWT(account.GetCredential("access_token")); jwtTier != "" { + usage.SubscriptionTier = jwtTier + usage.SubscriptionTierRaw = jwtTier return } - usage.SubscriptionTier = tier - usage.SubscriptionTierRaw = tier + signal := strings.TrimSpace(account.GetCredential("subscription_tier")) + if signal == "" && snapshot != nil { + signal = strings.TrimSpace(snapshot.SubscriptionTier) + } + if signal == "" && billing != nil { + signal = strings.TrimSpace(billing.Plan) + } + var limit *float64 + if billing != nil { + limit = billing.MonthlyLimitCents + } + if plan := xai.CanonicalGrokPlan(limit, signal, snapshot); plan != "" { + usage.SubscriptionTier = plan + if usage.SubscriptionTierRaw == "" { + usage.SubscriptionTierRaw = firstNonEmpty(signal, plan) + } + return + } + if usage.SubscriptionTier == "" && signal != "" { + usage.SubscriptionTier = signal + usage.SubscriptionTierRaw = signal + } } func grokBillingSnapshotFromExtra(extra map[string]any) (*xai.BillingSummary, error) { @@ -233,6 +249,23 @@ func grokBillingSnapshotFromExtra(extra map[string]any) (*xai.BillingSummary, er } } +func stampGrokQuotaSnapshotForPlan(account *Account, snapshot *xai.QuotaSnapshot, model string) { + if snapshot == nil { + return + } + if strings.TrimSpace(snapshot.Model) == "" { + model = strings.TrimSpace(model) + if model != "" { + snapshot.Model = xai.ResolveGrokTextResponsesModelID(model) + } + } + var prev *xai.QuotaSnapshot + if account != nil { + prev, _ = grokQuotaSnapshotFromExtra(account.Extra) + } + snapshot.ApplyGrok45ResponsesPlanSignal(prev) +} + func grokQuotaSnapshotFromExtra(extra map[string]any) (*xai.QuotaSnapshot, error) { if extra == nil { return nil, nil diff --git a/backend/internal/service/grok_quota_fetcher_test.go b/backend/internal/service/grok_quota_fetcher_test.go index 70ea5d09b3..c87a925e9e 100644 --- a/backend/internal/service/grok_quota_fetcher_test.go +++ b/backend/internal/service/grok_quota_fetcher_test.go @@ -24,6 +24,86 @@ func TestGrokQuotaFetcherBuildUsageInfoUnknownUntilFirstSnapshot(t *testing.T) { require.Contains(t, usage.Error, "unknown until billing is probed") } +func TestGrokQuotaFetcherDoesNotTreatGrok45ResponsesWindowAsHeavy(t *testing.T) { + t.Parallel() + + // 8300 / 53M is the grok-4.5 Responses rate-limit window, not a plan fingerprint. + reqLimit, tokLimit := int64(8300), int64(53_000_000) + fresh := time.Now().UTC().Format(time.RFC3339) + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "subscription_tier": "SuperGrokPro", + }, + Extra: map[string]any{ + grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{ + Requests: &xai.QuotaWindow{Limit: &reqLimit}, + Tokens: &xai.QuotaWindow{Limit: &tokLimit}, + LastHeadersSeenAt: fresh, + HeadersObserved: true, + UpdatedAt: fresh, + }, + }, + } + + usage := NewGrokQuotaFetcher().BuildUsageInfo(account) + require.Equal(t, "supergrok", usage.SubscriptionTier) +} + +func TestGrokQuotaFetcherUsesGrok45ResponsesWindowAsHeavy(t *testing.T) { + t.Parallel() + + reqLimit, tokLimit := int64(8300), int64(53_000_000) + fresh := time.Now().UTC().Format(time.RFC3339) + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "subscription_tier": "SuperGrokPro", + }, + Extra: map[string]any{ + grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{ + Model: "grok-4.5", + Requests: &xai.QuotaWindow{Limit: &reqLimit}, + Tokens: &xai.QuotaWindow{Limit: &tokLimit}, + LastHeadersSeenAt: fresh, + HeadersObserved: true, + UpdatedAt: fresh, + }, + }, + } + + usage := NewGrokQuotaFetcher().BuildUsageInfo(account) + require.Equal(t, "supergrok_heavy", usage.SubscriptionTier) +} + +func TestGrokQuotaFetcherJWTBeatsAmbiguousSuperGrokProQuota(t *testing.T) { + t.Parallel() + + heavyReq := int64(8300) + fresh := time.Now().UTC().Format(time.RFC3339) + account := &Account{ + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "access_token": makeGrokOAuthJWT(map[string]any{"tier": 1}), + "subscription_tier": "SuperGrokPro", + }, + Extra: map[string]any{ + grokQuotaSnapshotExtraKey: &xai.QuotaSnapshot{ + Requests: &xai.QuotaWindow{Limit: &heavyReq}, + LastHeadersSeenAt: fresh, + HeadersObserved: true, + UpdatedAt: fresh, + }, + }, + } + + usage := NewGrokQuotaFetcher().BuildUsageInfo(account) + require.Equal(t, "supergrok", usage.SubscriptionTier) +} + func TestGrokQuotaFetcherPrefersLiveJWTTierOverStaleBillingPlan(t *testing.T) { t.Parallel() @@ -70,7 +150,7 @@ func TestGrokQuotaFetcherUsesCredentialTierWhenBillingHasNoPlan(t *testing.T) { usage := NewGrokQuotaFetcher().BuildUsageInfo(account) require.NotNil(t, usage.GrokBilling) - require.Equal(t, "FREE", usage.SubscriptionTier) + require.Equal(t, "free", usage.SubscriptionTier) require.Equal(t, "FREE", usage.SubscriptionTierRaw) require.Equal(t, "active", usage.GrokEntitlementStatus) } diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 6658888a59..293dbcceab 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -181,6 +181,7 @@ func (s *GrokQuotaService) probeUsage(ctx context.Context, accountID int64) (*Gr defer func() { _ = resp.Body.Close() }() snapshot := xai.ObserveQuotaHeaders(resp.Header, resp.StatusCode, "active_probe") + stampGrokQuotaSnapshotForPlan(account, snapshot, probeModel) resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, time.Now()) if limited { normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now()) diff --git a/backend/internal/service/grok_upstream_failure.go b/backend/internal/service/grok_upstream_failure.go index 1348d068c5..e7bed08bf0 100644 --- a/backend/internal/service/grok_upstream_failure.go +++ b/backend/internal/service/grok_upstream_failure.go @@ -489,6 +489,9 @@ func (s *OpenAIGatewayService) applyGrokUpstreamFailureDecision( case GrokFailureEmptyUpstream: reason = "grok empty model output" case GrokFailureModelCapacity: + if persistGrokTransientModelCooldown(account, decision) { + return true + } reason = "grok model capacity" case GrokFailureRateLimit: // Pure 429 without free-usage language keeps the existing rate-limit diff --git a/backend/internal/service/grok_upstream_failure_test.go b/backend/internal/service/grok_upstream_failure_test.go index 6f0a916a78..0bd1063577 100644 --- a/backend/internal/service/grok_upstream_failure_test.go +++ b/backend/internal/service/grok_upstream_failure_test.go @@ -133,6 +133,22 @@ func TestHandleGrokAccountUpstreamError_EmptyOutputCoolsAccount(t *testing.T) { require.WithinDuration(t, before.Add(4*time.Minute), repo.lastTempUnschedUntil, time.Second) } +func TestHandleGrokAccountUpstreamError_MultiAgentCapacityBlocksOnlyThatModel(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9120, Platform: PlatformGrok, Type: AccountTypeOAuth} + ctx := withGrokTeamRateLimitModel(context.Background(), "grok-4.20-multi-agent-0309") + + svc.handleGrokAccountUpstreamError( + ctx, account, http.StatusBadGateway, nil, + []byte(`{"error":{"message":"engine_overloaded"}}`), + ) + + require.Zero(t, repo.tempUnschedCalls) + require.True(t, isGrokModelQuotaBlocked(account.ID, "grok-4.20-multi-agent-0309", time.Now())) + require.False(t, isGrokModelQuotaBlocked(account.ID, "grok-4.5", time.Now())) +} + func TestHandleGrokAccountUpstreamError_FreeUsageDoesNotCoolPoolMode(t *testing.T) { repo := &grokQuotaAccountRepo{} svc := &OpenAIGatewayService{accountRepo: repo} diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 010cd4df2a..34c281b28e 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -191,7 +191,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( Kind: kind, Message: upstreamMsg, }) - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + s.handleGrokAccountUpstreamError(withGrokTeamRateLimitModel(ctx, upstreamModel), account, resp.StatusCode, resp.Header, respBody) if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) { return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, @@ -209,7 +209,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( } if account.Platform == PlatformGrok { - s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) + s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, upstreamModel), account, resp.Header, resp.StatusCode) } // 8. Forward response diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 514f6bef07..d49c8697bd 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -1094,7 +1094,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( Kind: kind, Message: upstreamMsg, }) - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + s.handleGrokAccountUpstreamError(withGrokTeamRateLimitModel(ctx, grokComposerImageBridgeVisionModel), account, resp.StatusCode, resp.Header, respBody) if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) { return "", OpenAIUsage{}, &UpstreamFailoverError{ StatusCode: resp.StatusCode, @@ -1106,7 +1106,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg) } - s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) + s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, grokComposerImageBridgeVisionModel), account, resp.Header, resp.StatusCode) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil) if err != nil { return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err) @@ -1306,6 +1306,10 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco stateCtx, cancel = openAIAccountStateContext(ctx) defer cancel() } + if account.Extra == nil { + account.Extra = map[string]any{} + } + account.Extra[grokQuotaSnapshotExtraKey] = snapshot if s.accountRepo != nil { _ = s.accountRepo.UpdateExtra(stateCtx, accountID, updates) } @@ -1323,6 +1327,7 @@ func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, acco func (s *OpenAIGatewayService) updateGrokUsageFromResponse(ctx context.Context, account *Account, headers http.Header, statusCode int) { snapshot := parseGrokQuotaSnapshot(headers, statusCode, time.Now()) if snapshot != nil { + stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx)) s.updateGrokUsageSnapshot(ctx, account, snapshot) return } @@ -1627,6 +1632,35 @@ func withGrokTeamRateLimitModel(ctx context.Context, model string) context.Conte return context.WithValue(ctx, grokTeamRateLimitModelContextKey{}, model) } +func grokRequestedModelFromCtx(ctx context.Context) string { + if ctx == nil { + return "" + } + model, _ := ctx.Value(grokTeamRateLimitModelContextKey{}).(string) + return strings.TrimSpace(model) +} + +func isGrokHeavyTransientModel(requestedModel string) bool { + model := strings.ToLower(strings.TrimSpace(xai.ResolveGrokTextResponsesModelID(requestedModel))) + return strings.Contains(model, "multi-agent") +} + +func persistGrokTransientModelCooldown(account *Account, decision GrokUpstreamFailureDecision) bool { + if account == nil { + return false + } + model := strings.TrimSpace(decision.Model) + if model == "" || !isGrokHeavyTransientModel(model) { + return false + } + cooldown := decision.Cooldown + if cooldown <= 0 { + cooldown = 3 * time.Minute + } + markGrokModelTransientBlock(account.ID, model, time.Now().Add(cooldown)) + return true +} + func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) { if s == nil || account == nil { return @@ -1635,12 +1669,14 @@ func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Contex return } now := time.Now() - s.updateGrokUsageSnapshot(ctx, account, parseGrokQuotaSnapshot(headers, statusCode, now)) + snapshot := parseGrokQuotaSnapshot(headers, statusCode, now) + stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx)) + s.updateGrokUsageSnapshot(ctx, account, snapshot) // Body-first free-usage / empty / billing / capacity must run before the // status switch so non-429 free-usage bodies still cool the account. // Pool-mode still skips durable mutation unless an explicit temp rule matches. - decision := classifyGrokUpstreamFailure(statusCode, responseBody, "") + decision := classifyGrokUpstreamFailure(statusCode, responseBody, grokRequestedModelFromCtx(ctx)) if decision.ShouldCooldown && decision.Class != GrokFailureNone && decision.Class != GrokFailureRateLimit { if account.IsPoolMode() { // Allow configured temp rules (403) below; skip default body cools. diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge.go b/backend/internal/service/openai_gateway_grok_chat_bridge.go index 88b9b0f41b..e2b8f9a477 100644 --- a/backend/internal/service/openai_gateway_grok_chat_bridge.go +++ b/backend/internal/service/openai_gateway_grok_chat_bridge.go @@ -631,7 +631,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( Kind: kind, Message: upstreamMsg, }) - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + s.handleGrokAccountUpstreamError(withGrokTeamRateLimitModel(ctx, upstreamModel), account, resp.StatusCode, resp.Header, respBody) if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) { return nil, &UpstreamFailoverError{ StatusCode: resp.StatusCode, @@ -643,7 +643,7 @@ func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( return s.handleChatCompletionsErrorResponse(resp, c, account, billingModel) } - s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) + s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, upstreamModel), account, resp.Header, resp.StatusCode) var result *OpenAIForwardResult if clientStream { diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 2aa226dcd4..268aec7b5a 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -443,7 +443,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( return s.handleAnthropicErrorResponse(resp, c, account, billingModel) } if account.Platform == PlatformGrok && account.Type == AccountTypeOAuth && !account.IsShadow() { - s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) + s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, upstreamModel), account, resp.Header, resp.StatusCode) } if account.Type == AccountTypeOAuth && promptCacheKey != "" { diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index ac6a24e78a..87abf1ef04 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -250,7 +250,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( shouldFailover := s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody) if account.Platform == PlatformGrok { shouldFailover = s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + s.handleGrokAccountUpstreamError(withGrokTeamRateLimitModel(ctx, resolveGrokWSUpstreamModel(account, body, originalModel)), account, resp.StatusCode, resp.Header, respBody) if turn == 1 && shouldFailover { return nil, newOpenAIUpstreamFailoverError(resp.StatusCode, resp.Header, respBody, upstreamMsg, false) } @@ -265,7 +265,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg) } if account.Platform == PlatformGrok { - s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) + s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, resolveGrokWSUpstreamModel(account, body, originalModel)), account, resp.Header, resp.StatusCode) } responseID := "" diff --git a/frontend/src/views/admin/AccountsView.vue b/frontend/src/views/admin/AccountsView.vue index 96e07688b4..c33dcad0df 100644 --- a/frontend/src/views/admin/AccountsView.vue +++ b/frontend/src/views/admin/AccountsView.vue @@ -1481,15 +1481,79 @@ const { pause: pauseAutoRefresh, resume: resumeAutoRefresh } = useIntervalFn( { immediate: false } ) -// Persisted credentials.subscription_tier outranks lagging billing/quota snapshots. +const GROK_QUOTA_SIGNAL_MAX_AGE_MS = 24 * 60 * 60 * 1000 +const GROK_QUOTA_SIGNAL_MAX_FUTURE_SKEW_MS = 5 * 60 * 1000 + +function normalizeGrokPlanKey(value: unknown): string { + return String(value ?? '') + .trim() + .toLowerCase() + .replace(/[\s_-]+/g, '') +} + +function isGrokQuotaTimestampFresh(raw: unknown): boolean { + const value = String(raw || '').trim() + if (!value) return false + const observedAt = Date.parse(value) + if (!Number.isFinite(observedAt)) return false + const age = Date.now() - observedAt + return age <= GROK_QUOTA_SIGNAL_MAX_AGE_MS && age >= -GROK_QUOTA_SIGNAL_MAX_FUTURE_SKEW_MS +} + +function isGrok45ResponsesQuotaModel(model: unknown): boolean { + const value = String(model || '') + .trim() + .toLowerCase() + .replace(/^(x-ai|xai)\//, '') + return value === 'grok-4.5' || value.startsWith('grok-4.5-') +} + +function grokQuotaLooksHeavy(snapshot: Record | undefined): boolean { + const req = Number(snapshot?.requests?.limit ?? 0) + const tok = Number(snapshot?.tokens?.limit ?? 0) + return req >= 8300 || tok >= 53_000_000 +} + +function grok45ResponsesPlanIsHeavy(snapshot: Record | undefined): boolean { + if (!snapshot) return false + const hint = normalizeGrokPlanKey(snapshot.plan_from_45_responses) + if (hint === 'supergrokheavy' && isGrokQuotaTimestampFresh(snapshot.plan_from_45_responses_at)) { + return true + } + const observedAt = snapshot.last_headers_seen_at || snapshot.updated_at + return ( + isGrok45ResponsesQuotaModel(snapshot.model) && + isGrokQuotaTimestampFresh(observedAt) && + grokQuotaLooksHeavy(snapshot) + ) +} + +// JWT / unambiguous credentials outrank snapshots. SuperGrokPro is ambiguous +// (covers SuperGrok and Heavy). 8300/53M only upgrades when the window came +// from grok-4.5 Responses (or a carried 4.5 hint). function getAccountPlanType(row: any): string | undefined { if (!row) return undefined if (row.platform === 'grok') { const extra = (row.extra || {}) as Record const billing = extra.grok_billing_snapshot as Record | undefined - const quota = extra.grok_quota_snapshot as Record | undefined + const quota = (extra.grok_quota_snapshot || extra.grok_usage_snapshot) as Record | undefined + const cred = row.credentials?.subscription_tier + const credKey = normalizeGrokPlanKey(cred) + if (credKey && credKey !== 'supergrokpro') { + return cred + } + if ( + grok45ResponsesPlanIsHeavy(quota) && + (credKey === 'supergrokpro' || + normalizeGrokPlanKey(billing?.plan) === 'supergrok' || + normalizeGrokPlanKey(billing?.plan) === 'supergrokpro') + ) { + return 'SuperGrok Heavy' + } + if (credKey === 'supergrokpro') { + return billing?.plan || 'SuperGrok' + } return ( - row.credentials?.subscription_tier || billing?.plan || quota?.subscription_tier || extra.subscription_tier || diff --git a/frontend/src/views/admin/__tests__/AccountsView.sparkShadow.spec.ts b/frontend/src/views/admin/__tests__/AccountsView.sparkShadow.spec.ts index d532301692..54a37ed841 100644 --- a/frontend/src/views/admin/__tests__/AccountsView.sparkShadow.spec.ts +++ b/frontend/src/views/admin/__tests__/AccountsView.sparkShadow.spec.ts @@ -396,6 +396,38 @@ describe('admin AccountsView — 账号行展示', () => { type: 'oauth', credentials: { plan_type: 'SuperGrok' }, }, + { + id: 206, + name: 'supergrokpro-responses-quota', + platform: 'grok', + type: 'oauth', + credentials: { subscription_tier: 'SuperGrokPro' }, + extra: { + grok_billing_snapshot: { plan: 'SuperGrok' }, + grok_quota_snapshot: { + model: 'grok-4.5', + last_headers_seen_at: new Date().toISOString(), + requests: { limit: 8300 }, + tokens: { limit: 53_000_000 }, + }, + }, + }, + { + id: 207, + name: 'supergrokpro-other-model-quota', + platform: 'grok', + type: 'oauth', + credentials: { subscription_tier: 'SuperGrokPro' }, + extra: { + grok_billing_snapshot: { plan: 'SuperGrok' }, + grok_quota_snapshot: { + model: 'grok-4.6', + last_headers_seen_at: new Date().toISOString(), + requests: { limit: 8300 }, + tokens: { limit: 53_000_000 }, + }, + }, + }, ] listAccounts.mockResolvedValue({ @@ -416,6 +448,8 @@ describe('admin AccountsView — 账号行展示', () => { 'FREE', 'BASIC', 'SuperGrok', + 'SuperGrok Heavy', + 'SuperGrok', ]) wrapper.unmount()