Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/grok_team_rate_limit.go

@@ -0,0 +1,149 @@
+package service
+
+import (
+	"crypto/sha256"
+	"encoding/hex"
+	"strings"
+	"sync"
+	"time"
+)
+
+// In-memory team+model rate-limit overlay for Grok OAuth. When xAI rate-limits
+// one account in a team for a model, sibling accounts sharing team_id skip the
+// same model until the cooldown expires (mirrors grok2api teamModelRateLimit).
+//
+// Process-local only: multi-instance deployments each learn the block from their
+// own 429s. Prefer short TTLs so drift self-heals.
+type grokTeamModelRateLimit struct {
+	Until time.Time
+}
+
+type grokTeamModelRateLimitStore struct {
+	mu    sync.Mutex
+	items map[string]grokTeamModelRateLimit
+}
+
+var globalGrokTeamModelRateLimits = &grokTeamModelRateLimitStore{
+	items: make(map[string]grokTeamModelRateLimit),
+}
+
+const (
+	grokTeamRateLimitDefaultTTL = 10 * time.Minute
+	grokTeamRateLimitMaxTTL     = time.Hour
+	grokTeamRateLimitMinTTL     = 30 * time.Second
+)
+
+func grokTeamFingerprint(teamID string) string {
+	teamID = strings.TrimSpace(teamID)
+	if teamID == "" {
+		return ""
+	}
+	sum := sha256.Sum256([]byte(strings.ToLower(teamID)))
+	return hex.EncodeToString(sum[:8])
+}
+
+func grokTeamModelRateLimitKey(teamFingerprint, model string) string {
+	return teamFingerprint + "|" + strings.ToLower(strings.TrimSpace(model))
+}
+
+func accountGrokTeamID(account *Account) string {
+	if account == nil {
+		return ""
+	}
+	return strings.TrimSpace(account.GetCredential("team_id"))
+}
+
+// markGrokTeamModelRateLimit records that this team+model pair should be skipped
+// until until. No-op when team_id or model is empty.
+func markGrokTeamModelRateLimit(account *Account, model string, until time.Time) {
+	if account == nil || !account.IsGrokOAuth() {
+		return
+	}
+	fp := grokTeamFingerprint(accountGrokTeamID(account))
+	model = strings.TrimSpace(model)
+	if fp == "" || model == "" || until.IsZero() {
+		return
+	}
+	now := time.Now()
+	if !until.After(now) {
+		until = now.Add(grokTeamRateLimitDefaultTTL)
+	}
+	maxUntil := now.Add(grokTeamRateLimitMaxTTL)
+	if until.After(maxUntil) {
+		until = maxUntil
+	}
+	key := grokTeamModelRateLimitKey(fp, model)
+	globalGrokTeamModelRateLimits.mu.Lock()
+	defer globalGrokTeamModelRateLimits.mu.Unlock()
+	if cur, ok := globalGrokTeamModelRateLimits.items[key]; ok && cur.Until.After(until) {
+		return
+	}
+	globalGrokTeamModelRateLimits.items[key] = grokTeamModelRateLimit{Until: until}
+	// Opportunistic prune of expired entries.
+	for k, v := range globalGrokTeamModelRateLimits.items {
+		if !v.Until.After(now) {
+			delete(globalGrokTeamModelRateLimits.items, k)
+		}
+	}
+}
+
+// isGrokTeamModelRateLimited reports whether the account's team is currently
+// blocked for the requested model.
+func isGrokTeamModelRateLimited(account *Account, model string, now time.Time) bool {
+	if account == nil || !account.IsGrokOAuth() {
+		return false
+	}
+	fp := grokTeamFingerprint(accountGrokTeamID(account))
+	model = strings.TrimSpace(model)
+	if fp == "" || model == "" {
+		return false
+	}
+	key := grokTeamModelRateLimitKey(fp, model)
+	globalGrokTeamModelRateLimits.mu.Lock()
+	defer globalGrokTeamModelRateLimits.mu.Unlock()
+	cur, ok := globalGrokTeamModelRateLimits.items[key]
+	if !ok {
+		return false
+	}
+	if !cur.Until.After(now) {
+		delete(globalGrokTeamModelRateLimits.items, key)
+		return false
+	}
+	return true
+}
+
+// filterGrokTeamModelRateLimitedAccounts drops candidates whose team is under a
+// model-scoped rate-limit cool. Accounts without team_id pass through.
+func filterGrokTeamModelRateLimitedAccounts(accounts []Account, model string, now time.Time) []Account {
+	if len(accounts) == 0 || strings.TrimSpace(model) == "" {
+		return accounts
+	}
+	out := accounts[:0]
+	kept := false
+	for i := range accounts {
+		upstreamModel := canonicalOpenAIAccountSchedulingModel(&accounts[i], model)
+		if isGrokTeamModelRateLimited(&accounts[i], upstreamModel, now) {
+			continue
+		}
+		out = append(out, accounts[i])
+		kept = true
+	}
+	if !kept && len(out) == 0 {
+		// All filtered — return empty (caller treats as no capacity).
+		return nil
+	}
+	return out
+}
+
+// resolveGrokTeamRateLimitUntil derives a team cool window from an observed
+// account rate-limit reset, with sane clamps.
+func resolveGrokTeamRateLimitUntil(resetAt, now time.Time) time.Time {
+	if resetAt.After(now.Add(grokTeamRateLimitMinTTL)) {
+		maxUntil := now.Add(grokTeamRateLimitMaxTTL)
+		if resetAt.After(maxUntil) {
+			return maxUntil
+		}
+		return resetAt
+	}
+	return now.Add(grokTeamRateLimitDefaultTTL)
+}