starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 16:06:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: improve code quality

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-14 16:07:57 +08:00
co-authored by Codex
parent 0d86975ca6
commit 48230ae72a
18 changed files with 3302 additions and 2362 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/code-quality-audit.md
+403
View File
@@ -0,0 +1,403 @@
# 代码质量审计与改造计划
日期：2026-06-14
范围：当前仓库 `src/`、`tests/`，不包含 `designs/`。
目标：确认 `/tmp/code_quality_audit.md` 的结论，并基于当前代码给出后续改造顺序。项目处于早期阶段，不需要保留低质量兼容层；改造应以清晰边界、可测试、可维护为目标。
## 审计原则
1. 不为了“拆小文件”而拆文件。文件大小本身不是问题，问题是职责混杂、变更理由混在一起、审查和定位成本过高。
2. 不用“兼容历史”掩盖坏设计。内部 API 可以在测试保护下重构，外部行为必须由测试证明未回退。
3. Rust 代码遵循本项目的 `$rust-best-practices`：明确错误类型、少 clone、少分配、少隐式 panic，但性能优化要有证据。
4. 日志输出使用中文，保留必要英文专业术语，例如 WebSocket、SSE、OAuth、quota、token。
5. 每个改造切片先写或补齐测试，再改生产代码，最后跑目标测试、`cargo fmt --check`、`cargo clippy --all-targets --all-features --locked -- -D warnings`。
## 审计时量化事实
审计开始时 `src/` + `tests/` 约 37,055 行。生产代码里最需要关注的文件：
| 文件 | 行数 | 判断 |
|---|---:|---|
| `src/codex/serving/dispatch/mod.rs` | 1474 | 已有子模块，但主文件仍承担请求编排、上游调用、stream audit、fallback、限流处理等多种职责 |
| `src/codex/gateway/transport/websocket.rs` | 1175 | WebSocket pool、连接建立、请求编码、SSE 转发、错误分类、rate limit 捕获混在同一文件 |
| `src/codex/accounts/repository.rs` | 1103 | 账号 CRUD、token 更新、quota、usage、lease、pool read model 混在同一存储文件 |
| `src/codex/gateway/transport/client.rs` | 783 | 上游 transport 编排较重，但当前还可接受 |
| `src/codex/accounts/pool.rs` | 627 | 调度策略集中，复杂度主要来自业务策略，不是首要拆分对象 |
## Claude 审计确认
### 1. `update_tokens` / `update_from_claims` 重复
结论：成立。
当前 `AccountRepository::update_tokens` 和 `AccountRepository::update_from_claims` 都重复处理：
- `Utc::now().to_rfc3339()`
- access token 加密
- refresh token 可选加密
- expires_at 格式化
- “没有新 refresh_token 时保留旧值”的分支
需要修复，但不建议照搬报告里的 `QueryBuilder` 方案。当前只有几个固定 SQL 形态，动态构建 SQL 会把绑定顺序变成新的维护风险。更合适的第一步是：
- 补测试：`update_from_claims` 省略 refresh_token 时必须保留旧 refresh_token。
- 提取私有 token 写入准备 helper。
- 将 UPDATE SQL 改成命名多行常量。
- 保持 `disabled` / `banned` 不被 `update_tokens` 重新激活的现有语义。
### 2. SQL 字符串地狱
结论：部分成立。
报告里的状态已经有些过期：`LIST_POOL_ACCOUNTS_SQL`、`RECORD_USAGE_SQL`、`GET_USAGE_SQL` 等已是多行常量。但问题仍存在：
- `insert/get/list/list_metadata/update_tokens/update_from_claims/quota/lease/AccountUsageRepository::list/summary` 仍有较多长单行 SQL。
- 一部分 SQL 的业务语义不容易在 review 时看清，例如 token 更新时的状态保护、quota cooldown 的最大值保留、refresh lease 的抢占条件。
建议先使用命名多行常量。`sqlx::query_file!` 不是当前第一选择，因为项目当前使用运行时 `sqlx::query`，切换到宏式编译期校验会引入额外数据库元数据/离线缓存约束，改造面比可读性收益更大。
### 3. `AccountRepository` God Class
结论：部分成立。
`AccountUsageRepository` 已经存在，说明项目已经开始拆 usage 读模型；但 `AccountRepository` 里仍混有：
- 账号 CRUD
- token 更新
- pool read model
- quota JSON 与 cooldown
- usage record/get/window sync
- refresh lease
这确实不是理想边界。后续可以拆，但不应机械拆成很多小文件。推荐按变更理由拆：
- account 基础读写
- token 写入
- quota 状态
- usage 统计写入与窗口同步
- refresh lease
- pool read model
是否暴露多个 repository 类型，需要看调用侧是否真的收益明显。可以先做内部模块拆分，再决定是否调整服务层依赖。
### 4. 超大文件
结论：成立，但“拆文件”不是目标。
三个文件已经大到影响定位和审查：
- `dispatch/mod.rs`
- `websocket.rs`
- `repository.rs`
但文件拆分必须跟职责边界一致：
- `websocket.rs` 适合按 pool、connection、request codec、SSE forwarding、error classification 拆。
- `dispatch/mod.rs` 适合把 stream audit、fallback/rate-limit 应用、上游请求执行拆到更明确的模块。
- `repository.rs` 适合先把 SQL 与行映射收拢，再拆 token/quota/usage/lease。
不建议为了把每个文件压到某个行数而过度碎片化。
### 5. `Account` clone 性能开销
结论：部分成立，不是第一优先级。
`AccountPool` 确实会 clone `Account`，而 `Account` 包含 token 字符串。这个成本存在，但当前设计也有合理性：`AccountPool` 在 `Mutex` 后面，返回 owned `Account` 可以避免把锁生命周期带入 async 上游调用。
直接改成 `Arc<Account>` 不是无成本优化，会牵涉：
- pool 内账号状态更新如何替换 snapshot
- 调度策略如何避免共享可变状态
- token 刷新和 cooldown 更新如何保持一致
建议先完成结构性重构。性能优化应结合 profile 或压测数据，再考虑 immutable snapshot + `Arc<Account>`。
### 6. 字符串分配过度
结论：审计过泛，暂不作为直接改造项。
`.to_string()` / `.clone()` 数量本身不能证明问题。当前很多转换来自：
- JSON/SSE 协议需要 owned `String`
- secret 解密后写入运行时 pool
- header、metadata、测试数据构造
- async 任务跨 await / spawn 的所有权要求
后续可以用 Clippy、profile 和热点路径 review 定点优化，不能按数量批量替换成 `Cow` 或引用。
## 我的补充审计
### P0：token 更新重复和测试缺口
`update_tokens` 已有测试覆盖“没有新 refresh_token 时保留旧值”，但 `update_from_claims` 缺同类测试。由于导入已有账号时会走 `update_from_claims`，这是第一批代码改造前必须补齐的行为测试。
改造目标：
- 增加 `update_from_claims` 保留 refresh_token 测试。
- 提取 token 写入准备逻辑。
- 将对应 UPDATE SQL 改为命名多行常量。
- 保持现有行为，不顺手重构调用层。
### P0：Repository SQL 可读性
`repository.rs` 的 SQL 可读性是当前最明显的维护问题。建议第一阶段只处理账号 repository 内的长 SQL，不扩散到全项目。
改造目标：
- 账号表常用 SELECT 语句提取为常量。
- token/quota/lease/usage summary 的长 SQL 改成多行常量。
- SQL 常量命名表达业务语义，例如 `UPDATE_TOKENS_PRESERVING_DISABLED_STATUS_SQL`。
### P1：Repository 边界
存储层已经有一些 DDD 边界，但 `AccountRepository` 仍承载太多变更理由。推荐在 SQL 可读性改善后再拆，避免一边搬文件一边改 SQL。
推荐顺序：
1. 先把 SQL 常量和 row mapper 收拢。
2. 再拆 token/quota/lease/usage 内部模块。
3. 最后评估是否需要多个 public repository。
### P1：WebSocket transport 边界
`websocket.rs` 当前实现已经覆盖原版核心链路，但职责过密：
- pool lifecycle
- idle probe / GC
- handshake request 构造
- request body 转换
- WS -> SSE 转发
- rate limit 内部事件捕获
- WS 错误分类
建议后续按“协议转换”和“连接池”先拆，不改变行为：
- `websocket/pool.rs`
- `websocket/connection.rs`
- `websocket/codec.rs`
- `websocket/stream.rs`
- `websocket/error.rs`
这个拆分是为了边界清楚，不是为了追求小文件。
### P1：Dispatch 编排边界
`dispatch/mod.rs` 已有 `affinity/fallback/refresh/routing/stream/usage` 子模块，但主文件仍过重。建议把以下逻辑迁出：
- 上游 HTTP SSE / WebSocket stream audit
- quota/rate-limit header 应用
- fallback account retry 决策
- 上游请求执行 wrapper
目标是让 `mod.rs` 更像编排入口，而不是承载所有细节。
### P1：日志一致性
当前日志大部分已经是中文，且不少已使用结构化字段。但仍需继续统一：
- 非中文日志消息需要改成中文，例如 UpdateChecker 相关日志。
- `error = %e`、`account_id = %account_id` 这类结构化字段应保持。
- 不把 token、cookie、refresh_token 放进日志。
- HTTP trace 日志已存在中文入口，应继续检查是否有 request_id/span 传播缺口。
### P2：生产代码里的 `unwrap()`
大多数 `unwrap()` 位于测试模块。生产路径中发现 `usage_snapshots.rs` 的时间桶对齐使用 `and_hms_opt(...).unwrap()`。逻辑上这些时间值可证明有效，但风格上仍建议改成有语义的 helper 或 `expect` 文案，避免以后被复制到不可证明场景。
### P2：测试文件过大
`tests/codex_gateway/websocket.rs` 约 1436 行，已经接近生产大文件的认知成本。测试拆分应按场景拆，而不是一个 endpoint 一个文件：
- WebSocket pool 复用/淘汰
- WebSocket fallback
- WS -> SSE 转换
- rate limit internal event
- error classification
### P2：注释质量
项目里已有中文注释，但后续应保持“解释为什么”，避免解释显而易见的代码行为。适合保留的注释包括：
- token/refresh_token 保留策略
- SSE 响应头已发出后不能回滚 HTTP 状态
- 原版协议兼容的非显然参数
- 安全链路中的指纹/安装 ID/turn state 传递原因
## 推荐改造顺序
### Phase 1：账号 repository P0
目标：先处理最确定、最小风险的问题。
1. 补 `update_from_claims` refresh_token 保留测试。
2. 提取 token 更新准备 helper。
3. UPDATE token/claims SQL 改命名多行常量。
4. 账号 repository 的明显长 SQL 改命名常量。
5. 跑 `cargo test codex_accounts::repository` 或对应过滤测试，再跑 fmt/clippy。
验收条件：
- `update_tokens` 和 `update_from_claims` 不再各自重复 token 加密、expires_at 格式化和 updated_at 生成逻辑。
- `update_from_claims` 省略 refresh_token 时保留旧 refresh_token，有测试覆盖。
- `update_tokens` 仍保持现有保护：`disabled` / `banned` 账号不会因为 token 更新被重新激活。
- token/claims/quota/lease/usage summary 的长 UPDATE/SELECT SQL 不再以内联单行字符串形式散落在方法体里。
- 不引入 `QueryBuilder`，除非实际需要动态字段集合。当前固定 SQL 形态优先使用命名多行常量，降低绑定顺序风险。
- 不调整调用层、不拆 public repository API、不改数据库 schema。
- 验证命令至少包括：
- `cargo test --test codex_accounts repository --locked`
- `cargo fmt --check`
- `cargo clippy --all-targets --all-features --locked -- -D warnings`
Phase 1 不包含：
- `AccountRepository` public API 拆分。
- WebSocket/Dispatch 文件拆分。
- `Account` 改 `Arc<Account>`。
- 全局 `.clone()` / `.to_string()` 批量替换。
- 日志系统重构。
Phase 1 状态：已完成。
已完成内容：
- 增加 `update_from_claims` 省略 `refresh_token` 时保留旧 refresh token 的测试。
- 提取 token 写入准备逻辑，统一 access token 加密、可选 refresh token 加密、expires_at 格式化和 updated_at 生成。
- 将 token/claims 更新 SQL 改为命名多行常量，并保留 `disabled` / `banned` 不被 token 更新重新激活的语义。
- 将账号 repository 内 usage、quota、refresh lease 的长 SQL 和实现迁出到内部模块，public API 保持不变。
验证结果：
- `cargo test --test codex_accounts repository --locked`：13 passed。
- `cargo test prepare_token_write_should_encrypt_tokens_and_format_timestamps --locked`：1 passed。
- `cargo fmt --check`：通过。
- `cargo clippy --all-targets --all-features --locked -- -D warnings`：通过。
### Phase 2：WebSocket transport 边界
目标：在不改变 WebSocket 行为的前提下，按职责拆分 `src/codex/gateway/transport/websocket.rs`。
建议顺序：
1. 先拆错误类型和协议编解码 helper。
2. 再拆连接建立、握手请求构造、消息发送。
3. 最后拆连接池生命周期和 stream 转发。
验收条件：
- WebSocket pool 的 55 分钟 max age、25 秒 ping、串行复用语义保持不变。
- WS -> SSE 输出格式保持不变。
- rate limit internal event、turn state、set-cookie、rate-limit header 捕获保持不变。
- HTTP SSE fallback 行为保持不变。
- 不引入 IP proxy/VPN 相关能力。
Phase 2 状态：已完成。
已完成内容：
- 新增 `websocket/pool.rs`，集中 WebSocket pool key、pool config、55 分钟 max age、keepalive、GC、shutdown、account cap、busy key bypass 等池化生命周期逻辑。
- 新增 `websocket/codec.rs`，集中 upstream request body 转换、WS frame 文本提取、WS event 分类、WS error frame 分类、retry-after/header 提取。
- `websocket.rs` 保留对外入口、连接建立、active connection 归还/丢弃、WS -> SSE stream 编排。
- 没有引入 IP proxy/VPN 能力。
验证结果：
- `cargo test --test codex_gateway websocket --locked`：20 passed。
- `cargo fmt --check`：通过。
- `cargo clippy --all-targets --all-features --locked -- -D warnings`：通过。
### Phase 3：Dispatch 编排边界
目标：让 `src/codex/serving/dispatch/mod.rs` 更像请求编排入口，把可独立变化的细节迁到已有子模块或新的内部模块。
建议顺序：
1. 先确认已有 `affinity/fallback/refresh/routing/stream/usage` 子模块职责。
2. 优先迁出 stream audit 和 upstream response 转换逻辑。
3. 再评估 fallback/retry 决策、quota/rate-limit 应用是否需要继续迁出。
验收条件：
- 外部 API、SSE 输出、fallback 语义、quota cooldown、session affinity 保持不变。
- `dispatch/mod.rs` 不再同时承载所有 stream 细节。
- 每个迁移切片都有目标测试覆盖。
Phase 3 状态：已完成。
已完成内容：
- 新增 `dispatch/audit.rs`，集中 HTTP SSE / WebSocket stream 的审计、usage 记录、上游失败识别和账户 slot 释放。
- 扩展 `dispatch/fallback.rs`，将 fallback retry 的副作用迁入同一模块，包括 quota cooldown、Cloudflare cooldown、ban/quota 状态更新和备用账户获取。
- 新增 `dispatch/limits.rs`，集中 rate-limit header 解析、quota 缓存同步、window 同步和被动 cooldown 应用。
- `dispatch/mod.rs` 保留服务入口、请求编排和上游 client request 构造。上游 request 构造暂不硬拆，避免为了降行数引入复杂 lifetime/owned context 设计。
验证结果：
- `cargo test --test codex_serving --locked`：57 passed。
- `cargo fmt --check`：通过。
- `cargo clippy --all-targets --all-features --locked -- -D warnings`：通过。
### Phase 4：日志、panic 点和测试组织收尾
目标：处理不会改变业务行为但影响长期维护质量的问题。
检查项：
1. 日志消息使用中文，保留 WebSocket/SSE/OAuth/quota/token 等专业术语。
2. 生产代码中避免无语义的 `unwrap()` / `expect()`。
3. 测试文件按场景组织，不为了拆小而过度拆分。
4. 继续避免 token、cookie、refresh token 等敏感信息进入日志。
Phase 4 状态：已完成。
已完成内容：
- 将 fingerprint 更新检查器日志中的 `UpdateChecker` 英文前缀改为中文表达，保留 fingerprint 专业术语。
- 清理 `usage_snapshots.rs` 生产路径中的时间桶对齐 `unwrap()`，改为无 panic 的 helper；测试代码中的 unwrap 保留。
- 将 `tests/codex_gateway/websocket.rs` 中 WebSocket pool 场景迁到 `tests/codex_gateway/websocket/pool.rs`，保留共享 helper 和非 pool 场景测试在父文件。
- 不继续拆更细的测试文件，避免把一个协议场景拆成过多小模块。
验证结果：
- `cargo test usage_snapshots --locked`：4 passed。
- `cargo test --test codex_gateway websocket --locked`：20 passed。
- `cargo fmt --check`：通过。
- `cargo clippy --all-targets --all-features --locked -- -D warnings`：通过。
## 改造后边界
本轮改造没有追求小文件数量，而是按变更理由拆分。当前关键文件规模如下：
| 文件 | 行数 | 当前职责 |
|---|---:|---|
| `src/codex/accounts/repository.rs` | 705 | 账号基础读写、row mapper、pool read model |
| `src/codex/accounts/repository/token.rs` | 203 | token/claims 写入和 refresh_token 保留策略 |
| `src/codex/accounts/repository/quota.rs` | 125 | quota JSON、quota cooldown、Cloudflare cooldown |
| `src/codex/accounts/repository/usage.rs` | 403 | usage 累计、window 同步、usage 列表和汇总 |
| `src/codex/accounts/repository/lease.rs` | 51 | refresh lease 抢占和释放 |
| `src/codex/gateway/transport/websocket.rs` | 506 | WebSocket 对外入口、连接建立、WS -> SSE 编排 |
| `src/codex/gateway/transport/websocket/pool.rs` | 492 | WebSocket pool 生命周期、keepalive、GC、shutdown |
| `src/codex/gateway/transport/websocket/codec.rs` | 206 | upstream body 转换、frame/event/error/header 解析 |
| `src/codex/serving/dispatch/mod.rs` | 1013 | 服务入口、请求编排、上游 request 构造 |
| `src/codex/serving/dispatch/audit.rs` | 329 | HTTP SSE/WebSocket stream 审计和账户 slot 释放 |
| `src/codex/serving/dispatch/fallback.rs` | 255 | fallback retry 分类和副作用应用 |
| `src/codex/serving/dispatch/limits.rs` | 80 | rate-limit header、quota 缓存、window/cooldown 同步 |
| `tests/codex_gateway/websocket.rs` | 880 | WebSocket 非 pool 协议场景和共享测试 helper |
| `tests/codex_gateway/websocket/pool.rs` | 559 | WebSocket pool 复用、淘汰、cap、keepalive、shutdown 场景 |
`dispatch/mod.rs` 仍然偏大，但当前保留的是请求主流程和上游 request 构造。继续硬拆需要引入额外 context 类型或更复杂的生命周期设计，收益暂时低于风险。
## 后续观察项
### 性能优化
目标：只优化有证据的热点。
1. 跑压测或 profile，确认 clone/string 分配是否是瓶颈。
2. 如果 `Account` clone 是热点，再设计 immutable account snapshot 或 `Arc<Account>`。
3. 不做无证据的全局 `Cow`/引用替换。
### 后续日志治理
本轮只处理审计中确认的日志消息一致性问题，没有重构日志系统。后续如果继续做日志最佳实践，应单独落地 HTTP/span/request_id 字段治理，避免和结构重构混在一起。
src/codex/accounts/repository.rs
+221 -619
View File
@@ -9,6 +9,13 @@ use crate::{
utils::pagination::{decode_cursor, encode_cursor, Page},
};
mod lease;
mod quota;
mod token;
mod usage;
pub use usage::AccountUsageRepository;
#[derive(Debug, Error)]
pub enum AccountRepositoryError {
#[error("database error: {0}")]
@@ -152,6 +159,163 @@ pub struct AccountRepository {
secret_box: SecretBox,
}
const INSERT_ACCOUNT_SQL: &str = r"
insert into accounts (
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
status,
added_at,
updated_at
) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)";
const SELECT_STORED_ACCOUNT_BY_ID_SQL: &str = r"
select
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
status,
added_at,
updated_at
from accounts
where id = ?";
const SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL: &str = r"
select
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
status,
added_at,
updated_at
from accounts
where account_id = ?
and ((user_id is null and ? is null) or user_id = ?)
order by added_at asc
limit 1";
const LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL: &str = r"
select
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
status,
added_at,
updated_at
from accounts
where added_at < ?
or (added_at = ? and id < ?)
order by added_at desc, id desc
limit ?";
const LIST_STORED_ACCOUNTS_SQL: &str = r"
select
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
status,
added_at,
updated_at
from accounts
order by added_at desc, id desc
limit ?";
const LIST_ALL_STORED_ACCOUNTS_SQL: &str = r"
select
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
status,
added_at,
updated_at
from accounts
order by added_at desc, id desc";
const LIST_METADATA_AFTER_CURSOR_SQL: &str = r"
select
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_expires_at,
status,
added_at,
updated_at
from accounts
where added_at < ?
or (added_at = ? and id < ?)
order by added_at desc, id desc
limit ?";
const LIST_METADATA_SQL: &str = r"
select
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_expires_at,
status,
added_at,
updated_at
from accounts
order by added_at desc, id desc
limit ?";
const LIST_ALL_METADATA_SQL: &str = r"
select
id,
email,
account_id,
user_id,
label,
plan_type,
access_token_expires_at,
status,
added_at,
updated_at
from accounts
order by added_at desc, id desc";
const LIST_POOL_ACCOUNTS_SQL: &str = r"
select
accounts.id,
@@ -183,117 +347,6 @@ from accounts
left join account_usage on account_usage.account_id = accounts.id
order by added_at desc, accounts.id desc";
const RECORD_USAGE_SQL: &str = r"
insert into account_usage (
account_id,
request_count,
empty_response_count,
input_tokens,
output_tokens,
cached_tokens,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
last_used_at
) values (?, 1, ?, ?, ?, ?, 1, ?, ?, ?, ?)
on conflict(account_id) do update set
request_count = request_count + 1,
empty_response_count = empty_response_count + excluded.empty_response_count,
input_tokens = input_tokens + excluded.input_tokens,
output_tokens = output_tokens + excluded.output_tokens,
cached_tokens = cached_tokens + excluded.cached_tokens,
window_request_count = window_request_count + 1,
window_input_tokens = window_input_tokens + excluded.window_input_tokens,
window_output_tokens = window_output_tokens + excluded.window_output_tokens,
window_cached_tokens = window_cached_tokens + excluded.window_cached_tokens,
window_started_at = case
when account_usage.window_started_at is null
and (account_usage.window_reset_at is not null or account_usage.limit_window_seconds is not null)
then excluded.last_used_at
else account_usage.window_started_at
end,
last_used_at = excluded.last_used_at";
const GET_USAGE_SQL: &str = r"
select
account_id,
request_count,
empty_response_count,
input_tokens,
output_tokens,
cached_tokens,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_started_at,
window_reset_at,
limit_window_seconds,
last_used_at
from account_usage
where account_id = ?";
const SYNC_RATE_LIMIT_WINDOW_RESET_SQL: &str = r"
insert into account_usage (
account_id,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_started_at,
window_reset_at,
limit_window_seconds
) values (?, 0, 0, 0, 0, ?, ?, ?)
on conflict(account_id) do update set
window_request_count = 0,
window_input_tokens = 0,
window_output_tokens = 0,
window_cached_tokens = 0,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)";
const SYNC_RATE_LIMIT_WINDOW_SQL: &str = r"
insert into account_usage (
account_id,
window_reset_at,
limit_window_seconds
) values (?, ?, ?)
on conflict(account_id) do update set
window_reset_at = excluded.window_reset_at,
limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)";
const RESET_ACCOUNT_USAGE_SQL: &str = r"
insert into account_usage (
account_id,
request_count,
empty_response_count,
input_tokens,
output_tokens,
cached_tokens,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_started_at,
window_reset_at,
last_used_at
) values (?, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null, null)
on conflict(account_id) do update set
request_count = 0,
empty_response_count = 0,
input_tokens = 0,
output_tokens = 0,
cached_tokens = 0,
window_request_count = 0,
window_input_tokens = 0,
window_output_tokens = 0,
window_cached_tokens = 0,
window_started_at = null,
window_reset_at = null,
last_used_at = null";
impl AccountRepository {
pub fn new(pool: SqlitePool, secret_box: SecretBox) -> Self {
Self { pool, secret_box }
@@ -307,33 +360,33 @@ impl AccountRepository {
.as_ref()
.map(|token| self.secret_box.encrypt(token))
.transpose()?;
sqlx::query(
"insert into accounts (id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(account.id)
.bind(account.email)
.bind(account.account_id)
.bind(account.user_id)
.bind(account.label)
.bind(account.plan_type)
.bind(access_token_cipher)
.bind(refresh_token_cipher)
.bind(account.access_token_expires_at.map(|value| value.to_rfc3339()))
.bind(status_to_db(account.status))
.bind(&now)
.bind(now)
.execute(&self.pool)
.await?;
sqlx::query(INSERT_ACCOUNT_SQL)
.bind(account.id)
.bind(account.email)
.bind(account.account_id)
.bind(account.user_id)
.bind(account.label)
.bind(account.plan_type)
.bind(access_token_cipher)
.bind(refresh_token_cipher)
.bind(
account
.access_token_expires_at
.map(|value| value.to_rfc3339()),
)
.bind(status_to_db(account.status))
.bind(&now)
.bind(now)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn get(&self, id: &str) -> AccountRepositoryResult<Option<StoredAccount>> {
let row = sqlx::query(
"select id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at from accounts where id = ?",
)
.bind(id)
.fetch_optional(&self.pool)
.await?;
let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_ID_SQL)
.bind(id)
.fetch_optional(&self.pool)
.await?;
row.map(|row| self.account_from_row(&row)).transpose()
}
@@ -350,14 +403,12 @@ impl AccountRepository {
account_id: &str,
user_id: Option<&str>,
) -> AccountRepositoryResult<Option<StoredAccount>> {
let row = sqlx::query(
"select id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at from accounts where account_id = ? and ((user_id is null and ? is null) or user_id = ?) order by added_at asc limit 1",
)
.bind(account_id)
.bind(user_id)
.bind(user_id)
.fetch_optional(&self.pool)
.await?;
let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL)
.bind(account_id)
.bind(user_id)
.bind(user_id)
.fetch_optional(&self.pool)
.await?;
row.map(|row| self.account_from_row(&row)).transpose()
}
@@ -370,22 +421,18 @@ impl AccountRepository {
let rows = if let Some(cursor) = cursor {
let (created_at, id) =
decode_cursor(&cursor).ok_or(AccountRepositoryError::InvalidCursor)?;
sqlx::query(
"select id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at from accounts where added_at < ? or (added_at = ? and id < ?) order by added_at desc, id desc limit ?",
)
.bind(&created_at)
.bind(created_at)
.bind(id)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
sqlx::query(LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL)
.bind(&created_at)
.bind(created_at)
.bind(id)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
} else {
sqlx::query(
"select id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at from accounts order by added_at desc, id desc limit ?",
)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
sqlx::query(LIST_STORED_ACCOUNTS_SQL)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
};
let has_next = rows.len() > limit as usize;
@@ -405,11 +452,9 @@ impl AccountRepository {
}
pub async fn list_all(&self) -> AccountRepositoryResult<Vec<StoredAccount>> {
let rows = sqlx::query(
"select id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at from accounts order by added_at desc, id desc",
)
.fetch_all(&self.pool)
.await?;
let rows = sqlx::query(LIST_ALL_STORED_ACCOUNTS_SQL)
.fetch_all(&self.pool)
.await?;
let mut accounts = Vec::with_capacity(rows.len());
for row in rows {
accounts.push(self.account_from_row(&row)?);
@@ -426,22 +471,18 @@ impl AccountRepository {
let rows = if let Some(cursor) = cursor {
let (created_at, id) =
decode_cursor(&cursor).ok_or(AccountRepositoryError::InvalidCursor)?;
sqlx::query(
"select id, email, account_id, user_id, label, plan_type, access_token_expires_at, status, added_at, updated_at from accounts where added_at < ? or (added_at = ? and id < ?) order by added_at desc, id desc limit ?",
)
.bind(&created_at)
.bind(created_at)
.bind(id)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
sqlx::query(LIST_METADATA_AFTER_CURSOR_SQL)
.bind(&created_at)
.bind(created_at)
.bind(id)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
} else {
sqlx::query(
"select id, email, account_id, user_id, label, plan_type, access_token_expires_at, status, added_at, updated_at from accounts order by added_at desc, id desc limit ?",
)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
sqlx::query(LIST_METADATA_SQL)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
};
let has_next = rows.len() > limit as usize;
@@ -461,11 +502,9 @@ impl AccountRepository {
}
pub async fn list_all_metadata(&self) -> AccountRepositoryResult<Vec<StoredAccountMetadata>> {
let rows = sqlx::query(
"select id, email, account_id, user_id, label, plan_type, access_token_expires_at, status, added_at, updated_at from accounts order by added_at desc, id desc",
)
.fetch_all(&self.pool)
.await?;
let rows = sqlx::query(LIST_ALL_METADATA_SQL)
.fetch_all(&self.pool)
.await?;
let mut accounts = Vec::with_capacity(rows.len());
for row in rows {
accounts.push(metadata_from_row(&row)?);
@@ -533,19 +572,6 @@ impl AccountRepository {
Ok(accounts)
}
pub async fn list_quota_snapshots(&self) -> AccountRepositoryResult<Vec<AccountQuotaSnapshot>> {
let rows = sqlx::query(
"select id, email, quota_json, quota_fetched_at from accounts where quota_json is not null and trim(quota_json) <> '' order by coalesce(quota_fetched_at, '') desc, id desc",
)
.fetch_all(&self.pool)
.await?;
let mut snapshots = Vec::with_capacity(rows.len());
for row in rows {
snapshots.push(quota_snapshot_from_row(&row)?);
}
Ok(snapshots)
}
pub async fn set_status(
&self,
id: &str,
@@ -589,266 +615,6 @@ impl AccountRepository {
Ok(result.rows_affected())
}
pub async fn update_tokens(
&self,
id: &str,
update: TokenUpdate,
) -> AccountRepositoryResult<bool> {
let now = Utc::now().to_rfc3339();
let access_token_cipher = self.secret_box.encrypt(&update.access_token)?;
let expires_at = update
.access_token_expires_at
.map(|value| value.to_rfc3339());
let result = if let Some(refresh_token) = update.refresh_token {
let refresh_token_cipher = self.secret_box.encrypt(&refresh_token)?;
sqlx::query(
"update accounts set access_token_cipher = ?, refresh_token_cipher = ?, access_token_expires_at = ?, status = case when status in ('disabled', 'banned') then status else 'active' end, updated_at = ? where id = ?",
)
.bind(access_token_cipher)
.bind(refresh_token_cipher)
.bind(expires_at)
.bind(now)
.bind(id)
.execute(&self.pool)
.await?
} else {
// 刷新接口可能不返回新的 refresh_token；此时必须保留旧 RT，避免账号失去后续刷新能力。
sqlx::query(
"update accounts set access_token_cipher = ?, access_token_expires_at = ?, status = case when status in ('disabled', 'banned') then status else 'active' end, updated_at = ? where id = ?",
)
.bind(access_token_cipher)
.bind(expires_at)
.bind(now)
.bind(id)
.execute(&self.pool)
.await?
};
Ok(result.rows_affected() > 0)
}
pub async fn update_from_claims(
&self,
id: &str,
update: AccountClaimsUpdate,
) -> AccountRepositoryResult<bool> {
let now = Utc::now().to_rfc3339();
let access_token_cipher = self.secret_box.encrypt(&update.access_token)?;
let expires_at = update
.access_token_expires_at
.map(|value| value.to_rfc3339());
let result = if let Some(refresh_token) = update.refresh_token {
let refresh_token_cipher = self.secret_box.encrypt(&refresh_token)?;
sqlx::query(
"update accounts set email = ?, account_id = ?, user_id = ?, plan_type = ?, access_token_cipher = ?, refresh_token_cipher = ?, access_token_expires_at = ?, status = ?, updated_at = ? where id = ?",
)
.bind(update.email)
.bind(update.account_id)
.bind(update.user_id)
.bind(update.plan_type)
.bind(access_token_cipher)
.bind(refresh_token_cipher)
.bind(expires_at)
.bind(status_to_db(update.status))
.bind(now)
.bind(id)
.execute(&self.pool)
.await?
} else {
// OpenAI 刷新/导入未给新 RT 时保留原值，避免把可继续刷新的账号写坏。
sqlx::query(
"update accounts set email = ?, account_id = ?, user_id = ?, plan_type = ?, access_token_cipher = ?, access_token_expires_at = ?, status = ?, updated_at = ? where id = ?",
)
.bind(update.email)
.bind(update.account_id)
.bind(update.user_id)
.bind(update.plan_type)
.bind(access_token_cipher)
.bind(expires_at)
.bind(status_to_db(update.status))
.bind(now)
.bind(id)
.execute(&self.pool)
.await?
};
Ok(result.rows_affected() > 0)
}
pub async fn record_usage(
&self,
account_id: &str,
usage: UsageDelta,
) -> AccountRepositoryResult<()> {
let now = Utc::now().to_rfc3339();
sqlx::query(RECORD_USAGE_SQL)
.bind(account_id)
.bind(usage.empty_response_count)
.bind(usage.input_tokens)
.bind(usage.output_tokens)
.bind(usage.cached_tokens)
.bind(usage.input_tokens)
.bind(usage.output_tokens)
.bind(usage.cached_tokens)
.bind(now)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn get_usage(
&self,
account_id: &str,
) -> AccountRepositoryResult<Option<AccountUsageRecord>> {
let row = sqlx::query(GET_USAGE_SQL)
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
row.map(|row| usage_from_row(&row)).transpose()
}
pub async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> AccountRepositoryResult<()> {
let row = sqlx::query(
"select window_reset_at, limit_window_seconds from account_usage where account_id = ?",
)
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
let existing_reset_at = row
.as_ref()
.map(|row| parse_optional_rfc3339(row.get::<Option<String>, _>("window_reset_at")))
.transpose()?
.flatten();
let existing_limit_window_seconds = row
.as_ref()
.and_then(|row| optional_positive_i64_to_u64(row.get("limit_window_seconds")));
let limit_window_seconds_db = limit_window_seconds.map(u64_to_i64_saturating);
let reset_at_db = reset_at.to_rfc3339();
if should_reset_usage_window(
existing_reset_at,
existing_limit_window_seconds,
reset_at,
limit_window_seconds,
) {
sqlx::query(SYNC_RATE_LIMIT_WINDOW_RESET_SQL)
.bind(account_id)
.bind(Utc::now().to_rfc3339())
.bind(reset_at_db)
.bind(limit_window_seconds_db)
.execute(&self.pool)
.await?;
} else {
sqlx::query(SYNC_RATE_LIMIT_WINDOW_SQL)
.bind(account_id)
.bind(reset_at_db)
.bind(limit_window_seconds_db)
.execute(&self.pool)
.await?;
}
Ok(())
}
pub async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> AccountRepositoryResult<bool> {
let now = Utc::now().to_rfc3339();
let result = sqlx::query(
"update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?",
)
.bind(quota_json)
.bind(&now)
.bind(now)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
pub async fn get_quota_json(
&self,
account_id: &str,
) -> AccountRepositoryResult<Option<String>> {
let row = sqlx::query("select quota_json from accounts where id = ?")
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
Ok(row.and_then(|row| row.get("quota_json")))
}
pub async fn set_quota_cooldown_until(
&self,
id: &str,
cooldown_until: DateTime<Utc>,
) -> AccountRepositoryResult<bool> {
let cooldown_until = cooldown_until.to_rfc3339();
let result = sqlx::query(
"update accounts set quota_limit_reached = 1, quota_cooldown_until = case when quota_cooldown_until is not null and quota_cooldown_until > ? then quota_cooldown_until else ? end, updated_at = ? where id = ?",
)
.bind(&cooldown_until)
.bind(cooldown_until)
.bind(Utc::now().to_rfc3339())
.bind(id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
pub async fn set_cloudflare_cooldown_until(
&self,
id: &str,
cooldown_until: DateTime<Utc>,
) -> AccountRepositoryResult<bool> {
let result = sqlx::query(
"update accounts set cloudflare_cooldown_until = ?, updated_at = ? where id = ?",
)
.bind(cooldown_until.to_rfc3339())
.bind(Utc::now().to_rfc3339())
.bind(id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
pub async fn try_acquire_refresh_lease(
&self,
account_id: &str,
owner: &str,
lease_until: DateTime<Utc>,
) -> AccountRepositoryResult<bool> {
let now = Utc::now().to_rfc3339();
let result = sqlx::query(
"insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?) on conflict(account_id) do update set owner = excluded.owner, expires_at = excluded.expires_at, updated_at = excluded.updated_at where account_refresh_leases.expires_at <= ? or account_refresh_leases.owner = excluded.owner",
)
.bind(account_id)
.bind(owner)
.bind(lease_until.to_rfc3339())
.bind(&now)
.bind(now)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
pub async fn release_refresh_lease(
&self,
account_id: &str,
owner: &str,
) -> AccountRepositoryResult<bool> {
let result =
sqlx::query("delete from account_refresh_leases where account_id = ? and owner = ?")
.bind(account_id)
.bind(owner)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
fn account_from_row(
&self,
row: &sqlx::sqlite::SqliteRow,
@@ -919,170 +685,6 @@ fn optional_positive_i64_to_u64(value: Option<i64>) -> Option<u64> {
.filter(|value| *value > 0)
}
fn u64_to_i64_saturating(value: u64) -> i64 {
value.min(i64::MAX as u64) as i64
}
fn should_reset_usage_window(
existing_reset_at: Option<DateTime<Utc>>,
existing_limit_window_seconds: Option<u64>,
new_reset_at: DateTime<Utc>,
new_limit_window_seconds: Option<u64>,
) -> bool {
let Some(existing_reset_at) = existing_reset_at else {
return false;
};
if existing_reset_at == new_reset_at {
return false;
}
let drift = existing_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = new_limit_window_seconds
.or(existing_limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> AccountRepositoryResult<AccountUsageRecord> {
Ok(AccountUsageRecord {
account_id: row.get("account_id"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
window_request_count: row.get("window_request_count"),
window_input_tokens: row.get("window_input_tokens"),
window_output_tokens: row.get("window_output_tokens"),
window_cached_tokens: row.get("window_cached_tokens"),
window_started_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("window_started_at"),
)?,
window_reset_at: parse_optional_rfc3339(row.get::<Option<String>, _>("window_reset_at"))?,
limit_window_seconds: optional_positive_i64_to_u64(row.get("limit_window_seconds")),
last_used_at: parse_optional_rfc3339(row.get::<Option<String>, _>("last_used_at"))?,
})
}
fn quota_snapshot_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> AccountRepositoryResult<AccountQuotaSnapshot> {
Ok(AccountQuotaSnapshot {
account_id: row.get("id"),
email: row.get("email"),
quota_json: row.get("quota_json"),
quota_fetched_at: parse_optional_rfc3339(row.get::<Option<String>, _>("quota_fetched_at"))?,
})
}
#[derive(Clone)]
pub struct AccountUsageRepository {
pool: SqlitePool,
}
impl AccountUsageRepository {
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
pub async fn list(
&self,
cursor: Option<String>,
limit: u32,
) -> AccountRepositoryResult<Page<AccountUsageListRecord>> {
let fetch_limit = i64::from(limit) + 1;
let rows = if let Some(cursor) = cursor {
let (last_used_at, account_id) =
decode_cursor(&cursor).ok_or(AccountRepositoryError::InvalidCursor)?;
sqlx::query(
"select account_usage.account_id, accounts.email, accounts.label, accounts.plan_type, account_usage.request_count, account_usage.empty_response_count, account_usage.input_tokens, account_usage.output_tokens, account_usage.cached_tokens, account_usage.last_used_at from account_usage left join accounts on accounts.id = account_usage.account_id where coalesce(account_usage.last_used_at, '') < ? or (coalesce(account_usage.last_used_at, '') = ? and account_usage.account_id < ?) order by coalesce(account_usage.last_used_at, '') desc, account_usage.account_id desc limit ?",
)
.bind(&last_used_at)
.bind(last_used_at)
.bind(account_id)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
} else {
sqlx::query(
"select account_usage.account_id, accounts.email, accounts.label, accounts.plan_type, account_usage.request_count, account_usage.empty_response_count, account_usage.input_tokens, account_usage.output_tokens, account_usage.cached_tokens, account_usage.last_used_at from account_usage left join accounts on accounts.id = account_usage.account_id order by coalesce(account_usage.last_used_at, '') desc, account_usage.account_id desc limit ?",
)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
};
let has_next = rows.len() > limit as usize;
let take_count = rows.len().min(limit as usize);
let mut items = Vec::with_capacity(take_count);
for row in rows.into_iter().take(take_count) {
items.push(usage_list_from_row(&row)?);
}
let next_cursor = if has_next {
items.last().map(|usage| {
encode_cursor(
&usage
.last_used_at
.map(|value| value.to_rfc3339())
.unwrap_or_default(),
&usage.account_id,
)
})
} else {
None
};
Ok(Page { items, next_cursor })
}
pub async fn summary(&self) -> AccountRepositoryResult<AccountUsageSummary> {
let row = sqlx::query(
"select count(*) as account_count, coalesce(sum(request_count), 0) as request_count, coalesce(sum(empty_response_count), 0) as empty_response_count, coalesce(sum(input_tokens), 0) as input_tokens, coalesce(sum(output_tokens), 0) as output_tokens, coalesce(sum(cached_tokens), 0) as cached_tokens from account_usage",
)
.fetch_one(&self.pool)
.await?;
Ok(AccountUsageSummary {
account_count: row.get("account_count"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
})
}
pub async fn reset_account(&self, account_id: &str) -> AccountRepositoryResult<()> {
sqlx::query(RESET_ACCOUNT_USAGE_SQL)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(())
}
}
fn usage_list_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> AccountRepositoryResult<AccountUsageListRecord> {
Ok(AccountUsageListRecord {
account_id: row.get("account_id"),
email: row.get("email"),
label: row.get("label"),
plan_type: row.get("plan_type"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
last_used_at: parse_optional_rfc3339(row.get::<Option<String>, _>("last_used_at"))?,
})
}
fn metadata_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> AccountRepositoryResult<StoredAccountMetadata> {
src/codex/accounts/repository/lease.rs
+51
View File
@@ -0,0 +1,51 @@
use chrono::{DateTime, Utc};
use super::{AccountRepository, AccountRepositoryResult};
const TRY_ACQUIRE_REFRESH_LEASE_SQL: &str = r"
insert into account_refresh_leases (
account_id,
owner,
expires_at,
updated_at
) values (?, ?, ?, ?)
on conflict(account_id) do update set
owner = excluded.owner,
expires_at = excluded.expires_at,
updated_at = excluded.updated_at
where account_refresh_leases.expires_at <= ?
or account_refresh_leases.owner = excluded.owner";
impl AccountRepository {
pub async fn try_acquire_refresh_lease(
&self,
account_id: &str,
owner: &str,
lease_until: DateTime<Utc>,
) -> AccountRepositoryResult<bool> {
let now = Utc::now().to_rfc3339();
let result = sqlx::query(TRY_ACQUIRE_REFRESH_LEASE_SQL)
.bind(account_id)
.bind(owner)
.bind(lease_until.to_rfc3339())
.bind(&now)
.bind(now)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
pub async fn release_refresh_lease(
&self,
account_id: &str,
owner: &str,
) -> AccountRepositoryResult<bool> {
let result =
sqlx::query("delete from account_refresh_leases where account_id = ? and owner = ?")
.bind(account_id)
.bind(owner)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
}
src/codex/accounts/repository/quota.rs
+125
View File
@@ -0,0 +1,125 @@
use chrono::{DateTime, Utc};
use sqlx::Row;
use super::{
parse_optional_rfc3339, AccountQuotaSnapshot, AccountRepository, AccountRepositoryResult,
};
const LIST_QUOTA_SNAPSHOTS_SQL: &str = r"
select
id,
email,
quota_json,
quota_fetched_at
from accounts
where quota_json is not null
and trim(quota_json) <> ''
order by coalesce(quota_fetched_at, '') desc, id desc";
const UPDATE_QUOTA_JSON_SQL: &str = r"
update accounts
set
quota_json = ?,
quota_fetched_at = ?,
updated_at = ?
where id = ?";
const SET_QUOTA_COOLDOWN_UNTIL_SQL: &str = r"
update accounts
set
quota_limit_reached = 1,
quota_cooldown_until = case
when quota_cooldown_until is not null and quota_cooldown_until > ?
then quota_cooldown_until
else ?
end,
updated_at = ?
where id = ?";
const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r"
update accounts
set
cloudflare_cooldown_until = ?,
updated_at = ?
where id = ?";
impl AccountRepository {
pub async fn list_quota_snapshots(&self) -> AccountRepositoryResult<Vec<AccountQuotaSnapshot>> {
let rows = sqlx::query(LIST_QUOTA_SNAPSHOTS_SQL)
.fetch_all(&self.pool)
.await?;
let mut snapshots = Vec::with_capacity(rows.len());
for row in rows {
snapshots.push(quota_snapshot_from_row(&row)?);
}
Ok(snapshots)
}
pub async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> AccountRepositoryResult<bool> {
let now = Utc::now().to_rfc3339();
let result = sqlx::query(UPDATE_QUOTA_JSON_SQL)
.bind(quota_json)
.bind(&now)
.bind(now)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
pub async fn get_quota_json(
&self,
account_id: &str,
) -> AccountRepositoryResult<Option<String>> {
let row = sqlx::query("select quota_json from accounts where id = ?")
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
Ok(row.and_then(|row| row.get("quota_json")))
}
pub async fn set_quota_cooldown_until(
&self,
id: &str,
cooldown_until: DateTime<Utc>,
) -> AccountRepositoryResult<bool> {
let cooldown_until = cooldown_until.to_rfc3339();
let result = sqlx::query(SET_QUOTA_COOLDOWN_UNTIL_SQL)
.bind(&cooldown_until)
.bind(cooldown_until)
.bind(Utc::now().to_rfc3339())
.bind(id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
pub async fn set_cloudflare_cooldown_until(
&self,
id: &str,
cooldown_until: DateTime<Utc>,
) -> AccountRepositoryResult<bool> {
let result = sqlx::query(SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL)
.bind(cooldown_until.to_rfc3339())
.bind(Utc::now().to_rfc3339())
.bind(id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
}
fn quota_snapshot_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> AccountRepositoryResult<AccountQuotaSnapshot> {
Ok(AccountQuotaSnapshot {
account_id: row.get("id"),
email: row.get("email"),
quota_json: row.get("quota_json"),
quota_fetched_at: parse_optional_rfc3339(row.get::<Option<String>, _>("quota_fetched_at"))?,
})
}
src/codex/accounts/repository/token.rs
+203
View File
@@ -0,0 +1,203 @@
use chrono::{DateTime, Utc};
use secrecy::SecretString;
use super::{
status_to_db, AccountClaimsUpdate, AccountRepository, AccountRepositoryResult, TokenUpdate,
};
const UPDATE_TOKENS_WITH_REFRESH_SQL: &str = r"
update accounts
set
access_token_cipher = ?,
refresh_token_cipher = ?,
access_token_expires_at = ?,
status = case
when status in ('disabled', 'banned') then status
else 'active'
end,
updated_at = ?
where id = ?";
const UPDATE_TOKENS_PRESERVING_REFRESH_SQL: &str = r"
update accounts
set
access_token_cipher = ?,
access_token_expires_at = ?,
status = case
when status in ('disabled', 'banned') then status
else 'active'
end,
updated_at = ?
where id = ?";
const UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL: &str = r"
update accounts
set
email = ?,
account_id = ?,
user_id = ?,
plan_type = ?,
access_token_cipher = ?,
refresh_token_cipher = ?,
access_token_expires_at = ?,
status = ?,
updated_at = ?
where id = ?";
const UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL: &str = r"
update accounts
set
email = ?,
account_id = ?,
user_id = ?,
plan_type = ?,
access_token_cipher = ?,
access_token_expires_at = ?,
status = ?,
updated_at = ?
where id = ?";
#[derive(Debug)]
struct TokenWrite {
access_token_cipher: String,
refresh_token_cipher: Option<String>,
access_token_expires_at: Option<String>,
updated_at: String,
}
impl AccountRepository {
fn prepare_token_write(
&self,
access_token: &SecretString,
refresh_token: Option<&SecretString>,
access_token_expires_at: Option<DateTime<Utc>>,
) -> AccountRepositoryResult<TokenWrite> {
Ok(TokenWrite {
access_token_cipher: self.secret_box.encrypt(access_token)?,
refresh_token_cipher: refresh_token
.map(|token| self.secret_box.encrypt(token))
.transpose()?,
access_token_expires_at: access_token_expires_at.map(|value| value.to_rfc3339()),
updated_at: Utc::now().to_rfc3339(),
})
}
pub async fn update_tokens(
&self,
id: &str,
update: TokenUpdate,
) -> AccountRepositoryResult<bool> {
let TokenWrite {
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
updated_at,
} = self.prepare_token_write(
&update.access_token,
update.refresh_token.as_ref(),
update.access_token_expires_at,
)?;
let result = if let Some(refresh_token_cipher) = refresh_token_cipher {
sqlx::query(UPDATE_TOKENS_WITH_REFRESH_SQL)
.bind(access_token_cipher)
.bind(refresh_token_cipher)
.bind(access_token_expires_at)
.bind(updated_at)
.bind(id)
.execute(&self.pool)
.await?
} else {
// 刷新接口可能不返回新的 refresh_token；此时必须保留旧 RT，避免账号失去后续刷新能力。
sqlx::query(UPDATE_TOKENS_PRESERVING_REFRESH_SQL)
.bind(access_token_cipher)
.bind(access_token_expires_at)
.bind(updated_at)
.bind(id)
.execute(&self.pool)
.await?
};
Ok(result.rows_affected() > 0)
}
pub async fn update_from_claims(
&self,
id: &str,
update: AccountClaimsUpdate,
) -> AccountRepositoryResult<bool> {
let TokenWrite {
access_token_cipher,
refresh_token_cipher,
access_token_expires_at,
updated_at,
} = self.prepare_token_write(
&update.access_token,
update.refresh_token.as_ref(),
update.access_token_expires_at,
)?;
let status = status_to_db(update.status);
let result = if let Some(refresh_token_cipher) = refresh_token_cipher {
sqlx::query(UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL)
.bind(update.email)
.bind(update.account_id)
.bind(update.user_id)
.bind(update.plan_type)
.bind(access_token_cipher)
.bind(refresh_token_cipher)
.bind(access_token_expires_at)
.bind(status)
.bind(updated_at)
.bind(id)
.execute(&self.pool)
.await?
} else {
// OpenAI 刷新/导入未给新 RT 时保留原值，避免把可继续刷新的账号写坏。
sqlx::query(UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL)
.bind(update.email)
.bind(update.account_id)
.bind(update.user_id)
.bind(update.plan_type)
.bind(access_token_cipher)
.bind(access_token_expires_at)
.bind(status)
.bind(updated_at)
.bind(id)
.execute(&self.pool)
.await?
};
Ok(result.rows_affected() > 0)
}
}
#[cfg(test)]
mod tests {
use sqlx::SqlitePool;
use super::*;
use crate::platform::crypto::SecretBox;
#[tokio::test]
async fn prepare_token_write_should_encrypt_tokens_and_format_timestamps() {
let repository = AccountRepository::new(
SqlitePool::connect_lazy("sqlite::memory:").unwrap(),
SecretBox::new([55u8; 32]),
);
let access_token = SecretString::new("access-secret".to_string().into());
let refresh_token = SecretString::new("refresh-secret".to_string().into());
let expires_at = super::super::parse_rfc3339("2026-06-14T00:00:00Z").unwrap();
let write = repository
.prepare_token_write(&access_token, Some(&refresh_token), Some(expires_at))
.unwrap();
assert!(write.access_token_cipher.starts_with("v1:"));
assert!(!write.access_token_cipher.contains("access-secret"));
assert!(write.refresh_token_cipher.unwrap().starts_with("v1:"));
assert_eq!(
write.access_token_expires_at.as_deref(),
Some("2026-06-14T00:00:00+00:00")
);
assert!(super::super::parse_rfc3339(&write.updated_at).is_ok());
}
}
src/codex/accounts/repository/usage.rs
+403
View File
@@ -0,0 +1,403 @@
use chrono::{DateTime, Utc};
use sqlx::{Row, SqlitePool};
use crate::utils::pagination::{decode_cursor, encode_cursor, Page};
use super::{
optional_positive_i64_to_u64, parse_optional_rfc3339, AccountRepository,
AccountRepositoryError, AccountRepositoryResult, AccountUsageListRecord, AccountUsageRecord,
AccountUsageSummary, UsageDelta,
};
const RECORD_USAGE_SQL: &str = r"
insert into account_usage (
account_id,
request_count,
empty_response_count,
input_tokens,
output_tokens,
cached_tokens,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
last_used_at
) values (?, 1, ?, ?, ?, ?, 1, ?, ?, ?, ?)
on conflict(account_id) do update set
request_count = request_count + 1,
empty_response_count = empty_response_count + excluded.empty_response_count,
input_tokens = input_tokens + excluded.input_tokens,
output_tokens = output_tokens + excluded.output_tokens,
cached_tokens = cached_tokens + excluded.cached_tokens,
window_request_count = window_request_count + 1,
window_input_tokens = window_input_tokens + excluded.window_input_tokens,
window_output_tokens = window_output_tokens + excluded.window_output_tokens,
window_cached_tokens = window_cached_tokens + excluded.window_cached_tokens,
window_started_at = case
when account_usage.window_started_at is null
and (account_usage.window_reset_at is not null or account_usage.limit_window_seconds is not null)
then excluded.last_used_at
else account_usage.window_started_at
end,
last_used_at = excluded.last_used_at";
const GET_USAGE_SQL: &str = r"
select
account_id,
request_count,
empty_response_count,
input_tokens,
output_tokens,
cached_tokens,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_started_at,
window_reset_at,
limit_window_seconds,
last_used_at
from account_usage
where account_id = ?";
const SELECT_RATE_LIMIT_WINDOW_SQL: &str = r"
select
window_reset_at,
limit_window_seconds
from account_usage
where account_id = ?";
const SYNC_RATE_LIMIT_WINDOW_RESET_SQL: &str = r"
insert into account_usage (
account_id,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_started_at,
window_reset_at,
limit_window_seconds
) values (?, 0, 0, 0, 0, ?, ?, ?)
on conflict(account_id) do update set
window_request_count = 0,
window_input_tokens = 0,
window_output_tokens = 0,
window_cached_tokens = 0,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)";
const SYNC_RATE_LIMIT_WINDOW_SQL: &str = r"
insert into account_usage (
account_id,
window_reset_at,
limit_window_seconds
) values (?, ?, ?)
on conflict(account_id) do update set
window_reset_at = excluded.window_reset_at,
limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)";
const RESET_ACCOUNT_USAGE_SQL: &str = r"
insert into account_usage (
account_id,
request_count,
empty_response_count,
input_tokens,
output_tokens,
cached_tokens,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_started_at,
window_reset_at,
last_used_at
) values (?, 0, 0, 0, 0, 0, 0, 0, 0, 0, null, null, null)
on conflict(account_id) do update set
request_count = 0,
empty_response_count = 0,
input_tokens = 0,
output_tokens = 0,
cached_tokens = 0,
window_request_count = 0,
window_input_tokens = 0,
window_output_tokens = 0,
window_cached_tokens = 0,
window_started_at = null,
window_reset_at = null,
last_used_at = null";
const LIST_ACCOUNT_USAGE_AFTER_CURSOR_SQL: &str = r"
select
account_usage.account_id,
accounts.email,
accounts.label,
accounts.plan_type,
account_usage.request_count,
account_usage.empty_response_count,
account_usage.input_tokens,
account_usage.output_tokens,
account_usage.cached_tokens,
account_usage.last_used_at
from account_usage
left join accounts on accounts.id = account_usage.account_id
where coalesce(account_usage.last_used_at, '') < ?
or (
coalesce(account_usage.last_used_at, '') = ?
and account_usage.account_id < ?
)
order by coalesce(account_usage.last_used_at, '') desc, account_usage.account_id desc
limit ?";
const LIST_ACCOUNT_USAGE_SQL: &str = r"
select
account_usage.account_id,
accounts.email,
accounts.label,
accounts.plan_type,
account_usage.request_count,
account_usage.empty_response_count,
account_usage.input_tokens,
account_usage.output_tokens,
account_usage.cached_tokens,
account_usage.last_used_at
from account_usage
left join accounts on accounts.id = account_usage.account_id
order by coalesce(account_usage.last_used_at, '') desc, account_usage.account_id desc
limit ?";
const ACCOUNT_USAGE_SUMMARY_SQL: &str = r"
select
count(*) as account_count,
coalesce(sum(request_count), 0) as request_count,
coalesce(sum(empty_response_count), 0) as empty_response_count,
coalesce(sum(input_tokens), 0) as input_tokens,
coalesce(sum(output_tokens), 0) as output_tokens,
coalesce(sum(cached_tokens), 0) as cached_tokens
from account_usage";
impl AccountRepository {
pub async fn record_usage(
&self,
account_id: &str,
usage: UsageDelta,
) -> AccountRepositoryResult<()> {
let now = Utc::now().to_rfc3339();
sqlx::query(RECORD_USAGE_SQL)
.bind(account_id)
.bind(usage.empty_response_count)
.bind(usage.input_tokens)
.bind(usage.output_tokens)
.bind(usage.cached_tokens)
.bind(usage.input_tokens)
.bind(usage.output_tokens)
.bind(usage.cached_tokens)
.bind(now)
.execute(&self.pool)
.await?;
Ok(())
}
pub async fn get_usage(
&self,
account_id: &str,
) -> AccountRepositoryResult<Option<AccountUsageRecord>> {
let row = sqlx::query(GET_USAGE_SQL)
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
row.map(|row| usage_from_row(&row)).transpose()
}
pub async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> AccountRepositoryResult<()> {
let row = sqlx::query(SELECT_RATE_LIMIT_WINDOW_SQL)
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
let existing_reset_at = row
.as_ref()
.map(|row| parse_optional_rfc3339(row.get::<Option<String>, _>("window_reset_at")))
.transpose()?
.flatten();
let existing_limit_window_seconds = row
.as_ref()
.and_then(|row| optional_positive_i64_to_u64(row.get("limit_window_seconds")));
let limit_window_seconds_db = limit_window_seconds.map(u64_to_i64_saturating);
let reset_at_db = reset_at.to_rfc3339();
if should_reset_usage_window(
existing_reset_at,
existing_limit_window_seconds,
reset_at,
limit_window_seconds,
) {
sqlx::query(SYNC_RATE_LIMIT_WINDOW_RESET_SQL)
.bind(account_id)
.bind(Utc::now().to_rfc3339())
.bind(reset_at_db)
.bind(limit_window_seconds_db)
.execute(&self.pool)
.await?;
} else {
sqlx::query(SYNC_RATE_LIMIT_WINDOW_SQL)
.bind(account_id)
.bind(reset_at_db)
.bind(limit_window_seconds_db)
.execute(&self.pool)
.await?;
}
Ok(())
}
}
#[derive(Clone)]
pub struct AccountUsageRepository {
pool: SqlitePool,
}
impl AccountUsageRepository {
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
pub async fn list(
&self,
cursor: Option<String>,
limit: u32,
) -> AccountRepositoryResult<Page<AccountUsageListRecord>> {
let fetch_limit = i64::from(limit) + 1;
let rows = if let Some(cursor) = cursor {
let (last_used_at, account_id) =
decode_cursor(&cursor).ok_or(AccountRepositoryError::InvalidCursor)?;
sqlx::query(LIST_ACCOUNT_USAGE_AFTER_CURSOR_SQL)
.bind(&last_used_at)
.bind(last_used_at)
.bind(account_id)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
} else {
sqlx::query(LIST_ACCOUNT_USAGE_SQL)
.bind(fetch_limit)
.fetch_all(&self.pool)
.await?
};
let has_next = rows.len() > limit as usize;
let take_count = rows.len().min(limit as usize);
let mut items = Vec::with_capacity(take_count);
for row in rows.into_iter().take(take_count) {
items.push(usage_list_from_row(&row)?);
}
let next_cursor = if has_next {
items.last().map(|usage| {
encode_cursor(
&usage
.last_used_at
.map(|value| value.to_rfc3339())
.unwrap_or_default(),
&usage.account_id,
)
})
} else {
None
};
Ok(Page { items, next_cursor })
}
pub async fn summary(&self) -> AccountRepositoryResult<AccountUsageSummary> {
let row = sqlx::query(ACCOUNT_USAGE_SUMMARY_SQL)
.fetch_one(&self.pool)
.await?;
Ok(AccountUsageSummary {
account_count: row.get("account_count"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
})
}
pub async fn reset_account(&self, account_id: &str) -> AccountRepositoryResult<()> {
sqlx::query(RESET_ACCOUNT_USAGE_SQL)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(())
}
}
fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> AccountRepositoryResult<AccountUsageRecord> {
Ok(AccountUsageRecord {
account_id: row.get("account_id"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
window_request_count: row.get("window_request_count"),
window_input_tokens: row.get("window_input_tokens"),
window_output_tokens: row.get("window_output_tokens"),
window_cached_tokens: row.get("window_cached_tokens"),
window_started_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("window_started_at"),
)?,
window_reset_at: parse_optional_rfc3339(row.get::<Option<String>, _>("window_reset_at"))?,
limit_window_seconds: optional_positive_i64_to_u64(row.get("limit_window_seconds")),
last_used_at: parse_optional_rfc3339(row.get::<Option<String>, _>("last_used_at"))?,
})
}
fn usage_list_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> AccountRepositoryResult<AccountUsageListRecord> {
Ok(AccountUsageListRecord {
account_id: row.get("account_id"),
email: row.get("email"),
label: row.get("label"),
plan_type: row.get("plan_type"),
request_count: row.get("request_count"),
empty_response_count: row.get("empty_response_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
last_used_at: parse_optional_rfc3339(row.get::<Option<String>, _>("last_used_at"))?,
})
}
fn should_reset_usage_window(
existing_reset_at: Option<DateTime<Utc>>,
existing_limit_window_seconds: Option<u64>,
new_reset_at: DateTime<Utc>,
new_limit_window_seconds: Option<u64>,
) -> bool {
let Some(existing_reset_at) = existing_reset_at else {
return false;
};
if existing_reset_at == new_reset_at {
return false;
}
let drift = existing_reset_at
.signed_duration_since(new_reset_at)
.num_seconds()
.unsigned_abs();
let window_seconds = new_limit_window_seconds
.or(existing_limit_window_seconds)
.unwrap_or(0);
let threshold = if window_seconds > 0 {
window_seconds / 2
} else {
3_600
};
drift >= threshold
}
fn u64_to_i64_saturating(value: u64) -> i64 {
value.min(i64::MAX as u64) as i64
}
src/codex/accounts/usage_snapshots.rs
+10 -10
View File
@@ -45,21 +45,21 @@ impl Granularity {
Granularity::Raw => ts,
Granularity::FiveMin => {
let minutes = (ts.minute() / 5) * 5;
ts.date_naive()
.and_hms_opt(ts.hour(), minutes, 0)
.unwrap()
.and_utc()
align_utc_time(ts, ts.hour(), minutes, 0)
}
Granularity::Hourly => ts
.date_naive()
.and_hms_opt(ts.hour(), 0, 0)
.unwrap()
.and_utc(),
Granularity::Daily => ts.date_naive().and_hms_opt(0, 0, 0).unwrap().and_utc(),
Granularity::Hourly => align_utc_time(ts, ts.hour(), 0, 0),
Granularity::Daily => align_utc_time(ts, 0, 0, 0),
}
}
}
fn align_utc_time(ts: DateTime<Utc>, hour: u32, minute: u32, second: u32) -> DateTime<Utc> {
ts.date_naive()
.and_hms_opt(hour, minute, second)
.map(|value| value.and_utc())
.unwrap_or(ts)
}
#[derive(Debug, Clone)]
pub struct UsageSnapshotStore {
snapshots: VecDeque<UsageSnapshot>,
src/codex/gateway/fingerprint/update_checker.rs
+5 -5
View File
@@ -81,18 +81,18 @@ impl UpdateChecker {
info!(
interval_secs = POLL_INTERVAL.as_secs(),
"UpdateChecker 后台 fingerprint 版本检查器已启动"
"fingerprint 后台版本检查器已启动"
);
// 立即执行首次检查
if let Err(e) = self.check_and_apply_update().await {
warn!(error = %e, "UpdateChecker 首次检查失败");
warn!(error = %e, "fingerprint 首次版本检查失败");
}
loop {
ticker.tick().await;
if let Err(e) = self.check_and_apply_update().await {
warn!(error = %e, "UpdateChecker 定期检查失败");
warn!(error = %e, "fingerprint 定期版本检查失败");
}
}
})
@@ -119,7 +119,7 @@ impl UpdateChecker {
build = %build,
current_version = %current_version,
current_build = %current_build,
"UpdateChecker 发现新的 fingerprint 版本"
"发现新的 fingerprint 版本"
);
self.apply_version_update(version, build).await?;
@@ -131,7 +131,7 @@ impl UpdateChecker {
info!(
version = %version,
build = %build,
"UpdateChecker 已自动应用 fingerprint 版本"
"已自动应用 fingerprint 版本"
);
}
src/codex/gateway/transport/websocket.rs
+17 -686
View File
@@ -1,32 +1,36 @@
use std::{
collections::HashMap,
pin::Pin,
sync::Arc,
time::{Duration, Instant},
};
use std::{pin::Pin, sync::Arc, time::Instant};
use thiserror::Error;
use tokio::{net::TcpStream, sync::Mutex, time::timeout};
use tokio::sync::Mutex;
use tokio_tungstenite::{
connect_async,
tungstenite::{
client::IntoClientRequest,
http::{HeaderMap as WsHeaderMap, Request as WsRequest},
Error as WsError, Message,
client::IntoClientRequest, http::Request as WsRequest, Error as WsError, Message,
},
MaybeTlsStream, WebSocketStream,
};
use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
use reqwest::{header::HeaderMap, StatusCode};
use serde_json::{json, Value};
use crate::codex::gateway::transport::{
rate_limits::{parse_rate_limits_event_raw, rate_limits_to_header_pairs, ParsedRateLimits},
sse::encode_sse_event,
types::CodexResponsesRequest,
};
mod codec;
mod pool;
use codec::{
classify_ws_error_frame, codex_websocket_transport_error, is_internal_websocket_event,
is_terminal_websocket_event, retry_after_seconds_from_body, websocket_event_type,
websocket_message_text, websocket_request_body, websocket_sse_chunk,
};
use pool::{
CodexWebSocketConnectionMetadata, CodexWsStream, PooledWebSocketConnection,
WebSocketPoolAcquire,
};
pub use pool::{CodexWebSocketPool, CodexWebSocketPoolConfig, CodexWebSocketPoolKey};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodexTransport {
HttpSse,
@@ -80,480 +84,6 @@ pub struct CodexWebSocketStreamResponse {
pub rate_limit_updates: SharedRateLimitUpdates,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CodexWebSocketPoolKey {
base_url: String,
account_id: String,
conversation_id: String,
}
impl CodexWebSocketPoolKey {
pub fn new(
base_url: impl Into<String>,
account_id: impl Into<String>,
conversation_id: impl Into<String>,
) -> Self {
Self {
base_url: base_url.into().trim_end_matches('/').to_string(),
account_id: account_id.into(),
conversation_id: conversation_id.into(),
}
}
fn account_id(&self) -> &str {
&self.account_id
}
}
#[derive(Clone)]
pub struct CodexWebSocketPool {
inner: Arc<Mutex<WebSocketPoolState>>,
config: CodexWebSocketPoolConfig,
}
#[derive(Debug, Clone, Copy)]
pub struct CodexWebSocketPoolConfig {
pub enabled: bool,
pub max_age: Duration,
pub max_per_account: usize,
pub maintenance_interval: Option<Duration>,
pub ping_interval: Option<Duration>,
pub ping_timeout: Duration,
pub liveness_timeout: Option<Duration>,
}
impl CodexWebSocketPool {
const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60);
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5);
const DEFAULT_LIVENESS_TIMEOUT: Duration = Duration::from_millis(62_500);
pub fn with_default_max_age() -> Self {
Self::with_config(CodexWebSocketPoolConfig::default())
}
pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self {
Self::with_config(CodexWebSocketPoolConfig {
max_age,
max_per_account,
maintenance_interval: None,
ping_interval: None,
liveness_timeout: None,
..CodexWebSocketPoolConfig::default()
})
}
pub fn with_config(config: CodexWebSocketPoolConfig) -> Self {
let pool = Self {
inner: Arc::new(Mutex::new(WebSocketPoolState::default())),
config,
};
pool.spawn_maintenance_task();
pool
}
pub async fn evict_account(&self, account_id: &str) {
let mut idle_connections = Vec::new();
{
let mut state = self.inner.lock().await;
let keys = state
.slots
.keys()
.filter(|key| key.account_id() == account_id)
.cloned()
.collect::<Vec<_>>();
for key in keys {
if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) {
idle_connections.push(*connection);
}
}
}
close_idle_connections(idle_connections).await;
}
pub async fn shutdown(&self) {
let idle_connections = {
let mut state = self.inner.lock().await;
state.shutting_down = true;
state
.slots
.drain()
.filter_map(|(_, slot)| match slot {
WebSocketPoolSlot::Idle(connection) => Some(*connection),
WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None,
})
.collect::<Vec<_>>()
};
close_idle_connections(idle_connections).await;
}
pub async fn gc_sweep(&self) {
self.maintain_idle_connections().await;
}
async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
let mut expired_connection = None;
let acquire = {
let mut state = self.inner.lock().await;
if !self.config.enabled || state.shutting_down {
return WebSocketPoolAcquire::Bypass;
}
match state.slots.get(key) {
Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) => {
return WebSocketPoolAcquire::Bypass;
}
Some(WebSocketPoolSlot::Idle(_)) => {
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else {
return WebSocketPoolAcquire::Bypass;
};
if connection.created_at.elapsed() <= self.config.max_age {
state.slots.insert(key.clone(), WebSocketPoolSlot::Busy);
return WebSocketPoolAcquire::Reused(connection);
}
expired_connection = Some(*connection);
}
None => {}
}
if self.config.max_per_account == 0
|| account_slot_count(&state.slots, key.account_id()) >= self.config.max_per_account
{
WebSocketPoolAcquire::Bypass
} else {
state.slots.insert(key.clone(), WebSocketPoolSlot::Busy);
WebSocketPoolAcquire::FreshReserved
}
};
if let Some(connection) = expired_connection {
close_idle_connections(vec![connection]).await;
}
acquire
}
async fn put(&self, key: CodexWebSocketPoolKey, connection: PooledWebSocketConnection) {
let mut connection = Some(connection);
{
let mut state = self.inner.lock().await;
let expired = connection
.as_ref()
.is_some_and(|connection| connection.created_at.elapsed() > self.config.max_age);
if expired || state.shutting_down || !self.config.enabled {
state.slots.remove(&key);
} else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Busy)) {
if let Some(connection) = connection.take() {
state
.slots
.insert(key, WebSocketPoolSlot::Idle(Box::new(connection)));
}
}
}
if let Some(connection) = connection {
close_pooled_connection(connection).await;
}
}
async fn discard(&self, key: &CodexWebSocketPoolKey) {
let idle_connection = {
let mut state = self.inner.lock().await;
match state.slots.remove(key) {
Some(WebSocketPoolSlot::Idle(connection)) => Some(*connection),
Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) | None => None,
}
};
if let Some(connection) = idle_connection {
close_idle_connections(vec![connection]).await;
}
}
fn spawn_maintenance_task(&self) {
let Some(interval_duration) = self.config.maintenance_interval else {
return;
};
let inner = Arc::downgrade(&self.inner);
let config = self.config;
let Ok(_handle) = tokio::runtime::Handle::try_current() else {
return;
};
tokio::spawn(async move {
let mut interval = tokio::time::interval(interval_duration);
loop {
interval.tick().await;
let Some(inner) = inner.upgrade() else {
break;
};
let pool = CodexWebSocketPool { inner, config };
if pool.is_shutting_down().await {
break;
}
pool.maintain_idle_connections().await;
}
});
}
async fn is_shutting_down(&self) -> bool {
self.inner.lock().await.shutting_down
}
pub async fn maintain_idle_connections(&self) {
let maintenance = self.take_idle_connections_for_maintenance().await;
close_idle_connections(maintenance.close).await;
for (key, connection) in maintenance.probe {
let connection = probe_idle_connection(connection, self.config.ping_timeout).await;
self.return_maintained_connection(key, connection).await;
}
}
async fn take_idle_connections_for_maintenance(&self) -> WebSocketPoolMaintenance {
let mut close = Vec::new();
let mut probe = Vec::new();
let now = Instant::now();
let Some(ping_interval) = self.config.ping_interval else {
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
return WebSocketPoolMaintenance { close, probe };
}
let expired_keys = expired_idle_keys(&state.slots, self.config.max_age);
for key in expired_keys {
if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) {
close.push(*connection);
}
}
return WebSocketPoolMaintenance { close, probe };
};
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
return WebSocketPoolMaintenance { close, probe };
}
let keys = state
.slots
.iter()
.filter_map(|(key, slot)| match slot {
WebSocketPoolSlot::Idle(connection)
if should_close_idle_connection(
connection,
now,
self.config.max_age,
self.config.liveness_timeout,
) || should_probe_idle_connection(connection, now, ping_interval) =>
{
Some(key.clone())
}
WebSocketPoolSlot::Idle(_)
| WebSocketPoolSlot::Busy
| WebSocketPoolSlot::Checking => None,
})
.collect::<Vec<_>>();
for key in keys {
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) else {
continue;
};
if should_close_idle_connection(
&connection,
now,
self.config.max_age,
self.config.liveness_timeout,
) {
close.push(*connection);
} else {
state.slots.insert(key.clone(), WebSocketPoolSlot::Checking);
probe.push((key, *connection));
}
}
WebSocketPoolMaintenance { close, probe }
}
async fn return_maintained_connection(
&self,
key: CodexWebSocketPoolKey,
connection: Option<PooledWebSocketConnection>,
) {
let mut connection = connection;
{
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
state.slots.remove(&key);
} else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Checking)) {
if let Some(connection) = connection.take() {
state
.slots
.insert(key, WebSocketPoolSlot::Idle(Box::new(connection)));
} else {
state.slots.remove(&key);
}
}
}
if let Some(connection) = connection {
close_pooled_connection(connection).await;
}
}
}
impl Default for CodexWebSocketPool {
fn default() -> Self {
Self::with_default_max_age()
}
}
impl Default for CodexWebSocketPoolConfig {
fn default() -> Self {
Self {
enabled: true,
max_age: CodexWebSocketPool::DEFAULT_MAX_AGE,
max_per_account: CodexWebSocketPool::DEFAULT_MAX_PER_ACCOUNT,
maintenance_interval: Some(CodexWebSocketPool::DEFAULT_MAINTENANCE_INTERVAL),
ping_interval: Some(CodexWebSocketPool::DEFAULT_PING_INTERVAL),
ping_timeout: CodexWebSocketPool::DEFAULT_PING_TIMEOUT,
liveness_timeout: Some(CodexWebSocketPool::DEFAULT_LIVENESS_TIMEOUT),
}
}
}
#[derive(Default)]
struct WebSocketPoolState {
slots: HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
shutting_down: bool,
}
type CodexWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
#[derive(Debug, Clone, PartialEq, Eq)]
struct CodexWebSocketConnectionMetadata {
turn_state: Option<String>,
set_cookie_headers: Vec<String>,
rate_limit_headers: Vec<(String, String)>,
}
impl CodexWebSocketConnectionMetadata {
fn from_headers(headers: &WsHeaderMap) -> Self {
Self {
turn_state: turn_state(headers),
set_cookie_headers: set_cookie_headers(headers),
rate_limit_headers: rate_limit_headers(headers),
}
}
}
struct PooledWebSocketConnection {
websocket: CodexWsStream,
metadata: CodexWebSocketConnectionMetadata,
created_at: Instant,
last_activity_at: Instant,
last_ping_at: Option<Instant>,
}
enum WebSocketPoolSlot {
Idle(Box<PooledWebSocketConnection>),
Busy,
Checking,
}
enum WebSocketPoolAcquire {
Reused(Box<PooledWebSocketConnection>),
FreshReserved,
Bypass,
}
fn account_slot_count(
inner: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
account_id: &str,
) -> usize {
inner
.keys()
.filter(|key| key.account_id() == account_id)
.count()
}
async fn close_idle_connections(connections: Vec<PooledWebSocketConnection>) {
for connection in connections {
close_pooled_connection(connection).await;
}
}
async fn close_pooled_connection(connection: PooledWebSocketConnection) {
let mut websocket = connection.websocket;
let _ = websocket.close(None).await;
}
struct WebSocketPoolMaintenance {
close: Vec<PooledWebSocketConnection>,
probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>,
}
fn expired_idle_keys(
slots: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
max_age: Duration,
) -> Vec<CodexWebSocketPoolKey> {
slots
.iter()
.filter_map(|(key, slot)| match slot {
WebSocketPoolSlot::Idle(connection) if connection.created_at.elapsed() > max_age => {
Some(key.clone())
}
WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => {
None
}
})
.collect()
}
fn should_close_idle_connection(
connection: &PooledWebSocketConnection,
now: Instant,
max_age: Duration,
liveness_timeout: Option<Duration>,
) -> bool {
connection.created_at.elapsed() > max_age
|| liveness_timeout.is_some_and(|timeout| {
!timeout.is_zero() && now.duration_since(connection.last_activity_at) > timeout
})
}
fn should_probe_idle_connection(
connection: &PooledWebSocketConnection,
now: Instant,
ping_interval: Duration,
) -> bool {
!ping_interval.is_zero()
&& connection
.last_ping_at
.is_none_or(|last_ping_at| now.duration_since(last_ping_at) >= ping_interval)
}
async fn probe_idle_connection(
mut connection: PooledWebSocketConnection,
ping_timeout: Duration,
) -> Option<PooledWebSocketConnection> {
let now = Instant::now();
connection.last_ping_at = Some(now);
if connection
.websocket
.send(Message::Ping(Vec::new().into()))
.await
.is_err()
{
close_pooled_connection(connection).await;
return None;
}
if ping_timeout.is_zero() {
return Some(connection);
}
match timeout(ping_timeout, connection.websocket.next()).await {
Ok(Some(Ok(Message::Close(_)))) | Ok(None) | Ok(Some(Err(_))) | Err(_) => {
close_pooled_connection(connection).await;
None
}
Ok(Some(Ok(_message))) => {
connection.last_activity_at = Instant::now();
Some(connection)
}
}
}
struct ActiveWebSocket {
websocket: CodexWsStream,
metadata: CodexWebSocketConnectionMetadata,
@@ -920,77 +450,6 @@ fn websocket_url(base_url: &str) -> String {
}
}
fn websocket_request_body(request: &CodexResponsesRequest) -> Value {
let mut body = json!({
"type": "response.create",
"model": request.model,
"instructions": request.instructions,
"input": request.input,
"stream": true,
"store": false,
"tool_choice": request.tool_choice.clone().unwrap_or_else(|| json!("auto")),
"parallel_tool_calls": request.parallel_tool_calls.unwrap_or(true),
});
if let Some(previous_response_id) = &request.previous_response_id {
body["previous_response_id"] = Value::String(previous_response_id.clone());
}
if let Some(reasoning) = &request.reasoning {
body["reasoning"] = reasoning.clone();
}
if let Some(tools) = request.tools.as_ref().filter(|tools| !tools.is_empty()) {
body["tools"] = Value::Array(tools.clone());
}
if let Some(text) = &request.text {
body["text"] = text.clone();
}
if let Some(service_tier) = &request.service_tier {
body["service_tier"] = Value::String(service_tier.clone());
}
if let Some(prompt_cache_key) = &request.prompt_cache_key {
body["prompt_cache_key"] = Value::String(prompt_cache_key.clone());
}
if let Some(include) = request
.include
.as_ref()
.filter(|include| !include.is_empty())
{
body["include"] = Value::Array(include.iter().cloned().map(Value::String).collect());
}
if let Some(client_metadata) = &request.client_metadata {
body["client_metadata"] = client_metadata.clone();
}
body
}
fn websocket_message_text(message: Message) -> Option<String> {
match message {
Message::Text(text) => Some(text.to_string()),
Message::Binary(bytes) => String::from_utf8(bytes.to_vec()).ok(),
_ => None,
}
}
fn websocket_event_type(raw: &str) -> Option<String> {
serde_json::from_str::<Value>(raw).ok().and_then(|value| {
value
.get("type")
.and_then(Value::as_str)
.map(str::to_string)
})
}
fn websocket_sse_chunk(raw: &str, event: Option<&str>) -> String {
encode_sse_event(event.unwrap_or_default(), raw)
}
fn is_internal_websocket_event(raw: &str) -> bool {
websocket_event_type(raw).as_deref() == Some("codex.rate_limits")
}
fn is_terminal_websocket_event(event: &str) -> bool {
event == "response.completed" || event == "response.failed" || event == "error"
}
async fn forward_websocket_as_sse(
mut active: ActiveWebSocket,
tx: mpsc::UnboundedSender<Result<String, CodexWebSocketError>>,
@@ -1045,131 +504,3 @@ pub async fn append_rate_limit_updates(
headers.extend(rate_limits_to_header_pairs(update));
}
}
struct ClassifiedWebSocketError {
status: StatusCode,
connection_fatal: bool,
}
fn classify_ws_error_frame(raw: &str) -> Option<ClassifiedWebSocketError> {
let value = serde_json::from_str::<Value>(raw).ok()?;
let event_type = value.get("type").and_then(Value::as_str)?;
if event_type != "error" && event_type != "response.failed" {
return None;
}
let code = value
.pointer("/response/error/code")
.or_else(|| value.pointer("/response/error/type"))
.or_else(|| value.pointer("/error/code"))
.or_else(|| value.pointer("/error/type"))
.and_then(Value::as_str)?
.to_ascii_lowercase();
let status = rotatable_error_status(&code)?;
Some(ClassifiedWebSocketError {
status,
connection_fatal: code == "websocket_connection_limit_reached",
})
}
fn rotatable_error_status(code: &str) -> Option<StatusCode> {
match code {
"usage_limit_reached" | "rate_limit_exceeded" | "rate_limit_reached" => {
Some(StatusCode::TOO_MANY_REQUESTS)
}
"quota_exhausted" | "payment_required" => Some(StatusCode::PAYMENT_REQUIRED),
"unauthorized" | "token_invalid" | "token_expired" | "account_deactivated" => {
Some(StatusCode::UNAUTHORIZED)
}
"forbidden" | "account_banned" | "banned" => Some(StatusCode::FORBIDDEN),
"previous_response_not_found" => Some(StatusCode::BAD_REQUEST),
"websocket_connection_limit_reached" => Some(StatusCode::SERVICE_UNAVAILABLE),
_ => None,
}
}
fn codex_websocket_transport_error(error: WsError) -> CodexWebSocketError {
match error {
WsError::Http(response) => {
let (parts, body) = (*response).into_parts();
let body = body
.and_then(|bytes| String::from_utf8(bytes).ok())
.unwrap_or_default();
CodexWebSocketError::Upstream {
status: parts.status,
retry_after_seconds: retry_after_seconds(&parts.headers)
.or_else(|| retry_after_seconds_from_body(&body)),
body,
}
}
error => CodexWebSocketError::Transport(error),
}
}
fn turn_state(headers: &WsHeaderMap) -> Option<String> {
headers
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok())
.map(ToString::to_string)
}
fn set_cookie_headers(headers: &WsHeaderMap) -> Vec<String> {
headers
.get_all("set-cookie")
.iter()
.filter_map(|value| value.to_str().ok().map(ToString::to_string))
.collect()
}
fn rate_limit_headers(headers: &WsHeaderMap) -> Vec<(String, String)> {
headers
.iter()
.filter(|(name, _)| is_rate_limit_header(name.as_str()))
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_string(), value.to_string()))
})
.collect()
}
fn is_rate_limit_header(name: &str) -> bool {
let name = name.to_ascii_lowercase();
name == "retry-after"
|| name.contains("ratelimit")
|| name.contains("rate-limit")
|| name.starts_with("x-codex-primary-")
|| name.starts_with("x-codex-secondary-")
|| name.starts_with("x-codex-code-review-")
|| name.starts_with("x-codex-review-")
|| name.starts_with("x-code-review-")
}
fn retry_after_seconds(headers: &WsHeaderMap) -> Option<u64> {
headers
.get("retry-after")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.trim().parse::<u64>().ok())
.filter(|seconds| *seconds > 0)
}
fn retry_after_seconds_from_body(body: &str) -> Option<u64> {
let value = serde_json::from_str::<Value>(body).ok()?;
let error = value
.pointer("/response/error")
.or_else(|| value.get("error"))
.unwrap_or(&value);
if let Some(seconds) = error
.get("resets_in_seconds")
.and_then(Value::as_u64)
.filter(|seconds| *seconds > 0)
{
return Some(seconds);
}
let resets_at = error.get("resets_at").and_then(Value::as_u64)?;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.ok()?
.as_secs();
(resets_at > now).then_some(resets_at - now)
}
src/codex/gateway/transport/websocket/codec.rs
+206
View File
@@ -0,0 +1,206 @@
use reqwest::StatusCode;
use serde_json::{json, Value};
use tokio_tungstenite::tungstenite::{http::HeaderMap as WsHeaderMap, Error as WsError, Message};
use crate::codex::gateway::transport::{sse::encode_sse_event, types::CodexResponsesRequest};
use super::CodexWebSocketError;
pub(super) struct ClassifiedWebSocketError {
pub(super) status: StatusCode,
pub(super) connection_fatal: bool,
}
pub(super) fn websocket_request_body(request: &CodexResponsesRequest) -> Value {
let mut body = json!({
"type": "response.create",
"model": request.model,
"instructions": request.instructions,
"input": request.input,
"stream": true,
"store": false,
"tool_choice": request.tool_choice.clone().unwrap_or_else(|| json!("auto")),
"parallel_tool_calls": request.parallel_tool_calls.unwrap_or(true),
});
if let Some(previous_response_id) = &request.previous_response_id {
body["previous_response_id"] = Value::String(previous_response_id.clone());
}
if let Some(reasoning) = &request.reasoning {
body["reasoning"] = reasoning.clone();
}
if let Some(tools) = request.tools.as_ref().filter(|tools| !tools.is_empty()) {
body["tools"] = Value::Array(tools.clone());
}
if let Some(text) = &request.text {
body["text"] = text.clone();
}
if let Some(service_tier) = &request.service_tier {
body["service_tier"] = Value::String(service_tier.clone());
}
if let Some(prompt_cache_key) = &request.prompt_cache_key {
body["prompt_cache_key"] = Value::String(prompt_cache_key.clone());
}
if let Some(include) = request
.include
.as_ref()
.filter(|include| !include.is_empty())
{
body["include"] = Value::Array(include.iter().cloned().map(Value::String).collect());
}
if let Some(client_metadata) = &request.client_metadata {
body["client_metadata"] = client_metadata.clone();
}
body
}
pub(super) fn websocket_message_text(message: Message) -> Option<String> {
match message {
Message::Text(text) => Some(text.to_string()),
Message::Binary(bytes) => String::from_utf8(bytes.to_vec()).ok(),
_ => None,
}
}
pub(super) fn websocket_event_type(raw: &str) -> Option<String> {
serde_json::from_str::<Value>(raw).ok().and_then(|value| {
value
.get("type")
.and_then(Value::as_str)
.map(str::to_string)
})
}
pub(super) fn websocket_sse_chunk(raw: &str, event: Option<&str>) -> String {
encode_sse_event(event.unwrap_or_default(), raw)
}
pub(super) fn is_internal_websocket_event(raw: &str) -> bool {
websocket_event_type(raw).as_deref() == Some("codex.rate_limits")
}
pub(super) fn is_terminal_websocket_event(event: &str) -> bool {
event == "response.completed" || event == "response.failed" || event == "error"
}
pub(super) fn classify_ws_error_frame(raw: &str) -> Option<ClassifiedWebSocketError> {
let value = serde_json::from_str::<Value>(raw).ok()?;
let event_type = value.get("type").and_then(Value::as_str)?;
if event_type != "error" && event_type != "response.failed" {
return None;
}
let code = value
.pointer("/response/error/code")
.or_else(|| value.pointer("/response/error/type"))
.or_else(|| value.pointer("/error/code"))
.or_else(|| value.pointer("/error/type"))
.and_then(Value::as_str)?
.to_ascii_lowercase();
let status = rotatable_error_status(&code)?;
Some(ClassifiedWebSocketError {
status,
connection_fatal: code == "websocket_connection_limit_reached",
})
}
pub(super) fn codex_websocket_transport_error(error: WsError) -> CodexWebSocketError {
match error {
WsError::Http(response) => {
let (parts, body) = (*response).into_parts();
let body = body
.and_then(|bytes| String::from_utf8(bytes).ok())
.unwrap_or_default();
CodexWebSocketError::Upstream {
status: parts.status,
retry_after_seconds: retry_after_seconds(&parts.headers)
.or_else(|| retry_after_seconds_from_body(&body)),
body,
}
}
error => CodexWebSocketError::Transport(error),
}
}
pub(super) fn turn_state(headers: &WsHeaderMap) -> Option<String> {
headers
.get("x-codex-turn-state")
.and_then(|value| value.to_str().ok())
.map(ToString::to_string)
}
pub(super) fn set_cookie_headers(headers: &WsHeaderMap) -> Vec<String> {
headers
.get_all("set-cookie")
.iter()
.filter_map(|value| value.to_str().ok().map(ToString::to_string))
.collect()
}
pub(super) fn rate_limit_headers(headers: &WsHeaderMap) -> Vec<(String, String)> {
headers
.iter()
.filter(|(name, _)| is_rate_limit_header(name.as_str()))
.filter_map(|(name, value)| {
value
.to_str()
.ok()
.map(|value| (name.as_str().to_string(), value.to_string()))
})
.collect()
}
pub(super) fn retry_after_seconds_from_body(body: &str) -> Option<u64> {
let value = serde_json::from_str::<Value>(body).ok()?;
let error = value
.pointer("/response/error")
.or_else(|| value.get("error"))
.unwrap_or(&value);
if let Some(seconds) = error
.get("resets_in_seconds")
.and_then(Value::as_u64)
.filter(|seconds| *seconds > 0)
{
return Some(seconds);
}
let resets_at = error.get("resets_at").and_then(Value::as_u64)?;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.ok()?
.as_secs();
(resets_at > now).then_some(resets_at - now)
}
fn rotatable_error_status(code: &str) -> Option<StatusCode> {
match code {
"usage_limit_reached" | "rate_limit_exceeded" | "rate_limit_reached" => {
Some(StatusCode::TOO_MANY_REQUESTS)
}
"quota_exhausted" | "payment_required" => Some(StatusCode::PAYMENT_REQUIRED),
"unauthorized" | "token_invalid" | "token_expired" | "account_deactivated" => {
Some(StatusCode::UNAUTHORIZED)
}
"forbidden" | "account_banned" | "banned" => Some(StatusCode::FORBIDDEN),
"previous_response_not_found" => Some(StatusCode::BAD_REQUEST),
"websocket_connection_limit_reached" => Some(StatusCode::SERVICE_UNAVAILABLE),
_ => None,
}
}
fn is_rate_limit_header(name: &str) -> bool {
let name = name.to_ascii_lowercase();
name == "retry-after"
|| name.contains("ratelimit")
|| name.contains("rate-limit")
|| name.starts_with("x-codex-primary-")
|| name.starts_with("x-codex-secondary-")
|| name.starts_with("x-codex-code-review-")
|| name.starts_with("x-codex-review-")
|| name.starts_with("x-code-review-")
}
fn retry_after_seconds(headers: &WsHeaderMap) -> Option<u64> {
headers
.get("retry-after")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.trim().parse::<u64>().ok())
.filter(|seconds| *seconds > 0)
}
src/codex/gateway/transport/websocket/pool.rs
+492
View File
@@ -0,0 +1,492 @@
use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
use futures::{SinkExt, StreamExt};
use tokio::{net::TcpStream, sync::Mutex, time::timeout};
use tokio_tungstenite::{
tungstenite::{http::HeaderMap as WsHeaderMap, Message},
MaybeTlsStream, WebSocketStream,
};
use super::codec::{rate_limit_headers, set_cookie_headers, turn_state};
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CodexWebSocketPoolKey {
base_url: String,
account_id: String,
conversation_id: String,
}
impl CodexWebSocketPoolKey {
pub fn new(
base_url: impl Into<String>,
account_id: impl Into<String>,
conversation_id: impl Into<String>,
) -> Self {
Self {
base_url: base_url.into().trim_end_matches('/').to_string(),
account_id: account_id.into(),
conversation_id: conversation_id.into(),
}
}
fn account_id(&self) -> &str {
&self.account_id
}
}
#[derive(Clone)]
pub struct CodexWebSocketPool {
inner: Arc<Mutex<WebSocketPoolState>>,
config: CodexWebSocketPoolConfig,
}
#[derive(Debug, Clone, Copy)]
pub struct CodexWebSocketPoolConfig {
pub enabled: bool,
pub max_age: Duration,
pub max_per_account: usize,
pub maintenance_interval: Option<Duration>,
pub ping_interval: Option<Duration>,
pub ping_timeout: Duration,
pub liveness_timeout: Option<Duration>,
}
impl CodexWebSocketPool {
const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60);
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5);
const DEFAULT_LIVENESS_TIMEOUT: Duration = Duration::from_millis(62_500);
pub fn with_default_max_age() -> Self {
Self::with_config(CodexWebSocketPoolConfig::default())
}
pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self {
Self::with_config(CodexWebSocketPoolConfig {
max_age,
max_per_account,
maintenance_interval: None,
ping_interval: None,
liveness_timeout: None,
..CodexWebSocketPoolConfig::default()
})
}
pub fn with_config(config: CodexWebSocketPoolConfig) -> Self {
let pool = Self {
inner: Arc::new(Mutex::new(WebSocketPoolState::default())),
config,
};
pool.spawn_maintenance_task();
pool
}
pub async fn evict_account(&self, account_id: &str) {
let mut idle_connections = Vec::new();
{
let mut state = self.inner.lock().await;
let keys = state
.slots
.keys()
.filter(|key| key.account_id() == account_id)
.cloned()
.collect::<Vec<_>>();
for key in keys {
if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) {
idle_connections.push(*connection);
}
}
}
close_idle_connections(idle_connections).await;
}
pub async fn shutdown(&self) {
let idle_connections = {
let mut state = self.inner.lock().await;
state.shutting_down = true;
state
.slots
.drain()
.filter_map(|(_, slot)| match slot {
WebSocketPoolSlot::Idle(connection) => Some(*connection),
WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None,
})
.collect::<Vec<_>>()
};
close_idle_connections(idle_connections).await;
}
pub async fn gc_sweep(&self) {
self.maintain_idle_connections().await;
}
pub(super) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
let mut expired_connection = None;
let acquire = {
let mut state = self.inner.lock().await;
if !self.config.enabled || state.shutting_down {
return WebSocketPoolAcquire::Bypass;
}
match state.slots.get(key) {
Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) => {
return WebSocketPoolAcquire::Bypass;
}
Some(WebSocketPoolSlot::Idle(_)) => {
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else {
return WebSocketPoolAcquire::Bypass;
};
if connection.created_at.elapsed() <= self.config.max_age {
state.slots.insert(key.clone(), WebSocketPoolSlot::Busy);
return WebSocketPoolAcquire::Reused(connection);
}
expired_connection = Some(*connection);
}
None => {}
}
if self.config.max_per_account == 0
|| account_slot_count(&state.slots, key.account_id()) >= self.config.max_per_account
{
WebSocketPoolAcquire::Bypass
} else {
state.slots.insert(key.clone(), WebSocketPoolSlot::Busy);
WebSocketPoolAcquire::FreshReserved
}
};
if let Some(connection) = expired_connection {
close_idle_connections(vec![connection]).await;
}
acquire
}
pub(super) async fn put(
&self,
key: CodexWebSocketPoolKey,
connection: PooledWebSocketConnection,
) {
let mut connection = Some(connection);
{
let mut state = self.inner.lock().await;
let expired = connection
.as_ref()
.is_some_and(|connection| connection.created_at.elapsed() > self.config.max_age);
if expired || state.shutting_down || !self.config.enabled {
state.slots.remove(&key);
} else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Busy)) {
if let Some(connection) = connection.take() {
state
.slots
.insert(key, WebSocketPoolSlot::Idle(Box::new(connection)));
}
}
}
if let Some(connection) = connection {
close_pooled_connection(connection).await;
}
}
pub(super) async fn discard(&self, key: &CodexWebSocketPoolKey) {
let idle_connection = {
let mut state = self.inner.lock().await;
match state.slots.remove(key) {
Some(WebSocketPoolSlot::Idle(connection)) => Some(*connection),
Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) | None => None,
}
};
if let Some(connection) = idle_connection {
close_idle_connections(vec![connection]).await;
}
}
fn spawn_maintenance_task(&self) {
let Some(interval_duration) = self.config.maintenance_interval else {
return;
};
let inner = Arc::downgrade(&self.inner);
let config = self.config;
let Ok(_handle) = tokio::runtime::Handle::try_current() else {
return;
};
tokio::spawn(async move {
let mut interval = tokio::time::interval(interval_duration);
loop {
interval.tick().await;
let Some(inner) = inner.upgrade() else {
break;
};
let pool = CodexWebSocketPool { inner, config };
if pool.is_shutting_down().await {
break;
}
pool.maintain_idle_connections().await;
}
});
}
async fn is_shutting_down(&self) -> bool {
self.inner.lock().await.shutting_down
}
pub async fn maintain_idle_connections(&self) {
let maintenance = self.take_idle_connections_for_maintenance().await;
close_idle_connections(maintenance.close).await;
for (key, connection) in maintenance.probe {
let connection = probe_idle_connection(connection, self.config.ping_timeout).await;
self.return_maintained_connection(key, connection).await;
}
}
async fn take_idle_connections_for_maintenance(&self) -> WebSocketPoolMaintenance {
let mut close = Vec::new();
let mut probe = Vec::new();
let now = Instant::now();
let Some(ping_interval) = self.config.ping_interval else {
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
return WebSocketPoolMaintenance { close, probe };
}
let expired_keys = expired_idle_keys(&state.slots, self.config.max_age);
for key in expired_keys {
if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) {
close.push(*connection);
}
}
return WebSocketPoolMaintenance { close, probe };
};
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
return WebSocketPoolMaintenance { close, probe };
}
let keys = state
.slots
.iter()
.filter_map(|(key, slot)| match slot {
WebSocketPoolSlot::Idle(connection)
if should_close_idle_connection(
connection,
now,
self.config.max_age,
self.config.liveness_timeout,
) || should_probe_idle_connection(connection, now, ping_interval) =>
{
Some(key.clone())
}
WebSocketPoolSlot::Idle(_)
| WebSocketPoolSlot::Busy
| WebSocketPoolSlot::Checking => None,
})
.collect::<Vec<_>>();
for key in keys {
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) else {
continue;
};
if should_close_idle_connection(
&connection,
now,
self.config.max_age,
self.config.liveness_timeout,
) {
close.push(*connection);
} else {
state.slots.insert(key.clone(), WebSocketPoolSlot::Checking);
probe.push((key, *connection));
}
}
WebSocketPoolMaintenance { close, probe }
}
async fn return_maintained_connection(
&self,
key: CodexWebSocketPoolKey,
connection: Option<PooledWebSocketConnection>,
) {
let mut connection = connection;
{
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
state.slots.remove(&key);
} else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Checking)) {
if let Some(connection) = connection.take() {
state
.slots
.insert(key, WebSocketPoolSlot::Idle(Box::new(connection)));
} else {
state.slots.remove(&key);
}
}
}
if let Some(connection) = connection {
close_pooled_connection(connection).await;
}
}
}
impl Default for CodexWebSocketPool {
fn default() -> Self {
Self::with_default_max_age()
}
}
impl Default for CodexWebSocketPoolConfig {
fn default() -> Self {
Self {
enabled: true,
max_age: CodexWebSocketPool::DEFAULT_MAX_AGE,
max_per_account: CodexWebSocketPool::DEFAULT_MAX_PER_ACCOUNT,
maintenance_interval: Some(CodexWebSocketPool::DEFAULT_MAINTENANCE_INTERVAL),
ping_interval: Some(CodexWebSocketPool::DEFAULT_PING_INTERVAL),
ping_timeout: CodexWebSocketPool::DEFAULT_PING_TIMEOUT,
liveness_timeout: Some(CodexWebSocketPool::DEFAULT_LIVENESS_TIMEOUT),
}
}
}
#[derive(Default)]
struct WebSocketPoolState {
slots: HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
shutting_down: bool,
}
pub(super) type CodexWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct CodexWebSocketConnectionMetadata {
pub(super) turn_state: Option<String>,
pub(super) set_cookie_headers: Vec<String>,
pub(super) rate_limit_headers: Vec<(String, String)>,
}
impl CodexWebSocketConnectionMetadata {
pub(super) fn from_headers(headers: &WsHeaderMap) -> Self {
Self {
turn_state: turn_state(headers),
set_cookie_headers: set_cookie_headers(headers),
rate_limit_headers: rate_limit_headers(headers),
}
}
}
pub(super) struct PooledWebSocketConnection {
pub(super) websocket: CodexWsStream,
pub(super) metadata: CodexWebSocketConnectionMetadata,
pub(super) created_at: Instant,
pub(super) last_activity_at: Instant,
pub(super) last_ping_at: Option<Instant>,
}
enum WebSocketPoolSlot {
Idle(Box<PooledWebSocketConnection>),
Busy,
Checking,
}
pub(super) enum WebSocketPoolAcquire {
Reused(Box<PooledWebSocketConnection>),
FreshReserved,
Bypass,
}
fn account_slot_count(
inner: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
account_id: &str,
) -> usize {
inner
.keys()
.filter(|key| key.account_id() == account_id)
.count()
}
async fn close_idle_connections(connections: Vec<PooledWebSocketConnection>) {
for connection in connections {
close_pooled_connection(connection).await;
}
}
async fn close_pooled_connection(connection: PooledWebSocketConnection) {
let mut websocket = connection.websocket;
let _ = websocket.close(None).await;
}
struct WebSocketPoolMaintenance {
close: Vec<PooledWebSocketConnection>,
probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>,
}
fn expired_idle_keys(
slots: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
max_age: Duration,
) -> Vec<CodexWebSocketPoolKey> {
slots
.iter()
.filter_map(|(key, slot)| match slot {
WebSocketPoolSlot::Idle(connection) if connection.created_at.elapsed() > max_age => {
Some(key.clone())
}
WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => {
None
}
})
.collect()
}
fn should_close_idle_connection(
connection: &PooledWebSocketConnection,
now: Instant,
max_age: Duration,
liveness_timeout: Option<Duration>,
) -> bool {
connection.created_at.elapsed() > max_age
|| liveness_timeout.is_some_and(|timeout| {
!timeout.is_zero() && now.duration_since(connection.last_activity_at) > timeout
})
}
fn should_probe_idle_connection(
connection: &PooledWebSocketConnection,
now: Instant,
ping_interval: Duration,
) -> bool {
!ping_interval.is_zero()
&& connection
.last_ping_at
.is_none_or(|last_ping_at| now.duration_since(last_ping_at) >= ping_interval)
}
async fn probe_idle_connection(
mut connection: PooledWebSocketConnection,
ping_timeout: Duration,
) -> Option<PooledWebSocketConnection> {
let now = Instant::now();
connection.last_ping_at = Some(now);
if connection
.websocket
.send(Message::Ping(Vec::new().into()))
.await
.is_err()
{
close_pooled_connection(connection).await;
return None;
}
if ping_timeout.is_zero() {
return Some(connection);
}
match timeout(ping_timeout, connection.websocket.next()).await {
Ok(Some(Ok(Message::Close(_)))) | Ok(None) | Ok(Some(Err(_))) | Err(_) => {
close_pooled_connection(connection).await;
None
}
Ok(Some(Ok(_message))) => {
connection.last_activity_at = Instant::now();
Some(connection)
}
}
}
src/codex/serving/dispatch/audit.rs
+329
View File
@@ -0,0 +1,329 @@
use std::sync::Arc;
use axum::http::StatusCode;
use serde_json::json;
use tokio::sync::Mutex;
use crate::{
codex::gateway::transport::types::CodexResponsesRequest,
codex::{
accounts::pool::AccountPool,
gateway::transport::{
usage::extract_sse_usage,
websocket::{append_rate_limit_updates, CodexWebSocketError, SharedRateLimitUpdates},
},
logs::event::EventLevel,
},
};
use super::{
limits::apply_rate_limit_headers_with_deps,
log_codex_upstream_response_with_deps, record_response_affinity_with_deps,
record_usage_with_deps,
stream::{ensure_stream_metadata, responses_sse_failure},
CodexRequestLogContext, CodexUpstreamDependencies,
};
pub(super) struct StreamAudit {
deps: CodexUpstreamDependencies,
context: CodexRequestLogContext,
account_slot: AccountSlotGuard,
account_plan_type: Option<String>,
rate_limit_headers: Vec<(String, String)>,
}
impl StreamAudit {
pub(super) fn new(
deps: CodexUpstreamDependencies,
context: CodexRequestLogContext,
account_id: String,
account_plan_type: Option<String>,
rate_limit_headers: Vec<(String, String)>,
) -> Self {
let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account_id);
Self {
deps,
context,
account_slot,
account_plan_type,
rate_limit_headers,
}
}
pub(super) async fn complete(&mut self, body: &[u8]) {
apply_rate_limit_headers_with_deps(
&self.deps,
&self.context.account_id,
self.account_plan_type.as_deref(),
&self.rate_limit_headers,
)
.await;
let body = String::from_utf8_lossy(body);
let mut status = StatusCode::OK;
let mut level = EventLevel::Info;
let mut message = "v1 responses stream 已完成";
let mut metadata = match extract_sse_usage(&body) {
Ok(usage) => {
if let Some(usage) = usage {
if record_usage_with_deps(&self.deps, &self.context.account_id, usage)
.await
.is_err()
{
level = EventLevel::Warn;
message = "v1 responses stream 已完成但 usage 存储失败";
json!({"stream": true, "usage": usage, "usageStoreError": true})
} else {
json!({"stream": true, "usage": usage})
}
} else {
json!({"stream": true, "usage": null})
}
}
Err(error) => {
level = EventLevel::Warn;
message = "v1 responses stream 已完成但 SSE usage 无效";
json!({"stream": true, "sseParseError": error.to_string()})
}
};
match responses_sse_failure(&body) {
Ok(Some(failure)) => {
// SSE 响应头已发出，HTTP 状态不能回滚；用终止事件透传给客户端，并在审计里标记上游失败。
status = StatusCode::BAD_GATEWAY;
level = EventLevel::Error;
message = "v1 responses stream 上游 SSE 失败";
failure.extend_metadata(&mut metadata);
}
Ok(None) => {}
Err(error) => {
level = EventLevel::Warn;
message = "v1 responses stream 已完成但 SSE 失败 metadata 无效";
metadata = json!({"stream": true, "sseParseError": error.to_string()});
}
}
ensure_stream_metadata(&mut metadata, true);
log_codex_upstream_response_with_deps(
&self.deps,
&self.context,
status,
level,
message,
metadata,
)
.await;
self.account_slot.release().await;
}
pub(super) async fn log_transport_error(&mut self, error: &reqwest::Error) {
log_codex_upstream_response_with_deps(
&self.deps,
&self.context,
StatusCode::BAD_GATEWAY,
EventLevel::Error,
"v1 responses stream transport 失败",
json!({"stream": true, "transportError": error.to_string()}),
)
.await;
self.account_slot.release().await;
}
}
pub(super) struct WebSocketStreamAudit {
deps: CodexUpstreamDependencies,
context: CodexRequestLogContext,
account_slot: AccountSlotGuard,
account_plan_type: Option<String>,
request: CodexResponsesRequest,
turn_state: Option<String>,
rate_limit_headers: Vec<(String, String)>,
rate_limit_updates: SharedRateLimitUpdates,
}
impl WebSocketStreamAudit {
#[expect(
clippy::too_many_arguments,
reason = "stream audit captures one immutable snapshot from the upstream response boundary"
)]
pub(super) fn new(
deps: CodexUpstreamDependencies,
context: CodexRequestLogContext,
account_id: String,
account_plan_type: Option<String>,
request: CodexResponsesRequest,
turn_state: Option<String>,
rate_limit_headers: Vec<(String, String)>,
rate_limit_updates: SharedRateLimitUpdates,
) -> Self {
let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account_id);
Self {
deps,
context,
account_slot,
account_plan_type,
request,
turn_state,
rate_limit_headers,
rate_limit_updates,
}
}
pub(super) async fn complete(&mut self, body: &[u8]) {
let mut rate_limit_headers = self.rate_limit_headers.clone();
append_rate_limit_updates(&mut rate_limit_headers, &self.rate_limit_updates).await;
apply_rate_limit_headers_with_deps(
&self.deps,
&self.context.account_id,
self.account_plan_type.as_deref(),
&rate_limit_headers,
)
.await;
let body = String::from_utf8_lossy(body).into_owned();
let usage_result = extract_sse_usage(&body);
let response_usage = match &usage_result {
Ok(usage) => *usage,
Err(_) => None,
};
record_response_affinity_with_deps(
&self.deps,
&self.request,
&self.context.account_id,
&body,
self.turn_state.as_deref(),
response_usage,
)
.await;
let mut status = StatusCode::OK;
let mut level = EventLevel::Info;
let mut message = "v1 responses WebSocket stream 已完成";
let mut metadata = match usage_result {
Ok(Some(usage)) => {
if record_usage_with_deps(&self.deps, &self.context.account_id, usage)
.await
.is_err()
{
level = EventLevel::Warn;
message = "v1 responses WebSocket stream 已完成但 usage 存储失败";
json!({
"stream": true,
"transport": "websocket",
"usage": usage,
"rateLimitHeaders": self.rate_limit_headers.clone(),
"usageStoreError": true,
})
} else {
json!({
"stream": true,
"transport": "websocket",
"usage": usage,
"rateLimitHeaders": self.rate_limit_headers.clone(),
})
}
}
Ok(None) => json!({
"stream": true,
"transport": "websocket",
"usage": null,
"rateLimitHeaders": self.rate_limit_headers.clone(),
}),
Err(error) => {
level = EventLevel::Warn;
message = "v1 responses WebSocket stream 已完成但 SSE usage 无效";
json!({
"stream": true,
"transport": "websocket",
"rateLimitHeaders": self.rate_limit_headers.clone(),
"sseParseError": error.to_string(),
})
}
};
match responses_sse_failure(&body) {
Ok(Some(failure)) => {
// SSE 响应头已经发给客户端，HTTP 状态不能回滚，只能在审计中标记上游失败。
status = StatusCode::BAD_GATEWAY;
level = EventLevel::Error;
message = "v1 responses WebSocket stream 上游 SSE 失败";
failure.extend_metadata(&mut metadata);
}
Ok(None) => {}
Err(error) => {
level = EventLevel::Warn;
message = "v1 responses WebSocket stream 已完成但 SSE 失败 metadata 无效";
metadata = json!({
"stream": true,
"transport": "websocket",
"rateLimitHeaders": self.rate_limit_headers.clone(),
"sseParseError": error.to_string(),
});
}
}
ensure_stream_metadata(&mut metadata, true);
log_codex_upstream_response_with_deps(
&self.deps,
&self.context,
status,
level,
message,
metadata,
)
.await;
self.account_slot.release().await;
}
pub(super) async fn log_transport_error(&mut self, error: &CodexWebSocketError) {
log_codex_upstream_response_with_deps(
&self.deps,
&self.context,
StatusCode::BAD_GATEWAY,
EventLevel::Error,
"v1 responses WebSocket stream transport 失败",
json!({
"stream": true,
"transport": "websocket",
"rateLimitHeaders": self.rate_limit_headers.clone(),
"transportError": error.to_string(),
}),
)
.await;
self.account_slot.release().await;
}
}
struct AccountSlotGuard {
pool: Arc<Mutex<AccountPool>>,
account_id: String,
released: bool,
}
impl AccountSlotGuard {
fn new(pool: Arc<Mutex<AccountPool>>, account_id: String) -> Self {
Self {
pool,
account_id,
released: false,
}
}
async fn release(&mut self) {
if self.released {
return;
}
self.pool.lock().await.release(&self.account_id);
self.released = true;
}
}
impl Drop for AccountSlotGuard {
fn drop(&mut self) {
if self.released {
return;
}
let pool = self.pool.clone();
let account_id = self.account_id.clone();
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
handle.spawn(async move {
pool.lock().await.release(&account_id);
});
}
}
src/codex/serving/dispatch/fallback.rs
+124 -1
View File
@@ -1,7 +1,16 @@
use axum::http::StatusCode;
use chrono::{Duration, Utc};
use serde_json::{json, Value};
use crate::codex::gateway::transport::client::CodexClientError;
use crate::codex::{
accounts::{
model::{Account, AccountStatus},
pool::AccountAcquireRequest,
},
gateway::transport::{client::CodexClientError, rate_limits::cooldown_with_jitter},
};
use super::{usage::record_request_attempt, CodexUpstreamDependencies};
const DEFAULT_RATE_LIMIT_BACKOFF_SECONDS: u64 = 60;
const MAX_RATE_LIMIT_BACKOFF_SECONDS: u64 = 86_400 * 7;
@@ -130,3 +139,117 @@ fn is_cloudflare_challenge(body: &str) -> bool {
|| lower.contains("attention required")
|| lower.contains("just a moment")
}
pub(super) async fn apply_upstream_retry_and_acquire_fallback_with_deps(
deps: &CodexUpstreamDependencies,
account: &Account,
retry: UpstreamAccountRetry,
model: &str,
excluded_account_ids: &mut Vec<String>,
) -> Option<Account> {
apply_upstream_account_retry_with_deps(deps, account, retry).await;
excluded_account_ids.push(account.id.clone());
deps.account_pool
.lock()
.await
.acquire_with(
AccountAcquireRequest::new(model, Utc::now())
.with_exclude_account_ids(excluded_account_ids.iter().cloned()),
)
.map(|fallback| fallback.account)
}
pub(super) async fn apply_upstream_account_retry_with_deps(
deps: &CodexUpstreamDependencies,
account: &Account,
retry: UpstreamAccountRetry,
) {
match retry {
UpstreamAccountRetry::RateLimited {
retry_after_seconds,
} => {
let cooldown_until = Utc::now() + cooldown_with_jitter(retry_after_seconds, 2_000);
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo
.set_quota_cooldown_until(&account.id, cooldown_until)
.await
{
tracing::warn!(
error = %error,
account_id = %account.id,
cooldown_until = %cooldown_until,
"持久化 quota cooldown 失败"
);
}
}
deps.account_pool
.lock()
.await
.mark_quota_limited_until(&account.id, cooldown_until);
if let Err(error) = record_request_attempt(deps, &account.id).await {
tracing::warn!(
error = ?error,
account_id = %account.id,
"记录被 rate limit 的账户请求尝试失败"
);
}
}
UpstreamAccountRetry::QuotaExhausted => {
set_account_status(deps, account, AccountStatus::QuotaExhausted).await;
}
UpstreamAccountRetry::CloudflareChallenge { cooldown_seconds } => {
let cooldown_until = Utc::now() + Duration::seconds(cooldown_seconds as i64);
if let Some(cookie_repo) = deps.cookie_repository.as_ref() {
if let Err(error) = cookie_repo.delete_account_cookies(&account.id).await {
tracing::warn!(
error = %error,
account_id = %account.id,
"清理 Cloudflare 阻断账户 cookies 失败"
);
}
}
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo
.set_cloudflare_cooldown_until(&account.id, cooldown_until)
.await
{
tracing::warn!(
error = %error,
account_id = %account.id,
cooldown_until = %cooldown_until,
"持久化 Cloudflare cooldown 失败"
);
}
}
deps.account_pool
.lock()
.await
.set_cloudflare_cooldown_until(&account.id, cooldown_until);
}
UpstreamAccountRetry::Banned => {
set_account_status(deps, account, AccountStatus::Banned).await;
}
}
}
async fn set_account_status(
deps: &CodexUpstreamDependencies,
account: &Account,
status: AccountStatus,
) {
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo.set_status(&account.id, status).await {
tracing::warn!(
error = %error,
account_id = %account.id,
status = ?status,
"持久化上游账户状态失败"
);
}
}
deps.account_pool
.lock()
.await
.set_status(&account.id, status);
deps.websocket_pool.evict_account(&account.id).await;
}
src/codex/serving/dispatch/limits.rs
+80
View File
@@ -0,0 +1,80 @@
use chrono::Utc;
use serde_json::Value;
use crate::codex::gateway::transport::rate_limits::{parse_rate_limit_headers, rate_limit_quota};
use super::CodexUpstreamDependencies;
pub(super) async fn apply_rate_limit_headers_with_deps(
deps: &CodexUpstreamDependencies,
account_id: &str,
plan_type: Option<&str>,
rate_limit_headers: &[(String, String)],
) {
let Some(rate_limits) = parse_rate_limit_headers(rate_limit_headers) else {
return;
};
let existing_quota = existing_quota_json(deps, account_id).await;
let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref());
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo.update_quota_json(account_id, &quota.to_string()).await {
tracing::warn!(
error = %error,
account_id = %account_id,
"被动同步 quota 缓存失败"
);
}
}
let Some(reset_at) = rate_limits.primary_reset_at() else {
return;
};
deps.account_pool.lock().await.sync_rate_limit_window(
account_id,
reset_at,
rate_limits.primary_limit_window_seconds(),
);
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo
.sync_rate_limit_window(
account_id,
reset_at,
rate_limits.primary_limit_window_seconds(),
)
.await
{
tracing::warn!(
error = %error,
account_id = %account_id,
window_reset_at = %reset_at,
"持久化 rate-limit window 失败"
);
}
}
if !rate_limits.primary_limit_reached() || reset_at <= Utc::now() {
return;
}
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo.set_quota_cooldown_until(account_id, reset_at).await {
tracing::warn!(
error = %error,
account_id = %account_id,
cooldown_until = %reset_at,
"持久化被动 quota cooldown 失败"
);
}
}
deps.account_pool
.lock()
.await
.mark_quota_limited_until(account_id, reset_at);
deps.websocket_pool.evict_account(account_id).await;
}
async fn existing_quota_json(deps: &CodexUpstreamDependencies, account_id: &str) -> Option<Value> {
let repo = deps.account_repository.as_ref()?;
let raw = repo.get_quota_json(account_id).await.ok().flatten()?;
serde_json::from_str(&raw).ok()
}
src/codex/serving/dispatch/mod.rs
+21 -482
View File
@@ -1,5 +1,7 @@
pub mod affinity;
mod audit;
pub mod fallback;
mod limits;
pub mod refresh;
// 上游调度辅助命名为 routing，避免出现 dispatch::dispatch 的模块套娃。
pub mod routing;
@@ -17,7 +19,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use chrono::{Duration, Utc};
use chrono::Utc;
use futures::{stream as futures_stream, StreamExt};
use serde_json::{json, Value};
use tokio::sync::Mutex;
@@ -25,7 +27,7 @@ use tokio::sync::Mutex;
use crate::{
codex::accounts::cookies::repository::CookieRepository,
codex::accounts::{
model::{Account, AccountStatus},
model::Account,
pool::{AccountAcquireRequest, AccountPool},
repository::AccountRepository,
},
@@ -39,12 +41,10 @@ use crate::{
build_reqwest_client, CodexBackendClient, CodexBackendStream,
CodexBackendWebSocketStream, CodexClientError, CodexRequestContext,
},
rate_limits::{cooldown_with_jitter, parse_rate_limit_headers, rate_limit_quota},
types::CodexResponsesRequest,
usage::{extract_sse_usage, TokenUsage},
usage::TokenUsage,
websocket::{
append_rate_limit_updates, transport_for_request, CodexTransport, CodexWebSocketError,
CodexWebSocketPool, SharedRateLimitUpdates,
transport_for_request, CodexTransport, CodexWebSocketError, CodexWebSocketPool,
},
},
codex::logs::{
@@ -65,11 +65,16 @@ pub(crate) use self::{
};
use self::affinity::{compute_variant_hash, SessionAffinityMap, SessionAffinityRepository};
use self::audit::{StreamAudit, WebSocketStreamAudit};
use self::fallback::{
apply_upstream_account_retry_with_deps, apply_upstream_retry_and_acquire_fallback_with_deps,
};
use self::{
limits::apply_rate_limit_headers_with_deps,
refresh::refresh_account_after_unauthorized,
routing::request_domain,
stream::{completed_response_metadata, ensure_stream_metadata, responses_sse_failure},
usage::{record_empty_response_with_deps, record_request_attempt, record_usage_with_deps},
stream::{completed_response_metadata, ensure_stream_metadata},
usage::{record_empty_response_with_deps, record_usage_with_deps},
};
#[derive(Clone)]
@@ -757,194 +762,6 @@ async fn record_response_affinity_with_deps(
}
}
async fn apply_rate_limit_headers_with_deps(
deps: &CodexUpstreamDependencies,
account_id: &str,
plan_type: Option<&str>,
rate_limit_headers: &[(String, String)],
) {
let Some(rate_limits) = parse_rate_limit_headers(rate_limit_headers) else {
return;
};
let existing_quota = existing_quota_json(deps, account_id).await;
let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref());
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo.update_quota_json(account_id, &quota.to_string()).await {
tracing::warn!(
error = %error,
account_id = %account_id,
"被动同步 quota 缓存失败"
);
}
}
let Some(reset_at) = rate_limits.primary_reset_at() else {
return;
};
deps.account_pool.lock().await.sync_rate_limit_window(
account_id,
reset_at,
rate_limits.primary_limit_window_seconds(),
);
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo
.sync_rate_limit_window(
account_id,
reset_at,
rate_limits.primary_limit_window_seconds(),
)
.await
{
tracing::warn!(
error = %error,
account_id = %account_id,
window_reset_at = %reset_at,
"持久化 rate-limit window 失败"
);
}
}
if !rate_limits.primary_limit_reached() || reset_at <= Utc::now() {
return;
}
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo.set_quota_cooldown_until(account_id, reset_at).await {
tracing::warn!(
error = %error,
account_id = %account_id,
cooldown_until = %reset_at,
"持久化被动 quota cooldown 失败"
);
}
}
deps.account_pool
.lock()
.await
.mark_quota_limited_until(account_id, reset_at);
deps.websocket_pool.evict_account(account_id).await;
}
async fn existing_quota_json(deps: &CodexUpstreamDependencies, account_id: &str) -> Option<Value> {
let repo = deps.account_repository.as_ref()?;
let raw = repo.get_quota_json(account_id).await.ok().flatten()?;
serde_json::from_str(&raw).ok()
}
async fn apply_upstream_retry_and_acquire_fallback_with_deps(
deps: &CodexUpstreamDependencies,
account: &Account,
retry: UpstreamAccountRetry,
model: &str,
excluded_account_ids: &mut Vec<String>,
) -> Option<Account> {
apply_upstream_account_retry_with_deps(deps, account, retry).await;
excluded_account_ids.push(account.id.clone());
deps.account_pool
.lock()
.await
.acquire_with(
AccountAcquireRequest::new(model, Utc::now())
.with_exclude_account_ids(excluded_account_ids.iter().cloned()),
)
.map(|fallback| fallback.account)
}
async fn apply_upstream_account_retry_with_deps(
deps: &CodexUpstreamDependencies,
account: &Account,
retry: UpstreamAccountRetry,
) {
match retry {
UpstreamAccountRetry::RateLimited {
retry_after_seconds,
} => {
let cooldown_until = Utc::now() + cooldown_with_jitter(retry_after_seconds, 2_000);
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo
.set_quota_cooldown_until(&account.id, cooldown_until)
.await
{
tracing::warn!(
error = %error,
account_id = %account.id,
cooldown_until = %cooldown_until,
"持久化 quota cooldown 失败"
);
}
}
deps.account_pool
.lock()
.await
.mark_quota_limited_until(&account.id, cooldown_until);
if let Err(error) = record_request_attempt(deps, &account.id).await {
tracing::warn!(
error = ?error,
account_id = %account.id,
"记录被 rate limit 的账户请求尝试失败"
);
}
}
UpstreamAccountRetry::QuotaExhausted => {
set_account_status(deps, account, AccountStatus::QuotaExhausted).await;
}
UpstreamAccountRetry::CloudflareChallenge { cooldown_seconds } => {
let cooldown_until = Utc::now() + Duration::seconds(cooldown_seconds as i64);
if let Some(cookie_repo) = deps.cookie_repository.as_ref() {
if let Err(error) = cookie_repo.delete_account_cookies(&account.id).await {
tracing::warn!(
error = %error,
account_id = %account.id,
"清理 Cloudflare 阻断账户 cookies 失败"
);
}
}
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo
.set_cloudflare_cooldown_until(&account.id, cooldown_until)
.await
{
tracing::warn!(
error = %error,
account_id = %account.id,
cooldown_until = %cooldown_until,
"持久化 Cloudflare cooldown 失败"
);
}
}
deps.account_pool
.lock()
.await
.set_cloudflare_cooldown_until(&account.id, cooldown_until);
}
UpstreamAccountRetry::Banned => {
set_account_status(deps, account, AccountStatus::Banned).await;
}
}
}
async fn set_account_status(
deps: &CodexUpstreamDependencies,
account: &Account,
status: AccountStatus,
) {
if let Some(repo) = deps.account_repository.as_ref() {
if let Err(error) = repo.set_status(&account.id, status).await {
tracing::warn!(
error = %error,
account_id = %account.id,
status = ?status,
"持久化上游账户状态失败"
);
}
}
deps.account_pool
.lock()
.await
.set_status(&account.id, status);
deps.websocket_pool.evict_account(&account.id).await;
}
async fn responses_websocket_stream(
deps: CodexUpstreamDependencies,
request: CodexResponsesRequest,
@@ -1047,17 +864,16 @@ async fn responses_websocket_stream(
.into_response();
}
let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account.id);
let audit = WebSocketStreamAudit {
let audit = WebSocketStreamAudit::new(
deps,
context: log_context,
account_slot,
account_plan_type: account.plan_type,
log_context,
account.id,
account.plan_type,
request,
turn_state: stream_response.turn_state,
rate_limit_headers: stream_response.rate_limit_headers,
rate_limit_updates: stream_response.rate_limit_updates,
};
stream_response.turn_state,
stream_response.rate_limit_headers,
stream_response.rate_limit_updates,
);
let upstream = stream_response.body_stream;
use tokio::time::{interval, Duration};
@@ -1174,283 +990,6 @@ impl CodexRequestLogContext {
}
}
struct StreamAudit {
deps: CodexUpstreamDependencies,
context: CodexRequestLogContext,
account_slot: AccountSlotGuard,
account_plan_type: Option<String>,
rate_limit_headers: Vec<(String, String)>,
}
impl StreamAudit {
fn new(
deps: CodexUpstreamDependencies,
context: CodexRequestLogContext,
account_id: String,
account_plan_type: Option<String>,
rate_limit_headers: Vec<(String, String)>,
) -> Self {
let account_slot = AccountSlotGuard::new(deps.account_pool.clone(), account_id);
Self {
deps,
context,
account_slot,
account_plan_type,
rate_limit_headers,
}
}
async fn complete(&mut self, body: &[u8]) {
apply_rate_limit_headers_with_deps(
&self.deps,
&self.context.account_id,
self.account_plan_type.as_deref(),
&self.rate_limit_headers,
)
.await;
let body = String::from_utf8_lossy(body);
let mut status = StatusCode::OK;
let mut level = EventLevel::Info;
let mut message = "v1 responses stream 已完成";
let mut metadata = match extract_sse_usage(&body) {
Ok(usage) => {
if let Some(usage) = usage {
if record_usage_with_deps(&self.deps, &self.context.account_id, usage)
.await
.is_err()
{
level = EventLevel::Warn;
message = "v1 responses stream 已完成但 usage 存储失败";
json!({"stream": true, "usage": usage, "usageStoreError": true})
} else {
json!({"stream": true, "usage": usage})
}
} else {
json!({"stream": true, "usage": null})
}
}
Err(error) => {
level = EventLevel::Warn;
message = "v1 responses stream 已完成但 SSE usage 无效";
json!({"stream": true, "sseParseError": error.to_string()})
}
};
match responses_sse_failure(&body) {
Ok(Some(failure)) => {
// SSE 响应头已发出，HTTP 状态不能回滚；用终止事件透传给客户端，并在审计里标记上游失败。
status = StatusCode::BAD_GATEWAY;
level = EventLevel::Error;
message = "v1 responses stream 上游 SSE 失败";
failure.extend_metadata(&mut metadata);
}
Ok(None) => {}
Err(error) => {
level = EventLevel::Warn;
message = "v1 responses stream 已完成但 SSE 失败 metadata 无效";
metadata = json!({"stream": true, "sseParseError": error.to_string()});
}
}
ensure_stream_metadata(&mut metadata, true);
log_codex_upstream_response_with_deps(
&self.deps,
&self.context,
status,
level,
message,
metadata,
)
.await;
self.account_slot.release().await;
}
async fn log_transport_error(&mut self, error: &reqwest::Error) {
log_codex_upstream_response_with_deps(
&self.deps,
&self.context,
StatusCode::BAD_GATEWAY,
EventLevel::Error,
"v1 responses stream transport 失败",
json!({"stream": true, "transportError": error.to_string()}),
)
.await;
self.account_slot.release().await;
}
}
struct WebSocketStreamAudit {
deps: CodexUpstreamDependencies,
context: CodexRequestLogContext,
account_slot: AccountSlotGuard,
account_plan_type: Option<String>,
request: CodexResponsesRequest,
turn_state: Option<String>,
rate_limit_headers: Vec<(String, String)>,
rate_limit_updates: SharedRateLimitUpdates,
}
impl WebSocketStreamAudit {
async fn complete(&mut self, body: &[u8]) {
let mut rate_limit_headers = self.rate_limit_headers.clone();
append_rate_limit_updates(&mut rate_limit_headers, &self.rate_limit_updates).await;
apply_rate_limit_headers_with_deps(
&self.deps,
&self.context.account_id,
self.account_plan_type.as_deref(),
&rate_limit_headers,
)
.await;
let body = String::from_utf8_lossy(body).into_owned();
let usage_result = extract_sse_usage(&body);
let response_usage = match &usage_result {
Ok(usage) => *usage,
Err(_) => None,
};
record_response_affinity_with_deps(
&self.deps,
&self.request,
&self.context.account_id,
&body,
self.turn_state.as_deref(),
response_usage,
)
.await;
let mut status = StatusCode::OK;
let mut level = EventLevel::Info;
let mut message = "v1 responses WebSocket stream 已完成";
let mut metadata = match usage_result {
Ok(Some(usage)) => {
if record_usage_with_deps(&self.deps, &self.context.account_id, usage)
.await
.is_err()
{
level = EventLevel::Warn;
message = "v1 responses WebSocket stream 已完成但 usage 存储失败";
json!({
"stream": true,
"transport": "websocket",
"usage": usage,
"rateLimitHeaders": self.rate_limit_headers.clone(),
"usageStoreError": true,
})
} else {
json!({
"stream": true,
"transport": "websocket",
"usage": usage,
"rateLimitHeaders": self.rate_limit_headers.clone(),
})
}
}
Ok(None) => json!({
"stream": true,
"transport": "websocket",
"usage": null,
"rateLimitHeaders": self.rate_limit_headers.clone(),
}),
Err(error) => {
level = EventLevel::Warn;
message = "v1 responses WebSocket stream 已完成但 SSE usage 无效";
json!({
"stream": true,
"transport": "websocket",
"rateLimitHeaders": self.rate_limit_headers.clone(),
"sseParseError": error.to_string(),
})
}
};
match responses_sse_failure(&body) {
Ok(Some(failure)) => {
// SSE 响应头已经发给客户端，HTTP 状态不能回滚，只能在审计中标记上游失败。
status = StatusCode::BAD_GATEWAY;
level = EventLevel::Error;
message = "v1 responses WebSocket stream 上游 SSE 失败";
failure.extend_metadata(&mut metadata);
}
Ok(None) => {}
Err(error) => {
level = EventLevel::Warn;
message = "v1 responses WebSocket stream 已完成但 SSE 失败 metadata 无效";
metadata = json!({
"stream": true,
"transport": "websocket",
"rateLimitHeaders": self.rate_limit_headers.clone(),
"sseParseError": error.to_string(),
});
}
}
ensure_stream_metadata(&mut metadata, true);
log_codex_upstream_response_with_deps(
&self.deps,
&self.context,
status,
level,
message,
metadata,
)
.await;
self.account_slot.release().await;
}
async fn log_transport_error(&mut self, error: &CodexWebSocketError) {
log_codex_upstream_response_with_deps(
&self.deps,
&self.context,
StatusCode::BAD_GATEWAY,
EventLevel::Error,
"v1 responses WebSocket stream transport 失败",
json!({
"stream": true,
"transport": "websocket",
"rateLimitHeaders": self.rate_limit_headers.clone(),
"transportError": error.to_string(),
}),
)
.await;
self.account_slot.release().await;
}
}
struct AccountSlotGuard {
pool: Arc<Mutex<AccountPool>>,
account_id: String,
released: bool,
}
impl AccountSlotGuard {
fn new(pool: Arc<Mutex<AccountPool>>, account_id: String) -> Self {
Self {
pool,
account_id,
released: false,
}
}
async fn release(&mut self) {
if self.released {
return;
}
self.pool.lock().await.release(&self.account_id);
self.released = true;
}
}
impl Drop for AccountSlotGuard {
fn drop(&mut self) {
if self.released {
return;
}
let pool = self.pool.clone();
let account_id = self.account_id.clone();
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
handle.spawn(async move {
pool.lock().await.release(&account_id);
});
}
}
async fn log_codex_upstream_response_with_deps(
deps: &CodexUpstreamDependencies,
context: &CodexRequestLogContext,
tests/codex_accounts/repository.rs
+51 -1
View File
@@ -4,7 +4,7 @@ use secrecy::{ExposeSecret, SecretString};
use codex_proxy_rs::{
codex::accounts::{
model::AccountStatus,
repository::{AccountRepository, NewAccount, TokenUpdate, UsageDelta},
repository::{AccountClaimsUpdate, AccountRepository, NewAccount, TokenUpdate, UsageDelta},
},
platform::crypto::SecretBox,
platform::storage::db::connect_sqlite,
@@ -407,6 +407,56 @@ async fn account_repository_should_preserve_refresh_token_when_update_omits_one(
assert!(loaded.access_token_expires_at.is_some());
}
#[tokio::test]
async fn account_repository_should_preserve_refresh_token_when_claims_update_omits_one() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("accounts.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let repo = AccountRepository::new(pool, SecretBox::new([34u8; 32]));
repo.insert(NewAccount {
id: "acct_a".to_string(),
email: Some("old@example.com".to_string()),
account_id: Some("old-account".to_string()),
user_id: Some("old-user".to_string()),
label: Some("old-label".to_string()),
plan_type: Some("plus".to_string()),
access_token: SecretString::new("old-access".to_string().into()),
refresh_token: Some(SecretString::new("old-refresh".to_string().into())),
access_token_expires_at: None,
status: AccountStatus::Active,
})
.await
.unwrap();
assert!(repo
.update_from_claims(
"acct_a",
AccountClaimsUpdate {
email: Some("new@example.com".to_string()),
account_id: Some("new-account".to_string()),
user_id: Some("new-user".to_string()),
plan_type: Some("pro".to_string()),
access_token: SecretString::new("new-access".to_string().into()),
refresh_token: None,
access_token_expires_at: Some(Utc::now() + Duration::hours(2)),
status: AccountStatus::Active,
},
)
.await
.unwrap());
let loaded = repo.get("acct_a").await.unwrap().unwrap();
assert_eq!(loaded.access_token.expose_secret(), "new-access");
assert_eq!(loaded.refresh_token.unwrap().expose_secret(), "old-refresh");
assert_eq!(loaded.email.as_deref(), Some("new@example.com"));
assert_eq!(loaded.account_id.as_deref(), Some("new-account"));
assert_eq!(loaded.user_id.as_deref(), Some("new-user"));
assert_eq!(loaded.plan_type.as_deref(), Some("pro"));
assert!(loaded.access_token_expires_at.is_some());
}
async fn seed_repo_account(
pool: &sqlx::SqlitePool,
secret_box: &SecretBox,
tests/codex_gateway/websocket.rs
+2 -558
View File
@@ -25,6 +25,8 @@ use codex_proxy_rs::codex::gateway::transport::{
},
};
mod pool;
#[test]
fn transport_for_request_should_default_to_websocket_without_history() {
let request = base_request();
@@ -747,267 +749,6 @@ async fn websocket_success_should_capture_internal_rate_limit_events_without_for
.any(|(name, value)| { name == "x-codex-primary-reset-at" && value == "1893456300" }));
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (release_first_tx, release_first_rx) = oneshot::channel();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_ws =
accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _first_request = first_ws.next().await.unwrap().unwrap();
first_ws
.send(Message::Text(
json!({
"type": "response.output_text.delta",
"delta": "first connection is still busy"
})
.to_string()
.into(),
))
.await
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let mut second_ws =
accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _second_request = second_ws.next().await.unwrap().unwrap();
second_ws
.send(Message::Text(
websocket_completed_response("resp_busy_second", 2, 1).into(),
))
.await
.unwrap();
let third_on_new_connection = tokio::select! {
message = second_ws.next() => {
if matches!(message, Some(Ok(message)) if message.is_text()) {
false
} else {
let (third_stream, _) = listener.accept().await.unwrap();
let mut third_ws = accept_hdr_async(
third_stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _third_request = third_ws.next().await.unwrap().unwrap();
third_ws
.send(Message::Text(
websocket_completed_response("resp_busy_third", 2, 1).into(),
))
.await
.unwrap();
third_ws.close(None).await.unwrap();
true
}
}
accepted = listener.accept() => {
let (third_stream, _) = accepted.unwrap();
let mut third_ws = accept_hdr_async(
third_stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _third_request = third_ws.next().await.unwrap().unwrap();
third_ws
.send(Message::Text(
websocket_completed_response("resp_busy_third", 2, 1).into(),
))
.await
.unwrap();
third_ws.close(None).await.unwrap();
true
}
};
release_first_rx.await.unwrap();
first_ws
.send(Message::Text(
websocket_completed_response("resp_busy_first", 2, 1).into(),
))
.await
.unwrap();
first_ws.close(None).await.unwrap();
third_on_new_connection
});
let pool = Arc::new(CodexWebSocketPool::with_default_max_age());
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(pool, "entry-a");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-a".to_string());
let mut first_stream = client
.websocket_stream_response(&request, request_context("req_busy_first", Some("conv-a")))
.await
.unwrap()
.body_stream;
let first_chunk = first_stream.next().await.unwrap().unwrap();
assert!(first_chunk.contains("first connection is still busy"));
let second = client
.create_response(&request, request_context("req_busy_second", Some("conv-a")))
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_busy_second\""));
let third = client
.create_response(&request, request_context("req_busy_third", Some("conv-a")))
.await
.unwrap();
assert!(third.body.contains("\"id\":\"resp_busy_third\""));
release_first_tx.send(()).unwrap();
while first_stream.next().await.transpose().unwrap().is_some() {}
assert!(
server.await.unwrap(),
"busy pool key reused a second pooled websocket instead of one-shot bypass"
);
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_should_bypass_new_keys_after_account_cap() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_ws =
accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _first_request = first_ws.next().await.unwrap().unwrap();
first_ws
.send(Message::Text(
websocket_completed_response("resp_cap_first", 2, 1).into(),
))
.await
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let mut second_ws =
accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _second_request = second_ws.next().await.unwrap().unwrap();
second_ws
.send(Message::Text(
websocket_completed_response("resp_cap_second", 2, 1).into(),
))
.await
.unwrap();
let third_on_new_connection = tokio::select! {
message = second_ws.next() => {
if matches!(message, Some(Ok(message)) if message.is_text()) {
false
} else {
let (third_stream, _) = listener.accept().await.unwrap();
let mut third_ws = accept_hdr_async(
third_stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _third_request = third_ws.next().await.unwrap().unwrap();
third_ws
.send(Message::Text(
websocket_completed_response("resp_cap_third", 2, 1).into(),
))
.await
.unwrap();
third_ws.close(None).await.unwrap();
true
}
}
accepted = listener.accept() => {
let (third_stream, _) = accepted.unwrap();
let mut third_ws = accept_hdr_async(
third_stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _third_request = third_ws.next().await.unwrap().unwrap();
third_ws
.send(Message::Text(
websocket_completed_response("resp_cap_third", 2, 1).into(),
))
.await
.unwrap();
third_ws.close(None).await.unwrap();
true
}
};
first_ws.close(None).await.unwrap();
third_on_new_connection
});
let pool = Arc::new(CodexWebSocketPool::with_limits(Duration::from_secs(60), 1));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(pool, "entry-cap");
let mut first_request = base_request();
first_request.prompt_cache_key = Some("conversation-one".to_string());
let first = client
.create_response(
&first_request,
request_context("req_cap_first", Some("conv-one")),
)
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_cap_first\""));
let mut second_request = base_request();
second_request.prompt_cache_key = Some("conversation-two".to_string());
let second = client
.create_response(
&second_request,
request_context("req_cap_second", Some("conv-two")),
)
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_cap_second\""));
let third = client
.create_response(
&second_request,
request_context("req_cap_third", Some("conv-two")),
)
.await
.unwrap();
assert!(third.body.contains("\"id\":\"resp_cap_third\""));
assert!(
server.await.unwrap(),
"account cap bypass still allowed a capped conversation to be pooled"
);
}
#[tokio::test]
#[expect(
clippy::result_large_err,
@@ -1057,303 +798,6 @@ async fn websocket_stream_should_error_when_connection_closes_before_terminal_fr
server.await.unwrap();
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_keepalive_should_keep_idle_connection_when_pong_is_received() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _first_request = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
websocket_completed_response("resp_keepalive_first", 2, 1).into(),
))
.await
.unwrap();
let ping = websocket.next().await.unwrap().unwrap();
assert!(
ping.is_ping(),
"expected pool maintenance ping, got {ping:?}"
);
websocket
.send(Message::Pong(Vec::new().into()))
.await
.unwrap();
let reused = tokio::select! {
message = websocket.next() => {
let message = message.unwrap().unwrap();
if !message.is_text() {
return false;
}
websocket
.send(Message::Text(
websocket_completed_response("resp_keepalive_second", 2, 1).into(),
))
.await
.unwrap();
true
}
accepted = listener.accept() => {
let (stream, _) = accepted.unwrap();
let mut websocket = accept_hdr_async(
stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _request = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
websocket_completed_response("resp_keepalive_second", 2, 1).into(),
))
.await
.unwrap();
false
}
};
websocket.close(None).await.unwrap();
reused
});
let pool = Arc::new(CodexWebSocketPool::with_config(keepalive_pool_config(
Duration::from_millis(50),
)));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(Arc::clone(&pool), "entry-keepalive");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-keepalive".to_string());
let first = client
.create_response(
&request,
request_context("req_keepalive_first", Some("conv-keepalive")),
)
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_keepalive_first\""));
pool.maintain_idle_connections().await;
let second = client
.create_response(
&request,
request_context("req_keepalive_second", Some("conv-keepalive")),
)
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_keepalive_second\""));
assert!(
server.await.unwrap(),
"keepalive pong should leave the pooled websocket reusable"
);
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_keepalive_should_evict_idle_connection_without_pong() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_ws =
accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _first_request = first_ws.next().await.unwrap().unwrap();
first_ws
.send(Message::Text(
websocket_completed_response("resp_no_pong_first", 2, 1).into(),
))
.await
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let mut second_ws =
accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _second_request = second_ws.next().await.unwrap().unwrap();
second_ws
.send(Message::Text(
websocket_completed_response("resp_no_pong_second", 2, 1).into(),
))
.await
.unwrap();
second_ws.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::with_config(keepalive_pool_config(
Duration::from_millis(20),
)));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(Arc::clone(&pool), "entry-no-pong");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-no-pong".to_string());
let first = client
.create_response(
&request,
request_context("req_no_pong_first", Some("conv-no-pong")),
)
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_no_pong_first\""));
pool.maintain_idle_connections().await;
let second = client
.create_response(
&request,
request_context("req_no_pong_second", Some("conv-no-pong")),
)
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_no_pong_second\""));
server.await.unwrap();
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_gc_sweep_should_close_expired_idle_connections() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
for response_id in ["resp_gc_first", "resp_gc_second"] {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _request = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
websocket_completed_response(response_id, 2, 1).into(),
))
.await
.unwrap();
websocket.close(None).await.unwrap();
}
});
let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config(
Duration::from_millis(5),
8,
)));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(Arc::clone(&pool), "entry-gc");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-gc".to_string());
let first = client
.create_response(&request, request_context("req_gc_first", Some("conv-gc")))
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_gc_first\""));
sleep(Duration::from_millis(15)).await;
pool.gc_sweep().await;
let second = client
.create_response(&request, request_context("req_gc_second", Some("conv-gc")))
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_gc_second\""));
server.await.unwrap();
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_shutdown_should_close_idle_connections_and_disable_pooling() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
for response_id in ["resp_shutdown_first", "resp_shutdown_second"] {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _request = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
websocket_completed_response(response_id, 2, 1).into(),
))
.await
.unwrap();
websocket.close(None).await.unwrap();
}
});
let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config(
Duration::from_secs(60),
8,
)));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(Arc::clone(&pool), "entry-shutdown");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-shutdown".to_string());
let first = client
.create_response(
&request,
request_context("req_shutdown_first", Some("conv-shutdown")),
)
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_shutdown_first\""));
pool.shutdown().await;
let second = client
.create_response(
&request,
request_context("req_shutdown_second", Some("conv-shutdown")),
)
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_shutdown_second\""));
server.await.unwrap();
}
fn base_request() -> CodexResponsesRequest {
CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new())
}
tests/codex_gateway/websocket/pool.rs
+559
View File
@@ -0,0 +1,559 @@
use super::*;
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let (release_first_tx, release_first_rx) = oneshot::channel();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_ws =
accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _first_request = first_ws.next().await.unwrap().unwrap();
first_ws
.send(Message::Text(
json!({
"type": "response.output_text.delta",
"delta": "first connection is still busy"
})
.to_string()
.into(),
))
.await
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let mut second_ws =
accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _second_request = second_ws.next().await.unwrap().unwrap();
second_ws
.send(Message::Text(
websocket_completed_response("resp_busy_second", 2, 1).into(),
))
.await
.unwrap();
let third_on_new_connection = tokio::select! {
message = second_ws.next() => {
if matches!(message, Some(Ok(message)) if message.is_text()) {
false
} else {
let (third_stream, _) = listener.accept().await.unwrap();
let mut third_ws = accept_hdr_async(
third_stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _third_request = third_ws.next().await.unwrap().unwrap();
third_ws
.send(Message::Text(
websocket_completed_response("resp_busy_third", 2, 1).into(),
))
.await
.unwrap();
third_ws.close(None).await.unwrap();
true
}
}
accepted = listener.accept() => {
let (third_stream, _) = accepted.unwrap();
let mut third_ws = accept_hdr_async(
third_stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _third_request = third_ws.next().await.unwrap().unwrap();
third_ws
.send(Message::Text(
websocket_completed_response("resp_busy_third", 2, 1).into(),
))
.await
.unwrap();
third_ws.close(None).await.unwrap();
true
}
};
release_first_rx.await.unwrap();
first_ws
.send(Message::Text(
websocket_completed_response("resp_busy_first", 2, 1).into(),
))
.await
.unwrap();
first_ws.close(None).await.unwrap();
third_on_new_connection
});
let pool = Arc::new(CodexWebSocketPool::with_default_max_age());
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(pool, "entry-a");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-a".to_string());
let mut first_stream = client
.websocket_stream_response(&request, request_context("req_busy_first", Some("conv-a")))
.await
.unwrap()
.body_stream;
let first_chunk = first_stream.next().await.unwrap().unwrap();
assert!(first_chunk.contains("first connection is still busy"));
let second = client
.create_response(&request, request_context("req_busy_second", Some("conv-a")))
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_busy_second\""));
let third = client
.create_response(&request, request_context("req_busy_third", Some("conv-a")))
.await
.unwrap();
assert!(third.body.contains("\"id\":\"resp_busy_third\""));
release_first_tx.send(()).unwrap();
while first_stream.next().await.transpose().unwrap().is_some() {}
assert!(
server.await.unwrap(),
"busy pool key reused a second pooled websocket instead of one-shot bypass"
);
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_should_bypass_new_keys_after_account_cap() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_ws =
accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _first_request = first_ws.next().await.unwrap().unwrap();
first_ws
.send(Message::Text(
websocket_completed_response("resp_cap_first", 2, 1).into(),
))
.await
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let mut second_ws =
accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _second_request = second_ws.next().await.unwrap().unwrap();
second_ws
.send(Message::Text(
websocket_completed_response("resp_cap_second", 2, 1).into(),
))
.await
.unwrap();
let third_on_new_connection = tokio::select! {
message = second_ws.next() => {
if matches!(message, Some(Ok(message)) if message.is_text()) {
false
} else {
let (third_stream, _) = listener.accept().await.unwrap();
let mut third_ws = accept_hdr_async(
third_stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _third_request = third_ws.next().await.unwrap().unwrap();
third_ws
.send(Message::Text(
websocket_completed_response("resp_cap_third", 2, 1).into(),
))
.await
.unwrap();
third_ws.close(None).await.unwrap();
true
}
}
accepted = listener.accept() => {
let (third_stream, _) = accepted.unwrap();
let mut third_ws = accept_hdr_async(
third_stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _third_request = third_ws.next().await.unwrap().unwrap();
third_ws
.send(Message::Text(
websocket_completed_response("resp_cap_third", 2, 1).into(),
))
.await
.unwrap();
third_ws.close(None).await.unwrap();
true
}
};
first_ws.close(None).await.unwrap();
third_on_new_connection
});
let pool = Arc::new(CodexWebSocketPool::with_limits(Duration::from_secs(60), 1));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(pool, "entry-cap");
let mut first_request = base_request();
first_request.prompt_cache_key = Some("conversation-one".to_string());
let first = client
.create_response(
&first_request,
request_context("req_cap_first", Some("conv-one")),
)
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_cap_first\""));
let mut second_request = base_request();
second_request.prompt_cache_key = Some("conversation-two".to_string());
let second = client
.create_response(
&second_request,
request_context("req_cap_second", Some("conv-two")),
)
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_cap_second\""));
let third = client
.create_response(
&second_request,
request_context("req_cap_third", Some("conv-two")),
)
.await
.unwrap();
assert!(third.body.contains("\"id\":\"resp_cap_third\""));
assert!(
server.await.unwrap(),
"account cap bypass still allowed a capped conversation to be pooled"
);
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_keepalive_should_keep_idle_connection_when_pong_is_received() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _first_request = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
websocket_completed_response("resp_keepalive_first", 2, 1).into(),
))
.await
.unwrap();
let ping = websocket.next().await.unwrap().unwrap();
assert!(
ping.is_ping(),
"expected pool maintenance ping, got {ping:?}"
);
websocket
.send(Message::Pong(Vec::new().into()))
.await
.unwrap();
let reused = tokio::select! {
message = websocket.next() => {
let message = message.unwrap().unwrap();
if !message.is_text() {
return false;
}
websocket
.send(Message::Text(
websocket_completed_response("resp_keepalive_second", 2, 1).into(),
))
.await
.unwrap();
true
}
accepted = listener.accept() => {
let (stream, _) = accepted.unwrap();
let mut websocket = accept_hdr_async(
stream,
|_request: &WsRequest, response| Ok(response),
)
.await
.unwrap();
let _request = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
websocket_completed_response("resp_keepalive_second", 2, 1).into(),
))
.await
.unwrap();
false
}
};
websocket.close(None).await.unwrap();
reused
});
let pool = Arc::new(CodexWebSocketPool::with_config(keepalive_pool_config(
Duration::from_millis(50),
)));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(Arc::clone(&pool), "entry-keepalive");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-keepalive".to_string());
let first = client
.create_response(
&request,
request_context("req_keepalive_first", Some("conv-keepalive")),
)
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_keepalive_first\""));
pool.maintain_idle_connections().await;
let second = client
.create_response(
&request,
request_context("req_keepalive_second", Some("conv-keepalive")),
)
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_keepalive_second\""));
assert!(
server.await.unwrap(),
"keepalive pong should leave the pooled websocket reusable"
);
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_keepalive_should_evict_idle_connection_without_pong() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_ws =
accept_hdr_async(first_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _first_request = first_ws.next().await.unwrap().unwrap();
first_ws
.send(Message::Text(
websocket_completed_response("resp_no_pong_first", 2, 1).into(),
))
.await
.unwrap();
let (second_stream, _) = listener.accept().await.unwrap();
let mut second_ws =
accept_hdr_async(second_stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _second_request = second_ws.next().await.unwrap().unwrap();
second_ws
.send(Message::Text(
websocket_completed_response("resp_no_pong_second", 2, 1).into(),
))
.await
.unwrap();
second_ws.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::with_config(keepalive_pool_config(
Duration::from_millis(20),
)));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(Arc::clone(&pool), "entry-no-pong");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-no-pong".to_string());
let first = client
.create_response(
&request,
request_context("req_no_pong_first", Some("conv-no-pong")),
)
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_no_pong_first\""));
pool.maintain_idle_connections().await;
let second = client
.create_response(
&request,
request_context("req_no_pong_second", Some("conv-no-pong")),
)
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_no_pong_second\""));
server.await.unwrap();
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_gc_sweep_should_close_expired_idle_connections() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
for response_id in ["resp_gc_first", "resp_gc_second"] {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _request = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
websocket_completed_response(response_id, 2, 1).into(),
))
.await
.unwrap();
websocket.close(None).await.unwrap();
}
});
let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config(
Duration::from_millis(5),
8,
)));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(Arc::clone(&pool), "entry-gc");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-gc".to_string());
let first = client
.create_response(&request, request_context("req_gc_first", Some("conv-gc")))
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_gc_first\""));
sleep(Duration::from_millis(15)).await;
pool.gc_sweep().await;
let second = client
.create_response(&request, request_context("req_gc_second", Some("conv-gc")))
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_gc_second\""));
server.await.unwrap();
}
#[tokio::test]
#[expect(
clippy::result_large_err,
reason = "tokio-tungstenite handshake callbacks use a large error response type"
)]
async fn websocket_pool_shutdown_should_close_idle_connections_and_disable_pooling() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
for response_id in ["resp_shutdown_first", "resp_shutdown_second"] {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket =
accept_hdr_async(stream, |_request: &WsRequest, response| Ok(response))
.await
.unwrap();
let _request = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(
websocket_completed_response(response_id, 2, 1).into(),
))
.await
.unwrap();
websocket.close(None).await.unwrap();
}
});
let pool = Arc::new(CodexWebSocketPool::with_config(manual_pool_config(
Duration::from_secs(60),
8,
)));
let client = CodexBackendClient::new(
build_reqwest_client(false).unwrap(),
format!("http://{addr}"),
codex_proxy_rs::codex::gateway::fingerprint::model::Fingerprint::default_for_tests(),
)
.with_websocket_pool(Arc::clone(&pool), "entry-shutdown");
let mut request = base_request();
request.prompt_cache_key = Some("conversation-shutdown".to_string());
let first = client
.create_response(
&request,
request_context("req_shutdown_first", Some("conv-shutdown")),
)
.await
.unwrap();
assert!(first.body.contains("\"id\":\"resp_shutdown_first\""));
pool.shutdown().await;
let second = client
.create_response(
&request,
request_context("req_shutdown_second", Some("conv-shutdown")),
)
.await
.unwrap();
assert!(second.body.contains("\"id\":\"resp_shutdown_second\""));
server.await.unwrap();
}