fix: align least-used scheduling with TS

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-29 23:14:30 +08:00
co-authored by Codex
parent bccbe2a8c5
commit 8a415552b0
11 changed files with 666 additions and 82 deletions
+293
View File
@@ -0,0 +1,293 @@
# 调度与连接复用审计
本文档记录 Rust 版 `codex-proxy-rs` 与 TS 版 `/home/zyy/Codes/codex-proxy` 的调度审计过程。每个结论必须能回到源码、日志或数据库状态,不把猜测写成事实。
## 当前问题
- 用户在前端选择的是 `least_used` / 智能分配。
- 真实日志中,同一批连续自然文本请求在两个 active 账号之间来回切换:
- `acct_81aaba2a4e084162924b17d6f55e8a10`
- `acct_fa9f0172d2084eaf86ddbd0e5d0c1c16`
- 两个账号各自的 WebSocket 连接后续都能 `reuse`,但账号切换会让同一轮对话无法稳定复用同一个账号上的上游连接和缓存。
## TS 对照点
TS 调度入口:
- `/home/zyy/Codes/codex-proxy/src/auth/account-lifecycle.ts`
- `/home/zyy/Codes/codex-proxy/src/auth/rotation-strategy.ts`
- `/home/zyy/Codes/codex-proxy/src/proxy/ws-pool.ts`
TS `least_used` 比较顺序:
1. `cachedQuota.rate_limit.limit_reached` 的账号靠后。
2. 两个账号都有 `usage.window_reset_at` 时,更早 reset 的账号优先。
3. `usage.request_count` 更少的账号优先。
4. `usage.last_used` 更早的账号优先。
5. 只有完整比较结果相同的首组选项,才用 `roundRobinIndex` 轮转。
TS 前端文案是:
> 优先使用即将刷新额度的账号,最大化总使用量。
该文案依赖账号存在可比较的 `usage.window_reset_at`。如果候选账号都没有窗口重置时间,TS 会退化到 `request_count` + LRU,连续请求出现 A/B 交替是符合 TS 比较器的。
TS WebSocket pool 语义:
- pool key 是账号身份加 conversation identity。
- 注释写明 `${entryId}:${conversationId}`,实际目标是同账号同 conversation 复用同一条物理 WS。
- 因此智能分配切换账号时,`prompt_cache_key` 不会丢,但上游连接和缓存命中会跟着账号切换被打断。
## Rust 当前观察
Rust 调度入口:
- `src/upstream/accounts/pool.rs`
- `src/upstream/transport/client.rs`
- `src/proxy/dispatch/responses.rs`
Rust 当前 `least_used` 比较器与 TS 主要顺序一致:
- `quota_limit_reached`
- `window_reset_at`
- `request_count`
- `last_used_at`
Rust WebSocket pool key:
- `src/upstream/transport/client.rs`
- `websocket_pool_key()` 使用 `base_url + account_id + conversation_id`。
- `conversation_id` 来源优先级:`prompt_cache_key` -> `client_conversation_id` -> `previous_response_id`。
Rust 与 TS 一个高风险差异:
- TS 在 `AccountLifecycle.release()` 时调用 `registry.recordUsage()`,请求完成后才更新 `request_count` / `last_used`。
- Rust 当前在 `AccountPool::acquire_with_status_refresh()` 内调用 `mark_usage()`,获取账号时立刻更新 `request_count` / `last_used_at` / `window_request_count`,并且 `RuntimeAccountPoolService::acquire_with()` 立即持久化 `record_request()`。
- 这个差异会影响并发、失败请求、长流式请求中的排序时机。是否是用户看到 A/B 交替的主因还要继续用测试和 TS 行为确认。
## 2026-06-29 日志证据
日志:`.runtime/logs/codex-proxy-rs.2026-06-29.log`
22:11 后连续 `/v1/responses` 流式请求的账号序列:
- `req_5c9267ad...` -> `acct_81a...`,WS `new`
- `req_98cf5322...` -> `acct_fa9...`,WS `new`
- `req_c9babaa5...` -> `acct_81a...`,WS `new`
- `req_133b4d89...` -> `acct_fa9...`,WS `new`
- `req_5171359a...` -> `acct_81a...`,WS `reuse`
- `req_e62aae2e...` -> `acct_fa9...`,WS `reuse`
- `req_0936b7aa...` -> `acct_81a...`,WS `reuse`
- `req_a778e05a...` -> `acct_fa9...`,WS `reuse`
- `req_8c9647cc...` -> `acct_81a...`,WS `reuse`
- `req_960b167e...` -> `acct_81a...`,WS `reuse`
运行时 DB 当前状态:
- `runtime_settings.rotation_strategy = least_used`
- 两个 active 账号的 `window_reset_at` 都是 `null`。
- 两个 active 账号的 `request_count` 都是 `14`。
- 两个 active 账号的 `quota_limit_reached = 0`。
由此可解释:在没有 `window_reset_at` 可比较时,`least_used` 会退化到 `request_count` 和 LRU。连续请求容易在两个账号之间交替;这对缓存稳定不友好,但需要确认是否就是 TS 的预期语义。
## 待确认
- `websocket pool decision` 的 pool key hash 已补;最终流日志是否还需要重复写入 pool key hash,待真实链路日志再判断。
## 已确认偏差
### active 账号的 quota reset 没有同步到 window_reset_at
当前 DB 里两个 active 账号的 `quota_json.monthly_limit.reset_at` 都有值:
- `acct_81a...`:`monthly_limit.reset_at = 1785327099`,`window_minutes = 43200`
- `acct_fa9...`:`monthly_limit.reset_at = 1785325645`,`window_minutes = 43200`
但 `account_usage.window_reset_at` 都是 `null`,导致 `least_used` 无法执行“更早 reset 优先”的第二比较条件,只能退化到 `request_count` + LRU。
源码原因:
- `RuntimeAccountPoolService::apply_quota_snapshot()` 只有在 `quota_snapshot_reset_at()` 返回值时才调用 `sync_rate_limit_window()`。
- `src/upstream/accounts/quota/mod.rs` 当前 `quota_snapshot_reset_at()` 只有在 `monthly_limit_reached()` / `spend_control_limit_reached()` 时才读取 `/monthly_limit/reset_at`;非触顶 active 账号不会把 monthly/core reset 暴露给调度窗口。
- TS `src/routes/shared/proxy-rate-limit.ts` 对被动 rate-limit 同步是只要 `rateLimits.primary.reset_at != null` 就调用 `syncRateLimitWindow()`,不要求触顶。
修复目标:
- 非触顶 core rate-limit 窗口也要返回 reset/window,用于同步 `account_usage.window_reset_at` 和 `limit_window_seconds`。
- 当触顶来源是 monthly/spend control 时,冷却时间仍优先使用 monthly reset。
## 2026-06-29 修复记录
### quota reset 同步
已修改 `src/upstream/accounts/quota/mod.rs`:
- 非触顶 core bucket 只要包含 `reset_at`,`quota_snapshot_reset_at()` 就返回该窗口。
- 非触顶 core bucket 只要包含 `window_minutes`,`quota_snapshot_limit_window_seconds()` 就返回窗口秒数。
- 当 core 窗口不存在时,再回退到 `monthly_limit.reset_at` / `monthly_limit.window_minutes`。
- 当 monthly/spend-control 触顶时,仍优先使用 monthly reset 作为冷却时间。
已修改 `src/upstream/accounts/store.rs`:
- `list_pool_accounts()` / `get_pool_account()` 查询加入 `accounts.quota_json`。
- 恢复运行时账号池时,如果 `account_usage.window_reset_at` 或 `limit_window_seconds` 为空,会从 `quota_json` 推导窗口信息。
- 这不是兼容旧路径,而是把持久化 quota 快照纳入运行时调度输入,避免重启后智能分配退化。
已补测试:
- `quota_snapshot_reset_at_should_use_non_blocking_core_window`
- `quota_snapshot_reset_at_should_prefer_core_window_over_preserved_monthly_limit`
- `account_repository_should_restore_window_from_quota_json_when_usage_window_is_missing`
本轮已通过:
- `cargo test --test main quota_snapshot_reset_at -- --nocapture`
- `cargo test --test main account_pool -- --nocapture`
- `cargo test --test main account_repository -- --nocapture`
- `cargo check`
### 调度排查日志
已修改 `src/upstream/accounts/pool.rs`,账号选择时输出:
- `model`
- `rotation_strategy`
- `selection_source`
- `account_id`
- `candidate_count`
- `request_count`
- `window_request_count`
- `window_reset_at`
- `last_used_at`
- `quota_limit_reached`
- `quota_cooldown_until`
- `previous_slot_at`
- `rotation_cursor`
已修改 `src/upstream/transport/client.rs` / `src/upstream/transport/websocket_pool.rs`:
- `websocket pool decision` 增加 `conversation_id_hash`。
- `websocket pool decision` 增加 `ws_pool_key_hash`,用于判断同账号同 conversation 是否命中同一个 pool key。
- WebSocket fallback warn 已包含 `request_id` / `account_id` / `transport` / `fallback_transport` / `fallback_reason` / `error`。
### release 计数时机
已修改 `src/upstream/accounts/pool.rs`:
- `acquire` 只占用在途槽位,不再更新 `request_count` / `last_used_at` / `window_request_count`。
- 在途槽位记录 `created_at + model`,保留 TS 的 acquire lock 语义,同时让 release 能恢复对应 model。
- `release` 阶段更新 `request_count` / `last_used_at` / `window_request_count`。
- 如果长流式请求的槽位已被 stale cleanup 释放容量,`release` 仍按 TS 的 `recordUsage(entryId)` 语义更新账号级 request/LRU;此时只跳过无法恢复 model 的模型维度 request 计数。
已修改 `RuntimeAccountPoolService`:
- 去掉 `acquire_with()` 阶段的 `record_request()` 和 model request 计数。
- `release()` 阶段持久化账号 request 计数;存在 active slot model 时同步持久化 model request 计数。
这对齐 TS:
- TS `AccountLifecycle.acquire()` 只 `pushSlot()`。
- TS `AccountLifecycle.release()` 才 `popSlot()` 并调用 `registry.recordUsage()`。
本轮已通过:
- `cargo test --test main account_pool -- --nocapture`
- `cargo test --test main responses_should_record_request_count_when_5xx_retries_are_exhausted -- --nocapture`
- `cargo test --test main usage_logging -- --nocapture`
- `cargo test --test main account_repository -- --nocapture`
- `cargo check`
补充审计:
- TS `AccountLifecycle.release()` 先 `popSlot()`,随后无条件调用 `registry.recordUsage(entryId)`;`popSlot()` 找不到槽位不会阻止账号级 request/LRU 计数。
- Rust 已补齐该边界:slot 缺失只影响 model request 计数,不影响账号级 `request_count` / `last_used_at`。
- 已补测试:`release_should_mark_usage_after_stale_slot_cleanup`。
## 当前结论
`least_used` / 智能分配已经按 TS 对齐:
1. quota exhausted / quota limited 账号靠后或按配置跳过。
2. 两个账号都有 `window_reset_at` 时,优先选择更早 reset 的账号。
3. reset 无法比较时,选择更低 `request_count`。
4. request count 相同后,按更早 `last_used_at`。
5. 只有完整比较结果相同的首组选项才轮转。
6. request 计数和 LRU 更新时间已从 acquire 改为 release。
旧日志总是跳账号的原因:
- 当时两个 active 账号的 `account_usage.window_reset_at` 都是 `null`。
- 虽然 `accounts.quota_json.monthly_limit.reset_at` 有值,但旧 Rust 没把非触顶 reset 暴露给调度,也没在恢复账号池时从 `quota_json` 回填。
- 因此智能分配无法执行 TS 文案里的“优先使用即将刷新额度的账号”,只能退化到 `request_count + LRU`。
- 两个账号 `request_count` 又相同,连续请求自然容易 A/B/A/B。
- Rust 旧实现还在 acquire 时立即增加 request/LRU,进一步放大连续请求中的交替。
## 真实链路验证目标
- 修改后重启服务,运行同一组连续自然文本请求。
- 观察 `account selected for upstream request` 的 `window_reset_at` 是否不再为空。
- 在 `least_used` 下,如果两个账号都有不同 reset,应该优先选更早 reset 的账号,而不是仅靠 request_count/LRU 来回跳。
- 观察 `request_count` / `last_used_at` 是否只在 release 后影响下一次选择。
## 2026-06-29 真实链路验证
运行环境:
- 新二进制临时实例端口 `18080`,不影响旧 `8080` 进程。
- 审计目录:`.runtime/real-chain-scheduler-20260629_225737-least-used-release/`
- 请求:两条连续 `/v1/responses` stream,自然文本,`rotation_strategy=least_used`。
结果:
- `least_used_case=1`:HTTP 200,选中 `acct_fa9f0172d2084eaf86ddbd0e5d0c1c16`。
- `least_used_case=2`:HTTP 200,继续选中 `acct_fa9f0172d2084eaf86ddbd0e5d0c1c16`。
- 两次 `account selected for upstream request` 均为:
- `rotation_strategy=least_used`
- `selection_source=least_used`
- `candidate_count=2`
- `window_reset_at=Some(2026-07-29T11:47:25Z)`
- `quota_limit_reached=false`
- 第二次选择时,`request_count` 从第一次的 `19` 变为 `20`,`last_used_at` 变成第一次流结束后的时间 `2026-06-29T14:58:36.611782973+00:00`,验证 request/LRU 计数已经发生在 release 后。
- WebSocket pool 第一次 `new`,第二次 `reuse`,`conversation_id_hash=f1128ab74b73`、`ws_pool_key_hash=8899d9457e43` 一致。
结论:
- 修复后 `least_used` 没有再在两个账号之间 A/B 跳转。
- 当前选择 `acct_fa9...` 是符合 TS 文案和比较器的:它的 reset 早于 `acct_81a...`,应优先使用即将刷新额度的账号。
### release 边界补充验证
运行环境:
- 新二进制临时实例端口 `18080`,从既有 run config 启动,不修改 `8080` 配置。
- 使用管理登录创建临时 client API key,测试结束后已删除。
- 审计目录:`.runtime/real-chain-scheduler-20260629_225737-least-used-release/after-release-edge/`
- WebSocket 审计目录:`.runtime/real-chain-scheduler-20260629_225737-least-used-release/ws-audit-after-release-edge/`
- 请求:两条连续 `/v1/responses` stream,自然文本,同一个 `prompt_cache_key`,`rotation_strategy=least_used`。
结果:
- `least-used-after-1`:HTTP 200,选中 `acct_fa9f0172d2084eaf86ddbd0e5d0c1c16`。
- `request_id=req_aeecdf27-7319-43dd-9abc-2cdf14ea5e5a`
- `request_count=21`
- `window_request_count=21`
- `window_reset_at=Some(2026-07-29T11:47:25Z)`
- `websocket_pool_kind=Some("new")`
- `conversation_id_hash=31a3f63f6dd8`
- `ws_pool_key_hash=5e0f49abf9bb`
- `least-used-after-2`:HTTP 200,继续选中 `acct_fa9f0172d2084eaf86ddbd0e5d0c1c16`。
- `request_id=req_d60506fb-9553-4cd0-9fca-1f2ceef0d7ff`
- `request_count=22`
- `window_request_count=22`
- `last_used_at=2026-06-29T15:10:33.720793920+00:00`
- `window_reset_at=Some(2026-07-29T11:47:25Z)`
- `websocket_pool_kind=Some("reuse")`
- `conversation_id_hash=31a3f63f6dd8`
- `ws_pool_key_hash=5e0f49abf9bb`
补充结论:
- 第一条完成后,第二条选择前已看到 `request_count` 从 21 增至 22,说明 release 阶段计数生效。
- 两条连续请求没有 A/B 跳账号;在两个候选都有 reset 且 `acct_fa9...` reset 更早时,智能分配持续选它,符合 TS 文案。
- 两条请求的 `conversation_id_hash` 与 `ws_pool_key_hash` 一致,第二条 WS 复用成功。
+114 -52
View File
@@ -45,6 +45,16 @@ pub enum RotationStrategy {
Sticky,
}
impl RotationStrategy {
fn as_str(self) -> &'static str {
match self {
Self::LeastUsed => "least_used",
Self::RoundRobin => "round_robin",
Self::Sticky => "sticky",
}
}
}
/// 账号池配置。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountPoolOptions {
@@ -169,11 +179,24 @@ pub struct AccountWindowUsageDelta {
#[derive(Debug)]
pub struct AccountPool {
accounts: IndexMap<String, Account>,
slots: BTreeMap<String, VecDeque<DateTime<Utc>>>,
slots: BTreeMap<String, VecDeque<AccountSlot>>,
options: AccountPoolOptions,
rotation_cursor: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct AccountSlot {
created_at: DateTime<Utc>,
model: String,
}
/// 账号释放结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReleasedAccountUsage {
/// 被释放槽位对应的模型。槽位已被过期清理时为空。
pub model: Option<String>,
}
impl Default for AccountPool {
fn default() -> Self {
Self::with_options(AccountPoolOptions::default())
@@ -383,7 +406,13 @@ impl AccountPool {
.iter()
.find(|account| account.id == *preferred_account_id)
.cloned();
preferred.map(|account| (account, self.previous_slot_at(preferred_account_id)))
preferred.map(|account| {
(
account,
self.previous_slot_at(preferred_account_id),
"preferred",
)
})
} else {
None
}
@@ -392,35 +421,48 @@ impl AccountPool {
RotationStrategy::LeastUsed => {
let account = self.select_least_used(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous)
(account, previous, "least_used")
}
RotationStrategy::RoundRobin => {
let account = self.select_round_robin(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous)
(account, previous, "round_robin")
}
RotationStrategy::Sticky => {
let account = Self::select_sticky(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous)
(account, previous, "sticky")
}
})
});
let Some((selected, previous_slot_at)) = selected else {
let Some((selected, previous_slot_at, selection_source)) = selected else {
return AccountAcquireWithStatusRefresh {
acquired: None,
expired_account_ids,
};
};
tracing::info!(
model = %request.model,
rotation_strategy = self.options.rotation_strategy.as_str(),
selection_source,
account_id = %selected.id,
candidate_count = candidates.len(),
request_count = selected.request_count,
window_request_count = selected.window_request_count,
window_reset_at = ?selected.window_reset_at,
last_used_at = selected.last_used_at.as_deref().unwrap_or_default(),
quota_limit_reached = selected.quota_limit_reached,
quota_cooldown_until = ?selected.quota_cooldown_until,
previous_slot_at = ?previous_slot_at,
rotation_cursor = self.rotation_cursor,
"account selected for upstream request"
);
let selected_id = selected.id.clone();
self.push_slot(&selected_id, request.now);
let acquired = self
.mark_usage(&selected_id, request.now)
.unwrap_or(selected);
self.push_slot(&selected_id, request.now, &request.model);
AccountAcquireWithStatusRefresh {
acquired: Some(AcquiredAccount {
account: acquired,
account: selected,
previous_slot_at,
}),
expired_account_ids,
@@ -445,14 +487,22 @@ impl AccountPool {
}
/// 释放指定账号的一个在途槽位。
pub fn release(&mut self, account_id: &str) {
let Some(slots) = self.slots.get_mut(account_id) else {
return;
};
slots.pop_front();
if slots.is_empty() {
pub fn release(&mut self, account_id: &str) -> Option<ReleasedAccountUsage> {
self.accounts.get(account_id)?;
let mut remove_slots = false;
let slot = self.slots.get_mut(account_id).and_then(|slots| {
let slot = slots.pop_front();
remove_slots = slots.is_empty();
slot
});
if remove_slots {
self.slots.remove(account_id);
}
self.mark_request_usage(account_id, Utc::now());
Some(ReleasedAccountUsage {
model: slot.map(|slot| slot.model),
})
}
/// 计算账号池容量摘要。
@@ -563,17 +613,21 @@ impl AccountPool {
fn previous_slot_at(&self, account_id: &str) -> Option<DateTime<Utc>> {
self.slots
.get(account_id)
.and_then(|slots| slots.back().copied())
.and_then(|slots| slots.back())
.map(|slot| slot.created_at)
}
fn push_slot(&mut self, account_id: &str, now: DateTime<Utc>) {
fn push_slot(&mut self, account_id: &str, now: DateTime<Utc>, model: &str) {
self.slots
.entry(account_id.to_string())
.or_default()
.push_back(now);
.push_back(AccountSlot {
created_at: now,
model: model.to_string(),
});
}
fn mark_usage(&mut self, account_id: &str, now: DateTime<Utc>) -> Option<Account> {
fn mark_request_usage(&mut self, account_id: &str, now: DateTime<Utc>) -> Option<Account> {
let account = self.accounts.get_mut(account_id)?;
account.last_used_at = Some(now.to_rfc3339());
account.request_count = account.request_count.saturating_add(1);
@@ -612,7 +666,7 @@ impl AccountPool {
fn cleanup_stale_slots(&mut self, now: DateTime<Utc>) {
let ttl = self.options.stale_slot_ttl;
self.slots.retain(|_, slots| {
slots.retain(|slot_at| now.signed_duration_since(*slot_at) <= ttl);
slots.retain(|slot| now.signed_duration_since(slot.created_at) <= ttl);
!slots.is_empty()
});
}
@@ -732,7 +786,43 @@ impl RuntimeAccountPoolService {
/// 释放账号在途槽位。
pub async fn release(&self, account_id: &str) {
self.pool.lock().await.release(account_id);
let released = self.pool.lock().await.release(account_id);
let Some(released) = released else {
return;
};
if let Err(error) = self.store.record_request(account_id).await {
tracing::warn!(
account_id,
error = %error,
"failed to persist account request usage"
);
}
let Some(model) = released.model else {
tracing::debug!(
account_id,
"released account without active slot; skipped model request usage persistence"
);
return;
};
if let Err(error) = self
.store
.record_model_usage_delta(
account_id,
&model,
AccountModelUsageDelta {
requests: 1,
..AccountModelUsageDelta::default()
},
)
.await
{
tracing::warn!(
account_id,
model,
error = %error,
"failed to persist account model request usage"
);
}
}
/// 更新账号状态并同步内存池。
@@ -775,38 +865,10 @@ impl RuntimeAccountPoolService {
/// 按请求上下文获取可用账号。
pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option<AcquiredAccount> {
let model = request.model.clone();
let refresh = self.pool.lock().await.acquire_with_status_refresh(request);
self.persist_expired_statuses(refresh.expired_account_ids)
.await;
let acquired = refresh.acquired?;
if let Err(error) = self.store.record_request(&acquired.account.id).await {
tracing::warn!(
account_id = acquired.account.id,
error = %error,
"failed to persist account request usage"
);
}
if let Err(error) = self
.store
.record_model_usage_delta(
&acquired.account.id,
&model,
AccountModelUsageDelta {
requests: 1,
..AccountModelUsageDelta::default()
},
)
.await
{
tracing::warn!(
account_id = acquired.account.id,
model,
error = %error,
"failed to persist account model request usage"
);
}
Some(acquired)
refresh.acquired
}
async fn persist_expired_statuses(&self, account_ids: Vec<String>) {
+45 -6
View File
@@ -19,12 +19,13 @@ pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
/// 从配额快照读取 UTC 重置时间。
pub fn quota_snapshot_reset_at(quota: &Value) -> Option<DateTime<Utc>> {
let reset_at = if monthly_limit_reached(quota) || spend_control_limit_reached(quota) {
quota
.pointer("/monthly_limit/reset_at")
.and_then(positive_i64)
monthly_limit_reset_at(quota)
.or_else(|| blocking_core_bucket(quota).and_then(bucket_reset_at))
.or_else(|| core_bucket_with_reset_at(quota).and_then(bucket_reset_at))
} else {
blocking_core_bucket(quota).and_then(bucket_reset_at)
core_bucket_with_reset_at(quota)
.and_then(bucket_reset_at)
.or_else(|| monthly_limit_reset_at(quota))
}?;
DateTime::<Utc>::from_timestamp(reset_at, 0)
}
@@ -32,11 +33,12 @@ pub fn quota_snapshot_reset_at(quota: &Value) -> Option<DateTime<Utc>> {
/// 从配额快照读取窗口大小。
pub fn quota_snapshot_limit_window_seconds(quota: &Value) -> Option<u64> {
if monthly_limit_reached(quota) || spend_control_limit_reached(quota) {
return Some(MONTH_WINDOW_SECONDS);
return monthly_limit_window_seconds(quota).or(Some(MONTH_WINDOW_SECONDS));
}
blocking_core_bucket(quota)
core_bucket_with_window_minutes(quota)
.and_then(bucket_window_minutes)
.and_then(|minutes| minutes.checked_mul(60))
.or_else(|| monthly_limit_window_seconds(quota))
}
/// 从 Codex usage 响应中提取持久化配额快照。
@@ -291,6 +293,29 @@ fn blocking_core_bucket(quota: &Value) -> Option<&Value> {
})
}
fn core_bucket_with_reset_at(quota: &Value) -> Option<&Value> {
core_buckets(quota).find(|bucket| bucket_reset_at(bucket).is_some())
}
fn core_bucket_with_window_minutes(quota: &Value) -> Option<&Value> {
core_buckets(quota).find(|bucket| bucket_window_minutes(bucket).is_some())
}
fn core_buckets<'a>(quota: &'a Value) -> impl Iterator<Item = &'a Value> + 'a {
quota
.get("snapshots")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter(|snapshot| snapshot.get("source").and_then(Value::as_str) == Some("core"))
.flat_map(|snapshot| {
["primary", "secondary"]
.into_iter()
.filter_map(|key| snapshot.get(key))
})
.filter(|bucket| !bucket.is_null())
}
fn bucket_reset_at(bucket: &Value) -> Option<i64> {
bucket.get("reset_at").and_then(positive_i64)
}
@@ -309,6 +334,20 @@ fn credits_overage_limit_reached(quota: &Value) -> bool {
.unwrap_or(false)
}
fn monthly_limit_reset_at(quota: &Value) -> Option<i64> {
quota
.pointer("/monthly_limit/reset_at")
.and_then(positive_i64)
}
fn monthly_limit_window_seconds(quota: &Value) -> Option<u64> {
quota
.pointer("/monthly_limit/window_minutes")
.and_then(Value::as_u64)
.filter(|value| *value > 0)
.and_then(|minutes| minutes.checked_mul(60))
}
fn number_value(value: &Value) -> Option<f64> {
value
.as_f64()
+17 -2
View File
@@ -11,6 +11,9 @@ use crate::infra::json::{decode_cursor, page_offset, NumberedPage, Page};
use crate::upstream::accounts::model::{
Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta,
};
use crate::upstream::accounts::quota::{
quota_snapshot_limit_window_seconds, quota_snapshot_reset_at,
};
// ============================================================================
// SQL 常量
@@ -32,6 +35,7 @@ select
a.quota_limit_reached,
a.quota_verify_required,
a.quota_cooldown_until,
a.quota_json,
a.cloudflare_cooldown_until,
a.added_at,
coalesce(au.request_count, 0) as usage_request_count,
@@ -72,6 +76,7 @@ select
a.quota_limit_reached,
a.quota_verify_required,
a.quota_cooldown_until,
a.quota_json,
a.cloudflare_cooldown_until,
a.added_at,
coalesce(au.request_count, 0) as usage_request_count,
@@ -1632,6 +1637,14 @@ async fn get_pool_account(
}
fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult<Account> {
let quota_json = row
.get::<Option<String>, _>("quota_json")
.and_then(|quota_json| serde_json::from_str::<Value>(&quota_json).ok());
let quota_window_reset_at = quota_json.as_ref().and_then(quota_snapshot_reset_at);
let quota_limit_window_seconds = quota_json
.as_ref()
.and_then(quota_snapshot_limit_window_seconds);
Ok(Account {
id: row.get("id"),
email: row.get("email"),
@@ -1698,10 +1711,12 @@ fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult<Account> {
window_reset_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("usage_window_reset_at")
.as_deref(),
)?,
)?
.or(quota_window_reset_at),
limit_window_seconds: optional_positive_i64_to_u64(
row.get::<Option<i64>, _>("usage_limit_window_seconds"),
),
)
.or(quota_limit_window_seconds),
added_at: row.get("added_at"),
last_used_at: row.get("usage_last_used_at"),
})
+24
View File
@@ -581,6 +581,7 @@ impl CodexBackendClient {
tracing::warn!(error = %error, "failed to write Codex WebSocket audit artifact");
}
let pool_key = self.websocket_pool_key(upstream_request, context, pool_account_id);
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
let exchange = match (self.websocket_pool.as_deref(), pool_key) {
(Some(pool), Some(key)) => {
execute_response_create_request_with_pool(&prepared, Some((pool, key)), started_at)
@@ -592,6 +593,7 @@ impl CodexBackendClient {
log_websocket_pool_decision(
context.request_id,
pool_account_id.or(context.account_id),
pool_log_context.as_ref(),
exchange.pool_decision,
);
@@ -631,6 +633,7 @@ impl CodexBackendClient {
tracing::warn!(error = %error, "failed to write Codex WebSocket audit artifact");
}
let pool_key = self.websocket_pool_key(upstream_request, context, pool_account_id);
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
let exchange = match (self.websocket_pool.as_deref(), pool_key) {
(Some(pool), Some(key)) => {
execute_response_create_request_stream_with_pool(&prepared, Some((pool, key))).await
@@ -641,6 +644,7 @@ impl CodexBackendClient {
log_websocket_pool_decision(
context.request_id,
pool_account_id.or(context.account_id),
pool_log_context.as_ref(),
exchange.pool_decision,
);
@@ -961,6 +965,7 @@ fn update_first_token_ms(started_at: Instant, body_bytes: &[u8], first_token_ms:
fn log_websocket_pool_decision(
request_id: &str,
account_id: Option<&str>,
pool_context: Option<&WebSocketPoolLogContext>,
decision: Option<WebSocketPoolDecision>,
) {
let Some(decision) = decision else {
@@ -974,6 +979,8 @@ fn log_websocket_pool_decision(
account_id = account_id.unwrap_or_default(),
ws_pool = decision.kind(),
ws_pool_reason = reason,
conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()),
ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()),
"websocket pool decision"
);
} else {
@@ -982,11 +989,28 @@ fn log_websocket_pool_decision(
rid = %rid_short,
account_id = account_id.unwrap_or_default(),
ws_pool = decision.kind(),
conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()),
ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()),
"websocket pool decision"
);
}
}
#[derive(Debug, Clone)]
struct WebSocketPoolLogContext {
conversation_id_hash: String,
pool_key_hash: String,
}
impl WebSocketPoolLogContext {
fn from_key(key: &CodexWebSocketPoolKey) -> Self {
Self {
conversation_id_hash: key.conversation_id_hash(),
pool_key_hash: key.stable_hash(),
}
}
}
fn elapsed_millis_i64(started_at: Instant) -> i64 {
started_at.elapsed().as_millis().min(i64::MAX as u128) as i64
}
+22
View File
@@ -8,6 +8,7 @@ use std::{
use futures::{SinkExt, StreamExt};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use tokio::{net::TcpStream, sync::Mutex, time::timeout};
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
@@ -43,6 +44,18 @@ impl CodexWebSocketPoolKey {
pub(crate) fn account_id(&self) -> &str {
&self.account_id
}
pub(crate) fn conversation_id_hash(&self) -> String {
short_sha256([self.conversation_id.as_str()])
}
pub(crate) fn stable_hash(&self) -> String {
short_sha256([
self.base_url.as_str(),
self.account_id.as_str(),
self.conversation_id.as_str(),
])
}
}
/// WebSocket 连接池。
@@ -504,6 +517,15 @@ impl WebSocketPoolDecisionKind {
}
}
fn short_sha256<'a>(parts: impl IntoIterator<Item = &'a str>) -> String {
let mut hasher = Sha256::new();
for part in parts {
hasher.update(part.as_bytes());
hasher.update(b"\0");
}
hex::encode(hasher.finalize()).chars().take(12).collect()
}
fn account_slot_count(
slots: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
account_id: &str,
@@ -126,7 +126,7 @@ fn acquire_should_refresh_expired_cooldowns_before_selecting_account() {
assert_eq!(acquired.account.status, AccountStatus::Active);
assert!(acquired.account.quota_cooldown_until.is_none());
assert!(acquired.account.cloudflare_cooldown_until.is_none());
assert_eq!(acquired.account.window_request_count, 1);
assert_eq!(acquired.account.window_request_count, 0);
assert!(acquired
.account
.window_reset_at
@@ -209,13 +209,12 @@ fn sync_rate_limit_window_should_reset_window_counters_when_drift_crosses_thresh
let acquired = pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
pool.release(&acquired.account.id);
let released = pool.get(&acquired.account.id).unwrap();
assert_eq!(acquired.account.window_request_count, 1);
assert_eq!(acquired.account.window_input_tokens, 0);
assert_eq!(acquired.account.window_output_tokens, 0);
assert_eq!(acquired.account.window_cached_tokens, 0);
assert_eq!(
acquired.account.window_reset_at,
Some(now + Duration::seconds(121))
);
assert_eq!(released.window_request_count, 1);
assert_eq!(released.window_input_tokens, 0);
assert_eq!(released.window_output_tokens, 0);
assert_eq!(released.window_cached_tokens, 0);
assert_eq!(released.window_reset_at, Some(now + Duration::seconds(121)));
}
@@ -385,6 +385,7 @@ fn least_used_should_prefer_lru_before_tie_rotation() {
.id,
"seed"
);
pool.release("seed");
let acquired = pool
.acquire_with(&AccountAcquireRequest::new(
"gpt-5.5",
@@ -75,7 +75,7 @@ fn least_used_should_compare_request_count_when_only_one_account_has_window_rese
}
#[test]
fn acquire_should_mark_runtime_usage_window() {
fn release_should_mark_runtime_usage_window() {
let now = fixed_time();
let mut pool = AccountPool::default();
let mut account = crate::support::accounts::test_account("acct_a", AccountStatus::Active);
@@ -86,17 +86,42 @@ fn acquire_should_mark_runtime_usage_window() {
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
assert_eq!(acquired.account.request_count, 1);
assert_eq!(acquired.account.window_request_count, 1);
assert_eq!(acquired.account.window_started_at, Some(now));
assert_eq!(
acquired.account.window_reset_at,
Some(now + Duration::seconds(60))
);
assert_eq!(
acquired.account.last_used_at.as_deref(),
Some(now.to_rfc3339().as_str())
);
assert_eq!(acquired.account.request_count, 0);
assert_eq!(acquired.account.window_request_count, 0);
pool.release(&acquired.account.id);
let released = pool.get(&acquired.account.id).unwrap();
assert_eq!(released.request_count, 1);
assert_eq!(released.window_request_count, 1);
assert!(released.window_started_at.is_some());
assert!(released.window_reset_at.is_some());
assert!(released.last_used_at.is_some());
}
#[test]
fn release_should_mark_usage_after_stale_slot_cleanup() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
stale_slot_ttl: Duration::minutes(5),
..AccountPoolOptions::default()
});
pool.insert(crate::support::accounts::test_account(
"acct_a",
AccountStatus::Active,
));
let acquired = pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
let summary = pool.capacity_summary(now + Duration::minutes(6));
pool.release(&acquired.account.id);
let released = pool.get(&acquired.account.id).unwrap();
assert_eq!(summary.used_slots, 0);
assert_eq!(released.request_count, 1);
assert_eq!(released.window_request_count, 1);
}
#[test]
@@ -256,6 +256,53 @@ async fn account_repository_should_restore_window_usage_into_runtime_pool_accoun
assert_eq!(account.limit_window_seconds, Some(300));
}
#[tokio::test]
async fn account_repository_should_restore_window_from_quota_json_when_usage_window_is_missing() {
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 36).await;
let repo = SqliteAccountStore::new(pool.clone());
let reset_at = DateTime::<Utc>::from_timestamp(1_806_364_800, 0).unwrap();
seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await;
sqlx::query("update accounts set quota_json = ? where id = ?")
.bind(
r#"{
"snapshots": [{
"source": "core",
"blocked": false,
"primary": {
"used_percent": 11,
"remaining_percent": 89,
"reset_at": 1806364800,
"window_minutes": 43200,
"limit_reached": false
},
"secondary": null
}],
"monthly_limit": {
"key": "core-monthly",
"source": "rate_limit",
"used_percent": 11,
"remaining_percent": 89,
"reset_at": 1806364800,
"window_minutes": 43200,
"limit_reached": false
}
}"#,
)
.bind("acct_a")
.execute(&pool)
.await
.unwrap();
let account = AccountStore::list_pool_accounts(&repo)
.await
.unwrap()
.remove(0);
assert_eq!(account.window_reset_at, Some(reset_at));
assert_eq!(account.limit_window_seconds, Some(2_592_000));
}
#[tokio::test]
async fn account_repository_should_restore_quota_verify_required_into_runtime_pool_accounts() {
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 34).await;
+58 -1
View File
@@ -5,7 +5,8 @@ use codex_proxy_rs::infra::database::connect_sqlite;
use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore;
use codex_proxy_rs::upstream::accounts::model::AccountStatus;
use codex_proxy_rs::upstream::accounts::quota::{
quota_from_usage, quota_snapshot_limit_reached, RuntimeQuotaRefreshService,
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at, RuntimeQuotaRefreshService,
};
use codex_proxy_rs::upstream::accounts::store::{NewAccount, SqliteAccountStore};
use codex_proxy_rs::upstream::transport::CodexBackendClient;
@@ -156,6 +157,62 @@ fn quota_snapshot_limit_reached_should_keep_allowed_free_account_without_reset_c
assert!(!quota_snapshot_limit_reached(&quota));
}
#[test]
fn quota_snapshot_reset_at_should_use_non_blocking_core_window() {
let usage = json!({
"plan_type": "free",
"rate_limit": {
"allowed": true,
"limit_reached": false,
"primary_window": {
"used_percent": 6,
"reset_at": 1_806_364_800,
"limit_window_seconds": 2_592_000
}
}
});
let quota = quota_from_usage(&usage);
assert_eq!(
quota_snapshot_reset_at(&quota).map(|reset_at| reset_at.timestamp()),
Some(1_806_364_800)
);
assert_eq!(quota_snapshot_limit_window_seconds(&quota), Some(2_592_000));
}
#[test]
fn quota_snapshot_reset_at_should_prefer_core_window_over_preserved_monthly_limit() {
let quota = json!({
"snapshots": [{
"source": "core",
"blocked": false,
"primary": {
"used_percent": 25,
"remaining_percent": 75,
"reset_at": 1_893_456_300,
"window_minutes": 5,
"limit_reached": false
},
"secondary": null
}],
"monthly_limit": {
"key": "spend-control-monthly",
"source": "spend_control",
"used_percent": 52,
"remaining_percent": 48,
"reset_at": 1_896_048_000,
"window_minutes": 43200,
"limit_reached": false
}
});
assert_eq!(
quota_snapshot_reset_at(&quota).map(|reset_at| reset_at.timestamp()),
Some(1_893_456_300)
);
assert_eq!(quota_snapshot_limit_window_seconds(&quota), Some(300));
}
#[test]
fn quota_snapshot_limit_reached_should_block_explicit_credit_overage_limit() {
let quota = json!({