mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-07 13:48:43 +08:00
feat(grok): free 档本地用量软门禁与支付失败临时下线
对明确 free 的 OAuth 账号在调度路径应用可配置用量窗口门禁(统计失败 fail-open)。 402 payment required / 消费上限类 403 继续临时移出调度,管理端探测不走门禁。
This commit is contained in:
@@ -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/密码入口
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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] = ©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
|
||||
}
|
||||
@@ -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{
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user