feat(grok): free 档本地用量软门禁与支付失败临时下线

对明确 free 的 OAuth 账号在调度路径应用可配置用量窗口门禁(统计失败 fail-open)。
402 payment required / 消费上限类 403 继续临时移出调度,管理端探测不走门禁。
This commit is contained in:
IanShaw027
2026-08-07 13:04:05 +08:00
parent eb6c9663e7
commit ba58e74b33
9 changed files with 529 additions and 5 deletions
+1 -1
View File
@@ -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/密码入口
+48
View File
@@ -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)
}
+12
View File
@@ -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)
@@ -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
}
@@ -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] = &copyStats
}
}
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
}
@@ -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
@@ -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) {
@@ -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{
+9
View File
@@ -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