From 6b260392283f10a9b48d7fe0602473cc0e85c309 Mon Sep 17 00:00:00 2001 From: aghost Date: Sun, 18 Oct 2026 16:06:59 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat(antigravity):=20=E4=BC=98=E5=8C=96?= =?UTF-8?q?=E4=BC=9A=E8=AF=9D=E4=BA=B2=E5=92=8C=E8=B0=83=E5=BA=A6=E7=AD=96?= =?UTF-8?q?=E7=95=A5=E4=B8=8E=E5=86=B7=E5=8D=B4=E5=8D=95=E8=B0=83=E6=80=A7?= =?UTF-8?q?=E4=BF=9D=E6=8A=A4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 冷却键补齐与单调性保护： - 写入 public 冷却的同时，同步写入 wireModel 键（Gemini 配额耗尽时追加 family 键）； - 新增 SetModelCooldownUntilAtLeast / MarkModelCooldownUntilAtLeast，同账号共享 wire/family 冷却只延长不缩短； - 对 wireModel==model（如 gemini-4-argon-high）及兼任共享 backing 的模型统一实施 AtLeast 冷却保护，防止兄弟 tier 短 429 覆盖长冷却。 - 跨实例 Aware 过滤器与渠道感知： - 新增 WithAntigravityAwareCooldownFilterContext，auto 渠道下仅对 Antigravity 候选执行三键检查，非 AG 候选保持单键； - native Gemini 端点强制 Provider=antigravity，使 auto Key 能正确触发加权避险； - compact 执行器与 Responses WS 入口显式声明排除 AG，避免对 Codex relay 候选误套用加权。 - 会话亲和与 Fresh 加权避险调度： - 严格保持 Bounded/Strict 会话亲和语义，不做运行时主动迁移状态机，保护 Gemini implicit cache 效益； - 新会话首绑与请求重试时，按 load + 0.50*usage + 0.35*penalty 综合打分排序； - 在 FNV-1a tie-break 后引入 Murmur3 fmix64 算法，消除连续 ASCII 账号 ID 在全空闲同分时的离散偏斜； - 加权函数内部实现单轮 filter 求值保护与 handled 语义，号池全满时安全挂起，不回退重扫 filter。 - 完善覆盖与链路闭环： - 补齐 native countTokens 的 429/503 冷却写入及 admin 连通性测试惩罚隔离； - 补充全面的单元测试（包括别名联动、AtLeast 单调性、auto 渠道过滤、端到端 countTokens 加权、饱和池单次求值及雪崩离散测试等）。 --- auth/antigravity_balance.go | 225 +++++++ auth/antigravity_balance_test.go | 554 ++++++++++++++++++ auth/antigravity_ratelimit.go | 205 +++++++ auth/antigravity_ratelimit_test.go | 341 +++++++++++ auth/store.go | 278 +++++++-- docs/ANTIGRAVITY.md | 12 + docs/CONFIGURATION.md | 8 + ...-17-antigravity-affinity-balance-design.md | 266 +++++++++ proxy/antigravity_auto_channel_test.go | 234 ++++++++ proxy/antigravity_balance_test.go | 125 ++++ proxy/antigravity_compact_selection_test.go | 56 ++ proxy/antigravity_quota.go | 190 +++++- proxy/antigravity_quota_test.go | 495 ++++++++++++++++ proxy/antigravity_ws_selection_test.go | 63 ++ proxy/handler.go | 113 +++- proxy/handler_anthropic.go | 5 +- proxy/handler_gemini.go | 25 +- proxy/responses_ws.go | 13 +- proxy/retry_exclusions.go | 34 +- proxy/scheduler_queue_error_test.go | 4 +- proxy/selection_cancellation_test.go | 18 +- 21 files changed, 3180 insertions(+), 84 deletions(-) create mode 100644 auth/antigravity_balance.go create mode 100644 auth/antigravity_balance_test.go create mode 100644 auth/antigravity_ratelimit.go create mode 100644 auth/antigravity_ratelimit_test.go create mode 100644 docs/superpowers/specs/2026-10-17-antigravity-affinity-balance-design.md create mode 100644 proxy/antigravity_auto_channel_test.go create mode 100644 proxy/antigravity_balance_test.go create mode 100644 proxy/antigravity_compact_selection_test.go create mode 100644 proxy/antigravity_ws_selection_test.go diff --git a/auth/antigravity_balance.go b/auth/antigravity_balance.go new file mode 100644 index 000000000..6ca13c20c --- /dev/null +++ b/auth/antigravity_balance.go @@ -0,0 +1,225 @@ +package auth + +import ( + "context" + "hash/fnv" + "sort" + "strconv" + "strings" + "time" + + "github.com/codex2api/database" +) + +// SelectionOptions 描述一次选号的调度上下文。 +// +// Provider 是下游 Key 限定的上游渠道（database.UpstreamChannel*，空表示不指定/通用）， +// Model 是请求的 public model，Policy 是用量闸门（Standard/Spark）。首绑与请求重试的 +// Fresh 选号据此识别「这是一次 Antigravity 请求」，从而在 round_robin 下启用加权避险； +// 其余渠道与模式完全走既有选号逻辑，行为不变。 +type SelectionOptions struct { + Provider string // 例如 database.UpstreamChannelAntigravity；空表示不指定/通用 + Model string // public model + Policy DispatchPolicy +} + +// schedulerModeRoundRobin 是加权避险唯一生效的调度模式；remaining_quota / fill_first +// 保持既有用量优先语义，只继承冷却正确性修复。 +const schedulerModeRoundRobin = "round_robin" + +// isAntigravityProvider 报告 provider 是否指向 Antigravity 渠道。auth 侧的 +// UpstreamAntigravity 与 database.UpstreamChannelAntigravity 同为 "antigravity"， +// 这里两者都接受，大小写不敏感。 +func isAntigravityProvider(provider string) bool { + provider = strings.TrimSpace(provider) + if provider == "" { + return false + } + return strings.EqualFold(provider, database.UpstreamChannelAntigravity) || + strings.EqualFold(provider, UpstreamAntigravity) +} + +// antigravityBalancedSelectionEnabled 报告本次选号是否应走 Antigravity 加权避险： +// provider=antigravity、进程开关开启、且调度模式为 round_robin 时成立。 +func (s *Store) antigravityBalancedSelectionEnabled(opts SelectionOptions) bool { + if s == nil { + return false + } + if !isAntigravityProvider(opts.Provider) { + return false + } + if !GetAntigravityBalancedAffinity() { + return false + } + return s.GetSchedulerMode() == schedulerModeRoundRobin +} + +// fmix64 是 MurmurHash3 finalizer（也见于 SplitMix64）的 64 位雪崩混淆：三次 +// 「右移 33 位亦或 + 大素数乘法」把输入每一位的差异扩散到全部 64 位输出。 +// +// 纯算术实现，不依赖平台、字节序与编译器内联，任何 GOARCH 上恒得同一结果。 +func fmix64(h uint64) uint64 { + h ^= h >> 33 + h *= 0xff51afd7ed558ccd + h ^= h >> 33 + h *= 0xc4ceb9fe1a85ec53 + h ^= h >> 33 + return h +} + +// antigravityTieBreak 是加权选号同分时的确定性打散键：对 (亲和键, 账号ID) 做 +// FNV-1a，再用 fmix64 做终值雪崩混淆。 +// +// 同一亲和键在号池不变时恒得同一顺序（幂等），不同键把并列账号摊开。 +// +// 终值混淆不是可选项：FNV-1a 只有「乘 1099511628211 后亦或」一步扩散，输入尾部 +// 字节只影响输出的高位区间，低位比特几乎不混合。而这里的输入尾部是十进制 ASCII +// 的账号 ID——同批账号的 ID 是连续整数，后缀高度相似（"…1"、"…2"、"…10"、"…11"）， +// 直接取 FNV-1a 最小值会把并列账号系统性地偏向某几个 ID 前缀/后缀，同分打散退化成 +// 固定偏好。fmix64 之后每个输出比特都依赖全部输入比特，最小值选择才真正与账号身份 +// 无关。 +func antigravityTieBreak(key string, dbID int64) uint64 { + hasher := fnv.New64a() + _, _ = hasher.Write([]byte(key)) + _, _ = hasher.Write([]byte{':'}) + _, _ = hasher.Write([]byte(strconv.FormatInt(dbID, 10))) + return fmix64(hasher.Sum64()) +} + +// selectAntigravityBalancedAccount 是首绑 / Fresh 选号的 Antigravity 加权避险实现。 +// +// 候选按 (调度优先级, 健康档位) 从高到低分层，层内按 +// +// score = load + usage*0.50 + penalty*0.35 （load=occupied/limit，usage=已用配额/100） +// +// 升序打分，同分用 (key, DBID) 哈希做确定性打散；从最高层开始逐个尝试取号，命中即返回， +// 整层都拿不下（并发被抢占 / 命中 cached cooldown）时继续下一层。它只影响「下一个新会话 +// 优先选谁」，不改变任何冷却、绑定或并发语义。 +// +// 第二个返回值 handled 表示加权路径是否已接管本次选号：provider=antigravity、进程开关 +// 开启且调度模式为 round_robin 时成立。handled=false 表示本次不适用（调用方继续走既有 +// spread / HRW / round-robin 逻辑）；handled=true 时调用方必须直接采用本函数的结果， +// 包括返回 nil 的情况——号池确实没有可用容量，不得再回退重扫。 +// +// 单轮求值保护：本函数对每个候选账号的 filter(acc) 至多求值一次（快照循环内），取号循环 +// 只做 cached cooldown 与 tryAcquire，不再触碰 filter，杜绝闭包重复触发与 Redis 读放大。 +// 这也是 handled 存在的原因：一旦本函数已经消费过 filter，任何回退重扫都会二次求值。 +func (s *Store) selectAntigravityBalancedAccount(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, opts SelectionOptions) (*Account, bool) { + if !s.antigravityBalancedSelectionEnabled(opts) { + return nil, false + } + filter = s.withUsableEgressFilter(filter) + + type candidate struct { + acc *Account + schedulerPriority int64 + tierPriority int + limit int64 + score float64 + tie uint64 + } + + maxConcurrency := s.maxConcurrency.Load() + now := time.Now() + penaltyWeight, usageWeight := AntigravityThrottleWeights() + + accounts := s.accountSnapshotAccounts() + candidates := make([]candidate, 0, len(accounts)) + for _, acc := range accounts { + if exclude != nil && exclude[acc.DBID] { + continue + } + if !acc.dispatchableForPolicy(opts.Policy) { + continue + } + if !s.accountAllowedForAPIKey(acc, apiKeyID) { + continue + } + // filter 单次求值：整个选号只在这里对 acc 调一次 filter。 + if filter != nil && !filter(acc) { + continue + } + occupied := accountOccupiedRequests(acc) + tier, _, _, limit := acc.schedulerSnapshotForPolicy(maxConcurrency, opts.Policy) + if limit <= 0 || occupied >= limit { + continue + } + load := float64(occupied) / float64(limit) + usage := acc.usagePercentForScheduling() / 100.0 + penalty := acc.AntigravityPenalty(now) + score := load + usage*usageWeight + penalty*penaltyWeight + candidates = append(candidates, candidate{ + acc: acc, + schedulerPriority: acc.schedulerPriority(), + tierPriority: tierPriority(tier), + limit: limit, + score: score, + tie: antigravityTieBreak(key, acc.DBID), + }) + } + if len(candidates) == 0 { + // 加权路径已接管但号池没有可用容量：交给调用方直接返回 nil 进入等待， + // 不能再回退到会对同一批账号重扫 filter 的常规选号。 + return nil, true + } + + // 一次排序同时确定层序与层内序：层按 (调度优先级, 健康档位) 降序，层内按 score 升序， + // 同分用哈希 tie 升序做确定性打散。 + sort.Slice(candidates, func(i, j int) bool { + left, right := candidates[i], candidates[j] + if left.schedulerPriority != right.schedulerPriority { + return left.schedulerPriority > right.schedulerPriority + } + if left.tierPriority != right.tierPriority { + return left.tierPriority > right.tierPriority + } + if left.score != right.score { + return left.score < right.score + } + return left.tie < right.tie + }) + + // 逐层尝试：运营者设置的调度优先级与健康档位仍然高于低层账号，只有整层都拿不下才 + // 下沉到下一层，而不是直接放弃。 + for start := 0; start < len(candidates); { + end := start + 1 + for end < len(candidates) && + candidates[end].schedulerPriority == candidates[start].schedulerPriority && + candidates[end].tierPriority == candidates[start].tierPriority { + end++ + } + // cached cooldown 只在取号循环里对「即将获取」的账号做一次验证，不做前置 N 次 Redis 读。 + for _, c := range candidates[start:end] { + if s.accountHasBlockingCachedCooldown(c.acc, opts.Policy) { + continue + } + if s.tryAcquireAccount(c.acc, c.limit, true) { + return c.acc, true + } + } + start = end + } + return nil, true +} + +// NextForSessionWithDispatchGuardForProvider 是携带 SelectionOptions 的会话选号入口， +// 供 proxy 在已知请求渠道 / public model 时调用，使 Antigravity 首绑与 Fresh 重选走加权 +// 避险。非 Antigravity 渠道或未开启时与 NextForSessionWithDispatchGuard 行为完全一致。 +func (s *Store) NextForSessionWithDispatchGuardForProvider(opts SelectionOptions, key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter) (*Account, string, SessionAffinityGuard) { + return s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, false, opts) +} + +// NextForContinuationWithDispatchForProvider 是续链请求的 SelectionOptions 变体。续链 +// 只认绑定账号，加权避险不改变其绑定复用；opts 把 Policy 与渠道一起透传到底层选号—— +// 绑定不存在（首次绑定）时续链会落到 Fresh 选号，Provider 缺失会让 Antigravity 请求 +// 在重试链上丢掉加权避险上下文。 +func (s *Store) NextForContinuationWithDispatchForProvider(opts SelectionOptions, key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter) (*Account, string) { + account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, true, opts) + return account, proxyURL +} + +// WaitForDispatchAvailableForProvider 是 WaitForDispatchAvailable 的 SelectionOptions 变体， +// 让排队等待路径也能识别 Antigravity 请求并在恢复可用时走加权避险选号。 +func (s *Store) WaitForDispatchAvailableForProvider(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, opts SelectionOptions, preserveBinding bool, heartbeat ...SchedulerWaitHeartbeat) (*Account, string, SessionAffinityGuard, error) { + return s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, preserveBinding, opts, heartbeat...) +} diff --git a/auth/antigravity_balance_test.go b/auth/antigravity_balance_test.go new file mode 100644 index 000000000..61676b230 --- /dev/null +++ b/auth/antigravity_balance_test.go @@ -0,0 +1,554 @@ +package auth + +import ( + "fmt" + "hash/fnv" + "math/bits" + "strconv" + "sync" + "testing" + "time" + + "github.com/codex2api/database" +) + +// newAntigravityBalanceStore 构造一个 round_robin + 加权避险开启的号池，账号均为 +// 可调度的 Antigravity OAuth 账号。测试结束恢复全局开关默认值。 +func newAntigravityBalanceStore(t *testing.T, accounts int) *Store { + t.Helper() + store := NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 4, TestConcurrency: 1, TestModel: "gemini-3.8-flash-high"}) + store.SetSchedulerMode("round_robin") + for i := 1; i <= accounts; i++ { + store.AddAccount(&Account{ + DBID: int64(i), + UpstreamType: UpstreamAntigravity, + AccessToken: fmt.Sprintf("google-at-%d", i), + AntigravityProjectID: fmt.Sprintf("project-%d", i), + }) + } + t.Cleanup(func() { SetAntigravityBalancedAffinity(true) }) + SetAntigravityBalancedAffinity(true) + return store +} + +func antigravitySelectionOpts() SelectionOptions { + return SelectionOptions{ + Provider: database.UpstreamChannelAntigravity, + Model: "gemini-3.8-flash-high", + Policy: DispatchPolicyStandard, + } +} + +// 加权避险：4 个健康账号，账号 1 记录一次 429 惩罚后，多次不同 key 的首选必须显著 +// 避开账号 1，而其余健康账号仍被正常选中。 +func TestAntigravityBalancedSelectionAvoidsPenalizedAccount(t *testing.T) { + store := newAntigravityBalanceStore(t, 4) + defer store.Stop() + + opts := antigravitySelectionOpts() + const trials = 400 + + // 控制组：未记惩罚时四个健康账号应大致摊开，证明选号本身没有把账号 1 天然压到队尾。 + control := map[int64]int{} + for i := 0; i < trials; i++ { + acc := store.nextAccountForFreshAffinityWithDispatch(fmt.Sprintf("control-%d", i), 0, nil, nil, opts) + if acc == nil { + t.Fatal("control selection returned no account") + } + control[acc.DBID]++ + store.Release(acc) + } + if len(control) < 2 { + t.Fatalf("control selections did not spread across accounts: %v", control) + } + if control[1] == 0 { + t.Fatalf("control run already starves account 1 before any penalty: %v", control) + } + + penalized := store.FindByID(1) + if penalized == nil { + t.Fatal("penalized account 1 not found") + } + RecordAntigravityRateLimitEvent(penalized, AntigravityRateLimitEvent{ + StatusCode: 429, + Kind: "rate_limit", + Scope: AGScopeModelRateLimit, + Weight: AntigravityPenaltyWeightExplicitRateLimit, + ObservedAt: time.Now(), + }) + + treatment := map[int64]int{} + for i := 0; i < trials; i++ { + acc := store.nextAccountForFreshAffinityWithDispatch(fmt.Sprintf("treatment-%d", i), 0, nil, nil, opts) + if acc == nil { + t.Fatal("treatment selection returned no account") + } + treatment[acc.DBID]++ + store.Release(acc) + } + + minHealthy := -1 + for id, count := range treatment { + if id == 1 { + continue + } + if count == 0 { + t.Fatalf("healthy account %d was never selected after penalizing account 1: %v", id, treatment) + } + if minHealthy < 0 || count < minHealthy { + minHealthy = count + } + } + if minHealthy <= 0 { + t.Fatalf("no healthy account selected in treatment run: %v", treatment) + } + // 账号 1 的选中次数必须显著低于最少的健康账号（阈值：不足其 1/4）。 + if treatment[1]*4 >= minHealthy { + t.Fatalf("penalized account 1 not avoided enough: %v (want count[1] well below %d)", treatment, minHealthy) + } +} + +// 同分（全部空闲、无惩罚）时用 (key, DBID) 哈希做确定性打散：同一 key 恒命中同一 +// 账号，不同 key 摊到多个账号。 +func TestAntigravityBalancedSelectionDeterministicAndSpread(t *testing.T) { + store := newAntigravityBalanceStore(t, 4) + defer store.Stop() + opts := antigravitySelectionOpts() + + const key = "affinity-stable" + var first int64 + for i := 0; i < 6; i++ { + acc := store.nextAccountForFreshAffinityWithDispatch(key, 0, nil, nil, opts) + if acc == nil { + t.Fatal("no account selected") + } + if first == 0 { + first = acc.DBID + } else if acc.DBID != first { + t.Fatalf("same key drifted across selections: got %d want %d", acc.DBID, first) + } + store.Release(acc) + } + + seen := map[int64]bool{} + for i := 0; i < 60; i++ { + acc := store.nextAccountForFreshAffinityWithDispatch(fmt.Sprintf("spread-%d", i), 0, nil, nil, opts) + if acc == nil { + t.Fatal("no account selected") + } + seen[acc.DBID] = true + store.Release(acc) + } + if len(seen) < 2 { + t.Fatalf("distinct keys did not spread across accounts: %v", seen) + } +} + +// 单轮求值保护：加权选号对每个候选账号的 filter(acc) 至多求值一次。用「同一账号第二次 +// 求值即 panic」的 filter 直接验证——若发生重复求值，本测试会以 panic 失败。 +func TestAntigravityBalancedSelectionEvaluatesFilterOnce(t *testing.T) { + store := newAntigravityBalanceStore(t, 4) + defer store.Stop() + opts := antigravitySelectionOpts() + + var seen sync.Map + filter := func(acc *Account) bool { + if acc == nil { + return false + } + if _, loaded := seen.LoadOrStore(acc.DBID, struct{}{}); loaded { + panic(fmt.Sprintf("filter evaluated more than once for account %d", acc.DBID)) + } + return true + } + acc := store.nextAccountForFreshAffinityWithDispatch("single-eval", 0, nil, filter, opts) + if acc == nil { + t.Fatal("no account selected") + } + store.Release(acc) +} + +// 触发门控：仅 provider=antigravity + 开关开启 + round_robin 三者同时满足才走加权路径。 +// 其余组合必须返回 handled=false，让调用方继续走既有 spread / HRW / round-robin 逻辑。 +func TestAntigravityBalancedSelectionGating(t *testing.T) { + store := newAntigravityBalanceStore(t, 2) + defer store.Stop() + + if acc, handled := store.selectAntigravityBalancedAccount("k", 0, nil, nil, SelectionOptions{Provider: database.UpstreamChannelCodex, Policy: DispatchPolicyStandard}); acc != nil || handled { + t.Fatal("non-antigravity provider entered weighted selection") + } + if acc, handled := store.selectAntigravityBalancedAccount("k", 0, nil, nil, SelectionOptions{Policy: DispatchPolicyStandard}); acc != nil || handled { + t.Fatal("empty provider entered weighted selection") + } + + acc, handled := store.selectAntigravityBalancedAccount("k", 0, nil, nil, antigravitySelectionOpts()) + if !handled { + t.Fatal("antigravity provider did not take over the selection") + } + if acc == nil { + t.Fatal("antigravity provider did not select an account") + } + store.Release(acc) + + SetAntigravityBalancedAffinity(false) + if acc, handled := store.selectAntigravityBalancedAccount("k", 0, nil, nil, antigravitySelectionOpts()); acc != nil || handled { + t.Fatal("weighted selection stayed active after disabling the switch") + } + SetAntigravityBalancedAffinity(true) + + store.SetSchedulerMode("fill_first") + if acc, handled := store.selectAntigravityBalancedAccount("k", 0, nil, nil, antigravitySelectionOpts()); acc != nil || handled { + t.Fatal("weighted selection stayed active outside round_robin") + } + store.SetSchedulerMode("round_robin") +} + +// 加权路径必须尊重 exclude：被排除的账号即使分数最低也不能被选中。 +func TestAntigravityBalancedSelectionRespectsExclude(t *testing.T) { + store := newAntigravityBalanceStore(t, 3) + defer store.Stop() + opts := antigravitySelectionOpts() + + for i := 0; i < 40; i++ { + acc, handled := store.selectAntigravityBalancedAccount(fmt.Sprintf("excl-%d", i), 0, map[int64]bool{1: true, 2: true}, nil, opts) + if !handled { + t.Fatal("weighted selection did not take over") + } + if acc == nil { + t.Fatal("no account selected") + } + if acc.DBID == 1 || acc.DBID == 2 { + t.Fatalf("excluded account %d was selected", acc.DBID) + } + store.Release(acc) + } +} + +// saturateAntigravityAccount 把账号的占用/活跃并发顶到任何档位上限之上，模拟号池已满。 +func saturateAntigravityAccount(acc *Account) { + if acc == nil { + return + } + acc.ActiveRequests.Store(64) + acc.OccupiedRequests.Store(64) +} + +// 饱和池：饱和号池下加权路径必须自己接管本次选号——返回 nil 且不得让调用方回退重扫。 +// 用「同一账号第二次求值即 panic」的 filter 断言 filter 每个账号至多一次；spread 开与关 +// 两条回退路径都要覆盖（回退分别落在 NextExcludingWithDispatch 与 HRW 候选扫描）。 +func TestAntigravityBalancedSelectionSaturatedPoolDoesNotRefilter(t *testing.T) { + for _, spread := range []bool{false, true} { + name := "spread-off" + if spread { + name = "spread-on" + } + t.Run(name, func(t *testing.T) { + store := newAntigravityBalanceStore(t, 4) + defer store.Stop() + store.SetSessionAffinitySpread(spread) + for _, acc := range store.Accounts() { + saturateAntigravityAccount(acc) + } + + var mu sync.Mutex + evaluated := map[int64]int{} + filter := func(acc *Account) bool { + if acc == nil { + return false + } + mu.Lock() + defer mu.Unlock() + evaluated[acc.DBID]++ + if evaluated[acc.DBID] > 1 { + panic(fmt.Sprintf("filter evaluated %d times for account %d", evaluated[acc.DBID], acc.DBID)) + } + return true + } + + if acc := store.nextAccountForFreshAffinityWithDispatch("saturated", 0, nil, filter, antigravitySelectionOpts()); acc != nil { + t.Fatalf("a saturated pool must not hand out account %d", acc.DBID) + } + + mu.Lock() + defer mu.Unlock() + if len(evaluated) == 0 { + t.Fatal("the weighted path never evaluated the filter") + } + for dbID, count := range evaluated { + if count > 1 { + t.Fatalf("account %d evaluated the filter %d times, want at most 1", dbID, count) + } + } + }) + } +} + +// 饱和池：handled=true 但整层拿不下时返回 nil；加权路径内部按 (调度优先级, 健康档位) +// 逐层下沉，最高层全满必须继续尝试次高层，而不是直接返回 nil。 +func TestAntigravityBalancedSelectionFallsBackToLowerPriorityLayer(t *testing.T) { + store := newAntigravityBalanceStore(t, 3) + defer store.Stop() + opts := antigravitySelectionOpts() + + high := int64(10) + if !store.ApplyAccountSchedulerPriority(1, &high) { + t.Fatal("failed to raise the priority of account 1") + } + t.Cleanup(func() { store.ApplyAccountSchedulerPriority(1, nil) }) + + // 未饱和时最高层独占：同一批 key 只能命中账号 1。 + for i := 0; i < 20; i++ { + acc, handled := store.selectAntigravityBalancedAccount(fmt.Sprintf("prio-%d", i), 0, nil, nil, opts) + if !handled { + t.Fatal("weighted selection did not take over") + } + if acc == nil { + t.Fatal("no account selected") + } + if acc.DBID != 1 { + t.Fatalf("lower priority account %d won while the top layer was free", acc.DBID) + } + store.Release(acc) + } + + // 最高层全满：必须下沉到次高层，而不是返回 nil。 + saturateAntigravityAccount(store.FindByID(1)) + seen := map[int64]bool{} + for i := 0; i < 20; i++ { + acc, handled := store.selectAntigravityBalancedAccount(fmt.Sprintf("sunk-%d", i), 0, nil, nil, opts) + if !handled { + t.Fatal("weighted selection did not take over") + } + if acc == nil { + t.Fatal("a free lower layer must still be selected") + } + if acc.DBID == 1 { + t.Fatal("the saturated top layer handed out an account") + } + seen[acc.DBID] = true + store.Release(acc) + } + if len(seen) == 0 { + t.Fatal("no lower-layer account was reachable") + } +} + +// 饱和池：handled=false（非 Antigravity 渠道）时旧路径行为不变——散列绑定仍然生效， +// 同一亲和键幂等命中同一账号，不同键摊到多个账号。 +func TestAntigravityFreshAffinityLegacyPathForNonAntigravity(t *testing.T) { + store := newAntigravityBalanceStore(t, 4) + defer store.Stop() + store.SetSessionAffinitySpread(true) + t.Cleanup(func() { store.SetSessionAffinitySpread(false) }) + + opts := SelectionOptions{Provider: database.UpstreamChannelCodex, Model: "gpt-5.4", Policy: DispatchPolicyStandard} + + var first int64 + for i := 0; i < 6; i++ { + acc := store.nextAccountForFreshAffinityWithDispatch("legacy-stable", 0, nil, nil, opts) + if acc == nil { + t.Fatal("legacy path selected no account") + } + if first == 0 { + first = acc.DBID + } else if acc.DBID != first { + t.Fatalf("legacy hash binding drifted: got %d want %d", acc.DBID, first) + } + store.Release(acc) + } + + seen := map[int64]bool{} + for i := 0; i < 60; i++ { + acc := store.nextAccountForFreshAffinityWithDispatch(fmt.Sprintf("legacy-spread-%d", i), 0, nil, nil, opts) + if acc == nil { + t.Fatal("legacy path selected no account") + } + seen[acc.DBID] = true + store.Release(acc) + } + if len(seen) < 2 { + t.Fatalf("legacy path stopped spreading across accounts: %v", seen) + } +} + +// antigravityTieBreakRaw 是修复前的打散键：FNV-1a 终值不做雪崩混淆。测试用它证明 +// 下面的分布断言确实能抓到「去掉 fmix64」这一回归。 +func antigravityTieBreakRaw(key string, dbID int64) uint64 { + hasher := fnv.New64a() + _, _ = hasher.Write([]byte(key)) + _, _ = hasher.Write([]byte{':'}) + _, _ = hasher.Write([]byte(strconv.FormatInt(dbID, 10))) + return hasher.Sum64() +} + +// antigravityTieBreakWinner 统计「同分时取哈希最小值」在 accounts 个连续 DBID 账号上的 +// 实际胜率分布：这才是 tie-break 真正被使用的方式，也是十进制 ASCII 偏斜唯一会显形的地方。 +func antigravityTieBreakWinner(tie func(string, int64) uint64, accounts, trials int) []int { + wins := make([]int, accounts) + for i := 1; i <= trials; i++ { + key := fmt.Sprintf("user-session-%d", i) + best, bestHash := 0, tie(key, 1) + for account := int64(2); account <= int64(accounts); account++ { + if got := tie(key, account); got < bestHash { + best, bestHash = int(account-1), got + } + } + wins[best]++ + } + return wins +} + +func antigravityChiSquared(wins []int) float64 { + total, expected := 0, 0.0 + for _, win := range wins { + total += win + } + expected = float64(total) / float64(len(wins)) + chi := 0.0 + for _, win := range wins { + diff := float64(win) - expected + chi += diff * diff / expected + } + return chi +} + +// TestAntigravityTieBreakGoldenVectors 钉住跨平台确定性：纯算术 fmix64 在任何 GOARCH +// 上都必须给出同一批常量，否则同分顺序会随构建产物漂移。 +func TestAntigravityTieBreakGoldenVectors(t *testing.T) { + cases := []struct { + key string + dbID int64 + want uint64 + }{ + {"", 0, 0x43751afc069a4a45}, + {"user-session-1", 1, 0xf155dbb4071460da}, + {"gemini-key", 42, 0xcf28f3d8add60b65}, + {"中文亲和键", 1000003, 0x8dcf2c161bc0b644}, + } + for _, tc := range cases { + if got := antigravityTieBreak(tc.key, tc.dbID); got != tc.want { + t.Fatalf("antigravityTieBreak(%q, %d) = 0x%016x, want 0x%016x", tc.key, tc.dbID, got, tc.want) + } + } + if got := fmix64(0); got != 0 { + t.Fatalf("fmix64(0) = 0x%016x, want 0", got) + } + if got := fmix64(1); got != 0xb456bcfc34c2cb2c { + t.Fatalf("fmix64(1) = 0x%016x, want 0xb456bcfc34c2cb2c", got) + } + if got := fmix64(0xffffffffffffffff); got != 0x64b5720b4b825f21 { + t.Fatalf("fmix64(^0) = 0x%016x, want 0x64b5720b4b825f21", got) + } +} + +// TestAntigravityTieBreakDeterministicAndSensitive 覆盖打散键的基本性质：同输入恒等、 +// 键或账号变化时输出变化、大批量样本内无碰撞。 +func TestAntigravityTieBreakDeterministicAndSensitive(t *testing.T) { + if antigravityTieBreak("stable-key", 7) != antigravityTieBreak("stable-key", 7) { + t.Fatal("antigravityTieBreak is not deterministic") + } + seen := make(map[uint64]string, 20000) + for i := 1; i <= 10000; i++ { + key := fmt.Sprintf("affinity-%d", i) + for _, dbID := range []int64{int64(i), int64(i) + 1000000} { + hash := antigravityTieBreak(key, dbID) + sample := fmt.Sprintf("%s/%d", key, dbID) + if previous, ok := seen[hash]; ok { + t.Fatalf("hash collision between %s and %s", previous, sample) + } + seen[hash] = sample + } + } + // 相邻账号 ID（十进制 ASCII 只差最后一位）必须给出完全不同的哈希。 + similar := map[uint64]int{} + for dbID := int64(1); dbID <= 1000; dbID++ { + similar[antigravityTieBreak("same-key", dbID)]++ + } + if len(similar) != 1000 { + t.Fatalf("1000 consecutive account IDs produced %d distinct hashes", len(similar)) + } +} + +// TestAntigravityTieBreakAvalanche 覆盖雪崩性：输入单比特翻转后，输出平均应有约一半 +// 比特翻转。FNV-1a 单独使用时该值明显低于 32（尾部字节只影响高位区间），终值混淆把它 +// 拉回理论值。 +func TestAntigravityTieBreakAvalanche(t *testing.T) { + total, samples := 0.0, 0 + for dbID := int64(0); dbID < 2000; dbID++ { + base := antigravityTieBreak("avalanche-key", dbID) + for bit := uint64(0); bit < 64; bit++ { + flipped := antigravityTieBreak("avalanche-key", int64(uint64(dbID)^(1< 33 { + t.Fatalf("mean avalanche distance = %.2f bits, want ~32 of 64", mean) + } + // 对照：未混淆的 FNV-1a 明显偏离，说明本断言对终值混淆敏感。 + rawTotal, rawSamples := 0.0, 0 + for dbID := int64(0); dbID < 2000; dbID++ { + base := antigravityTieBreakRaw("avalanche-key", dbID) + for bit := uint64(0); bit < 64; bit++ { + flipped := antigravityTieBreakRaw("avalanche-key", int64(uint64(dbID)^(1<= mean-0.5 { + t.Fatalf("raw FNV-1a avalanche mean = %.2f, fmix64 mean = %.2f: the finalizer stopped contributing", rawMean, mean) + } +} + +// TestAntigravityTieBreakSpreadAcrossAccountIDs 是 P2-1 的核心断言：同分账号按哈希最小值 +// 打散时，每个账号的胜率必须接近 1/N。十进制 ASCII 后缀在纯 FNV-1a 下会把最小值系统性 +// 地推向某个账号（实测 8 个账号里有一个吃到 50% 的样本），fmix64 之后应回到均匀分布。 +func TestAntigravityTieBreakSpreadAcrossAccountIDs(t *testing.T) { + const accounts, trials = 8, 20000 + + wins := antigravityTieBreakWinner(antigravityTieBreak, accounts, trials) + expected := float64(trials) / float64(accounts) + for index, win := range wins { + share := float64(win) / expected + if share < 0.6 || share > 1.4 { + t.Fatalf("account %d won %.1f%% of ties, expected %.0f of %d, wins=%v", index+1, share*100, expected, trials, wins) + } + } + if chi := antigravityChiSquared(wins); chi > 60 { + t.Fatalf("tie-break winner chi-squared = %.1f (df=%d), wins=%v, want a uniform spread", chi, accounts-1, wins) + } + + // 对照：去掉终值混淆后同一批输入必须显著偏斜，否则本测试抓不到回归。 + rawWins := antigravityTieBreakWinner(antigravityTieBreakRaw, accounts, trials) + rawMax := 0 + for _, win := range rawWins { + if win > rawMax { + rawMax = win + } + } + if float64(rawMax) <= expected*1.5 { + t.Fatalf("raw FNV-1a winner distribution is unexpectedly flat (max %d of %d): this test no longer proves the need for fmix64", rawMax, trials) + } + + // 大 ID / 稀疏 ID 同样要摊开：真实部署里账号 ID 是数据库自增主键，可能已经很大。 + for _, ids := range [][]int64{{1000003, 1000004, 1000005, 1000006}, {1, 7, 999999, 1000000, 4294967296}} { + counts := make([]int, len(ids)) + for i := 1; i <= trials; i++ { + key := fmt.Sprintf("spread-%d", i) + best, bestHash := 0, antigravityTieBreak(key, ids[0]) + for index := 1; index < len(ids); index++ { + if got := antigravityTieBreak(key, ids[index]); got < bestHash { + best, bestHash = index, got + } + } + counts[best]++ + } + exp := float64(trials) / float64(len(ids)) + for index, count := range counts { + if share := float64(count) / exp; share < 0.55 || share > 1.45 { + t.Fatalf("ids %v: account %d won %.1f%% of ties, counts=%v", ids, index, share*100, counts) + } + } + } +} diff --git a/auth/antigravity_ratelimit.go b/auth/antigravity_ratelimit.go new file mode 100644 index 000000000..72b0aca56 --- /dev/null +++ b/auth/antigravity_ratelimit.go @@ -0,0 +1,205 @@ +package auth + +import ( + "math" + "os" + "strings" + "sync/atomic" + "time" +) + +// Antigravity 限流事件与首绑/Fresh 选号避险惩罚。 +// +// 上游 Cloud Code 在 429/503 上给出的信号粒度差别很大：RATE_LIMIT_EXCEEDED 带 +// RetryInfo 是「这个凭据此刻确实被限流」，MODEL_CAPACITY_EXHAUSTED 是 Google 全 +// 局容量不足（换号无用），QUOTA_EXHAUSTED 已经有长冷却 + Gemini family 冷却兜 +// 底。首绑加权避险只应该为前一类信号付出代价，所以这里把事件规范化成带权重的 +// 记录，并把权重压进一个进程内、随时间指数衰减的惩罚累加器（不落库、不进 +// Redis）：惩罚只影响「下一个新会话选谁」，不改变任何既有冷却与绑定语义。 + +// AntigravityRateLimitScope 是限流事件的归因范围。 +type AntigravityRateLimitScope string + +const ( + // AGScopeModelRateLimit：账号在该模型上被明确限流（含裸 RESOURCE_EXHAUSTED）。 + AGScopeModelRateLimit AntigravityRateLimitScope = "model_rate_limit" + // AGScopeSharedCapacity：MODEL_CAPACITY_EXHAUSTED，全模型共享容量短缺。 + AGScopeSharedCapacity AntigravityRateLimitScope = "shared_capacity" + // AGScopeAccountQuota：QUOTA_EXHAUSTED，已有长冷却与 family 冷却兜底。 + AGScopeAccountQuota AntigravityRateLimitScope = "account_quota" + // AGScopeUnknown：无法解析出 google.rpc 结构的 429。 + AGScopeUnknown AntigravityRateLimitScope = "unknown" +) + +// AntigravityRateLimitEvent 是一次归一化后的 429/503 观测。 +type AntigravityRateLimitEvent struct { + StatusCode int + Kind string + Scope AntigravityRateLimitScope + PublicModel string + WireModel string + Family string + Weight float64 + ResetAt time.Time + ObservedAt time.Time +} + +// 惩罚累加器的时间常数：半衰期 60s、窗口 180s、上限 4.0。 +// 权重常量供打分侧（首绑/Fresh 选号）与记录侧共用，避免两处各写一份魔数。 +const ( + antigravityThrottleWindow = 180 * time.Second + antigravityThrottleHalfLife = 60 * time.Second + antigravityThrottleMaxPenalty = 4.0 + antigravityThrottleWeight = 0.35 + antigravityUsageWeight = 0.50 +) + +// 事件权重（仅用于首绑/Fresh 选号避险打分，不参与冷却时长计算）。 +const ( + // AntigravityPenaltyWeightExplicitRateLimit：有明确 reason 或 RetryInfo 的限流。 + AntigravityPenaltyWeightExplicitRateLimit = 1.0 + // AntigravityPenaltyWeightUnclassifiedRateLimit：裸 RESOURCE_EXHAUSTED 或解析 + // 失败的非结构化 429——没有任何信息指向这个凭据，只记半个惩罚。 + AntigravityPenaltyWeightUnclassifiedRateLimit = 0.5 + // AntigravityPenaltyWeightNone：共享容量与账号配额不记短期惩罚。 + AntigravityPenaltyWeightNone = 0.0 +) + +// AntigravityThrottleWeights 返回首绑打分使用的惩罚/用量权重。 +func AntigravityThrottleWeights() (penaltyWeight, usageWeight float64) { + return antigravityThrottleWeight, antigravityUsageWeight +} + +// AntigravityThrottleWindow 返回惩罚累加器的有效窗口。 +func AntigravityThrottleWindow() time.Duration { + return antigravityThrottleWindow +} + +// antigravityThrottleDecay 是窗口内的指数衰减因子：elapsed=0 → 1， +// elapsed=半衰期 → 0.5。 +func antigravityThrottleDecay(elapsed time.Duration) float64 { + if elapsed <= 0 { + return 1 + } + if elapsed >= antigravityThrottleWindow { + return 0 + } + halfLives := float64(elapsed) / float64(antigravityThrottleHalfLife) + return math.Pow(0.5, halfLives) +} + +// RecordAntigravityRateLimitEvent 把一次限流观测累加进账号的进程内惩罚。 +// 权重 <= 0 的事件（共享容量、账号配额）不记账；累加值封顶 +// antigravityThrottleMaxPenalty，超过窗口后自然归零。 +func RecordAntigravityRateLimitEvent(acc *Account, event AntigravityRateLimitEvent) { + if acc == nil { + return + } + weight := event.Weight + if math.IsNaN(weight) || weight <= 0 { + return + } + if weight > AntigravityPenaltyWeightExplicitRateLimit { + weight = AntigravityPenaltyWeightExplicitRateLimit + } + now := event.ObservedAt + if now.IsZero() { + now = time.Now() + } + acc.mu.Lock() + defer acc.mu.Unlock() + penalty := acc.antigravityPenalty + if !acc.antigravityPenaltyAt.IsZero() { + elapsed := now.Sub(acc.antigravityPenaltyAt) + if elapsed < 0 { + elapsed = 0 + } + penalty *= antigravityThrottleDecay(elapsed) + } + penalty += weight + if math.IsNaN(penalty) || penalty < 0 { + penalty = 0 + } + if penalty > antigravityThrottleMaxPenalty { + penalty = antigravityThrottleMaxPenalty + } + acc.antigravityPenalty = penalty + acc.antigravityPenaltyAt = now +} + +// AntigravityPenalty 返回账号在当前时刻的衰减后惩罚（0 表示无近期限流）。 +func (a *Account) AntigravityPenalty(now time.Time) float64 { + if a == nil { + return 0 + } + if now.IsZero() { + now = time.Now() + } + a.mu.RLock() + defer a.mu.RUnlock() + if a.antigravityPenalty <= 0 || a.antigravityPenaltyAt.IsZero() { + return 0 + } + elapsed := now.Sub(a.antigravityPenaltyAt) + if elapsed < 0 { + elapsed = 0 + } + penalty := a.antigravityPenalty * antigravityThrottleDecay(elapsed) + if math.IsNaN(penalty) || penalty <= 0 { + return 0 + } + return penalty +} + +// ResetAntigravityPenalty 清空账号的限流惩罚（人工恢复、冷却清理时使用）。 +func (a *Account) ResetAntigravityPenalty() { + if a == nil { + return + } + a.mu.Lock() + defer a.mu.Unlock() + a.antigravityPenalty = 0 + a.antigravityPenaltyAt = time.Time{} +} + +// antigravityBalancedAffinityEnabled 是首绑/Fresh 选号加权避险的进程级开关。 +// 默认开启；ANTIGRAVITY_BALANCED_AFFINITY 取 0/false/off/none/no 时关闭，回退为 +// 原 round-robin。冷却正确性修复不受该开关影响，始终生效。 +var antigravityBalancedAffinityEnabled atomic.Bool + +func init() { + antigravityBalancedAffinityEnabled.Store(antigravityBalancedAffinityFromEnv()) +} + +func antigravityBalancedAffinityFromEnv() bool { + raw, ok := os.LookupEnv(antigravityBalancedAffinityEnv) + if !ok { + return true + } + switch strings.ToLower(strings.TrimSpace(raw)) { + case "", "1", "true", "yes", "y", "on": + return true + case "0", "false", "off", "none", "no": + return false + default: + return true + } +} + +// antigravityBalancedAffinityEnv 是开关的环境变量名。 +const antigravityBalancedAffinityEnv = "ANTIGRAVITY_BALANCED_AFFINITY" + +// SetAntigravityBalancedAffinity 覆盖加权避险开关（启动装配与测试使用）。 +func SetAntigravityBalancedAffinity(enabled bool) { + antigravityBalancedAffinityEnabled.Store(enabled) +} + +// GetAntigravityBalancedAffinity 报告首绑/Fresh 选号是否启用加权避险打分。 +func GetAntigravityBalancedAffinity() bool { + return antigravityBalancedAffinityEnabled.Load() +} + +// ReloadAntigravityBalancedAffinityFromEnv 重新读取环境变量（关闭需重启生效）。 +func ReloadAntigravityBalancedAffinityFromEnv() { + antigravityBalancedAffinityEnabled.Store(antigravityBalancedAffinityFromEnv()) +} diff --git a/auth/antigravity_ratelimit_test.go b/auth/antigravity_ratelimit_test.go new file mode 100644 index 000000000..6c88f1619 --- /dev/null +++ b/auth/antigravity_ratelimit_test.go @@ -0,0 +1,341 @@ +package auth + +import ( + "context" + "math" + "testing" + "time" + + "github.com/codex2api/cache" + "github.com/codex2api/database" +) + +// antigravityGeminiFamilyKey 与 proxy.antigravityGeminiFamilyCooldownKey 同一个 +// 合成键；auth 包不能反向依赖 proxy，这里按字面量断言。 +const antigravityGeminiFamilyKey = "antigravity:gemini" + +func antigravityPenaltyClose(t *testing.T, got, want float64, label string) { + t.Helper() + tolerance := 1e-6 + if want == 0 { + tolerance = 1e-9 + } + if math.Abs(got-want) > tolerance { + t.Fatalf("%s = %v, want %v", label, got, want) + } +} + +func TestAntigravityPenaltyDecaysWithHalfLife(t *testing.T) { + base := time.Date(2026, 10, 17, 12, 0, 0, 0, time.UTC) + acc := &Account{DBID: 8801, AccessToken: "token"} + if got := acc.AntigravityPenalty(base); got != 0 { + t.Fatalf("fresh account penalty = %v, want 0", got) + } + + RecordAntigravityRateLimitEvent(acc, AntigravityRateLimitEvent{ + StatusCode: 429, + Kind: "rate_limit", + Scope: AGScopeModelRateLimit, + PublicModel: "gemini-3.8-flash-high", + WireModel: "gemini-3.8-flash-tiered", + Family: antigravityGeminiFamilyKey, + Weight: AntigravityPenaltyWeightExplicitRateLimit, + ObservedAt: base, + }) + + antigravityPenaltyClose(t, acc.AntigravityPenalty(base), 1.0, "penalty at t0") + antigravityPenaltyClose(t, acc.AntigravityPenalty(base.Add(antigravityThrottleHalfLife)), 0.5, "penalty at one half-life") + antigravityPenaltyClose(t, acc.AntigravityPenalty(base.Add(2*antigravityThrottleHalfLife)), 0.25, "penalty at two half-lives") + // 窗口之外直接归零，不做「无限尾巴」的长尾惩罚。 + if got := acc.AntigravityPenalty(base.Add(antigravityThrottleWindow)); got != 0 { + t.Fatalf("penalty outside the window = %v, want 0", got) + } + if got := acc.AntigravityPenalty(base.Add(time.Hour)); got != 0 { + t.Fatalf("penalty one hour later = %v, want 0", got) + } + // 记录时刻之后的时间戳按 0 处理，不能让时钟回拨放大惩罚。 + if got := acc.AntigravityPenalty(base.Add(-time.Minute)); got != 1 { + t.Fatalf("penalty before the record = %v, want the un-decayed 1.0", got) + } + + acc.ResetAntigravityPenalty() + if got := acc.AntigravityPenalty(base); got != 0 { + t.Fatalf("penalty after reset = %v, want 0", got) + } +} + +func TestAntigravityPenaltyAccumulatesAndCaps(t *testing.T) { + base := time.Date(2026, 10, 17, 12, 0, 0, 0, time.UTC) + acc := &Account{DBID: 8802, AccessToken: "token"} + + RecordAntigravityRateLimitEvent(acc, AntigravityRateLimitEvent{Weight: 1.0, ObservedAt: base}) + RecordAntigravityRateLimitEvent(acc, AntigravityRateLimitEvent{Weight: 1.0, ObservedAt: base.Add(10 * time.Second)}) + + want := 1*antigravityThrottleDecay(10*time.Second) + 1 + antigravityPenaltyClose(t, acc.AntigravityPenalty(base.Add(10*time.Second)), want, "accumulated penalty") + if want <= 1 || want >= 2 { + t.Fatalf("accumulated penalty %v did not decay the previous record", want) + } + + // 距上次记录超过整个窗口后，旧累加值完全衰减，只留下本次事件。 + RecordAntigravityRateLimitEvent(acc, AntigravityRateLimitEvent{Weight: 0.5, ObservedAt: base.Add(10 * time.Second)}) + RecordAntigravityRateLimitEvent(acc, AntigravityRateLimitEvent{Weight: 1.0, ObservedAt: base.Add(10*time.Second + antigravityThrottleWindow)}) + antigravityPenaltyClose(t, acc.AntigravityPenalty(base.Add(10*time.Second+antigravityThrottleWindow)), 1.0, "penalty after the window expired") + + // 累加封顶 4.0：再多的 429 也不会把惩罚推成无界分数。 + for i := 0; i < 20; i++ { + RecordAntigravityRateLimitEvent(acc, AntigravityRateLimitEvent{Weight: 1.0, ObservedAt: base.Add(10*time.Second + antigravityThrottleWindow + time.Duration(i)*time.Millisecond)}) + } + if got := acc.AntigravityPenalty(base.Add(10*time.Second + antigravityThrottleWindow + 19*time.Millisecond)); got > antigravityThrottleMaxPenalty+1e-9 || got < 3.5 { + t.Fatalf("penalty = %v, want it to saturate at %v", got, antigravityThrottleMaxPenalty) + } +} + +func TestRecordAntigravityRateLimitEventWeights(t *testing.T) { + base := time.Date(2026, 10, 17, 12, 0, 0, 0, time.UTC) + + // 共享容量与账号配额不记短期惩罚。 + for _, scope := range []AntigravityRateLimitScope{AGScopeSharedCapacity, AGScopeAccountQuota} { + acc := &Account{DBID: 8810, AccessToken: "token"} + RecordAntigravityRateLimitEvent(acc, AntigravityRateLimitEvent{Scope: scope, Weight: AntigravityPenaltyWeightNone, ObservedAt: base}) + if got := acc.AntigravityPenalty(base); got != 0 { + t.Fatalf("scope %s penalty = %v, want 0", scope, got) + } + } + + // 裸 RESOURCE_EXHAUSTED 记半个惩罚，明确限流记一个。 + bare := &Account{DBID: 8811, AccessToken: "token"} + RecordAntigravityRateLimitEvent(bare, AntigravityRateLimitEvent{Scope: AGScopeModelRateLimit, Weight: AntigravityPenaltyWeightUnclassifiedRateLimit, ObservedAt: base}) + antigravityPenaltyClose(t, bare.AntigravityPenalty(base), 0.5, "bare RESOURCE_EXHAUSTED penalty") + + explicit := &Account{DBID: 8812, AccessToken: "token"} + RecordAntigravityRateLimitEvent(explicit, AntigravityRateLimitEvent{Scope: AGScopeModelRateLimit, Weight: AntigravityPenaltyWeightExplicitRateLimit, ObservedAt: base}) + antigravityPenaltyClose(t, explicit.AntigravityPenalty(base), 1.0, "explicit rate limit penalty") + + // 权重越界被夹到 [0, 1]，NaN 与负值直接丢弃。 + clamped := &Account{DBID: 8813, AccessToken: "token"} + RecordAntigravityRateLimitEvent(clamped, AntigravityRateLimitEvent{Weight: 9, ObservedAt: base}) + antigravityPenaltyClose(t, clamped.AntigravityPenalty(base), 1.0, "clamped weight") + + garbage := &Account{DBID: 8814, AccessToken: "token"} + RecordAntigravityRateLimitEvent(garbage, AntigravityRateLimitEvent{Weight: math.NaN(), ObservedAt: base}) + RecordAntigravityRateLimitEvent(garbage, AntigravityRateLimitEvent{Weight: -1, ObservedAt: base}) + if got := garbage.AntigravityPenalty(base); got != 0 { + t.Fatalf("invalid weights produced penalty %v, want 0", got) + } + + // 缺省 ObservedAt 落到当前时刻；nil 账号不 panic。 + fallback := &Account{DBID: 8815, AccessToken: "token"} + RecordAntigravityRateLimitEvent(fallback, AntigravityRateLimitEvent{Weight: 1.0}) + antigravityPenaltyClose(t, fallback.AntigravityPenalty(time.Now()), 1.0, "penalty with a default observation time") + RecordAntigravityRateLimitEvent(nil, AntigravityRateLimitEvent{Weight: 1.0}) + var nilAccount *Account + if got := nilAccount.AntigravityPenalty(time.Now()); got != 0 { + t.Fatalf("nil account penalty = %v, want 0", got) + } +} + +func TestAntigravityBalancedAffinitySwitch(t *testing.T) { + previous := GetAntigravityBalancedAffinity() + t.Cleanup(func() { SetAntigravityBalancedAffinity(previous) }) + + SetAntigravityBalancedAffinity(false) + if GetAntigravityBalancedAffinity() { + t.Fatal("switch stayed on after SetAntigravityBalancedAffinity(false)") + } + SetAntigravityBalancedAffinity(true) + if !GetAntigravityBalancedAffinity() { + t.Fatal("switch stayed off after SetAntigravityBalancedAffinity(true)") + } + + for value, want := range map[string]bool{ + "": true, + "1": true, + "true": true, + "on": true, + "weird": true, + "0": false, + "false": false, + "off": false, + "none": false, + "no": false, + "OFF": false, + } { + t.Setenv(antigravityBalancedAffinityEnv, value) + ReloadAntigravityBalancedAffinityFromEnv() + if got := GetAntigravityBalancedAffinity(); got != want { + t.Fatalf("ANTIGRAVITY_BALANCED_AFFINITY=%q -> %v, want %v", value, got, want) + } + } +} + +func newAntigravityRateLimitTestStore(t *testing.T, tokenCache cache.TokenCache, accounts ...*Account) *Store { + t.Helper() + store := NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4, TestConcurrency: 1, TestModel: "gpt-5.4"}) + t.Cleanup(store.Stop) + for _, acc := range accounts { + store.AddAccount(acc) + } + return store +} + +// TestWithAntigravityCooldownFilterContextSyncsSharedCooldowns 覆盖多实例部署： +// 实例 A 写入的 wireModel / Gemini family 冷却，实例 B 只按 publicModel 过滤时看 +// 不见，必须通过三键过滤器读到并同步进本地账号。 +func TestWithAntigravityCooldownFilterContextSyncsSharedCooldowns(t *testing.T) { + tokenCache := cache.NewMemory(8) + defer tokenCache.Close() + + account := &Account{DBID: 8820, UpstreamType: UpstreamAntigravity, AccessToken: "token", Status: StatusReady} + writing := newAntigravityRateLimitTestStore(t, tokenCache, account) + reading := newAntigravityRateLimitTestStore(t, tokenCache) + + writing.MarkModelCooldownUntil(account, "gemini-3.8-flash-tiered", "rate_limited_model", time.Now().Add(time.Minute)) + + // 抹掉本地内存态，模拟该账号在实例 B 上还没有任何冷却记录。 + account.mu.Lock() + account.ModelCooldowns = nil + account.mu.Unlock() + if account.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("local cooldown was not cleared") + } + + // 只按 publicModel 过滤：另一实例写入的 wire 冷却不可见，账号仍会被选中。 + plain := reading.WithAntigravityCooldownFilterContext(context.Background(), "gemini-3.8-flash-low", "", "", nil) + if !plain(account) { + t.Fatal("a public-model-only filter must pass the candidate") + } + + // 带上 wireModel 的三键过滤：命中另一实例的冷却，账号被过滤掉。 + filter := reading.WithAntigravityCooldownFilterContext(context.Background(), "gemini-3.8-flash-low", "gemini-3.8-flash-tiered", antigravityGeminiFamilyKey, nil) + if filter(account) { + t.Fatal("the wire-model cooldown written by the other instance must filter the candidate out") + } + if !account.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("the remote cooldown was not synchronized into the local account") + } + + // 下游 filter 仍然生效；nil 候选被拒绝。 + downstream := &Account{DBID: 8821, UpstreamType: UpstreamAntigravity, AccessToken: "token", Status: StatusReady} + rejecting := reading.WithAntigravityCooldownFilterContext(context.Background(), "gemini-3.8-flash-low", "gemini-3.8-flash-tiered", "", func(*Account) bool { return false }) + if rejecting(downstream) { + t.Fatal("the downstream filter must still be able to reject a candidate") + } + if filter(nil) { + t.Fatal("a nil candidate must be rejected") + } + + // 取消的上下文按 fail-closed 处理，不能在读缓存失败时放行。 + canceled, cancel := context.WithCancel(context.Background()) + cancel() + canceledFilter := reading.WithAntigravityCooldownFilterContext(canceled, "gemini-3.8-flash-low", "gemini-3.8-flash-tiered", "", nil) + if canceledFilter(&Account{DBID: 8822, AccessToken: "token"}) { + t.Fatal("a canceled context must reject the candidate") + } + + // 三个键都为空时等价于透传下游 filter。 + empty := reading.WithAntigravityCooldownFilterContext(context.Background(), "", "", "", func(*Account) bool { return true }) + if !empty(&Account{DBID: 8823, AccessToken: "token"}) { + t.Fatal("an empty key set must pass the candidate through") + } +} + +// TestWithAntigravityAwareCooldownFilterContextSplitsByAccountKind 覆盖 auto 渠道 +// （渠道未限定）的候选感知：Antigravity 候选走 public/wire/family 三键检查并 hydrate +// 远端冷却，非 Antigravity 候选只看 publicModel，与 WithModelCooldownFilterContext 的 +// 单键语义一致；下游 filter 对每个候选只调用一次；ctx 取消 fail-closed。 +func TestWithAntigravityAwareCooldownFilterContextSplitsByAccountKind(t *testing.T) { + tokenCache := cache.NewMemory(8) + defer tokenCache.Close() + + antigravity := &Account{DBID: 8830, UpstreamType: UpstreamAntigravity, AccessToken: "token", Status: StatusReady} + other := &Account{DBID: 8831, AccessToken: "token", Status: StatusReady} + writing := newAntigravityRateLimitTestStore(t, tokenCache, antigravity, other) + reading := newAntigravityRateLimitTestStore(t, tokenCache) + + // 另一实例只在共享 wire 模型上写了冷却：两个账号本地都没有任何冷却记录。 + writing.MarkModelCooldownUntil(antigravity, "gemini-3.8-flash-tiered", "rate_limited_model", time.Now().Add(time.Minute)) + writing.MarkModelCooldownUntil(other, "gemini-3.8-flash-tiered", "rate_limited_model", time.Now().Add(time.Minute)) + for _, acc := range []*Account{antigravity, other} { + acc.mu.Lock() + acc.ModelCooldowns = nil + acc.mu.Unlock() + } + + aware := reading.WithAntigravityAwareCooldownFilterContext(context.Background(), "gemini-3.8-flash-low", "gemini-3.8-flash-tiered", antigravityGeminiFamilyKey, nil) + + // Antigravity 候选：wire 冷却可见并拦截，且同步进本地内存。 + if aware(antigravity) { + t.Fatal("the aware filter must apply the three-key check to an Antigravity candidate") + } + if !antigravity.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("the aware filter did not hydrate the remote wire cooldown into the Antigravity account") + } + + // 非 Antigravity 候选：只看 publicModel，wire 冷却不可见 → 放行。 + if !aware(other) { + t.Fatal("a non-Antigravity candidate must keep the single public-model semantics") + } + + // 家族冷却同样只作用于 Antigravity 候选。 + familyOnly := &Account{DBID: 8832, AccessToken: "token", Status: StatusReady} + writing.MarkModelCooldownUntil(familyOnly, antigravityGeminiFamilyKey, "rate_limited_model", time.Now().Add(time.Minute)) + familyOnly.mu.Lock() + familyOnly.ModelCooldowns = nil + familyOnly.mu.Unlock() + if !aware(familyOnly) { + t.Fatal("a non-Antigravity candidate must not see the Antigravity family cooldown") + } + antigravityFamilyOnly := &Account{DBID: 8833, UpstreamType: UpstreamAntigravity, AccessToken: "token", Status: StatusReady} + writing.MarkModelCooldownUntil(antigravityFamilyOnly, antigravityGeminiFamilyKey, "rate_limited_model", time.Now().Add(time.Minute)) + antigravityFamilyOnly.mu.Lock() + antigravityFamilyOnly.ModelCooldowns = nil + antigravityFamilyOnly.mu.Unlock() + if aware(antigravityFamilyOnly) { + t.Fatal("the aware filter must apply the Gemini family cooldown to an Antigravity candidate") + } + + // 下游 filter 对每个候选只调用一次（两个分支共用同一次求值）。 + calls := map[int64]int{} + counting := reading.WithAntigravityAwareCooldownFilterContext(context.Background(), "gemini-3.8-flash-low", "gemini-3.8-flash-tiered", antigravityGeminiFamilyKey, func(acc *Account) bool { + calls[acc.DBID]++ + return true + }) + mixed := []*Account{ + {DBID: 8834, UpstreamType: UpstreamAntigravity, AccessToken: "token", Status: StatusReady}, + {DBID: 8835, AccessToken: "token", Status: StatusReady}, + } + for _, acc := range mixed { + if !counting(acc) { + t.Fatalf("the downstream filter rejected candidate %d unexpectedly", acc.DBID) + } + if calls[acc.DBID] != 1 { + t.Fatalf("the downstream filter ran %d times for candidate %d, want 1", calls[acc.DBID], acc.DBID) + } + } + if counting(nil) { + t.Fatal("a nil candidate must be rejected") + } + + // ctx 取消 fail-closed。 + canceled, cancel := context.WithCancel(context.Background()) + cancel() + canceledFilter := reading.WithAntigravityAwareCooldownFilterContext(canceled, "gemini-3.8-flash-low", "gemini-3.8-flash-tiered", "", nil) + if canceledFilter(&Account{DBID: 8836, AccessToken: "token"}) { + t.Fatal("a canceled context must reject the candidate") + } + + // publicModel 为空时：Antigravity 候选仍按 wire 检查，非 Antigravity 候选无键可查 → 放行。 + agNoPublic := &Account{DBID: 8837, UpstreamType: UpstreamAntigravity, AccessToken: "token", Status: StatusReady} + writing.MarkModelCooldownUntil(agNoPublic, "gemini-3.8-flash-tiered", "rate_limited_model", time.Now().Add(time.Minute)) + agNoPublic.mu.Lock() + agNoPublic.ModelCooldowns = nil + agNoPublic.mu.Unlock() + publicless := reading.WithAntigravityAwareCooldownFilterContext(context.Background(), "", "gemini-3.8-flash-tiered", "", nil) + if publicless(agNoPublic) { + t.Fatal("an Antigravity candidate must still be checked against the wire key without a public model") + } + if !publicless(&Account{DBID: 8838, AccessToken: "token", Status: StatusReady}) { + t.Fatal("a non-Antigravity candidate must pass when there is no public model to check") + } +} diff --git a/auth/store.go b/auth/store.go index 88b650d12..34d0c11b0 100644 --- a/auth/store.go +++ b/auth/store.go @@ -194,11 +194,17 @@ type Account struct { antigravityQuotaObservedAt time.Time antigravityQuotaValid bool antigravityModelMaxOutput map[string]int - BaseURL string - APIKey string - Models []string - ModelMapping string - CodexClientMetadataMode string + // antigravityPenalty 是 Antigravity 首绑/Fresh 选号避险使用的近期限流惩罚： + // 进程内、指数衰减、不落库（见 antigravity_ratelimit.go）。它只影响「下一个 + // 新会话优先选谁」，不参与冷却判定，也不改变任何既有绑定语义。 + // antigravityPenaltyAt 是最近一次记录惩罚的时刻。两者均受 a.mu 保护。 + antigravityPenalty float64 + antigravityPenaltyAt time.Time + BaseURL string + APIKey string + Models []string + ModelMapping string + CodexClientMetadataMode string // CodexPassthroughMode 是 OpenAI Responses 中转账号的 Codex 身份透传档位 // （off / auto / always），见 codex passthrough 常量定义。 CodexPassthroughMode string @@ -3097,19 +3103,22 @@ func normalizeModelCooldownKey(model string) string { return strings.ToLower(strings.TrimSpace(model)) } -func (a *Account) SetModelCooldownUntil(model, reason string, resetAt time.Time) ModelCooldown { +// normalizeModelCooldownArgs 归一化冷却键与原因。key 为空或 resetAt 为零值时返回 +// ok=false，调用方保持既有「什么都不写」的语义。 +func normalizeModelCooldownArgs(model, reason string, resetAt time.Time) (string, string, bool) { key := normalizeModelCooldownKey(model) if key == "" || resetAt.IsZero() { - return ModelCooldown{} + return "", "", false } reason = strings.TrimSpace(reason) if reason == "" { reason = "rate_limited" } - now := time.Now() + return key, reason, true +} - a.mu.Lock() - defer a.mu.Unlock() +// setModelCooldownUntilLocked 写入冷却并返回写入结果；调用方必须持有 a.mu 写锁。 +func (a *Account) setModelCooldownUntilLocked(key, reason string, resetAt, now time.Time) ModelCooldown { if a.ModelCooldowns == nil { a.ModelCooldowns = make(map[string]ModelCooldown) } @@ -3132,6 +3141,38 @@ func (a *Account) SetModelCooldownUntil(model, reason string, resetAt time.Time) return cooldown } +func (a *Account) SetModelCooldownUntil(model, reason string, resetAt time.Time) ModelCooldown { + key, reason, ok := normalizeModelCooldownArgs(model, reason, resetAt) + if !ok { + return ModelCooldown{} + } + + a.mu.Lock() + defer a.mu.Unlock() + return a.setModelCooldownUntilLocked(key, reason, resetAt, time.Now()) +} + +// SetModelCooldownUntilAtLeast 与 SetModelCooldownUntil 相同，但保证不缩短仍然 +// 有效的既有冷却：该 model 现有 ResetAt 仍在生效且晚于传入 resetAt 时保留既有冷却， +// 第二个返回值为 false 表示本次未写入。 +// +// 「读现有 ResetAt → 比较 → 写入」在同一次 a.mu 写锁临界区内完成：否则两个并发 +// 写入都会读到同一个旧值，后落盘的短 deadline 会把先写入的长 deadline 覆盖掉。 +func (a *Account) SetModelCooldownUntilAtLeast(model, reason string, resetAt time.Time) (ModelCooldown, bool) { + key, reason, ok := normalizeModelCooldownArgs(model, reason, resetAt) + if !ok { + return ModelCooldown{}, false + } + now := time.Now() + + a.mu.Lock() + defer a.mu.Unlock() + if existing, exists := a.ModelCooldowns[key]; exists && existing.ResetAt.After(now) && existing.ResetAt.After(resetAt) { + return existing, false + } + return a.setModelCooldownUntilLocked(key, reason, resetAt, now), true +} + func (a *Account) RestoreModelCooldown(model, reason string, resetAt, updatedAt time.Time) { key := normalizeModelCooldownKey(model) if key == "" || resetAt.IsZero() || !resetAt.After(time.Now()) { @@ -3183,6 +3224,26 @@ func (a *Account) ModelCooldownRemaining(model string) time.Duration { return remaining } +// ModelCooldownFor 返回账号上某个 model 当前仍在生效的冷却条目（含 ResetAt 与 +// Reason）；不存在或已过期时第二个返回值为 false。 +// +// 与 ModelCooldownRemaining 相比，它把 ResetAt 一起交出，调用方无需自己用 +// 「now + remaining」反推 deadline，因此可以在改写共享冷却键前拿到既有的精确 +// deadline（例如判断本次较短的冷却是否会把更长的共享冷却改短）。 +func (a *Account) ModelCooldownFor(model string) (ModelCooldown, bool) { + key := normalizeModelCooldownKey(model) + if key == "" { + return ModelCooldown{}, false + } + a.mu.RLock() + defer a.mu.RUnlock() + cooldown, ok := a.ModelCooldowns[key] + if !ok || !cooldown.ResetAt.After(time.Now()) { + return ModelCooldown{}, false + } + return cooldown, true +} + func (a *Account) ActiveModelCooldowns() []ModelCooldown { a.mu.RLock() defer a.mu.RUnlock() @@ -4190,6 +4251,120 @@ func (s *Store) WithModelCooldownFilterContext(ctx context.Context, model string } } +// WithAntigravityCooldownFilterContext 是 Antigravity 渠道的跨实例模型冷却过滤。 +// +// Antigravity 的一个公开档位（gemini-3.8-flash-high）可能与另一个档位 +// （gemini-3.8-flash-low）共用同一个上游 wire 模型（gemini-3.8-flash-tiered）， +// 而 Gemini 配额又按整个 family 计量。多实例部署时，另一实例写入的 wireModel 或 +// family 冷却在本实例的内存里不可见，只按 publicModel 过滤就会把已被上游限流的 +// 账号再次选中。这里对候选账号依次检查 publicModel、wireModel 与 familyKey 三个 +// 键的共享缓存冷却，命中即过滤，并把读到的远端冷却同步进本地账号。 +// +// familyKey 由调用方决定：非 Gemini 模型传空串即可。ctx 取消时按「拒绝候选」处理， +// 与 WithModelCooldownFilterContext 保持同一 fail-closed 语义。 +func (s *Store) WithAntigravityCooldownFilterContext(ctx context.Context, publicModel, wireModel, familyKey string, filter AccountFilter) AccountFilter { + if ctx == nil { + ctx = context.Background() + } + keys := antigravityCooldownFilterKeys(publicModel, wireModel, familyKey) + return func(acc *Account) bool { + if acc == nil || ctx.Err() != nil { + return false + } + if filter != nil && !filter(acc) { + return false + } + if ctx.Err() != nil { + return false + } + if s == nil || len(keys) == 0 { + return true + } + for _, key := range keys { + if s.accountHasCachedModelCooldownContext(ctx, acc, key) { + return false + } + if ctx.Err() != nil { + return false + } + } + return true + } +} + +// antigravityCooldownFilterKeys 把 publicModel / wireModel / familyKey 归一化成去重后 +// 的冷却键列表（保持该顺序），空键与重复键被丢弃。 +func antigravityCooldownFilterKeys(publicModel, wireModel, familyKey string) []string { + keys := make([]string, 0, 3) + seen := make(map[string]struct{}, 3) + for _, model := range []string{publicModel, wireModel, familyKey} { + key := normalizeModelCooldownKey(model) + if key == "" { + continue + } + if _, dup := seen[key]; dup { + continue + } + seen[key] = struct{}{} + keys = append(keys, key) + } + return keys +} + +// WithAntigravityAwareCooldownFilterContext 是「渠道未限定（auto Key）」请求的冷却过滤。 +// +// 下游 Key 没有指定上游渠道时，候选池里既可能有 Antigravity 账号（native Gemini +// generateContent/countTokens 被 antigravityOAuthChannelAccountFilter 硬限定为 +// Antigravity OAuth，Responses / Chat / Messages 也允许 Antigravity 账号），也可能有 +// 其它渠道的账号，而请求侧无法预先知道最终会落到哪一类。这里按候选账号自身的类型选择 +// 检查范围： +// - IsAntigravityAPI() 的候选：与 WithAntigravityCooldownFilterContext 完全一致的 +// publicModel / wireModel / familyKey 三键检查（去重），闭合「另一实例写入的 wire +// 或 family 冷却在本实例选号时看不见」的缺口； +// - 其它候选：与 WithModelCooldownFilterContext 完全一致的单一 publicModel 检查， +// 行为不变。 +// +// 下游 filter 对每个候选只调用一次（两个分支共用同一次求值）；ctx 取消时按「拒绝候选」 +// 处理，与上面两个过滤器保持同一 fail-closed 语义。 +func (s *Store) WithAntigravityAwareCooldownFilterContext(ctx context.Context, publicModel, wireModel, familyKey string, filter AccountFilter) AccountFilter { + if ctx == nil { + ctx = context.Background() + } + antigravityKeys := antigravityCooldownFilterKeys(publicModel, wireModel, familyKey) + publicKey := normalizeModelCooldownKey(publicModel) + return func(acc *Account) bool { + if acc == nil || ctx.Err() != nil { + return false + } + if filter != nil && !filter(acc) { + return false + } + if ctx.Err() != nil { + return false + } + if s == nil { + return true + } + keys := antigravityKeys + if !acc.IsAntigravityAPI() { + // 非 Antigravity 候选只看 publicModel，与既有单键语义完全一致。 + if publicKey == "" { + return true + } + keys = []string{publicKey} + } + for _, key := range keys { + if s.accountHasCachedModelCooldownContext(ctx, acc, key) { + return false + } + if ctx.Err() != nil { + return false + } + } + return true + } +} + func fastSchedulerEnabledFromEnv() bool { for _, key := range []string{"FAST_SCHEDULER_ENABLED", "CODEX_FAST_SCHEDULER"} { if truthyEnv(os.Getenv(key)) { @@ -7099,13 +7274,13 @@ func (s *Store) NextForSession(key string, apiKeyID int64, exclude map[int64]boo // 解除发生时绕过 binding 走完整挑号策略(NextExcludingWithFilter),后续 BindSessionAffinity // 会重新建立绑定。 func (s *Store) NextForSessionWithFilter(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter) (*Account, string) { - account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, false, DispatchPolicyStandard) + account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, false, SelectionOptions{Policy: DispatchPolicyStandard}) return account, proxyURL } // NextForSessionWithDispatch 优先复用绑定账号，并按用量策略选号。 func (s *Store) NextForSessionWithDispatch(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) (*Account, string) { - account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, false, policy) + account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, false, SelectionOptions{Policy: policy}) return account, proxyURL } @@ -7113,7 +7288,7 @@ func (s *Store) NextForSessionWithDispatch(key string, apiKeyID int64, exclude m // request paths. The returned guard must be passed to BindSessionAffinityWithGuard // after the attempt is selected or committed. func (s *Store) NextForSessionWithDispatchGuard(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) (*Account, string, SessionAffinityGuard) { - return s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, false, policy) + return s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, false, SelectionOptions{Policy: policy}) } // NextForContinuationWithFilter preserves an existing account binding for a @@ -7121,13 +7296,13 @@ func (s *Store) NextForSessionWithDispatchGuard(key string, apiKeyID int64, excl // to the account that created them, so bounded-affinity escape must not rotate // accounts. func (s *Store) NextForContinuationWithFilter(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter) (*Account, string) { - account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, true, DispatchPolicyStandard) + account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, true, SelectionOptions{Policy: DispatchPolicyStandard}) return account, proxyURL } // NextForContinuationWithDispatch preserves a bound turn and applies a usage policy. func (s *Store) NextForContinuationWithDispatch(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) (*Account, string) { - account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, true, policy) + account, proxyURL, _ := s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, true, SelectionOptions{Policy: policy}) return account, proxyURL } @@ -7141,10 +7316,11 @@ func (s *Store) NextForContinuationWithDispatch(key string, apiKeyID int64, excl // // 绑定本身不存在时仍走完整挑号，与普通请求一致；TTL 过期只影响普通请求， // preserveBinding=true 的续链请求仍保留原账号。 -func (s *Store) nextForSessionWithFilter(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, preserveBinding bool, policy DispatchPolicy) (*Account, string, SessionAffinityGuard) { +func (s *Store) nextForSessionWithFilter(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, preserveBinding bool, opts SelectionOptions) (*Account, string, SessionAffinityGuard) { if s == nil { return nil, "", SessionAffinityGuard{} } + policy := opts.Policy key = strings.TrimSpace(key) if key == "" { return s.NextExcludingWithDispatch(apiKeyID, exclude, filter, policy), "", SessionAffinityGuard{} @@ -7207,7 +7383,7 @@ func (s *Store) nextForSessionWithFilter(key string, apiKeyID int64, exclude map return nil, "", SessionAffinityGuard{} } if capacityFull { - fallback := s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, policy) + fallback := s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, opts) if fallback == nil { return nil, "", SessionAffinityGuard{} } @@ -7222,7 +7398,7 @@ func (s *Store) nextForSessionWithFilter(key string, apiKeyID int64, exclude map return s.takeByIDForContinuation(binding.accountID, apiKeyID, exclude, filter, key, policy), "", SessionAffinityGuard{} } s.UnbindSessionAffinity(key, binding.accountID) - return s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, policy), "", SessionAffinityGuard{} + return s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, opts), "", SessionAffinityGuard{} } // 跨进程缓存的 binding 也按 bounded 逻辑校验账号健康；Grok 账号套用 Grok 专属模式。 cacheMode := mode @@ -7246,7 +7422,7 @@ func (s *Store) nextForSessionWithFilter(key string, apiKeyID int64, exclude map return nil, "", SessionAffinityGuard{} } if capacityFull { - fallback := s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, policy) + fallback := s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, opts) if fallback == nil { return nil, "", SessionAffinityGuard{} } @@ -7256,7 +7432,7 @@ func (s *Store) nextForSessionWithFilter(key string, apiKeyID int64, exclude map } } - return s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, policy), "", SessionAffinityGuard{} + return s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, opts), "", SessionAffinityGuard{} } // nextAccountForFreshAffinity 为"新亲和键首次绑定"选号(issue #484)。 @@ -7269,13 +7445,22 @@ func (s *Store) nextForSessionWithFilter(key string, apiKeyID int64, exclude map // 下一名,账号增删只迁移受影响的键。层间仍严格尊重运营者设置的调度优先级 // 与健康档位,层内才忽略分数差异。 func (s *Store) nextAccountForFreshAffinity(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter) *Account { - return s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, DispatchPolicyStandard) + return s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, SelectionOptions{Policy: DispatchPolicyStandard}) } -func (s *Store) nextAccountForFreshAffinityWithDispatch(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) *Account { +func (s *Store) nextAccountForFreshAffinityWithDispatch(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, opts SelectionOptions) *Account { if s == nil { return nil } + policy := opts.Policy + // Antigravity 首绑 / Fresh 加权避险：handled=true 表示加权路径已接管本次选号（含号池 + // 无可用容量返回 nil 的情况），此处必须直接返回，不能再走下面的 spread / HRW / + // round-robin 重扫——那会对同一批账号二次求值 filter，破坏「单次选号每个账号 filter + // 至多一次」的约束。handled=false（未开启 / 非 Antigravity / 非 round_robin）时才 + // 平滑回退到既有散列逻辑。 + if acc, handled := s.selectAntigravityBalancedAccount(key, apiKeyID, exclude, filter, opts); handled { + return acc + } if !s.GetSessionAffinitySpread() || strings.TrimSpace(key) == "" { return s.NextExcludingWithDispatch(apiKeyID, exclude, filter, policy) } @@ -7844,44 +8029,45 @@ func (s *Store) hasContinuationCandidateWithDispatch(key string, apiKeyID int64, // WaitForSessionAvailableWithFilter waits for an account that satisfies the request-level filter. func (s *Store) WaitForSessionAvailableWithFilter(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter) (*Account, string) { - account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, false, DispatchPolicyStandard) + account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, false, SelectionOptions{Policy: DispatchPolicyStandard}) return account, proxyURL } func (s *Store) WaitForSessionAvailableWithDispatch(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) (*Account, string) { - account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, false, policy) + account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, false, SelectionOptions{Policy: policy}) return account, proxyURL } // WaitForSessionAvailableWithDispatchGuard is the binding-aware waiting path. // It preserves the capacity-spillover decision made by the successful retry. func (s *Store) WaitForSessionAvailableWithDispatchGuard(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) (*Account, string, SessionAffinityGuard) { - account, proxyURL, guard, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, false, policy) + account, proxyURL, guard, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, false, SelectionOptions{Policy: policy}) return account, proxyURL, guard } // WaitForContinuationAvailableWithFilter waits for the account already bound // to a stateful continuation instead of falling through to another account. func (s *Store) WaitForContinuationAvailableWithFilter(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter) (*Account, string) { - account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, true, DispatchPolicyStandard) + account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, true, SelectionOptions{Policy: DispatchPolicyStandard}) return account, proxyURL } func (s *Store) WaitForContinuationAvailableWithDispatch(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) (*Account, string) { - account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, true, policy) + account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, true, SelectionOptions{Policy: policy}) return account, proxyURL } // WaitForDispatchAvailable exposes queue admission failures to protocol handlers. // preserveBinding retains the same account-only semantics as continuation waits. func (s *Store) WaitForDispatchAvailable(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, preserveBinding bool, policy DispatchPolicy, heartbeat ...SchedulerWaitHeartbeat) (*Account, string, SessionAffinityGuard, error) { - return s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, preserveBinding, policy, heartbeat...) + return s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, preserveBinding, SelectionOptions{Policy: policy}, heartbeat...) } -func (s *Store) waitForSessionAvailableWithFilter(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, preserveBinding bool, policy DispatchPolicy, heartbeat ...SchedulerWaitHeartbeat) (*Account, string, SessionAffinityGuard, error) { +func (s *Store) waitForSessionAvailableWithFilter(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, preserveBinding bool, opts SelectionOptions, heartbeat ...SchedulerWaitHeartbeat) (*Account, string, SessionAffinityGuard, error) { if ctx == nil { ctx = context.Background() } + policy := opts.Policy if timeout <= 0 || ctx.Err() != nil { return nil, "", SessionAffinityGuard{}, ctx.Err() } @@ -7978,11 +8164,7 @@ func (s *Store) waitForSessionAvailableWithFilter(ctx context.Context, key strin var acc *Account var proxyURL string var guard SessionAffinityGuard - if preserveBinding { - acc, proxyURL = s.NextForContinuationWithDispatch(key, apiKeyID, exclude, filter, policy) - } else { - acc, proxyURL, guard = s.NextForSessionWithDispatchGuard(key, apiKeyID, exclude, filter, policy) - } + acc, proxyURL, guard = s.nextForSessionWithFilter(key, apiKeyID, exclude, filter, preserveBinding, opts) if acc != nil { if ctx.Err() != nil || !time.Now().Before(expires) { s.Release(acc) @@ -10265,6 +10447,33 @@ func (s *Store) MarkModelCooldownUntil(acc *Account, model, reason string, reset if cooldown.Model == "" { return cooldown } + s.afterModelCooldownMarked(acc, cooldown) + return cooldown +} + +// MarkModelCooldownUntilAtLeast 与 MarkModelCooldownUntil 相同，但不会缩短仍然 +// 有效的既有冷却：该 model 现有 ResetAt 晚于传入 resetAt 时保留既有冷却并直接返回， +// 不写内存、不刷健康度、不落库。 +// +// 用于多个公开档位共享同一上游 backing（wire 模型）或同一配额家族（Gemini family） +// 的冷却键：这类键会被同账号上的多个事件写入，若按最后写入者胜出，一个 5s 的裸 +// RESOURCE_EXHAUSTED 就能把同账号刚写入的 290s 共享冷却改短，短 deadline 之后兄弟 +// 档位又能选中刚被限流的凭据。比较与写入在 Account 侧同一次写锁内完成。 +func (s *Store) MarkModelCooldownUntilAtLeast(acc *Account, model, reason string, resetAt time.Time) ModelCooldown { + if acc == nil || resetAt.IsZero() || !resetAt.After(time.Now()) { + return ModelCooldown{} + } + cooldown, wrote := acc.SetModelCooldownUntilAtLeast(model, reason, resetAt) + if cooldown.Model == "" || !wrote { + return cooldown + } + s.afterModelCooldownMarked(acc, cooldown) + return cooldown +} + +// afterModelCooldownMarked 在冷却写入后统一刷新健康度、调度快照、跨实例冷却缓存 +// 与持久化。 +func (s *Store) afterModelCooldownMarked(acc *Account, cooldown ModelCooldown) { acc.mu.Lock() now := time.Now() acc.LastRateLimitedAt = now @@ -10284,7 +10493,6 @@ func (s *Store) MarkModelCooldownUntil(acc *Account, model, reason string, reset log.Printf("[账号 %d] 持久化模型冷却失败 model=%s: %v", acc.DBID, cooldown.Model, err) } } - return cooldown } func (s *Store) ClearModelCooldown(acc *Account, model string) { diff --git a/docs/ANTIGRAVITY.md b/docs/ANTIGRAVITY.md index c60282f17..7691ce73f 100644 --- a/docs/ANTIGRAVITY.md +++ b/docs/ANTIGRAVITY.md @@ -99,6 +99,18 @@ When creating or editing a Codex2API downstream API key, choose **Antigravity** Account groups are channel-isolated: an Antigravity account can only join an Antigravity group. +## Balanced affinity switch + +First binding and fresh session selection bias Antigravity candidates by concurrency occupancy, used quota, and the account's recent rate-limit penalty, so a credential Google just throttled sinks to the back of the queue while healthy ones keep serving new sessions. + +``` +ANTIGRAVITY_BALANCED_AFFINITY=false +``` + +The switch is on by default. `0`, `false`, `off`, `none`, and `no` (case-insensitive) turn the weighting off and restore plain round-robin; any other value keeps it on. It is read once when the process starts, so changing it requires a restart. + +Cooldown correctness is independent of this switch and always applies: the public and shared wire keys are cooled together (the Gemini-family key is added when the quota bucket is exhausted), a shared wire/family deadline is never shortened by a briefer event on the same account, and cooldowns written by another instance are honoured during selection. Entries whose candidate pool excludes Antigravity accounts (`/v1/responses/compact`, native remote compaction, the `/v1/responses` WebSocket upgrade path) never apply the weighting, even behind an Antigravity key. + ## Supported gateway endpoints Antigravity inference is admitted through `/v1/responses`, `/v1/chat/completions`, and `/v1/messages`; its models are exposed through `/v1/models` and the Codex model manifest. Chat and Messages translate their inbound body into a Responses payload before dispatch, which is exactly what the `v1internal` adapter consumes, so all three transports share one admission gate and one adapter. `/v1/responses/compact` still excludes Antigravity accounts because no compaction adapter exists, and the official Codex executors reject Antigravity credentials before any network request as a provider-boundary safeguard. diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 7f07ee404..62f01164e 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -100,6 +100,14 @@ Codex2API 采用三层配置架构： > `CODEX_UPSTREAM_TRANSPORT` 只控制 HTTP 入站请求转发到 Codex 上游时使用 `http` 还是 `ws`。客户端侧 WebSocket 入口独立可用：使用 `GET ws:///v1/responses` 建连，首帧发送 `response.create` JSON，服务端会通过 Codex 上游 WS 返回 Responses 事件帧。 +### Antigravity 调度开关 + +Antigravity 渠道的首绑 / Fresh 选号默认启用加权避险：候选账号按并发占用、已用配额与近期限流惩罚打分，刚被上游 429 的账号排到队尾，从而把新会话摊到健康账号上。 + +| 变量 | 必填 | 默认值 | 说明 | +|------|------|--------|------| +| `ANTIGRAVITY_BALANCED_AFFINITY` | 否 | `true` | 加权避险开关。取 `0` / `false` / `off` / `none` / `no`（大小写不敏感）时关闭，回退为原有 round-robin；其它值按开启处理。进程启动时读取一次，改动需重启生效。冷却正确性修复不受该开关影响、始终生效：public / wire 冷却写入（Gemini 配额耗尽时追加 family 键）、共享 wire / family 冷却不被更短事件缩短、跨实例冷却过滤（Antigravity 候选查三键） | + ### 数据库配置 #### PostgreSQL 模式 diff --git a/docs/superpowers/specs/2026-10-17-antigravity-affinity-balance-design.md b/docs/superpowers/specs/2026-10-17-antigravity-affinity-balance-design.md new file mode 100644 index 000000000..a2ba65cfe --- /dev/null +++ b/docs/superpowers/specs/2026-10-17-antigravity-affinity-balance-design.md @@ -0,0 +1,266 @@ +# Antigravity 调度优化与冷却正确性规格（收窄定稿版） + +> 状态：定稿（Phase 1/2/3 完成，复审通过）。 +> 实施记录：Phase 1 冷却正确性（canonical 双写 / countTokens / 事件累加器）；Phase 2 加权选号与跨实例过滤； +> 修复 Phase 2 审查 P1（auto Key -aware 过滤、handled 单次 filter）后复审通过；Phase 3 全量回归通过。 +> 收窄策略：仅做冷却正确性与首绑/Fresh 加权避险，不做运行时主动迁移状态机。 +> 核心策略：修复 canonical 冷却漏写/漏读真实漏洞 + 补齐 countTokens 429 链路 + 首绑与重试 Fresh 选号加权避险。 +> 架构取舍：**彻底移除运行时主动会话迁移状态机**（不做版本化 CAS、不做 pending ledger、不做 typed rejection 借用），完全保留既有 Bounded/Strict 语义与生命周期，杜绝跨实例写竞争与状态污染。 +> 代码行号基于 `main@5ee6c292`，实施时以实际代码为准。 + +## 1. 背景与根因收敛 + +### 1.1 现状与已确认事实 +- 线上配置已确认（均为默认值）：`affinity_mode=bounded`、`scheduler_mode=round_robin`、`skip_warm_tier=false`、`relay_model_cooldown_mode=off`、`session_affinity_spread=false`。 +- 请求内 HTTP 429 处理链：[Responses](proxy/handler.go#L4442-L4455)、[native Gemini](proxy/handler_gemini.go#L568-L577)、[Anthropic](proxy/handler_anthropic.go#L1167-L1177) 均在收到 429 时立即执行 `Release + UnbindSessionAffinity + retryExclusions.MarkHTTPFailure`，并触发下一轮重选。因此，**遇到 429 的请求在下一次尝试时本来就是 Fresh selection**。 +- 截图现象：错误摘要为 `rate_limited_model / 429 · Resource has been exhausted (e.g. check quota)`，持续约 122s，无首字。 +- 缓存现状：成功的 200 请求缓存命中率约 98%（cached≈176k / prompt≈178k），会话亲和正在有效命中 Gemini 上游 implicit cache，不能随意破坏活跃会话的粘性。 + +### 1.2 确认存在的真实代码漏洞 +1. **模型别名导致冷却漏写（P1-4）**： + - Antigravity 渠道中，`gemini-3.7/3.8-flash-low/medium/high` 共享相同的 upstream wire 模型 `gemini-3.7/3.8-flash-tiered`（[antigravity_models.go:50-60](proxy/antigravity_models.go#L50-L60)）。 + - 当前 `applyAntigravityCooldown`（[antigravity_quota.go:319-373](proxy/antigravity_quota.go#L319-L373)）仅针对调用方传入的 public model 写入冷却，**从未写入 wireModel 冷却**。 + - 过滤侧 `antigravityAccountModelRateLimited`（[antigravity_quota.go:282-293](proxy/antigravity_quota.go#L282-L293)）虽然检查了 wireModel，但因冷却表中没有 wireModel 记录，导致 high 被限流后，low/medium 的请求仍能选中该账号并继续触发 429。 +2. **跨实例共享缓存漏读**： + - `WithModelCooldownFilterContext`（[store.go:4170-4189](auth/store.go#L4170-L4189)）仅针对请求的 public model 查询与 hydrate 跨实例缓存；另一实例写入的 wireModel / Gemini family 冷却在当前实例不可见。 +3. **native Gemini countTokens 漏路径**： + - `handleGeminiCountTokens` 的非 200 分支（[handler_gemini.go:271-307](proxy/handler_gemini.go#L271-L307)）仅执行 `ReportFailure/Release/Unbind/exclude`，从未调用 `applyCooldownForModel`，完全漏掉了冷却写入与事件观测。 +4. **首绑与重试 Fresh 选号无负载与限流避险权重**： + - 默认配置下 `session_affinity_spread=false`，首绑走 round-robin； + - 请求内遇到 429 解绑后，或新会话到来时，选号器在同优先级层内不比对在途负载、配额用量以及近期限流历史，导致已被限流的账号在冷却短暂到期（或被跨 alias 请求）后极易再次被选中。 + +--- + +## 2. 目标与非目标 + +### 2.1 目标 +1. **修复冷却正确性**： + - 统一写入 public + wireModel（及配额耗尽时的 Gemini family）冷却； + - 跨实例共享过滤器增加对 wireModel 和 family 的查询与本地同步； + - 补齐 native countTokens 的 429/503 冷却写入。 +2. **首绑与重选加权避险**： + - 建立轻量级进程内限流惩罚累加器（`antigravityPenalty`）； + - 在新会话首绑及重选（Fresh selection）时，显式识别 `provider=antigravity`，按 `load + 0.50*usage + 0.35*penalty` 加权打分，大幅降低刚被限流账号的再次选中率。 +3. **安全与性能保障**： + - 单次选号中对候选账号的 filter 至多求值一次，杜绝测试闭包 panic 与 Redis 读放大； + - cached cooldown 检查保留在取号循环中，不执行前置 N 次 Redis 读。 + +### 2.2 非目标 +- **不做会话主动迁移状态机**：不修改 `nextForSessionWithFilter` 的绑定复用逻辑，不引入版本化 CAS、不引入复杂的分布式锁与 pending admission ledger。 +- **不覆盖 strict 模式**：完全尊重运营者的 strict 配置。 +- **不破坏现有调度模式**：`remaining_quota` 与 `fill_first` 模式保持原有用量优先语义，仅继承冷却正确性修复。 +- **不修改计费与错误码**：上游返回的错误状态码与计费逻辑保持完全一致。 + +--- + +## 3. 详细设计 + +### 3.1 统一限流事件与分类权重 + +在 `proxy/antigravity_quota.go` 与 `auth/antigravity_ratelimit.go` 中规范化 429/503 事件： + +```go +type AntigravityRateLimitScope string + +const ( + AGScopeModelRateLimit AntigravityRateLimitScope = "model_rate_limit" + AGScopeSharedCapacity AntigravityRateLimitScope = "shared_capacity" + AGScopeAccountQuota AntigravityRateLimitScope = "account_quota" + AGScopeUnknown AntigravityRateLimitScope = "unknown" +) + +type AntigravityRateLimitEvent struct { + StatusCode int + Kind string + Scope AntigravityRateLimitScope + PublicModel string + WireModel string + Family string + Weight float64 // 0.0, 0.5, 1.0 + ResetAt time.Time + ObservedAt time.Time +} +``` + +**权重规则（仅用于首绑/Fresh 选号避险打分）**： +- `RATE_LIMIT_EXCEEDED`（有明确 reason 或有 RetryInfo）：`Weight = 1.0`，`Scope = AGScopeModelRateLimit` +- 裸 `RESOURCE_EXHAUSTED`（无 reason 且无 RetryInfo，截图即此类）：`Weight = 0.5`，`Scope = AGScopeModelRateLimit` +- 非结构化 429（解析失败）：`Weight = 0.5`，`Scope = AGScopeUnknown` +- `MODEL_CAPACITY_EXHAUSTED`（共享模型容量）：`Weight = 0.0`，`Scope = AGScopeSharedCapacity`（不记惩罚） +- `QUOTA_EXHAUSTED`：`Weight = 0.0`，`Scope = AGScopeAccountQuota`（已有长冷却与 family 冷却，不记短期惩罚） +- 非结构化 503 / 未知错误：`Weight = 0.0` + +### 3.2 惩罚累加器（仅内存态） + +在 `auth.Account` 上增加轻量字段（受 `a.mu` 保护，不落库）： +```go +antigravityPenalty float64 // 衰减累加值，最大 4.0 +antigravityPenaltyAt time.Time // 最近一次记录时刻 +``` + +常量定义： +```go +const ( + antigravityThrottleWindow = 180 * time.Second + antigravityThrottleHalfLife = 60 * time.Second + antigravityThrottleMaxPenalty = 4.0 + antigravityThrottleWeight = 0.35 + antigravityUsageWeight = 0.50 +) +``` + +- 记录：`p = min(maxPenalty, oldPenalty*decay(now-at) + event.Weight)`，更新 `at = now`。 +- 读取：`p * decay(now-at)`，超过 180s 窗口直接返回 0。 + +### 3.3 Canonical 冷却与跨实例过滤修复 + +1. **写入侧（`proxy/antigravity_quota.go`）**： + - 在 `applyAntigravityCooldown` 中： + ```go + wireModel, _ := antigravityResolvePublicModelForAccount(account, model) + // 1. 写入 public model 冷却（保持既有「可被覆盖」语义） + store.MarkModelCooldownUntil(account, model, decision.Reason, decision.ResetAt) + // 2. 若 wireModel 存在且与 public model 不同，同步写入 wireModel 冷却 + // （共享 backing：只延长不缩短，避免兄弟 tier 的短事件把长冷却压短） + if wireModel != "" && wireModel != model { + store.MarkModelCooldownUntilAtLeast(account, wireModel, decision.Reason, decision.ResetAt) + } + // 3. 若为 Gemini 配额耗尽，写入 family 冷却（同样只延长不缩短） + if quota.Kind == antigravityQuotaKindQuotaExhausted && antigravityIsGeminiModel(model) { + store.MarkModelCooldownUntilAtLeast(account, antigravityGeminiFamilyCooldownKey, decision.Reason, decision.ResetAt) + } + ``` +2. **补齐 countTokens 路径（`proxy/handler_gemini.go`）**： + - 在 `handleGeminiCountTokens` 的非 200 分支中调用 `h.applyCooldownForModel(account, resp.StatusCode, errBody, resp, model)`，确保错误与事件被记录。 +3. **读取侧与跨实例同步（`auth/store.go`）**： + - 新增 `WithAntigravityCooldownFilterContext(ctx, publicModel, wireModel, familyKey, filter)`： + 对当前候选账号，依次检查 publicModel、wireModel 以及 familyKey（仅 Gemini）在共享缓存中的冷却状态并进行本地同步，任一命中冷却即过滤。 + - 新增 `WithAntigravityAwareCooldownFilterContext(...)`：供**渠道未限定（auto Key）**请求使用。 + `acc.IsAntigravityAPI()` 的候选走 public/wire/family 三键检查；其余候选只查 publicModel + （与既有 `WithModelCooldownFilterContext` 等价）。下游 filter 对每个候选只求值一次，ctx 取消 fail-closed。 + - `proxy/handler.go` 的 `withRequestModelCooldownFilter` 按渠道三档选择：antigravity→三键、 + auto（空串）→aware 版本、其他显式渠道→单键。这样 auto Key 选中 AG 账号时也能看到跨实例 + wire/family 冷却（P1-1），非 AG 候选行为不变。 + +### 3.4 首绑与重试 Fresh 选号加权均衡 + +1. **选号入口与上下文传递**： + ```go + type SelectionOptions struct { + Provider string // 显式传入，例如 database.UpstreamChannelAntigravity + Model string // public model + Policy DispatchPolicy + } + ``` + 在 proxy 的选号入口处（`nextAccountForSessionWithDispatchGuard` 等），通过 `requestUpstreamChannel(c)` 显式传入 `Provider` 与 `Model`。 + +2. **打分与排序逻辑（`auth/antigravity_balance.go`）**： + - 当 `opts.Provider == database.UpstreamChannelAntigravity` 且 `scheduler_mode == "round_robin"` 时： + - 候选按 `(schedulerPriority, tierPriority)` 从高到低分层，层内筛选符合条件的 Antigravity 候选； + - 候选打分公式（分数越小越优）： + ```go + load := float64(occupied) / float64(limit) // 0..1+ + usage := acc.usagePercentForScheduling() / 100.0 // 0..1 + penalty := acc.AntigravityPenalty(now) // 0..4 + score := load + usage*antigravityUsageWeight + penalty*antigravityThrottleWeight + ``` + - 排序：层内按 `score` 升序，同分时使用 FNV 哈希 tie-break； + - 取号：从最高层开始按顺序执行 `tryAcquireAccount(acc, limit, true)`，整层拿不下时下沉下一层，命中即返回。 + - **返回 `(*Account, handled bool)`**：`handled=true` 表示加权路径已接管本次选号（provider+开关+round_robin 成立）； + 此时即使返回 nil（号池全满/无可用候选），调用方也**不得**再回退到 spread/HRW/round-robin 重扫， + 否则会二次求值同一批账号的 filter。`handled=false` 时才平滑回退既有逻辑。 + - **native Gemini（generateContent/stream 与 countTokens）强制 `Provider=antigravity`**：这两个入口的候选池 + 本就被 `antigravityOAuthChannelAccountFilter` 硬限定为 AG OAuth，auto Key 下也必须启用加权避险； + 通用入口（Responses/Chat/Messages）在 auto Key 下 Provider 为空，不启用加权，只享受冷却过滤修复（规格边界）。 + +3. **单轮 Filter 求值保护**： + - 选号前对候选账号进行快照遍历，保证单次选号中每个账号的 `filter(acc)` 至多调用一次，避免测试中的闭包 channel panic 及生产环境中的 Redis 读重复。 + - `accountHasBlockingCachedCooldown` 的检查依然保留在遍历取号循环中，仅对即将获取的账号进行单次验证。 + - 加权路径接管后（handled=true）不回退重扫；号池饱和时返回 nil 交给上层等待/排队。**这是一次有意的行为变化**： + 原逻辑「宁可暂时聚集也不拒绝」；新逻辑优先保证 filter 单次求值，饱和时由上层 wait/queue 处理。 + +--- + +## 4. 行为变化与兼容性 + +| 场景 | 行为说明 | +| :--- | :--- | +| **已绑定且正常的会话** | 完全沿用原有绑定，100% 保护上游 prompt cache，零变更 | +| **请求内遇到 429** | 既有逻辑已解绑；随后的重试与下一轮请求走加权 Fresh 选号，自动避开刚被限流的账号 | +| **同 backing 跨 tier 请求** | high 触发冷却后，wireModel 冷却生效，low/medium 不会再误选同一受限账号 | +| **多实例部署** | 实例 A 写入的 wireModel/family 冷却，实例 B 在过滤时能即时读取并同步 | +| **auto Key（未限定渠道）** | 冷却过滤按候选账号类型分流：AG 候选查 public/wire/family 三键，非 AG 只查 public；native Gemini/countTokens 仍启用加权避险；通用入口不启用加权（规格边界） | +| **号池饱和（加权路径生效）** | 返回 nil 交给上层等待/排队，不再回退常规选号（换取 filter 单次求值；行为变化已记录） | +| **Strict 模式** | 完全尊重原语义，不执行强制解绑；无绑定时的首选享受加权避险 | +| **Remaining Quota / Fill First** | 保持原有用量优先排序，仅继承冷却与过滤修复 | +| **其他渠道（Codex / Claude / Grok）**| `opts.Provider != antigravity`，完全走原逻辑，零影响 | + +--- + +## 5. 开关与回滚 + +- 环境变量开关：`ANTIGRAVITY_BALANCED_AFFINITY` + - 默认值：`true`（开启加权避险）； + - 设置为 `0`, `false`, `off`, `none`, `no` 时关闭加权避险，回退为原 round-robin； + - 冷却修复（Canonical 双写与跨实例读）属于正确性修复，始终生效。 +- 关闭需重启生效。 + +--- + +## 6. 代码修改落点 + +| 文件 | 修改职责 | +| :--- | :--- | +| `auth/antigravity_balance.go`（新建） | `SelectionOptions`、加权打分与分层选号（`(*Account, handled)`）、FNV tie-break、filter 单次求值保护 | +| `auth/antigravity_ratelimit.go`（新建） | 限流事件数据结构、分类记录方法、惩罚累加器、开关管理 | +| `auth/store.go` | 字段挂载、`SelectionOptions` 贯穿选号/等待入口、`WithAntigravityCooldownFilterContext` 与 `WithAntigravityAwareCooldownFilterContext` 跨实例三键/分流过滤、handled 直接返回 | +| `proxy/antigravity_quota.go` | `applyAntigravityCooldown` 补齐 wireModel 冷却写入与单点事件记录；非结构化 429 补写 wire；共享 wire/family 冷却走 `MarkModelCooldownUntilAtLeast`（只延长不缩短）；admin 连接测试使用不记账模式 | +| `proxy/handler_gemini.go` | `handleGeminiCountTokens` 429/503 补齐调用 `applyCooldownForModel`；native 入口强制 `Provider=antigravity` | +| `proxy/handler.go` 及相关 handler | `withRequestModelCooldownFilter` 三档渠道分流；`SelectionOptions` 贯穿 Responses/Chat/Messages/WS/重试/等待/续链；候选池排除 AG 的 compact 入口改用 policy-only `SelectionOptions` | +| `auth/antigravity_balance_test.go`（新建） | 加权避险排序、分层 fallback、饱和池 filter 单次、handled 语义 | +| `proxy/antigravity_balance_test.go` / `proxy/antigravity_auto_channel_test.go`（新建） | 跨实例三键过滤、auto 渠道分流、native 强制 provider、端到端 countTokens 加权 | +| `proxy/antigravity_quota_test.go` | wireModel 同步写入、非结构化 429 wire、countTokens 冷却、admin 隔离 | + +--- + +## 7. 验收测试清单 + +1. **别名冷却联动测试**：请求 `gemini-3.8-flash-high` 触发 429 冷却后，同一账号立即请求 `gemini-3.8-flash-low`，断言被模型冷却拦截。 +2. **跨实例冷却同步测试**：实例 A 写入 `gemini-3.8-flash-tiered` 冷却，实例 B 使用带有共享缓存的 filter，断言成功过滤该账号。 +3. **countTokens 冷却写入测试**：`countTokens` 收到 429，断言该账号模型冷却被正确写入。 +4. **加权避险与首绑打分散列测试**： + - 4 个健康账号，其中账号 A 记录一次 429 惩罚；发起多次不同 key 的首选，断言账号 A 选中概率显著低于其他账号； + - 断言候选账号的 filter 闭包只被求值一次； + - 饱和池（所有候选并发满）+ 二次求值即 panic 的 filter：断言不 panic、每账号至多一次、返回 nil（handled=true 不回退）； + - 层间 fallback：最高层全满时能下沉到次高层取号。 +5. **auto Key 与非结构化 429**： + - auto 渠道 + AG 账号：跨实例 wire/family 冷却能拦截；auto 渠道 + 非 AG 账号仍只查 public； + - native Gemini/countTokens 在 auto Key 下仍强制 `Provider=antigravity`，端到端验证受罚账号不被选中； + - 非结构化 429 回退分支补写 wire 冷却（relay 冷却策略开启时）。 +6. **续链与等待路径**：preserveBinding/turnContinuationPinned 分支透传 `SelectionOptions`，不丢失加权上下文。 +7. **共享冷却 deadline 单调性**：同账号先写较长的 wire/family 冷却（如 290s）、再触发更短事件（如 5s）时，共享 deadline 保持不变（断言剩余 > 60s）；先短后长则延长。public 冷却保持既有覆盖语义不变。 +8. **compact 入口不套用加权**：compact 执行器与 native remote compaction 的候选池排除 AG 账号，Antigravity 渠道 Key 下这两个入口的 `SelectionOptions.Provider` 为空、只保留 `Policy`，不绕过 spread/HRW。 +9. **既有调度回归测试**： + - `go test ./auth/ -run 'Antigravity|Affinity|LockOrder' -count=1` + - `go test ./proxy/ -run 'Antigravity|GeminiNative|WithRequestModelCooldownFilter' -count=1` + - `go test -race ./auth/ ./proxy/ -count=1` + +--- + +## 8. 已知边界与后续优化（P2，不阻塞本次定稿） + +以下由 Phase 2 修复复审（deepseek-flash）确认，记录为已知行为/后续项： + +1. **auto Key 的通用入口不启用加权**：Responses/Chat/Messages 在 auto Key 下 Provider 为空，只享受三键冷却过滤修复，不启用加权避险；native Gemini/countTokens 因候选池只含 AG 而强制启用。若后续希望 auto Key 通用入口也加权，可按模型归属（`antigravityPublicModel`）推导 provider，属于后续增强。 +2. **capacityFull 绑定溢出路径重复 filter**：`takeByIDModeWithCapacity` 先 filter，容量满回退 fresh 选号时会再对同一账号求值一次 filter。这是改动前就存在的路径，不在本次「加权路径单次求值」范围内；后续可记录 filter 结果或前移容量检查。 +3. **加权路径未接入调度指标/索引**：未调用 `recordSchedulerSelection`，`selection_total` 会低估 AG fresh 选号；indexed 引擎下改为 O(N) 全量扫描 + 排序，大号池有性能/可观测性回归。后续可补指标与候选窗口。 +4. **非结构化 429 的 wire 冷却受 relay 策略约束**：默认 `relay_model_cooldown_mode=off` 时 `Apply429Cooldown` 不写任何冷却，因此非结构化 429 只记 0.5 惩罚、不写 public/wire；开启 relay 冷却后才会补写 wire。这尊重运营者配置。 +5. **测试缺口**：续链 Provider 透传入口、层间 fallback 的「整层 tryAcquire 失败」路径、load/usage 权重敏感性、capacityFull 重复 filter 尚未直接覆盖。 +6. **spread=true（非默认）行为变化**：AG 加权路径在 spread 检查之前接管，HRW 一致哈希被 FNV tie-break 取代；规格按 round_robin 触发，属于非默认配置的已接受变化。 +7. **`ResetAntigravityPenalty` 暂无生产调用点**：人工清冷却/恢复不会清惩罚，依赖自然衰减（180s 窗口）。 +8. **compact/WS 的冷却过滤仍按渠道三键**：这两个入口候选池不含 AG（且没有 `applyUpstreamChannelFilter`），但 `withRequestModelCooldownFilter` 的 antigravity 分支仍按三键检查；对 Codex 模型 `AntigravityWireModelID` 为空无影响，对 Gemini public model 理论上可能误过滤一个在 AG wire key 上持有冷却的 relay 账号。后续可改用 `WithAntigravityAwareCooldownFilterContext` 按账号类型分流。 +9. **`affinity_mode=off` 绕过加权（非默认）**：`nextForSessionWithFilter` 在 off 模式直接走 `NextExcludingWithDispatch`，丢弃 `SelectionOptions.Provider`，加权避险不生效；默认 bounded 不受影响。 +10. **`lazy_mode=true` 绕过 lazy 预热（非默认）**：加权路径使用 `dispatchableForPolicy`/实时快照，不调用 `accountLazySelectable`/`ensureLazyDispatchReady`，需要刷新凭据的账号可能被选中并多一次 401/刷新；默认 lazy=false 不受影响。 +11. **`SelectionOptions.Model` 目前 reserved**：auth 侧只读取 `Provider`/`Policy`，`Model` 由 handler 传入并保留给后续按模型族推导 scope/门控。 +12. **cache/DB 冷却写入为 last-writer-wins**：本机内存的 AtLeast 单调性成立；`afterModelCooldownMarked` 的 cache `SetRuntime` 与 DB upsert 在锁外，同账号兄弟 tier 并发时短事件可能在长事件之后把跨实例 cache/DB 改回短 deadline（窄窗口、跨实例 best-effort）。若要求严格跨实例单调，需要 max-aware 的 cache/DB 写入。 +13. **`TestForwardImagesCatchAllRetriesBeyondOrdinaryAttemptCap` 为既有 race flake**：5 次全量 `-race ./proxy/` 中 1 次失败（upstream calls=7/want 6），单独 `-race -count=20` 全过；未触碰 images/runtime settings，判断为既有 image 重试路径在 race 调度下的时序 flake。CI 建议隔离/记录该测试，不以全量 race 全绿作为本变更验收前提。 diff --git a/proxy/antigravity_auto_channel_test.go b/proxy/antigravity_auto_channel_test.go new file mode 100644 index 000000000..e0700deab --- /dev/null +++ b/proxy/antigravity_auto_channel_test.go @@ -0,0 +1,234 @@ +package proxy + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/codex2api/auth" + "github.com/codex2api/cache" + "github.com/codex2api/database" + "github.com/gin-gonic/gin" +) + +// auto 渠道：auto 渠道 Key（requestUpstreamChannel 为空）也可能选中 Antigravity 账号 +// （native Gemini 入口把候选池硬限定为 Antigravity OAuth，Responses / Chat / Messages +// 也允许），因此冷却过滤器必须对 Antigravity 候选做 public / wire / family 三键检查， +// 把另一实例写入的 wire 与 family 冷却一并拦下。 +func TestWithRequestModelCooldownFilterAutoChannelInterceptsAntigravityCandidate(t *testing.T) { + tokenCache := cache.NewMemory(8) + defer tokenCache.Close() + + writer := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer writer.Stop() + hotAccount := &auth.Account{ + DBID: 9110, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + Models: []string{"gemini-3.8-flash-tiered"}, + } + writer.AddAccount(hotAccount) + writer.MarkModelCooldownUntil(hotAccount, "gemini-3.8-flash-tiered", "rate_limited_model", time.Now().Add(time.Minute)) + + reader := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer reader.Stop() + handler := NewHandler(reader, nil, nil, nil) + remote := &auth.Account{ + DBID: 9110, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + } + if remote.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("the reading-side account must start with no local cooldown") + } + + // auto 渠道（未限定）+ Antigravity 候选：wire 冷却必须拦截并 hydrate 进本地内存。 + autoCtx := newChannelTestGin("") + filter := handler.withRequestModelCooldownFilter(autoCtx, context.Background(), "gemini-3.8-flash-low", nil) + if filter(remote) { + t.Fatal("the auto-channel filter must apply the three-key check to an Antigravity candidate") + } + if !remote.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("the cross-instance wire cooldown was not synchronized into the local account") + } + + // family 冷却同样要拦住 auto 渠道下的 Antigravity 候选。 + familyWriter := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer familyWriter.Stop() + quotaAccount := &auth.Account{ + DBID: 9111, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + } + familyWriter.AddAccount(quotaAccount) + familyWriter.MarkModelCooldownUntil(quotaAccount, antigravityGeminiFamilyCooldownKey, "quota_exhausted", time.Now().Add(time.Minute)) + remoteFamily := &auth.Account{ + DBID: 9111, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + } + if filter(remoteFamily) { + t.Fatal("the auto-channel filter must apply the Gemini family cooldown to an Antigravity candidate") + } +} + +// auto 渠道：auto 渠道下的非 Antigravity 候选必须保持单键语义——看不到 wire / family 冷却， +// 但 publicModel 的跨实例冷却仍然照常拦截。 +func TestWithRequestModelCooldownFilterAutoChannelKeepsSingleKeyForOtherAccounts(t *testing.T) { + tokenCache := cache.NewMemory(8) + defer tokenCache.Close() + + writer := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer writer.Stop() + wireHot := &auth.Account{DBID: 9112, AccessToken: "openai-token", Status: auth.StatusReady} + publicHot := &auth.Account{DBID: 9113, AccessToken: "openai-token", Status: auth.StatusReady} + writer.AddAccount(wireHot) + writer.AddAccount(publicHot) + writer.MarkModelCooldownUntil(wireHot, "gemini-3.8-flash-tiered", "rate_limited_model", time.Now().Add(time.Minute)) + writer.MarkModelCooldownUntil(publicHot, "gemini-3.8-flash-low", "rate_limited_model", time.Now().Add(time.Minute)) + writer.MarkModelCooldownUntil(publicHot, antigravityGeminiFamilyCooldownKey, "quota_exhausted", time.Now().Add(time.Minute)) + + reader := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer reader.Stop() + handler := NewHandler(reader, nil, nil, nil) + autoCtx := newChannelTestGin("") + filter := handler.withRequestModelCooldownFilter(autoCtx, context.Background(), "gemini-3.8-flash-low", nil) + + // 只有 wire 冷却的非 Antigravity 账号：单键语义下必须放行。 + if !filter(&auth.Account{DBID: 9112, AccessToken: "openai-token", Status: auth.StatusReady}) { + t.Fatal("a non-Antigravity candidate must not be filtered by the wire-model cooldown") + } + // public 冷却仍然拦截（family 冷却与它同批写入，不能成为放行的理由）。 + if filter(&auth.Account{DBID: 9113, AccessToken: "openai-token", Status: auth.StatusReady}) { + t.Fatal("a non-Antigravity candidate must still be filtered by its own public-model cooldown") + } + + // 显式渠道（codex）保持单键：既看不到 wire，也看 family 之外的键。 + codexCtx := newChannelTestGin(database.UpstreamChannelCodex) + codexFilter := handler.withRequestModelCooldownFilter(codexCtx, context.Background(), "gemini-3.8-flash-low", nil) + if !codexFilter(&auth.Account{DBID: 9112, AccessToken: "openai-token", Status: auth.StatusReady}) { + t.Fatal("an explicit non-Antigravity channel must keep the single public-model check") + } +} + +// auto 渠道：native Gemini 入口（generateContent / streamGenerateContent / countTokens）的 +// 候选池被硬限定为 Antigravity OAuth，SelectionOptions.Provider 必须恒为 antigravity， +// 否则 auto 渠道 Key 下加权避险整体失效。 +func TestGeminiNativeSelectionOptionsAlwaysPinAntigravityProvider(t *testing.T) { + cases := []struct { + name string + channel string + }{ + {name: "auto key", channel: ""}, + {name: "antigravity key", channel: database.UpstreamChannelAntigravity}, + {name: "codex key", channel: database.UpstreamChannelCodex}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + opts := geminiNativeSelectionOptions(newChannelTestGin(tc.channel), "gemini-3.8-flash-high", auth.DispatchPolicyStandard) + if opts.Provider != database.UpstreamChannelAntigravity { + t.Fatalf("Provider = %q, want %q", opts.Provider, database.UpstreamChannelAntigravity) + } + if opts.Model != "gemini-3.8-flash-high" || opts.Policy != auth.DispatchPolicyStandard { + t.Fatalf("SelectionOptions = %+v, want the model and policy preserved", opts) + } + }) + } + + // 对照：通用入口不覆盖渠道，auto Key 下 Provider 仍为空（加权避险由渠道判定关闭）。 + if opts := selectionOptionsForRequest(newChannelTestGin(""), "gemini-3.8-flash-high", auth.DispatchPolicyStandard); opts.Provider != "" { + t.Fatalf("the generic entry must keep the request channel, got %q", opts.Provider) + } +} + +// auto 渠道端到端：auto 渠道 Key 打到 native countTokens 时，候选池里的 Antigravity 账号 +// 必须参与加权避险——刚被限流（惩罚最高）的账号在号池空闲时一次都不该被选中。 +func TestGeminiNativeCountTokensAutoKeyUsesWeightedSelection(t *testing.T) { + gin.SetMode(gin.TestMode) + + var mu sync.Mutex + served := map[string]int{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // 上游按凭据识别账号：每个账号的 access token 不同，直接映射回账号。 + bearer := strings.TrimSpace(strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer")) + label := "unknown" + switch bearer { + case "at-hot": + label = "hot" + case "at-healthy": + label = "healthy" + } + mu.Lock() + served[label]++ + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `{"totalTokens":3}`) + })) + defer server.Close() + previous := antigravityOAuthEndpointBases + antigravityOAuthEndpointBases = []string{server.URL} + t.Cleanup(func() { antigravityOAuthEndpointBases = previous }) + + store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 4, TestConcurrency: 1, TestModel: "gemini-3.8-flash-high"}) + defer store.Stop() + store.SetSchedulerMode("round_robin") + store.SetMaxRetries(0) + store.SetMaxRateLimitRetries(0) + newAccount := func(id int64, project, token string) *auth.Account { + acc := &auth.Account{ + DBID: id, UpstreamType: auth.UpstreamAntigravity, AccessToken: token, + AntigravityProjectID: project, Status: auth.StatusReady, + Models: []string{"gemini-3.8-flash-tiered"}, + } + store.AddAccount(acc) + return acc + } + hot := newAccount(9120, "project-hot", "at-hot") + newAccount(9121, "project-healthy", "at-healthy") + + auth.SetAntigravityBalancedAffinity(true) + t.Cleanup(func() { auth.SetAntigravityBalancedAffinity(true) }) + auth.RecordAntigravityRateLimitEvent(hot, auth.AntigravityRateLimitEvent{ + StatusCode: 429, + Kind: "rate_limit", + Scope: auth.AGScopeModelRateLimit, + PublicModel: "gemini-3.8-flash-high", + WireModel: "gemini-3.8-flash-tiered", + Weight: auth.AntigravityPenaltyWeightExplicitRateLimit, + ObservedAt: time.Now(), + }) + if hot.AntigravityPenalty(time.Now()) <= 0 { + t.Fatal("the penalized account lost its penalty before the run") + } + + router := gin.New() + handler := NewHandler(store, nil, nil, nil) + router.Use(func(c *gin.Context) { + // auto 渠道 Key：不限定上游渠道，requestUpstreamChannel 返回空串。 + c.Set(contextAPIKeyRow, &database.APIKeyRow{ID: 7, Limits: database.APIKeyLimits{UpstreamChannel: database.UpstreamChannelAuto}}) + c.Next() + }) + router.POST("/v1beta/models/*action", handler.GeminiModelsAction) + + for i := 0; i < 12; i++ { + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3.8-flash-high:countTokens", + strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"probe-`+string(rune('a'+i))+`"}]}]}`)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("request %d: status = %d body=%s", i, rec.Code, rec.Body.String()) + } + } + + mu.Lock() + defer mu.Unlock() + if served["healthy"] == 0 { + t.Fatalf("no request reached the healthy account: %v", served) + } + if served["hot"] != 0 { + t.Fatalf("the penalized account still served %d of %d fresh sessions: %v", served["hot"], served["hot"]+served["healthy"], served) + } +} diff --git a/proxy/antigravity_balance_test.go b/proxy/antigravity_balance_test.go new file mode 100644 index 000000000..dda001e2c --- /dev/null +++ b/proxy/antigravity_balance_test.go @@ -0,0 +1,125 @@ +package proxy + +import ( + "context" + "net/http/httptest" + "testing" + "time" + + "github.com/codex2api/auth" + "github.com/codex2api/cache" + "github.com/codex2api/database" + "github.com/gin-gonic/gin" +) + +// newChannelTestGin 构造带指定上游渠道 API Key 的 gin 上下文；channel 为空表示不限渠道。 +func newChannelTestGin(channel string) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + if channel != "" { + c.Set(contextAPIKeyRow, &database.APIKeyRow{ID: 7, Limits: database.APIKeyLimits{UpstreamChannel: channel}}) + } + return c +} + +// TestWithRequestModelCooldownFilterInterceptsCrossInstanceAntigravity 覆盖 auto 渠道： +// Antigravity 渠道请求的冷却过滤器必须解析 publicModel 对应的 wireModel 与 Gemini +// familyKey，对「其它实例写入」的 wire / family 冷却一并拦截（跨 tier、跨实例）； +// 非 Antigravity 渠道仍只检查单一 publicModel，行为不变。 +func TestWithRequestModelCooldownFilterInterceptsCrossInstanceAntigravity(t *testing.T) { + tokenCache := cache.NewMemory(8) + defer tokenCache.Close() + + // 写入实例：账号 9100 在 gemini-3.8-flash-high 上被限流，冷却落在共享 wire 模型上。 + writer := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer writer.Stop() + hotAccount := &auth.Account{ + DBID: 9100, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + Models: []string{"gemini-3.8-flash-tiered"}, + } + writer.AddAccount(hotAccount) + writer.MarkModelCooldownUntil(hotAccount, "gemini-3.8-flash-tiered", "rate_limited_model", time.Now().Add(time.Minute)) + + // 读取实例：与写入实例共享缓存，但账号在本地内存里没有任何冷却记录（模拟另一实例视角）。 + reader := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer reader.Stop() + handler := NewHandler(reader, nil, nil, nil) + remote := &auth.Account{ + DBID: 9100, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + } + if remote.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("the reading-side account must start with no local cooldown") + } + + // Antigravity 渠道 + publicModel=gemini-3.8-flash-low（wire=gemini-3.8-flash-tiered）： + // 三键过滤器必须命中另一实例写入的 wire 冷却，把该账号过滤掉。 + antigravityCtx := newChannelTestGin(database.UpstreamChannelAntigravity) + filter := handler.withRequestModelCooldownFilter(antigravityCtx, context.Background(), "gemini-3.8-flash-low", nil) + if filter(remote) { + t.Fatal("the cross-instance wire-model cooldown must filter the Antigravity candidate") + } + if !remote.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("the remote wire cooldown was not synchronized into the local account") + } +} + +// TestWithRequestModelCooldownFilterInterceptsGeminiFamilyCooldown 覆盖 familyKey： +// 另一实例写入的 Gemini 家族共享冷却（antigravity:gemini）必须拦截任意 Gemini publicModel +// 的 Antigravity 候选，即使该 publicModel 与其 wire 都没有单独冷却。 +func TestWithRequestModelCooldownFilterInterceptsGeminiFamilyCooldown(t *testing.T) { + tokenCache := cache.NewMemory(8) + defer tokenCache.Close() + + writer := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer writer.Stop() + hotAccount := &auth.Account{ + DBID: 9101, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + } + writer.AddAccount(hotAccount) + writer.MarkModelCooldownUntil(hotAccount, antigravityGeminiFamilyCooldownKey, "rate_limited_model", time.Now().Add(time.Minute)) + + reader := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer reader.Stop() + handler := NewHandler(reader, nil, nil, nil) + remote := &auth.Account{ + DBID: 9101, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + } + + // publicModel 与 wire 都没有冷却，只有家族键有：Antigravity 过滤器仍应拦截。 + antigravityCtx := newChannelTestGin(database.UpstreamChannelAntigravity) + filter := handler.withRequestModelCooldownFilter(antigravityCtx, context.Background(), "gemini-3.8-flash-medium", nil) + if filter(remote) { + t.Fatal("the cross-instance Gemini family cooldown must filter the Antigravity candidate") + } +} + +// TestWithRequestModelCooldownFilterNonAntigravityKeepsSingleModelCheck 证明渠道感知： +// 非 Antigravity 渠道（无渠道限定）沿用单一 publicModel 检查，看不到 wire 冷却，账号放行。 +func TestWithRequestModelCooldownFilterNonAntigravityKeepsSingleModelCheck(t *testing.T) { + tokenCache := cache.NewMemory(8) + defer tokenCache.Close() + + writer := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer writer.Stop() + hotAccount := &auth.Account{ + DBID: 9102, AccessToken: "openai-token", Status: auth.StatusReady, + } + writer.AddAccount(hotAccount) + writer.MarkModelCooldownUntil(hotAccount, "gemini-3.8-flash-tiered", "rate_limited_model", time.Now().Add(time.Minute)) + + reader := auth.NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4}) + defer reader.Stop() + handler := NewHandler(reader, nil, nil, nil) + remote := &auth.Account{DBID: 9102, AccessToken: "openai-token", Status: auth.StatusReady} + + // 无渠道限定：只检查 publicModel=gemini-3.8-flash-low，wire 冷却不可见 → 放行。 + plainCtx := newChannelTestGin("") + filter := handler.withRequestModelCooldownFilter(plainCtx, context.Background(), "gemini-3.8-flash-low", nil) + if !filter(remote) { + t.Fatal("a non-Antigravity channel must not see the wire-model cooldown and must pass the candidate") + } +} diff --git a/proxy/antigravity_compact_selection_test.go b/proxy/antigravity_compact_selection_test.go new file mode 100644 index 000000000..ed79c2dab --- /dev/null +++ b/proxy/antigravity_compact_selection_test.go @@ -0,0 +1,56 @@ +package proxy + +import ( + "testing" + + "github.com/codex2api/auth" + "github.com/codex2api/database" +) + +// TestCompactSelectionOptionsSkipAntigravityWeighting 覆盖 compact 入口的选号上下文： +// compact 执行器与 native remote compaction 的候选池显式排除 Antigravity 账号，因此 +// Antigravity 渠道 Key 打到这两个入口时 SelectionOptions.Provider 必须为空，否则 +// Store 侧会对 Codex relay 候选套用 Antigravity 加权打分，并绕过 spread / HRW。 +func TestCompactSelectionOptionsSkipAntigravityWeighting(t *testing.T) { + c := newChannelTestGin(database.UpstreamChannelAntigravity) + + // 前置事实：compact 候选池确实不含 Antigravity 账号，Provider 为空才是正确语义。 + antigravityAccount := &auth.Account{ + DBID: 9130, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + Models: []string{"gemini-3.8-flash-tiered"}, + } + if accountFilterForCompactResponsesModelWithOriginal("gemini-3.8-flash-high", "gemini-3.8-flash-high", false)(antigravityAccount) { + t.Fatal("the compact executor pool must keep excluding Antigravity accounts") + } + if accountFilterForInlineCompactionModelWithOriginal("gemini-3.8-flash-high", "gemini-3.8-flash-high", false)(antigravityAccount) { + t.Fatal("the native remote compaction pool must keep excluding Antigravity accounts") + } + + // 对照：普通 Responses 请求（候选池可能含 Antigravity 账号）仍带渠道，加权避险照常启用。 + ordinary := selectionOptionsForCandidatePool(c, "gemini-3.8-flash-high", auth.DispatchPolicyStandard, false) + if ordinary.Provider != database.UpstreamChannelAntigravity { + t.Fatalf("ordinary Responses options Provider = %q, want the request channel", ordinary.Provider) + } + if ordinary.Model != "gemini-3.8-flash-high" || ordinary.Policy != auth.DispatchPolicyStandard { + t.Fatalf("ordinary Responses options = %+v, want the model and policy preserved", ordinary) + } + + // compact 执行器与 native remote compaction：只保留 Policy，不启用加权。 + cases := []struct { + name string + policy auth.DispatchPolicy + }{ + {name: "compact executor", policy: auth.DispatchPolicyStandard}, + {name: "native remote compaction", policy: auth.DispatchPolicySpark}, + } + for _, tc := range cases { + opts := selectionOptionsForCandidatePool(c, "gpt-5.6-sol", tc.policy, true) + if opts.Provider != "" { + t.Fatalf("%s: Provider = %q, want empty so the Antigravity weighting stays off", tc.name, opts.Provider) + } + if opts.Policy != tc.policy { + t.Fatalf("%s: Policy = %v, want it preserved", tc.name, opts.Policy) + } + } +} diff --git a/proxy/antigravity_quota.go b/proxy/antigravity_quota.go index 97c960fcf..7ff526a9c 100644 --- a/proxy/antigravity_quota.go +++ b/proxy/antigravity_quota.go @@ -292,6 +292,54 @@ func antigravityAccountModelRateLimited(account *auth.Account, model, wireModel return false } +// antigravitySharedWireCooldownKey 报告「把冷却写到公开档位名 model 上」这一动作是否 +// 兼任共享 wire 冷却键的写入。命中时这次写入必须遵守 deadline 单调性（AtLeast 变体）， +// 不能被最后写入者胜出地改短。 +// +// 两种情形： +// 1. 账号解析出的 wire 与请求的公开档位同名（wireModel == model）。典型就是 Argon： +// gemini-4-argon-low/medium/high 三个公开档位的 backing 都是 gemini-4-argon-high， +// 而 high 这一档的公开名与 wire 名完全相同，于是「写 public 冷却」与「写共享 wire +// 冷却」是同一次写入——public 键在这里就是共享键本身。 +// 2. 目录里另有别的公开档位 / 逻辑档位把 model 解析成自己的上游 backing（请求直接 +// 带原始 backing id，或原始 id 与某个公开档位同名）。 +// +// 两种情形下，一个 5s 的裸 RESOURCE_EXHAUSTED 都可能把兄弟档位刚写入的 290s 共享冷却 +// 改短：短 deadline 一过，同 backing 的兄弟档位又会选中刚被限流的凭据，继续收集 429。 +func antigravitySharedWireCooldownKey(model, wireModel string) bool { + trimmed := strings.ToLower(strings.TrimSpace(model)) + if trimmed == "" { + return false + } + if wireModel != "" && strings.EqualFold(strings.TrimSpace(wireModel), trimmed) { + return true + } + return antigravityModelSharedBySiblings(trimmed) +} + +// antigravityModelSharedBySiblings 报告固定档位目录或逻辑档位目录中是否存在「其它」公开 +// 档位，把 model 解析成自己的上游 backing。model 自身的定义不算：公开档位解析到自己的 +// backing 是常态，只有别的档位也骑上来才构成共享键（公开名与自己同名的情形由调用方的 +// wireModel == model 分支单独覆盖）。 +func antigravityModelSharedBySiblings(model string) bool { + for _, catalog := range [][]antigravityPublicModelDefinition{ + antigravityPublicModelCatalog, + antigravityLogicalCompatibilityCatalog, + } { + for _, definition := range catalog { + if strings.EqualFold(definition.id, model) { + continue + } + for _, wire := range AntigravityWireModelIDs(definition.id) { + if strings.EqualFold(wire, model) { + return true + } + } + } + } + return false +} + // antigravityNonPenalizingUpstreamFailure reports failures that say nothing // about the account's health: a model-wide capacity shortage is Google's // problem, not this credential's, so ReportRequestFailure must not count it. @@ -308,11 +356,107 @@ func antigravityNonPenalizingUpstreamFailure(account *auth.Account, statusCode i // fall back to the generic relay-style cooldown policy, and a recognised body // that carries no hint at all is backed off briefly instead of being charged // the full default bench. +// +// Live dispatch also feeds the scheduling penalty accumulator through here (see +// auth.RecordAntigravityRateLimitEvent); admin connectivity probes reuse the +// same cooldown mapping without that accounting via ApplyAntigravityCooldown. func applyAntigravityCooldown(store *auth.Store, account *auth.Account, statusCode int, body []byte, resp *http.Response, model string) codex429Decision { + return applyAntigravityCooldownWithPenalty(store, account, statusCode, body, resp, model, true) +} + +// antigravityRateLimitEvent normalises one 429/503 observation into the event +// shape the fresh-selection penalty understands. parsed is false when the body +// carried no recognisable google.rpc signal at all. +func antigravityRateLimitEvent(statusCode int, model, wireModel string, quota antigravityQuotaError, parsed bool, resetAt time.Time) auth.AntigravityRateLimitEvent { + event := auth.AntigravityRateLimitEvent{ + StatusCode: statusCode, + PublicModel: strings.TrimSpace(model), + WireModel: strings.TrimSpace(wireModel), + ResetAt: resetAt, + ObservedAt: time.Now(), + } + if antigravityIsGeminiModel(event.PublicModel) || antigravityIsGeminiModel(event.WireModel) { + event.Family = antigravityGeminiFamilyCooldownKey + } + if !parsed { + // A 429 with no google.rpc detail says nothing about this credential, + // so it is charged half a penalty instead of a full one. + event.Kind = "unstructured_rate_limit" + event.Scope = auth.AGScopeUnknown + event.Weight = auth.AntigravityPenaltyWeightUnclassifiedRateLimit + return event + } + event.Kind = string(quota.Kind) + switch quota.Kind { + case antigravityQuotaKindModelCapacity: + // A model-wide shortage is Google's problem; switching accounts cannot + // help, so the scheduler must not steer away from this credential. + event.Scope = auth.AGScopeSharedCapacity + event.Weight = auth.AntigravityPenaltyWeightNone + case antigravityQuotaKindQuotaExhausted: + // Already covered by the long per-model cooldown plus the Gemini + // family cooldown; a short-term penalty would only double-count it. + event.Scope = auth.AGScopeAccountQuota + event.Weight = auth.AntigravityPenaltyWeightNone + default: + event.Scope = auth.AGScopeModelRateLimit + if quota.Reason != "" || quota.HasRetryDelay { + event.Weight = auth.AntigravityPenaltyWeightExplicitRateLimit + } else { + event.Weight = auth.AntigravityPenaltyWeightUnclassifiedRateLimit + } + } + return event +} + +func applyAntigravityCooldownWithPenalty(store *auth.Store, account *auth.Account, statusCode int, body []byte, resp *http.Response, model string, recordPenalty bool) codex429Decision { quota, ok := parseAntigravityQuotaError(statusCode, body) if !ok { if statusCode == http.StatusTooManyRequests { - return Apply429Cooldown(store, account, body, resp, model) + // wire 解析要在通用 429 策略之前完成：它既决定下面的共享键判定，也决定 + // 兄弟档位的 wire 补写。 + wireModel := "" + if account != nil { + if wire, resolved := antigravityResolvePublicModelForAccount(account, model); resolved { + wireModel = strings.TrimSpace(wire) + } + } + // 公开档位名本身就是共享 wire 键时（gemini-4-argon-high），先记住既有 deadline： + // 通用 429 策略按「最后写入者胜出」覆盖 public 冷却，会把同账号兄弟档位刚写入的 + // 长冷却改短。 + sharedWireKey := store != nil && account != nil && antigravitySharedWireCooldownKey(model, wireModel) + var previous auth.ModelCooldown + if sharedWireKey { + previous, _ = account.ModelCooldownFor(model) + } + decision := Apply429Cooldown(store, account, body, resp, model) + // 非结构化 429（Cloud Code 负载抑制时只回 {"code":429,"status": + // "RESOURCE_EXHAUSTED"}）走通用 429 策略，而通用策略只按 public id 落冷却。 + // 多个公开档位共用一个 wire 模型时，同 wire 的兄弟档位仍会被选中，账号带着 + // 刚被限流的凭据再次上游。这里补写 wire 冷却；解析不到 wire（非目录模型 / + // 账号不具备该 wire）时跳过。wire 是共享冷却键，用 AtLeast 变体写入，避免本次 + // 更短的 deadline 把同账号上更长的 wire 冷却改短。 + // + // public 兼任共享 wire 键时（wireModel == decision.Model）没有单独的 wire 键可 + // 补写：通用策略那次覆盖写就是共享键写入，因此这里把被改短的长冷却写回。 + // 只处理 model 作用域且本次确实写了冷却的情况，账号作用域由账号冷却负责， + // 不应在这里被模型冷却复活。 + if !previous.ResetAt.IsZero() && decision.Scope == rateLimitScopeModel && !decision.ResetAt.IsZero() && + previous.ResetAt.After(decision.ResetAt) && strings.EqualFold(strings.TrimSpace(decision.Model), strings.TrimSpace(model)) { + if restored := store.MarkModelCooldownUntilAtLeast(account, decision.Model, previous.Reason, previous.ResetAt); !restored.ResetAt.IsZero() { + decision.ResetAt = restored.ResetAt + decision.Cooldown = time.Until(restored.ResetAt) + } + } + if store != nil && account != nil && decision.Model != "" && !decision.ResetAt.IsZero() { + if wireModel != "" && !strings.EqualFold(wireModel, decision.Model) { + store.MarkModelCooldownUntilAtLeast(account, wireModel, decision.Reason, decision.ResetAt) + } + } + if recordPenalty { + auth.RecordAntigravityRateLimitEvent(account, antigravityRateLimitEvent(statusCode, decision.Model, wireModel, antigravityQuotaError{}, false, decision.ResetAt)) + } + return decision } return codex429Decision{} } @@ -362,13 +506,46 @@ func applyAntigravityCooldown(store *auth.Store, account *auth.Account, statusCo if store == nil || account == nil || decision.Model == "" { return decision } - marked := store.MarkModelCooldownUntil(account, decision.Model, decision.Reason, resetAt) + // wire 解析必须在 public 冷却写入之前完成：公开档位名与 wire 同名时，下面这次 + // public 写入就是共享 wire 键的唯一一次写入，写入变体要据此选择。 + wireModel := "" + if wire, resolved := antigravityResolvePublicModelForAccount(account, decision.Model); resolved { + wireModel = strings.TrimSpace(wire) + } + // public 冷却默认保持「最后写入者胜出」的既有覆盖语义。唯一例外是 public 名本身 + // 兼任共享 wire 冷却键（gemini-4-argon-high 既是公开档位，也是 low/medium 的上游 + // backing）：此时覆盖写会把兄弟档位刚写入的 290s 共享冷却改成 5s，短窗口一过 + // low/medium 又选中刚被限流的凭据。这类键与 wire、family 一样属于同账号上的共享 + // 冷却键，必须用 AtLeast 变体，只延长、不改短。 + sharedWireKey := antigravitySharedWireCooldownKey(decision.Model, wireModel) + var marked auth.ModelCooldown + if sharedWireKey { + marked = store.MarkModelCooldownUntilAtLeast(account, decision.Model, decision.Reason, resetAt) + } else { + marked = store.MarkModelCooldownUntil(account, decision.Model, decision.Reason, resetAt) + } if !marked.ResetAt.IsZero() { decision.ResetAt = marked.ResetAt decision.Cooldown = time.Until(marked.ResetAt) } + // Public tiers may share one upstream backing (gemini-3.8-flash-low/medium/ + // high all ride gemini-3.8-flash-tiered). Cooling only the public id leaves + // the sibling tiers selectable on the very account the upstream just + // throttled, so the wire id is cooled with the same deadline. + // + // wire 与 family 都是同账号上的共享冷却键，会被多个公开档位的事件并发写入，因此 + // 用 AtLeast 变体：high 的 290s RATE_LIMIT_EXCEEDED 之后，同账号 low 的 5s 裸 + // RESOURCE_EXHAUSTED 只能延长、不能把共享 deadline 改短，否则短窗口一过兄弟档位 + // 又会选中刚被限流的凭据。public 与 wire 同名时没有第二个键可写，单调性由上面的 + // sharedWireKey 分支保证。 + if wireModel != "" && !strings.EqualFold(wireModel, decision.Model) { + store.MarkModelCooldownUntilAtLeast(account, wireModel, decision.Reason, decision.ResetAt) + } if quota.Kind == antigravityQuotaKindQuotaExhausted && (antigravityIsGeminiModel(decision.Model) || antigravityIsGeminiModel(quota.Model)) { - store.MarkModelCooldownUntil(account, antigravityGeminiFamilyCooldownKey, decision.Reason, resetAt) + store.MarkModelCooldownUntilAtLeast(account, antigravityGeminiFamilyCooldownKey, decision.Reason, decision.ResetAt) + } + if recordPenalty { + auth.RecordAntigravityRateLimitEvent(account, antigravityRateLimitEvent(statusCode, decision.Model, wireModel, quota, true, decision.ResetAt)) } return decision } @@ -378,10 +555,15 @@ func applyAntigravityCooldown(store *auth.Store, account *auth.Account, statusCo // sized from Google's structured retry hint, with the generic 429 policy as // fallback. It reports whether a cooldown was actually recorded so callers can // classify the failure as rate limiting instead of a broken account. +// +// A connectivity probe is not user traffic: it keeps the cooldown write but +// stays out of the live scheduling penalty accumulator, otherwise one manual +// test button would bias fresh selection away from an otherwise healthy +// account. func ApplyAntigravityCooldown(store *auth.Store, account *auth.Account, statusCode int, body []byte, resp *http.Response, model string) bool { if statusCode != http.StatusTooManyRequests && statusCode != http.StatusServiceUnavailable { return false } - decision := applyAntigravityCooldown(store, account, statusCode, body, resp, model) + decision := applyAntigravityCooldownWithPenalty(store, account, statusCode, body, resp, model, false) return !decision.ResetAt.IsZero() } diff --git a/proxy/antigravity_quota_test.go b/proxy/antigravity_quota_test.go index a5ca1b1a6..e98a4e949 100644 --- a/proxy/antigravity_quota_test.go +++ b/proxy/antigravity_quota_test.go @@ -3,6 +3,7 @@ package proxy import ( "context" "io" + "math" "net/http" "net/http/httptest" "strings" @@ -11,6 +12,8 @@ import ( "time" "github.com/codex2api/auth" + "github.com/codex2api/database" + "github.com/gin-gonic/gin" ) const antigravityQuotaExhaustedBody = `{"error":{"code":429,"message":"You have exhausted your capacity on this model.","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"QUOTA_EXHAUSTED","domain":"cloudcode-pa.googleapis.com","metadata":{"model":"gemini-3.7-flash-tiered"}},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"4m50s"}]}}` @@ -243,3 +246,495 @@ func TestApplyAntigravityCooldownPrefersRetryAfterHeaderOverTheShortBackOff(t *t t.Fatalf("cooldown = %s, want the 12s Retry-After header instead of the short back-off", decision.Cooldown) } } + +// TestApplyAntigravityCooldownCoolsSharedWireModel 覆盖别名冷却联动： +// gemini-3.8-flash-low/medium/high 共用上游 backing gemini-3.8-flash-tiered， +// high 被限流后必须同时冷却 wire 模型，否则 low/medium 仍会选中同一账号并继续 +// 收集 429。 +func TestApplyAntigravityCooldownCoolsSharedWireModel(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + account := &auth.Account{DBID: 7401, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project"} + + decision := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityReasonOnlyRateLimitBody), nil, "gemini-3.8-flash-high") + if decision.Model != "gemini-3.8-flash-high" || decision.Reason != "rate_limited_model" { + t.Fatalf("decision = %+v", decision) + } + if !account.IsModelRateLimited("gemini-3.8-flash-high") { + t.Fatal("the public tier was not cooled down") + } + if !account.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("the shared wire model was not cooled down alongside the public tier") + } + if remaining, wireRemaining := account.ModelCooldownRemaining("gemini-3.8-flash-high"), account.ModelCooldownRemaining("gemini-3.8-flash-tiered"); remaining <= 0 || wireRemaining > remaining+time.Second { + t.Fatalf("wire cooldown %s drifted from the public cooldown %s", wireRemaining, remaining) + } + // 同 backing 的兄弟档位被拦住，不同 backing 与 Claude 模型不受影响。 + if !antigravityAccountModelRateLimited(account, "gemini-3.8-flash-low", "gemini-3.8-flash-tiered") { + t.Fatal("a sibling tier on the same backing must be blocked") + } + if antigravityAccountModelRateLimited(account, "gemini-3.6-flash-low", "gemini-3.6-flash-low") { + t.Fatal("a different backing must stay schedulable") + } + if antigravityAccountModelRateLimited(account, "claude-sonnet-4-6", "claude-sonnet-4-6") { + t.Fatal("Claude models must stay schedulable") + } +} + +// TestApplyAntigravityCooldownCoolsSharedWireModelOnUnstructured429 覆盖非结构化 429： +// 解析不出 google.rpc 结构的 429 走通用 429 策略，而通用策略只按 public id 落冷却， +// 同 wire 的兄弟档位仍会选中刚被限流的账号。这里必须补写与 public 同 deadline 的 +// wire 冷却。 +func TestApplyAntigravityCooldownCoolsSharedWireModelOnUnstructured429(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + // Antigravity 账号属于 relay 风格，通用 429 策略默认关闭其模型冷却， + // 这里显式打开才能观察到 wire 补写。 + store.SetModelCooldownSettings(database.ModelCooldownSettings{ + RelayMode: database.ModelCooldownModeFixed, + RelaySeconds: 60, + OAuthMode: database.ModelCooldownModeFixed, + OAuthSeconds: 60, + }) + account := &auth.Account{DBID: 7451, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project"} + + if _, ok := parseAntigravityQuotaError(http.StatusTooManyRequests, []byte(`{"error":{"message":"slow down"}}`)); ok { + t.Fatal("the body must stay unparseable so the generic 429 path is exercised") + } + decision := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(`{"error":{"message":"slow down"}}`), &http.Response{Header: http.Header{}}, "gemini-3.8-flash-high") + if decision.Scope != rateLimitScopeModel || decision.Model != "gemini-3.8-flash-high" || decision.ResetAt.IsZero() { + t.Fatalf("decision = %+v, want a model-scoped cooldown", decision) + } + if !account.IsModelRateLimited("gemini-3.8-flash-high") { + t.Fatal("the public tier was not cooled down") + } + if !account.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("an unstructured 429 must cool the shared wire model alongside the public tier") + } + if remaining, wireRemaining := account.ModelCooldownRemaining("gemini-3.8-flash-high"), account.ModelCooldownRemaining("gemini-3.8-flash-tiered"); remaining <= 0 || wireRemaining > remaining+time.Second { + t.Fatalf("wire cooldown %s drifted from the public cooldown %s", wireRemaining, remaining) + } + if !antigravityAccountModelRateLimited(account, "gemini-3.8-flash-low", "gemini-3.8-flash-tiered") { + t.Fatal("a sibling tier on the same backing must be blocked") + } + if antigravityAccountModelRateLimited(account, "gemini-3.6-flash-low", "gemini-3.6-flash-low") { + t.Fatal("a different backing must stay schedulable") + } + // 非结构化 429 依旧只记半个惩罚，且事件带上解析出的 wire 模型。 + if got := account.AntigravityPenalty(time.Now()); got < 0.49 || got > 0.51 { + t.Fatalf("penalty = %v, want the unclassified 0.5", got) + } +} + +// TestApplyAntigravityCooldownPenaltyWeights 覆盖调度惩罚的事件权重：明确限流 +// 记满，裸 RESOURCE_EXHAUSTED 与非结构化 429 记一半，共享容量与账号配额不记。 +func TestApplyAntigravityCooldownPenaltyWeights(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + newAccount := func(id int64) *auth.Account { + return &auth.Account{DBID: id, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project"} + } + + cases := []struct { + name string + statusCode int + body string + want float64 + }{ + {"explicit rate limit", http.StatusTooManyRequests, antigravityReasonOnlyRateLimitBody, 1.0}, + {"bare resource exhausted", http.StatusTooManyRequests, antigravityUnclassifiedRateLimitBody, 0.5}, + {"unstructured 429", http.StatusTooManyRequests, `{"error":{"message":"slow down"}}`, 0.5}, + {"quota exhausted", http.StatusTooManyRequests, antigravityQuotaExhaustedBody, 0.0}, + {"shared model capacity", http.StatusServiceUnavailable, antigravityModelCapacityBody, 0.0}, + } + for _, tc := range cases { + account := newAccount(7410) + applyAntigravityCooldown(store, account, tc.statusCode, []byte(tc.body), nil, "gemini-3.8-flash-high") + if got := account.AntigravityPenalty(time.Now()); math.Abs(got-tc.want) > 0.01 { + t.Fatalf("%s penalty = %v, want %v", tc.name, got, tc.want) + } + } + + // 同一账号连续被限流时惩罚累加，首绑避险据此把它压到队尾。 + account := newAccount(7411) + applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityReasonOnlyRateLimitBody), nil, "gemini-3.8-flash-high") + applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityUnclassifiedRateLimitBody), nil, "gemini-3.8-flash-high") + if got := account.AntigravityPenalty(time.Now()); got < 1.49 || got > 1.51 { + t.Fatalf("accumulated penalty = %v, want ~1.5", got) + } +} + +// TestApplyAntigravityCooldownAdminProbeSkipsPenalty 覆盖 admin 连通性测试隔离： +// 冷却照常写入（否则测试按钮探到的限流会被线上流量无视），但不进线上调度惩罚 +// 累加器（否则一次手工测试就会把健康账号挤出首绑）。 +func TestApplyAntigravityCooldownAdminProbeSkipsPenalty(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + probed := &auth.Account{DBID: 7420, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project"} + + if !ApplyAntigravityCooldown(store, probed, http.StatusTooManyRequests, []byte(antigravityReasonOnlyRateLimitBody), nil, "gemini-3.8-flash-high") { + t.Fatal("the admin probe did not record a cooldown") + } + if !probed.IsModelRateLimited("gemini-3.8-flash-high") || !probed.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("the admin probe must keep writing the public and wire model cooldowns") + } + if got := probed.AntigravityPenalty(time.Now()); got != 0 { + t.Fatalf("admin probe penalty = %v, want 0", got) + } + + live := &auth.Account{DBID: 7421, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project"} + applyAntigravityCooldown(store, live, http.StatusTooManyRequests, []byte(antigravityReasonOnlyRateLimitBody), nil, "gemini-3.8-flash-high") + if got := live.AntigravityPenalty(time.Now()); got <= 0 { + t.Fatalf("live dispatch penalty = %v, want a recorded event", got) + } +} + +// TestGeminiNativeCountTokensAppliesModelCooldown 覆盖 native countTokens 的 +// 429 漏路径：该分支过去只计数与排除账号，从不写模型冷却。 +func TestGeminiNativeCountTokensAppliesModelCooldown(t *testing.T) { + gin.SetMode(gin.TestMode) + + var hits int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&hits, 1) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = io.WriteString(w, antigravityReasonOnlyRateLimitBody) + })) + defer server.Close() + previous := antigravityOAuthEndpointBases + antigravityOAuthEndpointBases = []string{server.URL} + t.Cleanup(func() { antigravityOAuthEndpointBases = previous }) + + store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 2}) + defer store.Stop() + store.SetMaxRetries(0) + store.SetMaxRateLimitRetries(0) + account := &auth.Account{ + DBID: 7430, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project", + Models: []string{"gemini-3.7-flash-tiered", "gemini-3.8-flash-tiered"}, + } + store.AddAccount(account) + + handler := NewHandler(store, nil, nil, nil) + row := &database.APIKeyRow{ID: 7, Limits: database.APIKeyLimits{UpstreamChannel: database.UpstreamChannelAntigravity}} + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(contextAPIKeyRow, row) + c.Next() + }) + router.POST("/v1beta/models/*action", handler.GeminiModelsAction) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3.8-flash-high:countTokens", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)) + req.Header.Set("Content-Type", "application/json") + router.ServeHTTP(rec, req) + + if atomic.LoadInt32(&hits) != 1 { + t.Fatalf("upstream hits = %d, want one attempt", hits) + } + if rec.Code != http.StatusTooManyRequests { + t.Fatalf("status = %d body=%s, want the upstream 429 surfaced", rec.Code, rec.Body.String()) + } + if !account.IsModelRateLimited("gemini-3.8-flash-high") { + t.Fatal("countTokens 429 did not cool down the public model") + } + if !account.IsModelRateLimited("gemini-3.8-flash-tiered") { + t.Fatal("countTokens 429 did not cool down the shared wire model") + } + if got := account.AntigravityPenalty(time.Now()); got <= 0 { + t.Fatalf("countTokens 429 penalty = %v, want a recorded event", got) + } +} + +// antigravityLongRateLimitBody 是 Google 明确告知 290s 恢复窗口的限流：低于 +// antigravityQuotaExhaustedThreshold，因此仍按 rate_limit 分类。 +const antigravityLongRateLimitBody = `{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"RATE_LIMIT_EXCEEDED","metadata":{"model":"gemini-3.8-flash-tiered"}},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"290s"}]}}` + +func newAntigravityCooldownTestAccount(id int64) *auth.Account { + return &auth.Account{DBID: id, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project"} +} + +// TestApplyAntigravityCooldownSharedWireDeadlineNotShortenedByBrieferEvent 覆盖共享 +// wire 冷却的 deadline 单调性：high 收到 290s 的 RATE_LIMIT_EXCEEDED 后，同账号 low +// 的裸 RESOURCE_EXHAUSTED（5s 回退）只能延长、不能把共享 wire 冷却改短，否则 5s 之后 +// low/medium 又会选中刚被限流的凭据。public 冷却保持既有覆盖语义不变。 +func TestApplyAntigravityCooldownSharedWireDeadlineNotShortenedByBrieferEvent(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + account := newAntigravityCooldownTestAccount(7461) + + long := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityLongRateLimitBody), nil, "gemini-3.8-flash-high") + if long.Cooldown < 280*time.Second { + t.Fatalf("long cooldown = %s, want the 290s retry hint", long.Cooldown) + } + if wireRemaining := account.ModelCooldownRemaining("gemini-3.8-flash-tiered"); wireRemaining < 280*time.Second { + t.Fatalf("wire cooldown = %s after the long event, want ~290s", wireRemaining) + } + + // 同账号兄弟档位的裸 RESOURCE_EXHAUSTED：只有 5s 的回退窗口。 + short := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityUnclassifiedRateLimitBody), nil, "gemini-3.8-flash-low") + if short.Cooldown > 10*time.Second { + t.Fatalf("short cooldown = %s, want the brief back-off", short.Cooldown) + } + if wireRemaining := account.ModelCooldownRemaining("gemini-3.8-flash-tiered"); wireRemaining < 60*time.Second { + t.Fatalf("wire cooldown = %s after a briefer event, want the longer deadline kept", wireRemaining) + } + if !antigravityAccountModelRateLimited(account, "gemini-3.8-flash-medium", "gemini-3.8-flash-tiered") { + t.Fatal("a sibling tier on the same backing must stay blocked") + } + // public 冷却保持既有语义：短事件照常覆盖同名 public 冷却，本次改动不触碰。 + if publicRemaining := account.ModelCooldownRemaining("gemini-3.8-flash-low"); publicRemaining <= 0 || publicRemaining > 10*time.Second { + t.Fatalf("public cooldown = %s, want the brief back-off to win on the public key", publicRemaining) + } +} + +// TestApplyAntigravityCooldownSharedWireDeadlineExtendsWithLongerEvent 是反向用例： +// 先短后长必须延长共享 wire 冷却，AtLeast 变体不是「先到先得」。 +func TestApplyAntigravityCooldownSharedWireDeadlineExtendsWithLongerEvent(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + account := newAntigravityCooldownTestAccount(7462) + + first := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityUnclassifiedRateLimitBody), nil, "gemini-3.8-flash-low") + if first.Cooldown > 10*time.Second { + t.Fatalf("first cooldown = %s, want the brief back-off", first.Cooldown) + } + if wireRemaining := account.ModelCooldownRemaining("gemini-3.8-flash-tiered"); wireRemaining > 10*time.Second { + t.Fatalf("wire cooldown = %s after the brief event, want ~5s", wireRemaining) + } + + applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityLongRateLimitBody), nil, "gemini-3.8-flash-medium") + if wireRemaining := account.ModelCooldownRemaining("gemini-3.8-flash-tiered"); wireRemaining < 280*time.Second { + t.Fatalf("wire cooldown = %s, want the longer deadline to extend the shared cooldown", wireRemaining) + } +} + +// TestApplyAntigravityCooldownGeminiFamilyDeadlineNotShortened 覆盖 Gemini 家族冷却 +// 键：短窗口的 QUOTA_EXHAUSTED 不能把家族冷却改短，家族键仍拦住其它 Gemini 档位。 +func TestApplyAntigravityCooldownGeminiFamilyDeadlineNotShortened(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + account := newAntigravityCooldownTestAccount(7463) + + longBody := strings.Replace(antigravityQuotaExhaustedBody, `"4m50s"`, `"10m"`, 1) + applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(longBody), nil, "gemini-3.7-flash-high") + if remaining := account.ModelCooldownRemaining(antigravityGeminiFamilyCooldownKey); remaining < 9*time.Minute { + t.Fatalf("family cooldown = %s after the long event, want ~10m", remaining) + } + + shortBody := strings.Replace(antigravityQuotaExhaustedBody, `"4m50s"`, `"60s"`, 1) + applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(shortBody), nil, "gemini-3.6-flash-low") + if remaining := account.ModelCooldownRemaining(antigravityGeminiFamilyCooldownKey); remaining < 9*time.Minute { + t.Fatalf("family cooldown = %s after a shorter quota window, want the longer deadline kept", remaining) + } + if !antigravityAccountModelRateLimited(account, "gemini-3.5-pro-high", "gemini-3.5-pro-high") { + t.Fatal("the Gemini family cooldown must keep blocking other Gemini tiers") + } +} + +// TestApplyAntigravityCooldownUnstructured429KeepsLongerWireDeadline 覆盖非结构化 429 +// 补写 wire 冷却的同一约束：通用 429 策略给出的较短 deadline 不得覆盖同账号上更长的 +// 共享 wire 冷却。 +func TestApplyAntigravityCooldownUnstructured429KeepsLongerWireDeadline(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + store.SetModelCooldownSettings(database.ModelCooldownSettings{ + RelayMode: database.ModelCooldownModeFixed, + RelaySeconds: 60, + OAuthMode: database.ModelCooldownModeFixed, + OAuthSeconds: 60, + }) + account := newAntigravityCooldownTestAccount(7464) + + // 先由 high 的长限流把共享 wire 冷却写到 ~290s。 + if _, ok := parseAntigravityQuotaError(http.StatusTooManyRequests, []byte(antigravityLongRateLimitBody)); !ok { + t.Fatal("the long rate limit body must stay parseable") + } + applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityLongRateLimitBody), nil, "gemini-3.8-flash-high") + + if _, ok := parseAntigravityQuotaError(http.StatusTooManyRequests, []byte(`{"error":{"message":"slow down"}}`)); ok { + t.Fatal("the body must stay unparseable so the generic 429 path is exercised") + } + decision := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(`{"error":{"message":"slow down"}}`), &http.Response{Header: http.Header{}}, "gemini-3.8-flash-low") + if decision.ResetAt.IsZero() { + t.Fatalf("decision = %+v, want the generic 429 policy to record a cooldown", decision) + } + if wireRemaining := account.ModelCooldownRemaining("gemini-3.8-flash-tiered"); wireRemaining < 120*time.Second { + t.Fatalf("wire cooldown = %s after an unstructured 429, want the longer deadline kept", wireRemaining) + } + if !antigravityAccountModelRateLimited(account, "gemini-3.8-flash-medium", "gemini-3.8-flash-tiered") { + t.Fatal("a sibling tier on the same backing must stay blocked") + } +} + +// antigravityArgonAccount 返回一个只发现到 gemini-4-argon-high 这块上游 backing 的账号： +// gemini-4-argon-low/medium/high 三个公开档位都骑在它上面，其中 high 的公开名与 wire 名 +// 完全相同，public 冷却键本身就是共享 wire 冷却键。 +func antigravityArgonAccount(id int64) *auth.Account { + return &auth.Account{ + DBID: id, + UpstreamType: auth.UpstreamAntigravity, + AccessToken: "google-token", + AntigravityProjectID: "google-project", + Models: []string{"gemini-4-argon-high"}, + } +} + +// TestAntigravitySharedWireCooldownKeyClassification 覆盖共享冷却键判定本身：公开名与 +// 解析出的 wire 同名、或目录里另有档位将 model 当作 backing 时成立；公开档位解析到自己 +// 独占且与公开名不同的 backing（或干脆解析不出 wire）时不成立——否则会把普通覆盖语义 +// 无差别改成单调写，短事件再也无法缩短任何模型冷却。 +func TestAntigravitySharedWireCooldownKeyClassification(t *testing.T) { + cases := []struct { + model string + wireModel string + want bool + }{ + // Argon：公开名即共享 backing 名。 + {"gemini-4-argon-high", "gemini-4-argon-high", true}, + {"gemini-4-argon-high", "", true}, + // 兄弟档位的公开名不是任何档位的 backing。 + {"gemini-4-argon-low", "gemini-4-argon-high", false}, + // 原始 backing id 被三个公开档位共享。 + {"gemini-3.8-flash-tiered", "", true}, + {"gemini-3.7-flash-tiered", "", true}, + // 逻辑档位的 backing 同时是固定档位的公开名。 + {"gemini-3.6-flash-low", "", true}, + // 解析出的 wire 与公开名同名：按规则一律按共享键写入，保证单调性。 + {"claude-sonnet-4-6", "claude-sonnet-4-6", true}, + // 解析不到 wire、目录里也没有别的档位以它为 backing：保持覆盖语义。 + {"claude-sonnet-4-6", "", false}, + {"some-discovered-raw-model", "", false}, + // 公开档位解析到自己独占且与公开名不同的 backing：保持覆盖语义。 + {"gemini-3.8-flash-low", "gemini-3.8-flash-tiered", false}, + {"gemini-3.1-pro-high", "gemini-pro-agent", false}, + {"", "", false}, + } + for _, tc := range cases { + if got := antigravitySharedWireCooldownKey(tc.model, tc.wireModel); got != tc.want { + t.Fatalf("antigravitySharedWireCooldownKey(%q, %q) = %v, want %v", tc.model, tc.wireModel, got, tc.want) + } + } +} + +// TestApplyAntigravityCooldownSharedNamePublicModelNotShortened 覆盖 P1-1： +// gemini-4-argon-high 的公开名与共享 wire 键同名，先由 low 的 290s 限流把共享冷却写到 +// ~290s，随后 high 自己的 5s 裸 RESOURCE_EXHAUSTED 走的是无条件覆盖的 public 写入路径， +// 必须同样受 AtLeast 单调性保护，不能把兄弟档位的长冷却改短。 +func TestApplyAntigravityCooldownSharedNamePublicModelNotShortened(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + account := antigravityArgonAccount(7471) + + if wire, ok := antigravityResolvePublicModelForAccount(account, "gemini-4-argon-high"); !ok || wire != "gemini-4-argon-high" { + t.Fatalf("gemini-4-argon-high resolved to %q, ok=%v, want the shared backing itself", wire, ok) + } + + first := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityLongRateLimitBody), nil, "gemini-4-argon-low") + if first.Cooldown < 280*time.Second { + t.Fatalf("sibling cooldown = %s, want the 290s retry hint", first.Cooldown) + } + if remaining := account.ModelCooldownRemaining("gemini-4-argon-high"); remaining < 280*time.Second { + t.Fatalf("shared key cooldown = %s after the sibling event, want ~290s", remaining) + } + + short := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityUnclassifiedRateLimitBody), nil, "gemini-4-argon-high") + if remaining := account.ModelCooldownRemaining("gemini-4-argon-high"); remaining < 60*time.Second { + t.Fatalf("shared key cooldown = %s after the same-name brief event, want the longer deadline kept", remaining) + } + // 决策本身也要报告生效的长 deadline，否则上层会按 5s 之后重新调度该账号。 + if short.ResetAt.IsZero() || time.Until(short.ResetAt) < 60*time.Second { + t.Fatalf("decision = %+v, want the kept longer deadline", short) + } + if !antigravityAccountModelRateLimited(account, "gemini-4-argon-medium", "gemini-4-argon-high") { + t.Fatal("a sibling tier on the same backing must stay blocked") + } +} + +// TestApplyAntigravityCooldownSharedNamePublicModelExtendsWithLongerEvent 是反向用例： +// high 先吃到 5s 短冷却，随后兄弟档位的 290s 限流必须能把同名共享键延长。 +func TestApplyAntigravityCooldownSharedNamePublicModelExtendsWithLongerEvent(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + account := antigravityArgonAccount(7472) + + first := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityUnclassifiedRateLimitBody), nil, "gemini-4-argon-high") + if first.Cooldown > 10*time.Second { + t.Fatalf("first cooldown = %s, want the brief back-off", first.Cooldown) + } + if remaining := account.ModelCooldownRemaining("gemini-4-argon-high"); remaining > 10*time.Second { + t.Fatalf("shared key cooldown = %s after the brief event, want ~5s", remaining) + } + + extended := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityLongRateLimitBody), nil, "gemini-4-argon-medium") + if extended.Cooldown < 280*time.Second { + t.Fatalf("extended decision = %+v, want the 290s retry hint", extended) + } + if remaining := account.ModelCooldownRemaining("gemini-4-argon-high"); remaining < 280*time.Second { + t.Fatalf("shared key cooldown = %s, want the longer deadline to extend it", remaining) + } +} + +// TestApplyAntigravityCooldownSharedNamePublicModelOnUnstructured429 覆盖非结构化 429 +// 的同一约束：通用 429 策略按覆盖语义写 public 冷却，而 public 名就是共享 wire 键时这次 +// 覆盖会把兄弟档位的长冷却改短，必须写回。 +func TestApplyAntigravityCooldownSharedNamePublicModelOnUnstructured429(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + store.SetModelCooldownSettings(database.ModelCooldownSettings{ + RelayMode: database.ModelCooldownModeFixed, + RelaySeconds: 60, + OAuthMode: database.ModelCooldownModeFixed, + OAuthSeconds: 60, + }) + account := antigravityArgonAccount(7473) + + applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityLongRateLimitBody), nil, "gemini-4-argon-low") + if remaining := account.ModelCooldownRemaining("gemini-4-argon-high"); remaining < 280*time.Second { + t.Fatalf("shared key cooldown = %s after the sibling event, want ~290s", remaining) + } + + if _, ok := parseAntigravityQuotaError(http.StatusTooManyRequests, []byte(`{"error":{"message":"slow down"}}`)); ok { + t.Fatal("the body must stay unparseable so the generic 429 path is exercised") + } + decision := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(`{"error":{"message":"slow down"}}`), &http.Response{Header: http.Header{}}, "gemini-4-argon-high") + if decision.Scope != rateLimitScopeModel || decision.Model != "gemini-4-argon-high" { + t.Fatalf("decision = %+v, want a model-scoped cooldown", decision) + } + if remaining := account.ModelCooldownRemaining("gemini-4-argon-high"); remaining < 120*time.Second { + t.Fatalf("shared key cooldown = %s after an unstructured 429, want the longer deadline kept", remaining) + } + if !antigravityAccountModelRateLimited(account, "gemini-4-argon-medium", "gemini-4-argon-high") { + t.Fatal("a sibling tier on the same backing must stay blocked") + } +} + +// TestApplyAntigravityCooldownUnsharedPublicModelKeepsOverwriteSemantics 是回归护栏： +// 公开档位解析到自己独占、且与公开名不同的 backing 时，短事件仍按既有语义覆盖同名 +// public 冷却，单调性保护只作用于共享键（wire / family / 与 wire 同名的公开名）。 +func TestApplyAntigravityCooldownUnsharedPublicModelKeepsOverwriteSemantics(t *testing.T) { + store := newProxyPremiumTestStore() + defer store.Stop() + account := newAntigravityCooldownTestAccount(7474) + + if wire, ok := antigravityResolvePublicModelForAccount(account, "gemini-3.8-flash-low"); !ok || wire != "gemini-3.8-flash-tiered" { + t.Fatalf("gemini-3.8-flash-low resolved to %q, ok=%v, want the shared backing", wire, ok) + } + + applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityLongRateLimitBody), nil, "gemini-3.8-flash-low") + if remaining := account.ModelCooldownRemaining("gemini-3.8-flash-low"); remaining < 280*time.Second { + t.Fatalf("public cooldown = %s, want the 290s retry hint", remaining) + } + short := applyAntigravityCooldown(store, account, http.StatusTooManyRequests, []byte(antigravityUnclassifiedRateLimitBody), nil, "gemini-3.8-flash-low") + if short.Cooldown > 10*time.Second { + t.Fatalf("decision = %+v, want the brief back-off", short) + } + if remaining := account.ModelCooldownRemaining("gemini-3.8-flash-low"); remaining <= 0 || remaining > 10*time.Second { + t.Fatalf("public cooldown = %s, want the brief event to keep overwriting a non-shared public key", remaining) + } + // 同一次短事件仍不得改短共享 wire 键上的长冷却。 + if remaining := account.ModelCooldownRemaining("gemini-3.8-flash-tiered"); remaining < 280*time.Second { + t.Fatalf("wire cooldown = %s, want the longer deadline kept on the shared key", remaining) + } +} diff --git a/proxy/antigravity_ws_selection_test.go b/proxy/antigravity_ws_selection_test.go new file mode 100644 index 000000000..10bb37408 --- /dev/null +++ b/proxy/antigravity_ws_selection_test.go @@ -0,0 +1,63 @@ +package proxy + +import ( + "testing" + + "github.com/codex2api/auth" + "github.com/codex2api/database" +) + +// TestResponsesWSAntigravityChannelSelectionSkipsWeighting 覆盖 Responses-over-WebSocket +// 入口（forwardResponsesWebSocketTurn）的选号上下文：该入口的候选池 +// accountFilterForResponsesWebSocket 恒不含 Antigravity 账号 +// （OpenAIResponsesUsesUpstreamWebsocket 对 AG 恒为 false，accountFilterForModel +// 直接拒绝 AG），因此 Antigravity 渠道 Key 打到该入口时 SelectionOptions.Provider +// 必须为空，否则 Store 侧会对 Codex relay 候选套用 Antigravity 加权打分，并绕过 +// spread / HRW 一致哈希。 +// +// 覆盖边界：驱动完整 WS handler 需要 websocket 升级与上游账号凭据，现有测试基础 +// 设施成本过高，这里对入口实际调用的 helper 组装做单元测试——入口处 +// selectionOpts 只来自 selectionOptionsForCandidatePool(..., true)，并原样传给 +// nextRetryAccountForContinuationWithSelection / +// nextRetryAccountForSessionWithDispatchGuard +// 两条选号路径，helper 断言即等价于入口传出的 SelectionOptions.Provider 为空。 +func TestResponsesWSAntigravityChannelSelectionSkipsWeighting(t *testing.T) { + c := newChannelTestGin(database.UpstreamChannelAntigravity) + + // 前置事实：Responses WS 候选池确实不含 Antigravity 账号，Provider 为空才是正确语义。 + antigravityAccount := &auth.Account{ + DBID: 9140, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", + AntigravityProjectID: "google-project", Status: auth.StatusReady, + Models: []string{"gemini-3.8-flash-tiered"}, + } + if antigravityAccount.OpenAIResponsesUsesUpstreamWebsocket() { + t.Fatal("Antigravity accounts must never dial the upstream Responses WebSocket") + } + if accountFilterForResponsesWebSocket("gpt-5.5")(antigravityAccount) { + t.Fatal("the Responses WS candidate pool must keep excluding Antigravity accounts") + } + + // 入口组装：Antigravity 渠道 Key + antigravityExcluded=true → Provider 为空、Policy 保留。 + cases := []struct { + name string + policy auth.DispatchPolicy + }{ + {name: "standard", policy: auth.DispatchPolicyStandard}, + {name: "spark", policy: auth.DispatchPolicySpark}, + } + for _, tc := range cases { + opts := selectionOptionsForCandidatePool(c, "gpt-5.5", tc.policy, true) + if opts.Provider != "" { + t.Fatalf("%s: Provider = %q, want empty so the Antigravity weighting stays off", tc.name, opts.Provider) + } + if opts.Policy != tc.policy { + t.Fatalf("%s: Policy = %v, want it preserved", tc.name, opts.Policy) + } + } + + // 对照：候选池可能含 Antigravity 账号的普通入口仍带渠道，加权避险照常启用。 + ordinary := selectionOptionsForCandidatePool(c, "gemini-3.8-flash-high", auth.DispatchPolicyStandard, false) + if ordinary.Provider != database.UpstreamChannelAntigravity { + t.Fatalf("ordinary Responses options Provider = %q, want the request channel", ordinary.Provider) + } +} diff --git a/proxy/handler.go b/proxy/handler.go index a0c7acb1a..c742f6cf1 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -110,15 +110,54 @@ func (h *Handler) nextAccountForSessionWithFilter(sessionID string, apiKeyID int } func (h *Handler) nextAccountForSessionWithDispatch(sessionID string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, policy auth.DispatchPolicy) (*auth.Account, string) { - account, proxyURL, _ := h.nextAccountForSessionWithDispatchGuard(sessionID, apiKeyID, exclude, filter, policy) + account, proxyURL, _ := h.nextAccountForSessionWithDispatchGuard(sessionID, apiKeyID, exclude, filter, auth.SelectionOptions{Policy: policy}) return account, proxyURL } -func (h *Handler) nextAccountForSessionWithDispatchGuard(sessionID string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, policy auth.DispatchPolicy) (*auth.Account, string, auth.SessionAffinityGuard) { +func (h *Handler) nextAccountForSessionWithDispatchGuard(sessionID string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, opts auth.SelectionOptions) (*auth.Account, string, auth.SessionAffinityGuard) { if h == nil || h.store == nil { return nil, "", auth.SessionAffinityGuard{} } - return h.store.NextForSessionWithDispatchGuard(sessionID, apiKeyID, exclude, filter, policy) + return h.store.NextForSessionWithDispatchGuardForProvider(opts, sessionID, apiKeyID, exclude, filter) +} + +// selectionOptionsForRequest 依据当前请求渠道与 public model 组装选号上下文：Antigravity +// 渠道请求据此在 Store 侧启用首绑 / Fresh 加权避险，其余渠道 Provider 非 antigravity， +// Store 走既有选号逻辑，行为不变。 +func selectionOptionsForRequest(c *gin.Context, model string, policy auth.DispatchPolicy) auth.SelectionOptions { + return selectionOptionsForRequestWithProvider(c, model, policy, "") +} + +// selectionOptionsForRequestWithProvider 与 selectionOptionsForRequest 相同，只是用 +// forceProvider 覆盖请求渠道作为选号渠道（forceProvider 为空表示沿用请求渠道）。 +// native Gemini（generateContent / countTokens）的候选池被 antigravityOAuthChannelAccountFilter +// 硬限定为 Antigravity OAuth：auto 渠道 Key 走到这两个入口时 requestUpstreamChannel 为空， +// 不覆盖就会把加权避险整个关掉，所以这两个入口按候选池的真实归属强制 +// Provider=antigravity。 +func selectionOptionsForRequestWithProvider(c *gin.Context, model string, policy auth.DispatchPolicy, forceProvider string) auth.SelectionOptions { + provider := strings.TrimSpace(forceProvider) + if provider == "" { + provider = requestUpstreamChannel(c) + } + return auth.SelectionOptions{ + Provider: provider, + Model: strings.TrimSpace(model), + Policy: policy, + } +} + +// selectionOptionsForCandidatePool 组装选号上下文，并由入口声明候选池是否已显式排除 +// Antigravity 账号（antigravityExcluded）。 +// +// 排除时只保留 Policy、不带 Provider：这类入口（compact 执行器、native remote +// compaction）的候选只剩 Codex 官方与 relay 账号，而 Antigravity 渠道 Key 打到这些 +// 入口时 requestUpstreamChannel 仍是 antigravity，沿用请求渠道就会让 Store 侧对 +// Codex relay 候选套用 Antigravity 加权打分，并绕过 spread / HRW 一致哈希。 +func selectionOptionsForCandidatePool(c *gin.Context, model string, policy auth.DispatchPolicy, antigravityExcluded bool) auth.SelectionOptions { + if antigravityExcluded { + return auth.SelectionOptions{Policy: policy} + } + return selectionOptionsForRequest(c, model, policy) } func dispatchPolicyForModel(model string) auth.DispatchPolicy { @@ -135,6 +174,46 @@ func (h *Handler) withModelCooldownFilter(ctx context.Context, model string, fil return h.store.WithModelCooldownFilterContext(ctx, model, filter) } +// antigravityRequestCooldownKeys 解析 public model 对应的 wireModel 与（Gemini 家族） +// 共享 familyKey，供跨实例三键冷却检查使用。wireModel 取自模型目录 +// （AntigravityWireModelID），与 antigravityResolvePublicModelForAccount 对受支持账号返回 +// 的 wire 完全一致；账号是否具备该 wire 支撑已由 antigravityChannelAccountFilter 单独把关， +// 这里只负责冷却键。非 Gemini 模型的 familyKey 为空串。 +func antigravityRequestCooldownKeys(model string) (wireModel string, familyKey string) { + wireModel, _ = AntigravityWireModelID(model) + if antigravityIsGeminiModel(model) || antigravityIsGeminiModel(wireModel) { + familyKey = antigravityGeminiFamilyCooldownKey + } + return wireModel, familyKey +} + +// withRequestModelCooldownFilter 应用跨实例模型冷却过滤器，按下游 Key 的渠道限定分三档： +// - antigravity：额外解析 public model 对应的 wireModel 及（Gemini 家族）共享 familyKey， +// 用 WithAntigravityCooldownFilterContext 对三个 key 一并做跨实例检查——任一 key 由其它 +// 实例写入的冷却都能在本实例选号时过滤掉该账号，闭合「本地内存已限流但跨实例缓存未拦截」 +// 的缺口； +// - auto（渠道未限定）：候选池里既可能有 Antigravity 账号（native Gemini 入口硬限定为 +// Antigravity OAuth，Responses / Chat / Messages 也允许），也可能有其它渠道账号，改用 +// WithAntigravityAwareCooldownFilterContext 按候选账号类型分别取三键 / 单键检查—— +// Antigravity 账号同样能看到跨实例 wire / family 冷却，非 Antigravity 账号行为不变； +// - 其它显式渠道（codex / grok / claude）：候选池不含 Antigravity 账号，沿用单一 +// publicModel 检查，行为不变。 +func (h *Handler) withRequestModelCooldownFilter(c *gin.Context, ctx context.Context, model string, filter auth.AccountFilter) auth.AccountFilter { + if h == nil || h.store == nil { + return filter + } + switch requestUpstreamChannel(c) { + case database.UpstreamChannelAntigravity: + wireModel, familyKey := antigravityRequestCooldownKeys(model) + return h.store.WithAntigravityCooldownFilterContext(ctx, model, wireModel, familyKey, filter) + case database.UpstreamChannelAuto: + wireModel, familyKey := antigravityRequestCooldownKeys(model) + return h.store.WithAntigravityAwareCooldownFilterContext(ctx, model, wireModel, familyKey, filter) + default: + return h.store.WithModelCooldownFilterContext(ctx, model, filter) + } +} + func (h *Handler) shouldUseWebsocketForHTTP() bool { if h == nil { return false @@ -4046,7 +4125,7 @@ func (h *Handler) Responses(c *gin.Context) { } else { accountFilter = accountFilterForResponsesModelWithOriginal(logModel, effectiveModel, allowCodexAccounts) } - accountFilter = h.withModelCooldownFilter(c.Request.Context(), effectiveModel, accountFilter) + accountFilter = h.withRequestModelCooldownFilter(c, c.Request.Context(), effectiveModel, accountFilter) if continuationUnavailable { accountFilter = responsesContinuationRelayFilter(accountFilter) } @@ -4100,6 +4179,10 @@ func (h *Handler) Responses(c *gin.Context) { capacityShedRetries := map[int64]int{} dispatchPolicy := dispatchPolicyForModel(effectiveModel) + // native remote compaction 的候选池被 accountFilterForInlineCompactionModelWithOriginal + // 硬排除 Antigravity 账号，此时不带渠道，避免 Antigravity 渠道 Key 对 Codex relay + // 候选套用加权避险；普通 Responses 请求的候选池仍可能含 Antigravity 账号，保持渠道。 + selectionOpts := selectionOptionsForCandidatePool(c, effectiveModel, dispatchPolicy, nativeRemoteCompactionV2) var affinityGuard auth.SessionAffinityGuard var selectionErr error grokQualityAttempts := 0 @@ -4113,11 +4196,12 @@ func (h *Handler) Responses(c *gin.Context) { if account != nil { stickyProxyURL = account.GetProxyURL() } else if continuationUnavailable && !relayContinuationAttempted { - account, stickyProxyURL, affinityGuard = h.nextAccountForSessionWithDispatchGuard(affinityKey, apiKeyID, retryExclusions.ForSelection(), accountFilter, dispatchPolicy) + account, stickyProxyURL, affinityGuard = h.nextAccountForSessionWithDispatchGuard(affinityKey, apiKeyID, retryExclusions.ForSelection(), accountFilter, selectionOpts) } else if turnContinuationPinned { - account, stickyProxyURL, selectionErr = h.nextRetryAccountForContinuationWithDispatch(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, dispatchPolicy) + // 续链重试同样传 SelectionOptions，避免重试把选号上下文丢掉。 + account, stickyProxyURL, selectionErr = h.nextRetryAccountForContinuationWithSelection(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts) } else { - account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, dispatchPolicy) + account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts) } } if account == nil { @@ -6055,7 +6139,7 @@ func (h *Handler) ResponsesCompact(c *gin.Context) { // compact 同时允许官方 Codex OAuth 账号与中转（OpenAI Responses API）账号： // 中转账号会命中上游自身的 /responses/compact，使仅接入中转的用户也能压缩（issue #174）。 accountFilter := accountFilterForCompactResponsesModelWithOriginal(routingModel, effectiveModel, modelIDInList(effectiveModel, SupportedModelIDs(c.Request.Context(), h.db))) - accountFilter = h.withModelCooldownFilter(c.Request.Context(), effectiveModel, accountFilter) + accountFilter = h.withRequestModelCooldownFilter(c, c.Request.Context(), effectiveModel, accountFilter) accountFilter = excludeClaudeAccountsFilter(accountFilter) if continuationUnavailable { accountFilter = responsesContinuationRelayFilter(accountFilter) @@ -6090,6 +6174,10 @@ func (h *Handler) ResponsesCompact(c *gin.Context) { relayContinuationAttempted := false dispatchPolicy := dispatchPolicyForModel(effectiveModel) + // compact 执行器的候选池被 accountFilterForCompactResponsesModelWithOriginal 硬排除 + // Antigravity 账号，因此不带渠道：Antigravity 渠道 Key 在这里也不该对 Codex relay + // 候选套用加权避险。 + selectionOpts := selectionOptionsForCandidatePool(c, effectiveModel, dispatchPolicy, true) for attempt := 0; ; attempt++ { var account *auth.Account var stickyProxyURL string @@ -6102,7 +6190,7 @@ func (h *Handler) ResponsesCompact(c *gin.Context) { } } if account == nil { - account, stickyProxyURL, affinityGuard = h.nextAccountForSessionWithDispatchGuard(affinityKey, apiKeyID, retryExclusions.ForSelection(), accountFilter, dispatchPolicy) + account, stickyProxyURL, affinityGuard = h.nextAccountForSessionWithDispatchGuard(affinityKey, apiKeyID, retryExclusions.ForSelection(), accountFilter, selectionOpts) } if account == nil { if continuousRetryCommitExpired(c, continuousRetryProtocolResponses) { @@ -6129,7 +6217,7 @@ func (h *Handler) ResponsesCompact(c *gin.Context) { sendResponseContextUnavailable(c, continuationStatus, continuationReason) return } - account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, dispatchPolicy) + account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts) if account == nil { if writeSchedulerQueueError(c, selectionErr, continuousRetryProtocolResponses) { return @@ -6905,7 +6993,7 @@ func (h *Handler) ChatCompletions(c *gin.Context) { // /v1/chat/completions 同时允许官方 Codex OAuth 账号与中转（OpenAI Responses API）账号： // 翻译后的请求体本身就是 Responses 形态，中转账号直接以 HTTP 转发（issue #181）。 accountFilter := accountFilterForResponsesModelWithOriginal(logModel, effectiveModel, modelIDInList(effectiveModel, SupportedModelIDs(c.Request.Context(), h.db))) - accountFilter = h.withModelCooldownFilter(c.Request.Context(), effectiveModel, accountFilter) + accountFilter = h.withRequestModelCooldownFilter(c, c.Request.Context(), effectiveModel, accountFilter) accountFilter = h.applyUpstreamChannelFilter(c, effectiveModel, accountFilter) accountFilter = excludeClaudeAccountsFilter(accountFilter) accountFilter = h.applyScopeBudgetFilter(c, accountFilter) @@ -6944,6 +7032,7 @@ func (h *Handler) ChatCompletions(c *gin.Context) { capacityShedRetries := map[int64]int{} dispatchPolicy := dispatchPolicyForModel(effectiveModel) + selectionOpts := selectionOptionsForRequest(c, effectiveModel, dispatchPolicy) var affinityGuard auth.SessionAffinityGuard var selectionErr error grokQualityAttempts := 0 @@ -6951,7 +7040,7 @@ func (h *Handler) ChatCompletions(c *gin.Context) { account, stickyProxyURL, retainedHTTPFallback := wsHTTPFallback.Take() if !retainedHTTPFallback { affinityGuard = auth.SessionAffinityGuard{} - account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, dispatchPolicy) + account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts) } if account == nil { if writeSchedulerQueueError(c, selectionErr, continuousRetryProtocolChat) { diff --git a/proxy/handler_anthropic.go b/proxy/handler_anthropic.go index f0fdde762..a6e0708a7 100644 --- a/proxy/handler_anthropic.go +++ b/proxy/handler_anthropic.go @@ -681,7 +681,7 @@ func (h *Handler) Messages(c *gin.Context) { return fixedTierFilter(account) } } - accountFilter = h.withModelCooldownFilter(c.Request.Context(), effectiveModel, accountFilter) + accountFilter = h.withRequestModelCooldownFilter(c, c.Request.Context(), effectiveModel, accountFilter) accountFilter = h.applyUpstreamChannelFilter(c, effectiveModel, accountFilter) accountFilter = h.applyScopeBudgetFilter(c, accountFilter) // scope 并发位在选中账号后才能占，请求退出时统一释放（issue #439 v2）。 @@ -729,6 +729,7 @@ func (h *Handler) Messages(c *gin.Context) { }() capacityShedRetries := map[int64]int{} + selectionOpts := selectionOptionsForRequest(c, effectiveModel, auth.DispatchPolicyStandard) var affinityGuard auth.SessionAffinityGuard var selectionErr error grokQualityAttempts := 0 @@ -737,7 +738,7 @@ func (h *Handler) Messages(c *gin.Context) { account, stickyProxyURL, retainedHTTPFallback := wsHTTPFallback.Take() if !retainedHTTPFallback { affinityGuard = auth.SessionAffinityGuard{} - account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter) + account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts) } if account == nil { if writeSchedulerQueueError(c, selectionErr, continuousRetryProtocolAnthropic) { diff --git a/proxy/handler_gemini.go b/proxy/handler_gemini.go index 9f5b4acb5..f5ac3991c 100644 --- a/proxy/handler_gemini.go +++ b/proxy/handler_gemini.go @@ -171,6 +171,17 @@ func geminiCountTokensInboundEndpoint() string { return "/v1beta/models:countTokens" } +// geminiNativeSelectionOptions 组装 native Gemini 入口（generateContent / +// streamGenerateContent / countTokens）的选号上下文。 +// +// 这两个入口的候选池被 antigravityOAuthChannelAccountFilter 硬限定为 Antigravity OAuth， +// 因此渠道必须按候选池的真实归属取 antigravity：auto 渠道 Key（requestUpstreamChannel 为 +// 空）走这里时，若沿用请求渠道，auth 侧 isAntigravityProvider("") 为 false，首绑 / Fresh +// 加权避险会整体失效。 +func geminiNativeSelectionOptions(c *gin.Context, model string, policy auth.DispatchPolicy) auth.SelectionOptions { + return selectionOptionsForRequestWithProvider(c, model, policy, database.UpstreamChannelAntigravity) +} + func (h *Handler) handleGeminiCountTokens(c *gin.Context, model string, rawBody []byte) { if h.enforceAPIKeyLimitsAndReply(c, model) { return @@ -184,7 +195,7 @@ func (h *Handler) handleGeminiCountTokens(c *gin.Context, model string, rawBody } accountFilter := antigravityOAuthChannelAccountFilter(model) - accountFilter = h.withModelCooldownFilter(c.Request.Context(), model, accountFilter) + accountFilter = h.withRequestModelCooldownFilter(c, c.Request.Context(), model, accountFilter) accountFilter = h.applyUpstreamChannelFilter(c, model, accountFilter) accountFilter = h.applyScopeBudgetFilter(c, accountFilter) defer h.ReleaseAPIKeyScopeConcurrency(c) @@ -203,6 +214,7 @@ func (h *Handler) handleGeminiCountTokens(c *gin.Context, model string, rawBody accountFilter = applyAffinityGroupRouting(c, sessionIdentity, accountFilter) affinityKey := sessionAffinityKey(sessionIdentity.affinityID, apiKeyID) dispatchPolicy := dispatchPolicyForModel(model) + selectionOpts := geminiNativeSelectionOptions(c, model, dispatchPolicy) inboundEndpoint := geminiCountTokensInboundEndpoint() upstreamEndpoint := "/v1internal:countTokens" @@ -213,7 +225,7 @@ func (h *Handler) handleGeminiCountTokens(c *gin.Context, model string, rawBody for attempt := 0; ; attempt++ { affinityGuard = auth.SessionAffinityGuard{} account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard( - c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, dispatchPolicy, + c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts, ) if account == nil { if writeSchedulerQueueError(c, selectionErr, continuousRetryProtocolOpenAI) { @@ -295,6 +307,10 @@ func (h *Handler) handleGeminiCountTokens(c *gin.Context, model string, rawBody h.store.Release(account) h.store.UnbindSessionAffinity(affinityKey, account.ID()) retryExclusions.MarkHTTPFailure(account.ID(), resp.StatusCode, errBody, maxRetries, attemptMaxRateLimitRetries, continuousRetryPolicy) + // countTokens 与 generateContent 打到的是同一个 Cloud Code 配额池：漏掉 + // 这里就等于 429/503 只被计数、从不写模型冷却，被上游限流的账号会立刻在 + // 下一次选号中重新入选（wireModel 联动与 family 冷却也一并丢失）。 + h.applyCooldownForModel(account, resp.StatusCode, errBody, resp, model) log.Printf("Gemini native countTokens upstream error (attempt %d, status %d): %s", attempt+1, resp.StatusCode, upstreamErrorConsoleBody(errBody)) if shouldRetryHTTPStatus(resp.StatusCode, errBody, &generalRetries, &rateLimitRetries, maxRetries, attemptMaxRateLimitRetries, continuousRetryPolicy) { retryOrdinal, retryLimit := retryStateForHTTPStatusWithBody(resp.StatusCode, errBody, generalRetries, rateLimitRetries, maxRetries, attemptMaxRateLimitRetries, continuousRetryPolicy) @@ -425,7 +441,7 @@ func (h *Handler) handleGeminiGenerateContent(c *gin.Context, model string, rawB } accountFilter := antigravityOAuthChannelAccountFilter(model) - accountFilter = h.withModelCooldownFilter(c.Request.Context(), model, accountFilter) + accountFilter = h.withRequestModelCooldownFilter(c, c.Request.Context(), model, accountFilter) accountFilter = h.applyUpstreamChannelFilter(c, model, accountFilter) accountFilter = h.applyScopeBudgetFilter(c, accountFilter) defer h.ReleaseAPIKeyScopeConcurrency(c) @@ -458,6 +474,7 @@ func (h *Handler) handleGeminiGenerateContent(c *gin.Context, model string, rawB accountFilter = applyAffinityGroupRouting(c, sessionIdentity, accountFilter) affinityKey := sessionAffinityKey(sessionIdentity.affinityID, apiKeyID) dispatchPolicy := dispatchPolicyForModel(model) + selectionOpts := geminiNativeSelectionOptions(c, model, dispatchPolicy) inboundEndpoint := geminiInboundEndpoint(stream) upstreamEndpoint := antigravityUpstreamEndpoint(stream) @@ -468,7 +485,7 @@ func (h *Handler) handleGeminiGenerateContent(c *gin.Context, model string, rawB for attempt := 0; ; attempt++ { affinityGuard = auth.SessionAffinityGuard{} account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard( - c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, dispatchPolicy, + c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts, ) if account == nil { if writeSchedulerQueueError(c, selectionErr, continuousRetryProtocolGemini) { diff --git a/proxy/responses_ws.go b/proxy/responses_ws.go index 66314bdf0..28fc70c13 100644 --- a/proxy/responses_ws.go +++ b/proxy/responses_ws.go @@ -487,7 +487,7 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. } accountFilter := accountFilterForResponsesWebSocket(effectiveModel) - accountFilter = h.withModelCooldownFilter(c.Request.Context(), effectiveModel, accountFilter) + accountFilter = h.withRequestModelCooldownFilter(c, c.Request.Context(), effectiveModel, accountFilter) accountFilter = applyAffinityGroupRouting(c, sessionIdentity, accountFilter) accountFilter = h.applyScopeBudgetFilter(c, accountFilter) // resolveCompactionAffinity 只在已知来源相互冲突时报错；缓存故障按未知 @@ -592,6 +592,12 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. }() dispatchPolicy := dispatchPolicyForModel(effectiveModel) + // Responses WS 入口的候选池（accountFilterForResponsesWebSocket）恒不含 Antigravity： + // OpenAIResponsesUsesUpstreamWebsocket 对 AG 账号恒为 false，accountFilterForModel + // 也直接拒绝 AG。Antigravity 渠道 Key 打到该入口时不得沿用请求渠道作为 Provider， + // 否则 Store 侧会对 Codex relay 候选套用 Antigravity 加权打分并绕过 spread / HRW， + // 与 compact 入口（selectionOptionsForCandidatePool(..., true)）保持一致。 + selectionOpts := selectionOptionsForCandidatePool(c, effectiveModel, dispatchPolicy, true) var affinityGuard auth.SessionAffinityGuard var selectionErr error for attempt := 0; ; attempt++ { @@ -620,9 +626,10 @@ func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket. if account != nil { stickyProxyURL = account.GetProxyURL() } else if continuationPinned { - account, stickyProxyURL, selectionErr = h.nextRetryAccountForContinuationWithDispatch(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, dispatchPolicy) + // 续链重试同样传 SelectionOptions，避免重试把渠道上下文丢掉。 + account, stickyProxyURL, selectionErr = h.nextRetryAccountForContinuationWithSelection(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts) } else { - account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, dispatchPolicy) + account, stickyProxyURL, affinityGuard, selectionErr = h.nextRetryAccountForSessionWithDispatchGuard(c.Request.Context(), affinityKey, apiKeyID, retryExclusions, accountFilter, selectionOpts) } } if account == nil { diff --git a/proxy/retry_exclusions.go b/proxy/retry_exclusions.go index 7b34444b3..5b9f15b10 100644 --- a/proxy/retry_exclusions.go +++ b/proxy/retry_exclusions.go @@ -456,11 +456,11 @@ var dispatchAccountWaitTimeout = 30 * time.Second // waitForRetryAccountAvailable keeps one queue admission for the normal // 30-second wait, including continuous-retry SSE/WebSocket heartbeats. func (h *Handler) waitForRetryAccountAvailable(ctx context.Context, affinityKey string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, preserveBinding bool, policy auth.DispatchPolicy) (*auth.Account, string, error) { - account, proxyURL, _, err := h.waitForRetryAccountAvailableWithGuard(ctx, affinityKey, apiKeyID, exclude, filter, preserveBinding, policy) + account, proxyURL, _, err := h.waitForRetryAccountAvailableWithGuard(ctx, affinityKey, apiKeyID, exclude, filter, preserveBinding, auth.SelectionOptions{Policy: policy}) return account, proxyURL, err } -func (h *Handler) waitForRetryAccountAvailableWithGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, preserveBinding bool, policy auth.DispatchPolicy) (*auth.Account, string, auth.SessionAffinityGuard, error) { +func (h *Handler) waitForRetryAccountAvailableWithGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, preserveBinding bool, opts auth.SelectionOptions) (*auth.Account, string, auth.SessionAffinityGuard, error) { if ctx == nil { ctx = context.Background() } @@ -482,7 +482,7 @@ func (h *Handler) waitForRetryAccountAvailableWithGuard(ctx context.Context, aff return step, nil } } - account, proxyURL, guard, err := h.store.WaitForDispatchAvailable(ctx, affinityKey, dispatchAccountWaitTimeout, apiKeyID, exclude, filter, preserveBinding, policy, heartbeat) + account, proxyURL, guard, err := h.store.WaitForDispatchAvailableForProvider(ctx, affinityKey, dispatchAccountWaitTimeout, apiKeyID, exclude, filter, opts, preserveBinding, heartbeat) account, proxyURL = guardRetryAccountContext(ctx, h.store.Release, account, proxyURL) if account == nil { guard = auth.SessionAffinityGuard{} @@ -498,25 +498,30 @@ func (h *Handler) nextRetryAccountForSessionWithDispatch(ctx context.Context, af return h.nextRetryAccount(ctx, affinityKey, apiKeyID, exclusions, filter, false, policy) } -func (h *Handler) nextRetryAccountForSessionWithDispatchGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, policy auth.DispatchPolicy) (*auth.Account, string, auth.SessionAffinityGuard, error) { - return h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, false, policy) +func (h *Handler) nextRetryAccountForSessionWithDispatchGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, opts auth.SelectionOptions) (*auth.Account, string, auth.SessionAffinityGuard, error) { + return h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, false, opts) } func (h *Handler) nextRetryAccountForSessionWithGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter) (*auth.Account, string, auth.SessionAffinityGuard, error) { - return h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, false, auth.DispatchPolicyStandard) + return h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) } func (h *Handler) nextRetryAccountForContinuation(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter) (*auth.Account, string) { return h.nextRetryAccount(ctx, affinityKey, apiKeyID, exclusions, filter, true, auth.DispatchPolicyStandard) } -func (h *Handler) nextRetryAccountForContinuationWithDispatch(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, policy auth.DispatchPolicy) (*auth.Account, string, error) { - account, proxyURL, _, err := h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, true, policy) +// nextRetryAccountForContinuationWithSelection 是续链（preserveBinding=true）重试的 +// SelectionOptions 变体：续链虽然只认已绑定的账号，但选号实现仍需要 Provider / Policy +// 才能走对渠道分支。传 policy-only 的 SelectionOptions 会让 auto 渠道下的 Antigravity +// 续链请求丢掉加权避险上下文。候选池已显式排除 Antigravity 的入口（compact 执行器 / +// native remote compaction）例外，见 selectionOptionsForCandidatePool。 +func (h *Handler) nextRetryAccountForContinuationWithSelection(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, opts auth.SelectionOptions) (*auth.Account, string, error) { + account, proxyURL, _, err := h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, true, opts) return account, proxyURL, err } func (h *Handler) nextRetryAccount(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, preserveBinding bool, policy auth.DispatchPolicy) (*auth.Account, string) { - account, proxyURL, _, _ := h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, preserveBinding, policy) + account, proxyURL, _, _ := h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, preserveBinding, auth.SelectionOptions{Policy: policy}) return account, proxyURL } @@ -535,13 +540,14 @@ func selectionFilterWithContext(ctx context.Context, filter auth.AccountFilter) } } -func (h *Handler) nextRetryAccountWithGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, preserveBinding bool, policy auth.DispatchPolicy) (*auth.Account, string, auth.SessionAffinityGuard, error) { +func (h *Handler) nextRetryAccountWithGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, preserveBinding bool, opts auth.SelectionOptions) (*auth.Account, string, auth.SessionAffinityGuard, error) { if ctx == nil { ctx = context.Background() } if h == nil || h.store == nil { return nil, "", auth.SessionAffinityGuard{}, nil } + policy := opts.Policy ctx, cancelSelection := context.WithTimeout(ctx, accountSelectionTimeout) defer cancelSelection() filter = selectionFilterWithContext(ctx, filter) @@ -554,9 +560,11 @@ func (h *Handler) nextRetryAccountWithGuard(ctx context.Context, affinityKey str var stickyProxyURL string var guard auth.SessionAffinityGuard if preserveBinding { - account, stickyProxyURL = h.store.NextForContinuationWithDispatch(affinityKey, apiKeyID, exclude, filter, policy) + // 续链同样带上 SelectionOptions：Provider 为空时与 policy-only 行为一致， + // 携带渠道时底层选号（含容量借用 / 首次绑定分支）不会丢掉加权避险上下文。 + account, stickyProxyURL = h.store.NextForContinuationWithDispatchForProvider(opts, affinityKey, apiKeyID, exclude, filter) } else { - account, stickyProxyURL, guard = h.nextAccountForSessionWithDispatchGuard(affinityKey, apiKeyID, exclude, filter, policy) + account, stickyProxyURL, guard = h.nextAccountForSessionWithDispatchGuard(affinityKey, apiKeyID, exclude, filter, opts) } if account != nil { if ctx.Err() != nil { @@ -578,7 +586,7 @@ func (h *Handler) nextRetryAccountWithGuard(ctx context.Context, affinityKey str // retry interval/backoff still paces the next upstream attempt. if !h.onlyResettableExclusionsBlock(affinityKey, apiKeyID, exclusions, exclude, filter, preserveBinding, policy) { var admissionErr error - account, stickyProxyURL, guard, admissionErr = h.waitForRetryAccountAvailableWithGuard(ctx, affinityKey, apiKeyID, exclude, filter, preserveBinding, policy) + account, stickyProxyURL, guard, admissionErr = h.waitForRetryAccountAvailableWithGuard(ctx, affinityKey, apiKeyID, exclude, filter, preserveBinding, opts) if account != nil { if ctx.Err() != nil { h.store.Release(account) diff --git a/proxy/scheduler_queue_error_test.go b/proxy/scheduler_queue_error_test.go index 08c9f3ba0..8e69c9a33 100644 --- a/proxy/scheduler_queue_error_test.go +++ b/proxy/scheduler_queue_error_test.go @@ -46,7 +46,7 @@ func TestSchedulerQueueOverloadStopsContinuousRetry(t *testing.T) { exclusions.MarkTransient(99) ctx, cancel := context.WithTimeout(context.Background(), time.Second) defer cancel() - _, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, exclusions, nil, false, auth.DispatchPolicyStandard) + _, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, exclusions, nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if !errors.Is(err, auth.ErrSchedulerQueueFull) || ctx.Err() != nil { t.Fatalf("overload was retried: %v (context %v)", err, ctx.Err()) } @@ -242,7 +242,7 @@ func TestSchedulerWaitHeartbeatPreservesOneAdmission(t *testing.T) { keepalive := &recordingContinuousRetryKeepalive{active: true} ctx, cancel := context.WithTimeout(contextWithContinuousRetryKeepalive(keepalive), 35*time.Millisecond) defer cancel() - _, _, _, _ = h.waitForRetryAccountAvailableWithGuard(ctx, "", 0, nil, nil, false, auth.DispatchPolicyStandard) + _, _, _, _ = h.waitForRetryAccountAvailableWithGuard(ctx, "", 0, nil, nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) m := store.GetSchedulerMetrics() if keepalive.writes < 2 || m.WaitStarted != 1 || m.Waiters != 0 || m.SelectionTotal != 1 { t.Fatalf("heartbeats disturbed queue: writes=%d metrics=%+v", keepalive.writes, m) diff --git a/proxy/selection_cancellation_test.go b/proxy/selection_cancellation_test.go index 0ab692189..149b64a24 100644 --- a/proxy/selection_cancellation_test.go +++ b/proxy/selection_cancellation_test.go @@ -36,7 +36,7 @@ func TestSelectionCancellationStopsInitialScan(t *testing.T) { return false } h := &Handler{store: store} - account, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, newRetryAccountExclusions(), filter, false, auth.DispatchPolicyStandard) + account, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, newRetryAccountExclusions(), filter, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if account != nil { store.Release(account) t.Fatal("selected an account after cancellation") @@ -75,7 +75,7 @@ func TestSelectionBudgetIncludesInitialScan(t *testing.T) { return false } h := &Handler{store: store} - account, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, newRetryAccountExclusions(), filter, false, auth.DispatchPolicyStandard) + account, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, newRetryAccountExclusions(), filter, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if account != nil || !errors.Is(err, context.DeadlineExceeded) || calls != 2 { t.Fatalf("selection returned account=%v, error=%v, filter calls=%d", account, err, calls) } @@ -94,7 +94,7 @@ func TestSelectionSuccessPreservesRequestLifetime(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() h := &Handler{store: store} - got, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, newRetryAccountExclusions(), nil, false, auth.DispatchPolicyStandard) + got, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, newRetryAccountExclusions(), nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if err != nil || got != account { t.Fatalf("selection = %v, %v", got, err) } @@ -113,7 +113,7 @@ func TestSelectionDeadlineStopsPoolWait(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) defer cancel() start := time.Now() - account, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, newRetryAccountExclusions(), nil, false, auth.DispatchPolicyStandard) + account, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, newRetryAccountExclusions(), nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if account != nil || !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("selection returned account=%v, error=%v", account, err) } @@ -145,7 +145,7 @@ func TestSelectionResetsRequestExclusionsCoveringWholePool(t *testing.T) { h := &Handler{store: store} start := time.Now() - got, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, exclusions, nil, false, auth.DispatchPolicyStandard) + got, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, exclusions, nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if err != nil || got != account { t.Fatalf("selection = %v, %v; want the %s-excluded account after reset", got, err, kind) } @@ -177,7 +177,7 @@ func TestSelectionTransientResetKeepsHardExclusions(t *testing.T) { h := &Handler{store: store} for i := 0; i < 3; i++ { - got, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, exclusions, nil, false, auth.DispatchPolicyStandard) + got, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, exclusions, nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if err != nil || got != transient { t.Fatalf("round %d selection = %v, %v; want only the transient account", i, got, err) } @@ -200,7 +200,7 @@ func TestSelectionTransientResetKeepsAccountCooldown(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) defer cancel() h := &Handler{store: store} - got, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, exclusions, nil, false, auth.DispatchPolicyStandard) + got, _, _, err := h.nextRetryAccountWithGuard(ctx, "", 0, exclusions, nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if got != nil { store.Release(got) t.Fatal("transient reset bypassed the account cooldown") @@ -223,7 +223,7 @@ func TestSelectionStillWaitsForSaturatedCandidate(t *testing.T) { exclusions.MarkTransient(failed.DBID) h := &Handler{store: store} - held, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, exclusions, nil, false, auth.DispatchPolicyStandard) + held, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, exclusions, nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) if err != nil || held != busy { t.Fatalf("setup selection = %v, %v", held, err) } @@ -233,7 +233,7 @@ func TestSelectionStillWaitsForSaturatedCandidate(t *testing.T) { store.Release(held) close(released) }() - got, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, exclusions, nil, false, auth.DispatchPolicyStandard) + got, _, _, err := h.nextRetryAccountWithGuard(context.Background(), "", 0, exclusions, nil, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) <-released if err != nil || got != busy { t.Fatalf("selection = %v, %v; want the released saturated account", got, err) From ea65d2aaa109ba8b2d575ff3603b2573b714834c Mon Sep 17 00:00:00 2001 From: aghost Date: Mon, 19 Oct 2026 01:27:21 +0800 Subject: [PATCH 2/2] =?UTF-8?q?docs:=20=E8=A1=A5=E5=85=A8=20diff=20?= =?UTF-8?q?=E8=A7=A6=E5=8F=8A=E5=87=BD=E6=95=B0=20docstring=20=E6=B3=A8?= =?UTF-8?q?=E9=87=8A=E4=BB=A5=E6=BB=A1=E8=B6=B3=20CI=20=E9=97=A8=E6=A7=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- auth/antigravity_balance_test.go | 4 ++++ auth/antigravity_ratelimit.go | 4 ++++ auth/antigravity_ratelimit_test.go | 14 ++++++++++++++ auth/store.go | 22 ++++++++++++++++++++++ proxy/antigravity_quota.go | 6 ++++++ proxy/antigravity_quota_test.go | 5 +++++ proxy/handler.go | 6 ++++++ proxy/handler_gemini.go | 9 +++++++++ proxy/responses_ws.go | 4 ++++ proxy/retry_exclusions.go | 15 +++++++++++++++ proxy/scheduler_queue_error_test.go | 5 +++++ proxy/selection_cancellation_test.go | 13 +++++++++++++ 12 files changed, 107 insertions(+) diff --git a/auth/antigravity_balance_test.go b/auth/antigravity_balance_test.go index 61676b230..9e7601657 100644 --- a/auth/antigravity_balance_test.go +++ b/auth/antigravity_balance_test.go @@ -31,6 +31,8 @@ func newAntigravityBalanceStore(t *testing.T, accounts int) *Store { return store } +// antigravitySelectionOpts 返回渠道为 Antigravity 的选号上下文，用于让测试进入 +// 首绑 / Fresh 加权避险分支。 func antigravitySelectionOpts() SelectionOptions { return SelectionOptions{ Provider: database.UpstreamChannelAntigravity, @@ -398,6 +400,8 @@ func antigravityTieBreakWinner(tie func(string, int64) uint64, accounts, trials return wins } +// antigravityChiSquared 以等概率为期望值计算各账号中奖次数的卡方统计量， +// 用于断言散列平票打散在账号维度上没有系统性偏置。 func antigravityChiSquared(wins []int) float64 { total, expected := 0, 0.0 for _, win := range wins { diff --git a/auth/antigravity_ratelimit.go b/auth/antigravity_ratelimit.go index 72b0aca56..519897074 100644 --- a/auth/antigravity_ratelimit.go +++ b/auth/antigravity_ratelimit.go @@ -167,10 +167,14 @@ func (a *Account) ResetAntigravityPenalty() { // 原 round-robin。冷却正确性修复不受该开关影响，始终生效。 var antigravityBalancedAffinityEnabled atomic.Bool +// init 在包加载时用环境变量初值打开或关闭首绑 / Fresh 选号加权避险开关。 func init() { antigravityBalancedAffinityEnabled.Store(antigravityBalancedAffinityFromEnv()) } +// antigravityBalancedAffinityFromEnv 解析 ANTIGRAVITY_BALANCED_AFFINITY：未设置、 +// 空值或无法识别的取值都按开启处理，只有 0/false/off/none/no（忽略大小写与空白） +// 才关闭加权避险。 func antigravityBalancedAffinityFromEnv() bool { raw, ok := os.LookupEnv(antigravityBalancedAffinityEnv) if !ok { diff --git a/auth/antigravity_ratelimit_test.go b/auth/antigravity_ratelimit_test.go index 6c88f1619..1b3486d37 100644 --- a/auth/antigravity_ratelimit_test.go +++ b/auth/antigravity_ratelimit_test.go @@ -14,6 +14,8 @@ import ( // 合成键；auth 包不能反向依赖 proxy，这里按字面量断言。 const antigravityGeminiFamilyKey = "antigravity:gemini" +// antigravityPenaltyClose 断言惩罚值落在容差内（期望 0 时收紧到 1e-9），避免指数 +// 衰减的二进制舍入把断言变成偶发失败。 func antigravityPenaltyClose(t *testing.T, got, want float64, label string) { t.Helper() tolerance := 1e-6 @@ -25,6 +27,8 @@ func antigravityPenaltyClose(t *testing.T, got, want float64, label string) { } } +// TestAntigravityPenaltyDecaysWithHalfLife 覆盖惩罚的指数衰减：每过一个半衰期减半、 +// 超出整个窗口直接归零、时钟回拨不放大惩罚，ResetAntigravityPenalty 立即清零。 func TestAntigravityPenaltyDecaysWithHalfLife(t *testing.T) { base := time.Date(2026, 10, 17, 12, 0, 0, 0, time.UTC) acc := &Account{DBID: 8801, AccessToken: "token"} @@ -64,6 +68,8 @@ func TestAntigravityPenaltyDecaysWithHalfLife(t *testing.T) { } } +// TestAntigravityPenaltyAccumulatesAndCaps 覆盖惩罚累加：新事件先按间隔衰减旧累加值 +// 再相加，跨过整个窗口后旧值完全归零，连续事件累加封顶在 antigravityThrottleMaxPenalty。 func TestAntigravityPenaltyAccumulatesAndCaps(t *testing.T) { base := time.Date(2026, 10, 17, 12, 0, 0, 0, time.UTC) acc := &Account{DBID: 8802, AccessToken: "token"} @@ -91,6 +97,9 @@ func TestAntigravityPenaltyAccumulatesAndCaps(t *testing.T) { } } +// TestRecordAntigravityRateLimitEventWeights 覆盖事件分类权重：共享容量与账号配额 +// 不记账，裸 RESOURCE_EXHAUSTED 记半个、明确限流记一个，权重越界夹回区间、NaN 与 +// 负值丢弃，缺省 ObservedAt 落到当前时刻，nil 账号不 panic。 func TestRecordAntigravityRateLimitEventWeights(t *testing.T) { base := time.Date(2026, 10, 17, 12, 0, 0, 0, time.UTC) @@ -135,6 +144,9 @@ func TestRecordAntigravityRateLimitEventWeights(t *testing.T) { } } +// TestAntigravityBalancedAffinitySwitch 覆盖加权避险开关：Set / Get 成对生效， +// ReloadAntigravityBalancedAffinityFromEnv 按 ANTIGRAVITY_BALANCED_AFFINITY 的取值 +// 重新判定开启与关闭，用例结束后恢复原值。 func TestAntigravityBalancedAffinitySwitch(t *testing.T) { previous := GetAntigravityBalancedAffinity() t.Cleanup(func() { SetAntigravityBalancedAffinity(previous) }) @@ -169,6 +181,8 @@ func TestAntigravityBalancedAffinitySwitch(t *testing.T) { } } +// newAntigravityRateLimitTestStore 构造仅用于冷却与惩罚断言的内存 Store：账号直接 +// AddAccount 注入，不经数据库，并注册 Stop 清理。 func newAntigravityRateLimitTestStore(t *testing.T, tokenCache cache.TokenCache, accounts ...*Account) *Store { t.Helper() store := NewStore(nil, tokenCache, &database.SystemSettings{MaxConcurrency: 4, TestConcurrency: 1, TestModel: "gpt-5.4"}) diff --git a/auth/store.go b/auth/store.go index 34d0c11b0..700ff9b71 100644 --- a/auth/store.go +++ b/auth/store.go @@ -3099,6 +3099,8 @@ func (a *Account) AllowsAPIKey(apiKeyID int64) bool { return ok } +// normalizeModelCooldownKey 归一化模型冷却键：去除首尾空白并转为小写。 +// 返回空串表示没有可用的模型标识，调用方据此跳过冷却写入与查询。 func normalizeModelCooldownKey(model string) string { return strings.ToLower(strings.TrimSpace(model)) } @@ -3141,6 +3143,10 @@ func (a *Account) setModelCooldownUntilLocked(key, reason string, resetAt, now t return cooldown } +// SetModelCooldownUntil 在账号内存中为指定模型写入「冷却直至 resetAt」并返回写入 +// 的条目（本方法仅更新内存 map，不负责持久化）；模型键为空或 resetAt 为零值时返回 +// 零值且不写入。wire、family 这类被多个公开档位共享的冷却键应改用 +// SetModelCooldownUntilAtLeast，避免更短的 deadline 把既有长冷却改短。 func (a *Account) SetModelCooldownUntil(model, reason string, resetAt time.Time) ModelCooldown { key, reason, ok := normalizeModelCooldownArgs(model, reason, resetAt) if !ok { @@ -3206,6 +3212,8 @@ func (a *Account) IsModelRateLimited(model string) bool { return ok && cooldown.ResetAt.After(time.Now()) } +// ModelCooldownRemaining 返回指定模型冷却距 ResetAt 的剩余时长；模型键为空、 +// 没有冷却条目或冷却已过期均返回 0。 func (a *Account) ModelCooldownRemaining(model string) time.Duration { key := normalizeModelCooldownKey(model) if key == "" { @@ -7448,6 +7456,10 @@ func (s *Store) nextAccountForFreshAffinity(key string, apiKeyID int64, exclude return s.nextAccountForFreshAffinityWithDispatch(key, apiKeyID, exclude, filter, SelectionOptions{Policy: DispatchPolicyStandard}) } +// nextAccountForFreshAffinityWithDispatch 是 nextAccountForFreshAffinity 的带 +// DispatchPolicy 变体：Antigravity 加权避险接管时直接返回其结果（含号池无可用 +// 容量返回 nil 的情况），避免再次回退重扫并重复求值 filter；未接管时才按散列开关 +// 回退到 HRW / 调度器 / round-robin 重扫。 func (s *Store) nextAccountForFreshAffinityWithDispatch(key string, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, opts SelectionOptions) *Account { if s == nil { return nil @@ -8033,6 +8045,9 @@ func (s *Store) WaitForSessionAvailableWithFilter(ctx context.Context, key strin return account, proxyURL } +// WaitForSessionAvailableWithDispatch 按指定 DispatchPolicy 排队等待满足过滤条件 +// 的会话账号，只返回账号与代理地址；需要保留容量溢出借用判定时改用 +// WaitForSessionAvailableWithDispatchGuard。 func (s *Store) WaitForSessionAvailableWithDispatch(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) (*Account, string) { account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, false, SelectionOptions{Policy: policy}) return account, proxyURL @@ -8052,6 +8067,9 @@ func (s *Store) WaitForContinuationAvailableWithFilter(ctx context.Context, key return account, proxyURL } +// WaitForContinuationAvailableWithDispatch 按指定 DispatchPolicy 排队等待续链亲和键 +// 绑定的账号可用（preserveBinding=true）：存在绑定时锁定该账号不顺延其他账号，无绑定时 +// 沿用底层 Fresh 选号语义。 func (s *Store) WaitForContinuationAvailableWithDispatch(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, policy DispatchPolicy) (*Account, string) { account, proxyURL, _, _ := s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, true, SelectionOptions{Policy: policy}) return account, proxyURL @@ -8063,6 +8081,10 @@ func (s *Store) WaitForDispatchAvailable(ctx context.Context, key string, timeou return s.waitForSessionAvailableWithFilter(ctx, key, timeout, apiKeyID, exclude, filter, preserveBinding, SelectionOptions{Policy: policy}, heartbeat...) } +// waitForSessionAvailableWithFilter 是各会话等待入口的共同底层实现：按 opts 在超时 +// 窗口内反复尝试选号并挂到可用性 hub 上等待唤醒，preserveBinding 为 true 时若存在绑定 +// 则锁定该账号不顺延其他账号（无绑定时走 Fresh 选号），heartbeat 非空时在等待期间按回调 +// 返回的间隔写出心跳。队列准入被拒（如队列已满）时把错误原样上抛，供协议处理器区分过载与无候选。 func (s *Store) waitForSessionAvailableWithFilter(ctx context.Context, key string, timeout time.Duration, apiKeyID int64, exclude map[int64]bool, filter AccountFilter, preserveBinding bool, opts SelectionOptions, heartbeat ...SchedulerWaitHeartbeat) (*Account, string, SessionAffinityGuard, error) { if ctx == nil { ctx = context.Background() diff --git a/proxy/antigravity_quota.go b/proxy/antigravity_quota.go index 7ff526a9c..985be1ddb 100644 --- a/proxy/antigravity_quota.go +++ b/proxy/antigravity_quota.go @@ -409,6 +409,12 @@ func antigravityRateLimitEvent(statusCode int, model, wireModel string, quota an return event } +// applyAntigravityCooldownWithPenalty 是 Antigravity 429/503 冷却落地的共同实现：解析 +// 结构化 google.rpc 信号定出冷却时长（优先尊重 RetryInfo 与 Retry-After 头，无明确提示 +// 时按 Kind 分类默认值处理，裸限流退到短回退；完全解析失败的非结构化 429 走通用 429 策略）， +// 写入公开模型键。对于共享的 wire 键统一只延长、不改短；Gemini family 冷却键仅在 +// QUOTA_EXHAUSTED 且归属 Gemini 模型时写入且同样只延长不改短。recordPenalty 为 false +// 时只落冷却、不进进程内避险惩罚累加器（admin 连接性探测走这一档）。 func applyAntigravityCooldownWithPenalty(store *auth.Store, account *auth.Account, statusCode int, body []byte, resp *http.Response, model string, recordPenalty bool) codex429Decision { quota, ok := parseAntigravityQuotaError(statusCode, body) if !ok { diff --git a/proxy/antigravity_quota_test.go b/proxy/antigravity_quota_test.go index e98a4e949..9fe9208c3 100644 --- a/proxy/antigravity_quota_test.go +++ b/proxy/antigravity_quota_test.go @@ -237,6 +237,9 @@ func TestApplyAntigravityCooldownBacksOffBrieflyWhenGoogleGivesNoHint(t *testing } } +// TestApplyAntigravityCooldownPrefersRetryAfterHeaderOverTheShortBackOff 覆盖 +// Retry-After 响应头优先：结构化 body 没有给出 RetryInfo 时，头部秒数必须取代裸 +// RESOURCE_EXHAUSTED 的短回退时长。 func TestApplyAntigravityCooldownPrefersRetryAfterHeaderOverTheShortBackOff(t *testing.T) { store := newProxyPremiumTestStore() account := &auth.Account{DBID: 7303, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project"} @@ -451,6 +454,8 @@ func TestGeminiNativeCountTokensAppliesModelCooldown(t *testing.T) { // antigravityQuotaExhaustedThreshold，因此仍按 rate_limit 分类。 const antigravityLongRateLimitBody = `{"error":{"code":429,"message":"Resource has been exhausted (e.g. check quota).","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"RATE_LIMIT_EXCEEDED","metadata":{"model":"gemini-3.8-flash-tiered"}},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"290s"}]}}` +// newAntigravityCooldownTestAccount 构造带项目凭据的 Antigravity 账号，供冷却用例 +// 复用同一组必需字段。 func newAntigravityCooldownTestAccount(id int64) *auth.Account { return &auth.Account{DBID: id, UpstreamType: auth.UpstreamAntigravity, AccessToken: "google-token", AntigravityProjectID: "google-project"} } diff --git a/proxy/handler.go b/proxy/handler.go index c742f6cf1..4a40614b7 100644 --- a/proxy/handler.go +++ b/proxy/handler.go @@ -109,11 +109,15 @@ func (h *Handler) nextAccountForSessionWithFilter(sessionID string, apiKeyID int return h.nextAccountForSessionWithDispatch(sessionID, apiKeyID, exclude, filter, auth.DispatchPolicyStandard) } +// nextAccountForSessionWithDispatch 按指定 DispatchPolicy 做会话选号，只返回账号与 +// 代理地址；需要容量溢出守卫时改用 nextAccountForSessionWithDispatchGuard。 func (h *Handler) nextAccountForSessionWithDispatch(sessionID string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, policy auth.DispatchPolicy) (*auth.Account, string) { account, proxyURL, _ := h.nextAccountForSessionWithDispatchGuard(sessionID, apiKeyID, exclude, filter, auth.SelectionOptions{Policy: policy}) return account, proxyURL } +// nextAccountForSessionWithDispatchGuard 用完整 SelectionOptions（渠道 + Policy）做 +// 会话选号，并返回容量溢出借用时携带的 SessionAffinityGuard。 func (h *Handler) nextAccountForSessionWithDispatchGuard(sessionID string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, opts auth.SelectionOptions) (*auth.Account, string, auth.SessionAffinityGuard) { if h == nil || h.store == nil { return nil, "", auth.SessionAffinityGuard{} @@ -167,6 +171,8 @@ func dispatchPolicyForModel(model string) auth.DispatchPolicy { return auth.DispatchPolicyStandard } +// withModelCooldownFilter 在既有过滤器上叠加单一公开模型的跨实例共享缓存冷却检查， +// 其它实例写入的冷却也能在本实例选号时滤掉该账号；store 未装配时原样返回入参过滤器。 func (h *Handler) withModelCooldownFilter(ctx context.Context, model string, filter auth.AccountFilter) auth.AccountFilter { if h == nil || h.store == nil { return filter diff --git a/proxy/handler_gemini.go b/proxy/handler_gemini.go index f5ac3991c..1ffcde37b 100644 --- a/proxy/handler_gemini.go +++ b/proxy/handler_gemini.go @@ -167,6 +167,8 @@ func (h *Handler) GeminiModelsAction(c *gin.Context) { } } +// geminiCountTokensInboundEndpoint 返回 native countTokens 入口的入站端点标识， +// 与上游端点（v1internal:countTokens）区分，供日志、错误与用量记录归因使用。 func geminiCountTokensInboundEndpoint() string { return "/v1beta/models:countTokens" } @@ -182,6 +184,9 @@ func geminiNativeSelectionOptions(c *gin.Context, model string, policy auth.Disp return selectionOptionsForRequestWithProvider(c, model, policy, database.UpstreamChannelAntigravity) } +// handleGeminiCountTokens 处理 native Gemini models:countTokens 请求：先过 Key 限额 +// 与并发额度，候选池被硬限定为 Antigravity OAuth 并叠加跨实例冷却、渠道与预算过滤， +// 随后在重试循环内带重试排除项与亲和守卫选号，转发到上游 v1internal:countTokens。 func (h *Handler) handleGeminiCountTokens(c *gin.Context, model string, rawBody []byte) { if h.enforceAPIKeyLimitsAndReply(c, model) { return @@ -428,6 +433,10 @@ func writeGeminiNativeError(c *gin.Context, status int, message string) { } } +// handleGeminiGenerateContent 处理 native Gemini generateContent / +// streamGenerateContent 请求：候选池被硬限定为 Antigravity OAuth 并叠加跨实例冷却、 +// 渠道与预算过滤，按 stream 决定是否安装 SSE 心跳，并在持续重试策略下带重试排除项 +// 与亲和守卫选号转发上游。 func (h *Handler) handleGeminiGenerateContent(c *gin.Context, model string, rawBody []byte, stream bool) { if h.enforceAPIKeyLimitsAndReply(c, model) { return diff --git a/proxy/responses_ws.go b/proxy/responses_ws.go index 28fc70c13..f478bd3ec 100644 --- a/proxy/responses_ws.go +++ b/proxy/responses_ws.go @@ -299,6 +299,10 @@ func stripNewAPIPolicyWebSocketEventID(payload []byte) ([]byte, string) { return cleaned, eventID } +// forwardResponsesWebSocketTurn 处理 WebSocket 传输上的一个 Responses 会话轮次：每个 +// response.create 都是独立的逻辑请求，先做内存准入、Key 绑定与配额复核并重置上一帧的 +// 安全上下文，再归一化请求体、在重试循环内带排除项重新选号转发；可恢复的失败按协议写回错误 +// 帧并保留连接，需要关闭连接的致命失败才返回 close error。 func (h *Handler) forwardResponsesWebSocketTurn(c *gin.Context, conn *websocket.Conn, rawPayload []byte, policyEventID string, options *responsesWSForwardOptions) (returnErr error) { defer releasePromptRequestFrameBody(c) reservation, admitted := security.TryAcquireRequestMemory(int64(len(rawPayload))) diff --git a/proxy/retry_exclusions.go b/proxy/retry_exclusions.go index 5b9f15b10..858813ff3 100644 --- a/proxy/retry_exclusions.go +++ b/proxy/retry_exclusions.go @@ -460,6 +460,8 @@ func (h *Handler) waitForRetryAccountAvailable(ctx context.Context, affinityKey return account, proxyURL, err } +// waitForRetryAccountAvailableWithGuard 按完整 SelectionOptions 排队等待一次队列准入， +// 持续重试心跳激活时把心跳挂进等待回调以保持连接活跃；未选中账号时一并丢弃亲和守卫。 func (h *Handler) waitForRetryAccountAvailableWithGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclude map[int64]bool, filter auth.AccountFilter, preserveBinding bool, opts auth.SelectionOptions) (*auth.Account, string, auth.SessionAffinityGuard, error) { if ctx == nil { ctx = context.Background() @@ -498,14 +500,21 @@ func (h *Handler) nextRetryAccountForSessionWithDispatch(ctx context.Context, af return h.nextRetryAccount(ctx, affinityKey, apiKeyID, exclusions, filter, false, policy) } +// nextRetryAccountForSessionWithDispatchGuard 是会话重试选号中同时携带 +// SelectionOptions 与亲和守卫的变体：不保留绑定（preserveBinding=false）， +// 排除项取自本次请求的重试排除集。 func (h *Handler) nextRetryAccountForSessionWithDispatchGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, opts auth.SelectionOptions) (*auth.Account, string, auth.SessionAffinityGuard, error) { return h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, false, opts) } +// nextRetryAccountForSessionWithGuard 以 standard 策略做会话重试选号，并返回亲和守卫。 func (h *Handler) nextRetryAccountForSessionWithGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter) (*auth.Account, string, auth.SessionAffinityGuard, error) { return h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, false, auth.SelectionOptions{Policy: auth.DispatchPolicyStandard}) } +// nextRetryAccountForContinuation 以 standard 策略做续链重试选号（preserveBinding=true）： +// 存在绑定时锁定已绑定账号，无绑定时沿用底层 Fresh 选号；需要渠道上下文时改用 +// nextRetryAccountForContinuationWithSelection。 func (h *Handler) nextRetryAccountForContinuation(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter) (*auth.Account, string) { return h.nextRetryAccount(ctx, affinityKey, apiKeyID, exclusions, filter, true, auth.DispatchPolicyStandard) } @@ -520,6 +529,8 @@ func (h *Handler) nextRetryAccountForContinuationWithSelection(ctx context.Conte return account, proxyURL, err } +// nextRetryAccount 在重试排除项与选号预算内完成一次重试选号，只返回账号与代理地址； +// 需要亲和守卫与队列准入错误时改用 nextRetryAccountWithGuard。 func (h *Handler) nextRetryAccount(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, preserveBinding bool, policy auth.DispatchPolicy) (*auth.Account, string) { account, proxyURL, _, _ := h.nextRetryAccountWithGuard(ctx, affinityKey, apiKeyID, exclusions, filter, preserveBinding, auth.SelectionOptions{Policy: policy}) return account, proxyURL @@ -540,6 +551,10 @@ func selectionFilterWithContext(ctx context.Context, filter auth.AccountFilter) } } +// nextRetryAccountWithGuard 是重试选号的底层实现：先在 accountSelectionTimeout 预算内 +// 做一次同步选号（过滤器同时受该预算约束），失败后触发调度状态对账并排队等待。阻塞者 +// 只剩本次请求自身的软 / 暂态排除项时跳过等待直接重置，硬排除、账号冷却与过滤器始终 +// 生效；预算耗尽返回上下文错误，队列准入被拒时上抛准入错误。 func (h *Handler) nextRetryAccountWithGuard(ctx context.Context, affinityKey string, apiKeyID int64, exclusions *retryAccountExclusions, filter auth.AccountFilter, preserveBinding bool, opts auth.SelectionOptions) (*auth.Account, string, auth.SessionAffinityGuard, error) { if ctx == nil { ctx = context.Background() diff --git a/proxy/scheduler_queue_error_test.go b/proxy/scheduler_queue_error_test.go index 8e69c9a33..511624131 100644 --- a/proxy/scheduler_queue_error_test.go +++ b/proxy/scheduler_queue_error_test.go @@ -39,6 +39,9 @@ func saturatedSchedulerQueue(t *testing.T) *auth.Store { return store } +// TestSchedulerQueueOverloadStopsContinuousRetry 断言调度队列准入被拒时选号立即带着 +// ErrSchedulerQueueFull 返回，不再进入持续重试循环，且既有上游重试排除项保持不变、 +// 拒绝计数各记一次。 func TestSchedulerQueueOverloadStopsContinuousRetry(t *testing.T) { s := saturatedSchedulerQueue(t) h := &Handler{store: s} @@ -232,6 +235,8 @@ func TestSchedulerQueueOverloadResponsesWebSocket(t *testing.T) { } } +// TestSchedulerWaitHeartbeatPreservesOneAdmission 断言等待期间心跳按间隔持续写出保活帧， +// 既保持连接活跃，也不会重复占用队列准入名额或遗留等待者。 func TestSchedulerWaitHeartbeatPreservesOneAdmission(t *testing.T) { previous := continuousRetryKeepaliveInterval continuousRetryKeepaliveInterval = 5 * time.Millisecond diff --git a/proxy/selection_cancellation_test.go b/proxy/selection_cancellation_test.go index 149b64a24..d5cdfa5ca 100644 --- a/proxy/selection_cancellation_test.go +++ b/proxy/selection_cancellation_test.go @@ -11,6 +11,8 @@ import ( "github.com/codex2api/database" ) +// TestSelectionCancellationStopsInitialScan 断言上下文取消后初次扫描立即停止：过滤器 +// 至多再被调用一次，不选入账号、不启动队列等待，并返回 context.Canceled。 func TestSelectionCancellationStopsInitialScan(t *testing.T) { for _, engine := range []string{"legacy", "indexed"} { for _, cancelBefore := range []bool{true, false} { @@ -59,6 +61,8 @@ func TestSelectionCancellationStopsInitialScan(t *testing.T) { } } +// TestSelectionBudgetIncludesInitialScan 断言选号预算把初次扫描自身算进去：慢过滤器 +// 返回后预算已到期，必须直接返回 DeadlineExceeded，不再启动任何队列等待。 func TestSelectionBudgetIncludesInitialScan(t *testing.T) { synctest.Test(t, func(t *testing.T) { store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1, SchedulerEngine: "legacy"}) @@ -85,6 +89,9 @@ func TestSelectionBudgetIncludesInitialScan(t *testing.T) { }) } +// TestSelectionSuccessPreservesRequestLifetime 断言选号超时只约束选号阶段：成功选号后 +// 即便超过 accountSelectionTimeout，父上下文不受影响、账号仍被持有，不误伤后续上游 +// 流式响应。 func TestSelectionSuccessPreservesRequestLifetime(t *testing.T) { synctest.Test(t, func(t *testing.T) { store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1}) @@ -106,6 +113,8 @@ func TestSelectionSuccessPreservesRequestLifetime(t *testing.T) { }) } +// TestSelectionDeadlineStopsPoolWait 断言请求截止时间到点即退出号池等待，不遗留队列 +// 等待者，也不会一直等到 accountSelectionTimeout。 func TestSelectionDeadlineStopsPoolWait(t *testing.T) { store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1, FastSchedulerEnabled: true}) t.Cleanup(store.Stop) @@ -162,6 +171,8 @@ func TestSelectionResetsRequestExclusionsCoveringWholePool(t *testing.T) { } } +// TestSelectionTransientResetKeepsHardExclusions 断言多轮暂态重置只会重新放开暂态排除 +// 的账号，MarkHard 标记的硬排除账号在每一轮都必须继续被拦住。 func TestSelectionTransientResetKeepsHardExclusions(t *testing.T) { for _, engine := range []string{"legacy", "indexed"} { t.Run(engine, func(t *testing.T) { @@ -188,6 +199,8 @@ func TestSelectionTransientResetKeepsHardExclusions(t *testing.T) { } } +// TestSelectionTransientResetKeepsAccountCooldown 断言暂态重置不会绕过账号级冷却：唯一 +// 候选仍处于冷却时只能等到上下文超时，不能借重置把它重新选入。 func TestSelectionTransientResetKeepsAccountCooldown(t *testing.T) { store := auth.NewStore(nil, nil, &database.SystemSettings{MaxConcurrency: 1, SchedulerEngine: "indexed"}) t.Cleanup(store.Stop)