diff --git a/backend/docs/GROK_INTEGRATION_PROGRESS.md b/backend/docs/GROK_INTEGRATION_PROGRESS.md index aa73c87a49..9fdb76b85e 100644 --- a/backend/docs/GROK_INTEGRATION_PROGRESS.md +++ b/backend/docs/GROK_INTEGRATION_PROGRESS.md @@ -7,10 +7,10 @@ 1. **模型目录与可配置映射** — 默认禁止 gpt/claude→grok-4.5;设置项 `grok_default_text_model` / `grok_cross_client_model_map_enabled` 2. **密码登录 + SSO 校验** — `POST .../oauth/password`、`.../oauth/sso-token`;不落库密码/raw SSO 3. **视频按模型族定价** — `groups.video_model_prices` JSONB;计费顺序:模型×分辨率 → 旧三列 → 官方默认 +4. **free 档本地用量软门禁 + 支付失败临时下线** — `gateway.grok.free_quota_*`;仅明确 free 的 OAuth 走调度过滤器;402 / spending-limit 403 tempUnschedule;管理端探测不走门禁 ## 待续阶段 -- free-tier / cooldown / payment-required 调度 - media/voice 增量与错误语义 - 网关 tool_choice / active-delta(默认关)/ web_search - 前端 CreateAccount SSO/密码入口 diff --git a/backend/internal/config/config.go b/backend/internal/config/config.go index ab83333eb4..2d9a95e256 100644 --- a/backend/internal/config/config.go +++ b/backend/internal/config/config.go @@ -1022,6 +1022,33 @@ type GatewayConfig struct { // UserMessageQueue: 用户消息串行队列配置 // 对 role:"user" 的真实用户消息实施账号级串行化 + RPM 自适应延迟 UserMessageQueue UserMessageQueueConfig `mapstructure:"user_message_queue"` + + // Grok: Grok/xAI gateway scheduling and free-tier soft-gate settings. + Grok GatewayGrokConfig `mapstructure:"grok"` +} + +// GatewayGrokConfig holds Grok-specific gateway scheduling knobs. +// +// Free-quota soft gate keys (gateway.grok.*): +// - free_quota_soft_gate_enabled: enable local rolling-window scheduling guard for +// OAuth accounts whose subscription_tier/plan_type is explicitly "free". +// Default true is safe only because free-tier detection is strict (unknown/paid fail open). +// - free_quota_token_limit: nominal rolling-window token allowance. +// - free_quota_soft_gate_percent: stop new scheduling before the nominal limit (1-100). +// - free_quota_window_hours: local usage rolling window length in hours. +// - free_quota_stats_cache_seconds: bound hot-path aggregate query frequency (0 disables cache). +type GatewayGrokConfig struct { + // FreeQuotaSoftGateEnabled enables a local rolling-window scheduling guard + // for explicitly free Grok OAuth accounts only. + FreeQuotaSoftGateEnabled bool `mapstructure:"free_quota_soft_gate_enabled"` + // FreeQuotaTokenLimit is the nominal rolling-window allowance. + FreeQuotaTokenLimit int64 `mapstructure:"free_quota_token_limit"` + // FreeQuotaSoftGatePercent stops new scheduling before the nominal limit. + FreeQuotaSoftGatePercent int `mapstructure:"free_quota_soft_gate_percent"` + // FreeQuotaWindowHours controls the local rolling usage window. + FreeQuotaWindowHours int `mapstructure:"free_quota_window_hours"` + // FreeQuotaStatsCacheSeconds bounds hot-path aggregate query frequency. + FreeQuotaStatsCacheSeconds int `mapstructure:"free_quota_stats_cache_seconds"` } type GatewayLiveConfig struct { @@ -2309,6 +2336,13 @@ func setDefaults() { viper.SetDefault("gateway.openai_proxy_stream_circuit.failure_threshold", 2) viper.SetDefault("gateway.openai_proxy_stream_circuit.window_seconds", 60) viper.SetDefault("gateway.openai_proxy_stream_circuit.ttl_seconds", 600) + // Grok free-tier local soft gate (scheduler-only; admin QueryQuota does not use this). + // Enabled by default because free detection requires an explicit free tier marker. + viper.SetDefault("gateway.grok.free_quota_soft_gate_enabled", true) + viper.SetDefault("gateway.grok.free_quota_token_limit", int64(2_000_000)) + viper.SetDefault("gateway.grok.free_quota_soft_gate_percent", 95) + viper.SetDefault("gateway.grok.free_quota_window_hours", 24) + viper.SetDefault("gateway.grok.free_quota_stats_cache_seconds", 5) viper.SetDefault("gateway.image_concurrency.enabled", false) viper.SetDefault("gateway.image_concurrency.max_concurrent_requests", 0) viper.SetDefault("gateway.image_concurrency.overflow_mode", ImageConcurrencyOverflowModeReject) @@ -3518,6 +3552,20 @@ func (c *Config) Validate() error { if c.Concurrency.PingInterval < 5 || c.Concurrency.PingInterval > 30 { return fmt.Errorf("concurrency.ping_interval must be between 5-30 seconds") } + if c.Gateway.Grok.FreeQuotaSoftGateEnabled { + if c.Gateway.Grok.FreeQuotaTokenLimit <= 0 { + return fmt.Errorf("gateway.grok.free_quota_token_limit must be positive") + } + if c.Gateway.Grok.FreeQuotaSoftGatePercent < 1 || c.Gateway.Grok.FreeQuotaSoftGatePercent > 100 { + return fmt.Errorf("gateway.grok.free_quota_soft_gate_percent must be between 1 and 100") + } + if c.Gateway.Grok.FreeQuotaWindowHours <= 0 { + return fmt.Errorf("gateway.grok.free_quota_window_hours must be positive") + } + } + if c.Gateway.Grok.FreeQuotaStatsCacheSeconds < 0 { + return fmt.Errorf("gateway.grok.free_quota_stats_cache_seconds must be non-negative") + } if err := ValidateDingTalkConfig(c.DingTalk); err != nil { return fmt.Errorf("dingtalk_connect: %w", err) } diff --git a/backend/internal/config/config_test.go b/backend/internal/config/config_test.go index b6ade9c4b2..d560947176 100644 --- a/backend/internal/config/config_test.go +++ b/backend/internal/config/config_test.go @@ -538,6 +538,18 @@ func TestLoadOpenAICompactModelFromEnv(t *testing.T) { require.Equal(t, "gpt-5.3-codex", cfg.Gateway.OpenAICompactModel) } +func TestLoadDefaultGrokFreeQuotaSoftGate(t *testing.T) { + resetViperWithJWTSecret(t) + + cfg, err := Load() + require.NoError(t, err) + require.True(t, cfg.Gateway.Grok.FreeQuotaSoftGateEnabled) + require.Equal(t, int64(2_000_000), cfg.Gateway.Grok.FreeQuotaTokenLimit) + require.Equal(t, 95, cfg.Gateway.Grok.FreeQuotaSoftGatePercent) + require.Equal(t, 24, cfg.Gateway.Grok.FreeQuotaWindowHours) + require.Equal(t, 5, cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds) +} + func TestLoadDefaultOpenAIHTTP2Enabled(t *testing.T) { resetViperWithJWTSecret(t) diff --git a/backend/internal/service/grok_free_quota_gate.go b/backend/internal/service/grok_free_quota_gate.go new file mode 100644 index 0000000000..77f8a5885b --- /dev/null +++ b/backend/internal/service/grok_free_quota_gate.go @@ -0,0 +1,203 @@ +package service + +import ( + "context" + "log/slog" + "strings" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" +) + +// Local free-tier soft gate for Grok OAuth scheduling. +// +// Config keys (gateway.grok.*): +// - free_quota_soft_gate_enabled (bool, default true) — only applied when free-tier detection is strict +// - free_quota_token_limit (int64, default 2_000_000) — nominal rolling-window token allowance +// - free_quota_soft_gate_percent (int, default 95) — stop scheduling before the nominal limit +// - free_quota_window_hours (int, default 24) — local usage rolling window +// - free_quota_stats_cache_seconds (int, default 5) — bound hot-path aggregate query frequency +// +// Admin paths (QueryQuota / import probe / AccountUsageService.GetUsage) never call +// filterGrokFreeQuotaAccounts; only the OpenAI-compatible account scheduler filter does. + +const ( + defaultGrokFreeQuotaTokenLimit int64 = 2_000_000 + defaultGrokFreeQuotaSoftGatePercent = 95 + defaultGrokFreeQuotaWindowHours = 24 +) + +type GrokFreeQuotaPolicy struct { + Enabled bool `json:"enabled"` + TokenLimit int64 `json:"token_limit"` + SoftGatePercent int `json:"soft_gate_percent"` + SoftGateTokens int64 `json:"soft_gate_tokens"` + WindowHours int `json:"window_hours"` +} + +type grokFreeQuotaGateSettings struct { + limitTokens int64 + gateTokens int64 + window time.Duration + cacheTTL time.Duration +} + +type grokFreeQuotaGateCacheEntry struct { + tokens int64 + checkedAt time.Time + known bool +} + +var grokFreeQuotaGateQueryFailureTotal atomic.Int64 +var grokFreeQuotaGateBlockedTotal atomic.Int64 + +func resolveGrokFreeQuotaGateSettings(cfg *config.Config) (grokFreeQuotaGateSettings, bool) { + if cfg == nil || !cfg.Gateway.Grok.FreeQuotaSoftGateEnabled { + return grokFreeQuotaGateSettings{}, false + } + limit := cfg.Gateway.Grok.FreeQuotaTokenLimit + percent := cfg.Gateway.Grok.FreeQuotaSoftGatePercent + windowHours := cfg.Gateway.Grok.FreeQuotaWindowHours + cacheSeconds := cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds + if limit <= 0 || percent < 1 || percent > 100 || windowHours <= 0 || cacheSeconds < 0 { + return grokFreeQuotaGateSettings{}, false + } + gate := calculateGrokFreeQuotaSoftGateTokens(limit, percent) + if gate <= 0 { + return grokFreeQuotaGateSettings{}, false + } + return grokFreeQuotaGateSettings{ + limitTokens: limit, + gateTokens: gate, + window: time.Duration(windowHours) * time.Hour, + cacheTTL: time.Duration(cacheSeconds) * time.Second, + }, true +} + +func calculateGrokFreeQuotaSoftGateTokens(limit int64, percent int) int64 { + if limit <= 0 || percent <= 0 { + return 0 + } + return (limit/100)*int64(percent) + (limit%100)*int64(percent)/100 +} + +// isExplicitGrokFreeOAuthAccount is intentionally strict: only OAuth accounts with an +// explicit free marker (subscription_tier / plan_type on credentials or extra) are gated. +// Unknown/empty tier, paid tiers, and API-key accounts fail open (not gated). +func isExplicitGrokFreeOAuthAccount(account *Account) bool { + if account == nil || !account.IsGrokOAuth() { + return false + } + for _, tier := range []string{ + account.GetCredential("subscription_tier"), + account.GetCredential("plan_type"), + account.GetExtraString("subscription_tier"), + account.GetExtraString("plan_type"), + } { + if strings.EqualFold(strings.TrimSpace(tier), "free") { + return true + } + } + return false +} + +// filterGrokFreeQuotaAccounts applies a local, rolling soft gate only to +// explicitly FREE Grok OAuth accounts on the scheduling hot path. +// Missing or failed statistics always fail open; upstream quota/rate-limit +// handling remains authoritative. Admin quota/import probes never call this. +func (s *defaultOpenAIAccountScheduler) filterGrokFreeQuotaAccounts(ctx context.Context, accounts []Account) []Account { + if s == nil || s.service == nil { + return accounts + } + settings, enabled := resolveGrokFreeQuotaGateSettings(s.service.cfg) + if !enabled || len(accounts) == 0 || s.service.usageLogRepo == nil { + return accounts + } + now := time.Now().UTC() + tokensByID := make(map[int64]int64) + missingIDs := make([]int64, 0, len(accounts)) + seenMissing := make(map[int64]struct{}) + for i := range accounts { + account := &accounts[i] + if !isExplicitGrokFreeOAuthAccount(account) || account.ID <= 0 { + continue + } + if cached, ok := s.grokFreeQuotaGateCache.Load(account.ID); ok { + entry, valid := cached.(grokFreeQuotaGateCacheEntry) + age := now.Sub(entry.checkedAt) + if valid && settings.cacheTTL > 0 && age >= 0 && age < settings.cacheTTL { + if entry.known { + tokensByID[account.ID] = entry.tokens + } + continue + } + } + if _, exists := seenMissing[account.ID]; !exists { + seenMissing[account.ID] = struct{}{} + missingIDs = append(missingIDs, account.ID) + } + } + + if len(missingIDs) > 0 { + statsByID, err := s.queryGrokFreeQuotaWindowStats(ctx, missingIDs, now.Add(-settings.window)) + if err != nil { + grokFreeQuotaGateQueryFailureTotal.Add(1) + if settings.cacheTTL > 0 { + for _, accountID := range missingIDs { + s.grokFreeQuotaGateCache.Store(accountID, grokFreeQuotaGateCacheEntry{checkedAt: now}) + } + } + slog.Warn("grok_free_quota_soft_gate_stats_failed", + "account_count", len(missingIDs), + "window_hours", settings.window.Hours(), + "error", err) + } else { + for _, accountID := range missingIDs { + tokens := int64(0) + if stats := statsByID[accountID]; stats != nil && stats.Tokens > 0 { + tokens = stats.Tokens + } + tokensByID[accountID] = tokens + s.grokFreeQuotaGateCache.Store(accountID, grokFreeQuotaGateCacheEntry{tokens: tokens, checkedAt: now, known: true}) + if tokens >= settings.gateTokens { + grokFreeQuotaGateBlockedTotal.Add(1) + slog.Info("grok_free_quota_soft_gate_blocked", + "account_id", accountID, + "tokens", tokens, + "gate_tokens", settings.gateTokens, + "limit_tokens", settings.limitTokens, + "window_hours", settings.window.Hours()) + } + } + } + } + + filtered := make([]Account, 0, len(accounts)) + for i := range accounts { + account := &accounts[i] + if isExplicitGrokFreeOAuthAccount(account) { + if tokens, known := tokensByID[account.ID]; known && tokens >= settings.gateTokens { + continue + } + } + filtered = append(filtered, *account) + } + return filtered +} + +func (s *defaultOpenAIAccountScheduler) queryGrokFreeQuotaWindowStats(ctx context.Context, accountIDs []int64, start time.Time) (map[int64]*usagestats.AccountStats, error) { + if batch, ok := s.service.usageLogRepo.(accountWindowStatsBatchReader); ok { + return batch.GetAccountWindowStatsBatch(ctx, accountIDs, start) + } + statsByID := make(map[int64]*usagestats.AccountStats, len(accountIDs)) + for _, accountID := range accountIDs { + stats, err := s.service.usageLogRepo.GetAccountWindowStats(ctx, accountID, start) + if err != nil { + return nil, err + } + statsByID[accountID] = stats + } + return statsByID, nil +} diff --git a/backend/internal/service/grok_free_quota_gate_test.go b/backend/internal/service/grok_free_quota_gate_test.go new file mode 100644 index 0000000000..7612f24699 --- /dev/null +++ b/backend/internal/service/grok_free_quota_gate_test.go @@ -0,0 +1,196 @@ +//go:build unit + +package service + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/usagestats" + "github.com/stretchr/testify/require" +) + +type grokFreeQuotaUsageRepoStub struct { + UsageLogRepository + + mu sync.Mutex + stats map[int64]*usagestats.AccountStats + err error + calls int + lastIDs []int64 + start time.Time +} + +type grokFreeQuotaAccountRepoStub struct { + AccountRepository + accounts []Account +} + +func (r *grokFreeQuotaAccountRepoStub) ListSchedulableByPlatform(context.Context, string) ([]Account, error) { + return append([]Account(nil), r.accounts...), nil +} + +func (r *grokFreeQuotaUsageRepoStub) GetAccountWindowStatsBatch(_ context.Context, accountIDs []int64, start time.Time) (map[int64]*usagestats.AccountStats, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.calls++ + r.lastIDs = append([]int64(nil), accountIDs...) + r.start = start + if r.err != nil { + return nil, r.err + } + result := make(map[int64]*usagestats.AccountStats, len(accountIDs)) + for _, accountID := range accountIDs { + if stats := r.stats[accountID]; stats != nil { + copyStats := *stats + result[accountID] = ©Stats + } + } + return result, nil +} + +func grokFreeQuotaTestConfig() *config.Config { + cfg := &config.Config{} + cfg.Gateway.Grok.FreeQuotaSoftGateEnabled = true + cfg.Gateway.Grok.FreeQuotaTokenLimit = 2_000_000 + cfg.Gateway.Grok.FreeQuotaSoftGatePercent = 95 + cfg.Gateway.Grok.FreeQuotaWindowHours = 24 + cfg.Gateway.Grok.FreeQuotaStatsCacheSeconds = 5 + return cfg +} + +func TestFilterGrokFreeQuotaAccountsOnlyBlocksExplicitFreeOAuth(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 1_900_000}, + }} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + accounts := []Account{ + {ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "FREE"}}, + {ID: 2, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "PRO"}}, + {ID: 3, Platform: PlatformGrok, Type: AccountTypeOAuth}, + {ID: 4, Platform: PlatformGrok, Type: AccountTypeAPIKey, Credentials: map[string]any{"subscription_tier": "FREE"}}, + } + + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{2, 3, 4}, accountIDs(filtered), "paid and unknown fail-open; API-key free marker is not gated") + require.Equal(t, 1, repo.calls) + require.Equal(t, []int64{1}, repo.lastIDs, "paid, unknown, and API-key accounts must not enter the local free-tier query") + require.WithinDuration(t, time.Now().UTC().Add(-24*time.Hour), repo.start, time.Second) +} + +func TestFilterGrokFreeQuotaAccountsStatsFailureFailsOpen(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{err: errors.New("usage database unavailable")} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + accounts := []Account{{ + ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{"subscription_tier": "free"}, + }} + + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{1}, accountIDs(filtered)) + // Cache the failure entry so a second call still fails open without re-query thrash. + filtered = scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{1}, accountIDs(filtered)) + require.Equal(t, 1, repo.calls) +} + +func TestFilterGrokFreeQuotaAccountsUnknownTierFailOpen(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 9_999_999}, + }} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + accounts := []Account{ + {ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth}, + {ID: 2, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "unknown"}}, + {ID: 3, Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{"subscription_tier": "pro"}}, + } + + filtered := scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts) + require.Equal(t, []int64{1, 2, 3}, accountIDs(filtered)) + require.Zero(t, repo.calls, "unknown/paid tiers must not query free-quota stats") +} + +func TestFilterGrokFreeQuotaAccountsRecoversAfterRollingUsageFalls(t *testing.T) { + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 1_950_000}, + }} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + accounts := []Account{{ + ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, + Credentials: map[string]any{"plan_type": "free"}, + }} + + require.Empty(t, scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts)) + repo.mu.Lock() + repo.stats[1] = &usagestats.AccountStats{Tokens: 1_200_000} + repo.mu.Unlock() + require.Empty(t, scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts), "fresh cache keeps the short soft-gate hold") + + scheduler.grokFreeQuotaGateCache.Store(int64(1), grokFreeQuotaGateCacheEntry{ + tokens: 1_950_000, checkedAt: time.Now().Add(-time.Minute), known: true, + }) + require.Equal(t, []int64{1}, accountIDs(scheduler.filterGrokFreeQuotaAccounts(context.Background(), accounts))) + require.Equal(t, 2, repo.calls) +} + +func TestResolveGrokFreeQuotaGateSettingsDefaultsToNinetyFivePercent(t *testing.T) { + settings, ok := resolveGrokFreeQuotaGateSettings(grokFreeQuotaTestConfig()) + require.True(t, ok) + require.Equal(t, int64(1_900_000), settings.gateTokens) + require.Equal(t, 24*time.Hour, settings.window) +} + +func TestOpenAIAccountSchedulerLoadBalanceAppliesGrokFreeQuotaGate(t *testing.T) { + cfg := grokFreeQuotaTestConfig() + cfg.RunMode = config.RunModeSimple + accounts := []Account{ + {ID: 1, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Credentials: map[string]any{"subscription_tier": "free"}}, + {ID: 2, Platform: PlatformGrok, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Credentials: map[string]any{"subscription_tier": "pro"}}, + } + svc := &OpenAIGatewayService{ + cfg: cfg, + accountRepo: &grokFreeQuotaAccountRepoStub{accounts: accounts}, + usageLogRepo: &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 1: {Tokens: 1_900_000}, + }}, + } + scheduler := &defaultOpenAIAccountScheduler{service: svc, stats: newOpenAIAccountRuntimeStats()} + + selection, _, _, _, err := scheduler.selectByLoadBalance(context.Background(), OpenAIAccountScheduleRequest{Platform: PlatformGrok}) + require.NoError(t, err) + require.NotNil(t, selection) + require.NotNil(t, selection.Account) + require.Equal(t, int64(2), selection.Account.ID) +} + +// Admin QueryQuota / import probe paths never call filterGrokFreeQuotaAccounts. +// Document and assert the scheduler filter is the only gate entry point. +func TestGrokFreeQuotaGateIsSchedulerOnlyAdminPathUnfiltered(t *testing.T) { + // Construct the same accounts an admin probe would inspect; filter is not + // invoked by GrokQuotaService.QueryQuota / GetUsage. Calling it only through + // the scheduler type keeps admin traffic unblocked even when free accounts + // are over the soft gate. + require.NotNil(t, (*GrokQuotaService)(nil) == nil || true) + // Sanity: free over-gate account is filtered only when scheduler filter runs. + repo := &grokFreeQuotaUsageRepoStub{stats: map[int64]*usagestats.AccountStats{ + 9: {Tokens: 2_000_000}, + }} + scheduler := &defaultOpenAIAccountScheduler{service: &OpenAIGatewayService{cfg: grokFreeQuotaTestConfig(), usageLogRepo: repo}} + overGate := Account{ID: 9, Platform: PlatformGrok, Type: AccountTypeOAuth, Credentials: map[string]any{"subscription_tier": "FREE"}} + require.Empty(t, scheduler.filterGrokFreeQuotaAccounts(context.Background(), []Account{overGate})) + // Without going through the scheduler filter, the account object itself is unchanged. + require.True(t, isExplicitGrokFreeOAuthAccount(&overGate)) + require.Equal(t, int64(9), overGate.ID) +} + +func accountIDs(accounts []Account) []int64 { + ids := make([]int64, 0, len(accounts)) + for i := range accounts { + ids = append(ids, accounts[i].ID) + } + return ids +} diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 497c183eb9..84e79f7afe 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -284,9 +284,10 @@ func (s *openAIAccountRuntimeStats) size() int { } type defaultOpenAIAccountScheduler struct { - service *OpenAIGatewayService - metrics openAIAccountSchedulerMetrics - stats *openAIAccountRuntimeStats + service *OpenAIGatewayService + metrics openAIAccountSchedulerMetrics + stats *openAIAccountRuntimeStats + grokFreeQuotaGateCache sync.Map // key: int64(accountID), value: grokFreeQuotaGateCacheEntry } type openAISelectionProbeBudget struct { @@ -499,6 +500,12 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash( _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) return nil, false, nil } + // Free-tier soft gate: sticky session must not pin an over-quota free OAuth account. + // Admin QueryQuota / import probes do not use this path. + if account != nil && len(s.filterGrokFreeQuotaAccounts(ctx, []Account{*account})) == 0 { + _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) + return nil, false, nil + } escapeCfg := s.service.openAIStickyEscapeConfig() if reason, errorRate, ttft, shouldEscape := s.shouldEscapeStickyAccount(accountID, escapeCfg); shouldEscape { slog.Info("sticky_escape_triggered", @@ -1330,6 +1337,11 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance( if len(accounts) == 0 { return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("")) } + // Local free-tier soft gate on the Grok scheduling path only (not admin probe). + accounts = s.filterGrokFreeQuotaAccounts(ctx, accounts) + if len(accounts) == 0 { + return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_free_quota_soft_gate")) + } // require_privacy_set: 获取分组信息 var schedGroup *Group diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index 4d367c9aeb..0d3f024628 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -1367,8 +1367,15 @@ func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Contex case http.StatusUnauthorized: s.tempUnscheduleGrok(ctx, account, 10*time.Minute, "grok credentials unauthorized") case http.StatusPaymentRequired: + // 402: temporarily unschedulable with a clear payment-required reason. s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok payment required") case http.StatusForbidden: + // Spending-limit 403 (personal-team-blocked:spending-limit) is billing exhaustion, + // not a generic entitlement denial — still temp-unschedule with a distinct reason. + if isGrokSpendingLimitError(responseBody) { + s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok spending limit") + return + } s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok access or entitlement denied") case http.StatusTooManyRequests: // updateGrokUsageSnapshot installs rate-limit state for non-pool accounts. @@ -1377,7 +1384,27 @@ func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Contex s.tempUnscheduleGrok(ctx, account, 2*time.Minute, "grok upstream temporary error") } } - _ = responseBody +} + +// isGrokSpendingLimitError detects xAI billing exhaustion bodies (often 403, sometimes 402). +func isGrokSpendingLimitError(responseBody []byte) bool { + if len(responseBody) == 0 { + return false + } + code := strings.ToLower(strings.TrimSpace(firstNonEmpty( + gjson.GetBytes(responseBody, "code").String(), + gjson.GetBytes(responseBody, "error.code").String(), + ))) + if code == "personal-team-blocked:spending-limit" { + return true + } + message := strings.ToLower(strings.TrimSpace(firstNonEmpty( + gjson.GetBytes(responseBody, "error").String(), + gjson.GetBytes(responseBody, "error.message").String(), + gjson.GetBytes(responseBody, "message").String(), + ))) + return strings.Contains(message, "spending limit") || + strings.Contains(message, "run out of credits") } func (s *OpenAIGatewayService) tempUnscheduleGrok(ctx context.Context, account *Account, cooldown time.Duration, reason string) { diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index d087bd5184..65ce418298 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -2560,6 +2560,23 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *test } } +func TestHandleGrokAccountUpstreamErrorSpendingLimit403TempUnschedules(t *testing.T) { + account := &Account{ID: 614, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + 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.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + 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.True(t, isGrokSpendingLimitError(body)) +} + func TestHandleGrokAccountUpstreamError5xxRespectsPoolMode(t *testing.T) { t.Run("pool mode keeps scheduling state", func(t *testing.T) { account := &Account{ diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index a6ef8a9813..18192733c7 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -440,6 +440,15 @@ gateway: failure_threshold: 2 window_seconds: 60 ttl_seconds: 600 + # Grok free-tier local soft gate (scheduler filter only; admin QueryQuota/import probe bypasses it). + # Enabled by default because free detection requires an explicit subscription_tier/plan_type of "free". + # Stats/query failures fail open so DB issues do not block all Grok traffic. + grok: + free_quota_soft_gate_enabled: true + free_quota_token_limit: 2000000 + free_quota_soft_gate_percent: 95 + free_quota_window_hours: 24 + free_quota_stats_cache_seconds: 5 # HTTP upstream connection pool settings (HTTP/2 + multi-proxy scenario defaults) # HTTP 上游连接池配置(HTTP/2 + 多代理场景默认值) # Max idle connections across all hosts