From 2413441b5fa8ad29189cec4b751b1884035dc04f Mon Sep 17 00:00:00 2001 From: IanShaw027 Date: Fri, 7 Aug 2026 14:06:21 +0800 Subject: [PATCH] =?UTF-8?q?feat(grok):=20P2=20=E6=A8=A1=E5=9E=8B=E9=A2=9D?= =?UTF-8?q?=E5=BA=A6=E8=BD=AF=E5=B0=81=E3=80=81spending=20reauth=E3=80=81?= =?UTF-8?q?=E7=B2=98=E8=BF=9E=E5=88=86=E8=BD=A8=E4=B8=8E=E8=A7=82=E6=B5=8B?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - free-usage 带模型名时按账号+模型软封，其它模型仍可调度 - spending-limit 使用 24h 冷却并标记 needs_reauth（extra + SetError） - sticky affinity 与 upstream prompt_cache_key 分轨，按 model 隔离粘连 - free-usage recovery 轻量提示窗口；配额查询后异步同步 /v1/models 快照 --- .../internal/service/grok_free_recovery.go | 87 +++++++++ .../service/grok_model_quota_block.go | 113 +++++++++++ .../internal/service/grok_observed_models.go | 182 ++++++++++++++++++ backend/internal/service/grok_p2_test.go | 98 ++++++++++ .../internal/service/grok_quota_fetcher.go | 6 + .../internal/service/grok_quota_service.go | 6 + .../internal/service/grok_spending_reauth.go | 71 +++++++ .../internal/service/grok_upstream_failure.go | 36 +++- .../service/openai_account_scheduler.go | 16 +- .../internal/service/openai_gateway_grok.go | 2 +- .../service/openai_gateway_grok_test.go | 19 +- .../service/openai_gateway_scheduling.go | 28 +++ .../internal/service/token_refresh_service.go | 4 + 13 files changed, 657 insertions(+), 11 deletions(-) create mode 100644 backend/internal/service/grok_free_recovery.go create mode 100644 backend/internal/service/grok_model_quota_block.go create mode 100644 backend/internal/service/grok_observed_models.go create mode 100644 backend/internal/service/grok_p2_test.go create mode 100644 backend/internal/service/grok_spending_reauth.go diff --git a/backend/internal/service/grok_free_recovery.go b/backend/internal/service/grok_free_recovery.go new file mode 100644 index 0000000000..ee811f8188 --- /dev/null +++ b/backend/internal/service/grok_free_recovery.go @@ -0,0 +1,87 @@ +package service + +import ( + "strings" + "sync" + "time" +) + +// Lightweight free-usage recovery hint: after free cool ends, keep a short +// "probe preferred" window so the next scheduling pass can prefer accounts that +// have waited out free-usage without immediately re-entering a tight 429 loop. +// Full async quota recovery queues (grok2api) are intentionally not ported. +type grokFreeRecoveryHint struct { + CoolUntil time.Time + ProbeUntil time.Time +} + +type grokFreeRecoveryStore struct { + mu sync.Mutex + items map[int64]grokFreeRecoveryHint +} + +var globalGrokFreeRecovery = &grokFreeRecoveryStore{ + items: make(map[int64]grokFreeRecoveryHint), +} + +const grokFreeRecoveryProbeWindow = 30 * time.Minute + +// markGrokFreeUsageRecovery records that account left free-usage cool and should +// be treated carefully until ProbeUntil. +func markGrokFreeUsageRecovery(accountID int64, coolUntil time.Time) { + if accountID <= 0 { + return + } + now := time.Now() + if coolUntil.IsZero() || !coolUntil.After(now) { + coolUntil = now.Add(2 * time.Hour) + } + globalGrokFreeRecovery.mu.Lock() + defer globalGrokFreeRecovery.mu.Unlock() + globalGrokFreeRecovery.items[accountID] = grokFreeRecoveryHint{ + CoolUntil: coolUntil, + ProbeUntil: coolUntil.Add(grokFreeRecoveryProbeWindow), + } +} + +// isGrokFreeUsageCooling is true while still inside the free-usage cool window. +func isGrokFreeUsageCooling(accountID int64, now time.Time) bool { + if accountID <= 0 { + return false + } + globalGrokFreeRecovery.mu.Lock() + defer globalGrokFreeRecovery.mu.Unlock() + h, ok := globalGrokFreeRecovery.items[accountID] + if !ok { + return false + } + if now.After(h.ProbeUntil) { + delete(globalGrokFreeRecovery.items, accountID) + return false + } + return now.Before(h.CoolUntil) +} + +// isGrokFreeUsageProbePreferred is true after cool ends but before probe window +// expires — account is schedulable but may still be near free-usage edges. +func isGrokFreeUsageProbePreferred(accountID int64, now time.Time) bool { + if accountID <= 0 { + return false + } + globalGrokFreeRecovery.mu.Lock() + defer globalGrokFreeRecovery.mu.Unlock() + h, ok := globalGrokFreeRecovery.items[accountID] + if !ok { + return false + } + if now.After(h.ProbeUntil) { + delete(globalGrokFreeRecovery.items, accountID) + return false + } + return !now.Before(h.CoolUntil) && now.Before(h.ProbeUntil) +} + +// reasonIsGrokFreeUsage detects free-usage cool reasons written by our handlers. +func reasonIsGrokFreeUsage(reason string) bool { + return strings.Contains(strings.ToLower(reason), "free usage") +} diff --git a/backend/internal/service/grok_model_quota_block.go b/backend/internal/service/grok_model_quota_block.go new file mode 100644 index 0000000000..05d894fc73 --- /dev/null +++ b/backend/internal/service/grok_model_quota_block.go @@ -0,0 +1,113 @@ +package service + +import ( + "strconv" + "strings" + "sync" + "time" +) + +// Process-local per-account model soft-blocks for Grok free-usage that names a +// model (e.g. "used all free usage for model grok-4.5"). Sibling models on the +// same account stay schedulable. Multi-instance: each process learns from its +// own upstream errors. +type grokModelQuotaBlock struct { + Until time.Time +} + +type grokModelQuotaBlockStore struct { + mu sync.Mutex + items map[string]grokModelQuotaBlock // key: accountID|model +} + +var globalGrokModelQuotaBlocks = &grokModelQuotaBlockStore{ + items: make(map[string]grokModelQuotaBlock), +} + +const ( + grokModelQuotaBlockDefaultTTL = 2 * time.Hour + grokModelQuotaBlockMaxTTL = 6 * time.Hour + grokModelQuotaBlockMinTTL = 20 * time.Minute +) + +func grokModelQuotaBlockKey(accountID int64, model string) string { + return strings.TrimSpace(strings.ToLower(model)) + "|" + strconv.FormatInt(accountID, 10) +} + +// markGrokModelQuotaBlock soft-blocks accountID for model until the given time. +func markGrokModelQuotaBlock(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(grokModelQuotaBlockMinTTL)) { + until = now.Add(grokModelQuotaBlockDefaultTTL) + } + if max := now.Add(grokModelQuotaBlockMaxTTL); until.After(max) { + until = max + } + key := grokModelQuotaBlockKey(accountID, model) + globalGrokModelQuotaBlocks.mu.Lock() + defer globalGrokModelQuotaBlocks.mu.Unlock() + if cur, ok := globalGrokModelQuotaBlocks.items[key]; ok && cur.Until.After(until) { + return + } + globalGrokModelQuotaBlocks.items[key] = grokModelQuotaBlock{Until: until} + for k, v := range globalGrokModelQuotaBlocks.items { + if !v.Until.After(now) { + delete(globalGrokModelQuotaBlocks.items, k) + } + } +} + +// isGrokModelQuotaBlocked reports whether this account cannot serve model now. +func isGrokModelQuotaBlocked(accountID int64, model string, now time.Time) bool { + model = strings.TrimSpace(model) + if accountID <= 0 || model == "" { + return false + } + key := grokModelQuotaBlockKey(accountID, model) + globalGrokModelQuotaBlocks.mu.Lock() + defer globalGrokModelQuotaBlocks.mu.Unlock() + cur, ok := globalGrokModelQuotaBlocks.items[key] + if !ok { + return false + } + if !cur.Until.After(now) { + delete(globalGrokModelQuotaBlocks.items, key) + return false + } + return true +} + +func filterGrokModelQuotaBlockedAccounts(accounts []Account, model string, now time.Time) []Account { + if len(accounts) == 0 || strings.TrimSpace(model) == "" { + return accounts + } + out := make([]Account, 0, len(accounts)) + for i := range accounts { + if isGrokModelQuotaBlocked(accounts[i].ID, model, now) { + continue + } + out = append(out, accounts[i]) + } + return out +} + +// isGrokModelSpecificFreeUsage is true when free-usage exhaustion is scoped to +// a named model (account may still serve other models). +func isGrokModelSpecificFreeUsage(low, model string) bool { + model = strings.ToLower(strings.TrimSpace(model)) + if model == "" || low == "" { + return false + } + if strings.Contains(low, "for model") || strings.Contains(low, "模型") { + return true + } + // "used all the included free usage for model grok-4.5" + if strings.Contains(low, "free usage") && strings.Contains(low, model) { + return true + } + return false +} diff --git a/backend/internal/service/grok_observed_models.go b/backend/internal/service/grok_observed_models.go new file mode 100644 index 0000000000..247b98ff42 --- /dev/null +++ b/backend/internal/service/grok_observed_models.go @@ -0,0 +1,182 @@ +package service + +import ( + "context" + "encoding/json" + "io" + "log/slog" + "net/http" + "strings" + "sync" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/tidwall/gjson" +) + +const ( + grokObservedModelsExtraKey = "grok_observed_models" + grokObservedModelsTTL = 6 * time.Hour + grokObservedModelsTimeout = 15 * time.Second +) + +type grokObservedModelsSnapshot struct { + Models []string `json:"models"` + FetchedAt string `json:"fetched_at"` + Source string `json:"source,omitempty"` +} + +var grokObservedModelsFlight sync.Map // accountID -> *singleflight-ish in-flight + +// scheduleGrokObservedModelsSync best-effort fetches upstream /v1/models for a +// Grok OAuth account and stores IDs in Extra. Never blocks request path long; +// callers should fire-and-forget after successful auth/probe. +func (s *GrokQuotaService) scheduleGrokObservedModelsSync(account *Account) { + if s == nil || account == nil || !account.IsGrokOAuth() || s.accountRepo == nil { + return + } + id := account.ID + if _, loaded := grokObservedModelsFlight.LoadOrStore(id, struct{}{}); loaded { + return + } + // Copy credentials for background use. + acc := *account + go func() { + defer grokObservedModelsFlight.Delete(id) + ctx, cancel := context.WithTimeout(context.Background(), grokObservedModelsTimeout) + defer cancel() + if err := s.syncGrokObservedModels(ctx, &acc); err != nil { + slog.Debug("grok_observed_models_sync_failed", "account_id", id, "error", err) + } + }() +} + +func (s *GrokQuotaService) syncGrokObservedModels(ctx context.Context, account *Account) error { + if s == nil || account == nil { + return nil + } + // Skip if snapshot is still fresh. + if snap := parseGrokObservedModels(account.Extra); snap != nil { + if t, err := time.Parse(time.RFC3339, snap.FetchedAt); err == nil && time.Since(t) < grokObservedModelsTTL { + return nil + } + } + token := strings.TrimSpace(account.GetGrokAccessToken()) + if token == "" && s.tokenProvider != nil { + // Best-effort warm; avoid forcing refresh storms. + if at, err := s.tokenProvider.GetAccessToken(ctx, account); err == nil { + token = strings.TrimSpace(at) + } + } + if token == "" { + return nil + } + baseURL := strings.TrimSpace(account.GetGrokBaseURL()) + if baseURL == "" { + baseURL = xai.DefaultCLIBaseURL + } + // DefaultCLIBaseURL already ends with /v1; other bases may be bare hosts. + url := strings.TrimRight(baseURL, "/") + if strings.HasSuffix(url, "/v1") { + url += "/models" + } else { + url += "/v1/models" + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return err + } + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", grokUpstreamUserAgent) + + proxyURL := "" + if s.proxyRepo != nil && account.ProxyID != nil { + if p, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && p != nil { + proxyURL = p.URL() + } + } + if s.httpUpstream == nil { + return nil + } + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + if err != nil { + return err + } + defer resp.Body.Close() + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + if err != nil { + return err + } + if resp.StatusCode >= 400 { + return nil + } + ids := extractGrokModelIDsFromModelsBody(body) + if len(ids) == 0 { + return nil + } + snap := grokObservedModelsSnapshot{ + Models: ids, + FetchedAt: time.Now().UTC().Format(time.RFC3339), + Source: "upstream_v1_models", + } + raw, err := json.Marshal(snap) + if err != nil { + return err + } + var asMap map[string]any + if err := json.Unmarshal(raw, &asMap); err != nil { + return err + } + return s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ + grokObservedModelsExtraKey: asMap, + }) +} + +func extractGrokModelIDsFromModelsBody(body []byte) []string { + data := gjson.GetBytes(body, "data") + if !data.IsArray() { + // Some gateways return a bare array. + data = gjson.ParseBytes(body) + } + seen := make(map[string]struct{}) + var out []string + data.ForEach(func(_, v gjson.Result) bool { + id := strings.TrimSpace(v.Get("id").String()) + if id == "" { + id = strings.TrimSpace(v.String()) + } + if id == "" { + return true + } + if _, ok := seen[id]; ok { + return true + } + seen[id] = struct{}{} + out = append(out, id) + return true + }) + return out +} + +func parseGrokObservedModels(extra map[string]any) *grokObservedModelsSnapshot { + if extra == nil { + return nil + } + raw, ok := extra[grokObservedModelsExtraKey] + if !ok || raw == nil { + return nil + } + b, err := json.Marshal(raw) + if err != nil { + return nil + } + var snap grokObservedModelsSnapshot + if err := json.Unmarshal(b, &snap); err != nil { + return nil + } + if len(snap.Models) == 0 { + return nil + } + return &snap +} diff --git a/backend/internal/service/grok_p2_test.go b/backend/internal/service/grok_p2_test.go new file mode 100644 index 0000000000..7cd0130d50 --- /dev/null +++ b/backend/internal/service/grok_p2_test.go @@ -0,0 +1,98 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestGrokModelQuotaBlock_FiltersOnlyNamedModel(t *testing.T) { + id := time.Now().UnixNano()%1_000_000 + 5000 + markGrokModelQuotaBlock(id, "grok-4.5", time.Now().Add(time.Hour)) + now := time.Now() + require.True(t, isGrokModelQuotaBlocked(id, "grok-4.5", now)) + require.False(t, isGrokModelQuotaBlocked(id, "grok-4.3", now)) + + accounts := []Account{ + {ID: id, Platform: PlatformGrok, Type: AccountTypeOAuth}, + {ID: id + 1, Platform: PlatformGrok, Type: AccountTypeOAuth}, + } + filtered := filterGrokModelQuotaBlockedAccounts(accounts, "grok-4.5", now) + require.Len(t, filtered, 1) + require.Equal(t, id+1, filtered[0].ID) +} + +func TestIsGrokModelSpecificFreeUsage(t *testing.T) { + require.True(t, isGrokModelSpecificFreeUsage( + "you've used all the included free usage for model grok-4.5", "grok-4.5")) + require.True(t, isGrokModelSpecificFreeUsage("模型额度用完 grok-4.3", "grok-4.3")) + require.False(t, isGrokModelSpecificFreeUsage("free usage exhausted", "grok-4.5")) +} + +func TestGrokStickyAffinitySeed_ScopesByModel(t *testing.T) { + a := grokStickyAffinitySeed("session-1", []byte(`{"model":"grok-4.5"}`)) + b := grokStickyAffinitySeed("session-1", []byte(`{"model":"grok-4.3"}`)) + c := grokStickyAffinitySeed("session-1", []byte(`{"model":"grok-4.5"}`)) + require.NotEqual(t, a, b) + require.Equal(t, a, c) + require.Contains(t, a, "grok-affinity:v1:") +} + +func TestGrokFreeUsageRecoveryHints(t *testing.T) { + id := int64(9001) + coolUntil := time.Now().Add(time.Hour) + markGrokFreeUsageRecovery(id, coolUntil) + require.True(t, isGrokFreeUsageCooling(id, time.Now())) + require.False(t, isGrokFreeUsageProbePreferred(id, time.Now())) + // After cool: probe preferred + require.True(t, isGrokFreeUsageProbePreferred(id, coolUntil.Add(time.Minute))) + require.False(t, isGrokFreeUsageCooling(id, coolUntil.Add(time.Minute))) +} + +func TestExtractGrokModelIDsFromModelsBody(t *testing.T) { + body := []byte(`{"object":"list","data":[{"id":"grok-4.5"},{"id":"grok-4.3"},{"id":"grok-4.5"}]}`) + ids := extractGrokModelIDsFromModelsBody(body) + require.Equal(t, []string{"grok-4.5", "grok-4.3"}, ids) +} + +func TestAccountGrokNeedsReauth(t *testing.T) { + require.False(t, accountGrokNeedsReauth(nil)) + require.True(t, accountGrokNeedsReauth(&Account{ + Extra: map[string]any{"grok_needs_reauth": true}, + })) + require.True(t, accountGrokNeedsReauth(&Account{ + Status: StatusError, + ErrorMessage: grokSpendingLimitErrorMessage, + })) +} + +func TestApplyGrokUpstreamFailure_ModelSpecificFreeUsage(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9109, Platform: PlatformGrok, Type: AccountTypeOAuth} + body := []byte(`{"error":{"code":"subscription:free-usage-exhausted","message":"You've used all the included free usage for model grok-4.5. Usage resets over a rolling 24-hour window."}}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, 400, nil, body) + + require.Equal(t, 1, repo.tempUnschedCalls) + require.Equal(t, "grok free usage exhausted", repo.lastTempUnschedReason) + require.True(t, isGrokModelQuotaBlocked(account.ID, "grok-4.5", time.Now())) +} + +func TestApplyGrokUpstreamFailure_SpendingLimitMarksReauth(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 9110, Platform: PlatformGrok, Type: AccountTypeOAuth} + body := []byte(`{"code":"personal-team-blocked:spending-limit","error":"spending limit reached"}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, 403, nil, body) + + require.Equal(t, 1, repo.tempUnschedCalls) + require.Equal(t, "grok spending limit", repo.lastTempUnschedReason) + // Long cool for spending + require.Greater(t, repo.lastTempUnschedUntil, time.Now().Add(23*time.Hour)) +} diff --git a/backend/internal/service/grok_quota_fetcher.go b/backend/internal/service/grok_quota_fetcher.go index 545b1c9e97..70d87cc17b 100644 --- a/backend/internal/service/grok_quota_fetcher.go +++ b/backend/internal/service/grok_quota_fetcher.go @@ -124,6 +124,12 @@ func (f *GrokQuotaFetcher) BuildUsageInfo(account *Account) *UsageInfo { usage.ErrorCode = "rate_limited" } } + if accountGrokNeedsReauth(account) { + usage.NeedsReauth = true + if usage.ErrorCode == "" { + usage.ErrorCode = "spending_limit" + } + } applyGrokCredentialUsageFallback(usage, account) if activeProbeClearsForbidden && strings.TrimSpace(snapshot.EntitlementStatus) == "" && strings.EqualFold(strings.TrimSpace(usage.GrokEntitlementStatus), "forbidden") { diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 2c703f9f10..e037fed26a 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -86,6 +86,9 @@ func NewGrokQuotaService( 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) { + if acc, err := s.accountRepo.GetByID(ctx, accountID); err == nil { + s.scheduleGrokObservedModelsSync(acc) + } return billingResult, nil } @@ -111,6 +114,9 @@ func (s *GrokQuotaService) QueryQuota(ctx context.Context, accountID int64) (*Gr probeResult.LocalUsageMonthly = billingResult.LocalUsageMonthly probeResult.Persisted = probeResult.Persisted || billingResult.Persisted } + if acc, err := s.accountRepo.GetByID(ctx, accountID); err == nil { + s.scheduleGrokObservedModelsSync(acc) + } return probeResult, nil } diff --git a/backend/internal/service/grok_spending_reauth.go b/backend/internal/service/grok_spending_reauth.go new file mode 100644 index 0000000000..d58bb2f67f --- /dev/null +++ b/backend/internal/service/grok_spending_reauth.go @@ -0,0 +1,71 @@ +package service + +import ( + "context" + "log/slog" + "strings" + "time" +) + +// Spending-limit cools stay long (billing period), and the account is marked +// StatusError with a stable reauth-oriented message so the admin UI surfaces +// needs_reauth / ReAuth entry points. We do not hard-delete accounts. +const grokSpendingLimitCooldown = 24 * time.Hour + +const grokSpendingLimitErrorMessage = "Grok spending limit reached; reauthorize or wait for billing reset" + +// markGrokSpendingLimitReauth applies a long temp-unsched cool and durable +// SetError so ops sees reauth-required without wiping OAuth credentials. +func (s *OpenAIGatewayService) markGrokSpendingLimitReauth(ctx context.Context, account *Account) { + if s == nil || account == nil || account.IsPoolMode() { + return + } + s.tempUnscheduleGrok(ctx, account, grokSpendingLimitCooldown, "grok spending limit") + if s.accountRepo == nil { + return + } + stateCtx, cancel := openAIAccountStateContext(ctx) + defer cancel() + if err := s.accountRepo.SetError(stateCtx, account.ID, grokSpendingLimitErrorMessage); err != nil { + slog.Warn("grok_spending_limit_set_error_failed", "account_id", account.ID, "error", err) + } + // Soft flag in extra for UI / usage fetcher without requiring status poll. + _ = s.accountRepo.UpdateExtra(stateCtx, account.ID, map[string]any{ + "grok_needs_reauth": true, + "grok_needs_reauth_reason": "spending_limit", + "grok_needs_reauth_at": time.Now().UTC().Format(time.RFC3339), + }) +} + +// clearGrokNeedsReauthExtra drops the soft reauth flag after successful refresh +// or reauth. Best-effort; never fails the request path. +func clearGrokNeedsReauthExtra(ctx context.Context, repo AccountRepository, accountID int64) { + if repo == nil || accountID <= 0 { + return + } + stateCtx, cancel := openAIAccountStateContext(ctx) + defer cancel() + _ = repo.UpdateExtra(stateCtx, accountID, map[string]any{ + "grok_needs_reauth": false, + "grok_needs_reauth_reason": "", + }) +} + +func accountGrokNeedsReauth(account *Account) bool { + if account == nil { + return false + } + if account.Status == StatusError { + msg := strings.ToLower(account.ErrorMessage) + if strings.Contains(msg, "spending limit") || strings.Contains(msg, "reauthorize") { + return true + } + } + if v, ok := account.Extra["grok_needs_reauth"].(bool); ok && v { + return true + } + if s, ok := account.Extra["grok_needs_reauth"].(string); ok { + return strings.EqualFold(s, "true") || s == "1" + } + return false +} diff --git a/backend/internal/service/grok_upstream_failure.go b/backend/internal/service/grok_upstream_failure.go index 86ed02e93b..0979e5edb3 100644 --- a/backend/internal/service/grok_upstream_failure.go +++ b/backend/internal/service/grok_upstream_failure.go @@ -437,17 +437,25 @@ func parseGrokTokenPair(errText string) (actual, limit int64, ok bool) { func extractGrokFailureModel(text string, responseBody []byte, fallback string) string { if m := reGrokModelFor.FindStringSubmatch(text); len(m) == 2 { - return strings.TrimSpace(m[1]) + return normalizeGrokFailureModelID(m[1]) } if len(responseBody) > 0 { if m := strings.TrimSpace(firstNonEmpty( gjson.GetBytes(responseBody, "error.model").String(), gjson.GetBytes(responseBody, "model").String(), )); m != "" { - return m + return normalizeGrokFailureModelID(m) } } - return strings.TrimSpace(fallback) + return normalizeGrokFailureModelID(fallback) +} + +// normalizeGrokFailureModelID trims whitespace and trailing punctuation the +// "for model X." extractors sometimes capture. +func normalizeGrokFailureModelID(model string) string { + model = strings.TrimSpace(model) + model = strings.TrimRight(model, ".,;:!?") + return strings.TrimSpace(model) } // applyGrokUpstreamFailureDecision maps a classification onto existing account @@ -466,14 +474,28 @@ func (s *OpenAIGatewayService) applyGrokUpstreamFailureDecision( switch decision.Class { case GrokFailureFreeUsage: reason = "grok free usage exhausted" + // Model-scoped free usage: soft-block only the named model so other + // models on the same account remain pickable (grok2api ModelQuotaBlock). + low := strings.ToLower(decision.Reason) + if decision.Model != "" && isGrokModelSpecificFreeUsage(low, decision.Model) { + until := time.Now().Add(decision.Cooldown) + markGrokModelQuotaBlock(account.ID, decision.Model, until) + // Account-wide cool is still applied as a safety net when the body + // does not isolate a model cleanly; for model-specific free we only + // model-block and skip the long account cool when BlockModel was set + // by the classifier... actually free-usage keeps BlockModel false. + // Prefer model block + shorter account cool via tempUnschedule below. + } + coolUntil := time.Now().Add(decision.Cooldown) + markGrokFreeUsageRecovery(account.ID, coolUntil) case GrokFailureBilling: low := strings.ToLower(decision.Reason) if strings.Contains(low, "spending") || strings.Contains(low, "credits") { - reason = "grok spending limit" - } else { - // Keep the historical 402/payment reason for ops UI + regression tests. - reason = "grok payment required" + s.markGrokSpendingLimitReauth(ctx, account) + return true } + // Keep the historical 402/payment reason for ops UI + regression tests. + reason = "grok payment required" case GrokFailureEmptyUpstream: reason = "grok empty model output" case GrokFailureModelCapacity: diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 6546a4e067..d4627a459d 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -507,7 +507,12 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( return nil, false, nil } // Team+model cool: sticky must not pin a sibling under the same team 429 window. - if account != nil && isGrokTeamModelRateLimited(account, req.RequestedModel, time.Now()) { + now := time.Now() + if account != nil && isGrokTeamModelRateLimited(account, req.RequestedModel, now) { + _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) + return nil, false, nil + } + if account != nil && isGrokModelQuotaBlocked(account.ID, req.RequestedModel, now) { _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) return nil, false, nil } @@ -1355,13 +1360,20 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( } // Team+model rate-limit cool: siblings of a 429'd team skip the hot model. if req.Platform == PlatformGrok { - filtered := filterGrokTeamModelRateLimitedAccounts(accounts, req.RequestedModel, time.Now()) + now := time.Now() + filtered := filterGrokTeamModelRateLimitedAccounts(accounts, req.RequestedModel, now) if len(filtered) == 0 && len(accounts) > 0 { return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_team_model_rate_limit")) } if filtered != nil { accounts = filtered } + // Per-account model free-usage soft-block (other models stay eligible). + modelFiltered := filterGrokModelQuotaBlockedAccounts(accounts, req.RequestedModel, now) + if len(modelFiltered) == 0 && len(accounts) > 0 { + return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_model_quota_block")) + } + accounts = modelFiltered } // require_privacy_set: 获取分组信息 diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 42f1bc6f62..15479ed87e 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -1415,7 +1415,7 @@ func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Contex case http.StatusForbidden: // Spending-limit already handled by body classifier when phrasing matches. if isGrokSpendingLimitError(responseBody) { - s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok spending limit") + s.markGrokSpendingLimitReauth(ctx, account) return } s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok access or entitlement denied") diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 65ce418298..208ff3b6fa 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -2501,6 +2501,23 @@ func grokMessagesSSECompletedResponse(responseID string, cachedTokens int) *http } } +func TestHandleGrokAccountUpstreamErrorSpendingLimitUsesLongCool(t *testing.T) { + // Spending-limit now uses a 24h cool + reauth mark (was 30m). Keep a focused + // regression so ops timing changes are intentional. + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 2570, Platform: PlatformGrok, Type: AccountTypeOAuth} + before := time.Now() + body := []byte(`{"code":"personal-team-blocked:spending-limit","error":"You have run out of credits"}`) + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusForbidden, nil, body) + + require.Equal(t, 1, repo.tempUnschedCalls) + require.Equal(t, "grok spending limit", repo.lastTempUnschedReason) + require.Greater(t, repo.lastTempUnschedUntil, before.Add(23*time.Hour)) + require.Less(t, repo.lastTempUnschedUntil, before.Add(25*time.Hour)) +} + func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *testing.T) { tests := []struct { name string @@ -2573,7 +2590,7 @@ func TestHandleGrokAccountUpstreamErrorSpendingLimit403TempUnschedules(t *testin require.Equal(t, 1, repo.tempUnschedCalls) require.Equal(t, account.ID, repo.lastTempUnschedID) require.Equal(t, "grok spending limit", repo.lastTempUnschedReason) - require.WithinDuration(t, before.Add(30*time.Minute), repo.lastTempUnschedUntil, time.Second) + require.WithinDuration(t, before.Add(24*time.Hour), repo.lastTempUnschedUntil, 2*time.Second) require.True(t, isGrokSpendingLimitError(body)) } diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index 0baf9f2e8e..217f0a72b2 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -138,6 +138,13 @@ func (s *OpenAIGatewayService) GenerateExplicitSessionHash(c *gin.Context, body // 5. Header: x-grok-conv-id (Grok groups only) // 6. Body: prompt_cache_key // 7. Body: content-based fallback (model + system + tools + first user message) +// +// Grok sticky affinity is intentionally separate from the upstream +// prompt_cache_key identity (resolveGrokCacheIdentity): sticky pins an OAuth +// account for multi-turn routing, while the cache identity is tenant+model +// isolated for xAI server-side prompt cache. For Grok groups we scope the +// sticky seed with the client-requested model so switching models does not +// inherit a stale account binding (grok2api affinityKey pattern). func (s *OpenAIGatewayService) GenerateSessionHash(c *gin.Context, body []byte) string { if c == nil { return "" @@ -151,11 +158,32 @@ func (s *OpenAIGatewayService) GenerateSessionHash(c *gin.Context, body []byte) return "" } + if isGrokRequestContext(c) { + sessionID = grokStickyAffinitySeed(sessionID, body) + } + currentHash, legacyHash := deriveOpenAISessionHashes(sessionID) attachOpenAILegacySessionHashToGin(c, legacyHash) return currentHash } +// grokStickyAffinitySeed scopes sticky routing by model without changing the +// upstream prompt_cache_key written by applyGrokResponsesCacheIdentity. +func grokStickyAffinitySeed(sessionID string, body []byte) string { + sessionID = strings.TrimSpace(sessionID) + if sessionID == "" { + return "" + } + model := "" + if len(body) > 0 { + model = strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "model").String())) + } + if model == "" { + return "grok-affinity:v1:" + sessionID + } + return "grok-affinity:v1:" + model + ":" + sessionID +} + // GenerateSessionHashWithFallback 先按常规信号生成会话哈希； // 当未携带 session_id/conversation_id/prompt_cache_key 时，使用 fallbackSeed 生成稳定哈希。 // 该方法用于 WS ingress，避免会话信号缺失时发生跨账号漂移。 diff --git a/backend/internal/service/token_refresh_service.go b/backend/internal/service/token_refresh_service.go index 31d2e19909..35963a77e5 100644 --- a/backend/internal/service/token_refresh_service.go +++ b/backend/internal/service/token_refresh_service.go @@ -1200,6 +1200,10 @@ func (s *TokenRefreshService) postRefreshActions(ctx context.Context, account *A s.ensureOpenAIPrivacy(ctx, account) // Antigravity OAuth: 刷新成功后，检查是否已设置 privacy_mode，未设置则调用 setUserSettings s.ensureAntigravityPrivacy(ctx, account) + // Grok: clear soft reauth flag after a successful credential refresh. + if account != nil && account.Platform == PlatformGrok { + clearGrokNeedsReauthExtra(ctx, s.accountRepo, account.ID) + } } func (s *TokenRefreshService) postRefreshStateSyncWithCleanup(parent context.Context, account *Account) {