fix: SuperGrokPro 用 4.5 窗口区分 Heavy,容量抖动只封单模型

JWT 的 SuperGrokPro 同时覆盖 SuperGrok 与 Heavy,账单月额度又滞后,
账号会误升/误降。multi-agent 容量抖动还会把整号提出调度。

- CanonicalGrokPlan:明确 JWT 优先;模糊档仅采新鲜的 grok-4.5 Responses 窗口
- 配额快照写入 plan_from_45_responses 时间戳,过期信号不用
- engine_overloaded 对 multi-agent 只封当前模型 0.5s–5min
This commit is contained in:
IanShaw027
2026-08-13 07:49:14 +08:00
parent f3d9491071
commit 363cc4994b
18 changed files with 601 additions and 35 deletions
+6
View File
@@ -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 {
+183 -1
View File
@@ -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 ""
}
@@ -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)
@@ -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 {
+2 -2
View File
@@ -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
}
@@ -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()
+47 -14
View File
@@ -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
@@ -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)
}
@@ -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())
@@ -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
@@ -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}
@@ -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
@@ -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.
@@ -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 {
@@ -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 != "" {
@@ -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 := ""
+67 -3
View File
@@ -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<string, any> | 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<string, any> | 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<string, any>
const billing = extra.grok_billing_snapshot as Record<string, any> | undefined
const quota = extra.grok_quota_snapshot as Record<string, any> | undefined
const quota = (extra.grok_quota_snapshot || extra.grok_usage_snapshot) as Record<string, any> | 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 ||
@@ -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()