From 0d86975ca696eb90992b0abbd3b5849e0345317a Mon Sep 17 00:00:00 2001 From: zyycn Date: Sun, 14 Jun 2026 14:22:35 +0800 Subject: [PATCH] refactor: align storage audits Co-authored-by: Codex --- docs/FINGERPRINT_AUTO_UPDATE.md | 24 +- docs/database-storage-audit.md | 87 +++++ docs/scheduler-completion-report.md | 342 +++++++----------- docs/scheduler-implementation.md | 212 ++--------- src/admin/auth/mod.rs | 1 + src/admin/auth/repository.rs | 94 +++++ src/admin/auth/service.rs | 74 +--- src/admin/http/logs.rs | 6 + src/admin/tasks/session_cleanup.rs | 18 +- src/codex/accounts/cookies/repository.rs | 21 +- src/codex/accounts/repository.rs | 317 ++++++++++++++-- src/codex/gateway/fingerprint/repository.rs | 34 +- .../gateway/fingerprint/update_checker.rs | 41 +-- src/codex/logs/repository.rs | 12 + src/codex/logs/service.rs | 52 ++- src/codex/serving/dispatch/affinity.rs | 235 +++++++++++- src/codex/serving/dispatch/mod.rs | 82 ++++- src/codex/serving/dispatch/refresh.rs | 5 +- src/codex/serving/responses.rs | 11 +- src/platform/storage/db.rs | 68 +++- src/platform/storage/schema.sql | 32 +- src/runtime/bootstrap.rs | 41 +-- src/runtime/state.rs | 31 +- src/runtime/tasks/coordinator.rs | 10 +- tests/admin.rs | 1 + tests/admin/auth_repository.rs | 57 +++ tests/codex_accounts/cookie_store.rs | 22 ++ tests/codex_accounts/repository.rs | 90 +++++ tests/codex_gateway/fingerprint_update.rs | 29 ++ tests/codex_operations/logs_pagination.rs | 91 +++++ tests/codex_serving/responses_http_sse.rs | 71 +++- tests/codex_serving/responses_websocket.rs | 12 + tests/codex_serving/upstream_errors.rs | 4 +- tests/codex_serving/upstream_fallback.rs | 4 +- tests/platform/storage_schema.rs | 120 ++++++ tests/runtime.rs | 2 + tests/runtime/startup.rs | 191 ++++++++++ tests/support/upstream.rs | 17 + 38 files changed, 1925 insertions(+), 636 deletions(-) create mode 100644 docs/database-storage-audit.md create mode 100644 src/admin/auth/repository.rs create mode 100644 tests/admin/auth_repository.rs diff --git a/docs/FINGERPRINT_AUTO_UPDATE.md b/docs/FINGERPRINT_AUTO_UPDATE.md index d7c80d6c8..bf2520f08 100644 --- a/docs/FINGERPRINT_AUTO_UPDATE.md +++ b/docs/FINGERPRINT_AUTO_UPDATE.md @@ -110,7 +110,7 @@ data/ ### 4. 数据库持久化 -自动更新 `fingerprints` 表： +自动更新通过 `FingerprintRepository::upsert_auto_update()` 写入 `fingerprints` 表的 `auto_updated` 当前记录： ```sql insert into fingerprints ( @@ -120,13 +120,14 @@ insert into fingerprints ( on conflict(id) do update set app_version = excluded.app_version, build_number = excluded.build_number, - chromium_version = excluded.chromium_version + chromium_version = excluded.chromium_version, + created_at = excluded.created_at ``` **好处：** -- 记录历史版本变更 -- 支持多指纹配置 -- 审计和回滚能力 +- 启动时可以稳定加载最新 `auto_update` fingerprint +- 自动更新逻辑不再直接持有业务表 SQL +- 历史记录写入仍由 `FingerprintRepository::insert_update()` 支持 ### 5. Chromium 版本匹配 @@ -260,19 +261,24 @@ pub enum UpdateError { --- -## 新增文件 +## 相关文件 1. **src/codex/gateway/fingerprint/update_checker.rs** - - 340 行核心实现 - UpdateChecker 结构体 - Appcast 解析逻辑 - 状态持久化 + - 通过 `FingerprintRepository` 写入数据库 -2. **src/platform/storage/paths.rs** +2. **src/codex/gateway/fingerprint/repository.rs** + - fingerprint 历史记录写入 + - `auto_updated` 当前记录 upsert + - 启动时加载 `auto_update` fingerprint + +3. **src/platform/storage/paths.rs** - 数据目录路径管理 - `data_dir()` 和 `ensure_data_dir()` -3. **修改文件：** +4. **修改文件：** - `src/codex/gateway/fingerprint/mod.rs` - 添加模块导出 - `src/platform/storage/mod.rs` - 添加 paths 模块 - `src/runtime/tasks/coordinator.rs` - 启动 update checker diff --git a/docs/database-storage-audit.md b/docs/database-storage-audit.md new file mode 100644 index 000000000..4a1cc02fb --- /dev/null +++ b/docs/database-storage-audit.md @@ -0,0 +1,87 @@ +# 数据库存储审计与优化计划 + +## 结论 + +当前数据库设计的主体边界是清晰的：账号密钥、cookie、quota 缓存、事件日志、管理会话和模型快照都有独立表承载。账号调度依赖的窗口用量和 session affinity 已从纯内存状态收敛到 SQLite，重启后的短期调度退化已经被消除。事件日志写入也已统一到 `LogService`，请求链路会遵守运行时日志策略。 + +最优先要修复的是账号用量窗口状态。以下字段属于 `account_usage`，不是 `accounts`： + +- `window_request_count` +- `window_input_tokens` +- `window_output_tokens` +- `window_cached_tokens` +- `window_started_at` +- `window_reset_at` +- `limit_window_seconds` + +原因是它们描述的是账号在当前 rate-limit window 内的用量快照，而不是账号身份、密钥或生命周期属性。放在 `account_usage` 可以让累计统计和窗口统计共享同一个存储边界，也能让 `least_used` 调度在重启后恢复更多上下文。 + +## 已验证问题 + +`Account` 和 `AccountPool` 已经有窗口字段，调度逻辑也会使用 `window_reset_at`、窗口内请求数和 token 计数。审计时 `account_usage` 表没有这些列，`AccountRepository::list_pool_accounts()` 也把窗口字段固定恢复为 0 或 `None`。当前实现已将这些字段落到 `account_usage`，并在旧库启动时自动补列。 + +这不会损坏 access token、refresh token 或 cookie，也不会造成永久性数据异常。真实影响是： + +- 重启后窗口计数归零，`least_used` 会丢失短期负载信息。 +- `window_reset_at` 丢失后，调度无法按上游 rate-limit window 做更精确排序。 +- 窗口内 token 消耗无法跨进程保留，统计和排障信息不完整。 + +Claude 审计中“窗口字段未持久化”的结论成立，但“配额控制失效”的表述偏重。项目已经持久化 `quota_json`、`quota_limit_reached`、`quota_cooldown_until` 和 Cloudflare cooldown；缺失的是窗口统计快照，而不是所有配额状态。 + +## 目标存储边界 + +`accounts` 继续只承载账号身份、密钥密文、状态和上游 quota/cooldown 摘要。 + +`account_usage` 承载所有本地统计： + +- 累计请求和 token 计数。 +- 空响应计数。 +- 当前窗口请求和 token 计数。 +- 当前窗口开始、重置时间和窗口长度。 + +`session_affinities` 承载 `previous_response_id` 到账号和会话上下文的映射。当前实现已补 `SessionAffinityRepository`，完成响应时同步写库，启动时清理过期记录并恢复未过期映射。表结构补充 `function_call_ids_json`，避免工具调用分支信息只存在于内存。 + +事件日志通过统一日志服务入口写入，避免请求链路绕过日志容量、开关和 body capture 策略。 + +## 优化阶段 + +### P0：持久化 account_usage 窗口状态（已落地） + +- 扩展 `account_usage` schema。 +- 对旧库启动时自动补列。 +- `record_usage` 同时更新累计计数和窗口计数。 +- `list_pool_accounts` 从数据库恢复窗口字段。 +- 收到 rate-limit header 时持久化 `window_reset_at` 和 `limit_window_seconds`。 + +### P1：持久化 session affinity（已落地） + +- 增加 `SessionAffinityRepository`。 +- 将内存 affinity map 的写入同步到 `session_affinities`。 +- 启动时清理过期记录并恢复未过期映射。 +- 补充 `function_call_ids_json`，恢复 `function_call_ids`。 + +### P1：统一事件日志入口（已落地） + +- v1 response 和 401 后账户刷新失败事件通过 `LogService::record` 写入。 +- 在一个入口执行 `enabled`、`capacity`、`capture_body` 和 metadata 规范。 +- 管理端 `/admin/logs/state` 和请求链路共享同一个 `LogService` 实例，运行时开关立即生效。 + +### P2：收敛 schema 约束（已落地） + +- 为账号状态、事件日志级别增加枚举 check constraints。 +- 为 enabled、quota_limit_reached 增加 0/1 check constraints。 +- 为 usage 计数、窗口计数、latency、session input_tokens 增加非负 check constraints。 +- 为 `limit_window_seconds` 增加正数约束。 + +### P2：收敛直接 SQL（已落地） + +- `admin_users` 和 `admin_sessions` 已收敛到 `AdminAuthRepository`。 +- 管理登录、session 校验、默认管理员初始化和过期 session 清理都通过 repository 访问数据库。 +- fingerprint 更新已从 updater 内部 SQL 迁移到 `FingerprintRepository`。 +- `fingerprints` 同时支持历史记录写入和 `auto_updated` 当前记录 upsert。 + +### P2：规范 cookie `expires_at` 读取语义（已落地） + +- `CookieRepository::cookie_header` 读取时过滤已过期 cookie。 +- 支持解析标准 `Expires` 时间和 RFC3339 时间。 +- 无法解析的历史值保守保留，避免误删仍可能有效的 cookie。 diff --git a/docs/scheduler-completion-report.md b/docs/scheduler-completion-report.md index c45a216bc..601236228 100644 --- a/docs/scheduler-completion-report.md +++ b/docs/scheduler-completion-report.md @@ -1,228 +1,130 @@ -# 🎉 后台调度器实现总结报告 - -## 任务完成状态 - -✅ **完成** - 已成功实现 TypeScript 原版 `codex-proxy` 项目的后台调度器核心功能 - -## 实现的功能 - -### 1. RefreshScheduler - OAuth 令牌刷新调度器 - -**文件**: `src/scheduler/refresh.rs` (458 行) - -**核心功能**: -- ✅ JWT 令牌过期时间解析 -- ✅ 在 `exp - margin` 时刻自动调度刷新 -- ✅ 指数退避重试策略：5s → 15s → 45s → 135s → 300s（最多5次） -- ✅ 永久失败检测（识别 `invalid_grant`、`invalid_token`、`account has been deactivated`、`refresh_token_reused` 等错误） -- ✅ 临时失败恢复调度（10分钟后重试） -- ✅ 崩溃恢复机制： - - `refreshing` 状态账户 → 立即重试 - - `expired` + 有 refreshToken 的账户 → 延迟重试（30秒起，间隔2秒） -- ✅ 并发控制（`in_flight` 标记防止重复刷新） -- ✅ 优雅关闭支持（通过 `SchedulerHandle`） - -**技术亮点**: -```rust -// 避免异步递归的设计 -- schedule_one() - 初始调度 -- schedule_next_refresh() - 刷新后调度，避免递归 - -// 并发控制 -in_flight: Arc>> - -// 永久错误分类 -const BAN_ERRORS: &[&str] = &["account has been deactivated", "refresh_token_reused"]; -const EXPIRED_ERRORS: &[&str] = &["invalid_grant", "invalid_token", "access_denied", ...]; -``` - -### 2. SessionCleanupScheduler - 会话清理调度器 - -**文件**: `src/scheduler/session_cleanup.rs` (62 行) - -**核心功能**: -- ✅ 定期删除过期的管理员会话 -- ✅ 可配置的清理间隔 -- ✅ 直接 SQL 操作（高效、轻量级） -- ✅ 优雅关闭支持 - -**实现**: -```rust -DELETE FROM admin_sessions WHERE expires_at < ? -``` - -### 3. 调度器基础设施 - -**文件**: `src/scheduler/types.rs` (27 行) - -**提供**: -- ✅ `SchedulerHandle` - 统一的调度器句柄 -- ✅ `SchedulerError` - 类型化错误处理 -- ✅ `SchedulerResult` - 统一的结果类型 - -## 统计数据 - -### 代码量 -- **新增 Rust 代码**: 553 行 -- **总计新增（含注释和空行）**: 574 行 -- **新增文件**: 4 个 Rust 源文件 - - `src/scheduler/mod.rs` - - `src/scheduler/types.rs` - - `src/scheduler/refresh.rs` - - `src/scheduler/session_cleanup.rs` - -### 文档 -- `SCHEDULER_IMPLEMENTATION.md` - 详细实现说明（200+ 行） -- `docs/scheduler-usage.md` - 使用指南和示例（220+ 行） - -### 测试状态 -- ✅ 所有 152 个现有测试保持通过 -- ✅ `cargo check --lib` 通过 -- ✅ `cargo clippy --all-targets --all-features --locked -- -D warnings` 通过 -- ✅ `cargo fmt --check` 通过 - -## 提交记录 - -``` -ce93de6 - feat: add background refresh and session cleanup schedulers -2025594 - docs: add scheduler implementation and usage documentation -``` - -## 与 TypeScript 原版的对比 - -| 特性 | TypeScript 版本 | Rust 版本 | -|------|----------------|----------| -| 异步处理 | 回调风格 | `async/await` | -| 定时器 | `setTimeout`/`setInterval` | `tokio::spawn` + `tokio::time` | -| 并发控制 | 文件锁 | `RwLock` | -| 错误处理 | 字符串/Error | 类型化 `thiserror` | -| 类型安全 | 弱类型 | 强类型 | - -## 架构设计亮点 - -### 1. 避免异步递归 -问题：`schedule_one` → `do_refresh` → `schedule_one` 会导致无限递归 - -解决：拆分为两个方法 -- `schedule_one()` - 初始调度，直接执行刷新 -- `schedule_next_refresh()` - 刷新后调度，避免递归 - -### 2. 并发安全 -使用 `Arc>` 实现线程安全的状态管理： -```rust -in_flight: Arc>> -timers: Arc>>> -destroyed: Arc> -``` - -### 3. 优雅关闭 -所有调度器通过 `mpsc::channel` 实现优雅关闭： -```rust -tokio::select! { - _ = ticker.tick() => { /* 执行任务 */ } - _ = shutdown_rx.recv() => { - info!("Scheduler shutting down"); - break; - } -} -``` - -## 集成到现有系统 - -### 新增的 AccountService 方法 - -```rust -impl AccountService { - /// 列出所有账户用于刷新调度器 - pub async fn list_all_for_refresh(&self) -> Result, AccountServiceError> { - self.repository()?.list_all().await.map_err(|_| AccountServiceError::List) - } -} - -impl std::fmt::Display for RefreshAccountError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // 实现 Display trait，用于调度器日志 - } -} -``` - -## 使用示例 - -```rust -// 1. 创建调度器 -let refresh_scheduler = RefreshScheduler::new(account_service, config); -let refresh_handle = refresh_scheduler.start().await; - -let session_cleanup = SessionCleanupScheduler::new(db, 3600); -let cleanup_handle = session_cleanup.start(); - -// 2. 运行应用 -// ... - -// 3. 优雅关闭 -refresh_handle.shutdown().await; -cleanup_handle.shutdown().await; -``` - -## 下一步建议 - -### 必要的集成工作 -1. **在 main.rs 中集成调度器** - - 启动时初始化调度器 - - 优雅关闭时停止调度器 - -2. **配置文件更新** - - 添加 `session_cleanup_interval_secs` 配置项 - -### 可选的增强功能 -3. **监控和可观测性** - - 添加调度器健康检查端点 - - 暴露刷新成功/失败指标 - - 活跃定时器数量监控 - -4. **高级功能** - - 跨进程刷新锁（用于多实例部署） - - 磁盘刷新令牌同步（防止一次性 RT 被多次消费） - - 配额刷新调度器 - - 模型刷新调度器 - - 指纹更新调度器 - -## 对应计划文档的完成情况 - -根据 `/home/zyy/桌面/Codes/codex-proxy-rs/docs/superpowers/plans/2026-06-11-codex-proxy-rs.md`: - -**Line 3002 - Refresh scheduler parity**: -- ✅ 长期运行的刷新定时器在 `exp - margin` -- ✅ 崩溃恢复处理 `refreshing`/`expired` 账户 -- ✅ 指数退避 -- ✅ 恢复调度 -- ✅ 永久失败阈值检测 -- ✅ per-account in-flight 抑制 -- ⏳ 跨进程刷新锁（可选） -- ⏳ 磁盘刷新令牌同步（可选） - -**Line 3005 - Background service lifecycle**: -- ✅ 账户刷新调度器实现 -- ✅ 会话清理调度器实现 -- ⏳ 启动时初始化（需集成到 main.rs） -- ⏳ 优雅关闭（需集成到 main.rs） -- ❌ 模型刷新（未实现，可选） -- ❌ 配额刷新（未实现，可选） -- ❌ 指纹轮询（未实现，可选） +# 账号调度与数据库审计落地报告 ## 结论 -✅ **核心功能已完成** - RefreshScheduler 和 SessionCleanupScheduler 的实现符合 TypeScript 原版的功能要求 +本轮把账号调度审计和数据库存储审计合并落地。当前实现已经覆盖审计中需要立即修复的调度、刷新、窗口统计、session affinity、日志入口、schema 约束、cookie 过期语义和直接 SQL 边界问题。 -🎯 **质量保证** - 所有测试通过，代码符合 Rust 最佳实践，完整文档 +不在本报告范围内的事项：IP 代理/VPN、旧项目数据迁移、真实 Codex Desktop TLS 指纹完全复刻。这些不是当前 Rust 服务要移植的能力。 -🚀 **可立即使用** - 提供了完整的使用示例和集成指南 +## 账号调度审计项 -📈 **下一步清晰** - 需要在 main.rs 中集成启动逻辑，其他功能为可选增强 +### least_used 策略 ---- +已实现完整优先级： -**实现时间**: 2026-06-13 -**实现者**: Claude Opus 4.8 -**代码审查**: ✅ 通过 clippy 和 rustfmt -**测试状态**: ✅ 所有测试通过 +1. quota-limited 账号降权。 +2. `window_reset_at` 更早的账号优先。 +3. `request_count` 更少的账号优先。 +4. 并列时轮换，避免固定账号长期被选中。 + +对应测试： + +- `least_used_should_deprioritize_quota_limited_accounts_when_skip_is_disabled` +- `least_used_should_prefer_earlier_rate_limit_window_reset` +- `least_used_should_prefer_lower_runtime_request_count` +- `account_pool_should_rotate_tied_least_used_accounts` + +### refresh_concurrency + +已通过 `tokio::sync::Semaphore` 增加全局刷新并发限制。配置来源为 `auth.refresh_concurrency`，并对 0 做最小值保护。 + +对应测试： + +- `refresh_scheduler_limits_refresh_concurrency` +- `refresh_concurrency_limit_should_never_be_zero` + +### 请求计数和窗口统计 + +`Account`、`AccountPool`、`AccountRepository` 和 `account_usage` 现在共同维护： + +- 累计 `request_count` +- 窗口内请求数 +- 窗口内 input/output/cached token 计数 +- `window_started_at` +- `window_reset_at` +- `limit_window_seconds` + +这些字段会在请求完成时累加，在收到 rate-limit header 时同步窗口边界，并在启动恢复账号池时重新加载。 + +对应测试： + +- `account_repository_should_accumulate_usage_window_counters` +- `account_repository_should_restore_window_usage_into_runtime_pool_accounts` +- `v1_responses_should_passively_cache_rate_limit_headers` +- `sqlite_schema_should_persist_account_usage_window_columns` + +## 数据库存储审计项 + +### account_usage 窗口状态 + +已落地到 `account_usage`，并为旧库启动补列。窗口统计属于 usage 域，不放入 `accounts`。 + +### session affinity + +已增加 `SessionAffinityRepository`，完成响应后写入 SQLite，启动时删除过期记录并恢复未过期映射。`function_call_ids_json` 已持久化，避免工具调用分支信息只存在于内存。 + +对应测试： + +- `app_state_should_restore_session_affinity_from_sqlite` +- `v1_responses_should_route_previous_response_id_to_recorded_account` +- `v1_responses_websocket_should_reuse_connection_for_recorded_conversation` + +### 事件日志入口 + +请求链路和管理端共享 `LogService`。写入事件日志时统一执行： + +- `enabled` +- `capacity` +- `capture_body` +- metadata 清理 + +对应测试： + +- `v1_responses_should_skip_event_log_when_logging_disabled` +- `log_service_should_trim_to_capacity_after_record` +- `log_service_should_remove_body_metadata_when_capture_body_disabled` + +### schema 约束 + +SQLite schema 已补充状态枚举、日志级别、布尔字段、非负计数、HTTP 状态码和 `limit_window_seconds` 约束。 + +对应测试： + +- `sqlite_schema_should_reject_invalid_account_status` +- `sqlite_schema_should_reject_invalid_event_log_level` +- `sqlite_schema_should_reject_non_boolean_flags` +- `sqlite_schema_should_reject_negative_account_usage_counts` + +### cookie 过期语义 + +`CookieRepository::cookie_header` 读取时过滤已过期 cookie，支持标准 `Expires` 时间和 RFC3339。无法解析的历史值会保守保留，避免误删仍可能有效的 cookie。 + +对应测试： + +- `cookie_repository_should_not_replay_expired_cookies` + +### 直接 SQL 边界 + +直接访问业务表的 SQL 已收敛到 repository： + +- `admin_users` / `admin_sessions`：`AdminAuthRepository` +- `fingerprints`：`FingerprintRepository` +- `session_affinities`：`SessionAffinityRepository` +- `event_logs`：`EventLogRepository`，由 `LogService` 统一入口调用 + +对应测试： + +- `admin_auth_repository_should_create_and_load_default_admin_once` +- `admin_auth_repository_should_create_validate_and_cleanup_sessions` +- `fingerprint_repository_should_upsert_auto_update_record` + +## 当前验证 + +最后一次完整验证命令： + +- `cargo fmt --check` +- `git diff --check` +- `cargo test` +- `cargo clippy --all-targets --all-features --locked -- -D warnings` + +以上命令均已通过。 diff --git a/docs/scheduler-implementation.md b/docs/scheduler-implementation.md index df9e4da5f..caaf8b055 100644 --- a/docs/scheduler-implementation.md +++ b/docs/scheduler-implementation.md @@ -1,193 +1,57 @@ -# 后台调度器实现完成 +# 后台调度器实现 -## 概述 +## 当前边界 -本次实现添加了两个关键的后台调度器，用于自动化账户令牌刷新和会话清理。这些调度器是 TypeScript 原版 `codex-proxy` 项目中核心功能的 Rust 原生重写。 +后台任务由 `src/runtime/tasks/coordinator.rs` 统一启动和关闭。运行时只保存任务句柄，具体业务逻辑留在各自领域模块： -## 已实现的功能 +- `src/codex/tasks/refresh.rs`：OAuth access token 刷新调度。 +- `src/codex/tasks/quota.rs`：quota 锁定账户的主动刷新。 +- `src/codex/tasks/model.rs`：Codex 后端模型列表刷新。 +- `src/admin/tasks/session_cleanup.rs`：管理员 session 过期清理。 +- `src/codex/gateway/fingerprint/update_checker.rs`：Codex Desktop fingerprint 版本轮询。 -### 1. RefreshScheduler - OAuth 刷新调度器 +调度器共享 `SchedulerHandle`，应用退出时由 `BackgroundTaskCoordinator::shutdown()` 逐个关闭。 -**位置**: `src/scheduler/refresh.rs` +## RefreshScheduler -**功能**: -- ✅ 在 JWT 令牌 `exp - margin` 时刻自动调度刷新 -- ✅ 指数退避重试策略（5次尝试：5s → 15s → 45s → 135s → 300s） -- ✅ 永久失败检测（识别 `invalid_grant`、`invalid_token`、账户封禁等错误） -- ✅ 临时失败的恢复调度（10分钟后重试） -- ✅ 崩溃恢复机制： - - `refreshing` 状态 → 立即重试 - - `expired` + 有 refreshToken → 延迟重试（30秒起，每个账户间隔2秒） -- ✅ 并发控制（防止同一账户重复刷新） -- ✅ JWT 解析和过期时间计算 -- ✅ 优雅关闭支持 +`RefreshScheduler` 对齐原版刷新调度的关键行为： -**关键实现细节**: -```rust -// 永久错误检测 -const BAN_ERRORS: &[&str] = &["account has been deactivated", "refresh_token_reused"]; -const EXPIRED_ERRORS: &[&str] = &[ - "invalid_grant", - "invalid_token", - "access_denied", - "refresh_token_expired", -]; +- 按 JWT `exp - refresh_margin_seconds` 调度刷新。 +- 使用指数退避，最多 5 次：5s、15s、45s、135s、300s。 +- 区分封禁类错误和 refresh token 失效类错误。 +- 进程重启后恢复 `refreshing` 和带 refresh token 的 `expired` 账户。 +- 用 per-account `in_flight` 抑制重复刷新。 +- 用 `auth.refresh_concurrency` 通过 `Semaphore` 限制全局刷新并发。 -// 指数退避算法 -let backoff = BASE_DELAY_MS.saturating_mul(3_u64.pow(attempt - 1)); -let backoff = backoff.min(300_000); // 最大5分钟 -``` +这部分的行为测试集中在 `tests/codex_accounts/refresh.rs`，包括刷新成功、401 触发刷新、错误状态映射和并发上限。 -**使用方式**: -```rust -let scheduler = RefreshScheduler::new(account_service, config); -let handle = scheduler.start().await; +## Account Scheduling -// 优雅关闭 -handle.shutdown().await; -``` +账号池调度逻辑在 `src/codex/accounts/pool.rs`： -### 2. SessionCleanupScheduler - 会话清理调度器 +- `least_used` 先降权 quota-limited 账号。 +- 再按 `window_reset_at` 排序，优先选择更早进入下一窗口的账号。 +- 再比较 `request_count`。 +- 最后在并列账号之间轮换，避免固定选择同一个账号。 -**位置**: `src/scheduler/session_cleanup.rs` +调度所需的 `request_count`、窗口 token 计数、`window_started_at`、`window_reset_at` 和 `limit_window_seconds` 已持久化到 `account_usage`。重启后 `AccountRepository::list_pool_accounts()` 会恢复这些字段，避免短期调度状态丢失。 -**功能**: -- ✅ 定期删除过期的管理员会话 -- ✅ 可配置的清理间隔 -- ✅ 优雅关闭支持 -- ✅ 直接操作 SQLite 数据库（高效、轻量级） +## Storage Boundaries -**实现**: -```rust -async fn cleanup_expired_sessions(&self) -> Result { - let now = chrono::Utc::now().to_rfc3339(); - let result = sqlx::query("DELETE FROM admin_sessions WHERE expires_at < ?") - .bind(&now) - .execute(&self.db) - .await?; - Ok(result.rows_affected()) -} -``` +后台任务不直接散落业务表 SQL： -### 3. 调度器基础设施 +- 默认管理员创建、管理员登录 session、session 校验和过期清理由 `AdminAuthRepository` 访问 `admin_users` / `admin_sessions`。 +- fingerprint 历史写入和 `auto_updated` 当前记录由 `FingerprintRepository` 访问 `fingerprints`。 +- session affinity 由 `SessionAffinityRepository` 访问 `session_affinities`。 +- 请求事件日志由 `LogService::record` 统一执行开关、容量和 body capture 策略。 -**位置**: `src/scheduler/types.rs`, `src/scheduler/mod.rs` +## Verification -**提供**: -- ✅ `SchedulerHandle` - 统一的调度器句柄，支持优雅关闭 -- ✅ `SchedulerError` - 类型化的错误处理 -- ✅ `SchedulerResult` - 统一的结果类型 +当前实现已通过： -## 与原 TypeScript 版本的对比 +- `cargo fmt --check` +- `git diff --check` +- `cargo test` +- `cargo clippy --all-targets --all-features --locked -- -D warnings` -### TypeScript 版本（参考） -- 使用 `setTimeout` 和 `setInterval` -- 回调风格的异步处理 -- 内存锁文件用于跨进程同步 - -### Rust 版本（新实现） -- 使用 `tokio::spawn` 和 `tokio::time` -- `async/await` 风格 -- `RwLock` 用于线程安全的状态管理 -- 类型安全的错误处理 - -## 架构设计决策 - -### 1. 避免异步递归 -问题：`schedule_one` 调用 `do_refresh`，`do_refresh` 又可能调用 `schedule_one`，导致无限递归。 - -解决方案： -- `schedule_one()` - 初始调度，直接在定时器中执行刷新 -- `schedule_next_refresh()` - 刷新成功后调度下次刷新，避免递归 - -### 2. 并发控制 -使用 `in_flight: Arc>>` 跟踪正在刷新的账户，防止重复刷新。 - -### 3. 优雅关闭 -所有调度器返回 `SchedulerHandle`，通过 `mpsc::channel` 实现优雅关闭： -```rust -tokio::select! { - _ = ticker.tick() => { /* 执行任务 */ } - _ = shutdown_rx.recv() => { /* 退出循环 */ } -} -``` - -## 集成到 AccountService - -新增方法： -```rust -impl AccountService { - /// 列出所有账户用于刷新调度器 - pub async fn list_all_for_refresh(&self) -> Result, AccountServiceError> { - self.repository()? - .list_all() - .await - .map_err(|_| AccountServiceError::List) - } -} - -// 为 RefreshAccountError 实现 Display trait -impl std::fmt::Display for RefreshAccountError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // ... 格式化输出 - } -} -``` - -## 测试状态 - -✅ **编译通过**: `cargo check --lib` -✅ **Clippy 通过**: `cargo clippy --all-targets --all-features --locked -- -D warnings` -✅ **格式化**: `cargo fmt` -✅ **现有测试通过**: 所有 152 个集成测试保持通过 - -## 下一步 - -根据计划文档 (line 3002, 3005)，仍需实现： - -1. **启动集成**: - - [ ] 在 `main.rs` 启动时初始化调度器 - - [ ] 在优雅关闭时停止调度器 - -2. **高级功能**（可选）: - - [ ] 跨进程刷新锁（用于多实例部署） - - [ ] 磁盘刷新令牌同步（防止一次性 RT 被多次消费） - - [ ] 配额刷新调度器 - - [ ] 模型刷新调度器 - - [ ] 指纹更新调度器 - -3. **监控和可观测性**: - - [ ] 调度器健康检查端点 - - [ ] 刷新成功/失败指标 - - [ ] 活跃定时器数量监控 - -## 提交信息 - -``` -feat: add background refresh and session cleanup schedulers - -- 实现 RefreshScheduler 用于在 JWT 过期前自动刷新访问令牌 -- 在 exp - margin 时刻调度刷新 -- 指数退避重试（5次尝试：5s → 15s → 45s → 135s → 300s） -- 永久失败检测（invalid_grant / invalid_token） -- 临时失败的恢复调度（10分钟） -- 崩溃恢复：refreshing → 立即重试，expired + refreshToken → 延迟重试 -- 实现 SessionCleanupScheduler 用于定期清理过期的管理员会话 -- 所有调度器支持优雅关闭 - -Co-Authored-By: Claude Opus 4.8 -``` - -## 文件清单 - -新增文件： -- `src/scheduler/mod.rs` - 模块导出 -- `src/scheduler/types.rs` - 调度器类型和错误定义 -- `src/scheduler/refresh.rs` - OAuth 刷新调度器（385 行） -- `src/scheduler/session_cleanup.rs` - 会话清理调度器（66 行） - -修改文件： -- `src/lib.rs` - 添加 `pub mod scheduler;` -- `src/codex/accounts/service/mod.rs` - 添加 `list_all_for_refresh()` 和 `Display` impl - -总计：554 行新代码 +本文件只描述当前 Rust 实现，不再保留旧 `src/scheduler/*` 路径或早期直接 SQL 方案。 diff --git a/src/admin/auth/mod.rs b/src/admin/auth/mod.rs index 5139f7cfd..b901e5d1b 100644 --- a/src/admin/auth/mod.rs +++ b/src/admin/auth/mod.rs @@ -1,2 +1,3 @@ pub mod api_key; +pub mod repository; pub mod service; diff --git a/src/admin/auth/repository.rs b/src/admin/auth/repository.rs new file mode 100644 index 000000000..2de2cb827 --- /dev/null +++ b/src/admin/auth/repository.rs @@ -0,0 +1,94 @@ +use chrono::{DateTime, Utc}; +use sqlx::{Row, SqlitePool}; +use uuid::Uuid; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminUserRecord { + pub id: String, + pub password_hash: String, +} + +#[derive(Clone)] +pub struct AdminAuthRepository { + pool: SqlitePool, +} + +impl AdminAuthRepository { + pub fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + pub async fn ensure_default_admin(&self, password_hash: &str) -> Result { + let mut tx = self.pool.begin().await?; + let count: (i64,) = sqlx::query_as("select count(*) from admin_users") + .fetch_one(&mut *tx) + .await?; + if count.0 > 0 { + tx.commit().await?; + return Ok(false); + } + + let admin_id = format!("admin_{}", Uuid::new_v4().simple()); + let now = Utc::now().to_rfc3339(); + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + ) + .bind(admin_id) + .bind(password_hash) + .bind(&now) + .bind(&now) + .execute(&mut *tx) + .await?; + tx.commit().await?; + Ok(true) + } + + pub async fn load_first_admin(&self) -> Result, sqlx::Error> { + let row = sqlx::query( + "select id, password_hash from admin_users order by created_at asc limit 1", + ) + .fetch_optional(&self.pool) + .await?; + Ok(row.map(|row| AdminUserRecord { + id: row.get("id"), + password_hash: row.get("password_hash"), + })) + } + + pub async fn create_session( + &self, + session_id: &str, + user_id: &str, + expires_at: DateTime, + ) -> Result<(), sqlx::Error> { + let now = Utc::now().to_rfc3339(); + sqlx::query( + "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", + ) + .bind(session_id) + .bind(user_id) + .bind(expires_at.to_rfc3339()) + .bind(now) + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn validate_session(&self, session_id: &str) -> Result { + let now = Utc::now().to_rfc3339(); + let row = sqlx::query("select 1 from admin_sessions where id = ? and expires_at > ?") + .bind(session_id) + .bind(now) + .fetch_optional(&self.pool) + .await?; + Ok(row.is_some()) + } + + pub async fn cleanup_expired_sessions(&self, now: DateTime) -> Result { + let result = sqlx::query("delete from admin_sessions where expires_at < ?") + .bind(now.to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } +} diff --git a/src/admin/auth/service.rs b/src/admin/auth/service.rs index a8cb3c245..664cf9b6c 100644 --- a/src/admin/auth/service.rs +++ b/src/admin/auth/service.rs @@ -1,11 +1,11 @@ use std::sync::Arc; use chrono::{DateTime, Duration, Utc}; -use sqlx::{Row, SqlitePool}; use tokio::sync::Mutex; use uuid::Uuid; use crate::{ + admin::auth::repository::AdminAuthRepository, codex::accounts::service::{AccountService, ValidatedAccountImportError}, codex::accounts::{ model::AccountStatus, @@ -22,7 +22,7 @@ use crate::{ #[derive(Clone)] pub struct AdminAuthService { config: Arc, - db: Option, + auth_repository: Option, account_repository: Option, account_pool: Arc>, oauth_client: Option>, @@ -122,16 +122,10 @@ pub enum AdminAuthPkceExchange { AlreadyCompleted, } -#[derive(Debug)] -struct AdminUserRow { - id: String, - password_hash: String, -} - impl AdminAuthService { pub fn new( config: Arc, - db: Option, + auth_repository: Option, account_repository: Option, account_pool: Arc>, oauth_client: Option>, @@ -140,7 +134,7 @@ impl AdminAuthService { ) -> Self { Self { config, - db, + auth_repository, account_repository, account_pool, oauth_client, @@ -150,8 +144,9 @@ impl AdminAuthService { } pub async fn login(&self, password: &str) -> Result { - let pool = self.pool()?; - let admin = load_first_admin(pool) + let auth_repository = self.auth_repository()?; + let admin = auth_repository + .load_first_admin() .await .map_err(|_| AdminAuthServiceError::LoadAdminUser)? .ok_or(AdminAuthServiceError::AdminPasswordInvalid)?; @@ -167,7 +162,8 @@ impl AdminAuthService { i64::try_from(ttl_minutes).map_err(|_| AdminAuthServiceError::InvalidSessionTtl)?; let expires_at = Utc::now() + Duration::minutes(ttl_minutes_i64); let session_id = format!("sess_{}", Uuid::new_v4().simple()); - create_admin_session(pool, &session_id, &admin.id, expires_at) + auth_repository + .create_session(&session_id, &admin.id, expires_at) .await .map_err(|_| AdminAuthServiceError::CreateSession)?; @@ -185,10 +181,11 @@ impl AdminAuthService { let Some(session_id) = session_id else { return Ok(false); }; - let pool = self - .pool() + let auth_repository = self + .auth_repository() .map_err(|_| AdminSessionValidationError::DatabaseUnavailable)?; - validate_admin_session(pool, session_id) + auth_repository + .validate_session(session_id) .await .map_err(|_| AdminSessionValidationError::ValidateSession) } @@ -291,8 +288,8 @@ impl AdminAuthService { .await } - fn pool(&self) -> Result<&SqlitePool, AdminAuthServiceError> { - self.db + fn auth_repository(&self) -> Result<&AdminAuthRepository, AdminAuthServiceError> { + self.auth_repository .as_ref() .ok_or(AdminAuthServiceError::DatabaseUnavailable) } @@ -359,47 +356,6 @@ impl AdminAuthService { } } -async fn load_first_admin(pool: &SqlitePool) -> Result, sqlx::Error> { - let row = - sqlx::query("select id, password_hash from admin_users order by created_at asc limit 1") - .fetch_optional(pool) - .await?; - Ok(row.map(|row| AdminUserRow { - id: row.get("id"), - password_hash: row.get("password_hash"), - })) -} - -async fn validate_admin_session(pool: &SqlitePool, session_id: &str) -> Result { - let now = Utc::now().to_rfc3339(); - let count: (i64,) = - sqlx::query_as("select count(*) from admin_sessions where id = ? and expires_at > ?") - .bind(session_id) - .bind(now) - .fetch_one(pool) - .await?; - Ok(count.0 > 0) -} - -async fn create_admin_session( - pool: &SqlitePool, - session_id: &str, - user_id: &str, - expires_at: DateTime, -) -> Result<(), sqlx::Error> { - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind(user_id) - .bind(expires_at.to_rfc3339()) - .bind(now) - .execute(pool) - .await?; - Ok(()) -} - fn account_auth_user(account: &StoredAccountMetadata) -> AdminAuthUser { AdminAuthUser { id: account.id.clone(), diff --git a/src/admin/http/logs.rs b/src/admin/http/logs.rs index f053e5935..fda072f77 100644 --- a/src/admin/http/logs.rs +++ b/src/admin/http/logs.rs @@ -231,6 +231,12 @@ fn log_service_error(error: LogServiceError, request_id: String) -> AdminError { "Failed to clear event logs", request_id, ), + LogServiceError::Write => AdminError::new( + StatusCode::INTERNAL_SERVER_ERROR, + 50001, + "Failed to write event log", + request_id, + ), LogServiceError::InvalidCapacity => AdminError::new( StatusCode::UNPROCESSABLE_ENTITY, 42201, diff --git a/src/admin/tasks/session_cleanup.rs b/src/admin/tasks/session_cleanup.rs index f8865a61c..c362b8f04 100644 --- a/src/admin/tasks/session_cleanup.rs +++ b/src/admin/tasks/session_cleanup.rs @@ -1,17 +1,22 @@ +use chrono::Utc; use tokio::time::{interval, Duration}; use tracing::{error, info}; +use crate::admin::auth::repository::AdminAuthRepository; use crate::runtime::tasks::types::SchedulerHandle; /// 会话清理调度器 - 定期删除过期的管理员会话 pub struct SessionCleanupScheduler { - db: sqlx::SqlitePool, + repository: AdminAuthRepository, interval_secs: u64, } impl SessionCleanupScheduler { - pub fn new(db: sqlx::SqlitePool, interval_secs: u64) -> Self { - Self { db, interval_secs } + pub fn new(repository: AdminAuthRepository, interval_secs: u64) -> Self { + Self { + repository, + interval_secs, + } } /// 启动会话清理调度器 @@ -49,11 +54,6 @@ impl SessionCleanupScheduler { } async fn cleanup_expired_sessions(&self) -> Result { - let now = chrono::Utc::now().to_rfc3339(); - let result = sqlx::query("DELETE FROM admin_sessions WHERE expires_at < ?") - .bind(&now) - .execute(&self.db) - .await?; - Ok(result.rows_affected()) + self.repository.cleanup_expired_sessions(Utc::now()).await } } diff --git a/src/codex/accounts/cookies/repository.rs b/src/codex/accounts/cookies/repository.rs index 9bdff24b3..e7a6ca327 100644 --- a/src/codex/accounts/cookies/repository.rs +++ b/src/codex/accounts/cookies/repository.rs @@ -1,4 +1,4 @@ -use chrono::Utc; +use chrono::{DateTime, Utc}; use secrecy::{ExposeSecret, SecretString}; use sqlx::{Row, SqlitePool}; use thiserror::Error; @@ -70,17 +70,21 @@ impl CookieRepository { request_domain: &str, ) -> CookieResult> { let rows = sqlx::query( - "select domain, name, value_cipher from account_cookies where account_id = ? order by name asc", + "select domain, name, value_cipher, expires_at from account_cookies where account_id = ? order by name asc", ) .bind(account_id) .fetch_all(&self.pool) .await?; let mut pairs = Vec::new(); + let now = Utc::now(); for row in rows { let domain = row.get::("domain"); if !domain_matches(request_domain, &domain) { continue; } + if cookie_is_expired(row.get::, _>("expires_at").as_deref(), now) { + continue; + } let name = row.get::("name"); let value_cipher = row.get::("value_cipher"); let value = self.secret_box.decrypt(&value_cipher)?; @@ -161,3 +165,16 @@ fn domain_matches(request_domain: &str, cookie_domain: &str) -> bool { .strip_suffix(cookie_domain) .is_some_and(|prefix| prefix.ends_with('.')) } + +fn cookie_is_expired(expires_at: Option<&str>, now: DateTime) -> bool { + expires_at + .and_then(parse_cookie_expires_at) + .is_some_and(|expires_at| expires_at <= now) +} + +fn parse_cookie_expires_at(value: &str) -> Option> { + DateTime::parse_from_rfc2822(value) + .or_else(|_| DateTime::parse_from_rfc3339(value)) + .map(|expires_at| expires_at.with_timezone(&Utc)) + .ok() +} diff --git a/src/codex/accounts/repository.rs b/src/codex/accounts/repository.rs index 8ba935e55..959d9619c 100644 --- a/src/codex/accounts/repository.rs +++ b/src/codex/accounts/repository.rs @@ -104,6 +104,13 @@ pub struct AccountUsageRecord { pub input_tokens: i64, pub output_tokens: i64, pub cached_tokens: i64, + pub window_request_count: i64, + pub window_input_tokens: i64, + pub window_output_tokens: i64, + pub window_cached_tokens: i64, + pub window_started_at: Option>, + pub window_reset_at: Option>, + pub limit_window_seconds: Option, pub last_used_at: Option>, } @@ -145,6 +152,148 @@ pub struct AccountRepository { secret_box: SecretBox, } +const LIST_POOL_ACCOUNTS_SQL: &str = r" +select + accounts.id, + email, + accounts.account_id, + user_id, + label, + plan_type, + access_token_cipher, + refresh_token_cipher, + access_token_expires_at, + status, + added_at, + updated_at, + quota_limit_reached, + quota_cooldown_until, + cloudflare_cooldown_until, + coalesce(account_usage.request_count, 0) as usage_request_count, + coalesce(account_usage.empty_response_count, 0) as usage_empty_response_count, + coalesce(account_usage.window_request_count, 0) as usage_window_request_count, + coalesce(account_usage.window_input_tokens, 0) as usage_window_input_tokens, + coalesce(account_usage.window_output_tokens, 0) as usage_window_output_tokens, + coalesce(account_usage.window_cached_tokens, 0) as usage_window_cached_tokens, + account_usage.window_started_at as usage_window_started_at, + account_usage.window_reset_at as usage_window_reset_at, + account_usage.limit_window_seconds as usage_limit_window_seconds, + account_usage.last_used_at as usage_last_used_at +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 } @@ -325,11 +474,9 @@ impl AccountRepository { } pub async fn list_pool_accounts(&self) -> AccountRepositoryResult> { - let rows = sqlx::query( - "select accounts.id, email, accounts.account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher, access_token_expires_at, status, added_at, updated_at, quota_limit_reached, quota_cooldown_until, cloudflare_cooldown_until, coalesce(account_usage.request_count, 0) as usage_request_count, coalesce(account_usage.empty_response_count, 0) as usage_empty_response_count, account_usage.last_used_at as usage_last_used_at from accounts left join account_usage on account_usage.account_id = accounts.id order by added_at desc, accounts.id desc", - ) - .fetch_all(&self.pool) - .await?; + let rows = sqlx::query(LIST_POOL_ACCOUNTS_SQL) + .fetch_all(&self.pool) + .await?; let mut accounts = Vec::with_capacity(rows.len()); for row in rows { let stored = self.account_from_row(&row)?; @@ -358,13 +505,27 @@ impl AccountRepository { )?, request_count: row.get::("usage_request_count").max(0) as u64, empty_response_count: row.get::("usage_empty_response_count").max(0) as u64, - window_request_count: 0, - window_input_tokens: 0, - window_output_tokens: 0, - window_cached_tokens: 0, - window_started_at: None, - window_reset_at: None, - limit_window_seconds: None, + window_request_count: nonnegative_i64_to_u64( + row.get::("usage_window_request_count"), + ), + window_input_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_input_tokens"), + ), + window_output_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_output_tokens"), + ), + window_cached_tokens: nonnegative_i64_to_u64( + row.get::("usage_window_cached_tokens"), + ), + window_started_at: parse_optional_rfc3339( + row.get::, _>("usage_window_started_at"), + )?, + window_reset_at: parse_optional_rfc3339( + row.get::, _>("usage_window_reset_at"), + )?, + limit_window_seconds: optional_positive_i64_to_u64( + row.get::, _>("usage_limit_window_seconds"), + ), added_at: stored.added_at.to_rfc3339(), last_used_at: row.get("usage_last_used_at"), }); @@ -518,17 +679,18 @@ impl AccountRepository { usage: UsageDelta, ) -> AccountRepositoryResult<()> { let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 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, last_used_at = excluded.last_used_at", - ) - .bind(account_id) - .bind(usage.empty_response_count) - .bind(usage.input_tokens) - .bind(usage.output_tokens) - .bind(usage.cached_tokens) - .bind(now) - .execute(&self.pool) - .await?; + 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(()) } @@ -536,13 +698,58 @@ impl AccountRepository { &self, account_id: &str, ) -> AccountRepositoryResult> { + 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, + limit_window_seconds: Option, + ) -> AccountRepositoryResult<()> { let row = sqlx::query( - "select account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at from account_usage where account_id = ?", + "select window_reset_at, limit_window_seconds from account_usage where account_id = ?", ) .bind(account_id) .fetch_optional(&self.pool) .await?; - row.map(|row| usage_from_row(&row)).transpose() + let existing_reset_at = row + .as_ref() + .map(|row| parse_optional_rfc3339(row.get::, _>("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( @@ -702,6 +909,47 @@ fn parse_rfc3339(value: &str) -> AccountRepositoryResult> { Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)) } +fn nonnegative_i64_to_u64(value: i64) -> u64 { + value.max(0) as u64 +} + +fn optional_positive_i64_to_u64(value: Option) -> Option { + value + .and_then(|value| u64::try_from(value).ok()) + .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>, + existing_limit_window_seconds: Option, + new_reset_at: DateTime, + new_limit_window_seconds: Option, +) -> 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 { Ok(AccountUsageRecord { account_id: row.get("account_id"), @@ -710,6 +958,15 @@ fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> AccountRepositoryResult, _>("window_started_at"), + )?, + window_reset_at: parse_optional_rfc3339(row.get::, _>("window_reset_at"))?, + limit_window_seconds: optional_positive_i64_to_u64(row.get("limit_window_seconds")), last_used_at: parse_optional_rfc3339(row.get::, _>("last_used_at"))?, }) } @@ -801,12 +1058,10 @@ impl AccountUsageRepository { } pub async fn reset_account(&self, account_id: &str) -> AccountRepositoryResult<()> { - sqlx::query( - "insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 0, 0, 0, 0, 0, null) on conflict(account_id) do update set request_count = 0, empty_response_count = 0, input_tokens = 0, output_tokens = 0, cached_tokens = 0, last_used_at = null", - ) - .bind(account_id) - .execute(&self.pool) - .await?; + sqlx::query(RESET_ACCOUNT_USAGE_SQL) + .bind(account_id) + .execute(&self.pool) + .await?; Ok(()) } } diff --git a/src/codex/gateway/fingerprint/repository.rs b/src/codex/gateway/fingerprint/repository.rs index f20a56df6..53c93db8d 100644 --- a/src/codex/gateway/fingerprint/repository.rs +++ b/src/codex/gateway/fingerprint/repository.rs @@ -7,6 +7,13 @@ use crate::codex::gateway::fingerprint::{ updater::{FingerprintUpdate, CODEX_DESKTOP_UPDATE_SOURCE}, }; +const AUTO_UPDATED_FINGERPRINT_ID: &str = "auto_updated"; +const AUTO_UPDATE_SOURCE: &str = "auto_update"; +const AUTO_UPDATE_PLATFORM: &str = "darwin"; +const AUTO_UPDATE_ARCH: &str = "arm64"; +const AUTO_UPDATE_USER_AGENT_TEMPLATE: &str = "Codex Desktop/{app_version} ({platform}; {arch})"; +const DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION: &str = "146"; + #[derive(Debug, Clone, PartialEq, Eq)] pub struct StoredFingerprint { pub app_version: String, @@ -45,6 +52,30 @@ impl FingerprintRepository { Ok(()) } + pub async fn upsert_auto_update( + &self, + app_version: &str, + build_number: &str, + chromium_version: Option<&str>, + ) -> Result<(), sqlx::Error> { + sqlx::query( + "insert into fingerprints (id, app_version, build_number, platform, arch, chromium_version, user_agent_template, source, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?) \ + on conflict(id) do update set app_version = excluded.app_version, build_number = excluded.build_number, platform = excluded.platform, arch = excluded.arch, chromium_version = excluded.chromium_version, user_agent_template = excluded.user_agent_template, source = excluded.source, created_at = excluded.created_at", + ) + .bind(AUTO_UPDATED_FINGERPRINT_ID) + .bind(app_version) + .bind(build_number) + .bind(AUTO_UPDATE_PLATFORM) + .bind(AUTO_UPDATE_ARCH) + .bind(chromium_version.unwrap_or(DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION)) + .bind(AUTO_UPDATE_USER_AGENT_TEMPLATE) + .bind(AUTO_UPDATE_SOURCE) + .bind(Utc::now().to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(()) + } + pub async fn latest(&self) -> Result, sqlx::Error> { let row = sqlx::query( "select app_version, build_number, source from fingerprints order by created_at desc, id desc limit 1", @@ -63,11 +94,12 @@ impl FingerprintRepository { r#" select app_version, build_number, platform, arch, chromium_version, user_agent_template from fingerprints - where source = 'auto_update' + where source = ? order by created_at desc limit 1 "#, ) + .bind(AUTO_UPDATE_SOURCE) .fetch_optional(&self.pool) .await?; diff --git a/src/codex/gateway/fingerprint/update_checker.rs b/src/codex/gateway/fingerprint/update_checker.rs index 713ea85fc..aa7b1e5a2 100644 --- a/src/codex/gateway/fingerprint/update_checker.rs +++ b/src/codex/gateway/fingerprint/update_checker.rs @@ -8,6 +8,7 @@ use tokio::sync::Mutex; use tokio::time::interval; use tracing::{info, warn}; +use crate::codex::gateway::fingerprint::repository::FingerprintRepository; use crate::platform::storage::paths; const APPCAST_URL: &str = "https://persistent.oaistatic.com/codex-app-prod/appcast.xml"; @@ -27,7 +28,7 @@ pub struct UpdateState { #[derive(Clone)] pub struct UpdateChecker { - db: Option, + repository: Option, state: Arc>, } @@ -39,7 +40,7 @@ struct InternalState { impl UpdateChecker { pub fn new(db: Option, current_version: String, current_build: String) -> Self { Self { - db, + repository: db.map(FingerprintRepository::new), state: Arc::new(Mutex::new(InternalState { current_version, current_build, @@ -140,8 +141,10 @@ impl UpdateChecker { async fn apply_version_update(&self, version: &str, build: &str) -> Result<(), UpdateError> { let chromium_version = load_matching_chromium_version(version, build); - if let Some(db) = &self.db { - update_fingerprint_in_db(db, version, build, chromium_version.as_deref()).await?; + if let Some(repository) = &self.repository { + repository + .upsert_auto_update(version, build, chromium_version.as_deref()) + .await?; } Ok(()) @@ -226,36 +229,6 @@ fn load_matching_chromium_version(version: &str, build: &str) -> Option } } -async fn update_fingerprint_in_db( - db: &SqlitePool, - version: &str, - build: &str, - chromium_version: Option<&str>, -) -> Result<(), UpdateError> { - sqlx::query( - r#" - insert into fingerprints ( - id, app_version, build_number, platform, arch, - chromium_version, user_agent_template, source, created_at - ) values (?, ?, ?, 'darwin', 'arm64', ?, ?, 'auto_update', ?) - on conflict(id) do update set - app_version = excluded.app_version, - build_number = excluded.build_number, - chromium_version = excluded.chromium_version - "#, - ) - .bind("auto_updated") - .bind(version) - .bind(build) - .bind(chromium_version.unwrap_or("146")) - .bind("Codex Desktop/{app_version} ({platform}; {arch})") - .bind(Utc::now().to_rfc3339()) - .execute(db) - .await?; - - Ok(()) -} - #[derive(Debug, thiserror::Error)] pub enum UpdateError { #[error("HTTP 请求失败: {0}")] diff --git a/src/codex/logs/repository.rs b/src/codex/logs/repository.rs index cfff12931..22cb18916 100644 --- a/src/codex/logs/repository.rs +++ b/src/codex/logs/repository.rs @@ -49,6 +49,18 @@ impl EventLogRepository { Ok(()) } + pub async fn trim_to_capacity(&self, capacity: u32) -> Result { + sqlx::query( + "delete from event_logs where id in ( + select id from event_logs order by created_at desc, id desc limit -1 offset ? + )", + ) + .bind(i64::from(capacity)) + .execute(&self.pool) + .await + .map(|result| result.rows_affected()) + } + pub async fn list( &self, cursor: Option, diff --git a/src/codex/logs/service.rs b/src/codex/logs/service.rs index 00e2aa9f2..298c2d89c 100644 --- a/src/codex/logs/service.rs +++ b/src/codex/logs/service.rs @@ -17,13 +17,21 @@ pub struct LogService { repository: Option, } -#[derive(Debug)] +#[derive(Debug, thiserror::Error)] pub enum LogServiceError { + #[error("事件日志仓储未初始化")] RepositoryUnavailable, + #[error("查询事件日志失败")] List, + #[error("读取事件日志失败")] Get, + #[error("统计事件日志失败")] Count, + #[error("清空事件日志失败")] Clear, + #[error("写入事件日志失败")] + Write, + #[error("日志容量必须大于 0")] InvalidCapacity, } @@ -111,10 +119,34 @@ impl LogService { config.capture_body = capture_body; } } + if let Some(capacity) = update.capacity { + self.repository()? + .trim_to_capacity(capacity) + .await + .map_err(|_| LogServiceError::Write)?; + } self.state().await } + pub async fn record(&self, mut event: EventLog) -> Result { + let config = self.config.read().await.clone(); + if !config.enabled { + return Ok(false); + } + apply_capture_body_policy(&mut event, config.capture_body); + let repository = self.repository()?; + repository + .insert(event) + .await + .map_err(|_| LogServiceError::Write)?; + repository + .trim_to_capacity(config.capacity) + .await + .map_err(|_| LogServiceError::Write)?; + Ok(true) + } + pub async fn get(&self, id: &str) -> Result, LogServiceError> { self.repository()? .get(id) @@ -137,6 +169,24 @@ impl LogService { } } +fn apply_capture_body_policy(event: &mut EventLog, capture_body: bool) { + if capture_body { + return; + } + let Some(metadata) = event.metadata.as_object_mut() else { + return; + }; + for key in [ + "body", + "rawBody", + "requestBody", + "responseBody", + "upstreamBody", + ] { + metadata.remove(key); + } +} + impl From for EventLogFilters { fn from(filter: LogListFilter) -> Self { Self { diff --git a/src/codex/serving/dispatch/affinity.rs b/src/codex/serving/dispatch/affinity.rs index 90c661246..5af6a8890 100644 --- a/src/codex/serving/dispatch/affinity.rs +++ b/src/codex/serving/dispatch/affinity.rs @@ -1,7 +1,11 @@ -use serde::Serialize; use std::collections::HashMap; use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use serde::Serialize; +use sqlx::{Row, SqlitePool}; +use thiserror::Error; use tokio::sync::RwLock; use tokio::time::interval; @@ -23,7 +27,38 @@ pub struct AffinityEntry { /// 变体哈希 - 用于区分并发对话分支 pub variant_hash: Option, /// 创建时间 - pub created_at: Instant, + pub created_at: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StoredSessionAffinity { + pub response_id: String, + pub account_id: String, + pub conversation_id: String, + pub turn_state: Option, + pub instructions_hash: Option, + pub input_tokens: Option, + pub function_call_ids: Vec, + pub variant_hash: Option, + pub expires_at: DateTime, + pub created_at: DateTime, +} + +#[derive(Debug, Error)] +pub enum SessionAffinityRepositoryError { + #[error("database error: {0}")] + Database(#[from] sqlx::Error), + #[error("invalid session affinity timestamp: {0}")] + InvalidTimestamp(#[from] chrono::ParseError), + #[error("invalid session affinity function call ids: {0}")] + InvalidFunctionCallIds(#[from] serde_json::Error), +} + +pub type SessionAffinityRepositoryResult = Result; + +#[derive(Clone)] +pub struct SessionAffinityRepository { + pool: SqlitePool, } /// 会话亲和性映射 - 将 previous_response_id 映射到创建它的账户 @@ -42,12 +77,103 @@ pub struct SessionAffinityMap { const DEFAULT_TTL: Duration = Duration::from_secs(4 * 60 * 60); // 4小时 const CLEANUP_INTERVAL: Duration = Duration::from_secs(10 * 60); // 10分钟 +const UPSERT_SESSION_AFFINITY_SQL: &str = r" +insert into session_affinities ( + response_id, + account_id, + conversation_id, + turn_state, + instructions_hash, + input_tokens, + function_call_ids_json, + variant_hash, + expires_at, + created_at +) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +on conflict(response_id) do update set + account_id = excluded.account_id, + conversation_id = excluded.conversation_id, + turn_state = excluded.turn_state, + instructions_hash = excluded.instructions_hash, + input_tokens = excluded.input_tokens, + function_call_ids_json = excluded.function_call_ids_json, + variant_hash = excluded.variant_hash, + expires_at = excluded.expires_at, + created_at = excluded.created_at"; + +const LIST_ACTIVE_SESSION_AFFINITIES_SQL: &str = r" +select + response_id, + account_id, + conversation_id, + turn_state, + instructions_hash, + input_tokens, + function_call_ids_json, + variant_hash, + expires_at, + created_at +from session_affinities +where expires_at > ? +order by created_at asc, response_id asc"; + pub(super) fn compute_variant_hash(value: &impl Serialize) -> Option { serde_json::to_string(value) .ok() .filter(|serialized| !serialized.is_empty()) } +impl SessionAffinityRepository { + pub fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + pub async fn upsert( + &self, + response_id: &str, + entry: &AffinityEntry, + ttl: Duration, + ) -> SessionAffinityRepositoryResult<()> { + let function_call_ids_json = serde_json::to_string(&entry.function_call_ids)?; + let expires_at = entry.created_at + chrono_duration_from_std(ttl); + sqlx::query(UPSERT_SESSION_AFFINITY_SQL) + .bind(response_id) + .bind(&entry.account_id) + .bind(&entry.conversation_id) + .bind(&entry.turn_state) + .bind(&entry.instructions_hash) + .bind(entry.input_tokens.map(u64_to_i64_saturating)) + .bind(function_call_ids_json) + .bind(&entry.variant_hash) + .bind(expires_at.to_rfc3339()) + .bind(entry.created_at.to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn list_active( + &self, + now: DateTime, + ) -> SessionAffinityRepositoryResult> { + let rows = sqlx::query(LIST_ACTIVE_SESSION_AFFINITIES_SQL) + .bind(now.to_rfc3339()) + .fetch_all(&self.pool) + .await?; + rows.into_iter() + .map(|row| stored_session_affinity_from_row(&row)) + .collect() + } + + pub async fn delete_expired(&self, now: DateTime) -> SessionAffinityRepositoryResult { + let result = sqlx::query("delete from session_affinities where expires_at <= ?") + .bind(now.to_rfc3339()) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } +} + impl SessionAffinityMap { pub fn new(ttl: Duration) -> Self { Self { @@ -74,9 +200,9 @@ impl SessionAffinityMap { let mut ticker = interval(CLEANUP_INTERVAL); loop { ticker.tick().await; - let now = Instant::now(); + let now = Utc::now(); let mut map_lock = map.write().await; - map_lock.retain(|_, entry| now.duration_since(entry.created_at) < ttl); + map_lock.retain(|_, entry| !entry_expired(entry.created_at, ttl, now)); } }); @@ -102,7 +228,7 @@ impl SessionAffinityMap { input_tokens: Option, function_call_ids: Option>, variant_hash: Option, - ) { + ) -> AffinityEntry { let instructions_hash = instructions.map(|s| { use sha2::{Digest, Sha256}; let mut hasher = Sha256::new(); @@ -119,10 +245,41 @@ impl SessionAffinityMap { input_tokens, function_call_ids: function_call_ids.unwrap_or_default(), variant_hash, - created_at: Instant::now(), + created_at: Utc::now(), }; - self.map.write().await.insert(response_id, entry); + self.map.write().await.insert(response_id, entry.clone()); + entry + } + + pub async fn restore(&self, records: Vec) -> usize { + let now = Utc::now(); + let mut restored = 0usize; + let mut map = self.map.write().await; + for record in records { + if record.expires_at <= now || entry_expired(record.created_at, self.ttl, now) { + continue; + } + map.insert( + record.response_id, + AffinityEntry { + account_id: record.account_id, + conversation_id: record.conversation_id, + turn_state: record.turn_state, + instructions_hash: record.instructions_hash, + input_tokens: record.input_tokens, + function_call_ids: record.function_call_ids, + variant_hash: record.variant_hash, + created_at: record.created_at, + }, + ); + restored = restored.saturating_add(1); + } + restored + } + + pub fn ttl(&self) -> Duration { + self.ttl } /// 查找响应ID对应的账户ID @@ -148,10 +305,10 @@ impl SessionAffinityMap { variant_hash: Option<&str>, ) -> Option { let map = self.map.read().await; - let now = Instant::now(); + let now = Utc::now(); let mut latest_response_id: Option = None; - let mut latest_created_at = Instant::now() - Duration::from_secs(u64::MAX / 2); + let mut latest_created_at: Option> = None; for (response_id, entry) in map.iter() { // 过滤对话ID @@ -167,20 +324,20 @@ impl SessionAffinityMap { } // 检查TTL - if now.duration_since(entry.created_at) >= self.ttl { + if entry_expired(entry.created_at, self.ttl, now) { continue; } // 检查最大年龄 if let Some(max) = max_age { - if now.duration_since(entry.created_at) > max { + if entry_age_exceeds(entry.created_at, now, max) { continue; } } // 保留最新的 - if entry.created_at > latest_created_at { - latest_created_at = entry.created_at; + if latest_created_at.is_none_or(|latest| entry.created_at > latest) { + latest_created_at = Some(entry.created_at); latest_response_id = Some(response_id.clone()); } } @@ -235,8 +392,8 @@ impl SessionAffinityMap { let map = self.map.read().await; let entry = map.get(response_id)?; - let now = Instant::now(); - if now.duration_since(entry.created_at) >= self.ttl { + let now = Utc::now(); + if entry_expired(entry.created_at, self.ttl, now) { return None; } @@ -244,6 +401,52 @@ impl SessionAffinityMap { } } +fn entry_expired(created_at: DateTime, ttl: Duration, now: DateTime) -> bool { + now.signed_duration_since(created_at) + .to_std() + .is_ok_and(|age| age >= ttl) +} + +fn entry_age_exceeds(created_at: DateTime, now: DateTime, max_age: Duration) -> bool { + now.signed_duration_since(created_at) + .to_std() + .is_ok_and(|age| age > max_age) +} + +fn chrono_duration_from_std(duration: Duration) -> chrono::Duration { + chrono::Duration::from_std(duration).unwrap_or_else(|_| chrono::Duration::seconds(i64::MAX)) +} + +fn u64_to_i64_saturating(value: u64) -> i64 { + value.min(i64::MAX as u64) as i64 +} + +fn optional_nonnegative_i64_to_u64(value: Option) -> Option { + value.and_then(|value| u64::try_from(value).ok()) +} + +fn parse_rfc3339(value: &str) -> SessionAffinityRepositoryResult> { + Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)) +} + +fn stored_session_affinity_from_row( + row: &sqlx::sqlite::SqliteRow, +) -> SessionAffinityRepositoryResult { + let function_call_ids_json = row.get::("function_call_ids_json"); + Ok(StoredSessionAffinity { + response_id: row.get("response_id"), + account_id: row.get("account_id"), + conversation_id: row.get("conversation_id"), + turn_state: row.get("turn_state"), + instructions_hash: row.get("instructions_hash"), + input_tokens: optional_nonnegative_i64_to_u64(row.get("input_tokens")), + function_call_ids: serde_json::from_str(&function_call_ids_json)?, + variant_hash: row.get("variant_hash"), + expires_at: parse_rfc3339(&row.get::("expires_at"))?, + created_at: parse_rfc3339(&row.get::("created_at"))?, + }) +} + impl Drop for SessionAffinityMap { fn drop(&mut self) { self.stop_cleanup(); diff --git a/src/codex/serving/dispatch/mod.rs b/src/codex/serving/dispatch/mod.rs index afbc2bb25..3865923d3 100644 --- a/src/codex/serving/dispatch/mod.rs +++ b/src/codex/serving/dispatch/mod.rs @@ -47,8 +47,10 @@ use crate::{ CodexWebSocketPool, SharedRateLimitUpdates, }, }, - codex::logs::event::{EventLevel, EventLog}, - codex::logs::repository::EventLogRepository, + codex::logs::{ + event::{EventLevel, EventLog}, + service::LogService, + }, config::AppConfig, }; @@ -62,7 +64,7 @@ pub(crate) use self::{ stream::{completed_response_json, CollectedResponse}, }; -use self::affinity::{compute_variant_hash, SessionAffinityMap}; +use self::affinity::{compute_variant_hash, SessionAffinityMap, SessionAffinityRepository}; use self::{ refresh::refresh_account_after_unauthorized, routing::request_domain, @@ -76,10 +78,11 @@ struct CodexUpstreamDependencies { account_pool: Arc>, account_repository: Option, cookie_repository: Option, - event_logs: Option, + logs: LogService, token_refresher: Option>, fingerprint: Fingerprint, // 用于实际请求的指纹 session_affinity: Arc, + session_affinity_repository: Option, websocket_pool: Arc, } @@ -88,13 +91,18 @@ pub(crate) struct CodexUpstreamService { deps: CodexUpstreamDependencies, } +pub(crate) struct CodexUpstreamRepositories { + pub(crate) account: Option, + pub(crate) cookie: Option, + pub(crate) session_affinity: Option, +} + impl CodexUpstreamService { pub(crate) fn new( config: Arc, account_pool: Arc>, - account_repository: Option, - cookie_repository: Option, - event_logs: Option, + repositories: CodexUpstreamRepositories, + logs: LogService, token_refresher: Option>, fingerprint: Fingerprint, ) -> Self { @@ -102,12 +110,13 @@ impl CodexUpstreamService { deps: CodexUpstreamDependencies { config, account_pool, - account_repository, - cookie_repository, - event_logs, + account_repository: repositories.account, + cookie_repository: repositories.cookie, + logs, token_refresher, fingerprint, session_affinity: Arc::new(SessionAffinityMap::with_default_ttl()), + session_affinity_repository: repositories.session_affinity, websocket_pool: Arc::new(CodexWebSocketPool::with_default_max_age()), }, } @@ -272,6 +281,18 @@ impl CodexUpstreamService { ) .await; } + + pub(crate) async fn reload_session_affinity_from_repository( + &self, + ) -> self::affinity::SessionAffinityRepositoryResult { + let Some(repository) = self.deps.session_affinity_repository.as_ref() else { + return Ok(0); + }; + let now = Utc::now(); + repository.delete_expired(now).await?; + let records = repository.list_active(now).await?; + Ok(self.deps.session_affinity.restore(records).await) + } } async fn responses_http_sse_stream( @@ -704,9 +725,11 @@ async fn record_response_affinity_with_deps( }; let variant_hash = compute_variant_hash(request); - deps.session_affinity + let response_id = metadata.response_id; + let entry = deps + .session_affinity .record( - metadata.response_id, + response_id.clone(), account_id.to_string(), conversation_id.to_string(), turn_state @@ -719,6 +742,19 @@ async fn record_response_affinity_with_deps( variant_hash, ) .await; + if let Some(repository) = deps.session_affinity_repository.as_ref() { + if let Err(error) = repository + .upsert(&response_id, &entry, deps.session_affinity.ttl()) + .await + { + tracing::warn!( + error = %error, + response_id = %response_id, + account_id = %account_id, + "持久化 session affinity 失败" + ); + } + } } async fn apply_rate_limit_headers_with_deps( @@ -751,6 +787,23 @@ async fn apply_rate_limit_headers_with_deps( 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; } @@ -1406,9 +1459,6 @@ async fn log_codex_upstream_response_with_deps( message: &str, mut metadata: Value, ) { - let Some(repo) = deps.event_logs.as_ref() else { - return; - }; ensure_stream_metadata(&mut metadata, context.stream); let mut event = EventLog::new("v1.response", level, message); event.request_id = Some(context.request_id.clone()); @@ -1418,7 +1468,7 @@ async fn log_codex_upstream_response_with_deps( event.status_code = Some(i64::from(status.as_u16())); event.latency_ms = Some(context.latency_ms()); event.metadata = metadata; - if let Err(error) = repo.insert(event).await { + if let Err(error) = deps.logs.record(event).await { tracing::warn!(error = %error, "写入 v1 response 事件日志失败"); } } diff --git a/src/codex/serving/dispatch/refresh.rs b/src/codex/serving/dispatch/refresh.rs index f3bbc3b87..0c654c20e 100644 --- a/src/codex/serving/dispatch/refresh.rs +++ b/src/codex/serving/dispatch/refresh.rs @@ -102,9 +102,6 @@ async fn log_account_refresh_failure( request_id: &str, model: &str, ) { - let Some(repo) = deps.event_logs.as_ref() else { - return; - }; let mut event = EventLog::new( "account.refresh", EventLevel::Warn, @@ -120,7 +117,7 @@ async fn log_account_refresh_failure( "failure": refresh_failure_value(failure), "accountStatus": status.map(account_status_value), }); - if let Err(error) = repo.insert(event).await { + if let Err(error) = deps.logs.record(event).await { tracing::warn!(error = %error, "写入账户刷新事件日志失败"); } } diff --git a/src/codex/serving/responses.rs b/src/codex/serving/responses.rs index d59d3a752..9fa77b941 100644 --- a/src/codex/serving/responses.rs +++ b/src/codex/serving/responses.rs @@ -16,7 +16,8 @@ use crate::{ codex::gateway::transport::types::CodexResponsesRequest, codex::logs::event::EventLevel, codex::serving::dispatch::{ - classify_upstream_account_retry, completed_response_json, no_available_accounts_response, + affinity::SessionAffinityRepositoryResult, classify_upstream_account_retry, + completed_response_json, no_available_accounts_response, normalize_service_tier_for_upstream, websocket_history_retry_metadata, CodexRequestLogContext, CodexUpstreamService, CollectedResponse, }, @@ -51,6 +52,14 @@ impl ResponsesService { self.upstream.fingerprint() } + pub async fn reload_session_affinity_from_repository( + &self, + ) -> SessionAffinityRepositoryResult { + self.upstream + .reload_session_affinity_from_repository() + .await + } + pub async fn handle( &self, request_id: &str, diff --git a/src/platform/storage/db.rs b/src/platform/storage/db.rs index 4cbc509af..aecdc5931 100644 --- a/src/platform/storage/db.rs +++ b/src/platform/storage/db.rs @@ -21,6 +21,7 @@ async fn initialize_schema(pool: &SqlitePool) -> Result<(), sqlx::Error> { let schema = include_str!("schema.sql"); sqlx::raw_sql(schema).execute(pool).await?; ensure_account_usage_schema(pool).await?; + ensure_session_affinity_schema(pool).await?; Ok(()) } @@ -28,19 +29,68 @@ async fn ensure_account_usage_schema(pool: &SqlitePool) -> Result<(), sqlx::Erro let rows = sqlx::query("pragma table_info(account_usage)") .fetch_all(pool) .await?; - let exists = rows - .iter() - .any(|row| row.get::("name") == "empty_response_count"); - if !exists { - sqlx::query( - "alter table account_usage add column empty_response_count integer not null default 0", - ) - .execute(pool) - .await?; + for (name, sql) in ACCOUNT_USAGE_ADDED_COLUMNS { + let exists = rows.iter().any(|row| row.get::("name") == *name); + if !exists { + sqlx::query(*sql).execute(pool).await?; + } } Ok(()) } +const ACCOUNT_USAGE_ADDED_COLUMNS: &[(&str, &str)] = &[ + ( + "empty_response_count", + "alter table account_usage add column empty_response_count integer not null default 0", + ), + ( + "window_request_count", + "alter table account_usage add column window_request_count integer not null default 0", + ), + ( + "window_input_tokens", + "alter table account_usage add column window_input_tokens integer not null default 0", + ), + ( + "window_output_tokens", + "alter table account_usage add column window_output_tokens integer not null default 0", + ), + ( + "window_cached_tokens", + "alter table account_usage add column window_cached_tokens integer not null default 0", + ), + ( + "window_started_at", + "alter table account_usage add column window_started_at text", + ), + ( + "window_reset_at", + "alter table account_usage add column window_reset_at text", + ), + ( + "limit_window_seconds", + "alter table account_usage add column limit_window_seconds integer", + ), +]; + +async fn ensure_session_affinity_schema(pool: &SqlitePool) -> Result<(), sqlx::Error> { + let rows = sqlx::query("pragma table_info(session_affinities)") + .fetch_all(pool) + .await?; + for (name, sql) in SESSION_AFFINITY_ADDED_COLUMNS { + let exists = rows.iter().any(|row| row.get::("name") == *name); + if !exists { + sqlx::query(*sql).execute(pool).await?; + } + } + Ok(()) +} + +const SESSION_AFFINITY_ADDED_COLUMNS: &[(&str, &str)] = &[( + "function_call_ids_json", + "alter table session_affinities add column function_call_ids_json text not null default '[]'", +)]; + fn ensure_sqlite_parent_dir(database_url: &str) -> Result<(), sqlx::Error> { let Some(path) = sqlite_file_path(database_url) else { return Ok(()); diff --git a/src/platform/storage/schema.sql b/src/platform/storage/schema.sql index b632b8bf5..881d21af4 100644 --- a/src/platform/storage/schema.sql +++ b/src/platform/storage/schema.sql @@ -31,7 +31,7 @@ create table if not exists client_api_keys ( prefix text not null, key_hash text not null, label text, - enabled integer not null default 1, + enabled integer not null default 1 check (enabled in (0, 1)), created_at text not null, last_used_at text ); @@ -51,10 +51,10 @@ create table if not exists accounts ( access_token_cipher text not null, refresh_token_cipher text, access_token_expires_at text, - status text not null, + status text not null check (status in ('active', 'expired', 'quota_exhausted', 'refreshing', 'disabled', 'banned')), quota_json text, quota_fetched_at text, - quota_limit_reached integer not null default 0, + quota_limit_reached integer not null default 0 check (quota_limit_reached in (0, 1)), quota_cooldown_until text, cloudflare_cooldown_until text, added_at text not null, @@ -81,11 +81,18 @@ create index if not exists idx_account_refresh_leases_expires on account_refresh -- ============================================ create table if not exists account_usage ( account_id text primary key references accounts(id) on delete cascade, - request_count integer not null default 0, - empty_response_count integer not null default 0, - input_tokens integer not null default 0, - output_tokens integer not null default 0, - cached_tokens integer not null default 0, + request_count integer not null default 0 check (request_count >= 0), + empty_response_count integer not null default 0 check (empty_response_count >= 0), + input_tokens integer not null default 0 check (input_tokens >= 0), + output_tokens integer not null default 0 check (output_tokens >= 0), + cached_tokens integer not null default 0 check (cached_tokens >= 0), + window_request_count integer not null default 0 check (window_request_count >= 0), + window_input_tokens integer not null default 0 check (window_input_tokens >= 0), + window_output_tokens integer not null default 0 check (window_output_tokens >= 0), + window_cached_tokens integer not null default 0 check (window_cached_tokens >= 0), + window_started_at text, + window_reset_at text, + limit_window_seconds integer check (limit_window_seconds is null or limit_window_seconds > 0), last_used_at text ); @@ -129,12 +136,12 @@ create table if not exists event_logs ( id text primary key, request_id text, kind text not null, - level text not null, + level text not null check (level in ('debug', 'info', 'warn', 'error')), account_id text, route text, model text, - status_code integer, - latency_ms integer, + status_code integer check (status_code is null or (status_code >= 100 and status_code <= 599)), + latency_ms integer check (latency_ms is null or latency_ms >= 0), message text not null, metadata_json text not null, created_at text not null @@ -163,7 +170,8 @@ create table if not exists session_affinities ( conversation_id text not null, turn_state text, instructions_hash text, - input_tokens integer, + input_tokens integer check (input_tokens is null or input_tokens >= 0), + function_call_ids_json text not null default '[]', variant_hash text, expires_at text not null, created_at text not null diff --git a/src/runtime/bootstrap.rs b/src/runtime/bootstrap.rs index 1678afda0..dcf0db0a3 100644 --- a/src/runtime/bootstrap.rs +++ b/src/runtime/bootstrap.rs @@ -1,9 +1,11 @@ use thiserror::Error; use crate::{ + admin::auth::repository::AdminAuthRepository, codex::accounts::repository::AccountRepositoryError, codex::gateway::fingerprint::{model::Fingerprint, repository::FingerprintRepository}, codex::gateway::oauth::OpenAiOAuthRefresher, + codex::serving::dispatch::affinity::SessionAffinityRepositoryError, config::AppConfig, platform::{ crypto::{CryptoError, SecretBox}, @@ -27,6 +29,8 @@ pub enum BootstrapError { Database(#[from] sqlx::Error), #[error("account repository error: {0}")] AccountRepository(#[from] AccountRepositoryError), + #[error("session affinity repository error: {0}")] + SessionAffinityRepository(#[from] SessionAffinityRepositoryError), } pub type BootstrapResult = Result; @@ -73,8 +77,8 @@ pub async fn build_state(config: AppConfig) -> BootstrapResult<(AppState, Sqlite } }; - // 初始化默认管理员账号（如果不存在） - ensure_default_admin_exists(&pool, &config).await?; + let admin_auth_repo = AdminAuthRepository::new(pool.clone()); + ensure_default_admin_exists(&admin_auth_repo, &config).await?; let state = AppState::with_pool_secret_api_key_hasher_oauth_client_and_fingerprint( config, @@ -85,36 +89,23 @@ pub async fn build_state(config: AppConfig) -> BootstrapResult<(AppState, Sqlite fingerprint, ); let restored_accounts = state.reload_account_pool_from_repository().await?; + let restored_affinities = state.reload_session_affinity_from_repository().await?; + tracing::info!( + account_count = restored_accounts, + session_affinity_count = restored_affinities, + "已从 SQLite 恢复运行时状态" + ); Ok((state, pool, restored_accounts)) } /// 确保默认管理员账号存在（首次启动时创建） async fn ensure_default_admin_exists( - pool: &SqlitePool, + repository: &AdminAuthRepository, config: &AppConfig, ) -> Result<(), sqlx::Error> { - // 检查是否已存在管理员 - let count: (i64,) = sqlx::query_as("select count(*) from admin_users") - .fetch_one(pool) - .await?; - - if count.0 == 0 { - // 创建默认管理员 - let admin_id = format!("admin_{}", uuid::Uuid::new_v4().simple()); - let password_hash = hash_admin_password(&config.admin.default_password) - .map_err(|e| sqlx::Error::Protocol(format!("failed to hash password: {}", e)))?; - let now = chrono::Utc::now().to_rfc3339(); - - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind(&admin_id) - .bind(&password_hash) - .bind(&now) - .bind(&now) - .execute(pool) - .await?; - + let password_hash = hash_admin_password(&config.admin.default_password) + .map_err(|e| sqlx::Error::Protocol(format!("failed to hash password: {e}")))?; + if repository.ensure_default_admin(&password_hash).await? { tracing::info!( username = %config.admin.default_username, "已创建默认管理员用户" diff --git a/src/runtime/state.rs b/src/runtime/state.rs index 05c79da8a..307f6f775 100644 --- a/src/runtime/state.rs +++ b/src/runtime/state.rs @@ -4,7 +4,7 @@ use sqlx::SqlitePool; use tokio::sync::Mutex; use crate::admin::{ - auth::{api_key::ApiKeyService, service::AdminAuthService}, + auth::{api_key::ApiKeyService, repository::AdminAuthRepository, service::AdminAuthService}, settings::SettingsService, }; use crate::codex::accounts::cookies::repository::CookieRepository; @@ -17,7 +17,10 @@ use crate::codex::accounts::{ }; use crate::codex::gateway::fingerprint::model::Fingerprint; use crate::codex::gateway::oauth::{OAuthClient, PkceSessionStore, TokenRefresher}; -use crate::codex::serving::dispatch::CodexUpstreamService; +use crate::codex::serving::dispatch::affinity::{ + SessionAffinityRepository, SessionAffinityRepositoryResult, +}; +use crate::codex::serving::dispatch::{CodexUpstreamRepositories, CodexUpstreamService}; use crate::codex::serving::{ chat::ChatService, diagnostics::DiagnosticsService, responses::ResponsesService, }; @@ -209,6 +212,15 @@ impl AppState { .await } + pub async fn reload_session_affinity_from_repository( + &self, + ) -> SessionAffinityRepositoryResult { + self.services + .responses + .reload_session_affinity_from_repository() + .await + } + /// 获取 AccountService 的引用（用于调度器） pub fn account_service(&self) -> Arc { Arc::new(self.services.accounts.clone()) @@ -268,6 +280,7 @@ impl AppState { config.clone(), pool_ref, secret_box_ref, + logs.clone(), token_refresher, account_pool, fingerprint, @@ -315,6 +328,7 @@ fn v1_services( config: Arc, pool: Option<&SqlitePool>, secret_box: Option<&SecretBox>, + logs: LogService, token_refresher: Option>, account_pool: Arc>, fingerprint: Fingerprint, @@ -323,6 +337,7 @@ fn v1_services( config.clone(), pool, secret_box, + logs, token_refresher, account_pool.clone(), fingerprint, @@ -361,7 +376,7 @@ fn admin_auth_service( ) -> AdminAuthService { AdminAuthService::new( config, - pool.cloned(), + pool.cloned().map(AdminAuthRepository::new), account_repository(pool, secret_box), account_pool, oauth_client, @@ -392,6 +407,7 @@ fn codex_upstream_service( config: Arc, pool: Option<&SqlitePool>, secret_box: Option<&SecretBox>, + logs: LogService, token_refresher: Option>, account_pool: Arc>, fingerprint: Fingerprint, @@ -399,9 +415,12 @@ fn codex_upstream_service( CodexUpstreamService::new( config, account_pool, - account_repository(pool, secret_box), - cookie_repository(pool, secret_box), - pool.cloned().map(EventLogRepository::new), + CodexUpstreamRepositories { + account: account_repository(pool, secret_box), + cookie: cookie_repository(pool, secret_box), + session_affinity: pool.cloned().map(SessionAffinityRepository::new), + }, + logs, token_refresher, fingerprint, ) diff --git a/src/runtime/tasks/coordinator.rs b/src/runtime/tasks/coordinator.rs index 205621e2c..15fca7fc7 100644 --- a/src/runtime/tasks/coordinator.rs +++ b/src/runtime/tasks/coordinator.rs @@ -1,7 +1,9 @@ use sqlx::SqlitePool; use crate::{ - admin::tasks::session_cleanup::SessionCleanupScheduler, + admin::{ + auth::repository::AdminAuthRepository, tasks::session_cleanup::SessionCleanupScheduler, + }, codex::{ gateway::fingerprint::{ model::Fingerprint, repository::FingerprintRepository, update_checker::UpdateChecker, @@ -84,8 +86,10 @@ pub async fn start_background_tasks( coordinator.push("refresh", refresh_scheduler.start().await); tracing::info!("token 刷新调度器已启动"); - let session_cleanup = - SessionCleanupScheduler::new(db_pool, config.admin.session_cleanup_interval_secs); + let session_cleanup = SessionCleanupScheduler::new( + AdminAuthRepository::new(db_pool), + config.admin.session_cleanup_interval_secs, + ); coordinator.push("session_cleanup", session_cleanup.start()); tracing::info!("会话清理调度器已启动"); diff --git a/tests/admin.rs b/tests/admin.rs index 54e979c4b..a074f38b9 100644 --- a/tests/admin.rs +++ b/tests/admin.rs @@ -4,6 +4,7 @@ mod admin { mod api_contract; mod api_keys_route; mod auth; + mod auth_repository; mod login_route; mod logs_route; mod models_route; diff --git a/tests/admin/auth_repository.rs b/tests/admin/auth_repository.rs new file mode 100644 index 000000000..5e045d887 --- /dev/null +++ b/tests/admin/auth_repository.rs @@ -0,0 +1,57 @@ +use chrono::{Duration, Utc}; + +use codex_proxy_rs::{ + admin::auth::repository::AdminAuthRepository, + platform::{identity::admin_session::hash_admin_password, storage::db::connect_sqlite}, +}; + +#[tokio::test] +async fn admin_auth_repository_should_create_and_load_default_admin_once() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-auth.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let repo = AdminAuthRepository::new(pool); + let password_hash = hash_admin_password("admin-password").unwrap(); + + let created = repo.ensure_default_admin(&password_hash).await.unwrap(); + let skipped = repo.ensure_default_admin("different-hash").await.unwrap(); + let admin = repo.load_first_admin().await.unwrap().unwrap(); + + assert!(created); + assert!(!skipped); + assert_eq!(admin.password_hash, password_hash); +} + +#[tokio::test] +async fn admin_auth_repository_should_create_validate_and_cleanup_sessions() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("admin-sessions.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let repo = AdminAuthRepository::new(pool); + let password_hash = hash_admin_password("admin-password").unwrap(); + repo.ensure_default_admin(&password_hash).await.unwrap(); + let admin = repo.load_first_admin().await.unwrap().unwrap(); + let future_session = "sess_future"; + let expired_session = "sess_expired"; + + repo.create_session(future_session, &admin.id, Utc::now() + Duration::minutes(5)) + .await + .unwrap(); + repo.create_session( + expired_session, + &admin.id, + Utc::now() - Duration::minutes(5), + ) + .await + .unwrap(); + + assert!(repo.validate_session(future_session).await.unwrap()); + assert!(!repo.validate_session(expired_session).await.unwrap()); + + let deleted = repo.cleanup_expired_sessions(Utc::now()).await.unwrap(); + + assert_eq!(deleted, 1); + assert!(repo.validate_session(future_session).await.unwrap()); +} diff --git a/tests/codex_accounts/cookie_store.rs b/tests/codex_accounts/cookie_store.rs index 726408f5d..7bb709839 100644 --- a/tests/codex_accounts/cookie_store.rs +++ b/tests/codex_accounts/cookie_store.rs @@ -65,6 +65,28 @@ async fn cookie_repository_encrypts_and_replays_account_scoped_cookies() { ); } +#[tokio::test] +async fn cookie_repository_should_not_replay_expired_cookies() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("expired-cookies.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + seed_account(&pool, "acct_a").await; + let repo = CookieRepository::new(pool, SecretBox::new([19u8; 32])); + + repo.capture_set_cookie( + "acct_a", + "cf_clearance=expired; Domain=.chatgpt.com; Path=/; Expires=Wed, 21 Oct 2015 07:28:00 GMT", + ) + .await + .unwrap(); + + assert_eq!( + repo.cookie_header("acct_a", "chatgpt.com").await.unwrap(), + None + ); +} + async fn seed_account(pool: &sqlx::SqlitePool, id: &str) { let now = Utc::now().to_rfc3339(); sqlx::query( diff --git a/tests/codex_accounts/repository.rs b/tests/codex_accounts/repository.rs index cf4789416..c752f76c4 100644 --- a/tests/codex_accounts/repository.rs +++ b/tests/codex_accounts/repository.rs @@ -193,6 +193,63 @@ async fn account_repository_should_accumulate_usage_counters() { assert!(usage.last_used_at.is_some()); } +#[tokio::test] +async fn account_repository_should_accumulate_usage_window_counters() { + 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.clone(), SecretBox::new([32u8; 32])); + + repo.insert(NewAccount { + id: "acct_a".to_string(), + email: None, + account_id: None, + user_id: None, + label: None, + plan_type: None, + access_token: SecretString::new("access-secret".to_string().into()), + refresh_token: None, + access_token_expires_at: None, + status: AccountStatus::Active, + }) + .await + .unwrap(); + + repo.record_usage( + "acct_a", + UsageDelta { + input_tokens: 10, + output_tokens: 4, + cached_tokens: 3, + empty_response_count: 0, + }, + ) + .await + .unwrap(); + repo.record_usage( + "acct_a", + UsageDelta { + input_tokens: 8, + output_tokens: 2, + cached_tokens: 1, + empty_response_count: 0, + }, + ) + .await + .unwrap(); + + let usage: (i64, i64, i64, i64) = sqlx::query_as( + "select window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens from account_usage where account_id = ?", + ) + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(usage, (2, 18, 6, 4)); +} + #[tokio::test] async fn account_repository_should_load_usage_count_into_runtime_pool_accounts() { let dir = tempfile::tempdir().unwrap(); @@ -232,6 +289,39 @@ async fn account_repository_should_load_usage_count_into_runtime_pool_accounts() assert_eq!(account.request_count, 1); } +#[tokio::test] +async fn account_repository_should_restore_window_usage_into_runtime_pool_accounts() { + 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 secret_box = SecretBox::new([33u8; 32]); + let repo = AccountRepository::new(pool.clone(), secret_box.clone()); + let window_started_at = Utc::now() - Duration::minutes(2); + let window_reset_at = Utc::now() + Duration::minutes(3); + + seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await; + sqlx::query( + "insert into account_usage (account_id, request_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds) values (?, 9, 7, 11, 13, 17, ?, ?, 300)", + ) + .bind("acct_a") + .bind(window_started_at.to_rfc3339()) + .bind(window_reset_at.to_rfc3339()) + .execute(&pool) + .await + .unwrap(); + + let account = repo.list_pool_accounts().await.unwrap().remove(0); + + assert_eq!(account.window_request_count, 7); + assert_eq!(account.window_input_tokens, 11); + assert_eq!(account.window_output_tokens, 13); + assert_eq!(account.window_cached_tokens, 17); + assert_eq!(account.window_started_at, Some(window_started_at)); + assert_eq!(account.window_reset_at, Some(window_reset_at)); + assert_eq!(account.limit_window_seconds, Some(300)); +} + #[tokio::test] async fn account_repository_should_exclude_refresh_lease_until_expiry() { let dir = tempfile::tempdir().unwrap(); diff --git a/tests/codex_gateway/fingerprint_update.rs b/tests/codex_gateway/fingerprint_update.rs index 959d87b77..8e0b69ffb 100644 --- a/tests/codex_gateway/fingerprint_update.rs +++ b/tests/codex_gateway/fingerprint_update.rs @@ -41,3 +41,32 @@ async fn fingerprint_updater_fetches_manifest_and_persists_history() { assert_eq!(latest.build_number, "5002"); assert_eq!(latest.source, CODEX_DESKTOP_UPDATE_SOURCE); } + +#[tokio::test] +async fn fingerprint_repository_should_upsert_auto_update_record() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("fingerprints.sqlite"); + let pool = connect_sqlite(&format!("sqlite://{}", db.display())) + .await + .unwrap(); + let repo = FingerprintRepository::new(pool.clone()); + + repo.upsert_auto_update("26.800.1", "6001", Some("147")) + .await + .unwrap(); + repo.upsert_auto_update("26.800.2", "6002", None) + .await + .unwrap(); + + let stored = repo.load_latest_auto_updated().await.unwrap().unwrap(); + let count: (i64,) = + sqlx::query_as("select count(*) from fingerprints where id = 'auto_updated'") + .fetch_one(&pool) + .await + .unwrap(); + + assert_eq!(count.0, 1); + assert_eq!(stored.app_version, "26.800.2"); + assert_eq!(stored.build_number, "6002"); + assert_eq!(stored.chromium_version, "146"); +} diff --git a/tests/codex_operations/logs_pagination.rs b/tests/codex_operations/logs_pagination.rs index 1c2943777..c5e1eac97 100644 --- a/tests/codex_operations/logs_pagination.rs +++ b/tests/codex_operations/logs_pagination.rs @@ -2,9 +2,12 @@ use codex_proxy_rs::{ codex::logs::{ event::{EventLevel, EventLog}, repository::EventLogRepository, + service::LogService, }, + config::LoggingConfig, platform::storage::db::connect_sqlite, }; +use serde_json::json; #[tokio::test] async fn event_logs_are_cursor_paginated() { @@ -32,3 +35,91 @@ async fn event_logs_are_cursor_paginated() { assert_eq!(second.items.len(), 1); assert!(second.next_cursor.is_none()); } + +#[tokio::test] +async fn log_service_should_skip_record_when_disabled() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("logs-disabled.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let service = LogService::new( + logging_config(false, 10, false), + Some(EventLogRepository::new(pool.clone())), + ); + + let recorded = service + .record(EventLog::new("request", EventLevel::Info, "disabled")) + .await + .unwrap(); + + let stored_count: (i64,) = sqlx::query_as("select count(*) from event_logs") + .fetch_one(&pool) + .await + .unwrap(); + assert!(!recorded); + assert_eq!(stored_count.0, 0); +} + +#[tokio::test] +async fn log_service_should_trim_to_capacity_after_record() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("logs-capacity.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let service = LogService::new( + logging_config(true, 1, false), + Some(EventLogRepository::new(pool.clone())), + ); + let mut first = EventLog::new("request", EventLevel::Info, "first"); + first.created_at = "2026-06-14T00:00:00Z".to_string(); + let mut second = EventLog::new("request", EventLevel::Info, "second"); + second.created_at = "2026-06-14T00:00:01Z".to_string(); + + service.record(first).await.unwrap(); + service.record(second).await.unwrap(); + + let messages: Vec<(String,)> = + sqlx::query_as("select message from event_logs order by created_at desc") + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(messages, vec![("second".to_string(),)]); +} + +#[tokio::test] +async fn log_service_should_remove_body_metadata_when_capture_body_disabled() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("logs-body-policy.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let service = LogService::new( + logging_config(true, 10, false), + Some(EventLogRepository::new(pool.clone())), + ); + let mut event = EventLog::new("request", EventLevel::Info, "body"); + event.metadata = json!({ + "body": "secret", + "requestBody": "secret", + "responseBody": "secret", + "safe": "visible" + }); + + service.record(event).await.unwrap(); + + let metadata_json: (String,) = sqlx::query_as("select metadata_json from event_logs") + .fetch_one(&pool) + .await + .unwrap(); + let metadata: serde_json::Value = serde_json::from_str(&metadata_json.0).unwrap(); + assert_eq!(metadata, json!({ "safe": "visible" })); +} + +fn logging_config(enabled: bool, capacity: u32, capture_body: bool) -> LoggingConfig { + LoggingConfig { + directory: "logs".to_string(), + retention_days: 14, + enabled, + capacity, + capture_body, + } +} diff --git a/tests/codex_serving/responses_http_sse.rs b/tests/codex_serving/responses_http_sse.rs index f572d203a..6d8eee6aa 100644 --- a/tests/codex_serving/responses_http_sse.rs +++ b/tests/codex_serving/responses_http_sse.rs @@ -2,7 +2,7 @@ use axum::{ body::Body, http::{Request, StatusCode}, }; -use chrono::Utc; +use chrono::{DateTime, Utc}; use serde_json::{json, Value}; use tower::ServiceExt; use wiremock::{ @@ -17,7 +17,7 @@ use codex_proxy_rs::{ use crate::support::{ response_json, response_text, - upstream::{build_imported_app, fetch_v1_event_log}, + upstream::{build_imported_app, enable_runtime_logging, fetch_v1_event_log}, }; const COMPLETED_USAGE_SSE: &str = @@ -34,6 +34,51 @@ const STREAM_USAGE_SSE: &str = include_str!("../fixtures/responses/http_sse/stre const DEFAULT_STREAM_SSE: &str = include_str!("../fixtures/responses/http_sse/default_stream.sse"); const EMPTY_COMPLETED_SSE: &str = "event: response.completed\ndata: {\"response\":{\"id\":\"resp_empty\",\"object\":\"response\",\"status\":\"completed\",\"output\":[],\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n"; +#[tokio::test] +async fn v1_responses_should_skip_event_log_when_logging_disabled() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(COMPLETED_USAGE_SSE), + ) + .mount(&server) + .await; + let imported = build_imported_app(server.uri()).await; + + let response = imported + .app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header( + "authorization", + format!("Bearer {}", imported.client_api_key), + ) + .header("content-type", "application/json") + .header("x-request-id", "req_logs_disabled") + .body(Body::from( + r#"{"model":"gpt-5.5","input":[],"stream":false}"#, + )) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let stored_count: (i64,) = sqlx::query_as( + "select count(*) from event_logs where request_id = ? and kind = 'v1.response'", + ) + .bind("req_logs_disabled") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(stored_count.0, 0); +} + #[tokio::test] async fn v1_responses_should_honor_explicit_http_sse_transport() { let server = MockServer::start().await; @@ -97,6 +142,7 @@ async fn v1_responses_should_use_imported_account_and_record_usage() { .mount(&server) .await; let imported = build_imported_app(server.uri()).await; + enable_runtime_logging(&imported.app).await; let response = imported .app @@ -203,6 +249,26 @@ async fn v1_responses_should_passively_cache_rate_limit_headers() { assert_eq!(quota["credits"]["balance"], 12); assert_eq!(stored.1, 1); assert!(stored.2.is_some()); + let window: (i64, i64, i64, i64, Option, Option, Option) = + sqlx::query_as( + "select window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds from account_usage where account_id = ?", + ) + .bind("acct_imported") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(window.0, 1); + assert_eq!(window.1, 7); + assert_eq!(window.2, 4); + assert_eq!(window.3, 2); + assert!(window.4.is_some()); + assert_eq!( + DateTime::parse_from_rfc3339(window.5.as_deref().unwrap()) + .unwrap() + .timestamp(), + reset_at + ); + assert_eq!(window.6, Some(300)); } #[tokio::test] @@ -509,6 +575,7 @@ async fn v1_responses_should_passthrough_stream_and_record_usage_and_log() { .mount(&server) .await; let imported = build_imported_app(server.uri()).await; + enable_runtime_logging(&imported.app).await; let response = imported .app diff --git a/tests/codex_serving/responses_websocket.rs b/tests/codex_serving/responses_websocket.rs index a34d6fd9a..11c1651d7 100644 --- a/tests/codex_serving/responses_websocket.rs +++ b/tests/codex_serving/responses_websocket.rs @@ -379,6 +379,18 @@ async fn v1_responses_should_route_previous_response_id_to_recorded_account() { assert_eq!(first_response.status(), StatusCode::OK); let first_body = response_text(first_response).await; assert!(first_body.contains("\"id\":\"resp_affinity_first\"")); + let stored_affinity: (String, String, String, Option, String) = sqlx::query_as( + "select account_id, conversation_id, function_call_ids_json, input_tokens, expires_at from session_affinities where response_id = ?", + ) + .bind("resp_affinity_first") + .fetch_one(&imported.pool) + .await + .unwrap(); + assert_eq!(stored_affinity.0, "acct_a"); + assert!(!stored_affinity.1.is_empty()); + assert_eq!(stored_affinity.2, "[]"); + assert_eq!(stored_affinity.3, Some(3)); + assert!(!stored_affinity.4.is_empty()); let second_response = imported .app diff --git a/tests/codex_serving/upstream_errors.rs b/tests/codex_serving/upstream_errors.rs index 1bd8ce187..1317ad8a8 100644 --- a/tests/codex_serving/upstream_errors.rs +++ b/tests/codex_serving/upstream_errors.rs @@ -10,7 +10,7 @@ use wiremock::{ use crate::support::{ response_json, response_text, - upstream::{build_imported_app, fetch_v1_event_log}, + upstream::{build_imported_app, enable_runtime_logging, fetch_v1_event_log}, }; const ERROR_EVENT_SSE: &str = include_str!("../fixtures/responses/http_sse/error_event.sse"); @@ -118,6 +118,7 @@ async fn v1_responses_stream_should_passthrough_error_event_and_log_failure() { .mount(&server) .await; let imported = build_imported_app(server.uri()).await; + enable_runtime_logging(&imported.app).await; let response = imported .app @@ -161,6 +162,7 @@ async fn v1_responses_stream_should_passthrough_response_failed_and_log_failure( .mount(&server) .await; let imported = build_imported_app(server.uri()).await; + enable_runtime_logging(&imported.app).await; let response = imported .app diff --git a/tests/codex_serving/upstream_fallback.rs b/tests/codex_serving/upstream_fallback.rs index a3cca0c74..10055ecbc 100644 --- a/tests/codex_serving/upstream_fallback.rs +++ b/tests/codex_serving/upstream_fallback.rs @@ -20,7 +20,8 @@ use crate::support::{ response_json, response_text, upstream::{ build_imported_app_with_accounts, build_imported_app_with_refresher, - build_imported_app_with_token_refresher, FailingTokenRefresher, ImportAccount, + build_imported_app_with_token_refresher, enable_runtime_logging, FailingTokenRefresher, + ImportAccount, }, }; @@ -604,6 +605,7 @@ async fn v1_responses_should_log_refresh_failure_after_401() { }, ) .await; + enable_runtime_logging(&imported.app).await; let response = imported .app diff --git a/tests/platform/storage_schema.rs b/tests/platform/storage_schema.rs index ef6682493..c05ddc3e8 100644 --- a/tests/platform/storage_schema.rs +++ b/tests/platform/storage_schema.rs @@ -16,6 +16,126 @@ async fn sqlite_schema_creates_accounts_and_event_tables() { assert_eq!(row.0, 4); } +#[tokio::test] +async fn sqlite_schema_should_persist_account_usage_window_columns() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("test.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let rows: Vec<(String,)> = sqlx::query_as( + "select name from pragma_table_info('account_usage') where name in ('window_request_count', 'window_input_tokens', 'window_output_tokens', 'window_cached_tokens', 'window_started_at', 'window_reset_at', 'limit_window_seconds') order by name", + ) + .fetch_all(&pool) + .await + .unwrap(); + + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "limit_window_seconds", + "window_cached_tokens", + "window_input_tokens", + "window_output_tokens", + "window_request_count", + "window_reset_at", + "window_started_at", + ] + ); +} + +#[tokio::test] +async fn sqlite_schema_should_persist_session_affinity_function_call_ids() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("test.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let row: Option<(String,)> = sqlx::query_as( + "select name from pragma_table_info('session_affinities') where name = 'function_call_ids_json'", + ) + .fetch_optional(&pool) + .await + .unwrap(); + + assert_eq!(row, Some(("function_call_ids_json".to_string(),))); +} + +#[tokio::test] +async fn sqlite_schema_should_reject_invalid_account_status() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("status-check.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let result = sqlx::query( + "insert into accounts (id, access_token_cipher, status, added_at, updated_at) values ('acct_bad', 'cipher', 'unknown', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn sqlite_schema_should_reject_invalid_event_log_level() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("event-level-check.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let result = sqlx::query( + "insert into event_logs (id, kind, level, message, metadata_json, created_at) values ('log_bad', 'request', 'fatal', 'bad', '{}', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn sqlite_schema_should_reject_non_boolean_flags() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("boolean-check.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let api_key_result = sqlx::query( + "insert into client_api_keys (id, name, prefix, key_hash, enabled, created_at) values ('key_bad', 'bad', 'cpr_x', 'hash', 2, '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + let account_result = sqlx::query( + "insert into accounts (id, access_token_cipher, status, quota_limit_reached, added_at, updated_at) values ('acct_bad', 'cipher', 'active', 2, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + + assert!(api_key_result.is_err()); + assert!(account_result.is_err()); +} + +#[tokio::test] +async fn sqlite_schema_should_reject_negative_account_usage_counts() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage-check.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + sqlx::query( + "insert into accounts (id, access_token_cipher, status, added_at, updated_at) values ('acct_a', 'cipher', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await + .unwrap(); + let result = + sqlx::query("insert into account_usage (account_id, request_count) values ('acct_a', -1)") + .execute(&pool) + .await; + + assert!(result.is_err()); +} + #[tokio::test] async fn connect_sqlite_creates_parent_directory_for_file_database() { let dir = tempfile::tempdir().unwrap(); diff --git a/tests/runtime.rs b/tests/runtime.rs index 143724dfc..fbf8392c4 100644 --- a/tests/runtime.rs +++ b/tests/runtime.rs @@ -1,3 +1,5 @@ +mod support; + mod runtime { mod http_trace; mod startup; diff --git a/tests/runtime/startup.rs b/tests/runtime/startup.rs index 91a13a7dc..ec4262bbe 100644 --- a/tests/runtime/startup.rs +++ b/tests/runtime/startup.rs @@ -1,5 +1,21 @@ +use std::time::Instant; + +use axum::{ + body::Bytes, + http::{HeaderMap, StatusCode}, +}; use chrono::{Duration, Utc}; +use futures::{SinkExt, StreamExt}; use secrecy::SecretString; +use serde_json::{json, Value}; +use tokio::net::TcpListener; +use tokio_tungstenite::{ + accept_hdr_async, + tungstenite::{ + handshake::server::{Request as WsRequest, Response as WsResponse}, + Message, + }, +}; use codex_proxy_rs::{ codex::accounts::{ @@ -16,6 +32,8 @@ use codex_proxy_rs::{ runtime::state::AppState, }; +use crate::support::response_text; + fn test_config(database_url: String) -> AppConfig { AppConfig { server: ServerConfig { @@ -174,3 +192,176 @@ async fn app_state_should_restore_persisted_cooldowns_into_runtime_pool() { .await .is_none()); } + +#[tokio::test] +async fn app_state_should_restore_session_affinity_from_sqlite() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + accept_successful_websocket_response( + &listener, + "Bearer access-a", + Some("turn_restore"), + "resp_restore_second", + ) + .await + }); + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("startup-affinity.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + let secret_box = SecretBox::new([33u8; 32]); + let repo = codex_proxy_rs::codex::accounts::repository::AccountRepository::new( + pool.clone(), + secret_box.clone(), + ); + repo.insert(NewAccount { + id: "acct_a".to_string(), + email: Some("a@example.com".to_string()), + account_id: Some("chatgpt-a".to_string()), + user_id: None, + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-a".to_string().into()), + refresh_token: Some(SecretString::new("refresh-a".to_string().into())), + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + status: AccountStatus::Active, + }) + .await + .unwrap(); + repo.insert(NewAccount { + id: "acct_b".to_string(), + email: Some("b@example.com".to_string()), + account_id: Some("chatgpt-b".to_string()), + user_id: None, + label: None, + plan_type: Some("plus".to_string()), + access_token: SecretString::new("access-b".to_string().into()), + refresh_token: Some(SecretString::new("refresh-b".to_string().into())), + access_token_expires_at: Some(Utc::now() + Duration::hours(1)), + status: AccountStatus::Active, + }) + .await + .unwrap(); + repo.record_usage( + "acct_a", + UsageDelta { + input_tokens: 1, + output_tokens: 0, + cached_tokens: 0, + empty_response_count: 0, + }, + ) + .await + .unwrap(); + let now = Utc::now(); + sqlx::query( + "insert into session_affinities (response_id, account_id, conversation_id, turn_state, instructions_hash, input_tokens, function_call_ids_json, variant_hash, expires_at, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind("resp_restore_first") + .bind("acct_a") + .bind("conv_restore") + .bind("turn_restore") + .bind(Option::::None) + .bind(3_i64) + .bind(r#"["call_restore"]"#) + .bind(Option::::None) + .bind((now + Duration::hours(3)).to_rfc3339()) + .bind(now.to_rfc3339()) + .execute(&pool) + .await + .unwrap(); + let mut config = test_config(url); + config.api.base_url = format!("http://{addr}"); + let state = AppState::with_pool_and_secret_box(config, pool, secret_box); + state.reload_account_pool_from_repository().await.unwrap(); + assert_eq!( + state + .reload_session_affinity_from_repository() + .await + .unwrap(), + 1 + ); + + let response = state + .services + .responses + .handle( + "req_restore_affinity", + HeaderMap::new(), + Bytes::from( + r#"{"model":"gpt-5.5","input":[],"previous_response_id":"resp_restore_first"}"#, + ), + Instant::now(), + ) + .await; + + assert_eq!(response.status(), StatusCode::OK); + let body = response_text(response).await; + assert!(body.contains("\"id\":\"resp_restore_second\"")); + let upstream_request = server.await.unwrap(); + assert_eq!( + upstream_request["previous_response_id"], + "resp_restore_first" + ); +} + +#[expect( + clippy::result_large_err, + reason = "tokio-tungstenite handshake callbacks use a large error response type" +)] +async fn accept_successful_websocket_response( + listener: &TcpListener, + expected_authorization: &str, + expected_turn_state: Option<&str>, + response_id: &str, +) -> Value { + let (stream, _) = listener.accept().await.unwrap(); + let expected_authorization = expected_authorization.to_string(); + let expected_turn_state = expected_turn_state.map(ToString::to_string); + let mut websocket = + accept_hdr_async(stream, move |request: &WsRequest, response: WsResponse| { + assert_eq!( + request + .headers() + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some(expected_authorization.as_str()) + ); + assert_eq!( + request + .headers() + .get("x-codex-turn-state") + .and_then(|value| value.to_str().ok()), + expected_turn_state.as_deref() + ); + Ok(response) + }) + .await + .unwrap(); + let message = websocket.next().await.unwrap().unwrap(); + let request = serde_json::from_str::(&message.into_text().unwrap()).unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "input_tokens_details": {"cached_tokens": 0} + } + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.close(None).await.unwrap(); + request +} diff --git a/tests/support/upstream.rs b/tests/support/upstream.rs index b9c3ebca7..ea10b25b0 100644 --- a/tests/support/upstream.rs +++ b/tests/support/upstream.rs @@ -317,3 +317,20 @@ pub async fn fetch_v1_event_log( serde_json::from_str(&row.4).unwrap(), ) } + +pub async fn enable_runtime_logging(app: &axum::Router) { + let response = app + .clone() + .oneshot( + Request::builder() + .method("PATCH") + .uri("/admin/logs/state") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({ "enabled": true }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); +}