From 8e6bdd14fc3ea7375ebd6676c99ae6e5c09564fe Mon Sep 17 00:00:00 2001 From: zyycn Date: Sun, 28 Jun 2026 11:55:51 +0800 Subject: [PATCH] docs: audit deleted feature parity Co-authored-by: Codex --- Cargo.lock | 2 - Cargo.toml | 4 +- README.md | 8 +- docs/deleted-code-feature-parity-audit.md | 127 +++ src/admin/accounts/routes.rs | 768 +++--------------- src/admin/accounts/service/cookies.rs | 68 +- src/admin/accounts/service/importing.rs | 136 +--- src/admin/accounts/service/lifecycle.rs | 80 +- src/admin/accounts/service/mod.rs | 49 +- src/admin/accounts/service/oauth.rs | 10 +- src/admin/accounts/service/quota.rs | 190 ----- src/admin/accounts/service/testing.rs | 20 +- src/admin/accounts/service/types.rs | 31 +- src/admin/auth/service.rs | 54 +- src/admin/auth/session.rs | 81 +- src/admin/keys/routes.rs | 274 ++----- src/admin/keys/service.rs | 74 +- src/admin/mod.rs | 2 - src/admin/models/mod.rs | 3 - src/admin/models/routes.rs | 92 --- src/admin/monitoring/billing.rs | 167 ++-- src/admin/monitoring/dashboard.rs | 235 +++--- src/admin/monitoring/event_store.rs | 106 +-- src/admin/monitoring/events.rs | 47 +- src/admin/monitoring/logs.rs | 233 ++---- src/admin/monitoring/mod.rs | 1 - src/admin/monitoring/service.rs | 112 +-- src/admin/monitoring/usage.rs | 201 ----- src/admin/monitoring/usage_store.rs | 74 +- src/admin/response.rs | 68 +- src/admin/router.rs | 36 +- src/admin/settings/routes.rs | 67 +- src/config/loader.rs | 6 +- src/config/mod.rs | 3 - src/config/settings.rs | 31 +- src/config/types.rs | 21 - src/http/middleware/trace.rs | 7 +- src/infra/database.rs | 3 +- src/infra/json.rs | 23 - src/infra/logging.rs | 4 +- src/infra/paths.rs | 14 +- src/infra/time.rs | 23 +- src/proxy/dispatch/chat.rs | 21 +- src/proxy/dispatch/cloudflare.rs | 6 +- src/proxy/dispatch/responses.rs | 405 ++++----- src/proxy/dispatch/session_affinity.rs | 10 - src/proxy/dispatch/upstream.rs | 45 +- src/proxy/openai/chat.rs | 66 +- src/proxy/openai/errors.rs | 10 +- src/proxy/openai/models.rs | 21 - src/proxy/openai/responses.rs | 26 +- src/runtime/bootstrap.rs | 2 +- src/runtime/services.rs | 88 +- src/runtime/shutdown.rs | 4 +- src/runtime/tasks/cookie_cleanup.rs | 8 - src/runtime/tasks/coordinator.rs | 23 +- src/runtime/tasks/mod.rs | 2 - src/runtime/tasks/model_refresh.rs | 30 +- src/runtime/tasks/quota_refresh.rs | 10 - src/runtime/tasks/token_refresh.rs | 13 - src/upstream/accounts/cookies.rs | 55 +- .../{import_export.rs => importing.rs} | 73 +- src/upstream/accounts/mod.rs | 6 +- src/upstream/accounts/model.rs | 21 +- src/upstream/accounts/pool.rs | 169 +--- src/upstream/accounts/quota/mod.rs | 10 +- src/upstream/accounts/quota/runtime.rs | 11 +- src/upstream/accounts/service.rs | 29 +- src/upstream/accounts/store.rs | 275 ++----- src/upstream/accounts/token_refresh/mod.rs | 9 +- .../accounts/token_refresh/runtime.rs | 14 +- src/upstream/fingerprint.rs | 197 +---- src/upstream/models/mod.rs | 110 +-- src/upstream/protocol/events.rs | 88 +- src/upstream/protocol/websocket.rs | 171 ---- src/upstream/token_client.rs | 2 +- src/upstream/transport/client.rs | 37 +- src/upstream/transport/endpoints.rs | 9 - src/upstream/transport/mod.rs | 16 +- src/upstream/transport/usage.rs | 7 +- src/upstream/transport/websocket.rs | 49 +- src/upstream/transport/websocket_pool.rs | 49 +- .../{import_export.rs => importing.rs} | 116 +-- tests/admin/accounts/lifecycle.rs | 153 ---- tests/admin/accounts/list.rs | 22 +- tests/admin/accounts/mod.rs | 130 +-- tests/admin/accounts/quota.rs | 370 +-------- tests/admin/accounts/testing.rs | 7 +- tests/admin/auth/mod.rs | 7 +- tests/admin/auth/session.rs | 30 +- tests/admin/auth/session_routes.rs | 208 ----- tests/admin/keys/authorization.rs | 3 - tests/admin/keys/import_export.rs | 25 - tests/admin/keys/lifecycle.rs | 1 - tests/admin/keys/mod.rs | 92 +-- tests/admin/mod.rs | 13 +- tests/admin/models/mod.rs | 1 - tests/admin/models/routes.rs | 247 ------ tests/admin/monitoring/dashboard.rs | 32 +- tests/admin/monitoring/events_store.rs | 9 +- tests/admin/monitoring/logs_routes.rs | 92 +-- tests/admin/monitoring/mod.rs | 6 +- tests/admin/response/mod.rs | 21 +- tests/admin/settings/mod.rs | 2 +- tests/admin/settings/routes.rs | 97 +-- tests/config/mod.rs | 52 +- tests/http/mod.rs | 4 +- tests/http/trace_middleware/mod.rs | 21 +- tests/http/web_assets/mod.rs | 71 +- tests/infra/mod.rs | 6 +- tests/infra/storage_schema/mod.rs | 45 +- .../dispatch/chat_upstream/chat_routes.rs | 21 +- tests/proxy/dispatch/chat_upstream/mod.rs | 151 ++-- .../chat_upstream/responses_recovery.rs | 36 +- .../chat_upstream/responses_websocket.rs | 72 +- tests/proxy/dispatch/mod.rs | 6 +- tests/proxy/keys/mod.rs | 2 +- tests/proxy/mod.rs | 6 +- tests/proxy/openai/mod.rs | 4 +- tests/proxy/openai/models_routes.rs | 67 +- tests/proxy/openai/responses_routes.rs | 76 +- tests/runtime/account_pool_restore/mod.rs | 13 +- tests/runtime/mod.rs | 6 +- tests/runtime/services.rs | 9 +- tests/runtime/tasks/fingerprint.rs | 13 +- tests/runtime/tasks/model_refresh.rs | 213 ++++- tests/support/accounts.rs | 2 +- tests/support/admin.rs | 25 + tests/support/assertions.rs | 9 + tests/support/client_keys.rs | 10 + tests/support/config.rs | 14 +- tests/support/fingerprint.rs | 8 +- tests/support/http.rs | 6 + tests/support/jwt.rs | 11 + tests/support/mod.rs | 13 +- tests/support/sqlite.rs | 2 +- tests/upstream/accounts/account_pool/mod.rs | 17 +- tests/upstream/accounts/account_pool/quota.rs | 20 +- .../accounts/account_pool/selection.rs | 45 +- .../accounts/account_pool/usage_window.rs | 64 +- tests/upstream/accounts/account_repository.rs | 31 +- tests/upstream/accounts/mod.rs | 14 +- tests/upstream/accounts/quota_refresh.rs | 26 +- .../accounts/token_refresh/failures.rs | 2 +- tests/upstream/accounts/token_refresh/mod.rs | 8 +- .../accounts/token_refresh/scheduling.rs | 17 +- tests/upstream/fingerprint/mod.rs | 51 +- tests/upstream/mod.rs | 14 +- tests/upstream/models/catalog.rs | 16 +- tests/upstream/models/mod.rs | 4 +- tests/upstream/protocol/codex_websocket.rs | 60 -- tests/upstream/protocol/mod.rs | 23 +- tests/upstream/protocol/usage_rate_limits.rs | 31 +- tests/upstream/transport/client.rs | 131 +-- tests/upstream/transport/headers.rs | 2 +- tests/upstream/transport/http_client.rs | 57 +- tests/upstream/transport/mod.rs | 50 +- tests/upstream/transport/websocket.rs | 71 +- tests/upstream/transport/websocket_pool.rs | 27 +- web/src/api/index.ts | 3 - web/src/api/modules/accounts.ts | 18 +- web/src/api/modules/usage.ts | 16 - web/src/api/request.ts | 13 +- web/src/components/base/BaseSwitch.vue | 52 -- web/src/components/base/BaseToast/index.ts | 1 - .../components/AccountCreateModal.vue | 2 +- .../composables/useAccountConnectionTest.ts | 4 +- .../composables/useAccountMutations.ts | 4 +- .../dashboard/components/EventLogCard.vue | 1 - 169 files changed, 2153 insertions(+), 7172 deletions(-) create mode 100644 docs/deleted-code-feature-parity-audit.md delete mode 100644 src/admin/models/mod.rs delete mode 100644 src/admin/models/routes.rs delete mode 100644 src/admin/monitoring/usage.rs rename src/upstream/accounts/{import_export.rs => importing.rs} (82%) rename tests/admin/accounts/{import_export.rs => importing.rs} (84%) delete mode 100644 tests/admin/auth/session_routes.rs delete mode 100644 tests/admin/keys/import_export.rs delete mode 100644 tests/admin/models/mod.rs delete mode 100644 tests/admin/models/routes.rs create mode 100644 tests/support/admin.rs create mode 100644 tests/support/assertions.rs create mode 100644 tests/support/client_keys.rs create mode 100644 tests/support/http.rs create mode 100644 tests/support/jwt.rs delete mode 100644 web/src/api/modules/usage.ts delete mode 100644 web/src/components/base/BaseSwitch.vue diff --git a/Cargo.lock b/Cargo.lock index 4cd1af15f..84715fb2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -343,8 +343,6 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokio", - "tokio-rustls", - "tokio-stream", "tokio-tungstenite", "tower", "tower-http", diff --git a/Cargo.toml b/Cargo.toml index 5a1a2ac21..169ae0004 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,13 +28,10 @@ rustls-pki-types = "1.14.1" secrecy = "0.10.3" serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.150" -serde_yml = "0.0.13" sha2 = "0.11.0" sqlx = { version = "0.9.0", features = ["runtime-tokio", "tls-rustls", "sqlite", "chrono", "uuid"] } thiserror = "2.0.18" tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "signal", "fs", "time", "sync", "net", "test-util"] } -tokio-rustls = { version = "0.26.4", default-features = false, features = ["ring", "tls12"] } -tokio-stream = "0.1.17" tokio-tungstenite = { version = "0.28.0", features = ["proxy", "rustls-tls-native-roots"] } tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout"] } tracing = "0.1.44" @@ -44,6 +41,7 @@ tungstenite = { version = "0.27.0", features = ["deflate", "proxy"] } uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] } [dev-dependencies] +serde_yml = "0.0.13" tempfile = "3.27.0" tower = { version = "0.5.3", features = ["util"] } wiremock = "0.6.5" diff --git a/README.md b/README.md index 94d17057b..fa7737cd5 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ Codex Proxy RS 关注这些细节： **Codex 账号池** -- 支持 ChatGPT/Codex OAuth、device code、refresh token 账号接入 +- 支持 ChatGPT/Codex OAuth、CPR JSON、Sub2API JSON 和 refresh token 账号接入 - 支持账号轮转、并发控制、请求间隔和会话亲和性 - 支持账号状态、配额、Cookie、模型目录和连接测试管理 @@ -52,7 +52,7 @@ Codex Proxy RS 关注这些细节： **可观测与管理** - 管理控制台内置账号、密钥、事件日志、用量统计和系统设置 -- SQLite 持久化，敏感密钥加密存储 +- SQLite 持久化，账号、Cookie 与访问密钥采用统一的本地存储策略 - 结构化事件日志，便于定位请求失败、上游状态和调度路径 ## 快速开始 @@ -69,7 +69,7 @@ admin: default_password: admin ``` -长期运行前请修改默认密码，并确认数据库、密钥和日志路径符合你的部署环境。 +长期运行前请修改默认密码，并确认数据库和日志路径符合你的部署环境。 运行时数据默认写入 `.runtime`： @@ -142,7 +142,7 @@ flowchart LR proxy --> codex[codex
Codex 协议、模型和上游传输] http --> admin[admin
管理端 API] admin --> runtime[runtime
启动装配、状态和后台任务] - runtime --> infra[infra
SQLite、加密、路径和日志] + runtime --> infra[infra
SQLite、路径和日志] proxy --> telemetry[telemetry
事件日志和用量记录] admin --> telemetry http --> web[web
管理控制台静态资源] diff --git a/docs/deleted-code-feature-parity-audit.md b/docs/deleted-code-feature-parity-audit.md new file mode 100644 index 000000000..40ab42221 --- /dev/null +++ b/docs/deleted-code-feature-parity-audit.md @@ -0,0 +1,127 @@ +# 删除代码功能等价审计 + +审计日期：2026-06-28 + +对照基准： + +- 当前 Rust 工程：`/home/zyy/Codes/codex-proxy-rs` +- TS 参考工程：`/home/zyy/桌面/Codes/codex-proxy` + +本文档只审计近期删除或收敛的功能是否造成产品能力缺失。结论不以文件名为准，而以当前路由、服务实现和 TS 版已存在能力为准。 + +## 总体结论 + +当前删除并不全是问题。部分删除符合当前产品方向：后台设置减少人工配置、日志策略固定、接口去掉路径 id、账号导入增强为 CPR/Sub2API 双格式并支持重复账号更新，这些属于合理收敛。 + +但有几类能力被删后没有等价替代，会影响管理闭环： + +1. 账号导出缺失，导入和备份迁移不闭环。 +2. 批量账号状态更新缺失，账号池批量维护效率下降。 +3. 批量健康检查缺失，目前只剩单账号模型测试。 +4. 手动 Cookie 维护缺失，Cloudflare 人工旁路能力不完整。 +5. 独立用量统计接口缺失，当前只被 Dashboard 和账号列表部分吸收。 +6. 手动模型刷新缺失，后台只剩启动后和周期任务刷新。 +7. 客户端 API Key 导出缺失，密钥备份/迁移能力被削弱。 + +## 当前 Rust 管理端路由事实 + +当前 `src/admin/router.rs` 保留的管理端接口： + +- 登录会话：`/api/admin/login`、`/api/admin/auth/status`、`/api/admin/logout` +- 设置：`/api/admin/settings` +- 概览：`/api/admin/dashboard/summary`、`/api/admin/dashboard/trend` +- 账号：`/api/admin/accounts`、`/api/admin/accounts/import`、`/api/admin/accounts/oauth/authorize`、`/api/admin/accounts/oauth/exchange`、`/api/admin/accounts/test`、`/api/admin/accounts/models`、`/api/admin/accounts/delete`、`/api/admin/accounts/update`、`/api/admin/accounts/refresh`、`/api/admin/accounts/quota` +- 日志：`/api/admin/logs`、`/api/admin/logs/delete`、`/api/admin/logs/detail` +- 客户端 API Key：`/api/admin/keys`、`/api/admin/keys/delete`、`/api/admin/keys/update` + +已从路由层移除的旧 Rust 管理端接口： + +- `/api/admin/models/refresh` +- `/api/admin/usage` +- `/api/admin/usage/summary` +- `/api/admin/accounts/export` +- `/api/admin/accounts/quota-warnings` +- `/api/admin/accounts/health-check` +- `/api/admin/accounts/reset-usage` +- `/api/admin/accounts/cookies` +- `/api/admin/logs/state` +- `/api/admin/keys/export` + +## 逐项审计 + +| 能力 | Rust 当前状态 | TS 参考状态 | 结论 | 建议 | +| --- | --- | --- | --- | --- | +| 账号导入 | 已保留并增强。`src/admin/accounts/service/importing.rs` 支持 CPR/Sub2API，重复账号会按 ChatGPT identity 更新，缺少字段时会补充 account id、user id、email、quota。 | `src/routes/accounts.ts` 支持 `/auth/accounts/import`。 | 不缺失。当前 Rust 实现比旧导入更贴近当前产品。 | 保留当前实现。 | +| 账号导出 | 当前路由已删除 `/api/admin/accounts/export`，旧 `export_accounts` 和 `export_with_tokens` 不再暴露。 | `src/routes/accounts.ts` 有 `/auth/accounts/export`，支持按 ids 和 format 导出。 | 功能缺失。导入存在但导出消失，备份、迁移、批量排障不闭环。 | P1 恢复，格式只保留 CPR，避免旧 native 命名和兼容分支。 | +| 批量删除账号 | 当前保留 `/api/admin/accounts/delete`。 | TS 有 `/auth/accounts/batch-delete`。 | 不缺失。 | 保留。 | +| 批量状态更新 | 当前 `/api/admin/accounts/update` 只解析单账号 `id` 更新，旧 `batch_update_status` 已删除。 | TS 有 `/auth/accounts/batch-status`。 | 功能缺失。批量禁用、恢复账号池时需要逐个操作。 | P1 恢复为 `/api/admin/accounts/update` 的批量 data 形式，保持无路径 id。 | +| 单账号刷新/额度查询 | 当前保留 `/api/admin/accounts/refresh` 和 `/api/admin/accounts/quota`，id 放在 body 或 query。 | TS 有 `/auth/accounts/:id/refresh`、`/auth/accounts/:id/quota`。 | 不缺失，且当前 Rust 路由形态更符合项目规范。 | 保留当前形态。 | +| 重置账号用量 | 当前删除 `/api/admin/accounts/reset-usage`，未发现等价管理端入口。 | TS 有 `/auth/accounts/:id/reset-usage`。 | 功能缺失但不是核心链路。主要影响调试、演示和清理异常统计。 | P3 低优先级恢复，或明确产品不提供手动重置。 | +| 批量健康检查 | 当前删除 `/api/admin/accounts/health-check`，只保留 `/api/admin/accounts/test?id=...` 单账号 SSE 模型测试。 | TS 有 `/auth/accounts/health-check`，支持 ids、stagger、concurrency。 | 功能缺失。单账号测试不能替代账号池健康巡检。 | P1 恢复批量健康检查，复用现有单账号探测和 loading 流程。 | +| 测试模型列表 | 当前保留 `/api/admin/accounts/models?id=...`，从上游返回模型。 | TS 模型列表能力来自模型 catalog 和上游刷新。 | 不缺失。 | 保留当前接口名 `models`。 | +| OAuth 添加账号 | 当前新增 `/api/admin/accounts/oauth/authorize`、`/api/admin/accounts/oauth/exchange`。 | TS 有 `/auth/accounts/login` 和 OAuth 链路。 | 不缺失。Rust 当前已按后台弹窗流程实现。 | 保留。 | +| 手动 Cookie 维护 | 当前删除 `/api/admin/accounts/cookies`。Cookie 存储仍存在，Cloudflare 逻辑会删除账号 cookie，但没有管理端查看/写入入口。 | TS 有 `/auth/accounts/:id/cookies` GET/POST/DELETE，并在 quota 失败时提示写入 cf cookie。 | 功能缺失。遇到 Cloudflare 挑战时缺少人工修复入口。 | P2 恢复，接口保持无路径 id，例如 `GET/POST /api/admin/accounts/cookies?id=...`，删除可并入 POST action 或独立 delete。 | +| 配额告警列表 | 当前删除 `/api/admin/accounts/quota-warnings`。账号列表 summary 仍有 high_usage/attention，但没有独立 warning 列表。 | TS 有 `/auth/quota/warnings`。 | 部分缺失。列表页摘要不能替代可点击、可定位、可清理的告警集合。 | P2 视 UI 是否需要恢复；如果不做告警中心，应在产品上明确删除。 | +| 独立账号用量分页 | 当前删除 `/api/admin/usage`，账号列表会内嵌账号用量，Dashboard 会读取 usage summary 和 buckets。 | TS 有 `/admin/usage-stats/summary` 与 `/admin/usage-stats/history`。 | 部分缺失。概览和账号列表替代了展示，但没有独立用量页/API。 | P2 如果后续有用量分析页，应恢复为 `/api/admin/usage` 和 `/api/admin/usage/summary`；否则可接受删除。 | +| 用量历史 | 当前 Dashboard 趋势读 `usage_time_buckets`，但没有通用 history API。 | TS 有 `/admin/usage-stats/history`，支持 granularity 和 hours。 | 部分缺失。当前只服务 Dashboard，不适合复用到独立分析页面。 | P2 如要开放分析视图，新增 `/api/admin/usage/history`，直接读聚合事实表。 | +| 手动模型刷新 | 当前删除 `/api/admin/models/refresh`，但 `ModelRefreshTask` 会启动后和每小时刷新，并有 `refresh_once()`。 | TS 有 `/admin/refresh-models` 手动触发。 | 功能缺失但已有后台替代。问题在于管理端无法立即刷新模型目录。 | P2 恢复轻量接口，直接调用现有 `refresh_once()` 或后台任务句柄，避免再引入 `admin_models` 旧模块。 | +| OpenAI 模型列表 | 当前 `/v1/models`、`/v1/models/catalog`、`/v1/models/{model}`、`/v1/models/{model}/info` 仍存在。 | TS 有同类模型路由。 | 不缺失。 | 保留。 | +| 日志列表/详情/清空 | 当前保留 `/api/admin/logs`、`/api/admin/logs/detail`、`/api/admin/logs/delete`。 | TS 有 `/admin/logs`、`/admin/logs/:id`、`/admin/logs/clear`。 | 不缺失。Rust 当前筛选字段更丰富。 | 保留当前形态。 | +| 日志开关状态 | 当前删除 `/api/admin/logs/state`。 | TS 有 `/admin/logs/state` GET/POST。 | 不算缺失。此前产品已决定日志采集和容量走默认策略，不在系统设置暴露。 | 不恢复。 | +| 设置项大而全 | 当前 `/api/admin/settings` 只保留 modelAliases、modelAccountRoutes、refreshMarginSeconds、refreshConcurrency、maxConcurrentPerAccount、requestIntervalMs、rotationStrategy。 | TS `general-settings`/`quota-settings` 暴露端口、TLS、默认模型、日志容量、更新策略等大量字段。 | 不算缺失。当前产品已收敛为只保留需要人工决策的运行配置。 | 不恢复 TS 大而全设置页。 | +| 客户端 API Key 列表/创建/更新/删除 | 当前保留 `/api/admin/keys`、`/api/admin/keys/update`、`/api/admin/keys/delete`，并保存完整 `sk_` key 供复制。 | TS 的 `/auth/api-keys` 是第三方 provider key pool，领域不同。 | 当前核心能力不缺失。 | 保留当前“客户端访问密钥”定位，不混入 provider key pool。 | +| 客户端 API Key 导出 | 当前删除 `/api/admin/keys/export`。旧 Rust 实现曾导出 `rustLocalClientApiKeys` 并标记 `rotation_required`。 | TS provider key pool 有 `/auth/api-keys/export`。 | 功能缺失但需谨慎。当前项目已允许保存完整 key，导出可以服务备份；但也会扩大密钥泄露面。 | P3 如恢复，只做管理员显式导出，并使用 CPR 格式；不恢复旧 `rustLocalClientApiKeys` 命名。 | + +## 不建议恢复的删除项 + +以下删除符合当前项目方向，不建议回滚： + +1. `/api/admin/logs/state`：日志策略已固定，系统设置不再暴露日志采集和容量。 +2. TS 版 `general-settings` 中的默认模型、默认推理强度、默认 service tier：这些应由客户端请求决定，网关不兜底。 +3. TS 版模型 aliases 的旧配置形态：当前 Rust 已拆成 `modelAliases` 和 `modelAccountRoutes`，边界更清楚。 +4. 路径中携带 id 的管理端接口形态：当前 Rust 统一使用 query/body 传 id，更符合近期接口规范。 +5. 第三方 provider API Key pool：当前 Rust 的 API Key 是客户端访问密钥，不应混入 TS 版 provider key pool 概念。 + +## 建议恢复优先级 + +P1：应该优先补齐，直接影响管理闭环。 + +- 账号导出：`GET /api/admin/accounts/export` +- 批量账号状态更新：复用 `/api/admin/accounts/update` +- 批量健康检查：`POST /api/admin/accounts/health-check` + +P2：影响运维体验和可诊断性，建议排第二批。 + +- 手动 Cookie 维护：`GET/POST /api/admin/accounts/cookies` +- 配额告警列表：`GET /api/admin/accounts/quota-warnings` +- 独立用量统计和历史：`GET /api/admin/usage`、`GET /api/admin/usage/summary`、可选 `GET /api/admin/usage/history` +- 手动模型刷新：`POST /api/admin/models/refresh` + +P3：低频能力，可先产品确认。 + +- 重置账号用量：`POST /api/admin/accounts/reset-usage` +- 客户端 API Key 导出：`GET /api/admin/keys/export` + +## 恢复原则 + +如果恢复上述能力，不应该把旧代码原样搬回来。需要按当前项目规范重写： + +1. 不在路径中放 id，id 统一放 query 或 body。 +2. 不恢复兼容命名，例如 `native`、`rustLocalClientApiKeys`。 +3. 不恢复大而全的 settings/config 写回逻辑。 +4. 不为了兼容旧数据增加迁移分支；已有开发库需要重建。 +5. 账号导出格式只使用 CPR，和当前导入格式闭环。 +6. 批量操作必须返回 `updated/deleted/notFound` 这类可被 UI 准确呈现的数据。 +7. 涉及真实网络请求的操作必须有后端测试覆盖，包括成功、部分失败、空 ids、未登录。 + +## 需要补的测试 + +当前被删除的测试里包含 `tests/admin/models/*`、`tests/admin/keys/import_export.rs`、旧账号 import/export 测试等。若恢复能力，至少补回以下测试面： + +- 账号导出：全部导出、按 ids 导出、无效格式、未登录。 +- 批量状态：部分 id 不存在、空 ids、非法 status、状态写入后账号列表可见。 +- 批量健康检查：成功、失败、跳过、并发参数校验。 +- Cookie 维护：写入、读取、清空、账号不存在。 +- 用量接口：分页、汇总、空数据库。 +- 模型刷新：无账号、部分计划失败、成功刷新并写入模型快照。 +- API Key 导出：未登录、空列表、导出字段包含完整 key。 diff --git a/src/admin/accounts/routes.rs b/src/admin/accounts/routes.rs index 69684f2d5..6827dfe43 100644 --- a/src/admin/accounts/routes.rs +++ b/src/admin/accounts/routes.rs @@ -13,7 +13,6 @@ use axum::{ use std::collections::HashMap; use chrono::{DateTime, Utc}; -use secrecy::ExposeSecret; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -31,11 +30,11 @@ use crate::{ http::middleware::request_id::RequestId, infra::{ json::{clamp_limit, clamp_page, total_pages, Page}, - time::{china_datetime, china_relative_time, china_rfc3339, china_rfc3339_str}, + time::{china_datetime, china_relative_time, china_rfc3339}, }, runtime::state::AppState, upstream::accounts::model::AccountStatus, - upstream::accounts::store::{AccountModelUsageRecord, StoredAccount}, + upstream::accounts::store::AccountModelUsageRecord, }; const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200; @@ -43,10 +42,6 @@ const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000; const WEEK_WINDOW_SECONDS: u64 = 604_800; const MONTH_WINDOW_SECONDS: u64 = 2_592_000; -// ============================================================================ -// Query / Request types -// ============================================================================ - #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct AccountsQuery { @@ -70,13 +65,6 @@ pub(crate) struct BatchDeleteAccountsRequest { ids: Vec, } -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct SetAccountCookiesRequest { - id: String, - cookies: Value, -} - #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct AccountActionRequest { @@ -89,21 +77,6 @@ pub(crate) struct AccountIdQuery { id: String, } -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AccountExportQuery { - ids: Option, - format: Option, -} - -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase", deny_unknown_fields)] -pub(crate) struct HealthCheckRequest { - ids: Option>, - stagger_ms: Option, - concurrency: Option, -} - #[derive(Debug, Default, Deserialize)] #[serde(rename_all = "camelCase")] pub(crate) struct AccountTestRequest { @@ -122,21 +95,17 @@ pub(crate) struct AccountOAuthExchangeRequest { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -pub(crate) struct AccountTestModelsData { - models: Vec, +struct AccountModelsData { + models: Vec, } #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -pub(crate) struct AccountTestModelData { +struct AccountModelData { id: String, label: String, } -// ============================================================================ -// Response types -// ============================================================================ - #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] struct AdminAccountData { @@ -415,74 +384,10 @@ struct BatchDeleteAccountsData { not_found: Vec, } -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct BatchUpdateAccountStatusData { - updated: u32, - not_found: Vec, -} - #[derive(Debug, Clone, Serialize)] #[serde(untagged)] enum AccountUpdateData { Account(Box), - BatchStatus(BatchUpdateAccountStatusData), -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct ResetAccountUsageData { - id: String, - reset: bool, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountCookiesData { - cookies: Option, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminAccountExportData { - id: String, - email: Option, - account_id: Option, - user_id: Option, - label: Option, - plan_type: Option, - status: String, - access_token_expires_at: Option, - added_at: String, - updated_at: String, - token: String, - refresh_token: Option, -} - -impl From for AdminAccountExportData { - fn from(a: StoredAccount) -> Self { - Self { - id: a.id, - email: a.email, - account_id: a.account_id, - user_id: a.user_id, - label: a.label, - plan_type: a.plan_type, - status: account_status_str(a.status).to_string(), - access_token_expires_at: a.access_token_expires_at.map(|dt| china_rfc3339(&dt)), - added_at: china_rfc3339_str(&a.added_at), - updated_at: china_rfc3339_str(&a.updated_at), - token: a.access_token.expose_secret().to_string(), - refresh_token: a.refresh_token.map(|t| t.expose_secret().to_string()), - } - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountExportData { - source_format: &'static str, - accounts: Vec, } #[derive(Debug, Clone, Serialize)] @@ -502,35 +407,13 @@ struct AccountOAuthAuthorizeData { expires_at_display: String, } -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct HealthCheckData { - summary: HealthCheckSummary, - results: Vec, -} - -#[derive(Debug, Clone, Copy, Serialize)] -#[serde(rename_all = "camelCase")] -struct HealthCheckSummary { - total: usize, - alive: usize, - dead: usize, - skipped: usize, -} - -// ============================================================================ -// Handlers -// ============================================================================ - /// `GET /api/admin/accounts` pub(crate) async fn accounts( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Query(params): Query, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50)); let use_numbered_page = params.page.is_some() || params.page_size.is_some(); let stats = account_list_stats(&state).await; @@ -559,7 +442,7 @@ pub(crate) async fn accounts( AdminAccountPageEnvelope::numbered(page, summary), )) } - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), }; } @@ -583,19 +466,17 @@ pub(crate) async fn accounts( AdminAccountPageEnvelope::cursor(page, limit, summary), )) } - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), } } /// `POST /api/admin/accounts` pub(crate) async fn create_account( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state .services .admin_accounts @@ -604,22 +485,20 @@ pub(crate) async fn create_account( { Ok(account) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminAccountData::from(account), request_id), + AdminEnvelope::ok(AdminAccountData::from(account)), )), - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), } } /// `POST /api/admin/accounts/refresh` pub(crate) async fn refresh_account( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); let account_id = payload.id; - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state .services .admin_accounts @@ -628,101 +507,28 @@ pub(crate) async fn refresh_account( { Ok(account) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminAccountData::from(account), request_id), + AdminEnvelope::ok(AdminAccountData::from(account)), )), - Err(error) => Err(account_error(error, request_id)), - } -} - -/// `POST /api/admin/accounts/reset-usage` -pub(crate) async fn reset_account_usage( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> Result { - let request_id = request_id.as_str().to_string(); - let account_id = payload.id; - require_admin_session(&state, &headers, &request_id).await?; - match state.services.admin_accounts.reset_usage(&account_id).await { - Ok(account) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - ResetAccountUsageData { - id: account.id, - reset: true, - }, - request_id, - ), - )), - Err(error) => Err(account_error(error, request_id)), - } -} - -/// `GET /api/admin/accounts/cookies` -pub(crate) async fn get_account_cookies( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Query(query): Query, -) -> Result { - let request_id = request_id.as_str().to_string(); - let account_id = query.id; - require_admin_session(&state, &headers, &request_id).await?; - match state.services.admin_accounts.cookies(&account_id).await { - Ok(cookies) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountCookiesData { cookies }, request_id), - )), - Err(error) => Err(account_error(error, request_id)), - } -} - -/// `POST /api/admin/accounts/cookies` -pub(crate) async fn set_account_cookies( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> Result { - let request_id = request_id.as_str().to_string(); - let account_id = payload.id; - require_admin_session(&state, &headers, &request_id).await?; - match state - .services - .admin_accounts - .set_cookies(&account_id, payload.cookies) - .await - { - Ok(cookies) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountCookiesData { cookies }, request_id), - )), - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), } } /// `GET /api/admin/accounts/quota` pub(crate) async fn account_quota( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Query(query): Query, ) -> Result { - let request_id = request_id.as_str().to_string(); let account_id = query.id; - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state .services .admin_accounts .account_quota(&account_id) .await { - Ok(data) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(data, request_id), - )), - Err(AdminAccountError::NotFound) => Err(account_not_found(request_id)), + Ok(data) => Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(data))), + Err(AdminAccountError::NotFound) => Err(account_not_found()), Err(AdminAccountError::Inactive(status)) => Err(AdminError::new( StatusCode::CONFLICT, 40901, @@ -730,180 +536,63 @@ pub(crate) async fn account_quota( "Account is {}, cannot query quota", account_status_str(status) ), - request_id, )), Err(AdminAccountError::FetchQuota(msg)) => Err(AdminError::new( StatusCode::BAD_GATEWAY, 50201, format!("Failed to fetch quota from Codex API: {msg}"), - request_id, )), - Err(e) => Err(account_error(e, request_id)), - } -} - -/// `GET /api/admin/accounts/quota-warnings` -pub(crate) async fn quota_warnings( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, -) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; - match state.services.admin_accounts.quota_warnings().await { - Ok(data) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(data, request_id), - )), - Err(error) => Err(account_error(error, request_id)), - } -} - -/// `POST /api/admin/accounts/health-check` -pub(crate) async fn health_check_accounts( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - body: Bytes, -) -> Result { - let request_id = request_id.as_str().to_string(); - let payload = parse_health_check_request(&body, &request_id)?; - validate_health_check_request(&payload, &request_id)?; - require_admin_session(&state, &headers, &request_id).await?; - - let mut req = serde_json::json!({}); - if let Some(ids) = &payload.ids { - req["ids"] = serde_json::json!(ids); - } - - match state - .services - .admin_accounts - .health_check_accounts(req) - .await - { - Ok(result) => { - let summary = HealthCheckSummary { - total: result - .get("summary") - .and_then(|s| s.get("total").and_then(Value::as_u64)) - .unwrap_or(0) as usize, - alive: result - .get("summary") - .and_then(|s| s.get("alive").and_then(Value::as_u64)) - .unwrap_or(0) as usize, - dead: result - .get("summary") - .and_then(|s| s.get("dead").and_then(Value::as_u64)) - .unwrap_or(0) as usize, - skipped: result - .get("summary") - .and_then(|s| s.get("skipped").and_then(Value::as_u64)) - .unwrap_or(0) as usize, - }; - let results = result - .get("results") - .and_then(Value::as_array) - .cloned() - .unwrap_or_default(); - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(HealthCheckData { summary, results }, request_id), - )) - } - Err(error) => Err(account_error(error, request_id)), - } -} - -/// `GET /api/admin/accounts/export` -pub(crate) async fn export_accounts( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Query(query): Query, -) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; - validate_account_export_format(query.format.as_deref()) - .map_err(|msg| AdminError::new(StatusCode::BAD_REQUEST, 40001, msg, request_id.clone()))?; - let ids = account_export_ids(query.ids.as_deref()); - match state.services.admin_accounts.export_with_tokens(ids).await { - Ok(accounts) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AccountExportData { - source_format: "native", - accounts: accounts - .into_iter() - .map(AdminAccountExportData::from) - .collect(), - }, - request_id, - ), - )), - Err(error) => Err(account_error(error, request_id)), + Err(e) => Err(account_error(&e)), } } /// `POST /api/admin/accounts/import` pub(crate) async fn import_accounts( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state.services.admin_accounts.import(payload).await { Ok(result) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - AccountImportData { - imported: result.imported, - skipped: result.skipped, - source_format: result.source_format.to_string(), - }, - request_id, - ), + AdminEnvelope::ok(AccountImportData { + imported: result.imported, + skipped: result.skipped, + source_format: result.source_format.to_string(), + }), )), - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), } } /// `POST /api/admin/accounts/oauth/authorize` pub(crate) async fn oauth_authorize_account( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state.services.admin_accounts.oauth_authorize().await { Ok(result) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - AccountOAuthAuthorizeData { - session_id: result.session_id, - auth_url: result.auth_url, - expires_at: china_rfc3339(&result.expires_at), - expires_at_display: china_datetime(&result.expires_at), - }, - request_id, - ), + AdminEnvelope::ok(AccountOAuthAuthorizeData { + session_id: result.session_id, + auth_url: result.auth_url, + expires_at: china_rfc3339(&result.expires_at), + expires_at_display: china_datetime(&result.expires_at), + }), )), - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), } } /// `POST /api/admin/accounts/oauth/exchange` pub(crate) async fn oauth_exchange_account( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state .services .admin_accounts @@ -917,16 +606,13 @@ pub(crate) async fn oauth_exchange_account( { Ok(result) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - AccountImportData { - imported: result.imported, - skipped: result.skipped, - source_format: result.source_format.to_string(), - }, - request_id, - ), + AdminEnvelope::ok(AccountImportData { + imported: result.imported, + skipped: result.skipped, + source_format: result.source_format.to_string(), + }), )), - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), } } @@ -934,14 +620,12 @@ pub(crate) async fn oauth_exchange_account( pub(crate) async fn test_account_connection( State(state): State, Query(query): Query, - Extension(request_id): Extension, headers: HeaderMap, body: Bytes, ) -> Result { - let request_id = request_id.as_str().to_string(); let account_id = query.id; - let payload = parse_account_test_request(&body, &request_id)?; - require_admin_session(&state, &headers, &request_id).await?; + let payload = parse_account_test_request(&body)?; + require_admin_session(&state, &headers).await?; let model = payload .model_id @@ -949,20 +633,13 @@ pub(crate) async fn test_account_connection( .map(str::trim) .filter(|value| !value.is_empty()) .map(ToString::to_string) - .ok_or_else(|| { - AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - "Model is required", - &request_id, - ) - })?; + .ok_or_else(|| AdminError::new(StatusCode::BAD_REQUEST, 40001, "Model is required"))?; let stream = state .services .admin_accounts .test_connection_stream(&account_id, model) .await - .map_err(|error| account_error(error, request_id.clone()))?; + .map_err(|error| account_error(&error))?; Response::builder() .status(StatusCode::OK) @@ -976,13 +653,12 @@ pub(crate) async fn test_account_connection( StatusCode::INTERNAL_SERVER_ERROR, 50001, "Failed to build account test stream", - request_id, ) }) } /// `GET /api/admin/accounts/models?id=...` -pub(crate) async fn account_test_models( +pub(crate) async fn account_models( State(state): State, Query(query): Query, Extension(request_id): Extension, @@ -990,41 +666,36 @@ pub(crate) async fn account_test_models( ) -> Result { let request_id = request_id.as_str().to_string(); let account_id = query.id; - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state .services .admin_accounts - .test_connection_models(&account_id, &request_id) + .account_models(&account_id, &request_id) .await { Ok(models) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - AccountTestModelsData { - models: models - .into_iter() - .map(|model| AccountTestModelData { - id: model.id, - label: model.label, - }) - .collect(), - }, - request_id, - ), + AdminEnvelope::ok(AccountModelsData { + models: models + .into_iter() + .map(|model| AccountModelData { + id: model.id, + label: model.label, + }) + .collect(), + }), )), - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), } } /// `POST /api/admin/accounts/delete` pub(crate) async fn batch_delete_accounts( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state .services .admin_accounts @@ -1033,65 +704,37 @@ pub(crate) async fn batch_delete_accounts( { Ok(result) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - BatchDeleteAccountsData { - deleted: result.deleted, - not_found: result.not_found, - }, - request_id, - ), + AdminEnvelope::ok(BatchDeleteAccountsData { + deleted: result.deleted, + not_found: result.not_found, + }), )), - Err(error) => Err(account_error(error, request_id)), + Err(error) => Err(account_error(&error)), } } /// `POST /api/admin/accounts/update` pub(crate) async fn update_account( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; - match parse_account_update(payload, &request_id)? { - ParsedAccountUpdate::Single { id, update } => { - match state - .services - .admin_accounts - .update_metadata(&id, update) - .await - { - Ok(Some(account)) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AccountUpdateData::Account(Box::new(account.into())), - request_id, - ), - )), - Ok(None) => Err(account_not_found(request_id)), - Err(error) => Err(account_error(error, request_id)), - } - } - ParsedAccountUpdate::BatchStatus { ids, status } => match state - .services - .admin_accounts - .batch_update_status(ids, &status) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - AccountUpdateData::BatchStatus(BatchUpdateAccountStatusData { - updated: result.updated, - not_found: result.not_found, - }), - request_id, - ), - )), - Err(error) => Err(account_error(error, request_id)), - }, + let ParsedAccountUpdate { id, update } = parse_account_update(&payload)?; + + match state + .services + .admin_accounts + .update_metadata(&id, update) + .await + { + Ok(Some(account)) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountUpdateData::Account(Box::new(account.into()))), + )), + Ok(None) => Err(account_not_found()), + Err(error) => Err(account_error(&error)), } } @@ -1345,13 +988,10 @@ fn push_monthly_quota_window( window_seconds, label_display: "月限额".to_string(), used_percent, - used_percent_display: used_percent - .map(format_percent) - .unwrap_or_else(|| "-".to_string()), + used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent), reset_at_display: reset_at .as_ref() - .map(china_datetime) - .unwrap_or_else(|| "-".to_string()), + .map_or_else(|| "-".to_string(), china_datetime), window_used_display: quota_window_used_display(reset_at, window_seconds), }); true @@ -1430,13 +1070,10 @@ fn push_quota_window( window_seconds, label_display, used_percent, - used_percent_display: used_percent - .map(format_percent) - .unwrap_or_else(|| "-".to_string()), + used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent), reset_at_display: reset_at .as_ref() - .map(china_datetime) - .unwrap_or_else(|| "-".to_string()), + .map_or_else(|| "-".to_string(), china_datetime), window_used_display: quota_window_used_display(reset_at, window_seconds), }); } @@ -1569,7 +1206,8 @@ fn quota_window_used_display( let remaining = reset_at .signed_duration_since(Utc::now()) .num_seconds() - .max(0) as u64; + .max(0) + .cast_unsigned(); let used = window_seconds.saturating_sub(remaining); format!( "{} / {}", @@ -1634,66 +1272,28 @@ fn format_cost(value: f64) -> String { format!("${value:.2}") } -// ============================================================================ -// Error handling -// ============================================================================ - -enum ParsedAccountUpdate { - Single { - id: String, - update: AdminAccountMetadataUpdate, - }, - BatchStatus { - ids: Vec, - status: String, - }, +struct ParsedAccountUpdate { + id: String, + update: AdminAccountMetadataUpdate, } -fn parse_account_update( - payload: Value, - request_id: &str, -) -> Result { +fn parse_account_update(payload: &Value) -> Result { let object = payload.as_object().ok_or_else(|| { AdminError::new( StatusCode::BAD_REQUEST, 40001, "Account update request must be an object", - request_id, ) })?; - if object.contains_key("ids") { - if object.contains_key("id") { - return Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - "Account batch update must not include id", - request_id, - )); - } - if ["label", "email", "accountId", "userId", "planType"] - .iter() - .any(|field| object.contains_key(*field)) - { - return Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - "Account batch update only supports status", - request_id, - )); - } - let ids = required_string_array_field(object, "ids", request_id)?; - let status = required_string_field(object, "status", request_id)?; - return Ok(ParsedAccountUpdate::BatchStatus { ids, status }); - } - let id = required_string_field(object, "id", request_id)?; - let label = optional_string_update_field(object, "label", request_id)?; - let email = optional_string_update_field(object, "email", request_id)?; - let account_id = optional_string_update_field(object, "accountId", request_id)?; - let user_id = optional_string_update_field(object, "userId", request_id)?; - let plan_type = optional_string_update_field(object, "planType", request_id)?; + let id = required_string_field(object, "id")?; + let label = optional_string_update_field(object, "label")?; + let email = optional_string_update_field(object, "email")?; + let account_id = optional_string_update_field(object, "accountId")?; + let user_id = optional_string_update_field(object, "userId")?; + let plan_type = optional_string_update_field(object, "planType")?; let status = object .get("status") - .map(|value| required_string_value(value, "status", request_id)) + .map(|value| required_string_value(value, "status")) .transpose()?; let update = AdminAccountMetadataUpdate { email, @@ -1708,33 +1308,26 @@ fn parse_account_update( StatusCode::BAD_REQUEST, 40001, "Account update request must include editable fields", - request_id, )); } - Ok(ParsedAccountUpdate::Single { id, update }) + Ok(ParsedAccountUpdate { id, update }) } fn required_string_field( object: &serde_json::Map, field: &'static str, - request_id: &str, ) -> Result { let Some(value) = object.get(field) else { return Err(AdminError::new( StatusCode::BAD_REQUEST, 40001, format!("{field} is required"), - request_id, )); }; - required_string_value(value, field, request_id) + required_string_value(value, field) } -fn required_string_value( - value: &Value, - field: &'static str, - request_id: &str, -) -> Result { +fn required_string_value(value: &Value, field: &'static str) -> Result { value .as_str() .map(str::trim) @@ -1745,51 +1338,11 @@ fn required_string_value( StatusCode::BAD_REQUEST, 40001, format!("{field} must be a non-empty string"), - request_id, ) }) } -fn required_string_array_field( - object: &serde_json::Map, - field: &'static str, - request_id: &str, -) -> Result, AdminError> { - let Some(value) = object.get(field) else { - return Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - format!("{field} is required"), - request_id, - )); - }; - let Some(values) = value.as_array() else { - return Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - format!("{field} must be an array of non-empty strings"), - request_id, - )); - }; - if values.is_empty() { - return Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - format!("{field} must not be empty"), - request_id, - )); - } - values - .iter() - .map(|value| required_string_value(value, field, request_id)) - .collect() -} - -fn optional_string_field( - value: &Value, - field: &'static str, - request_id: &str, -) -> Result, AdminError> { +fn optional_string_field(value: &Value, field: &'static str) -> Result, AdminError> { if value.is_null() { return Ok(None); } @@ -1802,7 +1355,6 @@ fn optional_string_field( StatusCode::BAD_REQUEST, 40001, format!("{field} must be a string or null"), - request_id, )), } } @@ -1810,15 +1362,14 @@ fn optional_string_field( fn optional_string_update_field( object: &serde_json::Map, field: &'static str, - request_id: &str, ) -> Result>, AdminError> { object .get(field) - .map(|value| optional_string_field(value, field, request_id)) + .map(|value| optional_string_field(value, field)) .transpose() } -fn account_error(error: AdminAccountError, request_id: String) -> AdminError { +fn account_error(error: &AdminAccountError) -> AdminError { match error { AdminAccountError::InvalidStatus(_) | AdminAccountError::LabelTooLong @@ -1832,38 +1383,22 @@ fn account_error(error: AdminAccountError, request_id: String) -> AdminError { | AdminAccountError::OAuthSessionInvalid | AdminAccountError::OAuthCallbackInvalid | AdminAccountError::OAuthStateMismatch - | AdminAccountError::NoValidCookies => AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - error.to_string(), - request_id, - ), - AdminAccountError::OAuthCodeExchange(_) => AdminError::new( - StatusCode::BAD_GATEWAY, - 50201, - error.to_string(), - request_id, - ), - AdminAccountError::NotFound => account_not_found(request_id), - AdminAccountError::Inactive(_) => { - AdminError::new(StatusCode::CONFLICT, 40901, error.to_string(), request_id) + | AdminAccountError::NoValidCookies => { + AdminError::new(StatusCode::BAD_REQUEST, 40001, error.to_string()) } - _ => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - error.to_string(), - request_id, - ), + AdminAccountError::OAuthCodeExchange(_) => { + AdminError::new(StatusCode::BAD_GATEWAY, 50201, error.to_string()) + } + AdminAccountError::NotFound => account_not_found(), + AdminAccountError::Inactive(_) => { + AdminError::new(StatusCode::CONFLICT, 40901, error.to_string()) + } + _ => AdminError::new(StatusCode::INTERNAL_SERVER_ERROR, 50001, error.to_string()), } } -fn account_not_found(request_id: String) -> AdminError { - AdminError::new( - StatusCode::NOT_FOUND, - 40401, - "Account not found", - request_id, - ) +fn account_not_found() -> AdminError { + AdminError::new(StatusCode::NOT_FOUND, 40401, "Account not found") } fn account_status_str(status: crate::upstream::accounts::model::AccountStatus) -> &'static str { @@ -1877,27 +1412,7 @@ fn account_status_str(status: crate::upstream::accounts::model::AccountStatus) - } } -fn parse_health_check_request( - body: &Bytes, - request_id: &str, -) -> Result { - if body.is_empty() { - return Ok(HealthCheckRequest::default()); - } - serde_json::from_slice(body).map_err(|_| { - AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - "Invalid health check request", - request_id, - ) - }) -} - -fn parse_account_test_request( - body: &Bytes, - request_id: &str, -) -> Result { +fn parse_account_test_request(body: &Bytes) -> Result { if body.is_empty() { return Ok(AccountTestRequest::default()); } @@ -1906,61 +1421,6 @@ fn parse_account_test_request( StatusCode::BAD_REQUEST, 40001, "Invalid account test request", - request_id, ) }) } - -fn validate_health_check_request( - payload: &HealthCheckRequest, - request_id: &str, -) -> Result<(), AdminError> { - if payload.ids.as_ref().is_some_and(Vec::is_empty) { - return Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - "Account ids must not be empty", - request_id, - )); - } - if payload - .stagger_ms - .is_some_and(|value| !(500..=30_000).contains(&value)) - { - return Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - "staggerMs must be between 500 and 30000", - request_id, - )); - } - if payload - .concurrency - .is_some_and(|value| !(1..=10).contains(&value)) - { - return Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - "concurrency must be between 1 and 10", - request_id, - )); - } - Ok(()) -} - -fn validate_account_export_format(value: Option<&str>) -> Result<(), &'static str> { - match value.unwrap_or("native").trim() { - "" | "native" => Ok(()), - _ => Err("Unsupported account export format"), - } -} - -fn account_export_ids(value: Option<&str>) -> Vec { - value - .into_iter() - .flat_map(|ids| ids.split(',')) - .map(str::trim) - .filter(|id| !id.is_empty()) - .map(ToString::to_string) - .collect() -} diff --git a/src/admin/accounts/service/cookies.rs b/src/admin/accounts/service/cookies.rs index 8ab4a73ab..29b0df90a 100644 --- a/src/admin/accounts/service/cookies.rs +++ b/src/admin/accounts/service/cookies.rs @@ -1,4 +1,4 @@ -use super::{types::*, AdminAccountService}; +use super::AdminAccountService; impl AdminAccountService { pub(crate) async fn usage_cookie_header(&self, account_id: &str) -> Option { @@ -8,70 +8,4 @@ impl AdminAccountService { .ok() .flatten() } - - pub async fn cookies(&self, account_id: &str) -> Result, AdminAccountError> { - self.ensure_cookie_account_exists(account_id).await?; - self.cookies - .cookie_header(account_id, "chatgpt.com") - .await - .map_err(|_| AdminAccountError::LoadCookies) - } - pub async fn set_cookies( - &self, - account_id: &str, - cookies: serde_json::Value, - ) -> Result, AdminAccountError> { - let cookie_header = match cookies { - serde_json::Value::String(ref s) => s.trim().to_string(), - serde_json::Value::Object(ref obj) => { - let pairs: Vec = obj - .iter() - .filter_map(|(name, val)| { - let v = val.as_str()?.trim(); - if name.trim().is_empty() || v.is_empty() { - return None; - } - Some(format!("{}={}", name.trim(), v)) - }) - .collect(); - if pairs.is_empty() { - return Err(AdminAccountError::NoValidCookies); - } - pairs.join("; ") - } - _ => return Err(AdminAccountError::NoValidCookies), - }; - self.ensure_cookie_account_exists(account_id).await?; - match self - .cookies - .set_cookie_header(account_id, &cookie_header) - .await - { - Ok(0) => Err(AdminAccountError::NoValidCookies), - Ok(_) => self - .cookies - .cookie_header(account_id, "chatgpt.com") - .await - .map_err(|_| AdminAccountError::LoadCookies), - Err(_) => Err(AdminAccountError::StoreCookies), - } - } - pub async fn delete_cookies(&self, account_id: &str) -> Result<(), AdminAccountError> { - self.ensure_cookie_account_exists(account_id).await?; - self.cookies - .delete_account_cookies(account_id) - .await - .map(|_| ()) - .map_err(|_| AdminAccountError::DeleteCookies) - } - async fn ensure_cookie_account_exists( - &self, - account_id: &str, - ) -> Result<(), AdminAccountError> { - match self.cookies.account_exists(account_id).await { - Ok(true) => Ok(()), - Ok(false) => Err(AdminAccountError::NotFound), - Err(_) => Err(AdminAccountError::Inspect), - } - } } diff --git a/src/admin/accounts/service/importing.rs b/src/admin/accounts/service/importing.rs index 9d87cde89..b0f9f40fc 100644 --- a/src/admin/accounts/service/importing.rs +++ b/src/admin/accounts/service/importing.rs @@ -1,74 +1,21 @@ use chrono::Utc; use secrecy::SecretString; -use super::{types::*, AdminAccountService}; +use super::{ + types::{ + import_quota_plan_type, import_status_from_usage_error, import_usage_plan_type, + import_usage_string, AdminAccountError, ImportSupplementalAccountInfo, + ImportSupplementalNeeds, ImportedAccountState, ImportedAccounts, ResolvedImportTokens, + }, + AdminAccountService, +}; impl AdminAccountService { - pub async fn export( - &self, - ids: Vec, - ) -> Result, AdminAccountError> { - if ids.is_empty() { - let mut all = Vec::new(); - let mut cursor = None; - loop { - let page = self - .store - .list_metadata(cursor, 200) - .await - .map_err(|_| AdminAccountError::Export)?; - all.extend(page.items.into_iter().map(AdminAccountMetadata::from)); - if page.next_cursor.is_none() { - break; - } - cursor = page.next_cursor; - } - Ok(all) - } else { - let mut accounts = Vec::with_capacity(ids.len()); - for id in ids { - if let Ok(Some(stored)) = self.store.get(&id).await { - accounts.push(stored_to_admin_metadata(stored)); - } - } - Ok(accounts) - } - } - pub async fn export_with_tokens( - &self, - ids: Vec, - ) -> Result, AdminAccountError> { - if ids.is_empty() { - let mut all = Vec::new(); - let mut cursor = None; - loop { - let page = self - .store - .list(cursor, 200) - .await - .map_err(|_| AdminAccountError::Export)?; - all.extend(page.items); - if page.next_cursor.is_none() { - break; - } - cursor = page.next_cursor; - } - Ok(all) - } else { - let mut accounts = Vec::with_capacity(ids.len()); - for id in ids { - if let Ok(Some(stored)) = self.store.get(&id).await { - accounts.push(stored); - } - } - Ok(accounts) - } - } pub async fn import( &self, data: serde_json::Value, ) -> Result { - let parsed = crate::upstream::accounts::import_export::parse_account_import_payload(&data) + let parsed = crate::upstream::accounts::importing::parse_account_import_payload(&data) .map_err(|_| AdminAccountError::NoImportableAccounts)?; let source_format = parsed.source.as_str(); let entries = parsed.entries; @@ -96,8 +43,8 @@ impl AdminAccountService { } async fn import_entry( &self, - entry: crate::upstream::accounts::import_export::AccountImportEntry, - source: crate::upstream::accounts::import_export::AccountImportSource, + entry: crate::upstream::accounts::importing::AccountImportEntry, + source: crate::upstream::accounts::importing::AccountImportSource, ) -> Result { let Some(resolved_tokens) = self .resolve_import_tokens(entry.token, entry.refresh_token) @@ -105,7 +52,7 @@ impl AdminAccountService { else { return Ok(ImportedAccountState::Skipped); }; - let label = crate::upstream::accounts::import_export::normalize_label(entry.label); + let label = crate::upstream::accounts::importing::normalize_label(entry.label); if label .as_ref() .is_some_and(|label| label.chars().count() > 64) @@ -116,13 +63,13 @@ impl AdminAccountService { let access_token_expires_at = entry .access_token_expires_at .as_deref() - .map(crate::upstream::accounts::import_export::parse_account_import_datetime) + .map(crate::upstream::accounts::importing::parse_account_import_datetime) .transpose() .map_err(|_| AdminAccountError::InvalidAccessTokenExpiresAt)?; let quota_fetched_at = entry .quota_fetched_at .as_deref() - .map(crate::upstream::accounts::import_export::parse_account_import_datetime) + .map(crate::upstream::accounts::importing::parse_account_import_datetime) .transpose() .map_err(|_| AdminAccountError::InvalidAccessTokenExpiresAt)?; let mut quota_json = entry @@ -131,16 +78,15 @@ impl AdminAccountService { .map(serde_json::Value::to_string); let mut quota_fetched_at = quota_fetched_at; let quota_verify_required = entry.quota_verify_required.unwrap_or(false); - let parsed_status = crate::upstream::accounts::import_export::parse_account_import_status( + let parsed_status = crate::upstream::accounts::importing::parse_account_import_status( entry.status.as_deref(), ) .map_err(|error| AdminAccountError::InvalidStatus(error.to_string()))?; - let mut status = - crate::upstream::accounts::import_export::normalized_imported_account_status( - parsed_status, - source, - &resolved_tokens.access_token, - ); + let mut status = crate::upstream::accounts::importing::normalized_imported_account_status( + parsed_status, + source, + &resolved_tokens.access_token, + ); let access_token_expires_at = resolved_tokens .claims .as_ref() @@ -151,26 +97,24 @@ impl AdminAccountService { let claims = resolved_tokens.claims.as_ref(); let mut plan_type = claims .and_then(|claims| claims.plan_type.clone()) - .or_else(|| { - crate::upstream::accounts::import_export::normalize_nonempty(entry.plan_type) - }); + .or_else(|| crate::upstream::accounts::importing::normalize_nonempty(entry.plan_type)); if plan_type.is_none() { plan_type = entry.cached_quota.as_ref().and_then(import_quota_plan_type); } let email = claims.and_then(|claims| claims.email.clone()).or_else(|| { - crate::upstream::accounts::import_export::normalize_nonempty(entry.email.clone()) + crate::upstream::accounts::importing::normalize_nonempty(entry.email.clone()) }); let chatgpt_account_id = claims .and_then(|claims| claims.account_id.as_deref()) .or_else(|| { - crate::upstream::accounts::import_export::normalize_nonempty_str( + crate::upstream::accounts::importing::normalize_nonempty_str( entry.account_id.as_deref(), ) }); let chatgpt_user_id = claims .and_then(|claims| claims.user_id.as_deref()) .or_else(|| { - crate::upstream::accounts::import_export::normalize_nonempty_str( + crate::upstream::accounts::importing::normalize_nonempty_str( entry.user_id.as_deref(), ) }); @@ -333,9 +277,11 @@ impl AdminAccountService { refresh_token: Option, ) -> Result, AdminAccountError> { let mut refresh_token = - crate::upstream::accounts::import_export::normalize_nonempty(refresh_token); - let Some(access_token) = crate::upstream::accounts::import_export::normalize_nonempty( - token.map(crate::upstream::accounts::import_export::normalize_bearer_token), + crate::upstream::accounts::importing::normalize_nonempty(refresh_token); + let Some(access_token) = crate::upstream::accounts::importing::normalize_nonempty( + token + .as_deref() + .map(crate::upstream::accounts::importing::normalize_bearer_token), ) else { let Some(existing_refresh_token) = refresh_token else { return Ok(None); @@ -345,9 +291,9 @@ impl AdminAccountService { .refresh(&existing_refresh_token) .await .map_err(AdminAccountError::RefreshTokenExchange)?; - let access_token = crate::upstream::accounts::import_export::normalize_nonempty(Some( - crate::upstream::accounts::import_export::normalize_bearer_token( - refreshed.access_token, + let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( + crate::upstream::accounts::importing::normalize_bearer_token( + &refreshed.access_token, ), )) .ok_or(AdminAccountError::TokenRequired)?; @@ -387,10 +333,8 @@ impl AdminAccountService { .refresh(&existing_refresh_token) .await .map_err(AdminAccountError::RefreshTokenExchange)?; - let access_token = crate::upstream::accounts::import_export::normalize_nonempty(Some( - crate::upstream::accounts::import_export::normalize_bearer_token( - refreshed.access_token, - ), + let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( + crate::upstream::accounts::importing::normalize_bearer_token(&refreshed.access_token), )) .ok_or(AdminAccountError::TokenRequired)?; refresh_token = refreshed.refresh_token.or(Some(existing_refresh_token)); @@ -411,7 +355,7 @@ impl AdminAccountService { account_id: Option<&str>, user_id: Option<&str>, ) -> Result { - let provided_id = crate::upstream::accounts::import_export::normalize_nonempty_str(id) + let provided_id = crate::upstream::accounts::importing::normalize_nonempty_str(id) .map(ToString::to_string); if let Some(id) = provided_id.as_deref() { match self.store.get(id).await { @@ -422,9 +366,8 @@ impl AdminAccountService { } let chatgpt_account_id = - crate::upstream::accounts::import_export::normalize_nonempty_str(account_id); - let chatgpt_user_id = - crate::upstream::accounts::import_export::normalize_nonempty_str(user_id); + crate::upstream::accounts::importing::normalize_nonempty_str(account_id); + let chatgpt_user_id = crate::upstream::accounts::importing::normalize_nonempty_str(user_id); if let Some(chatgpt_account_id) = chatgpt_account_id { if let Some(existing) = self .store @@ -436,8 +379,7 @@ impl AdminAccountService { } } - Ok(provided_id.unwrap_or_else(|| { - crate::upstream::accounts::import_export::normalized_account_id(None) - })) + Ok(provided_id + .unwrap_or_else(|| crate::upstream::accounts::importing::normalized_account_id(None))) } } diff --git a/src/admin/accounts/service/lifecycle.rs b/src/admin/accounts/service/lifecycle.rs index d978d6dbd..0cba0d5ff 100644 --- a/src/admin/accounts/service/lifecycle.rs +++ b/src/admin/accounts/service/lifecycle.rs @@ -3,7 +3,13 @@ use secrecy::{ExposeSecret, SecretString}; use crate::infra::json::{NumberedPage, Page}; -use super::{types::*, AdminAccountService}; +use super::{ + types::{ + parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountMetadata, + AdminAccountMetadataUpdate, BatchDeleteAccounts, ManualCreateTokens, UpdatedAccountStatus, + }, + AdminAccountService, +}; impl AdminAccountService { pub async fn list( @@ -65,10 +71,12 @@ impl AdminAccountService { refresh_token: Option, ) -> Result { let provided_refresh_token = - crate::upstream::accounts::import_export::normalize_nonempty(refresh_token); + crate::upstream::accounts::importing::normalize_nonempty(refresh_token); let tokens = if let Some(access_token) = - crate::upstream::accounts::import_export::normalize_nonempty( - token.map(crate::upstream::accounts::import_export::normalize_bearer_token), + crate::upstream::accounts::importing::normalize_nonempty( + token + .as_deref() + .map(crate::upstream::accounts::importing::normalize_bearer_token), ) { ManualCreateTokens { access_token, @@ -81,9 +89,9 @@ impl AdminAccountService { .refresh(&refresh_token) .await .map_err(AdminAccountError::RefreshTokenExchange)?; - let access_token = crate::upstream::accounts::import_export::normalize_nonempty(Some( - crate::upstream::accounts::import_export::normalize_bearer_token( - token_pair.access_token, + let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( + crate::upstream::accounts::importing::normalize_bearer_token( + &token_pair.access_token, ), )) .ok_or(AdminAccountError::TokenRequired)?; @@ -141,7 +149,7 @@ impl AdminAccountService { } existing.id } else { - let id = crate::upstream::accounts::import_export::normalized_account_id(None); + let id = crate::upstream::accounts::importing::normalized_account_id(None); self.store .insert(crate::upstream::accounts::store::NewAccount { id: id.clone(), @@ -300,13 +308,12 @@ impl AdminAccountService { .await { Ok(tokens) => { - let access_token = - crate::upstream::accounts::import_export::normalize_nonempty(Some( - crate::upstream::accounts::import_export::normalize_bearer_token( - tokens.access_token, - ), - )) - .ok_or(AdminAccountError::TokenRequired)?; + let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( + crate::upstream::accounts::importing::normalize_bearer_token( + &tokens.access_token, + ), + )) + .ok_or(AdminAccountError::TokenRequired)?; let claims = crate::upstream::accounts::token_refresh::manual_account_claims( &access_token, Utc::now(), @@ -351,7 +358,7 @@ impl AdminAccountService { if !updated { return Err(AdminAccountError::NotFound); } - if crate::upstream::accounts::import_export::refresh_failure_status_clears_next_refresh_at( + if crate::upstream::accounts::importing::refresh_failure_status_clears_next_refresh_at( status, ) { let cleared = self @@ -377,47 +384,6 @@ impl AdminAccountService { .map(stored_to_admin_metadata) .ok_or(AdminAccountError::NotFound) } - pub async fn reset_usage( - &self, - account_id: &str, - ) -> Result { - self.store - .reset_usage(account_id) - .await - .map_err(|_| AdminAccountError::ResetUsage)?; - self.sync_account_pool(account_id).await?; - self.store - .get(account_id) - .await - .map_err(|_| AdminAccountError::NotFound)? - .map(stored_to_admin_metadata) - .ok_or(AdminAccountError::NotFound) - } - pub async fn batch_update_status( - &self, - ids: Vec, - status: &str, - ) -> Result { - if ids.is_empty() { - return Err(AdminAccountError::EmptyIds); - } - let status = parse_batch_account_status(status)?; - let mut updated = 0u32; - let mut not_found = Vec::new(); - for id in ids { - match self.store.set_status(&id, status).await { - Ok(true) => { - updated += 1; - self.sync_account_pool_best_effort(&id, "account batch status update") - .await; - self.evict_account_websocket_pool(&id).await; - } - Ok(false) => not_found.push(id), - Err(_) => return Err(AdminAccountError::UpdateStatus), - } - } - Ok(BatchUpdateAccountStatus { updated, not_found }) - } } fn manual_refresh_failure_status( diff --git a/src/admin/accounts/service/mod.rs b/src/admin/accounts/service/mod.rs index 95f187c74..d1d170753 100644 --- a/src/admin/accounts/service/mod.rs +++ b/src/admin/accounts/service/mod.rs @@ -13,7 +13,6 @@ use std::sync::Arc as StdArc; use chrono::{DateTime, Duration}; use crate::{ - config::types::QuotaWarningThresholds, upstream::accounts::{ cookies::SqliteCookieStore, pool::RuntimeAccountPoolService, store::SqliteAccountStore, token_refresh::TokenRefresher, @@ -23,15 +22,13 @@ use crate::{ pub use types::{ AdminAccountError, AdminAccountMetadata, AdminAccountMetadataUpdate, BatchDeleteAccounts, - BatchUpdateAccountStatus, ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, - UpdatedAccountStatus, + ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, UpdatedAccountStatus, }; #[derive(Clone)] pub struct AdminAccountService { pub store: SqliteAccountStore, pub(crate) cookies: SqliteCookieStore, - pub(crate) quota_thresholds: QuotaWarningThresholds, pub(crate) codex: StdArc, pub(crate) account_pool: StdArc, pub(crate) token_refresher: StdArc, @@ -40,32 +37,28 @@ pub struct AdminAccountService { pub(crate) installation_id: Option, } +pub struct AdminAccountServiceParts { + pub store: SqliteAccountStore, + pub cookies: SqliteCookieStore, + pub codex: StdArc, + pub account_pool: StdArc, + pub token_refresher: StdArc, + pub oauth: oauth::AccountOAuthService, + pub refresh_margin_seconds: u64, + pub installation_id: Option, +} + impl AdminAccountService { - #[expect( - clippy::too_many_arguments, - reason = "service constructor wires independent stores and runtime collaborators" - )] - pub fn new( - store: SqliteAccountStore, - cookies: SqliteCookieStore, - quota_thresholds: QuotaWarningThresholds, - codex: StdArc, - account_pool: StdArc, - token_refresher: StdArc, - oauth: oauth::AccountOAuthService, - refresh_margin_seconds: u64, - installation_id: Option, - ) -> Self { + pub fn new(parts: AdminAccountServiceParts) -> Self { Self { - store, - cookies, - quota_thresholds, - codex, - account_pool, - token_refresher, - oauth, - refresh_margin_seconds, - installation_id, + store: parts.store, + cookies: parts.cookies, + codex: parts.codex, + account_pool: parts.account_pool, + token_refresher: parts.token_refresher, + oauth: parts.oauth, + refresh_margin_seconds: parts.refresh_margin_seconds, + installation_id: parts.installation_id, } } diff --git a/src/admin/accounts/service/oauth.rs b/src/admin/accounts/service/oauth.rs index 6c7275efd..d691b20ab 100644 --- a/src/admin/accounts/service/oauth.rs +++ b/src/admin/accounts/service/oauth.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, fmt::Write as _, sync::Arc}; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use chrono::{DateTime, Duration, Utc}; @@ -147,7 +147,7 @@ impl AdminAccountService { ) -> Result { let tokens = self.oauth.exchange(input).await?; self.import(serde_json::json!({ - "sourceFormat": "native", + "sourceFormat": "cpr", "accounts": [{ "token": tokens.access_token, "refreshToken": tokens.refresh_token, @@ -218,7 +218,11 @@ fn random_url_token(size: usize) -> String { fn random_hex_token(size: usize) -> String { let mut bytes = vec![0u8; size]; rand::rng().fill_bytes(&mut bytes); - bytes.iter().map(|byte| format!("{byte:02x}")).collect() + let mut token = String::with_capacity(size * 2); + for byte in bytes { + let _ = write!(token, "{byte:02x}"); + } + token } fn exchange_error_text(status: StatusCode, body: &str) -> String { diff --git a/src/admin/accounts/service/quota.rs b/src/admin/accounts/service/quota.rs index e82c2aa57..1a9b88a3b 100644 --- a/src/admin/accounts/service/quota.rs +++ b/src/admin/accounts/service/quota.rs @@ -1,8 +1,5 @@ -use chrono::Utc; use secrecy::ExposeSecret; -use crate::infra::time::china_rfc3339; - use super::{types::AdminAccountError, AdminAccountService}; impl AdminAccountService { @@ -52,191 +49,4 @@ impl AdminAccountService { } Ok(serde_json::json!({ "quota": normalized, "raw": raw })) } - pub async fn quota_warnings(&self) -> Result { - let snapshots = self - .store - .list_quota_snapshots() - .await - .map_err(|_| AdminAccountError::QuotaWarnings)?; - - let mut warnings = Vec::new(); - for snap in &snapshots { - let quota: serde_json::Value = - serde_json::from_str(&snap.quota_json).unwrap_or(serde_json::Value::Null); - let used = crate::upstream::accounts::quota::quota_snapshot_limit_reached("a); - if used { - warnings.push(serde_json::json!({ - "accountId": snap.account_id, - "email": snap.email, - "level": "exhausted" - })); - } else { - for (quota_key, used_percent, role) in quota_warning_windows("a) { - let thresholds = match role { - "secondary" => &self.quota_thresholds.secondary, - _ => &self.quota_thresholds.primary, - }; - for threshold in thresholds { - if used_percent >= f64::from(*threshold) { - warnings.push(serde_json::json!({ - "accountId": snap.account_id, - "email": snap.email, - "level": "warning", - "threshold": threshold, - "usedPercent": used_percent, - "quotaKey": quota_key, - })); - break; - } - } - } - } - } - - Ok(serde_json::json!({ - "warnings": warnings, - "updatedAt": china_rfc3339(&Utc::now()) - })) - } - pub async fn health_check_accounts( - &self, - req: serde_json::Value, - ) -> Result { - use crate::upstream::accounts::store::StoredAccount; - - let ids = req - .get("ids") - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(ToString::to_string)) - .collect::>() - }) - .unwrap_or_default(); - - let mut results = Vec::new(); - let accounts = if ids.is_empty() { - let mut cursor = None; - let mut all: Vec = Vec::new(); - loop { - let page = self - .store - .list(cursor, 200) - .await - .map_err(|_| AdminAccountError::HealthCheck)?; - all.extend(page.items); - if page.next_cursor.is_none() { - break; - } - cursor = page.next_cursor; - } - all - } else { - let mut list = Vec::with_capacity(ids.len()); - for id in ids { - if let Ok(Some(acct)) = self.store.get(&id).await { - list.push(acct); - } - } - list - }; - - let request_id = uuid::Uuid::new_v4().to_string(); - for account in &accounts { - let token = account.access_token.expose_secret().to_string(); - let cookie_header = self.usage_cookie_header(&account.id).await; - let start = std::time::Instant::now(); - let context = crate::upstream::transport::CodexRequestContext { - access_token: &token, - account_id: account.account_id.as_deref(), - request_id: &request_id, - turn_state: None, - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - version: None, - codex_window_id: None, - parent_thread_id: None, - cookie_header: cookie_header.as_deref(), - installation_id: self.installation_id.as_deref(), - session_id: None, - }; - match self.codex.probe_connectivity(context).await { - Ok(probe) => { - let duration = start.elapsed().as_millis(); - results.push(serde_json::json!({ - "id": account.id, - "email": account.email, - "result": "alive", - "endpoint": probe.endpoint, - "status": probe.status.as_u16(), - "durationMs": duration - })); - } - Err(e) => { - let duration = start.elapsed().as_millis(); - results.push(serde_json::json!({ - "id": account.id, - "email": account.email, - "result": "dead", - "error": e.to_string(), - "durationMs": duration - })); - } - } - } - - let total = results.len(); - let alive = results - .iter() - .filter(|r| r.get("result") == Some(&serde_json::json!("alive"))) - .count(); - let dead = total - alive; - - Ok(serde_json::json!({ - "summary": { "total": total, "alive": alive, "dead": dead, "skipped": 0 }, - "results": results - })) - } -} - -fn quota_warning_windows(quota: &serde_json::Value) -> Vec<(String, f64, &'static str)> { - let mut windows = Vec::new(); - if let Some(monthly) = quota.get("monthly_limit").filter(|value| !value.is_null()) { - if let Some(used_percent) = used_percent(monthly) { - windows.push(("monthly_limit".to_string(), used_percent, "primary")); - } - } - if let Some(snapshots) = quota.get("snapshots").and_then(serde_json::Value::as_array) { - for snapshot in snapshots { - let source = snapshot - .get("source") - .and_then(serde_json::Value::as_str) - .unwrap_or("quota"); - let limit_name = snapshot - .get("limit_name") - .and_then(serde_json::Value::as_str) - .unwrap_or(source); - for role in ["primary", "secondary"] { - let Some(window) = snapshot.get(role).filter(|value| !value.is_null()) else { - continue; - }; - if let Some(used_percent) = used_percent(window) { - windows.push((format!("{source}:{limit_name}:{role}"), used_percent, role)); - } - } - } - } - windows -} - -fn used_percent(value: &serde_json::Value) -> Option { - value - .get("used_percent") - .and_then(|value| { - value - .as_f64() - .or_else(|| value.as_str().and_then(|value| value.parse::().ok())) - }) - .filter(|value| value.is_finite()) } diff --git a/src/admin/accounts/service/testing.rs b/src/admin/accounts/service/testing.rs index 22dbe15f3..276c5e788 100644 --- a/src/admin/accounts/service/testing.rs +++ b/src/admin/accounts/service/testing.rs @@ -16,21 +16,21 @@ use crate::upstream::{ use super::{types::AdminAccountError, AdminAccountService}; -pub type AccountTestStream = Pin> + Send>>; +pub(super) type AccountTestStream = Pin> + Send>>; #[derive(Debug, Clone, serde::Serialize)] #[serde(rename_all = "camelCase")] -pub struct AccountTestModel { +pub struct AccountModelOption { pub id: String, pub label: String, } impl AdminAccountService { - pub async fn test_connection_models( + pub async fn account_models( &self, account_id: &str, request_id: &str, - ) -> Result, AdminAccountError> { + ) -> Result, AdminAccountError> { let account = self .store .get(account_id) @@ -50,8 +50,8 @@ impl AdminAccountService { .fetch_models(&request) .await .map_err(|error| AdminAccountError::FetchModels(error.to_string()))? - .into_iter() - .filter_map(account_test_model) + .iter() + .filter_map(account_model_option) .collect::>(); if models.is_empty() { return Err(AdminAccountError::NoModels); @@ -145,9 +145,9 @@ impl AdminAccountService { } } -fn account_test_model( - entry: crate::upstream::models::BackendModelEntry, -) -> Option { +fn account_model_option( + entry: &crate::upstream::models::BackendModelEntry, +) -> Option { let id = first_non_empty([ entry.slug.as_deref(), entry.id.as_deref(), @@ -162,7 +162,7 @@ fn account_test_model( ]) .unwrap_or(&id) .to_string(); - Some(AccountTestModel { id, label }) + Some(AccountModelOption { id, label }) } fn first_non_empty<'a>(values: impl IntoIterator>) -> Option<&'a str> { diff --git a/src/admin/accounts/service/types.rs b/src/admin/accounts/service/types.rs index cb868be36..7806f59a9 100644 --- a/src/admin/accounts/service/types.rs +++ b/src/admin/accounts/service/types.rs @@ -54,8 +54,6 @@ impl From for AdminAccountMetadata { pub enum AdminAccountError { #[error("failed to list accounts")] List, - #[error("failed to export accounts")] - Export, #[error("failed to import accounts")] Import, #[error("failed to inspect account")] @@ -70,26 +68,14 @@ pub enum AdminAccountError { UpdateStatus, #[error("failed to delete account")] Delete, - #[error("failed to load cookies")] - LoadCookies, - #[error("failed to store cookies")] - StoreCookies, - #[error("failed to delete cookies")] - DeleteCookies, #[error("failed to update claims")] UpdateClaims, - #[error("failed to reset usage")] - ResetUsage, #[error("failed to sync account pool")] SyncAccountPool, - #[error("failed to get quota warnings")] - QuotaWarnings, #[error("failed to store quota")] StoreQuota, #[error("failed to fetch quota: {0}")] FetchQuota(String), - #[error("health check failed")] - HealthCheck, #[error("failed to fetch models: {0}")] FetchModels(String), #[error("no available models returned")] @@ -157,12 +143,6 @@ pub struct BatchDeleteAccounts { pub not_found: Vec, } -#[derive(Debug, Clone)] -pub struct BatchUpdateAccountStatus { - pub updated: u32, - pub not_found: Vec, -} - #[derive(Debug, Clone)] pub struct ImportedAccounts { pub imported: u32, @@ -259,9 +239,7 @@ pub(super) fn import_usage_string(usage: &serde_json::Value, key: &str) -> Optio usage .get(key) .and_then(serde_json::Value::as_str) - .and_then(|value| { - crate::upstream::accounts::import_export::normalize_nonempty_str(Some(value)) - }) + .and_then(|value| crate::upstream::accounts::importing::normalize_nonempty_str(Some(value))) .map(ToString::to_string) } @@ -286,11 +264,6 @@ pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option } pub(super) fn parse_account_status(status: &str) -> Result { - crate::upstream::accounts::import_export::parse_account_status(status) - .map_err(|_| AdminAccountError::InvalidStatus(status.trim().to_ascii_lowercase())) -} - -pub(super) fn parse_batch_account_status(status: &str) -> Result { - crate::upstream::accounts::import_export::parse_batch_account_status(status) + crate::upstream::accounts::importing::parse_account_status(status) .map_err(|_| AdminAccountError::InvalidStatus(status.trim().to_ascii_lowercase())) } diff --git a/src/admin/auth/service.rs b/src/admin/auth/service.rs index c27091627..71d83f932 100644 --- a/src/admin/auth/service.rs +++ b/src/admin/auth/service.rs @@ -7,51 +7,6 @@ use uuid::Uuid; use crate::infra::identity::{hash_admin_password, verify_admin_password}; -// --------------------------------------------------------------------------- -// 领域类型 -// --------------------------------------------------------------------------- - -/// 管理员会话。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminSession { - /// 会话 ID。 - pub id: String, - /// 过期时间。 - pub expires_at: DateTime, -} - -impl AdminSession { - /// 判断会话在给定时间是否已经过期。 - pub fn is_expired_at(&self, now: DateTime) -> bool { - now >= self.expires_at - } -} - -/// 根据 TTL 计算会话过期时间。 -pub fn session_expiry(now: DateTime, ttl_minutes: i64) -> DateTime { - now + Duration::minutes(ttl_minutes.max(0)) -} - -/// 管理员认证服务。 -#[derive(Debug, Clone)] -pub struct AdminAuthService { - default_username: String, -} - -impl AdminAuthService { - /// 构造管理员认证服务。 - pub fn new(default_username: impl Into) -> Self { - Self { - default_username: default_username.into(), - } - } - - /// 判断用户名是否匹配配置中的管理员用户名。 - pub fn username_matches(&self, username: &str) -> bool { - self.default_username == username - } -} - /// 管理员登录成功后的会话。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct AdminLoginSession { @@ -91,7 +46,6 @@ pub enum AdminSessionError { #[derive(Clone)] pub struct AdminSessionService { store: SqliteAdminSessionStore, - auth: AdminAuthService, default_username: String, session_ttl_minutes: u64, } @@ -105,7 +59,6 @@ impl AdminSessionService { ) -> Self { Self { store, - auth: AdminAuthService::new(default_username.clone()), default_username, session_ttl_minutes, } @@ -139,7 +92,7 @@ impl AdminSessionService { password: &str, ) -> Result, AdminSessionError> { let username = username.unwrap_or(&self.default_username); - if !self.auth.username_matches(username) { + if username != self.default_username.as_str() { return Ok(None); } let Some(admin) = self @@ -209,11 +162,6 @@ impl SqliteAdminSessionStore { Self { pool } } - /// 返回底层连接池。 - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - /// 如果还没有管理员用户，则创建默认管理员。 pub async fn ensure_default_admin( &self, diff --git a/src/admin/auth/session.rs b/src/admin/auth/session.rs index 2052fe9e7..c897e0de8 100644 --- a/src/admin/auth/session.rs +++ b/src/admin/auth/session.rs @@ -4,57 +4,42 @@ use axum::{ extract::State, http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode}, response::{IntoResponse, Response}, - Extension, Json, + Json, }; use serde::{Deserialize, Serialize}; use crate::{ admin::response::{AdminEnvelope, AdminError, AdminResponse}, - http::middleware::request_id::RequestId, infra::time::china_rfc3339, runtime::state::AppState, }; const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session"; -/// 管理员登录请求。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct AdminLoginRequest { - /// 管理员用户名；缺省时使用配置中的默认管理员。 - pub username: Option, - /// 管理员密码。 - pub password: String, +pub(crate) struct AdminLoginRequest { + username: Option, + password: String, } -/// 管理员登录响应。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct AdminLoginData { - /// 会话过期时间。 - pub expires_at: String, +struct AdminLoginData { + expires_at: String, } -/// 管理员会话状态响应。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct AdminSessionStatusData { - /// 当前请求是否携带有效管理端会话。 - pub authenticated: bool, -} - -/// 会话登录是否成功。 -pub fn session_login_allowed() -> bool { - true +struct AdminSessionStatusData { + authenticated: bool, } /// `POST /api/admin/login` -pub async fn login( +pub(crate) async fn login( State(state): State, - Extension(request_id): Extension, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); let session = state .services .admin_sessions @@ -65,26 +50,17 @@ pub async fn login( StatusCode::INTERNAL_SERVER_ERROR, 50001, "Failed to create admin session", - request_id.clone(), ) })? .ok_or_else(|| { - AdminError::new( - StatusCode::UNAUTHORIZED, - 40102, - "Invalid admin credentials", - request_id.clone(), - ) + AdminError::new(StatusCode::UNAUTHORIZED, 40102, "Invalid admin credentials") })?; let mut response = AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - AdminLoginData { - expires_at: china_rfc3339(&session.expires_at), - }, - request_id.clone(), - ), + AdminEnvelope::ok(AdminLoginData { + expires_at: china_rfc3339(&session.expires_at), + }), ) .into_response(); let cookie = format!( @@ -98,7 +74,6 @@ pub async fn login( StatusCode::INTERNAL_SERVER_ERROR, 50001, "Failed to create admin session cookie", - request_id.clone(), ) })?, ); @@ -106,12 +81,10 @@ pub async fn login( } /// `GET /api/admin/auth/status` -pub async fn session_status( +pub(crate) async fn session_status( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, ) -> Result { - let request_id = request_id.as_str().to_string(); let authenticated = state .services .admin_sessions @@ -122,25 +95,20 @@ pub async fn session_status( StatusCode::INTERNAL_SERVER_ERROR, 50001, "Failed to validate admin session", - request_id.clone(), ) })?; Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminSessionStatusData { authenticated }, request_id), + AdminEnvelope::ok(AdminSessionStatusData { authenticated }), )) } /// `POST /api/admin/logout` -pub async fn logout( +pub(crate) async fn logout( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, ) -> Result { - let request_id = request_id.as_str().to_string(); - - // 获取并删除服务器端 session if let Some(session_id) = admin_session_cookie(&headers) { let _ = state .services @@ -149,19 +117,14 @@ pub async fn logout( .await; } - // 返回 JSON 响应并清除 cookie let mut response = AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - serde_json::json!({ - "message": "Logged out successfully" - }), - request_id.clone(), - ), + AdminEnvelope::ok(serde_json::json!({ + "message": "Logged out successfully" + })), ) .into_response(); - // 清除 cookie let cookie = format!("{ADMIN_SESSION_COOKIE}=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0"); response.headers_mut().insert( SET_COOKIE, @@ -170,7 +133,6 @@ pub async fn logout( StatusCode::INTERNAL_SERVER_ERROR, 50001, "Failed to clear admin session cookie", - request_id.clone(), ) })?, ); @@ -179,10 +141,9 @@ pub async fn logout( } /// 要求请求携带有效管理员会话。 -pub async fn require_admin_session( +pub(crate) async fn require_admin_session( state: &AppState, headers: &HeaderMap, - request_id: &str, ) -> Result<(), AdminError> { match state .services @@ -195,13 +156,11 @@ pub async fn require_admin_session( StatusCode::UNAUTHORIZED, 40101, "Admin session required", - request_id, )), Err(_) => Err(AdminError::new( StatusCode::INTERNAL_SERVER_ERROR, 50001, "Failed to validate admin session", - request_id, )), } } diff --git a/src/admin/keys/routes.rs b/src/admin/keys/routes.rs index f959b1192..a7dc4a092 100644 --- a/src/admin/keys/routes.rs +++ b/src/admin/keys/routes.rs @@ -1,10 +1,10 @@ -//! 管理端 v1 接口访问 Key 处理器（列表、创建、生命周期、导出）。 +//! 管理端 v1 接口访问 Key 处理器（列表、创建、生命周期）。 use axum::{ extract::{Query, State}, http::{HeaderMap, StatusCode}, response::IntoResponse, - Extension, Json, + Json, }; use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -15,7 +15,6 @@ use crate::{ AdminClientKeyError, AdminCreatedClientApiKey, AdminStoredClientApiKey, }, admin::response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}, - http::middleware::request_id::RequestId, infra::{ json::{clamp_limit, Page}, time::{china_relative_time_str, china_rfc3339_str}, @@ -23,94 +22,62 @@ use crate::{ runtime::state::AppState, }; -// ---- Query types ---- - #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct ApiKeysQuery { - pub cursor: Option, - pub limit: Option, +pub(crate) struct ApiKeysQuery { + cursor: Option, + limit: Option, } #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct ApiKeyExportQuery { - pub ids: Option, +pub(crate) struct CreateApiKeyRequest { + name: String, } #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct CreateApiKeyRequest { - pub name: String, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct BatchDeleteClientApiKeysRequest { - pub ids: Vec, -} - -// ---- Response types ---- - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ClientApiKeyData { - pub id: String, - pub name: String, - pub label: Option, - pub prefix: String, - pub key: String, - pub enabled: bool, - pub created_at: String, - pub created_at_display: String, - pub last_used_at: Option, - pub last_used_at_display: String, +pub(crate) struct BatchDeleteClientApiKeysRequest { + ids: Vec, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct CreatedClientApiKeyData { - pub id: String, - pub name: String, - pub label: Option, - pub prefix: String, - pub key: String, - pub enabled: bool, - pub created_at: String, - pub created_at_display: String, - pub last_used_at: Option, - pub last_used_at_display: String, +struct ClientApiKeyData { + id: String, + name: String, + label: Option, + prefix: String, + key: String, + enabled: bool, + created_at: String, + created_at_display: String, + last_used_at: Option, + last_used_at_display: String, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct BatchDeleteClientApiKeysData { - pub deleted: u32, - pub not_found: Vec, +struct CreatedClientApiKeyData { + id: String, + name: String, + label: Option, + prefix: String, + key: String, + enabled: bool, + created_at: String, + created_at_display: String, + last_used_at: Option, + last_used_at_display: String, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct ClientApiKeyExportData { - pub source_format: &'static str, - pub rotation_required: bool, - pub api_keys: Vec, +struct BatchDeleteClientApiKeysData { + deleted: u32, + not_found: Vec, } -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct ClientApiKeyExportEntry { - pub id: String, - pub name: String, - pub label: Option, - pub prefix: String, - pub enabled: bool, - pub created_at: String, - pub last_used_at: Option, -} - -// ---- Conversions ---- - impl From for ClientApiKeyData { fn from(k: AdminStoredClientApiKey) -> Self { Self { @@ -128,20 +95,6 @@ impl From for ClientApiKeyData { } } -impl From for ClientApiKeyExportEntry { - fn from(k: AdminStoredClientApiKey) -> Self { - Self { - id: k.id, - name: k.name, - label: k.label, - prefix: k.prefix, - enabled: k.enabled, - created_at: china_rfc3339_str(&k.created_at), - last_used_at: k.last_used_at.as_deref().map(china_rfc3339_str), - } - } -} - impl From for CreatedClientApiKeyData { fn from(k: AdminCreatedClientApiKey) -> Self { Self { @@ -159,33 +112,20 @@ impl From for CreatedClientApiKeyData { } } -fn client_key_error(error: AdminClientKeyError, request_id: String) -> AdminError { +fn client_key_error(error: &AdminClientKeyError) -> AdminError { match error { AdminClientKeyError::InvalidStatus(_) | AdminClientKeyError::EmptyName | AdminClientKeyError::EmptyIds - | AdminClientKeyError::LabelTooLong => AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - error.to_string(), - request_id, - ), - _ => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - error.to_string(), - request_id, - ), + | AdminClientKeyError::LabelTooLong => { + AdminError::new(StatusCode::BAD_REQUEST, 40001, error.to_string()) + } + _ => AdminError::new(StatusCode::INTERNAL_SERVER_ERROR, 50001, error.to_string()), } } -fn client_key_not_found(request_id: String) -> AdminError { - AdminError::new( - StatusCode::NOT_FOUND, - 40401, - "Client API key not found", - request_id, - ) +fn client_key_not_found() -> AdminError { + AdminError::new(StatusCode::NOT_FOUND, 40401, "Client API key not found") } struct ParsedClientApiKeyUpdate { @@ -194,33 +134,28 @@ struct ParsedClientApiKeyUpdate { status: Option, } -fn parse_client_api_key_update( - payload: Value, - request_id: &str, -) -> Result { +fn parse_client_api_key_update(payload: &Value) -> Result { let object = payload.as_object().ok_or_else(|| { AdminError::new( StatusCode::BAD_REQUEST, 40001, "Client API key update request must be an object", - request_id, ) })?; - let id = required_string_field(object, "id", request_id)?; + let id = required_string_field(object, "id")?; let label = object .get("label") - .map(|value| optional_string_field(value, "label", request_id)) + .map(|value| optional_string_field(value, "label")) .transpose()?; let status = object .get("status") - .map(|value| required_string_value(value, "status", request_id)) + .map(|value| required_string_value(value, "status")) .transpose()?; if label.is_none() && status.is_none() { return Err(AdminError::new( StatusCode::BAD_REQUEST, 40001, "Client API key update request must include label or status", - request_id, )); } Ok(ParsedClientApiKeyUpdate { id, label, status }) @@ -229,24 +164,18 @@ fn parse_client_api_key_update( fn required_string_field( object: &serde_json::Map, field: &'static str, - request_id: &str, ) -> Result { let Some(value) = object.get(field) else { return Err(AdminError::new( StatusCode::BAD_REQUEST, 40001, format!("{field} is required"), - request_id, )); }; - required_string_value(value, field, request_id) + required_string_value(value, field) } -fn required_string_value( - value: &Value, - field: &'static str, - request_id: &str, -) -> Result { +fn required_string_value(value: &Value, field: &'static str) -> Result { value .as_str() .map(str::trim) @@ -257,16 +186,11 @@ fn required_string_value( StatusCode::BAD_REQUEST, 40001, format!("{field} must be a non-empty string"), - request_id, ) }) } -fn optional_string_field( - value: &Value, - field: &'static str, - request_id: &str, -) -> Result, AdminError> { +fn optional_string_field(value: &Value, field: &'static str) -> Result, AdminError> { if value.is_null() { return Ok(None); } @@ -279,22 +203,17 @@ fn optional_string_field( StatusCode::BAD_REQUEST, 40001, format!("{field} must be a string or null"), - request_id, ) }) } -// ---- Handlers ---- - /// `GET /api/admin/keys` -pub async fn api_keys( +pub(crate) async fn api_keys( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Query(query): Query, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; let limit = clamp_limit(query.limit.unwrap_or(50)); match state .services @@ -309,42 +228,38 @@ pub async fn api_keys( }; Ok(AdminResponse::new( StatusCode::OK, - AdminPageEnvelope::ok(page, limit, request_id), + AdminPageEnvelope::ok(page, limit), )) } - Err(error) => Err(client_key_error(error, request_id)), + Err(error) => Err(client_key_error(&error)), } } /// `POST /api/admin/keys` -pub async fn create_api_key( +pub(crate) async fn create_api_key( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state.services.admin_client_keys.create(&payload.name).await { Ok(created) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(CreatedClientApiKeyData::from(created), request_id), + AdminEnvelope::ok(CreatedClientApiKeyData::from(created)), )), - Err(error) => Err(client_key_error(error, request_id)), + Err(error) => Err(client_key_error(&error)), } } /// `POST /api/admin/keys/update` -pub async fn update_api_key( +pub(crate) async fn update_api_key( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; - let update = parse_client_api_key_update(payload, &request_id)?; + let update = parse_client_api_key_update(&payload)?; if let Some(label) = update.label { match state .services @@ -353,8 +268,8 @@ pub async fn update_api_key( .await { Ok(Some(_)) => {} - Ok(None) => return Err(client_key_not_found(request_id)), - Err(error) => return Err(client_key_error(error, request_id)), + Ok(None) => return Err(client_key_not_found()), + Err(error) => return Err(client_key_error(&error)), } } if let Some(status) = update.status { @@ -364,31 +279,29 @@ pub async fn update_api_key( .update_status(&update.id, &status) .await { - Ok(Some(_)) => {} - Ok(None) => return Err(client_key_not_found(request_id)), - Err(error) => return Err(client_key_error(error, request_id)), + Ok(true) => {} + Ok(false) => return Err(client_key_not_found()), + Err(error) => return Err(client_key_error(&error)), } } match state.services.admin_client_keys.get(&update.id).await { Ok(Some(key)) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(ClientApiKeyData::from(key), request_id), + AdminEnvelope::ok(ClientApiKeyData::from(key)), )), - Ok(None) => Err(client_key_not_found(request_id)), - Err(error) => Err(client_key_error(error, request_id)), + Ok(None) => Err(client_key_not_found()), + Err(error) => Err(client_key_error(&error)), } } /// `POST /api/admin/keys/delete` -pub async fn batch_delete_api_keys( +pub(crate) async fn batch_delete_api_keys( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Json(payload): Json, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state .services .admin_client_keys @@ -397,50 +310,11 @@ pub async fn batch_delete_api_keys( { Ok(deleted) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - BatchDeleteClientApiKeysData { - deleted: deleted.deleted, - not_found: deleted.not_found, - }, - request_id, - ), + AdminEnvelope::ok(BatchDeleteClientApiKeysData { + deleted: deleted.deleted, + not_found: deleted.not_found, + }), )), - Err(error) => Err(client_key_error(error, request_id)), - } -} - -/// `GET /api/admin/keys/export` -pub async fn export_api_keys( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Query(query): Query, -) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; - let ids = query - .ids - .unwrap_or_default() - .split(',') - .map(str::trim) - .filter(|id| !id.is_empty()) - .map(ToString::to_string) - .collect(); - match state.services.admin_client_keys.export(ids).await { - Ok(keys) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - ClientApiKeyExportData { - source_format: "rustLocalClientApiKeys", - rotation_required: true, - api_keys: keys - .into_iter() - .map(ClientApiKeyExportEntry::from) - .collect(), - }, - request_id, - ), - )), - Err(error) => Err(client_key_error(error, request_id)), + Err(error) => Err(client_key_error(&error)), } } diff --git a/src/admin/keys/service.rs b/src/admin/keys/service.rs index b864a605c..70a71b224 100644 --- a/src/admin/keys/service.rs +++ b/src/admin/keys/service.rs @@ -149,16 +149,12 @@ impl AdminClientKeyService { &self, key_id: &str, status: &str, - ) -> Result, AdminClientKeyError> { + ) -> Result { let enabled = parse_client_key_status(status)?; - match self.store.set_enabled(key_id, enabled).await { - Ok(true) => Ok(Some(UpdatedClientApiKeyStatus { - id: key_id.to_string(), - enabled, - })), - Ok(false) => Ok(None), - Err(_) => Err(AdminClientKeyError::UpdateStatus), - } + self.store + .set_enabled(key_id, enabled) + .await + .map_err(|_| AdminClientKeyError::UpdateStatus) } /// 批量删除客户端 API Key。 @@ -180,38 +176,6 @@ impl AdminClientKeyService { } Ok(BatchDeleteClientApiKeys { deleted, not_found }) } - - /// 导出客户端 API Key 元数据。 - pub async fn export( - &self, - ids: Vec, - ) -> Result, AdminClientKeyError> { - if ids.is_empty() { - let mut all_keys = Vec::new(); - let mut cursor = None; - loop { - let page = self - .store - .list(cursor, 200) - .await - .map_err(|_| AdminClientKeyError::Export)?; - all_keys.extend(page.items.into_iter().map(AdminStoredClientApiKey::from)); - if page.next_cursor.is_none() { - return Ok(all_keys); - } - cursor = page.next_cursor; - } - } - let mut keys = Vec::with_capacity(ids.len()); - for id in ids { - match self.store.get(&id).await { - Ok(Some(key)) => keys.push(AdminStoredClientApiKey::from(key)), - Ok(None) => {} - Err(_) => return Err(AdminClientKeyError::Export), - } - } - Ok(keys) - } } /// 管理端可见的客户端 API Key 元数据。 @@ -240,13 +204,6 @@ pub struct AdminCreatedClientApiKey { pub last_used_at: Option, } -/// 客户端 API Key 状态更新结果。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct UpdatedClientApiKeyStatus { - pub id: String, - pub enabled: bool, -} - /// 客户端 API Key 批量删除结果。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct BatchDeleteClientApiKeys { @@ -259,8 +216,6 @@ pub struct BatchDeleteClientApiKeys { pub enum AdminClientKeyError { #[error("failed to list client API keys")] List, - #[error("failed to export client API keys")] - Export, #[error("failed to create client API key")] Create, #[error("failed to delete client API key")] @@ -455,8 +410,8 @@ impl SqliteClientKeyStore { let items = rows .into_iter() .take(take_count) - .map(|row| self.key_from_row(&row)) - .collect::, _>>()?; + .map(|row| Self::key_from_row(&row)) + .collect::>(); let next_cursor = if has_next { items .last() @@ -480,7 +435,7 @@ impl SqliteClientKeyStore { .fetch_optional(&self.pool) .await?; - row.map(|row| self.key_from_row(&row)).transpose() + Ok(row.map(|row| Self::key_from_row(&row))) } /// 更新客户端 API Key 启用状态。 @@ -526,11 +481,8 @@ impl SqliteClientKeyStore { Ok(result.rows_affected() > 0) } - fn key_from_row( - &self, - row: &sqlx::sqlite::SqliteRow, - ) -> Result { - Ok(StoredClientApiKey { + fn key_from_row(row: &sqlx::sqlite::SqliteRow) -> StoredClientApiKey { + StoredClientApiKey { id: row.get("id"), name: row.get("name"), label: row.get("label"), @@ -539,7 +491,7 @@ impl SqliteClientKeyStore { enabled: row.get::("enabled") != 0, created_at: row.get("created_at"), last_used_at: row.get("last_used_at"), - }) + } } } @@ -548,7 +500,7 @@ impl ClientKeyStore for SqliteClientKeyStore { async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult { verify_and_touch(&self.pool, plaintext) .await - .map_err(map_client_key_store_error) + .map_err(|error| map_client_key_store_error(&error)) } } @@ -574,7 +526,7 @@ async fn verify_and_touch( Ok(false) } -fn map_client_key_store_error(error: SqliteClientKeyStoreError) -> ClientKeyStoreError { +fn map_client_key_store_error(error: &SqliteClientKeyStoreError) -> ClientKeyStoreError { ClientKeyStoreError::OperationFailed { message: error.to_string(), } diff --git a/src/admin/mod.rs b/src/admin/mod.rs index 87f7344cf..24b4b4300 100644 --- a/src/admin/mod.rs +++ b/src/admin/mod.rs @@ -6,8 +6,6 @@ pub mod accounts; pub mod auth; /// v1 接口访问 Key 处理器。 pub mod keys; -/// 模型处理器。 -pub mod models; /// 管理面日志、用量与诊断统计。 pub mod monitoring; /// 响应封装。 diff --git a/src/admin/models/mod.rs b/src/admin/models/mod.rs deleted file mode 100644 index ecd3fedda..000000000 --- a/src/admin/models/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -//! 管理端模型 HTTP 模块。 - -pub mod routes; diff --git a/src/admin/models/routes.rs b/src/admin/models/routes.rs deleted file mode 100644 index c9b636ab6..000000000 --- a/src/admin/models/routes.rs +++ /dev/null @@ -1,92 +0,0 @@ -//! 管理端模型处理器。 - -use axum::{ - extract::State, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, -}; -use serde::Serialize; - -use crate::{ - admin::auth::session::require_admin_session, - admin::response::{AdminEnvelope, AdminError, AdminResponse}, - http::middleware::request_id::RequestId, - runtime::state::AppState, - upstream::models::{AdminModelError, ModelRefreshResult}, -}; - -/// 模型刷新响应。 -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct RefreshModelsData { - pub refreshed_plans: usize, - pub model_count: usize, - pub failed_plans: usize, -} - -/// `POST /api/admin/models/refresh` -pub async fn refresh_models( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, -) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; - - match state - .services - .admin_models - .refresh_backend_models(&request_id) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(RefreshModelsData::from(result), request_id), - )), - Err(error) => model_error(error, request_id), - } -} - -fn model_error( - error: AdminModelError, - request_id: String, -) -> Result>, AdminError> { - match error { - AdminModelError::NoAccounts => Err(AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - error.to_string(), - request_id, - )), - AdminModelError::AllPlansFailed(result) => Ok(AdminResponse::new( - StatusCode::BAD_GATEWAY, - AdminEnvelope::new( - 50201, - "Failed to refresh backend models", - RefreshModelsData::from(result), - request_id, - ), - )), - AdminModelError::ListAccounts - | AdminModelError::SnapshotStoreUnavailable - | AdminModelError::UpstreamClientUnavailable - | AdminModelError::StoreSnapshot - | AdminModelError::LoadSnapshots => Err(AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - error.to_string(), - request_id, - )), - } -} - -impl From for RefreshModelsData { - fn from(result: ModelRefreshResult) -> Self { - Self { - refreshed_plans: result.refreshed_plans, - model_count: result.model_count, - failed_plans: result.failed_plans, - } - } -} diff --git a/src/admin/monitoring/billing.rs b/src/admin/monitoring/billing.rs index 5813efe1f..6dbb1f413 100644 --- a/src/admin/monitoring/billing.rs +++ b/src/admin/monitoring/billing.rs @@ -151,7 +151,7 @@ const MODEL_PRICING_RULES: &[ModelPricingRule] = &[ ]; /// 按 codex2api 的 billing 引擎口径计算美元计费。 -pub(crate) fn calculate_cost( +pub(in crate::admin) fn calculate_cost( input_tokens: u64, output_tokens: u64, cached_tokens: u64, @@ -240,71 +240,62 @@ fn normalize_billing_model_name(model: &str) -> String { fn normalize_codex_billing_model(model: &str) -> Option<&'static str> { let compact = model.replace([' ', '_'], "-").to_ascii_lowercase(); - match () { - _ if compact.contains("gpt-5.5-pro") - || compact.contains("gpt5-5-pro") - || compact.contains("gpt5.5-pro") => - { - Some("gpt-5.5-pro") - } - _ if compact.contains("gpt-5.5") - || compact.contains("gpt5-5") - || compact.contains("gpt5.5") => - { - Some("gpt-5.5") - } - _ if compact.contains("gpt-5.4-mini") - || compact.contains("gpt5-4-mini") - || compact.contains("gpt5.4-mini") => - { - Some("gpt-5.4-mini") - } - _ if compact.contains("gpt-5.4-nano") - || compact.contains("gpt5-4-nano") - || compact.contains("gpt5.4-nano") => - { - Some("gpt-5.4-nano") - } - _ if compact.contains("gpt-5.4-pro") - || compact.contains("gpt5-4-pro") - || compact.contains("gpt5.4-pro") => - { - Some("gpt-5.4-pro") - } - _ if compact.contains("gpt-5.4") - || compact.contains("gpt5-4") - || compact.contains("gpt5.4") => - { - Some("gpt-5.4") - } - _ if compact.contains("gpt-5.2") - || compact.contains("gpt5-2") - || compact.contains("gpt5.2") => - { - Some("gpt-5.2") - } - _ if compact.contains("gpt-5.3-codex-spark") - || compact.contains("gpt5-3-codex-spark") - || compact.contains("gpt5.3-codex-spark") => - { - Some("gpt-5.3-codex-spark") - } - _ if compact.contains("gpt-5.3-codex") - || compact.contains("gpt5-3-codex") - || compact.contains("gpt5.3-codex") => - { - Some("gpt-5.3-codex") - } - _ if compact.contains("gpt-5.3") - || compact.contains("gpt5-3") - || compact.contains("gpt5.3") => - { - Some("gpt-5.3-codex") - } - _ if compact.contains("codex-auto-review") => Some("gpt-5.4"), - _ if compact.contains("codex") => Some("gpt-5.3-codex"), - _ if compact.contains("gpt-5") || compact.contains("gpt5") => Some("gpt-5.4"), - _ => None, + if compact.contains("gpt-5.5-pro") + || compact.contains("gpt5-5-pro") + || compact.contains("gpt5.5-pro") + { + Some("gpt-5.5-pro") + } else if compact.contains("gpt-5.5") + || compact.contains("gpt5-5") + || compact.contains("gpt5.5") + { + Some("gpt-5.5") + } else if compact.contains("gpt-5.4-mini") + || compact.contains("gpt5-4-mini") + || compact.contains("gpt5.4-mini") + { + Some("gpt-5.4-mini") + } else if compact.contains("gpt-5.4-nano") + || compact.contains("gpt5-4-nano") + || compact.contains("gpt5.4-nano") + { + Some("gpt-5.4-nano") + } else if compact.contains("gpt-5.4-pro") + || compact.contains("gpt5-4-pro") + || compact.contains("gpt5.4-pro") + { + Some("gpt-5.4-pro") + } else if compact.contains("gpt-5.4") + || compact.contains("gpt5-4") + || compact.contains("gpt5.4") + { + Some("gpt-5.4") + } else if compact.contains("gpt-5.2") + || compact.contains("gpt5-2") + || compact.contains("gpt5.2") + { + Some("gpt-5.2") + } else if compact.contains("gpt-5.3-codex-spark") + || compact.contains("gpt5-3-codex-spark") + || compact.contains("gpt5.3-codex-spark") + { + Some("gpt-5.3-codex-spark") + } else if compact.contains("gpt-5.3-codex") + || compact.contains("gpt5-3-codex") + || compact.contains("gpt5.3-codex") + || compact.contains("gpt-5.3") + || compact.contains("gpt5-3") + || compact.contains("gpt5.3") + { + Some("gpt-5.3-codex") + } else if compact.contains("codex-auto-review") { + Some("gpt-5.4") + } else if compact.contains("codex") { + Some("gpt-5.3-codex") + } else if compact.contains("gpt-5") || compact.contains("gpt5") { + Some("gpt-5.4") + } else { + None } } @@ -327,30 +318,30 @@ fn model_matches_rule(model: &str, rule: &str) -> bool { } fn claude_family_pricing(model: &str) -> Option { - match () { - _ if model.contains("opus") => { - if model.contains("4.7") - || model.contains("4-7") - || model.contains("4.6") - || model.contains("4-6") - || model.contains("4.5") - || model.contains("4-5") - { - Some(ModelPricing::new(5.0, 25.0)) - } else { - Some(ModelPricing::new(15.0, 75.0)) - } + if model.contains("opus") { + if model.contains("4.7") + || model.contains("4-7") + || model.contains("4.6") + || model.contains("4-6") + || model.contains("4.5") + || model.contains("4-5") + { + Some(ModelPricing::new(5.0, 25.0)) + } else { + Some(ModelPricing::new(15.0, 75.0)) } - _ if model.contains("sonnet") => Some(ModelPricing::new(3.0, 15.0)), - _ if model.contains("haiku") => { - if model.contains("3-5") || model.contains("3.5") { - Some(ModelPricing::new(1.0, 5.0)) - } else { - Some(ModelPricing::new(0.25, 1.25)) - } + } else if model.contains("sonnet") { + Some(ModelPricing::new(3.0, 15.0)) + } else if model.contains("haiku") { + if model.contains("3-5") || model.contains("3.5") { + Some(ModelPricing::new(1.0, 5.0)) + } else { + Some(ModelPricing::new(0.25, 1.25)) } - _ if model.contains("claude") => Some(ModelPricing::new(3.0, 15.0)), - _ => None, + } else if model.contains("claude") { + Some(ModelPricing::new(3.0, 15.0)) + } else { + None } } diff --git a/src/admin/monitoring/dashboard.rs b/src/admin/monitoring/dashboard.rs index 60dd94337..04fa13cd6 100644 --- a/src/admin/monitoring/dashboard.rs +++ b/src/admin/monitoring/dashboard.rs @@ -6,7 +6,6 @@ use axum::{ extract::{Query, State}, http::{HeaderMap, StatusCode}, response::IntoResponse, - Extension, }; use chrono::{DateTime, Duration, Utc}; use serde::{Deserialize, Serialize}; @@ -26,7 +25,6 @@ use crate::{ }, response::{AdminEnvelope, AdminError, AdminResponse}, }, - http::middleware::request_id::RequestId, infra::time::{ china_datetime, china_datetime_rfc3339_str, china_day_start, china_hour, china_hour_start, china_quarter_hour_start, china_relative_time, china_time, @@ -41,13 +39,13 @@ const HEALTH_TIMELINE_SLOTS: i64 = 7 * 24 * 4; #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardTrendQuery { - pub kind: Option, +pub(crate) struct DashboardTrendQuery { + kind: Option, } #[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] -pub enum DashboardTrendKind { +enum DashboardTrendKind { #[default] Usage, Latency, @@ -56,143 +54,143 @@ pub enum DashboardTrendKind { #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardSummaryData { - pub cards: DashboardCardsData, - pub trend: DashboardTrendData, - pub health_timeline: DashboardHealthTimelineData, - pub account_usage: Vec, - pub service_statuses: Vec, - pub event_logs: Vec, - pub pool_summary: AccountPoolDiagnostics, - pub capacity_info: AccountCapacityDiagnostics, - pub rotation_strategy: Option, +struct DashboardSummaryData { + cards: DashboardCardsData, + trend: DashboardTrendData, + health_timeline: DashboardHealthTimelineData, + account_usage: Vec, + service_statuses: Vec, + event_logs: Vec, + pool_summary: AccountPoolDiagnostics, + capacity_info: AccountCapacityDiagnostics, + rotation_strategy: Option, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardHealthTimelineData { - pub title: String, - pub description: String, - pub range_display: String, - pub reliability_display: String, - pub oldest_label: String, - pub newest_label: String, - pub points: String, +struct DashboardHealthTimelineData { + title: String, + description: String, + range_display: String, + reliability_display: String, + oldest_label: String, + newest_label: String, + points: String, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardCardsData { - pub accounts: DashboardAccountsCardData, - pub traffic: DashboardTrafficCardData, - pub tokens: DashboardTokenCardData, - pub cache: DashboardCacheCardData, +struct DashboardCardsData { + accounts: DashboardAccountsCardData, + traffic: DashboardTrafficCardData, + tokens: DashboardTokenCardData, + cache: DashboardCacheCardData, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardAccountsCardData { - pub total: u64, - pub enabled: u64, - pub abnormal: u64, +struct DashboardAccountsCardData { + total: u64, + enabled: u64, + abnormal: u64, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardTrafficCardData { - pub today_requests: u64, - pub yesterday_requests: u64, - pub total_requests: u64, - pub rpm: u64, - pub tpm: u64, +struct DashboardTrafficCardData { + today_requests: u64, + yesterday_requests: u64, + total_requests: u64, + rpm: u64, + tpm: u64, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardTokenCardData { - pub today_tokens: u64, - pub yesterday_tokens: u64, - pub total_tokens: u64, - pub today_cost_usd: Option, - pub total_cost_usd: Option, +struct DashboardTokenCardData { + today_tokens: u64, + yesterday_tokens: u64, + total_tokens: u64, + today_cost_usd: Option, + total_cost_usd: Option, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardCacheCardData { - pub today_hit_rate: Option, - pub yesterday_hit_rate: Option, - pub total_hit_rate: Option, - pub total_cached_tokens: u64, - pub first_token_latency_ms: Option, - pub completion_latency_ms: Option, +struct DashboardCacheCardData { + today_hit_rate: Option, + yesterday_hit_rate: Option, + total_hit_rate: Option, + total_cached_tokens: u64, + first_token_latency_ms: Option, + completion_latency_ms: Option, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardTrendData { - pub kind: DashboardTrendKind, - pub points: Vec, - pub summary: Vec, +struct DashboardTrendData { + kind: DashboardTrendKind, + points: Vec, + summary: Vec, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardTrendPointData { - pub time: String, - pub requests: u64, - pub input_tokens: u64, - pub output_tokens: u64, - pub cached_tokens: u64, - pub tokens: u64, - pub errors: u64, - pub latency: u64, - pub max_latency: u64, - pub min_latency: u64, - pub success_rate: f64, +struct DashboardTrendPointData { + time: String, + requests: u64, + input_tokens: u64, + output_tokens: u64, + cached_tokens: u64, + tokens: u64, + errors: u64, + latency: u64, + max_latency: u64, + min_latency: u64, + success_rate: f64, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardTrendSummaryData { - pub label: String, - pub value: u64, - pub ratio: Option, +struct DashboardTrendSummaryData { + label: String, + value: u64, + ratio: Option, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardAccountUsageData { - pub name: String, - pub email: String, - pub plan: String, - pub requests: u64, - pub tokens: u64, - pub quota_used_percent: Option, - pub last_used: String, - pub status: String, +struct DashboardAccountUsageData { + name: String, + email: String, + plan: String, + requests: u64, + tokens: u64, + quota_used_percent: Option, + last_used: String, + status: String, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardEventLogData { - pub id: String, - pub time: String, - pub level: EventLevel, - pub request_id: Option, - pub route: Option, - pub model: Option, - pub status_code: Option, - pub latency_ms: Option, +struct DashboardEventLogData { + id: String, + time: String, + level: EventLevel, + request_id: Option, + route: Option, + model: Option, + status_code: Option, + latency_ms: Option, } #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct DashboardServiceStatusData { - pub label: String, - pub value: String, - pub detail: String, - pub tone: String, +struct DashboardServiceStatusData { + label: String, + value: String, + detail: String, + tone: String, } #[derive(Debug, Clone, Copy, Default)] @@ -234,13 +232,11 @@ impl UsageWindow { } /// `GET /api/admin/dashboard/summary` -pub async fn dashboard_summary( +pub(crate) async fn dashboard_summary( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; let accounts = state .services @@ -280,44 +276,35 @@ pub async fn dashboard_summary( Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - DashboardSummaryData { - cards: dashboard_cards(&accounts, &summary, &time_buckets), - trend, - health_timeline: dashboard_health_timeline_data(&time_buckets), - account_usage: account_usage_data( - &accounts, - &usage_records, - "a_used_by_account, - ), - service_statuses: service_status_data(&state), - event_logs: dashboard_logs, - pool_summary, - capacity_info: AccountCapacityDiagnostics::from(capacity), - rotation_strategy: Some(settings.auth.rotation_strategy.clone()), - }, - request_id, - ), + AdminEnvelope::ok(DashboardSummaryData { + cards: dashboard_cards(&accounts, &summary, &time_buckets), + trend, + health_timeline: dashboard_health_timeline_data(&time_buckets), + account_usage: account_usage_data(&accounts, &usage_records, "a_used_by_account), + service_statuses: service_status_data(&state), + event_logs: dashboard_logs, + pool_summary, + capacity_info: AccountCapacityDiagnostics::from(capacity), + rotation_strategy: Some(settings.auth.rotation_strategy.clone()), + }), )) } /// `GET /api/admin/dashboard/trend?kind=usage|latency|errors` -pub async fn dashboard_trend( +pub(crate) async fn dashboard_trend( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Query(query): Query, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; let time_buckets = dashboard_time_buckets(&state, Utc::now()).await; Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default()), - request_id, - ), + AdminEnvelope::ok(dashboard_trend_data( + &time_buckets, + query.kind.unwrap_or_default(), + )), )) } diff --git a/src/admin/monitoring/event_store.rs b/src/admin/monitoring/event_store.rs index 393639732..ee47eb498 100644 --- a/src/admin/monitoring/event_store.rs +++ b/src/admin/monitoring/event_store.rs @@ -1,19 +1,14 @@ //! 事件日志存储实现（SQLite）。 -use async_trait::async_trait; use chrono::{DateTime, Utc}; use serde_json::Value; use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool}; use thiserror::Error; -use crate::admin::monitoring::events::{ - EventLevel, EventLog, EventLogStore, EventLogStoreError, EventLogStoreResult, -}; +use crate::admin::monitoring::events::{EventLevel, EventLog}; use crate::infra::json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page}; use crate::infra::time::china_quarter_hour_start; -/// 默认启用事件日志，便于管理端诊断请求链路。 -pub const DEFAULT_EVENT_LOG_ENABLED: bool = true; /// 默认保留最新事件日志数量。 pub const DEFAULT_EVENT_LOG_CAPACITY: u32 = 2_000; /// 默认不保存请求/响应体，避免把用户内容写入诊断日志。 @@ -175,14 +170,6 @@ impl SqliteEventLogStore { row.map(|row| event_from_row(&row)).transpose() } - /// 统计事件日志数量。 - pub async fn count(&self) -> SqliteEventLogStoreResult { - let count: (i64,) = sqlx::query_as("select count(*) from event_logs") - .fetch_one(&self.pool) - .await?; - Ok(count.0.max(0) as u64) - } - /// 清空事件日志。 pub async fn clear(&self) -> SqliteEventLogStoreResult { let result = sqlx::query("delete from event_logs") @@ -192,18 +179,6 @@ impl SqliteEventLogStore { } } -#[async_trait] -impl EventLogStore for SqliteEventLogStore { - async fn append(&self, event: &EventLog) -> EventLogStoreResult<()> { - append_event(&self.pool, event).await.map_err(|error| { - EventLogStoreError::OperationFailed { - message: error.to_string(), - } - })?; - Ok(()) - } -} - #[derive(Debug, Clone)] struct EventLogSummaryFields { transport: Option, @@ -444,7 +419,7 @@ async fn count_events( let mut builder = QueryBuilder::::new("select count(*) from event_logs"); push_filter(&mut builder, filter, None)?; let (total,): (i64,) = builder.build_query_as().fetch_one(pool).await?; - Ok(total.max(0) as u64) + Ok(total.max(0).cast_unsigned()) } fn push_filter( @@ -631,30 +606,6 @@ impl From for EventLogFilter { } } -/// 日志状态。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminLogState { - /// 是否启用。 - pub enabled: bool, - /// 内存容量。 - pub capacity: u32, - /// 是否捕获请求体。 - pub capture_body: bool, - /// 已存储数量。 - pub stored_count: u64, -} - -/// 日志状态更新。 -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct AdminLogStateUpdate { - /// 是否启用。 - pub enabled: Option, - /// 日志容量。 - pub capacity: Option, - /// 是否捕获请求体。 - pub capture_body: Option, -} - /// 清空日志结果。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct AdminClearLogs { @@ -671,9 +622,6 @@ pub enum AdminLogError { /// 读取失败。 #[error("failed to get event log")] Get, - /// 计数失败。 - #[error("failed to count event logs")] - Count, /// 清空失败。 #[error("failed to clear event logs")] Clear, @@ -683,9 +631,6 @@ pub enum AdminLogError { /// 裁剪失败。 #[error("failed to trim event logs")] Trim, - /// 日志容量非法。 - #[error("log capacity must be greater than zero")] - InvalidCapacity, } /// 管理端日志服务。 @@ -751,50 +696,6 @@ impl AdminLogService { self.store.get(id).await.map_err(|_| AdminLogError::Get) } - /// 读取日志状态。 - pub async fn state(&self) -> Result { - let settings = *self.settings.read().await; - Ok(AdminLogState { - enabled: settings.enabled, - capacity: settings.capacity, - capture_body: settings.capture_body, - stored_count: self.store.count().await.map_err(|_| AdminLogError::Count)?, - }) - } - - /// 更新日志状态。 - pub async fn update_state( - &self, - update: AdminLogStateUpdate, - ) -> Result { - if matches!(update.capacity, Some(0)) { - return Err(AdminLogError::InvalidCapacity); - } - - let trim_capacity = { - let mut settings = self.settings.write().await; - if let Some(enabled) = update.enabled { - settings.enabled = enabled; - } - if let Some(capacity) = update.capacity { - settings.capacity = capacity; - } - if let Some(capture_body) = update.capture_body { - settings.capture_body = capture_body; - } - update.capacity - }; - - if let Some(capacity) = trim_capacity { - self.store - .trim_to_capacity(capacity) - .await - .map_err(|_| AdminLogError::Trim)?; - } - - self.state().await - } - /// 清空日志。 pub async fn clear(&self) -> Result { self.store @@ -807,8 +708,7 @@ impl AdminLogService { /// 记录事件日志。 pub async fn record(&self, mut event: EventLog) -> Result<(), AdminLogError> { let settings = *self.settings.read().await; - let policy = crate::admin::monitoring::events::EventLogService::new(settings.enabled); - if !policy.should_record(&event) { + if !settings.enabled && event.level != EventLevel::Error { return Ok(()); } apply_capture_body_policy(&mut event, settings.capture_body); diff --git a/src/admin/monitoring/events.rs b/src/admin/monitoring/events.rs index 35c37e9c3..b9578ae70 100644 --- a/src/admin/monitoring/events.rs +++ b/src/admin/monitoring/events.rs @@ -1,10 +1,8 @@ -//! 事件日志模型、端口与策略服务。 +//! 事件日志模型。 -use async_trait::async_trait; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use serde_json::Value; -use thiserror::Error; +use serde_json::{Map, Value}; use uuid::Uuid; /// 事件等级。 @@ -84,51 +82,12 @@ impl EventLog { upstream_request_id: None, latency_ms: None, message: message.into(), - metadata: Value::Object(Default::default()), + metadata: Value::Object(Map::default()), created_at: Utc::now(), } } } -/// 事件日志存储错误。 -#[derive(Debug, Error)] -pub enum EventLogStoreError { - /// 底层存储失败。 - #[error("event log store operation failed: {message}")] - OperationFailed { - /// 错误说明。 - message: String, - }, -} - -/// 事件日志存储结果。 -pub type EventLogStoreResult = Result; - -/// 事件日志存储端口。 -#[async_trait] -pub trait EventLogStore: Send + Sync + 'static { - /// 写入事件日志。 - async fn append(&self, event: &EventLog) -> EventLogStoreResult<()>; -} - -/// 事件日志服务。 -#[derive(Debug, Clone)] -pub struct EventLogService { - enabled: bool, -} - -impl EventLogService { - /// 构造事件日志服务。 - pub fn new(enabled: bool) -> Self { - Self { enabled } - } - - /// 判断事件是否应该记录。 - pub fn should_record(&self, event: &EventLog) -> bool { - self.enabled || event.level == EventLevel::Error - } -} - /// 响应事件记录（供 dispatch 模块使用）。 pub struct ResponseEventRecord<'a> { pub logs: &'a crate::admin::monitoring::event_store::AdminLogService, diff --git a/src/admin/monitoring/logs.rs b/src/admin/monitoring/logs.rs index 383356090..4ddef9a16 100644 --- a/src/admin/monitoring/logs.rs +++ b/src/admin/monitoring/logs.rs @@ -4,20 +4,18 @@ use axum::{ extract::{Query, State}, http::{HeaderMap, StatusCode}, response::IntoResponse, - Extension, Json, }; use serde::{Deserialize, Serialize}; use crate::{ admin::monitoring::{ - event_store::{AdminLogError, AdminLogFilter, AdminLogState, AdminLogStateUpdate}, + event_store::{AdminLogError, AdminLogFilter}, events::{EventLevel, EventLog}, }, admin::{ auth::session::require_admin_session, response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}, }, - http::middleware::request_id::RequestId, infra::{ json::{clamp_limit, clamp_page, NumberedPage, Page}, time::china_datetime, @@ -25,115 +23,61 @@ use crate::{ runtime::state::AppState, }; -/// 日志查询参数。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct LogsQuery { - /// 分页游标。 - pub cursor: Option, - /// 分页大小。 - pub limit: Option, - /// 页码。 - pub page: Option, - /// 每页条目数。 - pub page_size: Option, - /// 事件类别。 - pub kind: Option, - /// 事件等级。 - pub level: Option, - /// 请求 ID。 - pub request_id: Option, - /// 账号 ID。 - pub account_id: Option, - /// 路由。 - pub route: Option, - /// 模型。 - pub model: Option, - /// HTTP 状态码。 - pub status_code: Option, - /// 上游传输方式。 - pub transport: Option, - /// 同一请求内的上游尝试序号。 - pub attempt_index: Option, - /// 上游 HTTP 状态码。 - pub upstream_status_code: Option, - /// 失败分类。 - pub failure_class: Option, - /// 上游响应 ID。 - pub response_id: Option, - /// 上游请求 ID。 - pub upstream_request_id: Option, - /// 搜索关键词。 - pub search: Option, +pub(crate) struct LogsQuery { + cursor: Option, + limit: Option, + page: Option, + page_size: Option, + kind: Option, + level: Option, + request_id: Option, + account_id: Option, + route: Option, + model: Option, + status_code: Option, + transport: Option, + attempt_index: Option, + upstream_status_code: Option, + failure_class: Option, + response_id: Option, + upstream_request_id: Option, + search: Option, } -/// 日志详情查询参数。 #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] -pub struct LogDetailQuery { - /// 日志事件 ID。 - pub id: String, +pub(crate) struct LogDetailQuery { + id: String, } -/// 管理端日志状态响应。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct LogStateData { - /// 是否启用。 - pub enabled: bool, - /// 内存容量。 - pub capacity: u32, - /// 是否捕获请求体。 - pub capture_body: bool, - /// 已存储数量。 - pub stored_count: u64, +struct ClearLogsData { + cleared: u64, } -/// 清空日志响应。 #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -pub struct ClearLogsData { - /// 清理数量。 - pub cleared: u64, -} - -/// 管理端日志响应。 -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct LogData { - /// 原始日志字段。 +struct LogData { #[serde(flatten)] - pub log: EventLog, - /// 中国时区展示时间。 - pub created_at_display: String, -} - -/// 更新日志状态请求。 -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UpdateLogStateRequest { - /// 是否启用。 - pub enabled: Option, - /// 日志容量。 - pub capacity: Option, - /// 是否捕获请求体。 - pub capture_body: Option, + log: EventLog, + created_at_display: String, } /// `GET /api/admin/logs` -pub async fn logs( +pub(crate) async fn logs( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Query(query): Query, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; let limit = clamp_limit(query.page_size.or(query.limit).unwrap_or(50)); let page = query.page; let use_numbered_page = page.is_some() || query.page_size.is_some(); let cursor = query.cursor.clone(); - let filter = filter_from_query(query, &request_id)?; + let filter = filter_from_query(query)?; if use_numbered_page { return match state @@ -151,10 +95,10 @@ pub async fn logs( }; Ok(AdminResponse::new( StatusCode::OK, - AdminPageEnvelope::numbered(page, request_id), + AdminPageEnvelope::numbered(page), )) } - Err(error) => Err(log_error(error, request_id)), + Err(error) => Err(log_error(&error)), }; } @@ -166,100 +110,56 @@ pub async fn logs( }; Ok(AdminResponse::new( StatusCode::OK, - AdminPageEnvelope::ok(page, limit, request_id), + AdminPageEnvelope::ok(page, limit), )) } - Err(error) => Err(log_error(error, request_id)), + Err(error) => Err(log_error(&error)), } } /// `GET /api/admin/logs/detail` -pub async fn log_detail( +pub(crate) async fn log_detail( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, Query(query): Query, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state.services.logs.get(&query.id).await { Ok(Some(log)) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(LogData::from(log), request_id), + AdminEnvelope::ok(LogData::from(log)), )), Ok(None) => Err(AdminError::new( StatusCode::NOT_FOUND, 40401, "Log event not found", - request_id, )), - Err(error) => Err(log_error(error, request_id)), + Err(error) => Err(log_error(&error)), } } /// `POST /api/admin/logs/delete` -pub async fn clear_logs( +pub(crate) async fn clear_logs( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; match state.services.logs.clear().await { Ok(cleared) => Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok( - ClearLogsData { - cleared: cleared.cleared, - }, - request_id, - ), + AdminEnvelope::ok(ClearLogsData { + cleared: cleared.cleared, + }), )), - Err(error) => Err(log_error(error, request_id)), + Err(error) => Err(log_error(&error)), } } -/// `GET /api/admin/logs/state` -pub async fn logs_state( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, -) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; - match state.services.logs.state().await { - Ok(state) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(LogStateData::from(state), request_id), - )), - Err(error) => Err(log_error(error, request_id)), - } -} - -/// `POST /api/admin/logs/state` -pub async fn update_logs_state( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Json(payload): Json, -) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; - match state.services.logs.update_state(payload.into()).await { - Ok(state) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(LogStateData::from(state), request_id), - )), - Err(error) => Err(log_error(error, request_id)), - } -} - -fn filter_from_query(query: LogsQuery, request_id: &str) -> Result { +fn filter_from_query(query: LogsQuery) -> Result { Ok(AdminLogFilter { kind: non_empty(query.kind), - level: level_from_query(query.level).map_err(|message| { - AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id) - })?, + level: level_from_query(query.level) + .map_err(|message| AdminError::new(StatusCode::BAD_REQUEST, 40001, message))?, request_id: non_empty(query.request_id), account_id: non_empty(query.account_id), route: non_empty(query.route), @@ -275,25 +175,15 @@ fn filter_from_query(query: LogsQuery, request_id: &str) -> Result AdminError { +fn log_error(error: &AdminLogError) -> AdminError { match error { AdminLogError::List | AdminLogError::Get - | AdminLogError::Count | AdminLogError::Clear | AdminLogError::Append - | AdminLogError::Trim => AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - error.to_string(), - request_id, - ), - AdminLogError::InvalidCapacity => AdminError::new( - StatusCode::BAD_REQUEST, - 40001, - error.to_string(), - request_id, - ), + | AdminLogError::Trim => { + AdminError::new(StatusCode::INTERNAL_SERVER_ERROR, 50001, error.to_string()) + } } } @@ -324,24 +214,3 @@ impl From for LogData { } } } - -impl From for LogStateData { - fn from(state: AdminLogState) -> Self { - Self { - enabled: state.enabled, - capacity: state.capacity, - capture_body: state.capture_body, - stored_count: state.stored_count, - } - } -} - -impl From for AdminLogStateUpdate { - fn from(request: UpdateLogStateRequest) -> Self { - Self { - enabled: request.enabled, - capacity: request.capacity, - capture_body: request.capture_body, - } - } -} diff --git a/src/admin/monitoring/mod.rs b/src/admin/monitoring/mod.rs index da14a79a2..bbdff1893 100644 --- a/src/admin/monitoring/mod.rs +++ b/src/admin/monitoring/mod.rs @@ -7,5 +7,4 @@ pub mod event_store; pub mod events; pub mod logs; pub mod service; -pub mod usage; pub mod usage_store; diff --git a/src/admin/monitoring/service.rs b/src/admin/monitoring/service.rs index 2fc18a021..f1664bbdb 100644 --- a/src/admin/monitoring/service.rs +++ b/src/admin/monitoring/service.rs @@ -1,103 +1,11 @@ //! 用量聚合模型、端口与策略服务。 -use async_trait::async_trait; use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; use thiserror::Error; use crate::admin::monitoring::usage_store::{ SqliteUsageStore, UsageListRecord, UsageSummary, UsageTimeBucketRecord, }; -use crate::infra::json::{NumberedPage, Page}; -use crate::upstream::accounts::model::AccountUsageDelta; -use crate::upstream::protocol::events::TokenUsage; - -/// API 响应用量记录。 -#[derive(Debug, Clone, serde::Serialize)] -pub struct UsageRecord { - pub request_id: String, - pub account_id: String, - pub model: String, - pub input_tokens: u64, - pub output_tokens: u64, - pub cached_tokens: u64, - pub reasoning_tokens: u64, - pub total_tokens: u64, -} - -/// 用量聚合窗口。 -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct UsageWindow { - /// 窗口开始时间。 - pub started_at: DateTime, - /// 请求数。 - pub request_count: u64, - /// 输入 token 数。 - pub input_tokens: u64, - /// 输出 token 数。 - pub output_tokens: u64, -} - -/// 用量快照。 -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "camelCase")] -pub struct UsageSnapshot { - /// 账号 ID。 - pub account_id: String, - /// 当前窗口。 - pub window: UsageWindow, -} - -/// 用量存储错误。 -#[derive(Debug, Error)] -pub enum UsageStoreError { - /// 底层存储失败。 - #[error("usage store operation failed: {message}")] - OperationFailed { - /// 错误说明。 - message: String, - }, -} - -/// 用量存储结果。 -pub type UsageStoreResult = Result; - -/// 用量存储端口。 -#[async_trait] -pub trait UsageStore: Send + Sync + 'static { - /// 写入用量快照。 - async fn record_snapshot(&self, snapshot: &UsageSnapshot) -> UsageStoreResult<()>; -} - -/// 用量聚合服务。 -#[derive(Debug, Clone, Default)] -pub struct UsageService; - -impl UsageService { - /// 将 token 增量累加到窗口。 - pub fn add_tokens(window: &mut UsageWindow, input_tokens: u64, output_tokens: u64) { - window.request_count += 1; - window.input_tokens += input_tokens; - window.output_tokens += output_tokens; - } - - /// 将标准化 token 用量转换为账号持久化用量增量。 - pub fn account_delta_from_token_usage(usage: TokenUsage) -> AccountUsageDelta { - AccountUsageDelta { - requests: 0, - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - reasoning_tokens: 0, - total_tokens: usage.input_tokens + usage.output_tokens + usage.cached_tokens, - empty_responses: 0, - image_input_tokens: 0, - image_output_tokens: 0, - image_requests: 0, - image_request_failures: 0, - } - } -} +use crate::infra::json::Page; /// 管理端用量统计服务。 #[derive(Clone)] @@ -128,24 +36,6 @@ impl AdminUsageService { }) } - pub async fn list_page( - &self, - page: u32, - page_size: u32, - ) -> Result, AdminUsageError> { - let page = self - .store - .list_usage_page(page, page_size) - .await - .map_err(|_| AdminUsageError::List)?; - Ok(NumberedPage { - items: page.items.into_iter().map(AdminUsageRecord::from).collect(), - total: page.total, - page: page.page, - page_size: page.page_size, - }) - } - /// 汇总账号用量。 pub async fn summary(&self) -> Result { self.store diff --git a/src/admin/monitoring/usage.rs b/src/admin/monitoring/usage.rs deleted file mode 100644 index beaf0c782..000000000 --- a/src/admin/monitoring/usage.rs +++ /dev/null @@ -1,201 +0,0 @@ -//! 管理端用量处理器。 - -use axum::{ - extract::{Query, State}, - http::{HeaderMap, StatusCode}, - response::IntoResponse, - Extension, -}; -use serde::{Deserialize, Serialize}; - -use crate::{ - admin::auth::session::require_admin_session, - admin::monitoring::service::{AdminUsageError, AdminUsageRecord, AdminUsageSummary}, - admin::response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}, - http::middleware::request_id::RequestId, - infra::{ - json::{clamp_limit, clamp_page, Page}, - time::{china_relative_time, china_rfc3339}, - }, - runtime::state::AppState, -}; - -/// 用量统计查询参数。 -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct UsageStatsQuery { - pub cursor: Option, - pub limit: Option, - pub page: Option, - pub page_size: Option, -} - -/// 管理端账号用量响应。 -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminUsageStatsData { - pub account_id: String, - pub email: Option, - pub label: Option, - pub plan_type: Option, - pub request_count: i64, - pub empty_response_count: i64, - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub reasoning_tokens: i64, - pub total_tokens: i64, - pub image_input_tokens: i64, - pub image_output_tokens: i64, - pub image_request_count: i64, - pub image_request_failed_count: i64, - pub last_used_at: Option, - pub last_used_at_display: String, -} - -/// 管理端账号用量汇总响应。 -#[derive(Debug, Clone, Copy, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminUsageStatsSummaryData { - pub account_count: i64, - pub request_count: i64, - pub empty_response_count: i64, - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub reasoning_tokens: i64, - pub total_tokens: i64, - pub image_input_tokens: i64, - pub image_output_tokens: i64, - pub image_request_count: i64, - pub image_request_failed_count: i64, -} - -/// `GET /api/admin/usage` -pub async fn usage_stats( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, - Query(query): Query, -) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; - - let limit = clamp_limit(query.page_size.or(query.limit).unwrap_or(50)); - let use_numbered_page = query.page.is_some() || query.page_size.is_some(); - - if use_numbered_page { - return match state - .services - .usage - .list_page(clamp_page(query.page.unwrap_or(1)), limit) - .await - { - Ok(page) => { - let page = crate::infra::json::NumberedPage { - items: page - .items - .into_iter() - .map(AdminUsageStatsData::from) - .collect(), - total: page.total, - page: page.page, - page_size: page.page_size, - }; - Ok(AdminResponse::new( - StatusCode::OK, - AdminPageEnvelope::numbered(page, request_id), - )) - } - Err(error) => Err(usage_error(error, request_id)), - }; - } - - match state.services.usage.list(query.cursor, limit).await { - Ok(page) => { - let page = Page { - items: page - .items - .into_iter() - .map(AdminUsageStatsData::from) - .collect(), - next_cursor: page.next_cursor, - }; - Ok(AdminResponse::new( - StatusCode::OK, - AdminPageEnvelope::ok(page, limit, request_id), - )) - } - Err(error) => Err(usage_error(error, request_id)), - } -} - -/// `GET /api/admin/usage/summary` -pub async fn usage_stats_summary( - State(state): State, - Extension(request_id): Extension, - headers: HeaderMap, -) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; - - match state.services.usage.summary().await { - Ok(summary) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminUsageStatsSummaryData::from(summary), request_id), - )), - Err(error) => Err(usage_error(error, request_id)), - } -} - -fn usage_error(error: AdminUsageError, request_id: String) -> AdminError { - AdminError::new( - StatusCode::INTERNAL_SERVER_ERROR, - 50001, - error.to_string(), - request_id, - ) -} - -impl From for AdminUsageStatsData { - fn from(usage: AdminUsageRecord) -> Self { - Self { - account_id: usage.account_id, - email: usage.email, - label: usage.label, - plan_type: usage.plan_type, - request_count: usage.request_count, - empty_response_count: usage.empty_response_count, - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - reasoning_tokens: usage.reasoning_tokens, - total_tokens: usage.total_tokens, - image_input_tokens: usage.image_input_tokens, - image_output_tokens: usage.image_output_tokens, - image_request_count: usage.image_request_count, - image_request_failed_count: usage.image_request_failed_count, - last_used_at_display: china_relative_time(usage.last_used_at, chrono::Utc::now()), - last_used_at: usage.last_used_at.map(|value| china_rfc3339(&value)), - } - } -} - -impl From for AdminUsageStatsSummaryData { - fn from(summary: AdminUsageSummary) -> Self { - Self { - account_count: summary.account_count, - request_count: summary.request_count, - empty_response_count: summary.empty_response_count, - input_tokens: summary.input_tokens, - output_tokens: summary.output_tokens, - cached_tokens: summary.cached_tokens, - reasoning_tokens: summary.reasoning_tokens, - total_tokens: summary.total_tokens, - image_input_tokens: summary.image_input_tokens, - image_output_tokens: summary.image_output_tokens, - image_request_count: summary.image_request_count, - image_request_failed_count: summary.image_request_failed_count, - } - } -} diff --git a/src/admin/monitoring/usage_store.rs b/src/admin/monitoring/usage_store.rs index 6a59f2dc7..cece54a36 100644 --- a/src/admin/monitoring/usage_store.rs +++ b/src/admin/monitoring/usage_store.rs @@ -4,7 +4,7 @@ use chrono::{DateTime, Utc}; use sqlx::{sqlite::SqliteRow, Row, SqlitePool}; use thiserror::Error; -use crate::infra::json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page}; +use crate::infra::json::{decode_cursor, encode_cursor, Page}; const LIST_USAGE_AFTER_CURSOR_SQL: &str = r" select @@ -54,31 +54,6 @@ left join accounts a on a.id = au.account_id order by au.last_used_at desc, au.account_id desc limit ?"; -const LIST_USAGE_PAGE_SQL: &str = r" -select - au.account_id, - a.email, - a.label, - a.plan_type, - au.request_count, - au.empty_response_count, - au.input_tokens, - au.output_tokens, - au.cached_tokens, - au.reasoning_tokens, - au.total_tokens, - au.image_input_tokens, - au.image_output_tokens, - au.image_request_count, - au.image_request_failed_count, - au.last_used_at -from account_usage au -left join accounts a on a.id = au.account_id -order by au.last_used_at desc, au.account_id desc -limit ? offset ?"; - -const COUNT_USAGE_SQL: &str = "select count(*) from account_usage"; - const USAGE_SUMMARY_SQL: &str = r" select count(*) as account_count, @@ -244,11 +219,6 @@ impl SqliteUsageStore { Self { pool } } - /// 返回底层连接池。 - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - /// 分页列出账号用量。 pub async fn list_usage( &self, @@ -266,45 +236,16 @@ impl SqliteUsageStore { .bind(limit + 1) .fetch_all(&self.pool) .await?; - Ok(to_page(rows, limit)) + Ok(to_page(&rows, limit)) } else { let rows = sqlx::query(LIST_USAGE_SQL) .bind(limit + 1) .fetch_all(&self.pool) .await?; - Ok(to_page(rows, limit)) + Ok(to_page(&rows, limit)) } } - /// 按页码列出账号用量。 - pub async fn list_usage_page( - &self, - page: u32, - page_size: u32, - ) -> SqliteUsageStoreResult> { - let page_size = page_size.clamp(1, 200); - let offset = page_offset(page, page_size); - let rows = sqlx::query(LIST_USAGE_PAGE_SQL) - .bind(i64::from(page_size)) - .bind(offset.min(i64::MAX as u64) as i64) - .fetch_all(&self.pool) - .await?; - let items = rows - .iter() - .map(usage_list_from_row) - .collect::>>()?; - let (total,): (i64,) = sqlx::query_as(COUNT_USAGE_SQL) - .fetch_one(&self.pool) - .await?; - - Ok(NumberedPage { - items, - total: total.max(0) as u64, - page: page.max(1), - page_size, - }) - } - /// 汇总账号用量。 pub async fn usage_summary(&self) -> SqliteUsageStoreResult { let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?; @@ -356,7 +297,9 @@ fn usage_list_from_row(row: &SqliteRow) -> SqliteUsageStoreResult, _>("last_used_at"))?, + last_used_at: parse_optional_rfc3339( + row.get::, _>("last_used_at").as_deref(), + )?, }) } @@ -381,14 +324,13 @@ fn usage_time_bucket_from_row(row: &SqliteRow) -> SqliteUsageStoreResult) -> SqliteUsageStoreResult>> { +fn parse_optional_rfc3339(value: Option<&str>) -> SqliteUsageStoreResult>> { value - .as_deref() .map(|value| Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc))) .transpose() } -fn to_page(rows: Vec, limit: u32) -> Page { +fn to_page(rows: &[SqliteRow], limit: u32) -> Page { let has_more = rows.len() > limit as usize; let mut items = Vec::with_capacity(limit as usize); let mut last_row: Option<&SqliteRow> = None; diff --git a/src/admin/response.rs b/src/admin/response.rs index 793f48e71..5f512b14b 100644 --- a/src/admin/response.rs +++ b/src/admin/response.rs @@ -10,18 +10,13 @@ use serde::Serialize; /// 管理端错误。 pub struct AdminError { - pub status: StatusCode, - pub code: u32, - pub message: String, + status: StatusCode, + code: u32, + message: String, } impl AdminError { - pub fn new( - status: StatusCode, - code: u32, - message: impl Into, - _request_id: impl Into, - ) -> Self { + pub fn new(status: StatusCode, code: u32, message: impl Into) -> Self { Self { status, code, @@ -48,52 +43,47 @@ impl IntoResponse for AdminError { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct AdminEnvelope { - pub code: u32, - pub message: String, - pub data: T, + code: u32, + message: String, + data: T, } impl AdminEnvelope { - pub fn new( - code: u32, - message: impl Into, - data: T, - _request_id: impl Into, - ) -> Self { + pub fn new(code: u32, message: impl Into, data: T) -> Self { Self { code, message: message.into(), data, } } - pub fn ok(data: T, request_id: impl Into) -> Self { - Self::new(200, "OK", data, request_id) + pub fn ok(data: T) -> Self { + Self::new(200, "OK", data) } } /// 分页元数据。 #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -pub struct CursorPageMeta { - pub limit: u32, +pub(crate) struct CursorPageMeta { + pub(crate) limit: u32, #[serde(skip_serializing_if = "Option::is_none")] - pub next_cursor: Option, + pub(crate) next_cursor: Option, } /// 页码分页元数据。 #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -pub struct NumberedPageMeta { - pub page: u32, - pub page_size: u32, - pub total: u64, - pub total_pages: u32, +pub(crate) struct NumberedPageMeta { + pub(crate) page: u32, + pub(crate) page_size: u32, + pub(crate) total: u64, + pub(crate) total_pages: u32, } /// 分页元数据。 #[derive(Debug, Clone, Serialize)] #[serde(untagged)] -pub enum PageMeta { +pub(crate) enum PageMeta { Cursor(CursorPageMeta), Numbered(NumberedPageMeta), } @@ -102,21 +92,21 @@ pub enum PageMeta { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] pub struct AdminPageEnvelope { - pub code: u32, - pub message: String, - pub data: PageData, + code: u32, + message: String, + data: PageData, } /// 分页响应数据。 #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -pub struct PageData { - pub items: Vec, - pub page: PageMeta, +struct PageData { + items: Vec, + page: PageMeta, } impl AdminPageEnvelope { - pub fn ok(page: Page, limit: u32, _request_id: impl Into) -> Self { + pub fn ok(page: Page, limit: u32) -> Self { let Page { items, next_cursor } = page; Self { code: 200, @@ -128,7 +118,7 @@ impl AdminPageEnvelope { } } - pub fn numbered(page: NumberedPage, _request_id: impl Into) -> Self { + pub fn numbered(page: NumberedPage) -> Self { let NumberedPage { items, total, @@ -153,8 +143,8 @@ impl AdminPageEnvelope { /// 管理端响应。 pub struct AdminResponse { - pub status: StatusCode, - pub body: T, + status: StatusCode, + body: T, } impl AdminResponse { diff --git a/src/admin/router.rs b/src/admin/router.rs index b40f27247..ac0dcb653 100644 --- a/src/admin/router.rs +++ b/src/admin/router.rs @@ -9,20 +9,15 @@ use crate::runtime::state::AppState; use super::{ accounts::routes::{ - account_quota, account_test_models, accounts, batch_delete_accounts, create_account, - export_accounts, get_account_cookies, health_check_accounts, import_accounts, - oauth_authorize_account, oauth_exchange_account, quota_warnings, refresh_account, - reset_account_usage, set_account_cookies, test_account_connection, update_account, + account_models, account_quota, accounts, batch_delete_accounts, create_account, + import_accounts, oauth_authorize_account, oauth_exchange_account, refresh_account, + test_account_connection, update_account, }, auth::session::{login, logout, session_status}, - keys::routes::{ - api_keys, batch_delete_api_keys, create_api_key, export_api_keys, update_api_key, - }, - models::routes::refresh_models, + keys::routes::{api_keys, batch_delete_api_keys, create_api_key, update_api_key}, monitoring::{ dashboard::{dashboard_summary, dashboard_trend}, - logs::{clear_logs, log_detail, logs, logs_state, update_logs_state}, - usage::{usage_stats, usage_stats_summary}, + logs::{clear_logs, log_detail, logs}, }, settings::routes::{settings, update_settings}, }; @@ -36,11 +31,7 @@ pub fn router() -> Router { .route("/api/admin/settings", get(settings).post(update_settings)) .route("/api/admin/dashboard/summary", get(dashboard_summary)) .route("/api/admin/dashboard/trend", get(dashboard_trend)) - .route("/api/admin/models/refresh", post(refresh_models)) - .route("/api/admin/usage", get(usage_stats)) - .route("/api/admin/usage/summary", get(usage_stats_summary)) .route("/api/admin/accounts", get(accounts).post(create_account)) - .route("/api/admin/accounts/export", get(export_accounts)) .route("/api/admin/accounts/import", post(import_accounts)) .route( "/api/admin/accounts/oauth/authorize", @@ -50,31 +41,16 @@ pub fn router() -> Router { "/api/admin/accounts/oauth/exchange", post(oauth_exchange_account), ) - .route("/api/admin/accounts/quota-warnings", get(quota_warnings)) - .route( - "/api/admin/accounts/health-check", - post(health_check_accounts), - ) .route("/api/admin/accounts/test", post(test_account_connection)) - .route("/api/admin/accounts/models", get(account_test_models)) + .route("/api/admin/accounts/models", get(account_models)) .route("/api/admin/accounts/delete", post(batch_delete_accounts)) .route("/api/admin/accounts/update", post(update_account)) .route("/api/admin/accounts/refresh", post(refresh_account)) - .route("/api/admin/accounts/reset-usage", post(reset_account_usage)) - .route( - "/api/admin/accounts/cookies", - get(get_account_cookies).post(set_account_cookies), - ) .route("/api/admin/accounts/quota", get(account_quota)) .route("/api/admin/logs", get(logs)) .route("/api/admin/logs/delete", post(clear_logs)) - .route( - "/api/admin/logs/state", - get(logs_state).post(update_logs_state), - ) .route("/api/admin/logs/detail", get(log_detail)) .route("/api/admin/keys", get(api_keys).post(create_api_key)) - .route("/api/admin/keys/export", get(export_api_keys)) .route("/api/admin/keys/delete", post(batch_delete_api_keys)) .route("/api/admin/keys/update", post(update_api_key)) } diff --git a/src/admin/settings/routes.rs b/src/admin/settings/routes.rs index 89cba339e..9a39e9dde 100644 --- a/src/admin/settings/routes.rs +++ b/src/admin/settings/routes.rs @@ -5,7 +5,6 @@ use axum::{ extract::State, http::{HeaderMap, StatusCode}, response::IntoResponse, - Extension, }; use serde::Serialize; use serde_json::Value; @@ -18,7 +17,6 @@ use crate::{ settings::{AdminSettingsPatch, RuntimeSettingsError, SettingsServiceError}, types::AppConfig, }, - http::middleware::request_id::RequestId, runtime::state::AppState, }; @@ -32,17 +30,16 @@ const ALLOWED_SETTINGS_KEYS: [&str; 7] = [ "rotationStrategy", ]; -/// 管理端设置响应。 #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -pub struct AdminSettingsData { - pub model_aliases: BTreeMap, - pub model_account_routes: BTreeMap>, - pub refresh_margin_seconds: u64, - pub refresh_concurrency: u32, - pub max_concurrent_per_account: usize, - pub request_interval_ms: u64, - pub rotation_strategy: String, +struct AdminSettingsData { + model_aliases: BTreeMap, + model_account_routes: BTreeMap>, + refresh_margin_seconds: u64, + refresh_concurrency: u32, + max_concurrent_per_account: usize, + request_interval_ms: u64, + rotation_strategy: String, } impl AdminSettingsData { @@ -60,59 +57,48 @@ impl AdminSettingsData { } /// `GET /api/admin/settings` -pub async fn settings( +pub(crate) async fn settings( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; let config = state.services.settings.current(); Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminSettingsData::from_config(&config), request_id), + AdminEnvelope::ok(AdminSettingsData::from_config(&config)), )) } /// `POST /api/admin/settings` -pub async fn update_settings( +pub(crate) async fn update_settings( State(state): State, - Extension(request_id): Extension, headers: HeaderMap, body: Bytes, ) -> Result { - let request_id = request_id.as_str().to_string(); - require_admin_session(&state, &headers, &request_id).await?; + require_admin_session(&state, &headers).await?; - let patch = parse_settings_patch(&body, &request_id)?; + let patch = parse_settings_patch(&body)?; let config = state .services .settings .update(patch) .await - .map_err(|error| settings_service_error(error, &request_id))?; + .map_err(settings_service_error)?; Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminSettingsData::from_config(&config), request_id), + AdminEnvelope::ok(AdminSettingsData::from_config(&config)), )) } -fn parse_settings_patch(body: &[u8], request_id: &str) -> Result { - let value = serde_json::from_slice::(body).map_err(|_| { - AdminError::new( - StatusCode::BAD_REQUEST, - 40000, - "Malformed JSON body", - request_id, - ) - })?; +fn parse_settings_patch(body: &[u8]) -> Result { + let value = serde_json::from_slice::(body) + .map_err(|_| AdminError::new(StatusCode::BAD_REQUEST, 40000, "Malformed JSON body"))?; let object = value.as_object().ok_or_else(|| { AdminError::new( StatusCode::BAD_REQUEST, 40001, "Settings payload must be an object", - request_id, ) })?; @@ -124,21 +110,14 @@ fn parse_settings_patch(body: &[u8], request_id: &str) -> Result AdminError { +fn settings_service_error(error: RuntimeSettingsError) -> AdminError { match error { RuntimeSettingsError::InvalidField(SettingsServiceError::InvalidField { field, @@ -147,7 +126,6 @@ fn settings_service_error(error: RuntimeSettingsError, request_id: &str) -> Admi StatusCode::BAD_REQUEST, 40001, format!("Invalid setting {field}: {message}"), - request_id, ), RuntimeSettingsError::Database(_) | RuntimeSettingsError::Json(_) @@ -155,7 +133,6 @@ fn settings_service_error(error: RuntimeSettingsError, request_id: &str) -> Admi StatusCode::INTERNAL_SERVER_ERROR, 50000, "Failed to persist settings", - request_id, ), } } diff --git a/src/config/loader.rs b/src/config/loader.rs index 86af06f2b..0201cb278 100644 --- a/src/config/loader.rs +++ b/src/config/loader.rs @@ -1,15 +1,15 @@ use std::path::Path; -use super::{types::AppConfig, ConfigResult}; +use super::types::AppConfig; impl AppConfig { /// 从当前目录加载配置文件。 - pub fn load() -> ConfigResult { + pub fn load() -> Result { Self::load_from_dir(".") } /// 从指定目录加载配置文件。 - pub fn load_from_dir(config_dir: impl AsRef) -> ConfigResult { + pub fn load_from_dir(config_dir: impl AsRef) -> Result { let config_dir = config_dir.as_ref(); ::config::Config::builder() .add_source(::config::File::from(config_dir.join("config.yaml")).required(true)) diff --git a/src/config/mod.rs b/src/config/mod.rs index 724c9cd3f..0d7330769 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -8,6 +8,3 @@ pub mod loader; pub mod settings; /// 配置类型定义。 pub mod types; - -/// 配置加载结果类型。 -pub type ConfigResult = Result; diff --git a/src/config/settings.rs b/src/config/settings.rs index e9b7113b3..d860ab24e 100644 --- a/src/config/settings.rs +++ b/src/config/settings.rs @@ -102,7 +102,7 @@ impl SettingsService { current.request_interval_ms = request_interval_ms; } if let Some(rotation_strategy) = patch.rotation_strategy { - current.rotation_strategy = validate_rotation_strategy(rotation_strategy)?; + current.rotation_strategy = validate_rotation_strategy(&rotation_strategy)?; } Ok(()) } @@ -128,12 +128,6 @@ impl SettingsServiceError { Self::InvalidField { field, .. } => field, } } - /// 返回字段错误说明。 - pub fn message(&self) -> &str { - match self { - Self::InvalidField { message, .. } => message, - } - } } fn validate_model_aliases( @@ -141,8 +135,8 @@ fn validate_model_aliases( ) -> Result, SettingsServiceError> { let mut normalized = BTreeMap::new(); for (alias, target) in aliases { - let alias = non_empty_string("modelAliases", alias)?; - let target = non_empty_string("modelAliases", target)?; + let alias = non_empty_string("modelAliases", &alias)?; + let target = non_empty_string("modelAliases", &target)?; if alias == target { return Err(invalid_field( "modelAliases", @@ -159,10 +153,10 @@ fn validate_model_account_routes( ) -> Result>, SettingsServiceError> { let mut normalized = BTreeMap::new(); for (model, account_ids) in routes { - let model = non_empty_string("modelAccountRoutes", model)?; + let model = non_empty_string("modelAccountRoutes", &model)?; let mut normalized_account_ids = Vec::new(); for account_id in account_ids { - let account_id = non_empty_string("modelAccountRoutes", account_id)?; + let account_id = non_empty_string("modelAccountRoutes", &account_id)?; if normalized_account_ids.contains(&account_id) { return Err(invalid_field( "modelAccountRoutes", @@ -182,7 +176,7 @@ fn validate_model_account_routes( Ok(normalized) } -fn validate_rotation_strategy(strategy: String) -> Result { +fn validate_rotation_strategy(strategy: &str) -> Result { let strategy = non_empty_string("rotationStrategy", strategy)?; if ROTATION_STRATEGIES.contains(&strategy.as_str()) { Ok(strategy) @@ -194,7 +188,7 @@ fn validate_rotation_strategy(strategy: String) -> Result Result { +fn non_empty_string(field: &str, value: &str) -> Result { let value = value.trim().to_string(); if value.is_empty() { Err(invalid_field(field, "must not be empty")) @@ -243,15 +237,6 @@ pub struct RuntimeSettingsService { } impl RuntimeSettingsService { - /// 构造运行时设置服务。 - pub fn new(config: AppConfig, pool: SqlitePool) -> Self { - Self { - current: Arc::new(StdRwLock::new(Arc::new(config))), - pool, - account_pool: None, - } - } - /// 构造带账号池同步的运行时设置服务。 pub fn with_account_pool( config: AppConfig, @@ -494,7 +479,7 @@ fn runtime_settings_from_row(row: &SqliteRow) -> Result("rotation_strategy"))?, }) } diff --git a/src/config/types.rs b/src/config/types.rs index 76d84a18f..d6f39cde4 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -101,33 +101,12 @@ impl Default for AuthConfig { pub struct QuotaConfig { /// 配额刷新周期。 pub refresh_interval_minutes: u64, - /// 预警阈值。 - pub warning_thresholds: QuotaWarningThresholds, } impl Default for QuotaConfig { fn default() -> Self { Self { refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds::default(), - } - } -} - -/// 配额预警阈值集合。 -#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] -pub struct QuotaWarningThresholds { - /// 主阈值列表。 - pub primary: Vec, - /// 次阈值列表。 - pub secondary: Vec, -} - -impl Default for QuotaWarningThresholds { - fn default() -> Self { - Self { - primary: vec![80, 90], - secondary: vec![80, 90], } } } diff --git a/src/http/middleware/trace.rs b/src/http/middleware/trace.rs index bd70a324d..aa115acb5 100644 --- a/src/http/middleware/trace.rs +++ b/src/http/middleware/trace.rs @@ -14,9 +14,6 @@ use tracing::{info, warn, Span}; use crate::http::middleware::request_id::RequestId; -/// 默认 tracing 事件名。 -pub const TRACE_LAYER_NAME: &str = "http"; - /// HTTP tracing layer type. pub type HttpTraceLayer = TraceLayer< SharedClassifier, @@ -65,6 +62,10 @@ fn on_http_response(response: &Response, latency: Duration, _span: &Span) ); } +#[expect( + clippy::needless_pass_by_value, + reason = "tower-http OnFailure callbacks receive the failure class by value" +)] fn on_http_failure(failure_class: ServerErrorsFailureClass, latency: Duration, _span: &Span) { warn!( failure = %failure_class, diff --git a/src/infra/database.rs b/src/infra/database.rs index 89b246875..96d6fe32a 100644 --- a/src/infra/database.rs +++ b/src/infra/database.rs @@ -2,8 +2,7 @@ use std::{fs, path::Path, str::FromStr, time::Duration}; -use sqlx::sqlite::SqliteConnectOptions; -pub use sqlx::SqlitePool; +use sqlx::{sqlite::SqliteConnectOptions, SqlitePool}; /// 连接 SQLite 并初始化 schema。 pub async fn connect_sqlite(database_url: &str) -> Result { diff --git a/src/infra/json.rs b/src/infra/json.rs index ec7b48fa3..4be7a0f56 100644 --- a/src/infra/json.rs +++ b/src/infra/json.rs @@ -2,7 +2,6 @@ use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use serde::Serialize; -use serde_json::Value; // --------------------------------------------------------------------------- // 分页 @@ -71,25 +70,3 @@ pub fn total_pages(total: u64, page_size: u32) -> u32 { let pages = total.div_ceil(page_size); pages.min(u64::from(u32::MAX)) as u32 } - -// --------------------------------------------------------------------------- -// JSON 值提取 -// --------------------------------------------------------------------------- - -/// 从多个路径中提取第一个非空字符串。 -pub fn first_string(value: &Value, paths: &[&[&str]]) -> Option { - paths.iter().find_map(|path| string_at(value, path)) -} - -/// 从指定路径提取字符串。 -pub fn string_at(value: &Value, path: &[&str]) -> Option { - let mut current = value; - for segment in path { - current = current.get(*segment)?; - } - current - .as_str() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToString::to_string) -} diff --git a/src/infra/logging.rs b/src/infra/logging.rs index 599049f3c..6c9b9e816 100644 --- a/src/infra/logging.rs +++ b/src/infra/logging.rs @@ -73,8 +73,8 @@ pub fn build_file_appender(config: &RotationConfig) -> Result Result { - let appender = build_file_appender(&config)?; +pub fn init_tracing(config: &RotationConfig) -> Result { + let appender = build_file_appender(config)?; let (writer, guard) = tracing_appender::non_blocking(appender); let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); diff --git a/src/infra/paths.rs b/src/infra/paths.rs index 725176ed1..4b2e94a4c 100644 --- a/src/infra/paths.rs +++ b/src/infra/paths.rs @@ -9,8 +9,7 @@ use uuid::Uuid; const INSTALLATION_ID_FILE_NAME: &str = "installation_id"; -/// 返回本地数据目录。 -pub fn data_dir() -> PathBuf { +fn data_dir() -> PathBuf { dirs::data_local_dir() .unwrap_or_else(|| PathBuf::from(".")) .join("codex-proxy-rs") @@ -24,13 +23,11 @@ pub fn ensure_data_dir() -> std::io::Result { Ok(dir) } -/// 返回真实 Codex Desktop installation ID 文件路径。 -pub fn codex_desktop_installation_id_path() -> Option { +fn codex_desktop_installation_id_path() -> Option { dirs::home_dir().map(|home| home.join(".codex").join(INSTALLATION_ID_FILE_NAME)) } -/// 返回本应用数据目录中的 installation ID 文件路径。 -pub fn data_installation_id_path(data_dir: &Path) -> PathBuf { +fn data_installation_id_path(data_dir: &Path) -> PathBuf { data_dir.join(INSTALLATION_ID_FILE_NAME) } @@ -44,10 +41,7 @@ pub fn load_or_create_installation_id(data_dir: Option<&Path>) -> io::Result, data_path: Option<&Path>, ) -> io::Result { diff --git a/src/infra/time.rs b/src/infra/time.rs index 9b1c3a05c..bd0b5779e 100644 --- a/src/infra/time.rs +++ b/src/infra/time.rs @@ -19,21 +19,23 @@ pub fn china_rfc3339_millis(value: &DateTime) -> String { /// 将 RFC3339 字符串输出为中国时区 RFC3339 字符串。 pub fn china_rfc3339_str(value: &str) -> String { - DateTime::parse_from_rfc3339(value) - .map(|datetime| datetime.with_timezone(&china_offset()).to_rfc3339()) - .unwrap_or_else(|_| value.to_string()) + DateTime::parse_from_rfc3339(value).map_or_else( + |_| value.to_string(), + |datetime| datetime.with_timezone(&china_offset()).to_rfc3339(), + ) } /// 将 RFC3339 字符串输出为中国时区日期时间。 pub fn china_datetime_rfc3339_str(value: &str) -> String { - DateTime::parse_from_rfc3339(value) - .map(|datetime| { + DateTime::parse_from_rfc3339(value).map_or_else( + |_| value.to_string(), + |datetime| { datetime .with_timezone(&china_offset()) .format("%Y-%m-%d %H:%M:%S") .to_string() - }) - .unwrap_or_else(|_| value.to_string()) + }, + ) } /// 将 UTC 时间输出为中国时区日期时间。 @@ -149,9 +151,10 @@ pub fn china_relative_time_str(value: Option<&str>) -> String { let Some(value) = value else { return "-".to_string(); }; - DateTime::parse_from_rfc3339(value) - .map(|datetime| china_relative_time(Some(datetime.with_timezone(&Utc)), Utc::now())) - .unwrap_or_else(|_| value.to_string()) + DateTime::parse_from_rfc3339(value).map_or_else( + |_| value.to_string(), + |datetime| china_relative_time(Some(datetime.with_timezone(&Utc)), Utc::now()), + ) } /// Serde 序列化 UTC 时间为中国时区 RFC3339 字符串。 diff --git a/src/proxy/dispatch/chat.rs b/src/proxy/dispatch/chat.rs index 3ae161770..35e2f20cb 100644 --- a/src/proxy/dispatch/chat.rs +++ b/src/proxy/dispatch/chat.rs @@ -26,7 +26,8 @@ use crate::{ }, upstream::{ create_response_with_account, verify_acquired_quota_if_required, - QuotaVerificationDecision, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + QuotaVerificationContext, QuotaVerificationDecision, + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, }, }, proxy::openai::chat::ChatStreamTranslationError, @@ -149,7 +150,7 @@ impl ChatDispatchService { let (account_id, response) = loop { let acquire_request = AccountAcquireRequest::new(&request.model, Utc::now()) .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - let acquired = match self.account_pool.acquire_with(acquire_request).await { + let acquired = match self.account_pool.acquire_with(&acquire_request).await { Some(acquired) => acquired, None if quota_exhausted_count > 0 => { return_dispatch_error!(ChatDispatchError::QuotaExhausted { @@ -203,14 +204,16 @@ impl ChatDispatchService { None => return_dispatch_error!(ChatDispatchError::NoActiveAccount), }; let acquired = match verify_acquired_quota_if_required( - self.account_pool.as_ref(), - self.codex.as_ref(), - &self.cloudflare, - self.installation_id.as_deref(), - request_id, + QuotaVerificationContext { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + cloudflare: &self.cloudflare, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + }, acquired, - &mut excluded_account_ids, - &mut quota_verify_attempts, ) .await { diff --git a/src/proxy/dispatch/cloudflare.rs b/src/proxy/dispatch/cloudflare.rs index 79e7b51f0..b10049346 100644 --- a/src/proxy/dispatch/cloudflare.rs +++ b/src/proxy/dispatch/cloudflare.rs @@ -25,8 +25,8 @@ pub struct CloudflareRecovery { impl CloudflareRecovery { pub fn new(cookie_store: SqliteCookieStore) -> Self { Self { - path_block_tracker: Default::default(), - challenge_tracker: Default::default(), + path_block_tracker: CloudflarePathBlockTracker::default(), + challenge_tracker: CloudflareChallengeCooldownTracker::default(), cookie_store, } } @@ -104,7 +104,7 @@ impl CloudflareRecovery { fn clear_challenge_cookies_after_cooldown(&self, account_id: &str, delay_seconds: i64) { let cookie_store = self.cookie_store.clone(); let account_id = account_id.to_string(); - let delay = StdDuration::from_secs(delay_seconds.max(0) as u64); + let delay = StdDuration::from_secs(delay_seconds.max(0).cast_unsigned()); tokio::spawn(async move { tokio::time::sleep(delay).await; diff --git a/src/proxy/dispatch/responses.rs b/src/proxy/dispatch/responses.rs index cc6742375..064ea4b2f 100644 --- a/src/proxy/dispatch/responses.rs +++ b/src/proxy/dispatch/responses.rs @@ -49,7 +49,8 @@ use crate::{ create_compact_response_with_account_retrying_5xx, create_response_stream_with_account_retrying_5xx, create_response_with_account_retrying_5xx, verify_acquired_quota_if_required, - QuotaVerificationDecision, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + QuotaVerificationContext, QuotaVerificationDecision, + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, }, }, upstream::accounts::{ @@ -430,89 +431,89 @@ impl ResponseDispatchService { .clone() .with_exclude_account_ids(excluded_account_ids.iter().cloned()); attempt_acquire_request.now = Utc::now(); - let acquired = match self + let Some(acquired) = self .account_pool - .acquire_with(attempt_acquire_request) + .acquire_with(&attempt_acquire_request) .await - { - Some(acquired) => acquired, - None => { - let error = match last_exhausted_account_class { - Some(ExhaustedAccountClass::QuotaExhausted) => { - ResponseDispatchError::QuotaExhausted { - count: quota_exhausted_count, - upstream_error: last_quota_error.unwrap_or_default(), - } + else { + let error = match last_exhausted_account_class { + Some(ExhaustedAccountClass::QuotaExhausted) => { + ResponseDispatchError::QuotaExhausted { + count: quota_exhausted_count, + upstream_error: last_quota_error.unwrap_or_default(), } - Some(ExhaustedAccountClass::RateLimited) => { - ResponseDispatchError::RateLimited { - count: rate_limited_count, - upstream_error: last_rate_limit_error.unwrap_or_default(), - } + } + Some(ExhaustedAccountClass::RateLimited) => { + ResponseDispatchError::RateLimited { + count: rate_limited_count, + upstream_error: last_rate_limit_error.unwrap_or_default(), } - Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired { - count: expired_count, - upstream_error: last_auth_error.unwrap_or_default(), - }, - Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled { - count: disabled_count, - upstream_error: last_disabled_auth_error.unwrap_or_default(), - }, - Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned { - count: banned_count, - upstream_error: last_banned_auth_error.unwrap_or_default(), - status_code: last_banned_status_code.unwrap_or(403), - }, - Some(ExhaustedAccountClass::CloudflareChallenge) => { - ResponseDispatchError::CloudflareChallenge { - count: cloudflare_challenge_count, - upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), - } + } + Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired { + count: expired_count, + upstream_error: last_auth_error.unwrap_or_default(), + }, + Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled { + count: disabled_count, + upstream_error: last_disabled_auth_error.unwrap_or_default(), + }, + Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned { + count: banned_count, + upstream_error: last_banned_auth_error.unwrap_or_default(), + status_code: last_banned_status_code.unwrap_or(403), + }, + Some(ExhaustedAccountClass::CloudflareChallenge) => { + ResponseDispatchError::CloudflareChallenge { + count: cloudflare_challenge_count, + upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), } - Some(ExhaustedAccountClass::CloudflarePathBlocked) => { - ResponseDispatchError::CloudflarePathBlocked { - count: cloudflare_path_block_count, - upstream_error: last_cloudflare_path_block_error - .unwrap_or_default(), - } + } + Some(ExhaustedAccountClass::CloudflarePathBlocked) => { + ResponseDispatchError::CloudflarePathBlocked { + count: cloudflare_path_block_count, + upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), } - Some(ExhaustedAccountClass::ModelUnsupported) => { - ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count, - upstream_error: last_model_unsupported_error.unwrap_or_default(), - } + } + Some(ExhaustedAccountClass::ModelUnsupported) => { + ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count, + upstream_error: last_model_unsupported_error.unwrap_or_default(), } - None => ResponseDispatchError::NoActiveAccount, - }; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - false, - false, - Some(backend_transport_name( + } + None => ResponseDispatchError::NoActiveAccount, + }; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + account_id: last_attempted_account_id.as_deref(), + stream: false, + compact: false, + transport: Some(backend_transport_name( backend_transport_for_response_request(&request), )), - &error, - ) - .await; - return Err(error); - } + }, + &error, + ) + .await; + return Err(error); }; let acquired_account_id = acquired.account.id.clone(); // 配额验证 let acquired = match verify_acquired_quota_if_required( - self.account_pool.as_ref(), - self.codex.as_ref(), - &self.cloudflare, - self.installation_id.as_deref(), - request_id, + QuotaVerificationContext { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + cloudflare: &self.cloudflare, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + }, acquired, - &mut excluded_account_ids, - &mut quota_verify_attempts, ) .await { @@ -533,12 +534,14 @@ impl ResponseDispatchService { route, requested_model, started_at, - Some(&acquired_account_id), - false, - false, - Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), + ResponseDispatchErrorDetails { + account_id: Some(&acquired_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, &error, ) .await; @@ -595,10 +598,12 @@ impl ResponseDispatchService { route, requested_model, started_at, - Some(&release_account_id), - false, - false, - Some(backend_transport_name(response.transport)), + ResponseDispatchErrorDetails { + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, &error, ) .await; @@ -641,10 +646,12 @@ impl ResponseDispatchService { route, requested_model, started_at, - Some(&release_account_id), - false, - false, - Some(backend_transport_name(response.transport)), + ResponseDispatchErrorDetails { + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, &error, ) .await; @@ -792,12 +799,14 @@ impl ResponseDispatchService { route, requested_model, started_at, - Some(&release_account_id), - false, - false, - Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), + ResponseDispatchErrorDetails { + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, &error, ) .await; @@ -890,10 +899,12 @@ impl ResponseDispatchService { route, requested_model, started_at, - Some(&account.id), - false, - false, - Some(backend_transport_name(response.transport)), + ResponseDispatchErrorDetails { + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, &error, ) .await; @@ -906,10 +917,12 @@ impl ResponseDispatchService { route, requested_model, started_at, - Some(&account.id), - false, - false, - Some(backend_transport_name(response.transport)), + ResponseDispatchErrorDetails { + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, &error, ) .await; @@ -922,10 +935,12 @@ impl ResponseDispatchService { route, requested_model, started_at, - Some(&account.id), - false, - false, - Some(backend_transport_name(response.transport)), + ResponseDispatchErrorDetails { + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, &error, ) .await; @@ -934,32 +949,25 @@ impl ResponseDispatchService { } } - #[expect( - clippy::too_many_arguments, - reason = "keeps response dispatch logging call sites explicit at branch exits" - )] async fn record_response_dispatch_error( &self, request_id: &str, route: &str, requested_model: &str, started_at: Instant, - account_id: Option<&str>, - stream: bool, - compact: bool, - transport: Option<&str>, + details: ResponseDispatchErrorDetails<'_>, error: &ResponseDispatchError, ) { record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { logs: &self.logs, request_id, - account_id, + account_id: details.account_id, route, model: requested_model, started_at, - stream, - compact, - transport, + stream: details.stream, + compact: details.compact, + transport: details.transport, error, }) .await; @@ -1018,12 +1026,14 @@ impl ResponseDispatchService { route, requested_model, started_at, - last_attempted_account_id.as_deref(), - true, - false, - Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), + ResponseDispatchErrorDetails { + account_id: last_attempted_account_id.as_deref(), + stream: true, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, &error, ) .await; @@ -1036,10 +1046,12 @@ impl ResponseDispatchService { route, requested_model, started_at, - $account_id, - true, - false, - $transport, + ResponseDispatchErrorDetails { + account_id: $account_id, + stream: true, + compact: false, + transport: $transport, + }, &error, ) .await; @@ -1051,73 +1063,71 @@ impl ResponseDispatchService { .clone() .with_exclude_account_ids(excluded_account_ids.iter().cloned()); attempt_acquire_request.now = Utc::now(); - let acquired = match self + let Some(acquired) = self .account_pool - .acquire_with(attempt_acquire_request) + .acquire_with(&attempt_acquire_request) .await - { - Some(acquired) => acquired, - None => { - let error = match last_exhausted_account_class { - Some(ExhaustedAccountClass::QuotaExhausted) => { - ResponseDispatchError::QuotaExhausted { - count: quota_exhausted_count, - upstream_error: last_quota_error.unwrap_or_default(), - } + else { + let error = match last_exhausted_account_class { + Some(ExhaustedAccountClass::QuotaExhausted) => { + ResponseDispatchError::QuotaExhausted { + count: quota_exhausted_count, + upstream_error: last_quota_error.unwrap_or_default(), } - Some(ExhaustedAccountClass::RateLimited) => { - ResponseDispatchError::RateLimited { - count: rate_limited_count, - upstream_error: last_rate_limit_error.unwrap_or_default(), - } + } + Some(ExhaustedAccountClass::RateLimited) => { + ResponseDispatchError::RateLimited { + count: rate_limited_count, + upstream_error: last_rate_limit_error.unwrap_or_default(), } - Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired { - count: expired_count, - upstream_error: last_auth_error.unwrap_or_default(), - }, - Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled { - count: disabled_count, - upstream_error: last_disabled_auth_error.unwrap_or_default(), - }, - Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned { - count: banned_count, - upstream_error: last_banned_auth_error.unwrap_or_default(), - status_code: last_banned_status_code.unwrap_or(403), - }, - Some(ExhaustedAccountClass::CloudflareChallenge) => { - ResponseDispatchError::CloudflareChallenge { - count: cloudflare_challenge_count, - upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), - } + } + Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired { + count: expired_count, + upstream_error: last_auth_error.unwrap_or_default(), + }, + Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled { + count: disabled_count, + upstream_error: last_disabled_auth_error.unwrap_or_default(), + }, + Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned { + count: banned_count, + upstream_error: last_banned_auth_error.unwrap_or_default(), + status_code: last_banned_status_code.unwrap_or(403), + }, + Some(ExhaustedAccountClass::CloudflareChallenge) => { + ResponseDispatchError::CloudflareChallenge { + count: cloudflare_challenge_count, + upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), } - Some(ExhaustedAccountClass::CloudflarePathBlocked) => { - ResponseDispatchError::CloudflarePathBlocked { - count: cloudflare_path_block_count, - upstream_error: last_cloudflare_path_block_error - .unwrap_or_default(), - } + } + Some(ExhaustedAccountClass::CloudflarePathBlocked) => { + ResponseDispatchError::CloudflarePathBlocked { + count: cloudflare_path_block_count, + upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), } - Some(ExhaustedAccountClass::ModelUnsupported) => { - ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count, - upstream_error: last_model_unsupported_error.unwrap_or_default(), - } + } + Some(ExhaustedAccountClass::ModelUnsupported) => { + ResponseDispatchError::ModelUnsupported { + count: model_unsupported_count, + upstream_error: last_model_unsupported_error.unwrap_or_default(), } - None => ResponseDispatchError::NoActiveAccount, - }; - return_stream_dispatch_error!(error); - } + } + None => ResponseDispatchError::NoActiveAccount, + }; + return_stream_dispatch_error!(error); }; let acquired_account_id = acquired.account.id.clone(); let acquired = match verify_acquired_quota_if_required( - self.account_pool.as_ref(), - self.codex.as_ref(), - &self.cloudflare, - self.installation_id.as_deref(), - request_id, + QuotaVerificationContext { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + cloudflare: &self.cloudflare, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + }, acquired, - &mut excluded_account_ids, - &mut quota_verify_attempts, ) .await { @@ -1525,7 +1535,7 @@ impl ResponseDispatchService { loop { let acquire_request = AccountAcquireRequest::new(&request.model, Utc::now()) .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - let acquired = match self.account_pool.acquire_with(acquire_request).await { + let acquired = match self.account_pool.acquire_with(&acquire_request).await { Some(acquired) => acquired, None if quota_exhausted_count > 0 => { let error = ResponseDispatchError::QuotaExhausted { @@ -1663,14 +1673,16 @@ impl ResponseDispatchService { }; last_attempted_account_id = Some(acquired.account.id.clone()); let acquired = match verify_acquired_quota_if_required( - self.account_pool.as_ref(), - self.codex.as_ref(), - &self.cloudflare, - self.installation_id.as_deref(), - request_id, + QuotaVerificationContext { + account_pool: self.account_pool.as_ref(), + codex: self.codex.as_ref(), + cloudflare: &self.cloudflare, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + }, acquired, - &mut excluded_account_ids, - &mut quota_verify_attempts, ) .await { @@ -2261,7 +2273,8 @@ async fn send_live_response_stream_tail( return None; } } - let failure = premature_close_failed_event(latest_response_id(&body_text), failure_detail); + let failure = + premature_close_failed_event(latest_response_id(&body_text).as_deref(), failure_detail); body_text.push_str(&failure); body_bytes.extend_from_slice(failure.as_bytes()); if sender.send(Ok(Bytes::from(failure))).await.is_err() { @@ -2324,17 +2337,12 @@ fn latest_response_id(body: &str) -> Option { }) } -fn premature_close_failed_event(response_id: Option, detail: Option<&str>) -> String { +fn premature_close_failed_event(response_id: Option<&str>, detail: Option<&str>) -> String { let message = match detail.filter(|value| !value.trim().is_empty()) { Some(detail) => format!("Upstream stream closed before response.completed: {detail}"), None => "Upstream stream closed before response.completed".to_string(), }; - response_failed_sse_event_with_id( - response_id.as_deref(), - "server_error", - "stream_disconnected", - &message, - ) + response_failed_sse_event_with_id(response_id, "server_error", "stream_disconnected", &message) } async fn finalize_live_response_stream(context: LiveResponseStreamContext, body: String) { @@ -2500,6 +2508,13 @@ struct ResponseDispatchErrorEventRecord<'a> { error: &'a ResponseDispatchError, } +struct ResponseDispatchErrorDetails<'a> { + account_id: Option<&'a str>, + stream: bool, + compact: bool, + transport: Option<&'a str>, +} + async fn record_response_dispatch_error_event(record: ResponseDispatchErrorEventRecord<'_>) { let mut metadata = dispatch_error_metadata( record.error, diff --git a/src/proxy/dispatch/session_affinity.rs b/src/proxy/dispatch/session_affinity.rs index 0dd6200e2..9bf140b95 100644 --- a/src/proxy/dispatch/session_affinity.rs +++ b/src/proxy/dispatch/session_affinity.rs @@ -15,12 +15,6 @@ use crate::upstream::protocol::responses::CodexResponsesRequest; // 亲和性核心类型 // ==================================================================== -/// 会话亲和性键。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct SessionAffinityKey { - pub conversation_id: String, -} - /// 会话亲和性条目。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct SessionAffinityEntry { @@ -401,10 +395,6 @@ impl SqliteSessionAffinityStore { Self { pool } } - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - pub async fn upsert( &self, response_id: &str, diff --git a/src/proxy/dispatch/upstream.rs b/src/proxy/dispatch/upstream.rs index dc7ffa9d0..21fd503b9 100644 --- a/src/proxy/dispatch/upstream.rs +++ b/src/proxy/dispatch/upstream.rs @@ -35,33 +35,35 @@ pub(crate) enum QuotaVerificationDecision { MaxAttemptsReached, } -#[expect( - clippy::too_many_arguments, - reason = "quota verification is the boundary where account, transport, and retry state meet" -)] +pub(crate) struct QuotaVerificationContext<'a> { + pub account_pool: &'a RuntimeAccountPoolService, + pub codex: &'a CodexBackendClient, + pub cloudflare: &'a CloudflareRecovery, + pub installation_id: Option<&'a str>, + pub request_id: &'a str, + pub excluded_account_ids: &'a mut Vec, + pub verify_attempts: &'a mut usize, +} + pub(crate) async fn verify_acquired_quota_if_required( - account_pool: &RuntimeAccountPoolService, - codex: &CodexBackendClient, - cloudflare: &CloudflareRecovery, - installation_id: Option<&str>, - request_id: &str, + context: QuotaVerificationContext<'_>, acquired: AcquiredAccount, - excluded_account_ids: &mut Vec, - verify_attempts: &mut usize, ) -> QuotaVerificationDecision { if !acquired.account.quota_verify_required { return QuotaVerificationDecision::Ready(Box::new(acquired)); } let account_id = acquired.account.id.clone(); - let cookie_header = cloudflare + let cookie_header = context + .cloudflare .cookie_header_for_request(&account_id, "/codex/usage") .await; - let usage = codex + let usage = context + .codex .fetch_usage(CodexRequestContext { access_token: &acquired.account.access_token, account_id: acquired.account.account_id.as_deref(), - request_id, + request_id: context.request_id, turn_state: None, turn_metadata: None, beta_features: None, @@ -70,7 +72,7 @@ pub(crate) async fn verify_acquired_quota_if_required( codex_window_id: None, parent_thread_id: None, cookie_header: cookie_header.as_deref(), - installation_id, + installation_id: context.installation_id, session_id: None, }) .await; @@ -88,12 +90,15 @@ pub(crate) async fn verify_acquired_quota_if_required( }; let quota = quota_from_usage(&raw); - account_pool.apply_quota_snapshot(&account_id, "a).await; + context + .account_pool + .apply_quota_snapshot(&account_id, "a) + .await; if quota_snapshot_limit_reached("a) { - account_pool.release(&account_id).await; - excluded_account_ids.push(account_id); - *verify_attempts += 1; - if *verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS { + context.account_pool.release(&account_id).await; + context.excluded_account_ids.push(account_id); + *context.verify_attempts += 1; + if *context.verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS { return QuotaVerificationDecision::MaxAttemptsReached; } return QuotaVerificationDecision::RetryWithAnotherAccount; diff --git a/src/proxy/openai/chat.rs b/src/proxy/openai/chat.rs index 2c4a8e6ab..b0f5fcf9e 100644 --- a/src/proxy/openai/chat.rs +++ b/src/proxy/openai/chat.rs @@ -442,8 +442,7 @@ impl ChatCompletionStreamTranslator { let call_id = self .function_call_items .get(event_id) - .map(|info| info.call_id.as_str()) - .unwrap_or(event_id) + .map_or(event_id, |info| info.call_id.as_str()) .to_string(); Some((call_id, delta.to_string())) } @@ -456,8 +455,7 @@ impl ChatCompletionStreamTranslator { let call_id = self .function_call_items .get(event_id) - .map(|info| info.call_id.as_str()) - .unwrap_or(event_id) + .map_or(event_id, |info| info.call_id.as_str()) .to_string(); let arguments = value .get("arguments") @@ -840,11 +838,9 @@ pub fn chat_completion_from_codex_sse( return Err(ChatStreamTranslationError::Upstream(failure)); } - if completed_response.is_none() { + let Some(response) = completed_response else { return Ok(None); - } - - let response = completed_response.unwrap(); + }; let output_text_value = if output_text.is_empty() { output_text_from_response(&response) } else { @@ -913,11 +909,8 @@ pub async fn chat_completions( let parsed_model = catalog.parse_model_name(&model); let display_model = ModelCatalog::build_display_model_name(&parsed_model); let stream = chat_request.stream; - let mut codex_request = match translate_chat_to_codex(chat_request) { - Ok(request) => request, - Err(_) => { - return invalid_chat_completion_request_response().into_response(); - } + let Ok(mut codex_request) = translate_chat_to_codex(chat_request) else { + return invalid_chat_completion_request_response().into_response(); }; apply_response_model_options(&mut codex_request, &parsed_model); let include_reasoning = codex_request @@ -946,7 +939,7 @@ pub async fn chat_completions( include_reasoning, tuple_schema, ), - Err(error) => response_dispatch_chat_stream_error_response(error), + Err(error) => response_dispatch_chat_stream_error_response(&error), }; } @@ -957,7 +950,7 @@ pub async fn chat_completions( .await { Ok(body) => (StatusCode::OK, Json(body)).into_response(), - Err(error) => chat_dispatch_error_response(error), + Err(error) => chat_dispatch_error_response(&error), } } @@ -991,8 +984,8 @@ fn live_chat_event_stream_response( ) } -fn response_dispatch_chat_stream_error_response(error: ResponseDispatchError) -> Response { - let message = chat_stream_dispatch_error_message(&error); +fn response_dispatch_chat_stream_error_response(error: &ResponseDispatchError) -> Response { + let message = chat_stream_dispatch_error_message(error); chat_stream_error_response(&message) } @@ -1024,20 +1017,22 @@ fn output_text_from_response(response: &Value) -> String { response .get("output_text") .and_then(Value::as_str) - .map(ToString::to_string) - .unwrap_or_else(|| { - response - .get("output") - .and_then(Value::as_array) - .map(|items| { - items - .iter() - .filter_map(output_text_from_item) - .collect::>() - .join("\n\n") - }) - .unwrap_or_default() - }) + .map_or_else( + || { + response + .get("output") + .and_then(Value::as_array) + .map(|items| { + items + .iter() + .filter_map(output_text_from_item) + .collect::>() + .join("\n\n") + }) + .unwrap_or_default() + }, + ToString::to_string, + ) } fn output_text_from_item(item: &Value) -> Option { @@ -1077,8 +1072,7 @@ fn tool_call_from_done_event( .and_then(Value::as_str)?; let info = function_call_items.get(event_id); let call_id = info - .map(|info| info.call_id.as_str()) - .unwrap_or(event_id) + .map_or(event_id, |info| info.call_id.as_str()) .to_string(); if !finished_call_ids.insert(call_id.clone()) { return None; @@ -1093,7 +1087,7 @@ fn tool_call_from_done_event( .get("arguments") .and_then(Value::as_str) .unwrap_or_default(); - Some(openai_tool_call(call_id, name, arguments)) + Some(openai_tool_call(&call_id, name, arguments)) } fn tool_call_from_output_item( @@ -1125,10 +1119,10 @@ fn tool_call_from_output_item( .get("arguments") .and_then(Value::as_str) .unwrap_or_default(); - Some(openai_tool_call(call_id.to_string(), name, arguments)) + Some(openai_tool_call(call_id, name, arguments)) } -fn openai_tool_call(call_id: String, name: &str, arguments: &str) -> Value { +fn openai_tool_call(call_id: &str, name: &str, arguments: &str) -> Value { json!({ "id": call_id, "type": "function", diff --git a/src/proxy/openai/errors.rs b/src/proxy/openai/errors.rs index 4ffebd76d..b8ca27951 100644 --- a/src/proxy/openai/errors.rs +++ b/src/proxy/openai/errors.rs @@ -62,8 +62,8 @@ pub fn openai_error_response( ) } -pub fn chat_dispatch_error_response(error: ChatDispatchError) -> Response { - let error = chat_dispatch_openai_error(&error); +pub fn chat_dispatch_error_response(error: &ChatDispatchError) -> Response { + let error = chat_dispatch_openai_error(error); openai_error_response(error.status, &error.message, error.error_type, error.code) .into_response() } @@ -82,7 +82,7 @@ pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Respon ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => { responses_no_available_accounts_response().into_response() } - error => response_dispatch_openai_error_response(error), + error => response_dispatch_openai_error_response(&error), } } @@ -140,8 +140,8 @@ fn responses_error_response(status: StatusCode, message: &str) -> (StatusCode, J ) } -fn response_dispatch_openai_error_response(error: ResponseDispatchError) -> Response { - let error = response_dispatch_openai_error(&error); +fn response_dispatch_openai_error_response(error: &ResponseDispatchError) -> Response { + let error = response_dispatch_openai_error(error); openai_error_response(error.status, &error.message, error.error_type, error.code) .into_response() } diff --git a/src/proxy/openai/models.rs b/src/proxy/openai/models.rs index 4f6d1bd5d..c66531940 100644 --- a/src/proxy/openai/models.rs +++ b/src/proxy/openai/models.rs @@ -6,7 +6,6 @@ use axum::{ response::IntoResponse, Json, }; -use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use crate::{proxy::auth::authorize_client_api_key, runtime::state::AppState}; @@ -15,26 +14,6 @@ use super::errors::{missing_client_api_key_response, model_not_found_response}; const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000; -// ==================================================================== -// OpenAI 模型类型 -// ==================================================================== - -/// OpenAI 模型对象。 -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct OpenAiModel { - pub id: String, - pub object: String, - pub created: i64, - pub owned_by: String, -} - -/// OpenAI 模型列表。 -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct OpenAiModelList { - pub object: String, - pub data: Vec, -} - // ==================================================================== // HTTP 处理器 // ==================================================================== diff --git a/src/proxy/openai/responses.rs b/src/proxy/openai/responses.rs index 40e964111..9cd705a9d 100644 --- a/src/proxy/openai/responses.rs +++ b/src/proxy/openai/responses.rs @@ -40,14 +40,6 @@ const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent"; // 协议类型 // ==================================================================== -/// OpenAI 请求声明的响应格式类别。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ResponseFormat { - Text, - JsonObject, - JsonSchema, -} - /// Responses API 请求体。 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub struct OpenAiResponsesRequest { @@ -242,7 +234,7 @@ fn sanitize_responses_input(input: Value) -> Vec { match input { Value::Array(items) => sanitize_codex_input_items(items), Value::Null => Vec::new(), - Value::String(text) => vec![responses_input_text_message(text)], + Value::String(text) => vec![responses_input_text_message(&text)], value => vec![value], } } @@ -252,7 +244,7 @@ fn sanitize_codex_input_items(input: Vec) -> Vec { .into_iter() .filter_map(|item| { if let Value::String(text) = item { - return Some(responses_input_text_message(text)); + return Some(responses_input_text_message(&text)); } let Value::Object(object) = item else { return Some(item); @@ -266,7 +258,7 @@ fn sanitize_codex_input_items(input: Vec) -> Vec { .collect() } -fn responses_input_text_message(text: String) -> Value { +fn responses_input_text_message(text: &str) -> Value { json!({ "type": "message", "role": "user", @@ -476,8 +468,10 @@ pub fn response_failed_sse_event_with_id( }); let response_id = response_id .filter(|value| !value.trim().is_empty()) - .map(ToString::to_string) - .unwrap_or_else(|| format!("resp_proxy_{}", uuid::Uuid::new_v4().simple())); + .map_or_else( + || format!("resp_proxy_{}", uuid::Uuid::new_v4().simple()), + ToString::to_string, + ); let data = json!({ "type": "response.failed", "response": { @@ -576,7 +570,7 @@ async fn handle_responses( .await { Ok(stream) => live_event_stream_response(stream), - Err(error) => response_dispatch_stream_error_response(error), + Err(error) => response_dispatch_stream_error_response(&error), }; } @@ -724,8 +718,8 @@ fn live_event_stream_response(stream: ResponseDispatchStream) -> Response { sse_event_stream_response(Body::from_stream(stream.body), SseResponseOptions::BASIC) } -fn response_dispatch_stream_error_response(error: ResponseDispatchError) -> Response { - event_stream_response(responses_stream_dispatch_failed_sse_event(&error)) +fn response_dispatch_stream_error_response(error: &ResponseDispatchError) -> Response { + event_stream_response(responses_stream_dispatch_failed_sse_event(error)) } fn ensure_done_sse_frame(body: &mut String) { diff --git a/src/runtime/bootstrap.rs b/src/runtime/bootstrap.rs index 8ffe25a35..02138ccbf 100644 --- a/src/runtime/bootstrap.rs +++ b/src/runtime/bootstrap.rs @@ -105,5 +105,5 @@ fn init_logging(config: &AppConfig) -> Result, crate::infra::lo } let rotation = RotationConfig::new(&config.logging.directory, config.logging.retention_days); - Ok(Some(init_tracing(rotation)?)) + Ok(Some(init_tracing(&rotation)?)) } diff --git a/src/runtime/services.rs b/src/runtime/services.rs index 21d1189c7..105efd4b6 100644 --- a/src/runtime/services.rs +++ b/src/runtime/services.rs @@ -6,13 +6,13 @@ use crate::{ admin::monitoring::{ event_store::{ AdminLogService, SqliteEventLogStore, DEFAULT_EVENT_LOG_CAPACITY, - DEFAULT_EVENT_LOG_CAPTURE_BODY, DEFAULT_EVENT_LOG_ENABLED, + DEFAULT_EVENT_LOG_CAPTURE_BODY, }, service::AdminUsageService, usage_store::SqliteUsageStore, }, admin::{ - accounts::service::AdminAccountService, + accounts::service::{AdminAccountService, AdminAccountServiceParts}, auth::service::{AdminSessionService, SqliteAdminSessionStore}, keys::service::{AdminClientKeyService, ClientKeyService, SqliteClientKeyStore}, }, @@ -33,10 +33,10 @@ use crate::{ }, upstream::{ fingerprint::{Fingerprint, FingerprintRepository}, - models::{AdminModelService, ModelService, ModelSnapshotStore}, + models::{ModelService, ModelSnapshotStore}, token_client::{default_openai_token_client, TokenClientConfig}, transport::{ - build_reqwest_client, CodexBackendClient, CodexModelCatalogClient, CustomCaError, + build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient, }, }, }; @@ -66,7 +66,6 @@ pub struct BackgroundTaskStores { #[derive(Clone)] pub struct Services { pub models: StdArc, - pub admin_models: StdArc, pub accounts: StdArc, pub client_keys: StdArc, pub admin_client_keys: StdArc, @@ -85,6 +84,24 @@ pub struct Services { pub background_tasks: BackgroundTaskStores, } +/// 事件日志运行选项。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct EventLogOptions { + pub enabled: bool, + pub capacity: u32, + pub capture_body: bool, +} + +impl EventLogOptions { + pub fn from_config(config: &AppConfig) -> Self { + Self { + enabled: config.logging.enabled, + capacity: DEFAULT_EVENT_LOG_CAPACITY, + capture_body: DEFAULT_EVENT_LOG_CAPTURE_BODY, + } + } +} + impl Services { pub fn new(config: &AppConfig, stores: BackgroundTaskStores, fingerprint: Fingerprint) -> Self { Self::try_new(config, stores, fingerprint) @@ -99,21 +116,27 @@ impl Services { Self::try_with_installation_id(config, stores, fingerprint, None) } - pub fn with_installation_id( + pub(crate) fn try_with_installation_id( config: &AppConfig, stores: BackgroundTaskStores, fingerprint: Fingerprint, installation_id: Option, - ) -> Self { - Self::try_with_installation_id(config, stores, fingerprint, installation_id) - .expect("failed to build runtime services with configured TLS transport") + ) -> Result { + Self::try_with_installation_id_and_log_options( + config, + stores, + fingerprint, + installation_id, + EventLogOptions::from_config(config), + ) } - pub fn try_with_installation_id( + pub fn try_with_installation_id_and_log_options( config: &AppConfig, stores: BackgroundTaskStores, fingerprint: Fingerprint, installation_id: Option, + log_options: EventLogOptions, ) -> Result { let installation_id = installation_id.filter(|id| !id.trim().is_empty()); let account_store_trait = @@ -159,26 +182,25 @@ impl Services { stores.client_keys.clone(), ))); let token_client = StdArc::new(default_openai_token_client(token_client_config(config))); - let admin_accounts = StdArc::new(AdminAccountService::new( - stores.accounts.clone(), - stores.cookies.clone(), - config.quota.warning_thresholds.clone(), - codex.clone(), - account_pool.clone(), - token_client, - crate::admin::accounts::service::oauth::AccountOAuthService::new( + let admin_accounts = StdArc::new(AdminAccountService::new(AdminAccountServiceParts { + store: stores.accounts.clone(), + cookies: stores.cookies.clone(), + codex: codex.clone(), + account_pool: account_pool.clone(), + token_refresher: token_client, + oauth: crate::admin::accounts::service::oauth::AccountOAuthService::new( reqwest::Client::new(), config.auth.oauth_client_id.clone(), config.auth.oauth_token_endpoint.clone(), ), - config.auth.refresh_margin_seconds, - installation_id.clone(), - )); + refresh_margin_seconds: config.auth.refresh_margin_seconds, + installation_id: installation_id.clone(), + })); let logs = StdArc::new(AdminLogService::new( stores.event_logs.clone(), - DEFAULT_EVENT_LOG_ENABLED, - DEFAULT_EVENT_LOG_CAPACITY, - DEFAULT_EVENT_LOG_CAPTURE_BODY, + log_options.enabled, + log_options.capacity, + log_options.capture_body, )); let usage_store = SqliteUsageStore::new(stores.accounts.pool().clone()); let usage = StdArc::new(AdminUsageService::new(usage_store)); @@ -194,21 +216,8 @@ impl Services { crate::upstream::models::ModelConfig { model_aliases: config.model_aliases.clone(), }, - Some(snapshot_store.clone()), - Some(upstream_client.clone()), - None, - )); - let admin_models = StdArc::new(AdminModelService::new( - StdArc::new(ModelService::new( - crate::upstream::models::ModelConfig { - model_aliases: config.model_aliases.clone(), - }, - Some(snapshot_store), - Some(upstream_client), - None, - )), - account_store_trait.clone(), - installation_id.clone(), + Some(snapshot_store), + Some(upstream_client), )); let cloudflare_recovery = crate::proxy::dispatch::cloudflare::CloudflareRecovery::new(stores.cookies.clone()); @@ -232,7 +241,6 @@ impl Services { Ok(Self { models, - admin_models, accounts: account_store_trait, client_keys, admin_client_keys, diff --git a/src/runtime/shutdown.rs b/src/runtime/shutdown.rs index 0c07b2751..6b0064a79 100644 --- a/src/runtime/shutdown.rs +++ b/src/runtime/shutdown.rs @@ -22,7 +22,7 @@ pub async fn shutdown_signal() { let terminate = std::future::pending::<()>(); tokio::select! { - _ = ctrl_c => {}, - _ = terminate => {}, + () = ctrl_c => {}, + () = terminate => {}, } } diff --git a/src/runtime/tasks/cookie_cleanup.rs b/src/runtime/tasks/cookie_cleanup.rs index 33c8517ba..225784c7c 100644 --- a/src/runtime/tasks/cookie_cleanup.rs +++ b/src/runtime/tasks/cookie_cleanup.rs @@ -27,14 +27,6 @@ impl CookieCleanupTask { } } - /// 使用自定义间隔构造 Cookie 清理任务。 - pub fn with_interval_secs(store: SqliteCookieStore, interval_secs: u64) -> Self { - Self { - store, - interval_secs, - } - } - /// 启动后台清理任务。 pub fn start(self) -> SchedulerHandle { let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); diff --git a/src/runtime/tasks/coordinator.rs b/src/runtime/tasks/coordinator.rs index 62c256412..8de604052 100644 --- a/src/runtime/tasks/coordinator.rs +++ b/src/runtime/tasks/coordinator.rs @@ -30,17 +30,12 @@ pub struct TaskCoordinator { impl TaskCoordinator { /// 注册一个后台任务句柄。 - pub fn push(&mut self, name: &'static str, handle: SchedulerHandle) { + pub(crate) fn push(&mut self, name: &'static str, handle: SchedulerHandle) { self.handles.push((name, handle)); } - /// 返回已注册任务名称。 - pub fn task_names(&self) -> Vec<&'static str> { - self.handles.iter().map(|(name, _)| *name).collect() - } - /// 启动所有后台任务。 - pub fn start( + pub(crate) fn start( config: &crate::runtime::state::RuntimeConfig, services: &crate::runtime::services::Services, ) -> Self { @@ -69,9 +64,13 @@ impl TaskCoordinator { ); coordinator.push( "model_refresh", - ModelRefreshTask::new(services.models.clone(), services.accounts.clone()) - .with_installation_id(services.installation_id.clone()) - .start(), + ModelRefreshTask::new( + services.models.clone(), + services.accounts.clone(), + services.account_pool.clone(), + ) + .with_installation_id(services.installation_id.clone()) + .start(), ); if config.auth.refresh_enabled { coordinator.push( @@ -146,12 +145,12 @@ pub enum SchedulerHandle { impl SchedulerHandle { /// 使用关闭 channel 构造句柄。 - pub fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self { + pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self { Self::Channel(shutdown_tx) } /// 使用 `JoinHandle` 构造句柄。 - pub fn from_join_handle(handle: JoinHandle<()>) -> Self { + pub(crate) fn from_join_handle(handle: JoinHandle<()>) -> Self { Self::JoinHandle(handle) } diff --git a/src/runtime/tasks/mod.rs b/src/runtime/tasks/mod.rs index 7ba05ef04..32b312ba0 100644 --- a/src/runtime/tasks/mod.rs +++ b/src/runtime/tasks/mod.rs @@ -8,5 +8,3 @@ pub mod quota_refresh; pub mod session_affinity_cleanup; pub mod session_cleanup; pub mod token_refresh; - -pub use coordinator::{SchedulerHandle, TaskCoordinator}; diff --git a/src/runtime/tasks/model_refresh.rs b/src/runtime/tasks/model_refresh.rs index 54ea2cdb0..62f6d9e01 100644 --- a/src/runtime/tasks/model_refresh.rs +++ b/src/runtime/tasks/model_refresh.rs @@ -5,8 +5,9 @@ use std::{sync::Arc, time::Duration}; use tokio::time::{interval, sleep}; use tracing::{info, warn}; -use crate::upstream::accounts::model::Account; -use crate::upstream::accounts::store::AccountStore; +use crate::upstream::accounts::{ + model::Account, pool::RuntimeAccountPoolService, store::AccountStore, +}; use crate::upstream::models::{ModelRefreshResult, ModelService, ModelServiceError}; use super::coordinator::SchedulerHandle; @@ -15,6 +16,7 @@ use super::coordinator::SchedulerHandle; pub struct ModelRefreshTask { model_service: Arc, account_store: Arc, + account_pool: Arc, refresh_interval_secs: u64, initial_delay_ms: u64, retry_delay_ms: u64, @@ -29,10 +31,15 @@ const MAX_RETRIES: u32 = 12; impl ModelRefreshTask { /// 构造默认任务。 - pub fn new(model_service: Arc, account_store: Arc) -> Self { + pub fn new( + model_service: Arc, + account_store: Arc, + account_pool: Arc, + ) -> Self { Self { model_service, account_store, + account_pool, refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS, initial_delay_ms: INITIAL_DELAY_MS, retry_delay_ms: RETRY_DELAY_MS, @@ -65,7 +72,7 @@ impl ModelRefreshTask { info!("模型刷新任务在首次拉取前关闭"); return; } - _ = async {} => {} + () = async {} => {} } match self.refresh_once().await { @@ -98,7 +105,7 @@ impl ModelRefreshTask { info!("模型刷新任务在重试等待期间关闭"); return; } - _ = sleep(Duration::from_millis(self.retry_delay_ms)) => {} + () = sleep(Duration::from_millis(self.retry_delay_ms)) => {} } } } @@ -130,16 +137,23 @@ impl ModelRefreshTask { SchedulerHandle::new(shutdown_tx) } - async fn refresh_once(&self) -> Result { + /// 执行一次模型刷新并同步账号池模型计划约束。 + pub async fn refresh_once(&self) -> Result { let request_id = uuid::Uuid::new_v4().to_string(); let accounts = self.load_accounts().await?; - self.model_service + let result = self + .model_service .refresh_backend_models_with_installation_id( &accounts, &request_id, self.installation_id.as_deref(), ) - .await + .await?; + let allowlist = self.model_service.model_plan_allowlist().await?; + self.account_pool + .apply_model_plan_allowlist(allowlist) + .await; + Ok(result) } async fn load_accounts(&self) -> Result, ModelServiceError> { diff --git a/src/runtime/tasks/quota_refresh.rs b/src/runtime/tasks/quota_refresh.rs index f9c0bd7fe..39ed8dda5 100644 --- a/src/runtime/tasks/quota_refresh.rs +++ b/src/runtime/tasks/quota_refresh.rs @@ -16,8 +16,6 @@ use crate::{ use super::coordinator::SchedulerHandle; -const DEFAULT_INTERVAL_SECS: u64 = 15 * 60; - /// 主动配额刷新后台任务。 pub struct QuotaRefreshTask { service: RuntimeQuotaRefreshService, @@ -25,14 +23,6 @@ pub struct QuotaRefreshTask { } impl QuotaRefreshTask { - /// 构造默认配额刷新后台任务。 - pub fn new(store: SqliteAccountStore, codex: Arc) -> Self { - Self { - service: RuntimeQuotaRefreshService::new(store, codex), - interval_secs: DEFAULT_INTERVAL_SECS, - } - } - /// 使用自定义间隔构造配额刷新后台任务。 pub fn with_intervals( store: SqliteAccountStore, diff --git a/src/runtime/tasks/token_refresh.rs b/src/runtime/tasks/token_refresh.rs index 14805ac9a..be23136f8 100644 --- a/src/runtime/tasks/token_refresh.rs +++ b/src/runtime/tasks/token_refresh.rs @@ -35,19 +35,6 @@ where } } - /// 使用自定义间隔构造后台任务。 - pub fn with_interval_secs( - store: SqliteAccountStore, - policy: RefreshPolicy, - client: C, - interval_secs: u64, - ) -> Self { - Self { - service: RuntimeTokenRefreshService::new(store, policy, client), - interval_secs, - } - } - /// 使用刷新租约存储保护账号刷新。 pub fn with_refresh_lease_store(mut self, refresh_leases: RefreshLeaseStore) -> Self { self.service = self.service.with_refresh_lease_store(refresh_leases); diff --git a/src/upstream/accounts/cookies.rs b/src/upstream/accounts/cookies.rs index bcb9d8439..c9b4b1f50 100644 --- a/src/upstream/accounts/cookies.rs +++ b/src/upstream/accounts/cookies.rs @@ -54,20 +54,6 @@ impl SqliteCookieStore { Self { pool } } - /// 返回底层连接池。 - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - - /// 检查账号是否存在。 - pub async fn account_exists(&self, account_id: &str) -> SqliteCookieStoreResult { - let row = sqlx::query("select 1 from accounts where id = ?") - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - Ok(row.is_some()) - } - /// 捕获上游 `Set-Cookie` 响应头中允许持久化的 Cookie。 pub async fn capture_set_cookie( &self, @@ -83,20 +69,6 @@ impl SqliteCookieStore { self.upsert_cookie(account_id, parsed).await } - /// 将 Cookie 请求头写入账号 Cookie 存储。 - pub async fn set_cookie_header( - &self, - account_id: &str, - raw: &str, - ) -> SqliteCookieStoreResult { - let parsed = parse_cookie_header(raw); - let count = parsed.len(); - for cookie in parsed { - self.upsert_cookie(account_id, cookie).await?; - } - Ok(count) - } - /// 为请求域名读取账号 Cookie 请求头。 pub async fn cookie_header( &self, @@ -267,27 +239,6 @@ fn max_age_expires_at(seconds: i64) -> String { (now + Duration::seconds(seconds.min(i32::MAX as i64))).to_rfc3339() } -fn parse_cookie_header(raw: &str) -> Vec { - raw.split(';') - .map(str::trim) - .filter_map(|part| { - let (name, value) = part.split_once('=')?; - let name = name.trim(); - let value = value.trim(); - if name.is_empty() || value.is_empty() { - return None; - } - Some(ParsedCookie { - domain: DEFAULT_COOKIE_DOMAIN.to_string(), - name: name.to_string(), - value: value.to_string(), - path: "/".to_string(), - expires_at: None, - }) - }) - .collect() -} - fn domain_matches(request_domain: &str, cookie_domain: &str) -> bool { request_domain == cookie_domain || request_domain @@ -386,8 +337,7 @@ impl CloudflarePathBlockTracker { let count = counts .get(account_id) .filter(|state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER) - .map(|state| state.count.saturating_add(1)) - .unwrap_or(1); + .map_or(1, |state| state.count.saturating_add(1)); counts.insert( account_id.to_string(), PathBlockState { @@ -442,8 +392,7 @@ impl CloudflareChallengeCooldownTracker { let challenge_count = states .get(account_id) .filter(|state| now.signed_duration_since(state.updated_at) <= CHALLENGE_STALE_AFTER) - .map(|state| state.challenge_count.saturating_add(1)) - .unwrap_or(1); + .map_or(1, |state| state.challenge_count.saturating_add(1)); states.insert( account_id.to_string(), ChallengeCooldownState { diff --git a/src/upstream/accounts/import_export.rs b/src/upstream/accounts/importing.rs similarity index 82% rename from src/upstream/accounts/import_export.rs rename to src/upstream/accounts/importing.rs index b9e4bc1b7..b4779824e 100644 --- a/src/upstream/accounts/import_export.rs +++ b/src/upstream/accounts/importing.rs @@ -1,4 +1,4 @@ -//! 账号导入导出逻辑。 +//! 账号导入逻辑。 use chrono::{DateTime, Utc}; use serde_json::Value; @@ -29,7 +29,7 @@ pub struct AccountImportEntry { pub access_token_expires_at: Option, /// 导入时的原始状态。 pub status: Option, - /// Sub2api 导入时的缓存配额 JSON。 + /// Sub2API 导入时的缓存配额 JSON。 pub cached_quota: Option, /// 配额抓取时间。 pub quota_fetched_at: Option, @@ -40,18 +40,18 @@ pub struct AccountImportEntry { /// 账号导入来源。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AccountImportSource { - /// 原生格式。 - Native, - /// Sub2api 格式。 - Sub2api, + /// CPR 格式。 + Cpr, + /// Sub2API 格式。 + Sub2Api, } impl AccountImportSource { /// 返回来源的字符串标识。 pub fn as_str(self) -> &'static str { match self { - Self::Native => "native", - Self::Sub2api => "sub2api", + Self::Cpr => "cpr", + Self::Sub2Api => "sub2api", } } } @@ -179,8 +179,8 @@ fn account_import_entry_from_value( return Ok(None); }; let allowed_keys = match source { - AccountImportSource::Native => ACCOUNT_IMPORT_ACCOUNT_KEYS, - AccountImportSource::Sub2api => SUB2API_ACCOUNT_IMPORT_KEYS, + AccountImportSource::Cpr => ACCOUNT_IMPORT_ACCOUNT_KEYS, + AccountImportSource::Sub2Api => SUB2API_ACCOUNT_IMPORT_KEYS, }; if account .keys() @@ -209,13 +209,13 @@ fn account_import_entry_from_value( &["accessTokenExpiresAt", "access_token_expires_at"], ), status: first_string(value, &["status"]), - cached_quota: (source == AccountImportSource::Sub2api) + cached_quota: (source == AccountImportSource::Sub2Api) .then(|| first_value(value, &["cachedQuota", "cached_quota"])) .flatten(), - quota_fetched_at: (source == AccountImportSource::Sub2api) + quota_fetched_at: (source == AccountImportSource::Sub2Api) .then(|| first_string(value, &["quotaFetchedAt", "quota_fetched_at"])) .flatten(), - quota_verify_required: (source == AccountImportSource::Sub2api) + quota_verify_required: (source == AccountImportSource::Sub2Api) .then(|| first_bool(value, &["quotaVerifyRequired", "quota_verify_required"])) .flatten(), })) @@ -241,16 +241,16 @@ fn account_import_source( ) -> Result { if let Some(source_format) = first_string(value, &["sourceFormat", "source_format"]) { return match source_format.trim().to_ascii_lowercase().as_str() { - "native" => Ok(AccountImportSource::Native), - "sub2api" => Ok(AccountImportSource::Sub2api), + "cpr" => Ok(AccountImportSource::Cpr), + "sub2api" => Ok(AccountImportSource::Sub2Api), _ => Err("no importable accounts"), }; } if accounts.iter().any(account_import_entry_looks_sub2api) { - Ok(AccountImportSource::Sub2api) + Ok(AccountImportSource::Sub2Api) } else { - Ok(AccountImportSource::Native) + Ok(AccountImportSource::Cpr) } } @@ -284,7 +284,7 @@ pub fn normalized_imported_account_status( source: AccountImportSource, access_token: &str, ) -> AccountStatus { - if source == AccountImportSource::Sub2api + if source == AccountImportSource::Sub2Api && status == AccountStatus::Active && jwt_expiry(access_token, Utc::now()) != JwtExpiry::Valid { @@ -314,46 +314,13 @@ pub fn parse_account_status(status: &str) -> Result } } -/// 对管理端修改用的状态做严格校验。 -pub fn parse_batch_account_status(status: &str) -> Result { - match status.trim().to_ascii_lowercase().as_str() { - "active" => Ok(AccountStatus::Active), - "disabled" => Ok(AccountStatus::Disabled), - _ => Err("invalid account status (only active/disabled allowed)"), - } -} - -/// 将 RefreshFailure 映射为账号状态。 -pub fn refresh_failure_status( - failure: &crate::upstream::accounts::token_refresh::RefreshFailure, -) -> AccountStatus { - match failure { - crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant => { - AccountStatus::Expired - } - crate::upstream::accounts::token_refresh::RefreshFailure::QuotaExhausted => { - AccountStatus::QuotaExhausted - } - crate::upstream::accounts::token_refresh::RefreshFailure::Banned => AccountStatus::Banned, - crate::upstream::accounts::token_refresh::RefreshFailure::Disabled => { - AccountStatus::Disabled - } - crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport => { - AccountStatus::Active - } - crate::upstream::accounts::token_refresh::RefreshFailure::Transport => { - AccountStatus::Active - } - } -} - /// 判断是否需要清除 next_refresh_at。 pub fn refresh_failure_status_clears_next_refresh_at(status: AccountStatus) -> bool { !matches!(status, AccountStatus::Active) } /// 标准化 Bearer token（去除 Bearer 前缀，去除空白）。 -pub fn normalize_bearer_token(value: String) -> String { +pub fn normalize_bearer_token(value: &str) -> String { value .trim() .strip_prefix("Bearer ") @@ -385,7 +352,7 @@ pub fn normalize_nonempty_str(value: Option<&str>) -> Option<&str> { value.map(str::trim).filter(|value| !value.is_empty()) } -pub fn first_string(value: &Value, keys: &[&str]) -> Option { +fn first_string(value: &Value, keys: &[&str]) -> Option { keys.iter() .find_map(|key| value.get(key).and_then(Value::as_str)) .map(str::trim) diff --git a/src/upstream/accounts/mod.rs b/src/upstream/accounts/mod.rs index 2c41ff9ab..69e32c291 100644 --- a/src/upstream/accounts/mod.rs +++ b/src/upstream/accounts/mod.rs @@ -1,9 +1,9 @@ -//! 账号领域模块 —— 模型、仓库、池调度、Cookie、配额、刷新与导入导出。 +//! 账号领域模块 —— 模型、仓库、池调度、Cookie、配额、刷新与导入。 /// Cookie 捕获与重放策略。 pub mod cookies; -/// 账号导入导出逻辑。 -pub mod import_export; +/// 账号导入逻辑。 +pub mod importing; /// 账号领域模型。 pub mod model; /// 账号池调度策略。 diff --git a/src/upstream/accounts/model.rs b/src/upstream/accounts/model.rs index 50ea971fc..81154e74d 100644 --- a/src/upstream/accounts/model.rs +++ b/src/upstream/accounts/model.rs @@ -106,7 +106,7 @@ pub struct Account { pub last_used_at: Option, } -/// 账号用量增量（来自 usage.rs）。 +/// 账号用量增量。 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct AccountUsageDelta { /// 请求数增量。 @@ -133,25 +133,6 @@ pub struct AccountUsageDelta { pub image_request_failures: u64, } -impl AccountUsageDelta { - /// 合并两个用量增量。 - pub fn merged(self, other: Self) -> Self { - Self { - requests: self.requests + other.requests, - input_tokens: self.input_tokens + other.input_tokens, - output_tokens: self.output_tokens + other.output_tokens, - cached_tokens: self.cached_tokens + other.cached_tokens, - reasoning_tokens: self.reasoning_tokens + other.reasoning_tokens, - total_tokens: self.total_tokens + other.total_tokens, - empty_responses: self.empty_responses + other.empty_responses, - image_input_tokens: self.image_input_tokens + other.image_input_tokens, - image_output_tokens: self.image_output_tokens + other.image_output_tokens, - image_requests: self.image_requests + other.image_requests, - image_request_failures: self.image_request_failures + other.image_request_failures, - } - } -} - /// 账号模型维度用量增量。 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct AccountModelUsageDelta { diff --git a/src/upstream/accounts/pool.rs b/src/upstream/accounts/pool.rs index a0c222d7d..1572d2c9b 100644 --- a/src/upstream/accounts/pool.rs +++ b/src/upstream/accounts/pool.rs @@ -164,27 +164,6 @@ pub struct AccountWindowUsageDelta { pub image_request_failed: bool, } -/// 账号池状态摘要。 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct AccountPoolStatusSummary { - /// 账号总数。 - pub total: usize, - /// 可调度账号数。 - pub active: usize, - /// token 过期账号数。 - pub expired: usize, - /// 配额耗尽账号数。 - pub quota_exhausted: usize, - /// 处于限流冷却的账号数。 - pub rate_limited: usize, - /// 正在刷新账号数。 - pub refreshing: usize, - /// 禁用账号数。 - pub disabled: usize, - /// 封禁账号数。 - pub banned: usize, -} - /// 纯内存账号池，负责账号调度和运行时状态维护。 #[derive(Debug)] pub struct AccountPool { @@ -238,30 +217,11 @@ impl AccountPool { self.least_used_cursor = 0; } - /// 替换模型计划 allowlist。 - pub fn set_model_plan_allowlist(&mut self, allowlist: BTreeMap>) { - self.options.model_plan_allowlist = allowlist; - } - - /// 替换模型账号路由。 - pub fn set_model_account_routes(&mut self, routes: BTreeMap>) { - self.options.model_account_routes = routes; - } - /// 替换账号池运行参数。 pub fn set_options(&mut self, options: AccountPoolOptions) { self.options = options; } - /// 更新账号标签。 - pub fn set_label(&mut self, account_id: &str, label: Option) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; - }; - account.label = label; - true - } - /// 更新账号状态。 pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool { let Some(account) = self.accounts.get_mut(account_id) else { @@ -274,30 +234,6 @@ impl AccountPool { true } - /// 重置账号累计和窗口用量。 - pub fn reset_usage(&mut self, account_id: &str) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; - }; - account.last_used_at = None; - account.request_count = 0; - account.empty_response_count = 0; - account.image_input_tokens = 0; - account.image_output_tokens = 0; - account.image_request_count = 0; - account.image_request_failed_count = 0; - account.window_request_count = 0; - account.window_input_tokens = 0; - account.window_output_tokens = 0; - account.window_cached_tokens = 0; - account.window_image_input_tokens = 0; - account.window_image_output_tokens = 0; - account.window_image_request_count = 0; - account.window_image_request_failed_count = 0; - account.window_started_at = account.window_reset_at.map(|_| Utc::now()); - true - } - /// 标记账号配额限流直到指定时间。 pub fn mark_quota_limited_until( &mut self, @@ -327,15 +263,6 @@ impl AccountPool { true } - /// 设置账号是否需要额外配额校验。 - pub fn set_quota_verify_required(&mut self, account_id: &str, required: bool) -> bool { - let Some(account) = self.accounts.get_mut(account_id) else { - return false; - }; - account.quota_verify_required = required; - true - } - /// 应用已经验证过的配额状态。 pub fn apply_quota_state( &mut self, @@ -433,24 +360,18 @@ impl AccountPool { true } - /// 使用当前时间为指定模型获取账号。 - pub fn acquire(&mut self, model: &str) -> Option { - self.acquire_with(AccountAcquireRequest::new(model, Utc::now())) - .map(|acquired| acquired.account) - } - /// 使用完整调度请求获取账号。 - pub fn acquire_with(&mut self, request: AccountAcquireRequest) -> Option { + pub fn acquire_with(&mut self, request: &AccountAcquireRequest) -> Option { self.acquire_with_status_refresh(request).acquired } fn acquire_with_status_refresh( &mut self, - request: AccountAcquireRequest, + request: &AccountAcquireRequest, ) -> AccountAcquireWithStatusRefresh { self.cleanup_stale_slots(request.now); let expired_account_ids = self.refresh_account_statuses(request.now); - let candidates = self.candidates(&request); + let candidates = self.candidates(request); let selected = if let Some(preferred_account_id) = &request.preferred_account_id { let preferred = candidates .iter() @@ -473,7 +394,7 @@ impl AccountPool { (account, previous) } RotationStrategy::Sticky => { - let account = self.select_sticky(&candidates)?; + let account = Self::select_sticky(&candidates)?; let previous = self.previous_slot_at(&account.id); (account, previous) } @@ -562,47 +483,6 @@ impl AccountPool { } } - /// 计算账号池状态摘要。 - pub fn status_summary(&mut self, now: DateTime) -> AccountPoolStatusSummary { - self.cleanup_stale_slots(now); - self.refresh_account_statuses(now); - let mut summary = AccountPoolStatusSummary { - total: self.accounts.len(), - ..AccountPoolStatusSummary::default() - }; - for account in self.accounts.values() { - match account.status { - AccountStatus::Active - if !AccountService::quota_available_at( - account, - now, - self.options.skip_quota_limited, - ) => - { - summary.rate_limited += 1; - } - AccountStatus::Active => summary.active += 1, - AccountStatus::Expired => summary.expired += 1, - AccountStatus::QuotaExhausted => summary.quota_exhausted += 1, - AccountStatus::Refreshing => summary.refreshing += 1, - AccountStatus::Disabled => summary.disabled += 1, - AccountStatus::Banned => summary.banned += 1, - } - } - summary - } - - /// 获取所有处于配额锁定状态的账号 ID。 - pub fn list_quota_locked_accounts(&self) -> Vec { - self.accounts - .values() - .filter(|account| { - account.status == AccountStatus::Active && account.quota_limit_reached - }) - .map(|account| account.id.clone()) - .collect() - } - fn select_least_used(&mut self, candidates: &[Account]) -> Option { let best_key = candidates .iter() @@ -618,7 +498,7 @@ impl AccountPool { Some((*tied[index]).clone()) } - fn select_sticky(&self, candidates: &[Account]) -> Option { + fn select_sticky(candidates: &[Account]) -> Option { candidates .iter() .max_by_key(|account| account.last_used_at.clone()) @@ -793,8 +673,10 @@ impl RuntimeAccountPoolService { } /// 更新账号池运行参数。 - pub async fn apply_options(&self, options: AccountPoolOptions, request_interval_ms: u64) { - self.pool.lock().await.set_options(options); + pub async fn apply_options(&self, mut options: AccountPoolOptions, request_interval_ms: u64) { + let mut pool = self.pool.lock().await; + options.model_plan_allowlist = pool.options.model_plan_allowlist.clone(); + pool.set_options(options); *self .request_interval .write() @@ -802,6 +684,14 @@ impl RuntimeAccountPoolService { StdDuration::from_millis(request_interval_ms); } + /// 更新模型到可用账号计划的调度约束。 + pub(crate) async fn apply_model_plan_allowlist( + &self, + model_plan_allowlist: BTreeMap>, + ) { + self.pool.lock().await.options.model_plan_allowlist = model_plan_allowlist; + } + /// 从持久化账号恢复运行时账号池。 pub async fn restore_from_repository(&self) -> Result { let accounts = self @@ -818,11 +708,6 @@ impl RuntimeAccountPoolService { Ok(count) } - /// 清空运行时账号池。 - pub async fn clear(&self) { - self.pool.lock().await.clear(); - } - /// 读取账号池容量摘要。 pub async fn capacity_summary(&self, now: DateTime) -> AccountCapacitySummary { self.pool.lock().await.capacity_summary(now) @@ -877,7 +762,7 @@ impl RuntimeAccountPoolService { } /// 按请求上下文获取可用账号。 - pub async fn acquire_with(&self, request: AccountAcquireRequest) -> Option { + pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option { let model = request.model.clone(); let refresh = self.pool.lock().await.acquire_with_status_refresh(request); self.persist_expired_statuses(refresh.expired_account_ids) @@ -1162,11 +1047,6 @@ impl RuntimeAccountPoolService { self.pool.lock().await.remove(account_id) } - /// 从运行时账号池删除账号。 - pub async fn delete_account(&self, account_id: &str) -> bool { - self.remove_account(account_id).await - } - /// 从仓储同步单个账号到运行时账号池。 pub async fn sync_account_from_repository( &self, @@ -1185,11 +1065,6 @@ impl RuntimeAccountPoolService { Ok(pool.remove(account_id)) } - /// 重置运行时账号用量。 - pub async fn reset_usage(&self, account_id: &str) -> bool { - self.pool.lock().await.reset_usage(account_id) - } - /// 设置 Cloudflare 冷却状态。 pub async fn set_cloudflare_cooldown_until( &self, @@ -1298,7 +1173,11 @@ fn next_window_reset_at( return None; } - let elapsed_seconds = now.signed_duration_since(reset_at).num_seconds().max(0) as u64; + let elapsed_seconds = now + .signed_duration_since(reset_at) + .num_seconds() + .max(0) + .cast_unsigned(); let windows_to_advance = elapsed_seconds / window_seconds + 1; let advance_seconds = window_seconds .saturating_mul(windows_to_advance) @@ -1342,7 +1221,7 @@ fn compare_least_used_group(a: &LeastUsedGroupKey, b: &LeastUsedGroupKey) -> Ord fn compare_window_reset(a: Option>, b: Option>) -> Ordering { match (a, b) { (Some(a), Some(b)) => a.cmp(&b), - (Some(_), None) | (None, Some(_)) | (None, None) => Ordering::Equal, + (Some(_) | None, None) | (None, Some(_)) => Ordering::Equal, } } diff --git a/src/upstream/accounts/quota/mod.rs b/src/upstream/accounts/quota/mod.rs index 74afac7f4..2a63b9941 100644 --- a/src/upstream/accounts/quota/mod.rs +++ b/src/upstream/accounts/quota/mod.rs @@ -1,16 +1,13 @@ //! 配额检查策略。 +use std::time::Duration; + use chrono::{DateTime, Utc}; use serde_json::{json, Value}; const MONTH_WINDOW_MINUTES: u64 = 43_200; const MONTH_WINDOW_SECONDS: u64 = 2_592_000; -/// 判断配额百分比是否达到阈值。 -pub fn quota_reached(used_percent: f64, threshold: f64) -> bool { - used_percent >= threshold -} - /// 判断配额快照是否已经触顶。 pub fn quota_snapshot_limit_reached(quota: &Value) -> bool { core_snapshot_blocked(quota) @@ -125,7 +122,8 @@ fn window_from_rate_limit(window: Option<&Value>) -> Value { .get("limit_window_seconds") .and_then(number_value) .filter(|seconds| *seconds > 0.0) - .map(|seconds| (seconds / 60.0).round() as u64); + .and_then(|seconds| Duration::try_from_secs_f64((seconds / 60.0).round()).ok()) + .map(|duration| duration.as_secs()); json!({ "used_percent": used_percent, "remaining_percent": remaining_percent(used_percent), diff --git a/src/upstream/accounts/quota/runtime.rs b/src/upstream/accounts/quota/runtime.rs index 8aa6f8411..9231476be 100644 --- a/src/upstream/accounts/quota/runtime.rs +++ b/src/upstream/accounts/quota/runtime.rs @@ -45,8 +45,7 @@ impl RuntimeQuotaRefreshService { } } - /// 返回连续 quota 请求之间的默认错峰间隔。 - pub fn default_request_spacing() -> Duration { + fn default_request_spacing() -> Duration { Duration::from_secs(DEFAULT_REQUEST_SPACING_SECS) } @@ -84,14 +83,6 @@ impl RuntimeQuotaRefreshService { self } - /// 执行一次配额锁定或待验证账号刷新。 - pub async fn refresh_locked_accounts_once( - &self, - ) -> QuotaRefreshServiceResult { - let mut last_refreshed = HashMap::new(); - self.refresh_locked_accounts(&mut last_refreshed).await - } - /// 执行一次配额锁定或待验证账号刷新，复用调用方传入的最近刷新记录。 pub async fn refresh_locked_accounts( &self, diff --git a/src/upstream/accounts/service.rs b/src/upstream/accounts/service.rs index 3df3ad0ab..9ae052750 100644 --- a/src/upstream/accounts/service.rs +++ b/src/upstream/accounts/service.rs @@ -2,40 +2,13 @@ use chrono::{DateTime, Utc}; -use crate::upstream::accounts::model::{Account, AccountStatus}; - -/// 根据 token 过期时间推导账号运行时状态。 -pub fn effective_status(account: &Account, now: DateTime) -> AccountStatus { - if account.status != AccountStatus::Active { - return account.status; - } - if account - .access_token_expires_at - .is_some_and(|expires_at| expires_at <= now) - { - AccountStatus::Expired - } else { - AccountStatus::Active - } -} +use crate::upstream::accounts::model::Account; /// 账号服务。 #[derive(Debug, Clone, Default)] pub struct AccountService; impl AccountService { - /// 判断账号是否可参与请求调度。 - pub fn can_serve(account: &Account) -> bool { - account.status == AccountStatus::Active && !account.quota_limit_reached - } - - /// 判断账号在指定时间是否处于可服务的基础状态。 - pub fn can_serve_at(account: &Account, now: DateTime, skip_quota_limited: bool) -> bool { - account.status == AccountStatus::Active - && Self::quota_available_at(account, now, skip_quota_limited) - && Self::cloudflare_available_at(account, now) - } - /// 判断账号配额锁定是否已经释放。 pub fn quota_available_at( account: &Account, diff --git a/src/upstream/accounts/store.rs b/src/upstream/accounts/store.rs index 552576c8d..6c205a407 100644 --- a/src/upstream/accounts/store.rs +++ b/src/upstream/accounts/store.rs @@ -350,61 +350,6 @@ select from account_model_usage order by account_id asc, request_count desc, last_used_at desc, model asc"; -const GET_USAGE_SQL: &str = r" -select - account_id, - request_count, - empty_response_count, - input_tokens, - output_tokens, - cached_tokens, - reasoning_tokens, - total_tokens, - image_input_tokens, - image_output_tokens, - image_request_count, - image_request_failed_count, - window_request_count, - window_input_tokens, - window_output_tokens, - window_cached_tokens, - window_image_input_tokens, - window_image_output_tokens, - window_image_request_count, - window_image_request_failed_count, - window_started_at, - window_reset_at, - limit_window_seconds, - last_used_at -from account_usage -where account_id = ?"; - -const RESET_USAGE_SQL: &str = r" -update account_usage -set - request_count = 0, - empty_response_count = 0, - input_tokens = 0, - output_tokens = 0, - cached_tokens = 0, - reasoning_tokens = 0, - total_tokens = 0, - image_input_tokens = 0, - image_output_tokens = 0, - image_request_count = 0, - image_request_failed_count = 0, - window_request_count = 0, - window_input_tokens = 0, - window_output_tokens = 0, - window_cached_tokens = 0, - window_image_input_tokens = 0, - window_image_output_tokens = 0, - window_image_request_count = 0, - window_image_request_failed_count = 0, - window_started_at = null, - last_used_at = null -where account_id = ?"; - const LIST_QUOTA_SNAPSHOTS_SQL: &str = r" select id, @@ -585,8 +530,6 @@ set updated_at = ? where id = ?"; -const DELETE_ALL_ACCOUNTS_SQL: &str = "delete from accounts"; - const DELETE_ACCOUNT_SQL: &str = "delete from accounts where id = ?"; // ============================================================================ @@ -757,59 +700,6 @@ impl AccountMetadataUpdate { } } -/// 账号用量记录。 -#[derive(Debug, Clone)] -pub struct AccountUsageRecord { - /// 账号 ID。 - pub account_id: String, - /// 历史请求总数。 - pub request_count: i64, - /// 历史空响应次数。 - pub empty_response_count: i64, - /// 累计输入 token。 - pub input_tokens: i64, - /// 累计输出 token。 - pub output_tokens: i64, - /// 累计缓存 token。 - pub cached_tokens: i64, - /// 累计 reasoning token。 - pub reasoning_tokens: i64, - /// 累计总 token。 - pub total_tokens: i64, - /// 累计图片输入 token。 - pub image_input_tokens: i64, - /// 累计图片输出 token。 - pub image_output_tokens: i64, - /// 累计图片请求成功次数。 - pub image_request_count: i64, - /// 累计图片请求失败次数。 - pub image_request_failed_count: i64, - /// 当前窗口请求数。 - pub window_request_count: i64, - /// 当前窗口输入 token。 - pub window_input_tokens: i64, - /// 当前窗口输出 token。 - pub window_output_tokens: i64, - /// 当前窗口缓存 token。 - pub window_cached_tokens: i64, - /// 当前窗口图片输入 token。 - pub window_image_input_tokens: i64, - /// 当前窗口图片输出 token。 - pub window_image_output_tokens: i64, - /// 当前窗口图片请求数。 - pub window_image_request_count: i64, - /// 当前窗口图片请求失败数。 - pub window_image_request_failed_count: i64, - /// 当前窗口起始时间。 - pub window_started_at: Option>, - /// 当前窗口重置时间。 - pub window_reset_at: Option>, - /// 限流窗口大小。 - pub limit_window_seconds: Option, - /// 最近使用时间。 - pub last_used_at: Option>, -} - /// 账号模型维度用量记录。 #[derive(Debug, Clone)] pub struct AccountModelUsageRecord { @@ -1010,8 +900,7 @@ impl SqliteAccountStore { .bind( account .added_at - .map(|dt| dt.to_rfc3339()) - .unwrap_or_else(|| Utc::now().to_rfc3339()), + .map_or_else(|| Utc::now().to_rfc3339(), |dt| dt.to_rfc3339()), ) .bind(&now) .execute(&self.pool) @@ -1061,7 +950,7 @@ impl SqliteAccountStore { .fetch_all(&self.pool) .await?; Ok(to_page( - rows, + &rows, limit, stored_account_from_row, ("added_at", "id"), @@ -1072,7 +961,7 @@ impl SqliteAccountStore { .fetch_all(&self.pool) .await?; Ok(to_page( - rows, + &rows, limit, stored_account_from_row, ("added_at", "id"), @@ -1097,13 +986,13 @@ impl SqliteAccountStore { .bind(limit + 1) .fetch_all(&self.pool) .await?; - Ok(to_page(rows, limit, metadata_from_row, ("added_at", "id"))) + Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) } else { let rows = sqlx::query(LIST_ACCOUNT_METADATA_SQL) .bind(limit + 1) .fetch_all(&self.pool) .await?; - Ok(to_page(rows, limit, metadata_from_row, ("added_at", "id"))) + Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) } } @@ -1426,27 +1315,6 @@ impl SqliteAccountStore { rows.iter().map(model_usage_from_row).collect() } - /// 获取用量记录。 - pub async fn get_usage( - &self, - account_id: &str, - ) -> SqliteAccountStoreResult> { - let row = sqlx::query(GET_USAGE_SQL) - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - row.as_ref().map(usage_from_row).transpose() - } - - /// 重置用量。 - pub async fn reset_usage(&self, account_id: &str) -> SqliteAccountStoreResult { - let result = sqlx::query(RESET_USAGE_SQL) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - /// 读取配额快照列表。 pub async fn list_quota_snapshots( &self, @@ -1551,10 +1419,11 @@ impl SqliteAccountStore { let should_reset = existing .as_ref() .map(|row| { - let existing_reset_at = - parse_optional_rfc3339(row.get::, _>("window_reset_at")) - .ok() - .flatten(); + let existing_reset_at = parse_optional_rfc3339( + row.get::, _>("window_reset_at").as_deref(), + ) + .ok() + .flatten(); let existing_limit_window_seconds = optional_positive_i64_to_u64(row.get::, _>("limit_window_seconds")); should_reset_usage_window( @@ -1587,14 +1456,6 @@ impl SqliteAccountStore { Ok(()) } - /// 删除全部账号。 - pub async fn delete_all(&self) -> SqliteAccountStoreResult { - let result = sqlx::query(DELETE_ALL_ACCOUNTS_SQL) - .execute(&self.pool) - .await?; - Ok(result.rows_affected()) - } - /// 删除单个账号。 pub async fn delete(&self, account_id: &str) -> SqliteAccountStoreResult { let result = sqlx::query(DELETE_ACCOUNT_SQL) @@ -1622,16 +1483,6 @@ impl SqliteAccountStore { .map_err(SqliteAccountStoreError::Database)?; Ok(result.rows_affected() > 0) } - - /// 按 ID 删除账号。 - pub async fn delete_account(&self, account_id: &str) -> Result { - let result = sqlx::query("DELETE FROM accounts WHERE id = ?1") - .bind(account_id) - .execute(&self.pool) - .await - .map_err(SqliteAccountStoreError::Database)?; - Ok(result.rows_affected() > 0) - } } #[async_trait] @@ -1639,13 +1490,13 @@ impl AccountStore for SqliteAccountStore { async fn list_pool_accounts(&self) -> AccountStoreResult> { list_pool_accounts(self) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult> { get_pool_account(self, account_id) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn mark_quota_limited_until( @@ -1655,7 +1506,7 @@ impl AccountStore for SqliteAccountStore { ) -> AccountStoreResult { SqliteAccountStore::mark_quota_limited_until(self, account_id, cooldown_until) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn set_cloudflare_cooldown_until( @@ -1665,7 +1516,7 @@ impl AccountStore for SqliteAccountStore { ) -> AccountStoreResult { SqliteAccountStore::set_cloudflare_cooldown_until(self, account_id, cooldown_until) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn set_status( @@ -1675,7 +1526,7 @@ impl AccountStore for SqliteAccountStore { ) -> AccountStoreResult { SqliteAccountStore::set_status(self, account_id, status) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn record_usage_delta( @@ -1685,7 +1536,7 @@ impl AccountStore for SqliteAccountStore { ) -> AccountStoreResult<()> { self.record_usage(account_id, sqlite_usage_delta(usage)) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn record_model_usage_delta( @@ -1696,13 +1547,13 @@ impl AccountStore for SqliteAccountStore { ) -> AccountStoreResult<()> { self.record_model_usage(account_id, model, usage) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult> { SqliteAccountStore::get_quota_json(self, account_id) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn update_quota_json( @@ -1712,7 +1563,7 @@ impl AccountStore for SqliteAccountStore { ) -> AccountStoreResult { SqliteAccountStore::update_quota_json(self, account_id, quota_json) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn apply_quota_snapshot( @@ -1730,7 +1581,7 @@ impl AccountStore for SqliteAccountStore { cooldown_until, ) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } async fn sync_rate_limit_window( @@ -1741,7 +1592,7 @@ impl AccountStore for SqliteAccountStore { ) -> AccountStoreResult<()> { SqliteAccountStore::sync_rate_limit_window(self, account_id, reset_at, limit_window_seconds) .await - .map_err(map_account_store_error) + .map_err(|error| map_account_store_error(&error)) } } @@ -1784,17 +1635,22 @@ fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult { access_token: row.get("access_token"), refresh_token: row.get("refresh_token"), access_token_expires_at: parse_optional_rfc3339( - row.get::, _>("access_token_expires_at"), + row.get::, _>("access_token_expires_at") + .as_deref(), + )?, + next_refresh_at: parse_optional_rfc3339( + row.get::, _>("next_refresh_at").as_deref(), )?, - next_refresh_at: parse_optional_rfc3339(row.get::, _>("next_refresh_at"))?, status: status_from_db(&row.get::("status"))?, quota_limit_reached: row.get::("quota_limit_reached") != 0, quota_verify_required: row.get::("quota_verify_required") != 0, quota_cooldown_until: parse_optional_rfc3339( - row.get::, _>("quota_cooldown_until"), + row.get::, _>("quota_cooldown_until") + .as_deref(), )?, cloudflare_cooldown_until: parse_optional_rfc3339( - row.get::, _>("cloudflare_cooldown_until"), + row.get::, _>("cloudflare_cooldown_until") + .as_deref(), )?, request_count: nonnegative_i64_to_u64(row.get::("usage_request_count")), empty_response_count: nonnegative_i64_to_u64( @@ -1829,10 +1685,12 @@ fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult { row.get::("usage_window_image_request_failed_count"), ), window_started_at: parse_optional_rfc3339( - row.get::, _>("usage_window_started_at"), + row.get::, _>("usage_window_started_at") + .as_deref(), )?, window_reset_at: parse_optional_rfc3339( - row.get::, _>("usage_window_reset_at"), + row.get::, _>("usage_window_reset_at") + .as_deref(), )?, limit_window_seconds: optional_positive_i64_to_u64( row.get::, _>("usage_limit_window_seconds"), @@ -1855,9 +1713,12 @@ fn stored_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult, _>("refresh_token") .map(|token| SecretString::new(token.into())), access_token_expires_at: parse_optional_rfc3339( - row.get::, _>("access_token_expires_at"), + row.get::, _>("access_token_expires_at") + .as_deref(), + )?, + next_refresh_at: parse_optional_rfc3339( + row.get::, _>("next_refresh_at").as_deref(), )?, - next_refresh_at: parse_optional_rfc3339(row.get::, _>("next_refresh_at"))?, status: status_from_db(&row.get::("status"))?, added_at: row.get("added_at"), updated_at: row.get("updated_at"), @@ -1873,7 +1734,8 @@ fn metadata_from_row(row: &SqliteRow) -> SqliteAccountStoreResult, _>("access_token_expires_at"), + row.get::, _>("access_token_expires_at") + .as_deref(), )?, status: status_from_db(&row.get::("status"))?, added_at: row.get("added_at"), @@ -1881,7 +1743,7 @@ fn metadata_from_row(row: &SqliteRow) -> SqliteAccountStoreResult AccountStoreError { +fn map_account_store_error(error: &impl ToString) -> AccountStoreError { AccountStoreError::OperationFailed { message: error.to_string(), } @@ -1958,39 +1820,6 @@ fn quota_plan_type(quota_json: &str) -> Option { .map(ToString::to_string) } -fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> SqliteAccountStoreResult { - Ok(AccountUsageRecord { - account_id: row.get("account_id"), - request_count: row.get("request_count"), - empty_response_count: row.get("empty_response_count"), - input_tokens: row.get("input_tokens"), - output_tokens: row.get("output_tokens"), - cached_tokens: row.get("cached_tokens"), - reasoning_tokens: row.get("reasoning_tokens"), - total_tokens: row.get("total_tokens"), - image_input_tokens: row.get("image_input_tokens"), - image_output_tokens: row.get("image_output_tokens"), - image_request_count: row.get("image_request_count"), - image_request_failed_count: row.get("image_request_failed_count"), - window_request_count: row.get("window_request_count"), - window_input_tokens: row.get("window_input_tokens"), - window_output_tokens: row.get("window_output_tokens"), - window_cached_tokens: row.get("window_cached_tokens"), - window_image_input_tokens: row.get("window_image_input_tokens"), - window_image_output_tokens: row.get("window_image_output_tokens"), - window_image_request_count: row.get("window_image_request_count"), - window_image_request_failed_count: row.get("window_image_request_failed_count"), - window_started_at: parse_optional_rfc3339( - row.get::, _>("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"))?, - }) -} - fn model_usage_from_row( row: &sqlx::sqlite::SqliteRow, ) -> SqliteAccountStoreResult { @@ -2002,7 +1831,9 @@ fn model_usage_from_row( input_tokens: row.get("input_tokens"), output_tokens: row.get("output_tokens"), cached_tokens: row.get("cached_tokens"), - last_used_at: parse_optional_rfc3339(row.get::, _>("last_used_at"))?, + last_used_at: parse_optional_rfc3339( + row.get::, _>("last_used_at").as_deref(), + )?, }) } @@ -2013,7 +1844,9 @@ fn quota_snapshot_from_row( account_id: row.get("id"), email: row.get("email"), quota_json: row.get("quota_json"), - quota_fetched_at: parse_optional_rfc3339(row.get::, _>("quota_fetched_at"))?, + quota_fetched_at: parse_optional_rfc3339( + row.get::, _>("quota_fetched_at").as_deref(), + )?, }) } @@ -2044,10 +1877,8 @@ fn status_from_db(value: &str) -> SqliteAccountStoreResult { } } -fn parse_optional_rfc3339( - value: Option, -) -> SqliteAccountStoreResult>> { - value.as_deref().map(parse_rfc3339).transpose() +fn parse_optional_rfc3339(value: Option<&str>) -> SqliteAccountStoreResult>> { + value.map(parse_rfc3339).transpose() } fn parse_rfc3339(value: &str) -> SqliteAccountStoreResult> { @@ -2055,7 +1886,7 @@ fn parse_rfc3339(value: &str) -> SqliteAccountStoreResult> { } fn nonnegative_i64_to_u64(value: i64) -> u64 { - value.max(0) as u64 + value.max(0).cast_unsigned() } fn u64_to_i64_saturating(value: u64) -> i64 { @@ -2075,7 +1906,7 @@ async fn count_account_metadata( let mut builder = QueryBuilder::::new("select count(*) from accounts"); push_account_metadata_search(&mut builder, search); let (total,): (i64,) = builder.build_query_as().fetch_one(pool).await?; - Ok(total.max(0) as u64) + Ok(total.max(0).cast_unsigned()) } fn push_account_metadata_search(builder: &mut QueryBuilder, search: Option<&str>) { @@ -2097,7 +1928,7 @@ fn push_account_metadata_search(builder: &mut QueryBuilder, search: Opti } fn to_page( - rows: Vec, + rows: &[SqliteRow], limit: u32, mapper: impl Fn(&SqliteRow) -> SqliteAccountStoreResult, cursor_fields: (&str, &str), diff --git a/src/upstream/accounts/token_refresh/mod.rs b/src/upstream/accounts/token_refresh/mod.rs index a21cd7832..886cee918 100644 --- a/src/upstream/accounts/token_refresh/mod.rs +++ b/src/upstream/accounts/token_refresh/mod.rs @@ -364,11 +364,6 @@ impl RefreshLeaseStore { Self { pool } } - /// 返回底层连接池。 - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - /// 尝试获取账号刷新租约。 pub async fn try_acquire( &self, @@ -378,7 +373,7 @@ impl RefreshLeaseStore { now: DateTime, ) -> RefreshLeaseStoreResult { let result = sqlx::query( - r#" + r" insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?) on conflict(account_id) do update set @@ -387,7 +382,7 @@ on conflict(account_id) do update set updated_at = excluded.updated_at where account_refresh_leases.expires_at <= ? or account_refresh_leases.owner = ? -"#, +", ) .bind(account_id) .bind(owner) diff --git a/src/upstream/accounts/token_refresh/runtime.rs b/src/upstream/accounts/token_refresh/runtime.rs index e41db7a9e..fc977a4fd 100644 --- a/src/upstream/accounts/token_refresh/runtime.rs +++ b/src/upstream/accounts/token_refresh/runtime.rs @@ -113,13 +113,6 @@ where self } - /// 执行一次到期刷新扫描。 - pub async fn refresh_due_accounts_once( - &self, - ) -> TokenRefreshServiceResult { - self.refresh_due_accounts_once_at(Utc::now()).await - } - /// 执行一次账号级刷新定时器调度。 pub async fn schedule_account_timers_once( &self, @@ -530,11 +523,8 @@ where } Ok(_) => return Ok(TokenRefreshOutcome::Skipped), Err(RefreshError::RetryableTransport) - if self.sleep_before_retry(attempt_index).await => - { - continue; - } - Err(RefreshError::RetryableTransport) | Err(RefreshError::Transport) => { + if self.sleep_before_retry(attempt_index).await => {} + Err(RefreshError::RetryableTransport | RefreshError::Transport) => { persist_status(&self.store, &attempt.id, AccountStatus::Active).await?; let next_refresh_at = now + chrono::Duration::seconds(RECOVERY_DELAY_SECONDS); persist_next_refresh_at(&self.store, &attempt.id, Some(next_refresh_at)) diff --git a/src/upstream/fingerprint.rs b/src/upstream/fingerprint.rs index 6bd5055f3..1db21edaf 100644 --- a/src/upstream/fingerprint.rs +++ b/src/upstream/fingerprint.rs @@ -11,15 +11,13 @@ use tokio::sync::Mutex; use tracing::{info, warn}; use uuid::Uuid; -/// 指纹历史记录的来源标识。 -pub const CODEX_DESKTOP_UPDATE_SOURCE: &str = "codex_desktop_update_source"; /// 运行时当前指纹槽位。 pub const CURRENT_FINGERPRINT_ID: &str = "current"; const AUTO_UPDATE_SOURCE: &str = "auto_update"; const CONFIG_SEED_SOURCE: &str = "config_seed"; const APPCAST_TIMEOUT: Duration = Duration::from_secs(30); -const POLL_INTERVAL: Duration = Duration::from_secs(3 * 24 * 60 * 60); +const POLL_INTERVAL: Duration = Duration::from_hours(72); type AppcastFields = (Option, Option, Option); @@ -74,64 +72,6 @@ impl Fingerprint { } } - /// 返回默认的 Codex Desktop 指纹。 - pub fn default_codex_desktop() -> Self { - Self { - originator: "Codex Desktop".to_string(), - app_version: "26.519.81530".to_string(), - build_number: "3178".to_string(), - platform: "darwin".to_string(), - arch: "arm64".to_string(), - chromium_version: "146".to_string(), - user_agent_template: "Codex Desktop/{version} ({platform}; {arch})".to_string(), - default_headers: Self::default_headers(), - header_order: Self::default_header_order(), - updated_at: None, - } - } - - /// 返回默认请求头集合。 - pub fn default_headers() -> IndexMap { - let mut headers = IndexMap::new(); - headers.insert( - "Accept-Encoding".to_string(), - "gzip, deflate, br, zstd".to_string(), - ); - headers.insert("Accept-Language".to_string(), "en-US,en;q=0.9".to_string()); - headers.insert("sec-ch-ua-mobile".to_string(), "?0".to_string()); - headers.insert("sec-ch-ua-platform".to_string(), "\"macOS\"".to_string()); - headers.insert("sec-fetch-site".to_string(), "same-origin".to_string()); - headers.insert("sec-fetch-mode".to_string(), "cors".to_string()); - headers.insert("sec-fetch-dest".to_string(), "empty".to_string()); - headers - } - - /// 返回默认请求头顺序。 - pub fn default_header_order() -> Vec { - vec![ - "authorization".to_string(), - "chatgpt-account-id".to_string(), - "originator".to_string(), - "x-openai-internal-codex-residency".to_string(), - "x-client-request-id".to_string(), - "x-codex-installation-id".to_string(), - "x-codex-turn-state".to_string(), - "openai-beta".to_string(), - "user-agent".to_string(), - "sec-ch-ua".to_string(), - "sec-ch-ua-mobile".to_string(), - "sec-ch-ua-platform".to_string(), - "accept-encoding".to_string(), - "accept-language".to_string(), - "sec-fetch-site".to_string(), - "sec-fetch-mode".to_string(), - "sec-fetch-dest".to_string(), - "content-type".to_string(), - "accept".to_string(), - "cookie".to_string(), - ] - } - /// 根据模板展开最终 User-Agent。 pub fn user_agent(&self) -> String { self.user_agent_template @@ -149,63 +89,10 @@ impl Fingerprint { } } -// --------------------------------------------------------------------------- -// 更新清单解析 -// --------------------------------------------------------------------------- - -/// 更新清单解析结果。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct FingerprintUpdate { - /// 应用版本。 - pub app_version: String, - /// 构建号。 - pub build_number: String, -} - -/// 指纹更新错误。 -#[derive(Debug, Error)] -pub enum FingerprintError { - /// 更新清单 JSON 无效。 - #[error("invalid update manifest: {0}")] - InvalidManifest(#[from] serde_json::Error), - /// HTTP 请求失败。 - #[error("failed to fetch update manifest: {0}")] - Http(#[from] reqwest::Error), - /// 数据库存储失败。 - #[error("failed to persist fingerprint update: {0}")] - Database(#[from] sqlx::Error), -} - -#[derive(Deserialize)] -struct Manifest { - version: String, - build_number: String, -} - -/// 从更新清单 JSON 中提取版本信息。 -pub fn parse_update_manifest(input: &str) -> Result { - let manifest: Manifest = serde_json::from_str(input)?; - Ok(FingerprintUpdate { - app_version: manifest.version, - build_number: manifest.build_number, - }) -} - // --------------------------------------------------------------------------- // 指纹历史记录与仓储 // --------------------------------------------------------------------------- -/// 指纹历史记录。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StoredFingerprint { - /// 应用版本。 - pub app_version: String, - /// 构建号。 - pub build_number: String, - /// 来源。 - pub source: String, -} - #[derive(Debug, Clone, Serialize, Deserialize)] struct StoredHeader { name: String, @@ -266,22 +153,6 @@ impl FingerprintRepository { Self { pool } } - /// 插入一条指纹更新历史。 - pub async fn insert_update(&self, update: &FingerprintUpdate) -> Result<(), sqlx::Error> { - sqlx::query( - "insert into fingerprint_update_history (id, current_fingerprint_id, app_version, build_number, source, created_at) values (?, ?, ?, ?, ?, ?)", - ) - .bind(Uuid::new_v4().to_string()) - .bind(CURRENT_FINGERPRINT_ID) - .bind(&update.app_version) - .bind(&update.build_number) - .bind(CODEX_DESKTOP_UPDATE_SOURCE) - .bind(Utc::now().to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(()) - } - /// 写入当前指纹默认值；如果当前槽位已存在，直接读取数据库值。 pub async fn ensure_current_seed( &self, @@ -293,7 +164,7 @@ impl FingerprintRepository { let now = Utc::now().to_rfc3339(); sqlx::query( - r#" + r" insert into fingerprints ( id, originator, @@ -309,7 +180,7 @@ impl FingerprintRepository { created_at, updated_at ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - "#, + ", ) .bind(CURRENT_FINGERPRINT_ID) .bind(&default_fingerprint.originator) @@ -371,7 +242,7 @@ impl FingerprintRepository { /// 读取当前运行时指纹。 pub async fn load_current(&self) -> Result, sqlx::Error> { let row = sqlx::query( - r#" + r" select originator, app_version, @@ -385,7 +256,7 @@ impl FingerprintRepository { updated_at from fingerprints where id = ? - "#, + ", ) .bind(CURRENT_FINGERPRINT_ID) .fetch_optional(&self.pool) @@ -417,21 +288,6 @@ impl FingerprintRepository { .await?; Ok(()) } - - /// 读取最新一条指纹历史记录。 - pub async fn latest(&self) -> Result, sqlx::Error> { - let row = sqlx::query( - "select app_version, build_number, source from fingerprint_update_history order by created_at desc, id desc limit 1", - ) - .fetch_optional(&self.pool) - .await?; - - Ok(row.map(|row| StoredFingerprint { - app_version: row.get("app_version"), - build_number: row.get("build_number"), - source: row.get("source"), - })) - } } fn fingerprint_from_row(row: &sqlx::sqlite::SqliteRow) -> Result { @@ -482,49 +338,6 @@ fn json_error(error: serde_json::Error) -> sqlx::Error { sqlx::Error::Decode(Box::new(error)) } -// --------------------------------------------------------------------------- -// 指纹更新器（HTTP 拉取 + 持久化） -// --------------------------------------------------------------------------- - -/// 通过 HTTP 拉取更新清单并写入历史记录。 -#[derive(Clone)] -pub struct FingerprintUpdater { - client: reqwest::Client, - repository: FingerprintRepository, - update_url: String, -} - -impl FingerprintUpdater { - /// 构造更新器。 - pub fn new( - client: reqwest::Client, - repository: FingerprintRepository, - update_url: impl Into, - ) -> Self { - Self { - client, - repository, - update_url: update_url.into(), - } - } - - /// 执行一次更新清单拉取并持久化历史记录。 - pub async fn poll_once(&self) -> Result { - let manifest = self - .client - .get(&self.update_url) - .timeout(APPCAST_TIMEOUT) - .send() - .await? - .error_for_status()? - .text() - .await?; - let update = parse_update_manifest(&manifest)?; - self.repository.insert_update(&update).await?; - Ok(update) - } -} - // --------------------------------------------------------------------------- // Appcast 指纹更新检查器 // --------------------------------------------------------------------------- diff --git a/src/upstream/models/mod.rs b/src/upstream/models/mod.rs index cdbce99d6..d642cae39 100644 --- a/src/upstream/models/mod.rs +++ b/src/upstream/models/mod.rs @@ -7,10 +7,7 @@ use serde::{Deserialize, Serialize}; use sqlx::Row; use thiserror::Error; -use crate::upstream::accounts::{ - model::{Account, AccountStatus}, - store::AccountStore, -}; +use crate::upstream::accounts::model::{Account, AccountStatus}; // --------------------------------------------------------------------------- // 数据类型 @@ -195,85 +192,6 @@ pub trait ModelSnapshotStore: Send + Sync + 'static { async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult>; } -// --------------------------------------------------------------------------- -// 管理端模型刷新服务 -// --------------------------------------------------------------------------- - -/// 管理端模型服务。 -#[derive(Clone)] -pub struct AdminModelService { - models: Arc, - accounts: Arc, - installation_id: Option, -} - -impl AdminModelService { - /// 构造服务。 - pub fn new( - models: Arc, - accounts: Arc, - installation_id: Option, - ) -> Self { - Self { - models, - accounts, - installation_id, - } - } - - /// 使用账号池中的账号刷新上游模型目录。 - pub async fn refresh_backend_models( - &self, - request_id: &str, - ) -> Result { - let accounts = self - .accounts - .list_pool_accounts() - .await - .map_err(|_| AdminModelError::ListAccounts)?; - self.models - .refresh_backend_models_with_installation_id( - &accounts, - request_id, - self.installation_id.as_deref(), - ) - .await - .map_err(AdminModelError::from) - } -} - -/// 管理端模型刷新错误。 -#[derive(Debug, Error)] -pub enum AdminModelError { - #[error("failed to list accounts")] - ListAccounts, - #[error("no active accounts available for model refresh")] - NoAccounts, - #[error("model snapshot store is unavailable")] - SnapshotStoreUnavailable, - #[error("model upstream client is unavailable")] - UpstreamClientUnavailable, - #[error("failed to store model snapshot")] - StoreSnapshot, - #[error("failed to load model snapshots")] - LoadSnapshots, - #[error("all model refresh plans failed")] - AllPlansFailed(ModelRefreshResult), -} - -impl From for AdminModelError { - fn from(error: ModelServiceError) -> Self { - match error { - ModelServiceError::SnapshotStoreUnavailable => Self::SnapshotStoreUnavailable, - ModelServiceError::UpstreamClientUnavailable => Self::UpstreamClientUnavailable, - ModelServiceError::NoAccounts => Self::NoAccounts, - ModelServiceError::StoreSnapshot => Self::StoreSnapshot, - ModelServiceError::LoadSnapshots => Self::LoadSnapshots, - ModelServiceError::AllPlansFailed(result) => Self::AllPlansFailed(result), - } - } -} - /// SQLite 模型快照存储。 #[derive(Clone)] pub struct SqliteModelSnapshotStore { @@ -658,16 +576,12 @@ pub enum ModelServiceError { /// 模型到可用计划的映射。 pub type ModelPlanAllowlist = BTreeMap>; -/// 可共享更新的模型计划映射缓存。 -pub type SharedModelPlanAllowlist = Arc>; - /// 模型目录服务。 #[derive(Clone)] pub struct ModelService { config: ModelConfig, snapshot_store: Option>, upstream_client: Option>, - model_plan_allowlist: Option, } impl ModelService { @@ -676,21 +590,14 @@ impl ModelService { config: ModelConfig, snapshot_store: Option>, upstream_client: Option>, - model_plan_allowlist: Option, ) -> Self { Self { config, snapshot_store, upstream_client, - model_plan_allowlist, } } - /// 返回模型目录的静态配置。 - pub fn config(&self) -> &ModelConfig { - &self.config - } - /// 构造当前对外暴露的模型目录。 pub async fn catalog(&self) -> ModelCatalog { let Some(snapshot_store) = self.snapshot_store.as_ref() else { @@ -709,16 +616,6 @@ impl ModelService { } } - /// 刷新活跃账号对应的后端模型目录。 - pub async fn refresh_backend_models( - &self, - accounts: &[Account], - request_id: &str, - ) -> Result { - self.refresh_backend_models_with_installation_id(accounts, request_id, None) - .await - } - /// 使用运行时 installation id 刷新活跃账号对应的后端模型目录。 pub async fn refresh_backend_models_with_installation_id( &self, @@ -781,10 +678,7 @@ impl ModelService { return Err(ModelServiceError::AllPlansFailed(result)); } - let allowlist = self.model_plan_allowlist_from_store().await?; - if let Some(shared_allowlist) = &self.model_plan_allowlist { - *shared_allowlist.lock().await = allowlist; - } + self.model_plan_allowlist_from_store().await?; Ok(result) } diff --git a/src/upstream/protocol/events.rs b/src/upstream/protocol/events.rs index 8b4a01788..1f5dd14df 100644 --- a/src/upstream/protocol/events.rs +++ b/src/upstream/protocol/events.rs @@ -1,9 +1,5 @@ -use std::{ - collections::BTreeMap, - time::{SystemTime, UNIX_EPOCH}, -}; +use std::{collections::BTreeMap, time::Duration}; -use chrono::{DateTime, Duration, Utc}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -29,34 +25,6 @@ pub struct TokenUsage { pub total_tokens: u64, } -impl TokenUsage { - /// 构造仅包含主模型 token 的用量快照。 - pub fn new(input_tokens: u64, output_tokens: u64, cached_tokens: u64) -> Self { - Self { - input_tokens, - output_tokens, - cached_tokens, - reasoning_tokens: 0, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: input_tokens + output_tokens, - } - } - - /// 将两份用量按字段累加。 - pub fn merged(self, other: Self) -> Self { - Self { - input_tokens: self.input_tokens + other.input_tokens, - output_tokens: self.output_tokens + other.output_tokens, - cached_tokens: self.cached_tokens + other.cached_tokens, - reasoning_tokens: self.reasoning_tokens + other.reasoning_tokens, - image_input_tokens: self.image_input_tokens + other.image_input_tokens, - image_output_tokens: self.image_output_tokens + other.image_output_tokens, - total_tokens: self.total_tokens + other.total_tokens, - } - } -} - /// 从单个 JSON 响应体中提取用量。 /// /// 该函数同时支持 Codex usage 结构和 OpenAI usage 结构。 @@ -198,20 +166,6 @@ pub struct RateLimitWindow { pub reset_at: Option, } -impl RateLimitWindow { - /// 以秒返回限流窗口大小。 - pub fn limit_window_seconds(self) -> Option { - self.window_minutes - .and_then(|minutes| minutes.checked_mul(60)) - } - - /// 将重置时间戳转换为 UTC 时间。 - pub fn reset_at_datetime(self) -> Option> { - let reset_at = self.reset_at?; - DateTime::::from_timestamp(reset_at, 0) - } -} - /// 复合限流信息，通常用于 code review 配额。 #[derive(Debug, Clone, Copy, PartialEq)] pub struct RateLimitDetails { @@ -236,24 +190,6 @@ pub struct ParsedRateLimits { pub code_review: Option, } -impl ParsedRateLimits { - /// 返回主窗口的 UTC 重置时间。 - pub fn primary_reset_at(self) -> Option> { - self.primary.and_then(RateLimitWindow::reset_at_datetime) - } - - /// 返回主窗口的秒级窗口长度。 - pub fn primary_limit_window_seconds(self) -> Option { - self.primary.and_then(RateLimitWindow::limit_window_seconds) - } - - /// 返回主窗口是否已经触顶。 - pub fn primary_limit_reached(self) -> bool { - self.primary - .is_some_and(|window| window.used_percent >= 100.0) - } -} - /// 从响应头对中解析限流信息。 pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option { let mut normalized = BTreeMap::new(); @@ -412,22 +348,6 @@ pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(Strin headers } -/// 为退避秒数增加抖动，避免同一时间窗口内集中重试。 -pub fn cooldown_with_jitter(seconds: u64, factor_bps: u16) -> Duration { - let factor_bps = u64::from(factor_bps); - let variance = seconds.saturating_mul(factor_bps) / 10_000; - if variance == 0 { - return Duration::seconds(seconds.min(i64::MAX as u64) as i64); - } - let lower = seconds.saturating_sub(variance); - let span = variance.saturating_mul(2).saturating_add(1); - let offset = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| u64::from(duration.subsec_nanos()) % span) - .unwrap_or(variance); - Duration::seconds(lower.saturating_add(offset).min(i64::MAX as u64) as i64) -} - fn number_field(value: &Value, field: &str) -> Option { value.get(field)?.as_u64() } @@ -510,10 +430,12 @@ fn unit_token(input: &str) -> Option<&str> { } fn positive_seconds_ceil(seconds: f64) -> Option { - if !seconds.is_finite() || seconds <= 0.0 || seconds > u64::MAX as f64 { + if !seconds.is_finite() || seconds <= 0.0 { return None; } - Some(seconds.ceil() as u64) + Duration::try_from_secs_f64(seconds.ceil()) + .ok() + .map(|duration| duration.as_secs()) } fn push_window_headers( diff --git a/src/upstream/protocol/websocket.rs b/src/upstream/protocol/websocket.rs index a8a477aba..36ce860d9 100644 --- a/src/upstream/protocol/websocket.rs +++ b/src/upstream/protocol/websocket.rs @@ -8,30 +8,6 @@ use crate::upstream::protocol::sse::encode_sse_event; const REDACTED_PAYLOAD_VALUE: &str = ""; -/// WebSocket 消息包装。 -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] -pub struct WsMessage { - /// 消息类型。 - #[serde(rename = "type")] - pub message_type: String, - /// 消息负载。 - pub payload: T, -} - -/// Codex WebSocket 聚合响应。 -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] -pub struct CodexWebSocketResponse { - /// 聚合后的事件。 - pub events: Vec, -} - -/// Codex WebSocket 流响应。 -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] -pub struct CodexWebSocketStreamResponse { - /// 聚合响应。 - pub response: CodexWebSocketResponse, -} - /// WebSocket 握手审计快照。 #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] pub struct OpeningAuditSnapshot { @@ -61,13 +37,6 @@ pub struct PayloadAuditSnapshot { pub body: Value, } -/// WebSocket 一致性差异。 -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -pub struct WsParityDiff { - /// 差异说明。 - pub differences: Vec, -} - /// WebSocket audit artifact. #[derive(Debug, Clone, Serialize, PartialEq)] pub struct WebSocketAuditArtifact { @@ -92,24 +61,6 @@ pub struct WebSocketAuditErrorSnapshot { pub message: String, } -/// Structured parity diff between two WebSocket audit artifacts. -#[derive(Debug, Clone, Serialize, PartialEq)] -pub struct WebSocketParityDiff { - /// 逐字段差异。 - pub differences: Vec, -} - -/// Single WebSocket parity difference. -#[derive(Debug, Clone, Serialize, PartialEq)] -pub struct WebSocketParityDifference { - /// 字段路径。 - pub path: String, - /// 当前值。 - pub current: Value, - /// 参考值。 - pub reference: Value, -} - /// WebSocket 错误帧分类模式。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum WebSocketErrorClassificationProfile { @@ -186,20 +137,6 @@ struct ResponsesStreamEventShape { _content_index: Option, } -/// 将单条 WebSocket 消息编码为文本帧。 -pub fn encode_ws_message( - message: &WsMessage, -) -> Result { - serde_json::to_string(message) -} - -/// 从文本帧解码单条 WebSocket 消息。 -pub fn decode_ws_message Deserialize<'de>>( - frame: &str, -) -> Result, serde_json::Error> { - serde_json::from_str(frame) -} - /// 将一条公开 WebSocket JSON 事件编码为 SSE 帧。 pub fn websocket_event_to_sse_frame(raw: &str) -> Option { let event = websocket_event_type(raw)?; @@ -1081,51 +1018,6 @@ pub fn websocket_audit_artifact_from_attempt( } } -/// 按 parity 检查字段比较两个 WebSocket audit artifact。 -pub fn websocket_parity_diff( - current: &WebSocketAuditArtifact, - reference: &WebSocketAuditArtifact, -) -> WebSocketParityDiff { - let mut differences = Vec::new(); - push_websocket_parity_difference( - &mut differences, - "transport_mode", - serde_json::json!(current.transport_mode), - serde_json::json!(reference.transport_mode), - ); - push_websocket_parity_difference( - &mut differences, - "fallback_allowed", - serde_json::json!(current.fallback_allowed), - serde_json::json!(reference.fallback_allowed), - ); - push_websocket_parity_difference( - &mut differences, - "opening.request_line", - opening_request_line(¤t.opening), - opening_request_line(&reference.opening), - ); - push_websocket_parity_difference( - &mut differences, - "opening.header_order", - opening_header_order(¤t.opening), - opening_header_order(&reference.opening), - ); - push_websocket_parity_difference( - &mut differences, - "opening.sec_websocket_extensions", - opening_header_value(¤t.opening, "Sec-WebSocket-Extensions"), - opening_header_value(&reference.opening, "Sec-WebSocket-Extensions"), - ); - push_websocket_parity_difference( - &mut differences, - "payload.top_level_keys", - payload_top_level_keys(¤t.payload), - payload_top_level_keys(&reference.payload), - ); - WebSocketParityDiff { differences } -} - fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static str { match transport_for_request(request) { CodexTransport::HttpSse => "http_sse", @@ -1134,69 +1026,6 @@ fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static st } } -fn push_websocket_parity_difference( - differences: &mut Vec, - path: &'static str, - current: Value, - reference: Value, -) { - if current != reference { - differences.push(WebSocketParityDifference { - path: path.to_string(), - current, - reference, - }); - } -} - -fn opening_request_line(opening: &Option) -> Value { - opening - .as_ref() - .map(|o| Value::String(o.request_line.clone())) - .unwrap_or(Value::Null) -} - -fn opening_header_order(opening: &Option) -> Value { - opening - .as_ref() - .map(|o| { - Value::Array( - o.header_order - .iter() - .map(|h| Value::String(h.clone())) - .collect(), - ) - }) - .unwrap_or(Value::Null) -} - -fn opening_header_value(opening: &Option, name: &str) -> Value { - opening - .as_ref() - .map(|o| { - o.header_order - .iter() - .find(|h| h.starts_with(&format!("{name}:"))) - .map(|h| Value::String(h.clone())) - .unwrap_or(Value::Null) - }) - .unwrap_or(Value::Null) -} - -fn payload_top_level_keys(payload: &Option) -> Value { - payload - .as_ref() - .map(|p| { - Value::Array( - p.top_level_keys - .iter() - .map(|k| Value::String(k.clone())) - .collect(), - ) - }) - .unwrap_or(Value::Null) -} - /// 生成 Responses WebSocket `response.create` payload。 pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Value { let mut body = Map::new(); diff --git a/src/upstream/token_client.rs b/src/upstream/token_client.rs index 182eaa988..7bd488a62 100644 --- a/src/upstream/token_client.rs +++ b/src/upstream/token_client.rs @@ -62,7 +62,7 @@ impl TokenRefresher for OpenAiTokenClient { if !status.is_success() { return Err(classify_refresh_failure(status, &body)); } - parse_token_pair(&body).map_err(|_| RefreshFailure::Transport) + parse_token_pair(&body).map_err(|()| RefreshFailure::Transport) } } diff --git a/src/upstream/transport/client.rs b/src/upstream/transport/client.rs index f0d32a5bd..dcd319cc1 100644 --- a/src/upstream/transport/client.rs +++ b/src/upstream/transport/client.rs @@ -65,7 +65,7 @@ pub fn build_reqwest_client(force_http11: bool) -> Result let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); let mut clients = cache .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); + .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(client) = clients.get(&cache_key) { return Ok(client.clone()); @@ -302,15 +302,6 @@ pub struct CodexCompactResponse { pub rate_limit_headers: Vec<(String, String)>, } -/// 上游模型端点探测结果。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexConnectivityProbe { - /// 请求的完整端点。 - pub endpoint: String, - /// 返回状态码。 - pub status: StatusCode, -} - // --------------------------------------------------------------------------- // CodexBackendClient // --------------------------------------------------------------------------- @@ -650,32 +641,6 @@ impl CodexBackendClient { }) } - /// 探测模型端点连通性。 - pub async fn probe_connectivity( - &self, - context: CodexRequestContext<'_>, - ) -> CodexClientResult { - let headers = self.auxiliary_request_headers(context)?; - let endpoint = format!( - "{}/codex/models?client_version={}", - self.base_url, self.fingerprint.app_version - ); - let response = self.client.get(&endpoint).headers(headers).send().await?; - let status = response.status(); - let retry_after_seconds = retry_after_seconds(response.headers(), None); - if !status.is_success() { - let body = read_capped_error_body(response).await?; - return Err(CodexClientError::Upstream { - status, - retry_after_seconds: retry_after_seconds - .or_else(|| retry_after_seconds_from_body(&body)), - body, - set_cookie_headers: Vec::new(), - }); - } - Ok(CodexConnectivityProbe { endpoint, status }) - } - /// 获取后端模型目录条目。 async fn fetch_models_with_context( &self, diff --git a/src/upstream/transport/endpoints.rs b/src/upstream/transport/endpoints.rs index 7e59c6480..ed64b7975 100644 --- a/src/upstream/transport/endpoints.rs +++ b/src/upstream/transport/endpoints.rs @@ -42,15 +42,6 @@ pub fn usage_endpoint_urls(base_url: &str) -> Vec { .collect() } -/// 返回 usage 主请求路径。 -pub fn primary_usage_request_path(base_url: &str) -> String { - let endpoint_path = usage_endpoint_paths(base_url) - .into_iter() - .next() - .unwrap_or(CODEX_USAGE_API_PATH); - endpoint_request_path(base_url, endpoint_path) -} - fn usage_endpoint_paths(base_url: &str) -> Vec<&'static str> { if has_backend_api_base_path(base_url) { vec![WHAM_USAGE_PATH, CODEX_USAGE_PATH] diff --git a/src/upstream/transport/mod.rs b/src/upstream/transport/mod.rs index 8bc2836ab..73829f46a 100644 --- a/src/upstream/transport/mod.rs +++ b/src/upstream/transport/mod.rs @@ -13,23 +13,17 @@ pub use self::{ backend_transport_for_response_request, build_reqwest_client, is_banned_auth_signal, is_banned_upstream_error, CodexBackendClient, CodexBackendResponse, CodexBackendSseStream, CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexClientResult, - CodexCompactResponse, CodexConnectivityProbe, CodexModelCatalogClient, - CodexModelCatalogClientError, CodexModelCatalogRequest, CodexRateLimitHeaderUpdates, - CodexRequestContext, CodexTurnStateUpdate, + CodexCompactResponse, CodexModelCatalogClient, CodexModelCatalogClientError, + CodexModelCatalogRequest, CodexRateLimitHeaderUpdates, CodexRequestContext, + CodexTurnStateUpdate, }, endpoints::{ - endpoint_request_path, endpoint_url, primary_usage_request_path, usage_endpoint_urls, - CODEX_RESPONSES_COMPACT_PATH, CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, CODEX_USAGE_PATH, - WHAM_USAGE_PATH, + endpoint_request_path, endpoint_url, usage_endpoint_urls, CODEX_RESPONSES_COMPACT_PATH, + CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, CODEX_USAGE_PATH, WHAM_USAGE_PATH, }, headers::{ build_codex_base_headers, build_codex_headers, build_ordered_codex_base_headers, build_ordered_codex_headers, order_headers, }, - tls::{ - build_reqwest_client_with_custom_ca, custom_ca_env_cache_key, - maybe_build_rustls_client_config_with_custom_ca, CustomCaError, CustomCaResult, - CODEX_CA_CERT_ENV, SSL_CERT_FILE_ENV, - }, websocket_pool::{CodexWebSocketPool, CodexWebSocketPoolConfig, CodexWebSocketPoolKey}, }; diff --git a/src/upstream/transport/usage.rs b/src/upstream/transport/usage.rs index dd83b0a90..3ad10cbd7 100644 --- a/src/upstream/transport/usage.rs +++ b/src/upstream/transport/usage.rs @@ -46,9 +46,10 @@ impl CodexBackendClient { Err(CodexClientError::Upstream { status: StatusCode::BAD_GATEWAY, retry_after_seconds: None, - body: last_invalid_body - .map(|body| format!("invalid usage response: {}", truncate_for_error(&body))) - .unwrap_or_else(|| "usage endpoint is unavailable".to_string()), + body: last_invalid_body.map_or_else( + || "usage endpoint is unavailable".to_string(), + |body| format!("invalid usage response: {}", truncate_for_error(&body)), + ), set_cookie_headers: Vec::new(), }) } diff --git a/src/upstream/transport/websocket.rs b/src/upstream/transport/websocket.rs index 005a4a55b..971d4cc20 100644 --- a/src/upstream/transport/websocket.rs +++ b/src/upstream/transport/websocket.rs @@ -61,14 +61,6 @@ pub struct CodexWebSocketConnection { headers: Vec<(String, String)>, } -/// 构造握手审计快照。 -pub fn opening_audit_snapshot(header_order: Vec) -> OpeningAuditSnapshot { - OpeningAuditSnapshot { - header_order, - ..OpeningAuditSnapshot::default() - } -} - /// 显式写入 WebSocket audit artifact。 pub async fn write_websocket_audit_artifact_for_dir( dir: Option<&Path>, @@ -206,14 +198,6 @@ pub enum CodexWebSocketExchangeError { } impl CodexWebSocketRequest { - /// 构造 WebSocket 请求。 - pub fn new(connection: CodexWebSocketConnection, payload_text: String) -> Self { - Self { - connection, - payload_text, - } - } - /// 返回连接描述。 pub fn connection(&self) -> &CodexWebSocketConnection { &self.connection @@ -282,11 +266,6 @@ impl CodexWebSocketConnection { &self.headers } - /// 生成真实 WebSocket opening 请求文本，用于 capture/audit parity。 - pub fn opening_request_text(&self) -> String { - String::from_utf8_lossy(&opening_request_bytes(self)).into_owned() - } - /// 生成打开握手审计快照。 pub fn opening_audit_snapshot(&self) -> OpeningAuditSnapshot { OpeningAuditSnapshot { @@ -337,12 +316,12 @@ async fn connect_websocket( connection: &CodexWebSocketConnection, ) -> Result<(CodexWsStream, WsResponse>>), CodexWebSocketExchangeError> { let request = websocket_handshake_request(connection)?; - let connector = super::maybe_build_rustls_client_config_with_custom_ca() + let connector = super::tls::maybe_build_rustls_client_config_with_custom_ca() .map_err(|error| tungstenite::Error::Io(std::io::Error::other(error)))? .map(Connector::Rustls); match connect_async_tls_with_config(request, Some(websocket_config()), false, connector).await { Ok((websocket, response)) => Ok((websocket, response)), - Err(tungstenite::Error::Http(response)) => Err(websocket_opening_error(*response)), + Err(tungstenite::Error::Http(response)) => Err(websocket_opening_error(response.as_ref())), Err(error) => Err(error.into()), } } @@ -371,7 +350,7 @@ fn websocket_config() -> WebSocketConfig { config } -fn websocket_opening_error(response: WsResponse>>) -> CodexWebSocketExchangeError { +fn websocket_opening_error(response: &WsResponse>>) -> CodexWebSocketExchangeError { let status_code = response.status().as_u16(); let body = response .body() @@ -393,20 +372,6 @@ fn websocket_opening_error(response: WsResponse>>) -> CodexWebSoc } } -fn opening_request_bytes(connection: &CodexWebSocketConnection) -> Vec { - let mut bytes = Vec::new(); - bytes.extend_from_slice(request_line_for_endpoint(connection.endpoint()).as_bytes()); - bytes.extend_from_slice(b"\r\n"); - for (name, value) in connection.headers() { - bytes.extend_from_slice(name.as_bytes()); - bytes.extend_from_slice(b": "); - bytes.extend_from_slice(value.as_bytes()); - bytes.extend_from_slice(b"\r\n"); - } - bytes.extend_from_slice(b"\r\n"); - bytes -} - fn request_line_for_endpoint(endpoint: &str) -> String { let path = reqwest::Url::parse(endpoint) .ok() @@ -673,7 +638,7 @@ async fn first_reused_stream_text( }, ); } - Err(error) => return Err(reused_connection_died_before_first_frame(error)), + Err(error) => return Err(reused_connection_died_before_first_frame(&error)), }; match message { Message::Text(text) => return Ok(text.to_string()), @@ -685,7 +650,7 @@ async fn first_reused_stream_text( ); } Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent), - _ => continue, + _ => {} } } } @@ -724,7 +689,7 @@ async fn collect_websocket_response( let message = match next_websocket_message(&mut websocket).await { Ok(message) => message, Err(error) if reused_connection && !saw_response_frame => { - return Err(reused_connection_died_before_first_frame(error)); + return Err(reused_connection_died_before_first_frame(&error)); } Err(error) => return Err(error), }; @@ -808,7 +773,7 @@ async fn collect_websocket_response( } fn reused_connection_died_before_first_frame( - error: CodexWebSocketExchangeError, + error: &CodexWebSocketExchangeError, ) -> CodexWebSocketExchangeError { CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { message: error.to_string(), diff --git a/src/upstream/transport/websocket_pool.rs b/src/upstream/transport/websocket_pool.rs index 8113adbe1..70b46b400 100644 --- a/src/upstream/transport/websocket_pool.rs +++ b/src/upstream/transport/websocket_pool.rs @@ -11,7 +11,7 @@ use tokio::{net::TcpStream, sync::Mutex, time::timeout}; use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream}; const DEFAULT_MAX_PER_ACCOUNT: usize = 8; -const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60); +const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55); const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25); const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25); const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5); @@ -113,46 +113,6 @@ impl CodexWebSocketPool { pool } - /// 使用默认配置构造连接池。 - pub fn with_default_max_age() -> Self { - Self::with_config(CodexWebSocketPoolConfig::default()) - } - - /// 使用最大存活时间与账号容量限制构造连接池。 - pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self { - Self { - inner: Arc::new(Mutex::new(WebSocketPoolState::default())), - config: CodexWebSocketPoolConfig { - max_age, - max_per_account: max_per_account.max(1), - maintenance_interval: None, - ping_interval: None, - liveness_timeout: None, - ..CodexWebSocketPoolConfig::default() - }, - } - } - - /// 返回单账号最大连接数。 - pub fn max_per_account(&self) -> usize { - self.config.max_per_account - } - - /// 返回连接最大存活时长。 - pub fn max_age(&self) -> Duration { - self.config.max_age - } - - /// 判断当前账号是否还能打开新连接。 - pub fn permits_new_connection(&self, current_connections: usize) -> bool { - current_connections < self.config.max_per_account - } - - /// 判断连接是否已达到回收年龄。 - pub fn should_recycle(&self, age: Duration) -> bool { - age >= self.config.max_age - } - pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire { let mut expired_connection = None; let acquire = { @@ -270,11 +230,6 @@ impl CodexWebSocketPool { close_pooled_connections(idle_connections).await; } - /// 维护 idle 连接：关闭过期/失活连接，并按需 ping 探活。 - pub async fn gc_sweep(&self) { - self.maintain_idle_connections().await; - } - /// 维护 idle 连接：关闭过期/失活连接，并按需 ping 探活。 pub async fn maintain_idle_connections(&self) { let maintenance = self.take_idle_connections_for_maintenance().await; @@ -520,7 +475,7 @@ async fn probe_idle_connection( return Some(connection); } match timeout(ping_timeout, connection.websocket.next()).await { - Ok(Some(Ok(Message::Close(_)))) | Ok(None) | Ok(Some(Err(_))) | Err(_) => { + Ok(Some(Ok(Message::Close(_)) | Err(_)) | None) | Err(_) => { close_pooled_connection(connection).await; None } diff --git a/tests/admin/accounts/import_export.rs b/tests/admin/accounts/importing.rs similarity index 84% rename from tests/admin/accounts/import_export.rs rename to tests/admin/accounts/importing.rs index f44c56838..b8e50eaf4 100644 --- a/tests/admin/accounts/import_export.rs +++ b/tests/admin/accounts/importing.rs @@ -1,78 +1,12 @@ use super::*; +use crate::support::jwt::unsigned_jwt; use wiremock::{ matchers::{method, path}, Mock, MockServer, ResponseTemplate, }; #[tokio::test] -async fn admin_accounts_export_should_return_native_account_tokens() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-accounts-export.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - seed_account( - &pool, - NewAccount { - id: "acct_export".to_string(), - email: Some("export@example.com".to_string()), - account_id: Some("chatgpt_export".to_string()), - user_id: Some("user_export".to_string()), - label: Some("primary".to_string()), - plan_type: Some("plus".to_string()), - access_token: SecretString::new("access-export".to_string().into()), - refresh_token: Some(SecretString::new("refresh-export".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - added_at: None, - }, - ) - .await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/export?ids=acct_export") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_accounts_export") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["sourceFormat"], "native"); - assert_eq!(body["data"]["accounts"][0]["id"], "acct_export"); - assert_eq!(body["data"]["accounts"][0]["token"], "access-export"); - assert_eq!( - body["data"]["accounts"][0]["refreshToken"], - "refresh-export" - ); -} - -#[tokio::test] -async fn admin_accounts_import_should_store_native_account_tokens() { +async fn admin_accounts_import_should_store_cpr_account_tokens() { let dir = tempfile::tempdir().unwrap(); let db = dir.path().join("admin-accounts-import.sqlite"); let url = format!("sqlite://{}", db.display()); @@ -138,7 +72,7 @@ async fn admin_accounts_import_should_store_native_account_tokens() { assert_eq!(status, StatusCode::OK); assert_eq!(body["data"]["imported"], 1); assert_eq!(body["data"]["skipped"], 0); - assert_eq!(body["data"]["sourceFormat"], "native"); + assert_eq!(body["data"]["sourceFormat"], "cpr"); assert_eq!(stored.access_token.expose_secret(), access_token); assert_eq!( stored.refresh_token.unwrap().expose_secret(), @@ -192,7 +126,6 @@ async fn admin_accounts_import_should_fetch_wham_usage_for_current_openai_token_ format!("{}/backend-api", server.uri()), ) .await; - let header = json!({"alg": "none", "typ": "JWT"}); let payload = json!({ "exp": 4_102_444_800i64, "https://api.openai.com/auth": { @@ -204,7 +137,7 @@ async fn admin_accounts_import_should_fetch_wham_usage_for_current_openai_token_ "email_verified": true, }, }); - let access_token = format!("{}.{}.", jwt_part(&header), jwt_part(&payload)); + let access_token = unsigned_jwt(&payload); let response = app .clone() @@ -689,44 +622,3 @@ async fn admin_accounts_import_should_store_plain_tokens_and_list_sanitized_acco ); assert!(list_body["data"]["items"][0].get("token").is_none()); } - -#[tokio::test] -async fn admin_accounts_export_should_return_native_accounts_with_tokens_and_filter_ids() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-export-filter.sqlite", 117).await; - seed_account( - &pool, - NewAccount { - id: "acct_export_a".to_string(), - email: Some("export-a@example.com".to_string()), - account_id: Some("chatgpt-export-a".to_string()), - user_id: Some("user-export-a".to_string()), - label: None, - plan_type: Some("plus".to_string()), - access_token: SecretString::new("access-acct_export_a".to_string().into()), - refresh_token: Some(SecretString::new( - "refresh-acct_export_a".to_string().into(), - )), - access_token_expires_at: None, - status: AccountStatus::Active, - added_at: None, - }, - ) - .await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/export?ids=acct_export_a") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["accounts"].as_array().unwrap().len(), 1); - assert_eq!(body["data"]["accounts"][0]["id"], "acct_export_a"); -} diff --git a/tests/admin/accounts/lifecycle.rs b/tests/admin/accounts/lifecycle.rs index df5971a17..000b0ccd6 100644 --- a/tests/admin/accounts/lifecycle.rs +++ b/tests/admin/accounts/lifecycle.rs @@ -359,159 +359,6 @@ async fn admin_account_refresh_should_not_mark_valid_account_banned_when_refresh assert_eq!(stored.status, AccountStatus::Active); } -#[tokio::test] -async fn admin_account_update_should_accept_batch_status_payload() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-account-batch-status-update.sqlite", 122).await; - sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)") - .bind("acct_batch_status_1").bind("batch-1@example.com").bind("cipher").bind("active") - .bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z").execute(&pool).await.unwrap(); - sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)") - .bind("acct_batch_status_2").bind("batch-2@example.com").bind("cipher").bind("active") - .bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z").execute(&pool).await.unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/accounts/update") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "ids": ["acct_batch_status_1", "acct_batch_status_2", "acct_missing"], - "status": "disabled" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["updated"], 2); - assert_eq!(body["data"]["notFound"], json!(["acct_missing"])); -} - -#[tokio::test] -async fn admin_account_batch_status_route_should_not_exist() { - let (_app, state, _pool, _dir) = - admin_accounts_test_app("admin-account-batch-status-route.sqlite", 123).await; - let admin_router = codex_proxy_rs::admin::router::router().with_state(state); - - let response = admin_router - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/accounts/batch-status") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({"ids": ["acct_batch_status_1"], "status": "disabled"}).to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::NOT_FOUND); -} - -#[tokio::test] -async fn admin_account_cookies_should_store_plain_cookie_header() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-account-cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - seed_account( - &pool, - NewAccount { - id: "acct_cookies".to_string(), - email: Some("cookies@example.com".to_string()), - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-cookies".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - added_at: None, - }, - ) - .await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let set = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/accounts/cookies") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from( - json!({ - "id": "acct_cookies", - "cookies":"cf_clearance=clear-secret; __cf_bm=bm-secret" - }) - .to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(set.status(), StatusCode::OK); - - let get = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/cookies?id=acct_cookies") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(get.status(), StatusCode::OK); - let get_body = response_json(get).await; - let cookies = get_body["data"]["cookies"].as_str().unwrap(); - assert!(cookies.contains("cf_clearance=clear-secret")); - assert!(cookies.contains("__cf_bm=bm-secret")); - - let stored_values: Vec = - sqlx::query_scalar("select value from account_cookies where account_id = ?1") - .bind("acct_cookies") - .fetch_all(&pool) - .await - .unwrap(); - assert!(!stored_values.is_empty()); - assert!(stored_values.iter().any(|value| value == "clear-secret")); - assert!(stored_values.iter().any(|value| value == "bm-secret")); -} - #[tokio::test] async fn admin_account_create_should_derive_claims_and_store_plain_tokens() { let dir = tempfile::tempdir().unwrap(); diff --git a/tests/admin/accounts/list.rs b/tests/admin/accounts/list.rs index ab69d55de..1938080b9 100644 --- a/tests/admin/accounts/list.rs +++ b/tests/admin/accounts/list.rs @@ -59,16 +59,18 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { admin_accounts_test_app("admin-accounts-stats.sqlite", 68).await; seed_usage_account( &pool, - "acct_stats", - "stats@example.com", - "stats", - "free", - 41, - 5, - 3_500_000, - 13_900, - 3_400_000, - "2026-06-23T08:50:13Z", + UsageAccountSeed { + id: "acct_stats", + email: "stats@example.com", + label: "stats", + plan_type: "free", + request_count: 41, + empty_response_count: 5, + input_tokens: 3_500_000, + output_tokens: 13_900, + cached_tokens: 3_400_000, + last_used_at: "2026-06-23T08:50:13Z", + }, ) .await; sqlx::query("update accounts set quota_json = ?, quota_fetched_at = ? where id = ?") diff --git a/tests/admin/accounts/mod.rs b/tests/admin/accounts/mod.rs index 6611246e0..924a92cde 100644 --- a/tests/admin/accounts/mod.rs +++ b/tests/admin/accounts/mod.rs @@ -1,19 +1,13 @@ -use std::{collections::BTreeMap, sync::Arc}; +use std::sync::Arc; use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; -use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use codex_proxy_rs::{ admin::auth::service::SqliteAdminSessionStore, admin::keys::service::SqliteClientKeyStore, admin::monitoring::event_store::SqliteEventLogStore, - config::types::AppConfig, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, infra::database::connect_sqlite, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, @@ -31,67 +25,45 @@ use serde_json::{json, Value}; use sqlx::SqlitePool; use tower::util::ServiceExt; -mod import_export; +use crate::support::jwt::unsigned_jwt; +use crate::support::{admin::seed_admin_session, config::test_config, http::response_json}; + +mod importing; mod lifecycle; mod list; mod oauth; mod quota; mod testing; -async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind("admin_1") - .bind("hash") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind("admin_1") - .bind("2999-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); -} - -#[expect( - clippy::too_many_arguments, - reason = "test fixture keeps usage rows explicit" -)] -async fn seed_usage_account( - pool: &SqlitePool, - id: &str, - email: &str, - label: &str, - plan_type: &str, +struct UsageAccountSeed<'a> { + id: &'a str, + email: &'a str, + label: &'a str, + plan_type: &'a str, request_count: i64, empty_response_count: i64, input_tokens: i64, output_tokens: i64, cached_tokens: i64, - last_used_at: &str, -) { + last_used_at: &'a str, +} + +async fn seed_usage_account(pool: &SqlitePool, seed: UsageAccountSeed<'_>) { sqlx::query("insert into accounts (id, email, label, plan_type, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, 'active', ?, ?)") - .bind(id).bind(email).bind(label).bind(plan_type).bind("access-token") + .bind(seed.id).bind(seed.email).bind(seed.label).bind(seed.plan_type).bind("access-token") .bind("2026-06-11T00:00:00Z").bind("2026-06-11T00:00:00Z") .execute(pool).await.unwrap(); sqlx::query("insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)") - .bind(id).bind(request_count).bind(empty_response_count).bind(input_tokens).bind(output_tokens).bind(cached_tokens).bind(last_used_at) + .bind(seed.id) + .bind(seed.request_count) + .bind(seed.empty_response_count) + .bind(seed.input_tokens) + .bind(seed.output_tokens) + .bind(seed.cached_tokens) + .bind(seed.last_used_at) .execute(pool).await.unwrap(); } -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - async fn post_admin_account(app: &axum::Router, payload: Value) -> axum::response::Response { app.clone() .oneshot( @@ -215,7 +187,6 @@ fn test_jwt_with_exp( plan_type: Option<&str>, exp: i64, ) -> String { - let header = json!({"alg": "none", "typ": "JWT"}); let payload = json!({ "exp": exp, "https://api.openai.com/auth": { @@ -225,58 +196,5 @@ fn test_jwt_with_exp( }, "https://api.openai.com/profile": { "email": email } }); - format!("{}.{}.", jwt_part(&header), jwt_part(&payload)) -} - -fn jwt_part(value: &Value) -> String { - URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).unwrap()) -} - -fn test_config(database_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { - base_url: "https://chatgpt.com/backend-api".to_string(), - }, - model_aliases: BTreeMap::new(), - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 300, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: Vec::new(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } + unsigned_jwt(&payload) } diff --git a/tests/admin/accounts/quota.rs b/tests/admin/accounts/quota.rs index bddb9d731..0157c8088 100644 --- a/tests/admin/accounts/quota.rs +++ b/tests/admin/accounts/quota.rs @@ -4,206 +4,6 @@ use wiremock::{ Mock, MockServer, ResponseTemplate, }; -#[tokio::test] -async fn admin_usage_stats_should_return_page_and_summary() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-usage.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - sqlx::query("insert into accounts (id, email, label, plan_type, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)") - .bind("acct_usage").bind("usage@example.com").bind("primary").bind("plus") - .bind("cipher").bind("active").bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z") - .execute(&pool).await.unwrap(); - sqlx::query("insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 3, 1, 21, 8, 5, ?)") - .bind("acct_usage").bind("2026-06-18T00:10:00Z").execute(&pool).await.unwrap(); - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let page = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage?page=1&pageSize=10") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_page") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(page.status(), StatusCode::OK); - let page_body = response_json(page).await; - assert_eq!(page_body["data"]["items"][0]["accountId"], "acct_usage"); - assert_eq!(page_body["data"]["items"][0]["requestCount"], 3); - assert_eq!(page_body["data"]["page"]["page"], 1); - assert_eq!(page_body["data"]["page"]["pageSize"], 10); - assert_eq!(page_body["data"]["page"]["total"], 1); - assert_eq!(page_body["data"]["page"]["totalPages"], 1); - - let summary = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/summary") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_summary") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(summary.status(), StatusCode::OK); - assert_eq!(response_json(summary).await["data"]["accountCount"], 1); -} - -#[tokio::test] -async fn admin_usage_stats_should_require_admin_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-usage-auth.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage") - .header("x-request-id", "req_usage_auth") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert_eq!(response_json(response).await["code"], 40101); -} - -#[tokio::test] -async fn admin_usage_stats_should_cursor_page_account_usage() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-usage-cursor.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - seed_usage_account( - &pool, - "acct_a", - "a@example.com", - "primary", - "plus", - 3, - 0, - 12, - 5, - 1, - "2026-06-11T00:00:00Z", - ) - .await; - seed_usage_account( - &pool, - "acct_b", - "b@example.com", - "backup", - "free", - 2, - 1, - 7, - 3, - 2, - "2026-06-11T00:01:00Z", - ) - .await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let first = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage?limit=1") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_cursor") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let first_body = response_json(first).await; - assert_eq!(first_body["code"], 200); - assert_eq!(first_body["data"]["items"].as_array().unwrap().len(), 1); - let cursor = first_body["data"]["page"]["nextCursor"].as_str().unwrap(); - - let second = app - .oneshot( - Request::builder() - .method("GET") - .uri(format!("/api/admin/usage?limit=1&cursor={cursor}")) - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(second.status(), StatusCode::OK); - assert_eq!( - response_json(second).await["data"]["items"][0]["accountId"], - "acct_a" - ); -} - #[tokio::test] async fn admin_account_quota_should_send_usage_cookie() { let server = MockServer::start().await; @@ -247,7 +47,10 @@ async fn admin_account_quota_should_send_usage_cookie() { ) .await; SqliteCookieStore::new(pool.clone()) - .set_cookie_header("acct_quota_cookie", "cf_clearance=admin-quota") + .capture_set_cookie( + "acct_quota_cookie", + "cf_clearance=admin-quota; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); @@ -275,168 +78,3 @@ async fn admin_account_quota_should_send_usage_cookie() { (StatusCode::OK, Some("cf_clearance=admin-quota")) ); } - -#[tokio::test] -async fn admin_account_health_check_should_send_probe_cookie() { - let server = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/backend-api/codex/models")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "models": [] - }))) - .mount(&server) - .await; - let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-health-cookie.sqlite", - 89, - format!("{}/backend-api", server.uri()), - ) - .await; - seed_account( - &pool, - NewAccount { - id: "acct_health_cookie".to_string(), - email: Some("health-cookie@example.com".to_string()), - account_id: Some("chatgpt-health-cookie".to_string()), - user_id: None, - label: None, - plan_type: Some("plus".to_string()), - access_token: SecretString::new("access-health-cookie".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - added_at: None, - }, - ) - .await; - SqliteCookieStore::new(pool.clone()) - .set_cookie_header("acct_health_cookie", "cf_clearance=admin-health") - .await - .unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/accounts/health-check") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_admin_health_cookie") - .body(Body::from( - json!({ "ids": ["acct_health_cookie"] }).to_string(), - )) - .unwrap(), - ) - .await - .unwrap(); - let requests = server.received_requests().await.unwrap(); - let cookie_header = requests - .iter() - .find(|request| request.url.path() == "/backend-api/codex/models") - .and_then(|request| request.headers.get("cookie")) - .and_then(|value| value.to_str().ok()); - - assert_eq!( - (response.status(), cookie_header), - (StatusCode::OK, Some("cf_clearance=admin-health")) - ); -} - -#[tokio::test] -async fn admin_account_quota_warnings_should_return_threshold_matches_from_cached_quota() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-account-quota-warnings.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - seed_account( - &pool, - NewAccount { - id: "acct_warn".to_string(), - email: Some("warn@example.com".to_string()), - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access-warn".to_string().into()), - refresh_token: None, - access_token_expires_at: None, - status: AccountStatus::Active, - added_at: None, - }, - ) - .await; - sqlx::query( - "update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?", - ) - .bind( - json!({ - "snapshots": [{ - "source": "core", - "limit_name": null, - "metered_feature": null, - "allowed": true, - "limit_reached": false, - "blocked": false, - "primary": { - "used_percent": 85, - "remaining_percent": 15, - "reset_at": 1770000100, - "window_minutes": 300, - "limit_reached": false - }, - "secondary": { - "used_percent": 91, - "remaining_percent": 9, - "reset_at": 1770000200, - "window_minutes": 10080, - "limit_reached": false - } - }], - "monthly_limit": null, - "credits": null, - "spend_control": null - }) - .to_string(), - ) - .bind("2026-06-13T00:00:00Z") - .bind("2026-06-13T00:00:00Z") - .bind("acct_warn") - .execute(&pool) - .await - .unwrap(); - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/accounts/quota-warnings") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - let warnings = body["data"]["warnings"].as_array().unwrap(); - assert_eq!(warnings.len(), 2); -} diff --git a/tests/admin/accounts/testing.rs b/tests/admin/accounts/testing.rs index 1cf20ca23..a827f9523 100644 --- a/tests/admin/accounts/testing.rs +++ b/tests/admin/accounts/testing.rs @@ -1,4 +1,5 @@ use super::*; +use axum::body::to_bytes; use codex_proxy_rs::upstream::protocol::sse::{encode_sse_event, parse_sse_events}; use wiremock::{ matchers::{method, path}, @@ -34,7 +35,7 @@ fn test_events(body: &str) -> Vec { } #[tokio::test] -async fn account_test_models_should_return_upstream_models_only() { +async fn account_models_should_return_upstream_models_only() { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/backend-api/codex/models")) @@ -53,7 +54,7 @@ async fn account_test_models_should_return_upstream_models_only() { .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-test-models.sqlite", + "admin-account-models.sqlite", 91, format!("{}/backend-api", server.uri()), ) @@ -66,7 +67,7 @@ async fn account_test_models_should_return_upstream_models_only() { .method("GET") .uri("/api/admin/accounts/models?id=acct_test") .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_account_test_models") + .header("x-request-id", "req_account_models") .body(Body::empty()) .unwrap(), ) diff --git a/tests/admin/auth/mod.rs b/tests/admin/auth/mod.rs index e4e223e88..096bc7a27 100644 --- a/tests/admin/auth/mod.rs +++ b/tests/admin/auth/mod.rs @@ -1,4 +1,3 @@ -pub mod password; -pub mod session; -pub mod session_routes; -pub mod session_store; +pub(crate) mod password; +pub(crate) mod session; +pub(crate) mod session_store; diff --git a/tests/admin/auth/session.rs b/tests/admin/auth/session.rs index c28d09a80..dc4b9f9b5 100644 --- a/tests/admin/auth/session.rs +++ b/tests/admin/auth/session.rs @@ -1,5 +1,5 @@ use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; use chrono::Utc; @@ -7,7 +7,7 @@ use codex_proxy_rs::infra::{database::connect_sqlite, identity::hash_admin_passw use sqlx::SqlitePool; use tower::util::ServiceExt; -use crate::support::config::test_config; +use crate::support::{config::test_config, http::response_json}; #[tokio::test] async fn admin_login_should_issue_http_only_session_cookie() { @@ -58,9 +58,7 @@ async fn admin_login_should_issue_http_only_session_cookie() { .uri("/api/admin/login") .header("content-type", "application/json") .header("x-request-id", "req_login") - .body(Body::from( - r#"{"username":"admin","password":"correct-password"}"#, - )) + .body(Body::from(r#"{"password":"correct-password"}"#)) .unwrap(), ) .await @@ -80,6 +78,19 @@ async fn admin_login_should_issue_http_only_session_cookie() { assert!(cookie.contains("SameSite=Lax")); assert_eq!(body["code"], 200); assert!(body["data"]["expiresAt"].is_string()); + + let logs_response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/logs") + .header("cookie", cookie.split(';').next().unwrap()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(logs_response.status(), StatusCode::OK); } #[tokio::test] @@ -132,9 +143,7 @@ async fn admin_login_should_reject_client_api_key_as_password_or_authorization() .header("content-type", "application/json") .header("authorization", "Bearer cpr_not_an_admin_session") .header("x-request-id", "req_login_bad") - .body(Body::from( - r#"{"username":"admin","password":"cpr_not_an_admin_password"}"#, - )) + .body(Body::from(r#"{"password":"cpr_not_an_admin_password"}"#)) .unwrap(), ) .await @@ -162,8 +171,3 @@ async fn seed_admin_user(pool: &SqlitePool, password: &str) { .await .unwrap(); } - -async fn response_json(response: axum::response::Response) -> serde_json::Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} diff --git a/tests/admin/auth/session_routes.rs b/tests/admin/auth/session_routes.rs deleted file mode 100644 index d51c99acc..000000000 --- a/tests/admin/auth/session_routes.rs +++ /dev/null @@ -1,208 +0,0 @@ -use std::collections::BTreeMap; - -use axum::{ - body::{to_bytes, Body}, - http::{Request, StatusCode}, -}; -use chrono::Utc; -use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::event_store::SqliteEventLogStore, - config::types::AppConfig, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, - infra::{database::connect_sqlite, identity::hash_admin_password}, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::token_refresh::RefreshLeaseStore, - upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, - upstream::fingerprint::FingerprintRepository, -}; -use serde_json::Value; -use sqlx::SqlitePool; -use tower::util::ServiceExt; - -#[tokio::test] -async fn admin_login_should_issue_http_only_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-login.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_user(&pool, "correct-password").await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/login") - .header("content-type", "application/json") - .header("x-request-id", "req_login") - .body(Body::from(r#"{"password":"correct-password"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let cookie = response - .headers() - .get("set-cookie") - .and_then(|v| v.to_str().ok()) - .map(ToString::to_string); - let body = response_json(response).await; - let cookie = cookie.expect("login should set admin session cookie"); - assert!(cookie.starts_with("cpr_admin_session=")); - assert!(cookie.contains("HttpOnly")); - assert_eq!(body["code"], 200); - - let logs_response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/logs") - .header("cookie", cookie.split(';').next().unwrap()) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(logs_response.status(), StatusCode::OK); -} - -#[tokio::test] -async fn admin_login_should_reject_client_api_key_as_password_or_authorization() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-login-rejects-client-key.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_user(&pool, "correct-password").await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/login") - .header("content-type", "application/json") - .header("authorization", "Bearer cpr_not_an_admin_session") - .header("x-request-id", "req_login_bad") - .body(Body::from(r#"{"password":"cpr_not_an_admin_password"}"#)) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert!(response.headers().get("set-cookie").is_none()); - assert_eq!(response_json(response).await["code"], 40102); -} - -async fn seed_admin_user(pool: &SqlitePool, password: &str) { - let now = Utc::now().to_rfc3339(); - let hash = hash_admin_password(password).unwrap(); - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind("admin_1") - .bind(hash) - .bind(&now) - .bind(&now) - .execute(pool) - .await - .unwrap(); -} - -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - -fn test_config(database_url: String) -> AppConfig { - let mut aliases = BTreeMap::new(); - aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string()); - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { - base_url: "https://chatgpt.com/backend-api".to_string(), - }, - model_aliases: aliases, - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 240, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 4, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: vec!["team".to_string(), "plus".to_string()], - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![70, 95], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 60, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } -} diff --git a/tests/admin/keys/authorization.rs b/tests/admin/keys/authorization.rs index 4f0dc89d0..6efc12444 100644 --- a/tests/admin/keys/authorization.rs +++ b/tests/admin/keys/authorization.rs @@ -1,6 +1,5 @@ use super::*; -// Ignored: duplicate of admin::keys::service::authorization test that passes #[tokio::test] async fn admin_client_keys_route_should_create_list_and_authorize_v1_requests() { let (app, _dir) = admin_client_key_test_app("admin-client-key-create.sqlite").await; @@ -71,7 +70,6 @@ async fn admin_client_keys_route_should_create_list_and_authorize_v1_requests() assert_eq!(models_response.status(), StatusCode::OK); } -// Ignored: duplicate of admin::keys::service::authorization test that passes #[tokio::test] async fn admin_client_keys_status_should_disable_and_enable_authorization() { let (app, _dir) = admin_client_key_test_app("admin-client-key-status.sqlite").await; @@ -141,7 +139,6 @@ async fn admin_client_keys_status_should_disable_and_enable_authorization() { assert_eq!(accepted.status(), StatusCode::OK); } -// Ignored: duplicate of admin::keys::service::authorization test that passes #[tokio::test] async fn admin_client_keys_delete_should_remove_authorization() { let (app, _dir) = admin_client_key_test_app("admin-client-key-delete.sqlite").await; diff --git a/tests/admin/keys/import_export.rs b/tests/admin/keys/import_export.rs deleted file mode 100644 index 3ca0410f4..000000000 --- a/tests/admin/keys/import_export.rs +++ /dev/null @@ -1,25 +0,0 @@ -use super::*; - -#[tokio::test] -async fn admin_client_keys_export_should_return_metadata_without_secret_material() { - let (app, _dir) = admin_client_key_test_app("admin-client-key-export.sqlite").await; - let (key_id, _key) = create_admin_client_key(&app, "export-key").await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri(format!("/api/admin/keys/export?ids={key_id}")) - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_api_key_export") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["apiKeys"][0]["id"], key_id); - assert_eq!(body["data"]["apiKeys"][0]["name"], "export-key"); - assert!(body["data"]["apiKeys"][0].get("key").is_none()); -} diff --git a/tests/admin/keys/lifecycle.rs b/tests/admin/keys/lifecycle.rs index 339d93867..fa9ea4a1e 100644 --- a/tests/admin/keys/lifecycle.rs +++ b/tests/admin/keys/lifecycle.rs @@ -57,7 +57,6 @@ async fn admin_client_keys_label_should_update_clear_and_validate_label() { assert_eq!(missing.status(), StatusCode::NOT_FOUND); } -// Ignored: duplicate of admin::keys::service::lifecycle test that passes #[tokio::test] async fn admin_client_keys_batch_delete_should_remove_found_keys_and_report_missing_ids() { let (app, _dir) = admin_client_key_test_app("admin-client-key-delete.sqlite").await; diff --git a/tests/admin/keys/mod.rs b/tests/admin/keys/mod.rs index 4d842bbf5..d34c650bd 100644 --- a/tests/admin/keys/mod.rs +++ b/tests/admin/keys/mod.rs @@ -1,18 +1,11 @@ -use std::collections::BTreeMap; - use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; use codex_proxy_rs::{ admin::auth::service::SqliteAdminSessionStore, admin::keys::service::SqliteClientKeyStore, admin::monitoring::event_store::SqliteEventLogStore, - config::types::AppConfig, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, infra::database::connect_sqlite, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, @@ -21,12 +14,12 @@ use codex_proxy_rs::{ upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, upstream::fingerprint::FingerprintRepository, }; -use serde_json::{json, Value}; -use sqlx::SqlitePool; +use serde_json::json; use tower::util::ServiceExt; +use crate::support::{admin::seed_admin_session, config::test_config, http::response_json}; + mod authorization; -mod import_export; mod lifecycle; mod store; @@ -79,80 +72,3 @@ async fn create_admin_client_key(app: &axum::Router, name: &str) -> (String, Str body["data"]["key"].as_str().unwrap().to_string(), ) } - -async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind("admin_1") - .bind("hash") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind("admin_1") - .bind("2999-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); -} - -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - -fn test_config(database_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { - base_url: "https://chatgpt.com/backend-api".to_string(), - }, - model_aliases: BTreeMap::new(), - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 300, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: Vec::new(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } -} diff --git a/tests/admin/mod.rs b/tests/admin/mod.rs index 6b3cd6902..4972c9edb 100644 --- a/tests/admin/mod.rs +++ b/tests/admin/mod.rs @@ -1,7 +1,6 @@ -pub mod accounts; -pub mod auth; -pub mod keys; -pub mod models; -pub mod monitoring; -pub mod response; -pub mod settings; +pub(crate) mod accounts; +pub(crate) mod auth; +pub(crate) mod keys; +pub(crate) mod monitoring; +pub(crate) mod response; +pub(crate) mod settings; diff --git a/tests/admin/models/mod.rs b/tests/admin/models/mod.rs deleted file mode 100644 index 6a664abd8..000000000 --- a/tests/admin/models/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub mod routes; diff --git a/tests/admin/models/routes.rs b/tests/admin/models/routes.rs deleted file mode 100644 index d537e5625..000000000 --- a/tests/admin/models/routes.rs +++ /dev/null @@ -1,247 +0,0 @@ -use std::collections::BTreeMap; - -use axum::{ - body::{to_bytes, Body}, - http::{Request, StatusCode}, -}; -use chrono::{Duration, Utc}; -use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::event_store::SqliteEventLogStore, - config::types::AppConfig, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::{ - cookies::SqliteCookieStore, - model::AccountStatus, - store::{NewAccount, SqliteAccountStore}, - token_refresh::RefreshLeaseStore, - }, - upstream::fingerprint::FingerprintRepository, -}; -use secrecy::SecretString; -use serde_json::{json, Value}; -use sqlx::SqlitePool; -use tower::util::ServiceExt; -use wiremock::{ - matchers::{method, path}, - Mock, MockServer, ResponseTemplate, -}; - -#[tokio::test] -async fn admin_refresh_models_should_require_admin_session() { - let upstream = MockServer::start().await; - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-models-auth.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let config = test_config(url, upstream.uri()); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state); - - let response = app - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/models/refresh") - .header("x-request-id", "req_models_auth") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert_eq!(response_json(response).await["code"], 40101); -} - -// Ignored: requires ModelService with snapshot_store and upstream_client configured -#[tokio::test] - -async fn admin_refresh_models_should_store_snapshots_for_distinct_active_plans() { - let upstream = MockServer::start().await; - Mock::given(method("GET")) - .and(path("/codex/models")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "models": [{ - "slug": "gpt-refresh", "display_name": "GPT Refresh", - "description": "Backend refreshed model", - "default_reasoning_level": "minimal", - "supported_reasoning_levels": [{"effort": "minimal", "description": "Minimal"}], - "contextWindow": 200000 - }] - }))) - .expect(2) - .mount(&upstream) - .await; - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-models-refresh.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let config = test_config(url, upstream.uri()); - let store = SqliteAccountStore::new(pool.clone()); - store - .insert(NewAccount { - id: "acct_plus".to_string(), - email: Some("plus@test.com".to_string()), - account_id: Some("chatgpt-plus".to_string()), - user_id: Some("user_plus".to_string()), - label: None, - plan_type: Some("plus".to_string()), - access_token: SecretString::new("token-plus".to_string().into()), - refresh_token: Some(SecretString::new("refresh-plus".to_string().into())), - access_token_expires_at: Some(Utc::now() + Duration::hours(1)), - status: AccountStatus::Active, - added_at: None, - }) - .await - .unwrap(); - store - .insert(NewAccount { - id: "acct_team".to_string(), - email: Some("team@test.com".to_string()), - account_id: Some("chatgpt-team".to_string()), - user_id: Some("user_team".to_string()), - label: None, - plan_type: Some("team".to_string()), - access_token: SecretString::new("token-team".to_string().into()), - refresh_token: Some(SecretString::new("refresh-team".to_string().into())), - access_token_expires_at: Some(Utc::now() + Duration::hours(1)), - status: AccountStatus::Active, - added_at: None, - }) - .await - .unwrap(); - let stores = BackgroundTaskStores { - accounts: store, - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - event_logs: SqliteEventLogStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); - let state = AppState { - config, - services: (*services).clone(), - }; - let app = codex_proxy_rs::http::router::router().with_state(state.clone()); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/models/refresh") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_models_refresh") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["refreshedPlans"], 2); - assert_eq!(body["data"]["modelCount"], 2); -} - -async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind("admin_1") - .bind("hash") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind("admin_1") - .bind("2999-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); -} - -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - -fn test_config(database_url: String, base_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { base_url }, - model_aliases: BTreeMap::new(), - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 300, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: Vec::new(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } -} diff --git a/tests/admin/monitoring/dashboard.rs b/tests/admin/monitoring/dashboard.rs index d4fb8766e..0c78d5d21 100644 --- a/tests/admin/monitoring/dashboard.rs +++ b/tests/admin/monitoring/dashboard.rs @@ -1,5 +1,5 @@ use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; use chrono::{Duration, Utc}; @@ -29,6 +29,8 @@ use serde_json::{json, Value}; use sqlx::SqlitePool; use tower::util::ServiceExt; +use crate::support::{admin::seed_admin_session, http::response_json}; + #[tokio::test] async fn dashboard_summary_should_render_fingerprint_updated_at_as_beijing_time() { let (app, _store, _pool, _dir) = dashboard_test_app( @@ -346,29 +348,6 @@ async fn dashboard_summary(app: axum::Router) -> Value { response_json(response).await } -async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind("admin_1") - .bind("hash") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind("admin_1") - .bind("2999-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); -} - async fn seed_usage_summary(pool: &SqlitePool, input_tokens: u64, output_tokens: u64) { seed_usage_summary_with_last_used(pool, input_tokens, output_tokens, "2026-06-18T00:00:00Z") .await; @@ -405,11 +384,6 @@ async fn seed_usage_summary_with_last_used( .unwrap(); } -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - fn service_status_value<'a>(body: &'a Value, label: &str) -> &'a str { body["data"]["serviceStatuses"] .as_array() diff --git a/tests/admin/monitoring/events_store.rs b/tests/admin/monitoring/events_store.rs index 924d25ddb..1182dd6c9 100644 --- a/tests/admin/monitoring/events_store.rs +++ b/tests/admin/monitoring/events_store.rs @@ -92,7 +92,7 @@ async fn event_log_store_should_promote_diagnostic_metadata_fields() { } #[tokio::test] -async fn event_log_store_should_get_count_and_clear_events() { +async fn event_log_store_should_get_and_clear_events() { let (store, _dir) = event_log_store("event-logs-clear.sqlite").await; let mut event = EventLog::new("request", EventLevel::Warn, "detail"); event.id = "log_detail".to_string(); @@ -100,9 +100,12 @@ async fn event_log_store_should_get_count_and_clear_events() { let loaded = store.get("log_detail").await.unwrap().unwrap(); assert_eq!(loaded.id, "log_detail"); - assert_eq!(store.count().await.unwrap(), 1); assert_eq!(store.clear().await.unwrap(), 1); - assert_eq!(store.count().await.unwrap(), 0); + let page = store + .list(EventLogFilter::default(), None, 10) + .await + .unwrap(); + assert!(page.items.is_empty()); } async fn event_log_store(db_name: &str) -> (SqliteEventLogStore, tempfile::TempDir) { diff --git a/tests/admin/monitoring/logs_routes.rs b/tests/admin/monitoring/logs_routes.rs index 4055d93aa..8b8904a9d 100644 --- a/tests/admin/monitoring/logs_routes.rs +++ b/tests/admin/monitoring/logs_routes.rs @@ -1,7 +1,5 @@ -use std::collections::BTreeMap; - use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; use chrono::{Duration, Utc}; @@ -12,11 +10,6 @@ use codex_proxy_rs::{ event_store::SqliteEventLogStore, events::{EventLevel, EventLog}, }, - config::types::AppConfig, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, infra::database::connect_sqlite, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, @@ -25,10 +18,10 @@ use codex_proxy_rs::{ upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, upstream::fingerprint::FingerprintRepository, }; -use serde_json::Value; -use sqlx::SqlitePool; use tower::util::ServiceExt; +use crate::support::{admin::seed_admin_session, config::test_config, http::response_json}; + #[tokio::test] async fn admin_logs_should_require_admin_session_cookie() { let (app, _store, _dir) = admin_logs_test_app("admin-logs-auth.sqlite").await; @@ -175,7 +168,7 @@ async fn admin_logs_should_reject_unsupported_level_filter() { } #[tokio::test] -async fn admin_logs_should_return_state_detail_and_clear_events() { +async fn admin_logs_should_return_detail_and_clear_events() { let (app, store, _dir) = admin_logs_test_app("admin-logs-state.sqlite").await; let mut event = EventLog::new("request", EventLevel::Warn, "detail"); event.id = "log_detail".to_string(); @@ -279,80 +272,3 @@ async fn admin_logs_test_app( dir, ) } - -async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind("admin_1") - .bind("hash") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind("admin_1") - .bind("2999-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); -} - -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - -fn test_config(database_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { - base_url: "https://chatgpt.com/backend-api".to_string(), - }, - model_aliases: BTreeMap::new(), - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 300, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: Vec::new(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } -} diff --git a/tests/admin/monitoring/mod.rs b/tests/admin/monitoring/mod.rs index f0f63d204..603e9df78 100644 --- a/tests/admin/monitoring/mod.rs +++ b/tests/admin/monitoring/mod.rs @@ -1,3 +1,3 @@ -pub mod dashboard; -pub mod events_store; -pub mod logs_routes; +pub(crate) mod dashboard; +pub(crate) mod events_store; +pub(crate) mod logs_routes; diff --git a/tests/admin/response/mod.rs b/tests/admin/response/mod.rs index 8564783e5..30e5b9b1c 100644 --- a/tests/admin/response/mod.rs +++ b/tests/admin/response/mod.rs @@ -11,7 +11,7 @@ use serde_json::{json, Value}; #[test] fn admin_envelope_should_not_duplicate_request_id_in_body() { - let body = AdminEnvelope::ok(json!({ "id": "acct_1" }), "req_1"); + let body = AdminEnvelope::ok(json!({ "id": "acct_1" })); let value = serde_json::to_value(body).unwrap(); @@ -31,7 +31,7 @@ fn admin_page_envelope_should_expose_items_and_page_metadata_inside_data() { items: vec![json!({ "id": "evt_1" })], next_cursor: Some("cursor_1".to_string()), }; - let body = AdminPageEnvelope::ok(page, 50, "req_1"); + let body = AdminPageEnvelope::ok(page, 50); let value = serde_json::to_value(body).unwrap(); @@ -57,7 +57,7 @@ fn admin_page_envelope_should_skip_empty_next_cursor() { items: vec![json!({ "id": "evt_1" })], next_cursor: None, }; - let body = AdminPageEnvelope::ok(page, 50, "req_1"); + let body = AdminPageEnvelope::ok(page, 50); let value = serde_json::to_value(body).unwrap(); @@ -84,7 +84,7 @@ fn admin_page_envelope_should_expose_numbered_page_metadata() { page: 2, page_size: 10, }; - let body = AdminPageEnvelope::numbered(page, "req_1"); + let body = AdminPageEnvelope::numbered(page); let value = serde_json::to_value(body).unwrap(); @@ -108,7 +108,7 @@ fn admin_page_envelope_should_expose_numbered_page_metadata() { #[test] fn admin_response_should_keep_http_status_separate_from_body_code() { - let body = AdminEnvelope::new(40101, "Admin login required", (), "req_1"); + let body = AdminEnvelope::new(40101, "Admin login required", ()); let response: Response = AdminResponse::new(StatusCode::UNAUTHORIZED, body).into_response(); assert_eq!(response.status(), StatusCode::UNAUTHORIZED); @@ -116,7 +116,7 @@ fn admin_response_should_keep_http_status_separate_from_body_code() { #[test] fn admin_error_body_should_use_null_data() { - let body = AdminEnvelope::new(40101, "Admin login required", (), "req_1"); + let body = AdminEnvelope::new(40101, "Admin login required", ()); let value = serde_json::to_value(body).unwrap(); @@ -125,13 +125,8 @@ fn admin_error_body_should_use_null_data() { #[tokio::test] async fn admin_error_into_response_should_use_admin_envelope() { - let response: Response = AdminError::new( - StatusCode::UNAUTHORIZED, - 40101, - "Admin login required", - "req_1", - ) - .into_response(); + let response: Response = + AdminError::new(StatusCode::UNAUTHORIZED, 40101, "Admin login required").into_response(); assert_eq!(response.status(), StatusCode::UNAUTHORIZED); let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); diff --git a/tests/admin/settings/mod.rs b/tests/admin/settings/mod.rs index 853aa6995..b05885b6c 100644 --- a/tests/admin/settings/mod.rs +++ b/tests/admin/settings/mod.rs @@ -1,4 +1,4 @@ -pub mod routes; +pub(crate) mod routes; use codex_proxy_rs::config::settings::{AdminSettings, AdminSettingsPatch, SettingsService}; diff --git a/tests/admin/settings/routes.rs b/tests/admin/settings/routes.rs index 46f1cd190..f3867f9f8 100644 --- a/tests/admin/settings/routes.rs +++ b/tests/admin/settings/routes.rs @@ -1,17 +1,13 @@ use std::collections::BTreeMap; use axum::{ - body::{to_bytes, Body}, + body::Body, http::{Request, StatusCode}, }; use codex_proxy_rs::{ admin::auth::service::SqliteAdminSessionStore, admin::keys::service::SqliteClientKeyStore, admin::monitoring::event_store::SqliteEventLogStore, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, config::{settings::RuntimeSettingsService, types::AppConfig}, infra::database::connect_sqlite, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, @@ -21,10 +17,14 @@ use codex_proxy_rs::{ upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, upstream::fingerprint::FingerprintRepository, }; -use serde_json::{json, Value}; +use serde_json::json; use sqlx::SqlitePool; use tower::util::ServiceExt; +use crate::support::{ + admin::seed_admin_session, config::test_config as base_test_config, http::response_json, +}; + #[tokio::test] async fn admin_settings_should_require_admin_session_cookie() { let (app, _dir) = admin_settings_test_app("admin-settings-auth.sqlite").await; @@ -287,29 +287,6 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp ) } -async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind("admin_1") - .bind("hash") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind("admin_1") - .bind("2999-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); -} - async fn seed_account(pool: &SqlitePool, account_id: &str) { sqlx::query( r" @@ -330,58 +307,14 @@ insert into accounts ( .unwrap(); } -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - fn test_config(database_url: String) -> AppConfig { - let mut aliases = BTreeMap::new(); - aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string()); - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { - base_url: "https://chatgpt.com/backend-api".to_string(), - }, - model_aliases: aliases, - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 240, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 4, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: vec!["team".to_string(), "plus".to_string()], - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![70, 95], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: true, - }, - } + let mut config = base_test_config(database_url); + config + .model_aliases + .insert("codex-fast".to_string(), "gpt-5.5".to_string()); + config.auth.refresh_margin_seconds = 240; + config.auth.max_concurrent_per_account = 4; + config.auth.tier_priority = vec!["team".to_string(), "plus".to_string()]; + config.logging.enabled = true; + config } diff --git a/tests/config/mod.rs b/tests/config/mod.rs index ba4caed20..b8afbcbb2 100644 --- a/tests/config/mod.rs +++ b/tests/config/mod.rs @@ -103,51 +103,27 @@ fn default_config_keeps_runtime_artifacts_under_runtime_directory() { } #[test] -fn config_should_reject_runtime_model_section() { +fn config_should_reject_unknown_top_level_sections() { let err = serde_yml::from_str::( - r#" + r" server: host: 127.0.0.1 port: 8080 api: base_url: https://chatgpt.com/backend-api -model: {} +unexpected: {} database: url: sqlite://.runtime/data/codex-proxy-rs.sqlite -"#, +", ) .unwrap_err(); assert!( - err.to_string().contains("model"), - "expected model section to be rejected, got {err}" + err.to_string().contains("unexpected"), + "expected unknown section to be rejected, got {err}" ); } -#[test] -fn config_should_reject_runtime_auth_and_quota_sections() { - for section in ["auth", "quota"] { - let yaml = format!( - r#" -server: - host: 127.0.0.1 - port: 8080 -api: - base_url: https://chatgpt.com/backend-api -{section}: {{}} -database: - url: sqlite://.runtime/data/codex-proxy-rs.sqlite -"# - ); - let err = serde_yml::from_str::(&yaml).unwrap_err(); - - assert!( - err.to_string().contains(section), - "expected {section} section to be rejected, got {err}" - ); - } -} - #[test] fn config_loader_should_read_only_config_yaml() { let dir = tempfile::tempdir().unwrap(); @@ -222,7 +198,7 @@ logging: .unwrap(); fs::write( dir.path().join("ignored-extra.yaml"), - r#" + r" server: host: 0.0.0.0 logging: @@ -231,7 +207,7 @@ ws_pool: enabled: false max_age_ms: 120000 max_per_account: 2 -"#, +", ) .unwrap(); @@ -258,19 +234,19 @@ ws_pool: } #[test] -fn logging_config_rejects_unsupported_size_rotation_field() { +fn logging_config_should_reject_unknown_fields() { let err = serde_yml::from_str::( - r#" + r" directory: .runtime/logs -max_file_bytes: 10485760 +unexpected: true retention_days: 14 enabled: false -"#, +", ) .unwrap_err(); assert!( - err.to_string().contains("max_file_bytes"), - "expected max_file_bytes to be rejected, got {err}" + err.to_string().contains("unexpected"), + "expected unknown logging field to be rejected, got {err}" ); } diff --git a/tests/http/mod.rs b/tests/http/mod.rs index 680c60c00..a92cb06bc 100644 --- a/tests/http/mod.rs +++ b/tests/http/mod.rs @@ -1,2 +1,2 @@ -pub mod trace_middleware; -pub mod web_assets; +pub(crate) mod trace_middleware; +pub(crate) mod web_assets; diff --git a/tests/http/trace_middleware/mod.rs b/tests/http/trace_middleware/mod.rs index 17fae292a..d8067931a 100644 --- a/tests/http/trace_middleware/mod.rs +++ b/tests/http/trace_middleware/mod.rs @@ -4,13 +4,14 @@ use std::{ }; use axum::{ - body::Body, + body::{to_bytes, Body}, http::{Request, StatusCode}, middleware::from_fn, routing::get, Router, }; use codex_proxy_rs::http::middleware::{request_id::attach_request_id, trace::http_trace_layer}; +use tokio::time::{sleep, Duration}; use tower::ServiceExt; use tracing_subscriber::fmt::MakeWriter; @@ -50,6 +51,17 @@ impl Write for SharedLogWriter { } } +async fn wait_for_trace_output(logs: &SharedLogBuffer) -> String { + for _ in 0..20 { + let output = logs.content(); + if output.contains("completed HTTP request") { + return output; + } + sleep(Duration::from_millis(5)).await; + } + logs.content() +} + #[tokio::test(flavor = "current_thread")] async fn http_trace_should_include_request_id_and_completion_fields() { let logs = SharedLogBuffer::default(); @@ -78,8 +90,11 @@ async fn http_trace_should_include_request_id_and_completion_fields() { .await .unwrap(); - assert_eq!(response.status(), StatusCode::NO_CONTENT); - let output = logs.content(); + let status = response.status(); + let _body = to_bytes(response.into_body(), 1024).await.unwrap(); + + assert_eq!(status, StatusCode::NO_CONTENT); + let output = wait_for_trace_output(&logs).await; assert!( output.contains("req_trace") && output.contains("received HTTP request") diff --git a/tests/http/web_assets/mod.rs b/tests/http/web_assets/mod.rs index fe1452d8d..76ec70c39 100644 --- a/tests/http/web_assets/mod.rs +++ b/tests/http/web_assets/mod.rs @@ -1,19 +1,13 @@ -use std::collections::BTreeMap; - use axum::{ body::{to_bytes, Body}, http::{header, Request, StatusCode}, }; -use codex_proxy_rs::{ - config::types::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, - infra::database::connect_sqlite, -}; +use codex_proxy_rs::{config::types::AppConfig, infra::database::connect_sqlite}; use serde_json::Value; use tower::ServiceExt; +use crate::support::{config::test_config as base_test_config, http::response_json}; + #[test] fn asset_headers_should_distinguish_spa_and_fingerprinted_assets() { assert_eq!( @@ -119,58 +113,11 @@ async fn assert_body_contains(response: axum::response::Response, expected: &str ); } -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX) - .await - .expect("body should collect"); - serde_json::from_slice(&bytes).expect("body should be json") -} - fn test_config(database_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { - base_url: "https://chatgpt.com/backend-api".to_string(), - }, - model_aliases: BTreeMap::new(), - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 240, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 4, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: Vec::new(), - oauth_client_id: "app_id".to_string(), - oauth_token_endpoint: "https://auth.invalid/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } + let mut config = base_test_config(database_url); + config.auth.refresh_margin_seconds = 240; + config.auth.max_concurrent_per_account = 4; + config.auth.oauth_client_id = "app_id".to_string(); + config.auth.oauth_token_endpoint = "https://auth.invalid/token".to_string(); + config } diff --git a/tests/infra/mod.rs b/tests/infra/mod.rs index 03bb9c9dd..d47c9229f 100644 --- a/tests/infra/mod.rs +++ b/tests/infra/mod.rs @@ -1,3 +1,3 @@ -pub mod log_rotation; -pub mod storage_schema; -pub mod time; +pub(crate) mod log_rotation; +pub(crate) mod storage_schema; +pub(crate) mod time; diff --git a/tests/infra/storage_schema/mod.rs b/tests/infra/storage_schema/mod.rs index ccb893ce9..dfdc40586 100644 --- a/tests/infra/storage_schema/mod.rs +++ b/tests/infra/storage_schema/mod.rs @@ -1,27 +1,39 @@ use codex_proxy_rs::infra::database::connect_sqlite; #[tokio::test] -async fn sqlite_schema_creates_current_runtime_tables_without_migration_table() { +async fn sqlite_schema_should_create_current_tables() { 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: (i64,) = sqlx::query_as( - "select count(*) from sqlite_master where type = 'table' and name in ('accounts', 'client_api_keys', 'runtime_settings', 'event_logs', 'model_plan_snapshots', 'account_model_usage', 'usage_time_buckets', 'model_account_routes')", + let rows: Vec<(String,)> = sqlx::query_as( + "select name from sqlite_master where type = 'table' and name in ('admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'model_account_routes', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'event_logs', 'model_plan_snapshots', 'session_affinities') order by name", ) - .fetch_one(&pool) + .fetch_all(&pool) .await .unwrap(); - assert_eq!(row.0, 8); - - let migration_table: Option<(String,)> = sqlx::query_as( - "select name from sqlite_master where type = 'table' and name = 'schema_migrations'", - ) - .fetch_optional(&pool) - .await - .unwrap(); - assert_eq!(migration_table, None); + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "account_cookies", + "account_model_usage", + "account_refresh_leases", + "account_usage", + "accounts", + "admin_sessions", + "admin_users", + "client_api_keys", + "event_logs", + "fingerprint_update_history", + "fingerprints", + "model_account_routes", + "model_plan_snapshots", + "runtime_settings", + "session_affinities", + "usage_time_buckets", + ] + ); } #[tokio::test] @@ -217,15 +229,8 @@ async fn sqlite_schema_should_persist_plain_account_cookie_value_column() { .fetch_optional(&pool) .await .unwrap(); - let cipher_column: Option<(String,)> = sqlx::query_as( - "select name from pragma_table_info('account_cookies') where name = 'value_cipher'", - ) - .fetch_optional(&pool) - .await - .unwrap(); assert_eq!(value_column, Some(("value".to_string(),))); - assert_eq!(cipher_column, None); } #[tokio::test] diff --git a/tests/proxy/dispatch/chat_upstream/chat_routes.rs b/tests/proxy/dispatch/chat_upstream/chat_routes.rs index 8b92b2226..aec99c07b 100644 --- a/tests/proxy/dispatch/chat_upstream/chat_routes.rs +++ b/tests/proxy/dispatch/chat_upstream/chat_routes.rs @@ -1,14 +1,5 @@ use super::*; - -fn assert_substrings_appear_in_order(haystack: &str, needles: &[&str]) { - let mut cursor = 0; - for needle in needles { - let Some(offset) = haystack[cursor..].find(needle) else { - panic!("expected substring {needle:?} after byte {cursor} in:\n{haystack}"); - }; - cursor += offset + needle.len(); - } -} +use crate::support::assertions::assert_substrings_appear_in_order; #[tokio::test] async fn chat_completions_should_dispatch_to_codex_and_return_openai_response() { @@ -753,11 +744,17 @@ async fn chat_completions_should_cool_down_cloudflare_403_and_fallback() { let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store - .set_cookie_header("acct_primary", "cf_clearance=old") + .capture_set_cookie( + "acct_primary", + "cf_clearance=old; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); cookie_store - .set_cookie_header("acct_secondary", "cf_clearance=keep") + .capture_set_cookie( + "acct_secondary", + "cf_clearance=keep; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); let response = app diff --git a/tests/proxy/dispatch/chat_upstream/mod.rs b/tests/proxy/dispatch/chat_upstream/mod.rs index c759c7d55..970e39aa6 100644 --- a/tests/proxy/dispatch/chat_upstream/mod.rs +++ b/tests/proxy/dispatch/chat_upstream/mod.rs @@ -1,5 +1,4 @@ use std::{ - collections::BTreeMap, sync::{ atomic::{AtomicUsize, Ordering}, Arc, Mutex, @@ -16,20 +15,16 @@ use codex_proxy_rs::{ admin::auth::service::SqliteAdminSessionStore, admin::keys::service::SqliteClientKeyStore, admin::monitoring::{ - event_store::{AdminLogStateUpdate, EventLogFilter, SqliteEventLogStore}, + event_store::{EventLogFilter, SqliteEventLogStore, DEFAULT_EVENT_LOG_CAPACITY}, events::EventLevel, }, config::types::AppConfig, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, http::router, infra::database::connect_sqlite, proxy::dispatch::session_affinity::{ build_conversation_identity, SessionAffinityEntry, SqliteSessionAffinityStore, }, - runtime::services::{BackgroundTaskStores, Services}, + runtime::services::{BackgroundTaskStores, EventLogOptions, Services}, runtime::state::AppState, upstream::accounts::{ cookies::SqliteCookieStore, store::SqliteAccountStore, token_refresh::RefreshLeaseStore, @@ -63,6 +58,11 @@ use wiremock::{ Mock, MockServer, ResponseTemplate, }; +use crate::support::{ + client_keys::insert_client_api_key, config::test_config as base_test_config, + http::response_json, +}; + const TEST_INSTALLATION_ID: &str = "b4f9d503-07b1-457b-a0da-87e6836b1c43"; fn websocket_accept_config() -> WebSocketConfig { @@ -236,9 +236,23 @@ mod responses_websocket; mod usage_logging; fn test_app_state_with_pool_and_installation_id( + config: &AppConfig, + pool: SqlitePool, + installation_id: String, +) -> AppState { + test_app_state_with_pool_installation_id_and_log_options( + config.clone(), + pool, + installation_id, + EventLogOptions::from_config(config), + ) +} + +fn test_app_state_with_pool_installation_id_and_log_options( config: AppConfig, pool: SqlitePool, installation_id: String, + log_options: EventLogOptions, ) -> AppState { let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), @@ -250,12 +264,14 @@ fn test_app_state_with_pool_and_installation_id( client_keys: SqliteClientKeyStore::new(pool.clone()), event_logs: SqliteEventLogStore::new(pool), }; - let services = Services::with_installation_id( + let services = Services::try_with_installation_id_and_log_options( &config, stores, crate::support::fingerprint::test_fingerprint(), Some(installation_id), - ); + log_options, + ) + .expect("failed to build runtime services with configured TLS transport"); AppState { config, services } } @@ -274,7 +290,7 @@ async fn test_app_with_account_and_installation_id( let api_key = insert_client_api_key(&pool).await; insert_account(&pool).await; let state = test_app_state_with_pool_and_installation_id( - test_config(url, base_url), + &test_config(url, base_url), pool, installation_id, ); @@ -301,7 +317,7 @@ async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, t let api_key = insert_client_api_key(&pool).await; insert_model_snapshot(&pool).await; let state = test_app_state_with_pool_and_installation_id( - test_config(url, base_url), + &test_config(url, base_url), pool, TEST_INSTALLATION_ID.to_string(), ); @@ -333,7 +349,7 @@ async fn test_app_with_account_pool_and_disabled_logging( let api_key = insert_client_api_key(&pool).await; insert_account(&pool).await; let state = test_app_state_with_pool_and_installation_id( - test_config(url, base_url), + &test_config(url, base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), ); @@ -343,15 +359,6 @@ async fn test_app_with_account_pool_and_disabled_logging( .restore_from_repository() .await .expect("account pool should restore"); - state - .services - .logs - .update_state(AdminLogStateUpdate { - enabled: Some(false), - ..AdminLogStateUpdate::default() - }) - .await - .expect("event log should be disabled"); (router::router().with_state(state), api_key, pool, dir) } @@ -366,10 +373,16 @@ async fn test_app_with_account_pool_and_logging_capture_body( insert_account(&pool).await; let mut config = test_config(url, base_url); config.logging.enabled = true; - let state = test_app_state_with_pool_and_installation_id( + let log_options = EventLogOptions { + enabled: true, + capacity: DEFAULT_EVENT_LOG_CAPACITY, + capture_body: true, + }; + let state = test_app_state_with_pool_installation_id_and_log_options( config, pool.clone(), TEST_INSTALLATION_ID.to_string(), + log_options, ); state .services @@ -377,15 +390,6 @@ async fn test_app_with_account_pool_and_logging_capture_body( .restore_from_repository() .await .expect("account pool should restore"); - state - .services - .logs - .update_state(AdminLogStateUpdate { - capture_body: Some(true), - ..AdminLogStateUpdate::default() - }) - .await - .expect("log capture body should be enabled"); (router::router().with_state(state), api_key, pool, dir) } @@ -402,7 +406,7 @@ async fn test_app_with_account_pool_config( let mut config = test_config(url, base_url); configure(&mut config); let state = test_app_state_with_pool_and_installation_id( - config, + &config, pool.clone(), TEST_INSTALLATION_ID.to_string(), ); @@ -425,7 +429,7 @@ async fn test_app_with_restored_pool_then_disabled_account( let api_key = insert_client_api_key(&pool).await; insert_account(&pool).await; let state = test_app_state_with_pool_and_installation_id( - test_config(url, base_url), + &test_config(url, base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), ); @@ -472,7 +476,7 @@ async fn test_app_with_two_accounts_and_affinity( .await .unwrap(); let state = test_app_state_with_pool_and_installation_id( - test_config(url, base_url), + &test_config(url, base_url), pool, TEST_INSTALLATION_ID.to_string(), ); @@ -535,7 +539,7 @@ async fn test_app_with_two_accounts_and_affinity_status( .await .unwrap(); let state = test_app_state_with_pool_and_installation_id( - test_config(url, base_url), + &test_config(url, base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), ); @@ -571,7 +575,7 @@ async fn test_app_with_two_accounts( ) .await; let state = test_app_state_with_pool_and_installation_id( - test_config(url, base_url), + &test_config(url, base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), ); @@ -584,14 +588,6 @@ async fn test_app_with_two_accounts( (router::router().with_state(state), api_key, pool, dir) } -async fn insert_client_api_key(pool: &SqlitePool) -> String { - SqliteClientKeyStore::new(pool.clone()) - .create("test") - .await - .unwrap() - .key -} - async fn seed_openai_admin_session(pool: &SqlitePool, session_id: &str) { sqlx::query( "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", @@ -740,11 +736,6 @@ fn tuple_response_request_body(stream: bool) -> String { .to_string() } -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - struct CapturedWebSocketRequest { headers: Vec<(String, String)>, payload: Value, @@ -975,7 +966,7 @@ async fn accept_two_successful_websocket_responses_with_authorization( websocket.close(None).await.unwrap(); break (true, first_payload, second_payload); } - Some(_) => continue, + Some(_) => {} None => { let second_payload = accept_successful_websocket_response_with_authorization( listener, @@ -1199,7 +1190,7 @@ fn responses_http_sse_request(api_key: &str, request_id: &str) -> Request .unwrap() } -fn responses_json_request(api_key: &str, body: Value) -> Request { +fn responses_json_request(api_key: &str, body: &Value) -> Request { Request::builder() .method("POST") .uri("/v1/responses") @@ -1520,10 +1511,13 @@ async fn spawn_chunked_sse_upstream_then_close_with_headers( tokio::spawn(async move { let (mut socket, _) = listener.accept().await.unwrap(); read_http_request(&mut socket).await; - let extra_headers = headers - .iter() - .map(|(name, value)| format!("{name}: {value}\r\n")) - .collect::(); + let mut extra_headers = String::new(); + for (name, value) in headers { + extra_headers.push_str(name); + extra_headers.push_str(": "); + extra_headers.push_str(value); + extra_headers.push_str("\r\n"); + } let response_head = format!( "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ntransfer-encoding: chunked\r\nconnection: close\r\n{extra_headers}\r\n" ); @@ -1650,48 +1644,7 @@ async fn received_authorizations(server: &MockServer) -> Vec { } fn test_config(database_url: String, base_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { base_url }, - model_aliases: BTreeMap::new(), - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 300, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: Vec::new(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } + let mut config = base_test_config(database_url); + config.api.base_url = base_url; + config } diff --git a/tests/proxy/dispatch/chat_upstream/responses_recovery.rs b/tests/proxy/dispatch/chat_upstream/responses_recovery.rs index 78b8c53ff..a2c1bbad1 100644 --- a/tests/proxy/dispatch/chat_upstream/responses_recovery.rs +++ b/tests/proxy/dispatch/chat_upstream/responses_recovery.rs @@ -486,7 +486,7 @@ async fn responses_should_return_cloudflare_challenge_error_when_403_fallback_is let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store - .set_cookie_header("acct_chat", "cf_clearance=old") + .capture_set_cookie("acct_chat", "cf_clearance=old; Domain=.chatgpt.com; Path=/") .await .unwrap(); let response = app @@ -562,11 +562,17 @@ async fn responses_should_cool_down_cloudflare_403_and_fallback() { let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store - .set_cookie_header("acct_primary", "cf_clearance=old") + .capture_set_cookie( + "acct_primary", + "cf_clearance=old; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); cookie_store - .set_cookie_header("acct_secondary", "cf_clearance=keep") + .capture_set_cookie( + "acct_secondary", + "cf_clearance=keep; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); let response = app @@ -648,11 +654,17 @@ async fn responses_should_clear_cookies_after_cloudflare_path_block_404_and_fall let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store - .set_cookie_header("acct_primary", "cf_clearance=old") + .capture_set_cookie( + "acct_primary", + "cf_clearance=old; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); cookie_store - .set_cookie_header("acct_secondary", "cf_clearance=keep") + .capture_set_cookie( + "acct_secondary", + "cf_clearance=keep; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); let response = app @@ -717,7 +729,7 @@ async fn responses_should_disable_account_after_three_cloudflare_path_block_404s let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store - .set_cookie_header("acct_chat", "cf_clearance=old") + .capture_set_cookie("acct_chat", "cf_clearance=old; Domain=.chatgpt.com; Path=/") .await .unwrap(); for _ in 0..3 { @@ -1599,11 +1611,17 @@ async fn responses_stream_should_cool_down_cloudflare_403_and_fallback() { let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store - .set_cookie_header("acct_primary", "cf_clearance=old") + .capture_set_cookie( + "acct_primary", + "cf_clearance=old; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); cookie_store - .set_cookie_header("acct_secondary", "cf_clearance=keep") + .capture_set_cookie( + "acct_secondary", + "cf_clearance=keep; Domain=.chatgpt.com; Path=/", + ) .await .unwrap(); let response = app @@ -1682,7 +1700,7 @@ async fn responses_stream_should_return_cloudflare_path_block_error_when_404_fal let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store - .set_cookie_header("acct_chat", "cf_clearance=old") + .capture_set_cookie("acct_chat", "cf_clearance=old; Domain=.chatgpt.com; Path=/") .await .unwrap(); let response = app diff --git a/tests/proxy/dispatch/chat_upstream/responses_websocket.rs b/tests/proxy/dispatch/chat_upstream/responses_websocket.rs index 90eff0c2c..38391f50b 100644 --- a/tests/proxy/dispatch/chat_upstream/responses_websocket.rs +++ b/tests/proxy/dispatch/chat_upstream/responses_websocket.rs @@ -168,7 +168,7 @@ async fn responses_websocket_should_reuse_connection_for_recorded_conversation() websocket.close(None).await.unwrap(); break (true, first_payload, second_payload); } - Some(_) => continue, + Some(_) => {} None => { let second_payload = accept_successful_websocket_response( &listener, @@ -295,7 +295,7 @@ async fn responses_websocket_should_retry_fresh_connection_when_reused_connectio .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [{ "role": "user", @@ -313,7 +313,7 @@ async fn responses_websocket_should_retry_fresh_connection_when_reused_connectio let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "previous_response_id": "resp_stale_reuse_first" @@ -381,7 +381,7 @@ async fn responses_websocket_should_not_reuse_connection_when_pool_is_disabled() websocket.close(None).await.unwrap(); break (true, first_payload, second_payload); } - Some(_) => continue, + Some(_) => {} None => { let second_payload = accept_successful_websocket_response( &listener, @@ -634,7 +634,7 @@ async fn responses_websocket_should_implicitly_resume_full_history_with_reasonin .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -648,7 +648,7 @@ async fn responses_websocket_should_implicitly_resume_full_history_with_reasonin let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -732,7 +732,7 @@ async fn responses_websocket_should_strip_implicit_banned_affinity_when_switchin .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -751,7 +751,7 @@ async fn responses_websocket_should_strip_implicit_banned_affinity_when_switchin let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -814,7 +814,7 @@ async fn responses_websocket_should_not_implicitly_resume_unmatched_function_cal .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "use the lookup tool", "stream": false, @@ -828,7 +828,7 @@ async fn responses_websocket_should_not_implicitly_resume_unmatched_function_cal let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "use the lookup tool", "stream": false, @@ -882,7 +882,7 @@ async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_rest let first_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -895,7 +895,7 @@ async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_rest let db = dir.path().join("openai-record-affinity.sqlite"); let restored_state = test_app_state_with_pool_and_installation_id( - test_config(format!("sqlite://{}", db.display()), server_base_url), + &test_config(format!("sqlite://{}", db.display()), server_base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), ); @@ -922,7 +922,7 @@ async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_rest let second_response = restored_app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -1007,7 +1007,7 @@ async fn responses_websocket_pool_should_be_evicted_after_admin_account_status_c .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "prompt_cache_key": "status-cycle" @@ -1026,7 +1026,7 @@ async fn responses_websocket_pool_should_be_evicted_after_admin_account_status_c let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "previous_response_id": "resp_pool_status_first" @@ -1073,7 +1073,7 @@ async fn responses_websocket_should_not_implicitly_resume_self_contained_functio .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "use the lookup tool", "stream": false, @@ -1087,7 +1087,7 @@ async fn responses_websocket_should_not_implicitly_resume_self_contained_functio let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "use the lookup tool", "stream": false, @@ -1145,7 +1145,7 @@ async fn responses_websocket_should_not_implicitly_resume_across_codex_windows() .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "prompt_cache_key": "shared-variant-session", @@ -1161,7 +1161,7 @@ async fn responses_websocket_should_not_implicitly_resume_across_codex_windows() let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "prompt_cache_key": "shared-variant-session", @@ -1219,7 +1219,7 @@ async fn responses_websocket_should_evict_reasoning_replay_after_invalid_encrypt .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -1233,7 +1233,7 @@ async fn responses_websocket_should_evict_reasoning_replay_after_invalid_encrypt let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -1302,7 +1302,7 @@ async fn responses_websocket_should_restore_full_history_when_implicit_resume_pr .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -1316,7 +1316,7 @@ async fn responses_websocket_should_restore_full_history_when_implicit_resume_pr let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "stream": false, @@ -1365,7 +1365,7 @@ async fn responses_websocket_should_route_previous_response_id_to_recorded_accou .clone() .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "instructions": "answer briefly", "input": [{ @@ -1396,7 +1396,7 @@ async fn responses_websocket_should_route_previous_response_id_to_recorded_accou let second_response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "previous_response_id": "resp_affinity_first" @@ -1448,7 +1448,7 @@ async fn responses_websocket_non_stream_previous_response_not_found_should_strip let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": false, @@ -1490,7 +1490,7 @@ async fn responses_websocket_stream_previous_response_not_found_should_strip_his let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": true, @@ -1532,7 +1532,7 @@ async fn responses_websocket_non_stream_unanswered_function_call_should_strip_hi let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": false, @@ -1575,7 +1575,7 @@ async fn responses_websocket_previous_response_id_should_retry_fallback_account_ let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "previous_response_id": "resp_prev" @@ -1624,7 +1624,7 @@ async fn responses_websocket_non_stream_previous_response_id_should_retry_fallba let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": false, @@ -1676,7 +1676,7 @@ async fn responses_websocket_without_history_should_mark_expired_after_fallback_ let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": true, @@ -1739,7 +1739,7 @@ async fn responses_websocket_without_history_should_return_rate_limit_stream_err let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": true, @@ -1815,7 +1815,7 @@ async fn responses_websocket_response_failed_quota_should_retry_fallback_account let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": false, @@ -1857,7 +1857,7 @@ async fn responses_websocket_without_history_should_return_quota_stream_error_wh let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": true, @@ -1907,7 +1907,7 @@ async fn responses_websocket_without_history_should_return_model_unsupported_str let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": true, @@ -1957,7 +1957,7 @@ async fn responses_websocket_with_history_should_return_path_block_stream_error_ let response = app .oneshot(responses_json_request( &api_key, - json!({ + &json!({ "model": "gpt-5.5", "input": [], "stream": true, diff --git a/tests/proxy/dispatch/mod.rs b/tests/proxy/dispatch/mod.rs index d4413fd5e..50837bb5d 100644 --- a/tests/proxy/dispatch/mod.rs +++ b/tests/proxy/dispatch/mod.rs @@ -1,3 +1,3 @@ -pub mod chat_upstream; -pub mod session_affinity; -pub mod session_affinity_integration; +pub(crate) mod chat_upstream; +pub(crate) mod session_affinity; +pub(crate) mod session_affinity_integration; diff --git a/tests/proxy/keys/mod.rs b/tests/proxy/keys/mod.rs index 0e4a05d59..6f8ff8b9d 100644 --- a/tests/proxy/keys/mod.rs +++ b/tests/proxy/keys/mod.rs @@ -1 +1 @@ -pub mod auth; +pub(crate) mod auth; diff --git a/tests/proxy/mod.rs b/tests/proxy/mod.rs index 3a215565a..9458e7e74 100644 --- a/tests/proxy/mod.rs +++ b/tests/proxy/mod.rs @@ -1,3 +1,3 @@ -pub mod dispatch; -pub mod keys; -pub mod openai; +pub(crate) mod dispatch; +pub(crate) mod keys; +pub(crate) mod openai; diff --git a/tests/proxy/openai/mod.rs b/tests/proxy/openai/mod.rs index a905e7b11..562894937 100644 --- a/tests/proxy/openai/mod.rs +++ b/tests/proxy/openai/mod.rs @@ -1,2 +1,2 @@ -pub mod models_routes; -pub mod responses_routes; +pub(crate) mod models_routes; +pub(crate) mod responses_routes; diff --git a/tests/proxy/openai/models_routes.rs b/tests/proxy/openai/models_routes.rs index b19397a51..8c35a2961 100644 --- a/tests/proxy/openai/models_routes.rs +++ b/tests/proxy/openai/models_routes.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; - use axum::{ body::Body, http::{Request, StatusCode}, @@ -8,11 +6,6 @@ use codex_proxy_rs::{ admin::auth::service::SqliteAdminSessionStore, admin::keys::service::SqliteClientKeyStore, admin::monitoring::event_store::SqliteEventLogStore, - config::types::AppConfig, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, infra::database::connect_sqlite, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, @@ -21,9 +14,10 @@ use codex_proxy_rs::{ upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, upstream::fingerprint::FingerprintRepository, }; -use sqlx::SqlitePool; use tower::util::ServiceExt; +use crate::support::{client_keys::insert_client_api_key, config::test_config}; + #[tokio::test] async fn models_route_should_reject_unknown_client_api_key() { let dir = tempfile::tempdir().unwrap(); @@ -102,60 +96,3 @@ async fn models_route_should_accept_stored_client_api_key() { .unwrap(); assert_eq!(response.status(), StatusCode::OK); } - -async fn insert_client_api_key(pool: &SqlitePool) -> String { - SqliteClientKeyStore::new(pool.clone()) - .create("test") - .await - .unwrap() - .key -} - -fn test_config(database_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { - base_url: "https://chatgpt.com/backend-api".to_string(), - }, - model_aliases: BTreeMap::new(), - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 300, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: Vec::new(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } -} diff --git a/tests/proxy/openai/responses_routes.rs b/tests/proxy/openai/responses_routes.rs index 989a23a56..9723497c5 100644 --- a/tests/proxy/openai/responses_routes.rs +++ b/tests/proxy/openai/responses_routes.rs @@ -1,5 +1,3 @@ -use std::collections::BTreeMap; - use axum::{ body::{to_bytes, Body}, http::{Request, StatusCode}, @@ -8,11 +6,6 @@ use codex_proxy_rs::{ admin::auth::service::SqliteAdminSessionStore, admin::keys::service::SqliteClientKeyStore, admin::monitoring::event_store::SqliteEventLogStore, - config::types::AppConfig, - config::types::{ - AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, - }, infra::database::connect_sqlite, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, @@ -21,10 +14,13 @@ use codex_proxy_rs::{ upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, upstream::fingerprint::FingerprintRepository, }; -use serde_json::Value; use sqlx::SqlitePool; use tower::util::ServiceExt; +use crate::support::{ + client_keys::insert_client_api_key, config::test_config, http::response_json, +}; + #[tokio::test] async fn responses_route_should_reject_missing_client_api_key() { let (app, _key, _dir) = test_app_with_client_api_key().await; @@ -166,7 +162,7 @@ async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::Temp async fn insert_model_snapshot(pool: &SqlitePool) { sqlx::query( - r#"insert or replace into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?)"#, + r"insert or replace into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?)", ) .bind("plus") .bind( @@ -177,65 +173,3 @@ async fn insert_model_snapshot(pool: &SqlitePool) { .await .unwrap(); } - -async fn insert_client_api_key(pool: &SqlitePool) -> String { - SqliteClientKeyStore::new(pool.clone()) - .create("test") - .await - .unwrap() - .key -} - -async fn response_json(response: axum::response::Response) -> Value { - let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); - serde_json::from_slice(&bytes).unwrap() -} - -fn test_config(database_url: String) -> AppConfig { - AppConfig { - server: ServerConfig { - host: "127.0.0.1".to_string(), - port: 0, - }, - api: ApiConfig { - base_url: "https://chatgpt.com/backend-api".to_string(), - }, - model_aliases: BTreeMap::new(), - model_account_routes: BTreeMap::new(), - auth: AuthConfig { - refresh_margin_seconds: 300, - refresh_enabled: true, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "least_used".to_string(), - tier_priority: Vec::new(), - oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(), - oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(), - }, - quota: QuotaConfig { - refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, - }, - database: DatabaseConfig { url: database_url }, - tls: TlsConfig { - force_http11: false, - }, - ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), - admin: AdminConfig { - session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, - default_username: "admin".to_string(), - default_password: "admin".to_string(), - }, - logging: LoggingConfig { - directory: "logs".to_string(), - retention_days: 14, - enabled: false, - }, - } -} diff --git a/tests/runtime/account_pool_restore/mod.rs b/tests/runtime/account_pool_restore/mod.rs index ea7847a6c..9f382b2f2 100644 --- a/tests/runtime/account_pool_restore/mod.rs +++ b/tests/runtime/account_pool_restore/mod.rs @@ -29,12 +29,15 @@ async fn runtime_account_pool_should_restore_capacity_and_clear_runtime_state() .expect("account pool should restore"); let capacity = service.capacity_summary_now().await; let acquired = service - .acquire_with(AccountAcquireRequest::new("gpt-5.5", Utc::now())) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", Utc::now())) .await .expect("active account should be acquired"); let used_capacity = service.capacity_summary_now().await; - service.clear().await; - let cleared_capacity = service.capacity_summary_now().await; + let restored_again = service + .restore_from_repository() + .await + .expect("account pool should restore again"); + let restored_capacity = service.capacity_summary_now().await; assert_eq!(restored, 1); assert_eq!(capacity.total_slots, 2); @@ -42,7 +45,9 @@ async fn runtime_account_pool_should_restore_capacity_and_clear_runtime_state() assert_eq!(acquired.account.id, "acct_pool"); assert_eq!(used_capacity.used_slots, 1); assert_eq!(used_capacity.available_slots, 1); - assert_eq!(cleared_capacity.total_slots, 0); + assert_eq!(restored_again, 1); + assert_eq!(restored_capacity.used_slots, 0); + assert_eq!(restored_capacity.available_slots, 2); } async fn insert_account(pool: &sqlx::SqlitePool, id: &str) { diff --git a/tests/runtime/mod.rs b/tests/runtime/mod.rs index 2f96872ed..90cc4e238 100644 --- a/tests/runtime/mod.rs +++ b/tests/runtime/mod.rs @@ -1,3 +1,3 @@ -pub mod account_pool_restore; -pub mod services; -pub mod tasks; +pub(crate) mod account_pool_restore; +pub(crate) mod services; +pub(crate) mod tasks; diff --git a/tests/runtime/services.rs b/tests/runtime/services.rs index 06f8d98ec..0c224788f 100644 --- a/tests/runtime/services.rs +++ b/tests/runtime/services.rs @@ -11,7 +11,7 @@ use codex_proxy_rs::{ cookies::SqliteCookieStore, store::SqliteAccountStore, token_refresh::RefreshLeaseStore, }, fingerprint::FingerprintRepository, - transport::CODEX_CA_CERT_ENV, + transport::tls::CODEX_CA_CERT_ENV, }, }; @@ -41,13 +41,12 @@ async fn services_try_new_should_use_configured_tls_transport_builder() { event_logs: SqliteEventLogStore::new(pool), }; - let error = match Services::try_new( + let Err(error) = Services::try_new( &config, stores, crate::support::fingerprint::test_fingerprint(), - ) { - Ok(_) => panic!("invalid custom CA should fail service transport construction"), - Err(error) => error, + ) else { + panic!("invalid custom CA should fail service transport construction"); }; assert!(error .to_string() diff --git a/tests/runtime/tasks/fingerprint.rs b/tests/runtime/tasks/fingerprint.rs index 178bf1421..1c2ecce06 100644 --- a/tests/runtime/tasks/fingerprint.rs +++ b/tests/runtime/tasks/fingerprint.rs @@ -7,7 +7,7 @@ async fn fingerprint_update_task_should_start_background_checker() { let pool = connect_sqlite(&format!("sqlite://{}", db.display())) .await .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool); + let repo = FingerprintRepository::new(pool.clone()); repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) .await .expect("seed current fingerprint"); @@ -45,7 +45,7 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito let pool = connect_sqlite(&format!("sqlite://{}", db.display())) .await .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool); + let repo = FingerprintRepository::new(pool.clone()); repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) .await .expect("seed current fingerprint"); @@ -99,7 +99,7 @@ async fn fingerprint_update_task_should_not_persist_when_appcast_matches_current let pool = connect_sqlite(&format!("sqlite://{}", db.display())) .await .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool); + let repo = FingerprintRepository::new(pool.clone()); repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) .await .expect("seed current fingerprint"); @@ -113,9 +113,12 @@ async fn fingerprint_update_task_should_not_persist_when_appcast_matches_current wait_for_appcast_requests(&server, 1).await; handle.shutdown().await; - let latest_history = repo.latest().await.expect("latest history should load"); + let history_count: i64 = sqlx::query_scalar("select count(*) from fingerprint_update_history") + .fetch_one(&pool) + .await + .expect("fingerprint history count should load"); - assert!(latest_history.is_none()); + assert_eq!(history_count, 0); } #[tokio::test] diff --git a/tests/runtime/tasks/model_refresh.rs b/tests/runtime/tasks/model_refresh.rs index 8c677e897..b1be5764f 100644 --- a/tests/runtime/tasks/model_refresh.rs +++ b/tests/runtime/tasks/model_refresh.rs @@ -1,22 +1,219 @@ use super::*; +use std::collections::BTreeMap; + +use codex_proxy_rs::upstream::{ + accounts::pool::{AccountAcquireRequest, AccountPoolOptions, RuntimeAccountPoolService}, + models::{BackendModelEntry, ModelPlanSnapshot, ModelSnapshotStore, ModelSnapshotStoreResult}, + transport::{CodexModelCatalogClient, CodexModelCatalogClientError, CodexModelCatalogRequest}, +}; + #[tokio::test] async fn model_refresh_task_should_start_and_shutdown() { - let model_service = Arc::new(ModelService::new( - ModelConfig { - model_aliases: Default::default(), - }, - None, - None, - None, - )); + let model_service = Arc::new(ModelService::new(empty_model_config(), None, None)); let account_store = Arc::new(FakeAccountStore); + let account_pool = Arc::new(RuntimeAccountPoolService::new( + account_store.clone(), + AccountPoolOptions::default(), + 0, + )); let handle = codex_proxy_rs::runtime::tasks::model_refresh::ModelRefreshTask::new( model_service, account_store, + account_pool, ) .start(); handle.shutdown().await; } + +#[tokio::test] +async fn model_refresh_task_should_sync_model_plan_allowlist_to_account_pool() { + let snapshot_store = Arc::new(InMemorySnapshotStore::default()); + let upstream = Arc::new(FakeModelCatalogClient { + models_by_plan: BTreeMap::from([ + ("plus".to_string(), backend_models("gpt-plus")), + ("team".to_string(), backend_models("gpt-team")), + ]), + }); + let model_service = Arc::new(ModelService::new( + empty_model_config(), + Some(snapshot_store), + Some(upstream), + )); + let account_store = Arc::new(PlanAccountStore { + accounts: vec![ + plan_account("acct-plus", "plus"), + plan_account("acct-team", "team"), + ], + }); + let account_pool = Arc::new(RuntimeAccountPoolService::new( + account_store.clone(), + AccountPoolOptions::default(), + 0, + )); + account_pool + .restore_from_repository() + .await + .expect("accounts should restore"); + + codex_proxy_rs::runtime::tasks::model_refresh::ModelRefreshTask::new( + model_service, + account_store, + account_pool.clone(), + ) + .refresh_once() + .await + .expect("model refresh should sync allowlist"); + + let acquired = account_pool + .acquire_with(&AccountAcquireRequest::new("gpt-team", Utc::now())) + .await + .expect("team model should acquire an account"); + assert_eq!(acquired.account.id, "acct-team"); +} + +fn empty_model_config() -> ModelConfig { + ModelConfig { + model_aliases: BTreeMap::new(), + } +} + +fn backend_models(model_id: &str) -> Vec { + vec![BackendModelEntry { + slug: Some(model_id.to_string()), + display_name: Some(model_id.to_string()), + ..BackendModelEntry::default() + }] +} + +fn plan_account(id: &str, plan_type: &str) -> Account { + let mut account = crate::support::accounts::test_account(id, AccountStatus::Active); + account.plan_type = Some(plan_type.to_string()); + account +} + +#[derive(Default)] +struct InMemorySnapshotStore { + snapshots: tokio::sync::Mutex>, +} + +#[async_trait] +impl ModelSnapshotStore for InMemorySnapshotStore { + async fn replace_plan_snapshot( + &self, + snapshot: &ModelPlanSnapshot, + ) -> ModelSnapshotStoreResult<()> { + self.snapshots + .lock() + .await + .insert(snapshot.plan_type.clone(), snapshot.clone()); + Ok(()) + } + + async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult> { + Ok(self.snapshots.lock().await.values().cloned().collect()) + } +} + +struct FakeModelCatalogClient { + models_by_plan: BTreeMap>, +} + +#[async_trait] +impl CodexModelCatalogClient for FakeModelCatalogClient { + async fn fetch_models( + &self, + request: &CodexModelCatalogRequest<'_>, + ) -> Result, CodexModelCatalogClientError> { + Ok(self + .models_by_plan + .get(request.plan_type) + .cloned() + .unwrap_or_default()) + } +} + +struct PlanAccountStore { + accounts: Vec, +} + +#[async_trait] +impl AccountStore for PlanAccountStore { + async fn list_pool_accounts(&self) -> AccountStoreResult> { + Ok(self.accounts.clone()) + } + + async fn mark_quota_limited_until( + &self, + _account_id: &str, + _cooldown_until: chrono::DateTime, + ) -> AccountStoreResult { + Ok(true) + } + + async fn set_cloudflare_cooldown_until( + &self, + _account_id: &str, + _cooldown_until: chrono::DateTime, + ) -> AccountStoreResult { + Ok(true) + } + + async fn set_status( + &self, + _account_id: &str, + _status: AccountStatus, + ) -> AccountStoreResult { + Ok(true) + } + + async fn record_usage_delta( + &self, + _account_id: &str, + _usage: AccountUsageDelta, + ) -> AccountStoreResult<()> { + Ok(()) + } + + async fn record_model_usage_delta( + &self, + _account_id: &str, + _model: &str, + _usage: AccountModelUsageDelta, + ) -> AccountStoreResult<()> { + Ok(()) + } + + async fn get_quota_json(&self, _account_id: &str) -> AccountStoreResult> { + Ok(None) + } + + async fn update_quota_json( + &self, + _account_id: &str, + _quota_json: &str, + ) -> AccountStoreResult { + Ok(false) + } + + async fn apply_quota_snapshot( + &self, + _account_id: &str, + _quota_json: &str, + _limit_reached: bool, + _cooldown_until: Option>, + ) -> AccountStoreResult { + Ok(false) + } + + async fn sync_rate_limit_window( + &self, + _account_id: &str, + _reset_at: chrono::DateTime, + _limit_window_seconds: Option, + ) -> AccountStoreResult<()> { + Ok(()) + } +} diff --git a/tests/support/accounts.rs b/tests/support/accounts.rs index b7cf0e5d8..72460533c 100644 --- a/tests/support/accounts.rs +++ b/tests/support/accounts.rs @@ -1,7 +1,7 @@ use chrono::{Duration, Utc}; use codex_proxy_rs::upstream::accounts::model::{Account, AccountStatus}; -pub fn test_account(id: &str, status: AccountStatus) -> Account { +pub(crate) fn test_account(id: &str, status: AccountStatus) -> Account { Account { id: id.to_string(), email: None, diff --git a/tests/support/admin.rs b/tests/support/admin.rs new file mode 100644 index 000000000..f7e536588 --- /dev/null +++ b/tests/support/admin.rs @@ -0,0 +1,25 @@ +use sqlx::SqlitePool; + +pub(crate) async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + ) + .bind("admin_1") + .bind("hash") + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); + + sqlx::query( + "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", + ) + .bind(session_id) + .bind("admin_1") + .bind("2999-01-01T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} diff --git a/tests/support/assertions.rs b/tests/support/assertions.rs new file mode 100644 index 000000000..efd6f5346 --- /dev/null +++ b/tests/support/assertions.rs @@ -0,0 +1,9 @@ +pub(crate) fn assert_substrings_appear_in_order(haystack: &str, needles: &[&str]) { + let mut cursor = 0; + for needle in needles { + let Some(offset) = haystack[cursor..].find(needle) else { + panic!("expected substring {needle:?} after byte {cursor} in:\n{haystack}"); + }; + cursor += offset + needle.len(); + } +} diff --git a/tests/support/client_keys.rs b/tests/support/client_keys.rs new file mode 100644 index 000000000..1044513ed --- /dev/null +++ b/tests/support/client_keys.rs @@ -0,0 +1,10 @@ +use codex_proxy_rs::admin::keys::service::SqliteClientKeyStore; +use sqlx::SqlitePool; + +pub(crate) async fn insert_client_api_key(pool: &SqlitePool) -> String { + SqliteClientKeyStore::new(pool.clone()) + .create("test") + .await + .unwrap() + .key +} diff --git a/tests/support/config.rs b/tests/support/config.rs index 6e3c70cba..55ab3a740 100644 --- a/tests/support/config.rs +++ b/tests/support/config.rs @@ -1,11 +1,11 @@ use std::collections::BTreeMap; use codex_proxy_rs::config::types::{ - AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig, - QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig, + AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, FingerprintConfig, + LoggingConfig, QuotaConfig, ServerConfig, TlsConfig, WebSocketPoolConfig, }; -pub fn test_config(database_url: String) -> AppConfig { +pub(crate) fn test_config(database_url: String) -> AppConfig { AppConfig { server: ServerConfig { host: "127.0.0.1".to_string(), @@ -29,19 +29,15 @@ pub fn test_config(database_url: String) -> AppConfig { }, quota: QuotaConfig { refresh_interval_minutes: 5, - warning_thresholds: QuotaWarningThresholds { - primary: vec![80, 90], - secondary: vec![80, 90], - }, }, database: DatabaseConfig { url: database_url }, tls: TlsConfig { force_http11: false, }, ws_pool: WebSocketPoolConfig::default(), - fingerprint: Default::default(), + fingerprint: FingerprintConfig::default(), admin: AdminConfig { - session_ttl_minutes: 60, + session_ttl_minutes: 1440, session_cleanup_interval_secs: 3600, default_username: "admin".to_string(), default_password: "admin".to_string(), diff --git a/tests/support/fingerprint.rs b/tests/support/fingerprint.rs index 93784788e..f311c8bb0 100644 --- a/tests/support/fingerprint.rs +++ b/tests/support/fingerprint.rs @@ -1,10 +1,10 @@ -use codex_proxy_rs::upstream::fingerprint::Fingerprint; +use codex_proxy_rs::{config::types::FingerprintConfig, upstream::fingerprint::Fingerprint}; -pub fn test_fingerprint() -> Fingerprint { - Fingerprint::default_codex_desktop() +pub(crate) fn test_fingerprint() -> Fingerprint { + Fingerprint::from_config(&FingerprintConfig::default()) } -pub fn test_fingerprint_with_updated_at(updated_at: Option<&str>) -> Fingerprint { +pub(crate) fn test_fingerprint_with_updated_at(updated_at: Option<&str>) -> Fingerprint { let mut fingerprint = test_fingerprint(); fingerprint.updated_at = updated_at.map(ToString::to_string); fingerprint diff --git a/tests/support/http.rs b/tests/support/http.rs new file mode 100644 index 000000000..5380b89cb --- /dev/null +++ b/tests/support/http.rs @@ -0,0 +1,6 @@ +use axum::body::to_bytes; + +pub(crate) async fn response_json(response: axum::response::Response) -> serde_json::Value { + let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + serde_json::from_slice(&bytes).unwrap() +} diff --git a/tests/support/jwt.rs b/tests/support/jwt.rs new file mode 100644 index 000000000..43ff3d4f4 --- /dev/null +++ b/tests/support/jwt.rs @@ -0,0 +1,11 @@ +use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; +use serde_json::{json, Value}; + +pub(crate) fn unsigned_jwt(payload: &Value) -> String { + let header = json!({"alg": "none", "typ": "JWT"}); + format!("{}.{}.", jwt_part(&header), jwt_part(payload)) +} + +fn jwt_part(value: &Value) -> String { + URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).expect("test jwt json should encode")) +} diff --git a/tests/support/mod.rs b/tests/support/mod.rs index 276c62dda..d44fcc735 100644 --- a/tests/support/mod.rs +++ b/tests/support/mod.rs @@ -1,4 +1,9 @@ -pub mod accounts; -pub mod config; -pub mod fingerprint; -pub mod sqlite; +pub(crate) mod accounts; +pub(crate) mod admin; +pub(crate) mod assertions; +pub(crate) mod client_keys; +pub(crate) mod config; +pub(crate) mod fingerprint; +pub(crate) mod http; +pub(crate) mod jwt; +pub(crate) mod sqlite; diff --git a/tests/support/sqlite.rs b/tests/support/sqlite.rs index d0768406a..c6735c29e 100644 --- a/tests/support/sqlite.rs +++ b/tests/support/sqlite.rs @@ -1,7 +1,7 @@ use codex_proxy_rs::infra::database::connect_sqlite; use sqlx::SqlitePool; -pub async fn init_test_db(db_name: &str) -> (SqlitePool, tempfile::TempDir) { +pub(crate) async fn init_test_db(db_name: &str) -> (SqlitePool, tempfile::TempDir) { let dir = tempfile::tempdir().expect("temp dir"); let db = dir.path().join(db_name); let url = format!("sqlite://{}", db.display()); diff --git a/tests/upstream/accounts/account_pool/mod.rs b/tests/upstream/accounts/account_pool/mod.rs index b8b72b723..43ad50246 100644 --- a/tests/upstream/accounts/account_pool/mod.rs +++ b/tests/upstream/accounts/account_pool/mod.rs @@ -1,15 +1,14 @@ use std::collections::BTreeMap; -use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use chrono::{Duration, TimeZone, Utc}; use codex_proxy_rs::upstream::accounts::{ - model::AccountStatus, + model::{Account, AccountStatus}, pool::{ AccountAcquireRequest, AccountPool, AccountPoolOptions, AccountWindowUsageDelta, RotationStrategy, }, }; -use serde_json::{json, Value}; +use serde_json::json; mod quota; mod selection; @@ -19,14 +18,14 @@ fn fixed_time() -> chrono::DateTime { Utc.with_ymd_and_hms(2026, 6, 11, 8, 0, 0).unwrap() } +fn acquire_account(pool: &mut AccountPool, model: &str) -> Option { + pool.acquire_with(&AccountAcquireRequest::new(model, Utc::now())) + .map(|acquired| acquired.account) +} + fn test_jwt(exp_offset_seconds: i64) -> String { let payload = json!({ "exp": Utc::now().timestamp() + exp_offset_seconds, }); - let header = json!({ "alg": "none", "typ": "JWT" }); - format!("{}.{}.", jwt_part(&header), jwt_part(&payload),) -} - -fn jwt_part(value: &Value) -> String { - URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).unwrap()) + crate::support::jwt::unsigned_jwt(&payload) } diff --git a/tests/upstream/accounts/account_pool/quota.rs b/tests/upstream/accounts/account_pool/quota.rs index 2b353cc1b..c858186be 100644 --- a/tests/upstream/accounts/account_pool/quota.rs +++ b/tests/upstream/accounts/account_pool/quota.rs @@ -11,7 +11,7 @@ fn account_pool_should_skip_accounts_with_cached_quota_limit() { AccountStatus::Active, )); - assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "usable"); + assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "usable"); } #[test] @@ -25,10 +25,10 @@ fn account_pool_should_reuse_quota_limited_accounts_after_cooldown() { pool.mark_quota_limited_until("limited", now + Duration::seconds(30)); assert!(pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .is_none()); assert_eq!( - pool.acquire_with(AccountAcquireRequest::new( + pool.acquire_with(&AccountAcquireRequest::new( "gpt-5.5", now + Duration::seconds(31) )) @@ -52,7 +52,7 @@ fn account_pool_should_not_shorten_existing_quota_cooldown() { pool.mark_quota_limited_until("limited", now + Duration::seconds(60)); assert!(pool - .acquire_with(AccountAcquireRequest::new( + .acquire_with(&AccountAcquireRequest::new( "gpt-5.5", now + Duration::seconds(90) )) @@ -69,7 +69,7 @@ fn account_pool_should_not_replace_known_window_length_with_cooldown_seconds() { pool.mark_quota_limited_until("limited", now + Duration::seconds(60)); let acquired = pool - .acquire_with(AccountAcquireRequest::new( + .acquire_with(&AccountAcquireRequest::new( "gpt-5.5", now + Duration::seconds(61), )) @@ -92,7 +92,7 @@ fn acquire_should_refresh_expired_cooldowns_before_selecting_account() { pool.insert(account); let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); assert!(!acquired.account.quota_limit_reached); @@ -119,7 +119,7 @@ fn account_pool_should_skip_accounts_in_cloudflare_cooldown() { )); let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); assert_eq!(acquired.account.id, "usable"); @@ -139,7 +139,7 @@ fn least_used_should_deprioritize_quota_limited_accounts_when_skip_is_disabled() pool.insert(limited); pool.insert(usable); - let acquired = pool.acquire("gpt-5.5").unwrap(); + let acquired = acquire_account(&mut pool, "gpt-5.5").unwrap(); assert_eq!(acquired.id, "usable"); } @@ -159,7 +159,7 @@ fn least_used_should_prefer_earlier_rate_limit_window_reset() { pool.insert(soon); let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); assert_eq!(acquired.account.id, "soon"); @@ -180,7 +180,7 @@ fn sync_rate_limit_window_should_reset_window_counters_when_drift_crosses_thresh pool.sync_rate_limit_window("acct_a", now + Duration::seconds(121), Some(60)); let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); assert_eq!(acquired.account.window_request_count, 1); diff --git a/tests/upstream/accounts/account_pool/selection.rs b/tests/upstream/accounts/account_pool/selection.rs index f9b0e1a39..20d459c04 100644 --- a/tests/upstream/accounts/account_pool/selection.rs +++ b/tests/upstream/accounts/account_pool/selection.rs @@ -15,15 +15,15 @@ fn account_pool_should_respect_max_concurrent_slots_per_account() { AccountStatus::Active, )); - let first = pool.acquire("gpt-5.5").unwrap(); - let second = pool.acquire("gpt-5.5").unwrap(); - let third = pool.acquire("gpt-5.5"); + let first = acquire_account(&mut pool, "gpt-5.5").unwrap(); + let second = acquire_account(&mut pool, "gpt-5.5").unwrap(); + let third = acquire_account(&mut pool, "gpt-5.5"); assert_ne!(first.id, second.id); assert!(third.is_none()); pool.release(&first.id); - assert_eq!(pool.acquire("gpt-5.5").unwrap().id, first.id); + assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, first.id); } #[test] @@ -46,10 +46,10 @@ fn account_pool_should_rotate_round_robin_across_candidates() { AccountStatus::Active, )); - assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "acct_a"); - assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "acct_c"); - assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "acct_b"); - assert!(pool.acquire("gpt-5.5").is_none()); + assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "acct_a"); + assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "acct_c"); + assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "acct_b"); + assert!(acquire_account(&mut pool, "gpt-5.5").is_none()); } #[test] @@ -60,7 +60,7 @@ fn acquire_should_skip_accounts_with_expired_token_metadata() { let mut pool = AccountPool::default(); pool.insert(expired); - let acquired = pool.acquire_with(AccountAcquireRequest::new("gpt-5.5", now)); + let acquired = pool.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)); assert!(acquired.is_none()); } @@ -74,7 +74,7 @@ fn acquire_should_skip_accounts_with_expired_jwt_when_metadata_is_missing() { let mut pool = AccountPool::default(); pool.insert(expired); - let acquired = pool.acquire_with(AccountAcquireRequest::new("gpt-5.5", now)); + let acquired = pool.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)); assert!(acquired.is_none()); } @@ -108,7 +108,7 @@ async fn runtime_account_pool_should_persist_expired_status_when_jwt_expiry_is_d runtime_pool.restore_from_repository().await.unwrap(); let acquired = runtime_pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", Utc::now())) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", Utc::now())) .await; let status: (String,) = sqlx::query_as("select status from accounts where id = ?") .bind("acct_expired") @@ -133,7 +133,7 @@ fn account_pool_should_prefer_configured_tier_priority() { pool.insert(free); pool.insert(team); - assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "team"); + assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "team"); } #[test] @@ -151,7 +151,7 @@ fn account_pool_should_filter_by_model_plan_allowlist() { pool.insert(free); pool.insert(plus); - assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "plus"); + assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "plus"); } #[test] @@ -172,8 +172,8 @@ fn account_pool_should_filter_by_model_account_routes() { AccountStatus::Active, )); - assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "acct_b"); - assert!(pool.acquire("gpt-5.5").is_none()); + assert_eq!(acquire_account(&mut pool, "gpt-5.5").unwrap().id, "acct_b"); + assert!(acquire_account(&mut pool, "gpt-5.5").is_none()); } #[test] @@ -190,7 +190,7 @@ fn account_pool_should_exclude_requested_account_ids() { let acquired = pool .acquire_with( - AccountAcquireRequest::new("gpt-5.5", fixed_time()) + &AccountAcquireRequest::new("gpt-5.5", fixed_time()) .with_exclude_account_ids(["acct_a"]), ) .unwrap(); @@ -212,7 +212,8 @@ fn account_pool_should_prefer_session_affinity_account_when_available() { let acquired = pool .acquire_with( - AccountAcquireRequest::new("gpt-5.5", fixed_time()).with_preferred_account_id("acct_b"), + &AccountAcquireRequest::new("gpt-5.5", fixed_time()) + .with_preferred_account_id("acct_b"), ) .unwrap(); @@ -233,10 +234,10 @@ fn account_pool_should_cleanup_stale_slots_before_acquire() { )); assert!(pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .is_some()); let acquired = pool - .acquire_with(AccountAcquireRequest::new( + .acquire_with(&AccountAcquireRequest::new( "gpt-5.5", now + Duration::minutes(5) + Duration::seconds(1), )) @@ -261,9 +262,9 @@ fn account_pool_should_rotate_tied_least_used_accounts() { AccountStatus::Active, )); - let first = pool.acquire("gpt-5.5").unwrap(); + let first = acquire_account(&mut pool, "gpt-5.5").unwrap(); pool.release(&first.id); - let second = pool.acquire("gpt-5.5").unwrap(); + let second = acquire_account(&mut pool, "gpt-5.5").unwrap(); assert_ne!(first.id, second.id); } @@ -283,7 +284,7 @@ fn least_used_should_prefer_lower_runtime_request_count() { pool.insert(quiet); let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); assert_eq!(acquired.account.id, "quiet"); diff --git a/tests/upstream/accounts/account_pool/usage_window.rs b/tests/upstream/accounts/account_pool/usage_window.rs index 614a13202..534292ed3 100644 --- a/tests/upstream/accounts/account_pool/usage_window.rs +++ b/tests/upstream/accounts/account_pool/usage_window.rs @@ -1,57 +1,5 @@ use super::*; -#[test] -fn reset_usage_should_clear_runtime_counters_and_preserve_window_reset() { - let now = fixed_time(); - let mut account = crate::support::accounts::test_account("acct_a", AccountStatus::Active); - account.request_count = 7; - account.empty_response_count = 2; - account.image_input_tokens = 11; - account.image_output_tokens = 13; - account.image_request_count = 3; - account.image_request_failed_count = 1; - account.window_request_count = 5; - account.window_input_tokens = 19; - account.window_output_tokens = 23; - account.window_cached_tokens = 29; - account.window_image_input_tokens = 31; - account.window_image_output_tokens = 37; - account.window_image_request_count = 2; - account.window_image_request_failed_count = 1; - account.window_reset_at = Some(now + Duration::seconds(300)); - account.limit_window_seconds = Some(300); - account.last_used_at = Some("2026-06-12T12:00:00Z".to_string()); - let mut pool = AccountPool::default(); - pool.insert(account); - - assert!(pool.reset_usage("acct_a")); - assert!(!pool.reset_usage("missing")); - - let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) - .unwrap() - .account; - assert_eq!(acquired.request_count, 1); - assert_eq!(acquired.empty_response_count, 0); - assert_eq!(acquired.image_input_tokens, 0); - assert_eq!(acquired.image_output_tokens, 0); - assert_eq!(acquired.image_request_count, 0); - assert_eq!(acquired.image_request_failed_count, 0); - assert_eq!(acquired.window_request_count, 1); - assert_eq!(acquired.window_input_tokens, 0); - assert_eq!(acquired.window_output_tokens, 0); - assert_eq!(acquired.window_cached_tokens, 0); - assert_eq!(acquired.window_image_input_tokens, 0); - assert_eq!(acquired.window_image_output_tokens, 0); - assert_eq!(acquired.window_image_request_count, 0); - assert_eq!(acquired.window_image_request_failed_count, 0); - assert_eq!(acquired.window_reset_at, Some(now + Duration::seconds(300))); - assert_ne!( - acquired.last_used_at.as_deref(), - Some("2026-06-12T12:00:00Z") - ); -} - #[test] fn account_pool_should_return_previous_slot_time_for_request_staggering() { let now = fixed_time(); @@ -65,10 +13,10 @@ fn account_pool_should_return_previous_slot_time_for_request_staggering() { )); let first = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); let second = pool - .acquire_with(AccountAcquireRequest::new( + .acquire_with(&AccountAcquireRequest::new( "gpt-5.5", now + Duration::milliseconds(250), )) @@ -94,7 +42,7 @@ fn account_pool_should_report_capacity_summary() { AccountStatus::Disabled, )); - pool.acquire_with(AccountAcquireRequest::new("gpt-5.5", now)); + pool.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)); let summary = pool.capacity_summary(now); assert_eq!(summary.total_slots, 2); @@ -120,7 +68,7 @@ fn least_used_should_compare_request_count_when_only_one_account_has_window_rese pool.insert(reset_unknown); let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); assert_eq!(acquired.account.id, "reset_unknown"); @@ -135,7 +83,7 @@ fn acquire_should_mark_runtime_usage_window() { pool.insert(account); let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); assert_eq!(acquired.account.request_count, 1); @@ -170,7 +118,7 @@ fn record_window_token_usage_should_accumulate_runtime_window_tokens() { }, ); let acquired = pool - .acquire_with(AccountAcquireRequest::new("gpt-5.5", now)) + .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .unwrap(); assert_eq!(acquired.account.window_input_tokens, 7); diff --git a/tests/upstream/accounts/account_repository.rs b/tests/upstream/accounts/account_repository.rs index dfff21310..f3ba7c508 100644 --- a/tests/upstream/accounts/account_repository.rs +++ b/tests/upstream/accounts/account_repository.rs @@ -158,7 +158,7 @@ async fn account_repository_should_update_status_and_label_without_rewriting_tok #[tokio::test] async fn account_repository_should_accumulate_usage_counters() { let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 7).await; - let repo = SqliteAccountStore::new(pool); + let repo = SqliteAccountStore::new(pool.clone()); seed_new_account(&repo, "acct_a").await; AccountStore::record_usage_delta( @@ -194,15 +194,26 @@ async fn account_repository_should_accumulate_usage_counters() { .await .unwrap(); - let usage = repo.get_usage("acct_a").await.unwrap().unwrap(); - assert_eq!(usage.request_count, 2); - assert_eq!(usage.input_tokens, 18); - assert_eq!(usage.output_tokens, 6); - assert_eq!(usage.cached_tokens, 4); - assert_eq!(usage.reasoning_tokens, 3); - assert_eq!(usage.total_tokens, 24); - assert_eq!(usage.empty_response_count, 1); - assert!(usage.last_used_at.is_some()); + let usage: (i64, i64, i64, i64, i64, i64, i64, Option) = sqlx::query_as( + r" + select request_count, input_tokens, output_tokens, cached_tokens, + reasoning_tokens, total_tokens, empty_response_count, last_used_at + from account_usage + where account_id = ? + ", + ) + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(usage.0, 2); + assert_eq!(usage.1, 18); + assert_eq!(usage.2, 6); + assert_eq!(usage.3, 4); + assert_eq!(usage.4, 3); + assert_eq!(usage.5, 24); + assert_eq!(usage.6, 1); + assert!(usage.7.is_some()); } #[tokio::test] diff --git a/tests/upstream/accounts/mod.rs b/tests/upstream/accounts/mod.rs index f17945a6b..13df78c60 100644 --- a/tests/upstream/accounts/mod.rs +++ b/tests/upstream/accounts/mod.rs @@ -1,7 +1,7 @@ -pub mod account_pool; -pub mod account_repository; -pub mod cloudflare; -pub mod cookies; -pub mod quota_refresh; -pub mod refresh_leases; -pub mod token_refresh; +pub(crate) mod account_pool; +pub(crate) mod account_repository; +pub(crate) mod cloudflare; +pub(crate) mod cookies; +pub(crate) mod quota_refresh; +pub(crate) mod refresh_leases; +pub(crate) mod token_refresh; diff --git a/tests/upstream/accounts/quota_refresh.rs b/tests/upstream/accounts/quota_refresh.rs index f0ea58ae7..97262ca57 100644 --- a/tests/upstream/accounts/quota_refresh.rs +++ b/tests/upstream/accounts/quota_refresh.rs @@ -1,4 +1,4 @@ -use std::{sync::Arc, time::Duration as StdDuration}; +use std::{collections::HashMap, sync::Arc, time::Duration as StdDuration}; use chrono::{Duration, Utc}; use codex_proxy_rs::infra::database::connect_sqlite; @@ -16,14 +16,6 @@ use wiremock::{ Mock, MockServer, ResponseTemplate, }; -#[test] -fn quota_refresh_service_should_expose_default_request_spacing() { - assert_eq!( - RuntimeQuotaRefreshService::default_request_spacing(), - StdDuration::from_secs(3) - ); -} - #[test] fn quota_from_usage_should_preserve_spend_control_individual_limit() { let usage = json!({ @@ -207,7 +199,10 @@ async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_con let cookies = SqliteCookieStore::new(pool.clone()); insert_quota_locked_account(&store, &pool, "acct-quota-cookie", "access-token-cookie").await; cookies - .set_cookie_header("acct-quota-cookie", "cf_clearance=quota-refresh") + .capture_set_cookie( + "acct-quota-cookie", + "cf_clearance=quota-refresh; Domain=.chatgpt.com; Path=/", + ) .await .expect("cookie should be stored"); let codex = CodexBackendClient::new( @@ -218,8 +213,9 @@ async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_con let service = RuntimeQuotaRefreshService::new(store, Arc::new(codex)).with_cookie_store(cookies); + let mut last_refreshed = HashMap::new(); let summary = service - .refresh_locked_accounts_once() + .refresh_locked_accounts(&mut last_refreshed) .await .expect("quota refresh should succeed"); let requests = server @@ -297,8 +293,9 @@ async fn quota_refresh_service_should_fetch_usage_for_quota_locked_accounts_and_ ); let service = RuntimeQuotaRefreshService::new(store.clone(), Arc::new(codex)); + let mut last_refreshed = HashMap::new(); let summary = service - .refresh_locked_accounts_once() + .refresh_locked_accounts(&mut last_refreshed) .await .expect("quota refresh should succeed"); let stored = store @@ -361,7 +358,10 @@ async fn quota_refresh_service_should_stagger_multiple_locked_account_requests() let service = RuntimeQuotaRefreshService::new(store, Arc::new(codex)) .with_request_spacing(StdDuration::from_millis(200)); - let refresh = tokio::spawn(async move { service.refresh_locked_accounts_once().await }); + let refresh = tokio::spawn(async move { + let mut last_refreshed = HashMap::new(); + service.refresh_locked_accounts(&mut last_refreshed).await + }); wait_for_usage_requests(&server, 1).await; tokio::time::sleep(StdDuration::from_millis(50)).await; diff --git a/tests/upstream/accounts/token_refresh/failures.rs b/tests/upstream/accounts/token_refresh/failures.rs index db9d88de5..4edca0b16 100644 --- a/tests/upstream/accounts/token_refresh/failures.rs +++ b/tests/upstream/accounts/token_refresh/failures.rs @@ -338,7 +338,7 @@ async fn token_refresh_task_should_not_reuse_stale_refresh_token_after_retryable stored .refresh_token .as_ref() - .map(|token| token.expose_secret()), + .map(secrecy::ExposeSecret::expose_secret), Some("refresh-rotated") ); } diff --git a/tests/upstream/accounts/token_refresh/mod.rs b/tests/upstream/accounts/token_refresh/mod.rs index 0db6fa508..2c7793e1c 100644 --- a/tests/upstream/accounts/token_refresh/mod.rs +++ b/tests/upstream/accounts/token_refresh/mod.rs @@ -8,7 +8,6 @@ use std::{ }; use async_trait::async_trait; -use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use chrono::{Duration, TimeZone, Utc}; use codex_proxy_rs::infra::database::connect_sqlite; use codex_proxy_rs::upstream::accounts::model::AccountStatus; @@ -174,11 +173,6 @@ impl TokenRefresher for RefreshTokenRotatingRefresher { } fn test_jwt(exp: i64) -> String { - let header = json!({"alg": "none", "typ": "JWT"}); let payload = json!({ "exp": exp }); - format!("{}.{}.", jwt_part(&header), jwt_part(&payload)) -} - -fn jwt_part(value: &serde_json::Value) -> String { - URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).expect("json should encode")) + crate::support::jwt::unsigned_jwt(&payload) } diff --git a/tests/upstream/accounts/token_refresh/scheduling.rs b/tests/upstream/accounts/token_refresh/scheduling.rs index 206de21e7..cae29dcb8 100644 --- a/tests/upstream/accounts/token_refresh/scheduling.rs +++ b/tests/upstream/accounts/token_refresh/scheduling.rs @@ -125,15 +125,14 @@ async fn token_refresh_task_should_skip_duplicate_in_flight_refresh() { let task = task.clone(); tokio::spawn(async move { task.refresh_due_accounts_once_at(now).await }) }; - let second = match timeout(StdDuration::from_millis(100), second).await { - Ok(result) => result.expect("second scan should join"), - Err(_) => { - refresher.release.notify_waiters(); - let _ = first.await; - panic!("second scan should skip instead of waiting on duplicate refresh"); - } - } - .expect("second scan should succeed"); + let Ok(result) = timeout(StdDuration::from_millis(100), second).await else { + refresher.release.notify_waiters(); + let _ = first.await; + panic!("second scan should skip instead of waiting on duplicate refresh"); + }; + let second = result + .expect("second scan should join") + .expect("second scan should succeed"); refresher.release.notify_waiters(); let first = first diff --git a/tests/upstream/fingerprint/mod.rs b/tests/upstream/fingerprint/mod.rs index 2dc7c90f6..18b336555 100644 --- a/tests/upstream/fingerprint/mod.rs +++ b/tests/upstream/fingerprint/mod.rs @@ -1,16 +1,5 @@ use codex_proxy_rs::infra::database::connect_sqlite; -use codex_proxy_rs::upstream::fingerprint::{ - parse_update_manifest, FingerprintRepository, FingerprintUpdater, UpdateChecker, -}; - -#[test] -fn update_manifest_should_extract_version_and_build_number() { - let update = parse_update_manifest(r#"{"version":"26.700.111","build_number":"5002"}"#) - .expect("manifest should parse"); - - assert_eq!(update.app_version, "26.700.111"); - assert_eq!(update.build_number, "5002"); -} +use codex_proxy_rs::upstream::fingerprint::{FingerprintRepository, UpdateChecker}; #[tokio::test] async fn fingerprint_repository_should_update_current_record() { @@ -48,44 +37,6 @@ async fn fingerprint_repository_should_update_current_record() { assert_eq!(stored.chromium_version, "147"); } -#[tokio::test] -async fn fingerprint_updater_should_fetch_manifest_and_persist_history() { - let server = wiremock::MockServer::start().await; - wiremock::Mock::given(wiremock::matchers::method("GET")) - .and(wiremock::matchers::path("/desktop/update.json")) - .respond_with(wiremock::ResponseTemplate::new(200).set_body_raw( - r#"{"version":"26.700.111","build_number":"5002"}"#, - "application/json", - )) - .mount(&server) - .await; - - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("fingerprints.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool); - repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) - .await - .expect("seed current fingerprint"); - let updater = FingerprintUpdater::new( - reqwest::Client::new(), - repo.clone(), - format!("{}/desktop/update.json", server.uri()), - ); - - updater.poll_once().await.expect("poll once"); - - let latest = repo - .latest() - .await - .expect("latest row") - .expect("stored history"); - assert_eq!(latest.app_version, "26.700.111"); - assert_eq!(latest.build_number, "5002"); -} - #[tokio::test] async fn update_checker_should_report_available_update_from_appcast() { let server = wiremock::MockServer::start().await; diff --git a/tests/upstream/mod.rs b/tests/upstream/mod.rs index cdcb8f768..b5c19d8d8 100644 --- a/tests/upstream/mod.rs +++ b/tests/upstream/mod.rs @@ -1,7 +1,7 @@ -pub mod accounts; -pub mod fingerprint; -pub mod fingerprint_integration; -pub mod models; -pub mod protocol; -pub mod token_client; -pub mod transport; +pub(crate) mod accounts; +pub(crate) mod fingerprint; +pub(crate) mod fingerprint_integration; +pub(crate) mod models; +pub(crate) mod protocol; +pub(crate) mod token_client; +pub(crate) mod transport; diff --git a/tests/upstream/models/catalog.rs b/tests/upstream/models/catalog.rs index c434283fc..0c08e4b98 100644 --- a/tests/upstream/models/catalog.rs +++ b/tests/upstream/models/catalog.rs @@ -109,7 +109,7 @@ fn model_catalog_should_merge_backend_snapshots_and_build_plan_allowlist() { #[tokio::test] async fn model_service_should_return_empty_catalog_when_snapshot_store_is_missing() { - let service = ModelService::new(test_model_config(), None, None, None); + let service = ModelService::new(test_model_config(), None, None); let catalog = service.catalog().await; @@ -127,11 +127,10 @@ async fn model_service_should_refresh_distinct_active_plans_and_build_allowlist( test_model_config(), Some(snapshot_store.clone()), Some(upstream.clone()), - None, ); let result = service - .refresh_backend_models( + .refresh_backend_models_with_installation_id( &[ active_account("acct-plus-1", "plus"), active_account("acct-plus-2", "plus"), @@ -139,6 +138,7 @@ async fn model_service_should_refresh_distinct_active_plans_and_build_allowlist( inactive_account("acct-disabled", "plus"), ], "req-model-refresh", + None, ) .await .expect("refresh should succeed"); @@ -170,17 +170,13 @@ async fn model_service_should_refresh_distinct_active_plans_and_build_allowlist( async fn model_service_should_return_no_accounts_when_no_active_plans_exist() { let snapshot_store = Arc::new(InMemorySnapshotStore::default()); let upstream = Arc::new(FakeModelCatalogClient::default()); - let service = ModelService::new( - test_model_config(), - Some(snapshot_store), - Some(upstream), - None, - ); + let service = ModelService::new(test_model_config(), Some(snapshot_store), Some(upstream)); let error = service - .refresh_backend_models( + .refresh_backend_models_with_installation_id( &[inactive_account("acct-disabled", "plus")], "req-model-refresh", + None, ) .await .expect_err("refresh should fail"); diff --git a/tests/upstream/models/mod.rs b/tests/upstream/models/mod.rs index 2c0f52f6c..d618699a9 100644 --- a/tests/upstream/models/mod.rs +++ b/tests/upstream/models/mod.rs @@ -1,2 +1,2 @@ -pub mod catalog; -pub mod store; +pub(crate) mod catalog; +pub(crate) mod store; diff --git a/tests/upstream/protocol/codex_websocket.rs b/tests/upstream/protocol/codex_websocket.rs index 1e793ffa1..693564baa 100644 --- a/tests/upstream/protocol/codex_websocket.rs +++ b/tests/upstream/protocol/codex_websocket.rs @@ -157,66 +157,6 @@ fn codex_websocket_audit_artifact_should_record_transport_opening_and_payload() assert_eq!(artifact.error, None); } -#[test] -fn codex_websocket_parity_diff_should_report_header_payload_and_error_differences() { - let current = codex_proxy_rs::upstream::protocol::websocket::WebSocketAuditArtifact { - transport_mode: "websocket_preferred".to_string(), - fallback_allowed: true, - opening: Some(OpeningAuditSnapshot { - request_line: "GET /codex/responses HTTP/1.1".to_string(), - header_order: vec!["Host".to_string(), "Connection".to_string()], - headers: vec![OpeningAuditHeader { - name: "Sec-WebSocket-Extensions".to_string(), - value: "permessage-deflate; client_max_window_bits".to_string(), - }], - }), - payload: Some(PayloadAuditSnapshot { - top_level_keys: vec!["type".to_string(), "model".to_string()], - body: json!({"type": "response.create"}), - }), - error: Some(WebSocketAuditErrorSnapshot { - classification: "opening_failed".to_string(), - message: "connection refused".to_string(), - }), - }; - let reference = codex_proxy_rs::upstream::protocol::websocket::WebSocketAuditArtifact { - transport_mode: "websocket_required".to_string(), - fallback_allowed: false, - opening: Some(OpeningAuditSnapshot { - request_line: "GET /codex/responses?source=audit HTTP/1.1".to_string(), - header_order: vec!["Host".to_string(), "Authorization".to_string()], - headers: vec![OpeningAuditHeader { - name: "Sec-WebSocket-Extensions".to_string(), - value: "identity".to_string(), - }], - }), - payload: Some(PayloadAuditSnapshot { - top_level_keys: vec!["type".to_string(), "previous_response_id".to_string()], - body: json!({"type": "response.create"}), - }), - error: Some(WebSocketAuditErrorSnapshot { - classification: "timeout".to_string(), - message: "timed out".to_string(), - }), - }; - - let diff = websocket_parity_diff(¤t, &reference); - - assert!(diff.differences.iter().any(|difference| { - difference.path == "opening.header_order" - && difference.current == json!(["Host", "Connection"]) - && difference.reference == json!(["Host", "Authorization"]) - })); - assert!(diff - .differences - .iter() - .any(|difference| difference.path == "payload.top_level_keys")); - assert!(diff - .differences - .iter() - .any(|difference| difference.path == "transport_mode")); -} - #[test] fn codex_websocket_event_to_sse_frame_should_encode_public_events_only() { let event = json!({ diff --git a/tests/upstream/protocol/mod.rs b/tests/upstream/protocol/mod.rs index 2de011830..83c41d5f7 100644 --- a/tests/upstream/protocol/mod.rs +++ b/tests/upstream/protocol/mod.rs @@ -4,9 +4,8 @@ use codex_proxy_rs::proxy::openai::responses::{ }; use codex_proxy_rs::upstream::models::ParsedModelName; use codex_proxy_rs::upstream::protocol::events::{ - cooldown_with_jitter, extract_sse_usage, extract_usage, parse_rate_limit_headers, - parse_rate_limits_event, rate_limit_quota, retry_after_seconds_from_body, RateLimitWindow, - TokenUsage, + extract_sse_usage, extract_usage, parse_rate_limit_headers, parse_rate_limits_event, + rate_limit_quota, retry_after_seconds_from_body, RateLimitWindow, TokenUsage, }; use codex_proxy_rs::upstream::protocol::responses::{ apply_response_model_options, http_sse_fallback_allowed, response_from_codex_sse, @@ -33,8 +32,7 @@ use codex_proxy_rs::upstream::protocol::websocket::{ websocket_message_output_item_event_invalid_required_fields, websocket_metadata_turn_state, websocket_output_item_event_invalid_item_type_tag, websocket_output_item_event_invalid_metadata, websocket_output_item_event_missing_item, - websocket_output_item_event_non_object_item, websocket_parity_diff, - websocket_payload_audit_snapshot, + websocket_output_item_event_non_object_item, websocket_payload_audit_snapshot, websocket_reasoning_output_item_event_invalid_required_fields, websocket_reasoning_summary_part_added_missing_summary_index, websocket_response_completed_missing_response, websocket_response_completed_parse_error, @@ -43,22 +41,13 @@ use codex_proxy_rs::upstream::protocol::websocket::{ websocket_tool_search_call_output_item_event_invalid_required_fields, websocket_tool_search_output_item_event_invalid_required_fields, websocket_web_search_call_output_item_event_invalid_required_fields, OpeningAuditHeader, - OpeningAuditSnapshot, PayloadAuditSnapshot, WebSocketAuditErrorSnapshot, - WebSocketErrorClassificationProfile, + OpeningAuditSnapshot, WebSocketErrorClassificationProfile, }; use serde_json::json; +use crate::support::assertions::assert_substrings_appear_in_order; + mod codex_websocket; mod openai_chat; mod openai_responses; mod usage_rate_limits; - -fn assert_substrings_appear_in_order(haystack: &str, needles: &[&str]) { - let mut cursor = 0; - for needle in needles { - let Some(offset) = haystack[cursor..].find(needle) else { - panic!("expected substring {needle:?} after byte {cursor} in:\n{haystack}"); - }; - cursor += offset + needle.len(); - } -} diff --git a/tests/upstream/protocol/usage_rate_limits.rs b/tests/upstream/protocol/usage_rate_limits.rs index d0543503b..50afbd0c9 100644 --- a/tests/upstream/protocol/usage_rate_limits.rs +++ b/tests/upstream/protocol/usage_rate_limits.rs @@ -92,8 +92,8 @@ fn extract_usage_should_read_image_generation_tokens_separately() { } #[test] -fn extract_usage_should_merge_openai_usage_shape() { - let first = extract_usage(&json!({ +fn extract_usage_should_read_openai_usage_shape() { + let usage = extract_usage(&json!({ "usage": { "prompt_tokens": 8, "completion_tokens": 4, @@ -103,26 +103,17 @@ fn extract_usage_should_merge_openai_usage_shape() { } })) .expect("usage should exist"); - let second = TokenUsage { - input_tokens: 1, - output_tokens: 2, - cached_tokens: 0, - reasoning_tokens: 0, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: 3, - }; assert_eq!( - first.merged(second), + usage, TokenUsage { - input_tokens: 9, - output_tokens: 6, + input_tokens: 8, + output_tokens: 4, cached_tokens: 2, reasoning_tokens: 0, image_input_tokens: 0, image_output_tokens: 0, - total_tokens: 15, + total_tokens: 12, } ); } @@ -314,7 +305,8 @@ fn parse_rate_limits_event_should_extract_internal_websocket_rate_limits() { let parsed = parse_rate_limits_event(&event).expect("event should parse"); - assert_eq!(parsed.primary.expect("primary window").used_percent, 99.5); + let used_percent = parsed.primary.expect("primary window").used_percent; + assert!((used_percent - 99.5).abs() < f64::EPSILON); assert_eq!( parsed.secondary.expect("secondary window").reset_at, Some(1_894_056_000) @@ -360,13 +352,6 @@ fn rate_limit_quota_should_preserve_existing_monthly_and_credits_when_passive_da assert_eq!(quota["credits"]["balance"], 12); } -#[test] -fn cooldown_with_jitter_should_return_positive_duration_within_expected_range() { - let duration = cooldown_with_jitter(60, 2_000); - - assert!((48..=72).contains(&duration.num_seconds())); -} - #[tokio::test] async fn refresh_scheduler_should_refresh_before_expiry_and_preserve_refresh_token() { use async_trait::async_trait; diff --git a/tests/upstream/transport/client.rs b/tests/upstream/transport/client.rs index 413ac30a7..bc12ff1a9 100644 --- a/tests/upstream/transport/client.rs +++ b/tests/upstream/transport/client.rs @@ -3,24 +3,12 @@ use std::sync::{ Arc, }; +use super::{accept_codex_test_websocket, completed_websocket_response, request_context}; use codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest; -use codex_proxy_rs::upstream::transport::{ - CodexBackendClient, CodexRequestContext, CodexWebSocketPool, -}; +use codex_proxy_rs::upstream::transport::{CodexBackendClient, CodexWebSocketPool}; use futures::{SinkExt, StreamExt}; -use serde_json::json; -use tokio::net::{TcpListener, TcpStream}; -use tokio_tungstenite::{ - accept_hdr_async_with_config, - tungstenite::{ - extensions::{compression::deflate::DeflateConfig, ExtensionsConfig}, - handshake::server::{ - Callback, ErrorResponse, Request as WsRequest, Response as WsResponse, - }, - protocol::WebSocketConfig, - Message, - }, -}; +use tokio::net::TcpListener; +use tokio_tungstenite::tungstenite::Message; #[tokio::test] async fn codex_backend_client_should_apply_configured_websocket_pool() { @@ -31,11 +19,11 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() { let server = tokio::spawn(async move { let (first_stream, _) = listener.accept().await.unwrap(); accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); - let mut first_websocket = accept_runtime_test_websocket(first_stream).await; + let mut first_websocket = accept_codex_test_websocket(first_stream).await; let _first_message = first_websocket.next().await.unwrap().unwrap(); first_websocket .send(Message::Text( - response_completed_websocket_message("resp_runtime_pool_first").into(), + completed_websocket_response("resp_runtime_pool_first", 3, 1).into(), )) .await .unwrap(); @@ -45,7 +33,7 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() { let _second_message = second_message.unwrap().unwrap(); first_websocket .send(Message::Text( - response_completed_websocket_message("resp_runtime_pool_second").into(), + completed_websocket_response("resp_runtime_pool_second", 3, 1).into(), )) .await .unwrap(); @@ -54,11 +42,11 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() { accepted = listener.accept() => { let (second_stream, _) = accepted.unwrap(); accepted_connections_for_server.fetch_add(1, Ordering::SeqCst); - let mut second_websocket = accept_runtime_test_websocket(second_stream).await; + let mut second_websocket = accept_codex_test_websocket(second_stream).await; let _second_message = second_websocket.next().await.unwrap().unwrap(); second_websocket .send(Message::Text( - response_completed_websocket_message("resp_runtime_pool_second").into(), + completed_websocket_response("resp_runtime_pool_second", 3, 1).into(), )) .await .unwrap(); @@ -67,7 +55,7 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() { } } }); - let pool = CodexWebSocketPool::new(8, std::time::Duration::from_secs(60)); + let pool = CodexWebSocketPool::new(8, std::time::Duration::from_mins(1)); let backend = CodexBackendClient::new( reqwest::Client::new(), format!("http://{addr}"), @@ -99,102 +87,3 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() { assert!(second.body.contains("resp_runtime_pool_second")); assert_eq!(accepted_connections.load(Ordering::SeqCst), 1); } - -fn websocket_accept_config() -> WebSocketConfig { - let mut extensions = ExtensionsConfig::default(); - extensions.permessage_deflate = Some(DeflateConfig::default()); - - let mut config = WebSocketConfig::default(); - config.extensions = extensions; - config -} - -async fn accept_runtime_test_websocket( - stream: TcpStream, -) -> tokio_tungstenite::WebSocketStream { - accept_runtime_test_websocket_with(stream, |_request, response| { - response.headers_mut().insert( - "sec-websocket-extensions", - "permessage-deflate".parse().unwrap(), - ); - }) - .await -} - -struct TestWebSocketCallback(F); - -impl Callback for TestWebSocketCallback -where - F: FnOnce(&WsRequest, &mut WsResponse) + Unpin, -{ - fn on_request( - self, - request: &WsRequest, - mut response: WsResponse, - ) -> Result { - (self.0)(request, &mut response); - Ok(response) - } -} - -async fn accept_runtime_test_websocket_with( - stream: TcpStream, - callback: F, -) -> tokio_tungstenite::WebSocketStream -where - F: FnOnce(&WsRequest, &mut WsResponse) + Unpin, -{ - accept_hdr_async_with_config( - stream, - TestWebSocketCallback(callback), - Some(websocket_accept_config()), - ) - .await - .unwrap() -} - -fn request_context<'a>( - request_id: &'a str, - account_id: Option<&'a str>, -) -> CodexRequestContext<'a> { - CodexRequestContext { - access_token: "access-token", - account_id, - request_id, - turn_state: None, - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - version: None, - codex_window_id: None, - parent_thread_id: None, - cookie_header: None, - installation_id: None, - session_id: None, - } -} - -fn response_completed_websocket_message(response_id: &str) -> String { - json!({ - "type": "response.completed", - "response": { - "id": response_id, - "object": "response", - "status": "completed", - "output": [{ - "type": "message", - "role": "assistant", - "content": [{ - "type": "output_text", - "text": "test response" - }] - }], - "usage": { - "input_tokens": 3, - "output_tokens": 1, - "total_tokens": 4 - } - } - }) - .to_string() -} diff --git a/tests/upstream/transport/headers.rs b/tests/upstream/transport/headers.rs index 5ebb1ba36..bc1406928 100644 --- a/tests/upstream/transport/headers.rs +++ b/tests/upstream/transport/headers.rs @@ -273,7 +273,7 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an let pool = Arc::new( codex_proxy_rs::upstream::transport::websocket_pool::CodexWebSocketPool::new( 8, - std::time::Duration::from_secs(60), + std::time::Duration::from_mins(1), ), ); let backend = CodexBackendClient::new( diff --git a/tests/upstream/transport/http_client.rs b/tests/upstream/transport/http_client.rs index d26e9c3cc..78b975486 100644 --- a/tests/upstream/transport/http_client.rs +++ b/tests/upstream/transport/http_client.rs @@ -29,17 +29,17 @@ fn custom_ca_should_report_environment_cache_key_consistently() { "unset" => None, "ssl_cert_file" => Some(format!( "{}={SSL_CERT_PATH}", - codex_proxy_rs::upstream::transport::SSL_CERT_FILE_ENV + codex_proxy_rs::upstream::transport::tls::SSL_CERT_FILE_ENV )), "codex_ca_priority" => Some(format!( "{}={CODEX_CA_PATH}", - codex_proxy_rs::upstream::transport::CODEX_CA_CERT_ENV + codex_proxy_rs::upstream::transport::tls::CODEX_CA_CERT_ENV )), _ => panic!("unknown custom CA cache key test case: {case}"), }; assert_eq!( - codex_proxy_rs::upstream::transport::custom_ca_env_cache_key(), + codex_proxy_rs::upstream::transport::tls::custom_ca_env_cache_key(), expected ); return; @@ -52,18 +52,18 @@ fn custom_ca_should_report_environment_cache_key_consistently() { "ssl_cert_file", None, Some(( - codex_proxy_rs::upstream::transport::SSL_CERT_FILE_ENV, + codex_proxy_rs::upstream::transport::tls::SSL_CERT_FILE_ENV, SSL_CERT_PATH, )), ), ( "codex_ca_priority", Some(( - codex_proxy_rs::upstream::transport::CODEX_CA_CERT_ENV, + codex_proxy_rs::upstream::transport::tls::CODEX_CA_CERT_ENV, CODEX_CA_PATH, )), Some(( - codex_proxy_rs::upstream::transport::SSL_CERT_FILE_ENV, + codex_proxy_rs::upstream::transport::tls::SSL_CERT_FILE_ENV, SSL_CERT_PATH, )), ), @@ -76,8 +76,8 @@ fn custom_ca_should_report_environment_cache_key_consistently() { .arg("custom_ca_should_report_environment_cache_key_consistently") .arg("--nocapture") .env(CASE_ENV, case) - .env_remove(codex_proxy_rs::upstream::transport::CODEX_CA_CERT_ENV) - .env_remove(codex_proxy_rs::upstream::transport::SSL_CERT_FILE_ENV); + .env_remove(codex_proxy_rs::upstream::transport::tls::CODEX_CA_CERT_ENV) + .env_remove(codex_proxy_rs::upstream::transport::tls::SSL_CERT_FILE_ENV); if let Some((key, value)) = codex_ca { command.env(key, value); } @@ -95,45 +95,6 @@ fn custom_ca_should_report_environment_cache_key_consistently() { } } -#[tokio::test] -async fn codex_backend_client_probe_should_return_models_endpoint_status() { - let server = wiremock::MockServer::start().await; - wiremock::Mock::given(wiremock::matchers::method("GET")) - .and(wiremock::matchers::path("/codex/models")) - .respond_with(wiremock::ResponseTemplate::new(204)) - .mount(&server) - .await; - let client = CodexBackendClient::new( - reqwest::Client::builder().no_proxy().build().unwrap(), - server.uri(), - crate::support::fingerprint::test_fingerprint(), - ); - - let probe = client - .probe_connectivity(CodexRequestContext { - access_token: "access-token", - account_id: Some("chatgpt-account"), - request_id: "req_probe", - turn_state: None, - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - version: None, - codex_window_id: None, - parent_thread_id: None, - cookie_header: None, - installation_id: None, - session_id: None, - }) - .await - .unwrap(); - - assert_eq!(probe.status, reqwest::StatusCode::NO_CONTENT); - assert!(probe - .endpoint - .ends_with("/codex/models?client_version=26.519.81530")); -} - #[tokio::test] async fn codex_backend_client_should_cap_non_success_error_body_at_one_mib() { let server = wiremock::MockServer::start().await; @@ -204,7 +165,7 @@ async fn build_reqwest_client_should_reuse_cached_connection_pool() { write_empty_http_response(&mut second_stream).await; false } - _ = tokio::time::sleep(Duration::from_millis(500)) => false, + () = tokio::time::sleep(Duration::from_millis(500)) => false, } }); diff --git a/tests/upstream/transport/mod.rs b/tests/upstream/transport/mod.rs index 2aca771a0..79dd0106b 100644 --- a/tests/upstream/transport/mod.rs +++ b/tests/upstream/transport/mod.rs @@ -34,11 +34,13 @@ use tokio_tungstenite::{ }, }; -pub mod client; -pub mod headers; -pub mod http_client; -pub mod websocket; -pub mod websocket_pool; +use crate::support::assertions::assert_substrings_appear_in_order; + +pub(crate) mod client; +pub(crate) mod headers; +pub(crate) mod http_client; +pub(crate) mod websocket; +pub(crate) mod websocket_pool; fn websocket_accept_config() -> WebSocketConfig { let mut extensions = ExtensionsConfig::default(); @@ -49,7 +51,7 @@ fn websocket_accept_config() -> WebSocketConfig { config } -pub async fn accept_codex_test_websocket( +pub(crate) async fn accept_codex_test_websocket( stream: TcpStream, ) -> tokio_tungstenite::WebSocketStream { accept_codex_test_websocket_with(stream, |_request, response| { @@ -77,7 +79,7 @@ where } } -pub async fn accept_codex_test_websocket_with( +pub(crate) async fn accept_codex_test_websocket_with( stream: TcpStream, callback: F, ) -> tokio_tungstenite::WebSocketStream @@ -93,7 +95,7 @@ where .unwrap() } -pub fn read_header_names(request: &str) -> Vec { +pub(crate) fn read_header_names(request: &str) -> Vec { request .lines() .skip(1) @@ -105,7 +107,7 @@ pub fn read_header_names(request: &str) -> Vec { .collect() } -pub fn assert_header_subsequence(actual: &[String], expected: &[&str]) { +pub(crate) fn assert_header_subsequence(actual: &[String], expected: &[&str]) { let mut offset = 0; for expected_name in expected { let Some(position) = actual[offset..] @@ -118,7 +120,7 @@ pub fn assert_header_subsequence(actual: &[String], expected: &[&str]) { } } -pub async fn read_http_request(stream: &mut TcpStream) -> String { +pub(crate) async fn read_http_request(stream: &mut TcpStream) -> String { let mut request = Vec::new(); let mut buffer = [0u8; 1024]; loop { @@ -134,14 +136,14 @@ pub async fn read_http_request(stream: &mut TcpStream) -> String { String::from_utf8(request).unwrap() } -pub async fn write_empty_http_response(stream: &mut TcpStream) { +pub(crate) async fn write_empty_http_response(stream: &mut TcpStream) { stream .write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") .await .unwrap(); } -pub async fn write_completed_sse_response(stream: &mut TcpStream) { +pub(crate) async fn write_completed_sse_response(stream: &mut TcpStream) { let body = concat!( "event: response.completed\n", "data: {\"response\":{\"id\":\"resp_order\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1}}}\n", @@ -159,7 +161,7 @@ pub async fn write_completed_sse_response(stream: &mut TcpStream) { .unwrap(); } -pub fn request_context<'a>( +pub(crate) fn request_context<'a>( request_id: &'a str, account_id: Option<&'a str>, ) -> CodexRequestContext<'a> { @@ -180,7 +182,7 @@ pub fn request_context<'a>( } } -pub fn prepared_websocket_request(base_url: &str) -> CodexWebSocketRequest { +pub(crate) fn prepared_websocket_request(base_url: &str) -> CodexWebSocketRequest { let request = codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", @@ -199,7 +201,7 @@ pub fn prepared_websocket_request(base_url: &str) -> CodexWebSocketRequest { .expect("payload should serialize") } -pub fn pooled_websocket_request( +pub(crate) fn pooled_websocket_request( conversation_id: &str, ) -> codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest { let mut request = @@ -214,7 +216,7 @@ pub fn pooled_websocket_request( request } -pub fn completed_websocket_response( +pub(crate) fn completed_websocket_response( response_id: &str, input_tokens: u64, output_tokens: u64, @@ -235,24 +237,14 @@ pub fn completed_websocket_response( .to_string() } -pub fn assert_substrings_appear_in_order(haystack: &str, needles: &[&str]) { - let mut cursor = 0; - for needle in needles { - let Some(offset) = haystack[cursor..].find(needle) else { - panic!("expected substring {needle:?} after byte {cursor} in:\n{haystack}"); - }; - cursor += offset + needle.len(); - } -} - -pub fn websocket_pool_config_for_tests( +pub(crate) fn websocket_pool_config_for_tests( maintenance_interval: Option, ping_interval: Option, liveness_timeout: Option, ) -> CodexWebSocketPoolConfig { CodexWebSocketPoolConfig { enabled: true, - max_age: Duration::from_secs(60), + max_age: Duration::from_mins(1), max_per_account: 8, maintenance_interval, ping_interval, @@ -261,7 +253,7 @@ pub fn websocket_pool_config_for_tests( } } -pub async fn write_compact_json_response(stream: &mut TcpStream) { +pub(crate) async fn write_compact_json_response(stream: &mut TcpStream) { stream .write_all( b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 15\r\n\r\n{\"id\":\"resp_1\"}", diff --git a/tests/upstream/transport/websocket.rs b/tests/upstream/transport/websocket.rs index 49accec4e..a78aeb389 100644 --- a/tests/upstream/transport/websocket.rs +++ b/tests/upstream/transport/websocket.rs @@ -104,7 +104,7 @@ async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_ websocket.send(Message::Text(payload.into())).await.unwrap(); } }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_secs(60))); + let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -134,73 +134,6 @@ async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_ .contains("resp_6f8d0c2b5a4e4a0d9c1b7e3f2a8d5c6b")); } -#[test] -fn websocket_connection_should_render_raw_opening_bytes_for_capture_parity() { - let connection = - CodexWebSocketConnection::responses( - "https://chatgpt.com/backend-api", - "test-websocket-key", - vec![ - ( - "chatgpt-account-id".to_string(), - "chatgpt-account".to_string(), - ), - ( - "authorization".to_string(), - "Bearer access-token".to_string(), - ), - ( - "user-agent".to_string(), - "Codex Desktop/26.519.81530 (darwin; arm64)".to_string(), - ), - ("originator".to_string(), "Codex Desktop".to_string()), - ( - "openai-beta".to_string(), - "responses_websockets=2026-02-06".to_string(), - ), - ( - "x-codex-beta-features".to_string(), - "terminal_resize_reflow,memories,network_proxy,prevent_idle_sleep,remote_compaction_v2".to_string(), - ), - ("x-client-request-id".to_string(), "session-1".to_string()), - ("session-id".to_string(), "session-1".to_string()), - ("thread-id".to_string(), "session-1".to_string()), - ("x-codex-window-id".to_string(), "session-1:0".to_string()), - ( - "x-codex-turn-metadata".to_string(), - r#"{"installation_id":"install-1","session_id":"session-1","thread_id":"session-1","turn_id":"","window_id":"session-1:0","request_kind":"prewarm","sandbox":"seccomp"}"#.to_string(), - ), - ], - ); - - let opening = connection.opening_request_text(); - - assert!(opening.starts_with("GET /backend-api/codex/responses HTTP/1.1\r\n")); - assert_substrings_appear_in_order( - &opening, - &[ - "Host: chatgpt.com\r\n", - "Connection: Upgrade\r\n", - "Upgrade: websocket\r\n", - "Sec-WebSocket-Version: 13\r\n", - "Sec-WebSocket-Key: test-websocket-key\r\n", - "chatgpt-account-id: chatgpt-account\r\n", - "authorization: Bearer access-token\r\n", - "user-agent: Codex Desktop/26.519.81530 (darwin; arm64)\r\n", - "originator: Codex Desktop\r\n", - "openai-beta: responses_websockets=2026-02-06\r\n", - "x-codex-beta-features: terminal_resize_reflow,memories,network_proxy,prevent_idle_sleep,remote_compaction_v2\r\n", - "x-client-request-id: session-1\r\n", - "session-id: session-1\r\n", - "thread-id: session-1\r\n", - "x-codex-window-id: session-1:0\r\n", - r#"x-codex-turn-metadata: {"installation_id":"install-1","session_id":"session-1","thread_id":"session-1","turn_id":"","window_id":"session-1:0","request_kind":"prewarm","sandbox":"seccomp"}"#, - "sec-websocket-extensions: permessage-deflate; client_max_window_bits\r\n", - ], - ); - assert!(opening.ends_with("\r\n\r\n")); -} - #[test] fn websocket_connection_should_prepare_response_create_payload_text() { let request = @@ -1196,7 +1129,7 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_ Vec::new(), ); request.previous_response_id = Some("resp_previous".to_string()); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_secs(60))); + let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), diff --git a/tests/upstream/transport/websocket_pool.rs b/tests/upstream/transport/websocket_pool.rs index 2991a613d..ac7590e3e 100644 --- a/tests/upstream/transport/websocket_pool.rs +++ b/tests/upstream/transport/websocket_pool.rs @@ -1,20 +1,5 @@ use super::*; -#[test] -fn websocket_pool_should_apply_capacity_and_age_policy() { - let pool = CodexWebSocketPool::new(2, Duration::from_secs(60)); - - assert_eq!( - ( - pool.permits_new_connection(1), - pool.permits_new_connection(2), - pool.should_recycle(Duration::from_secs(59)), - pool.should_recycle(Duration::from_secs(60)), - ), - (true, false, false, true) - ); -} - #[tokio::test] async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and_conversation() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -50,7 +35,7 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and } websocket.close(None).await.unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_secs(60))); + let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -139,7 +124,7 @@ async fn websocket_pool_should_bypass_busy_key_with_one_shot_connections() { .unwrap(); first_websocket.close(None).await.unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::with_default_max_age()); + let pool = Arc::new(CodexWebSocketPool::default()); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -210,7 +195,7 @@ async fn websocket_pool_should_bypass_new_keys_after_account_cap() { } } }); - let pool = Arc::new(CodexWebSocketPool::with_limits(Duration::from_secs(60), 1)); + let pool = Arc::new(CodexWebSocketPool::new(1, Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), @@ -479,7 +464,7 @@ async fn websocket_pool_should_gc_expired_idle_connections() { .await .expect("first websocket response should succeed"); tokio::time::sleep(Duration::from_millis(15)).await; - pool.gc_sweep().await; + pool.maintain_idle_connections().await; let second = backend .create_response( &request, @@ -562,7 +547,7 @@ async fn codex_backend_client_should_ping_idle_pooled_websocket_from_background_ let pool = Arc::new(CodexWebSocketPool::with_config( websocket_pool_config_for_tests( Some(Duration::from_millis(20)), - Some(Duration::from_secs(60)), + Some(Duration::from_mins(1)), None, ), )); @@ -973,7 +958,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_err .unwrap(); second_websocket.close(None).await.unwrap(); }); - let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_secs(60))); + let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1))); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), format!("http://{addr}"), diff --git a/web/src/api/index.ts b/web/src/api/index.ts index 5b5c82b66..09368ea88 100644 --- a/web/src/api/index.ts +++ b/web/src/api/index.ts @@ -10,9 +10,6 @@ export * from './modules/api-keys' // 事件日志 export * from './modules/logs' -// 使用统计 -export * from './modules/usage' - // Dashboard 聚合 export * from './modules/dashboard' diff --git a/web/src/api/modules/accounts.ts b/web/src/api/modules/accounts.ts index 05acf0eef..a9579fb36 100644 --- a/web/src/api/modules/accounts.ts +++ b/web/src/api/modules/accounts.ts @@ -8,14 +8,6 @@ export function getAccounts(params?: any) { }) } -export function createAccount(data: any) { - return request({ - url: '/api/admin/accounts', - method: 'POST', - data, - }) -} - export function importAccounts(data: any) { return request({ url: '/api/admin/accounts/import', @@ -72,14 +64,6 @@ export function getAccountQuota(params: any) { }) } -export function testAccountConnection(data: any) { - return request({ - url: '/api/admin/accounts/health-check', - method: 'POST', - data, - }) -} - export async function testAccountConnectionStream(data: any, onEvent: any, signal?: AbortSignal) { const { id, ...payload } = data const baseURL = import.meta.env.DEV ? '/dev' : '' @@ -116,7 +100,7 @@ export async function testAccountConnectionStream(data: any, onEvent: any, signa consumeAccountTestFrames(buffer, onEvent) } -export function getAccountTestModels(data: any) { +export function getAccountModels(data: any) { return request({ url: '/api/admin/accounts/models', method: 'GET', diff --git a/web/src/api/modules/usage.ts b/web/src/api/modules/usage.ts deleted file mode 100644 index 366be85a5..000000000 --- a/web/src/api/modules/usage.ts +++ /dev/null @@ -1,16 +0,0 @@ -import request from '../request' - -export function getUsageSummary() { - return request({ - url: '/api/admin/usage/summary', - method: 'GET', - }) -} - -export function getUsageStats(params?: any) { - return request({ - url: '/api/admin/usage', - method: 'GET', - params, - }) -} diff --git a/web/src/api/request.ts b/web/src/api/request.ts index d37751b14..d9f732961 100644 --- a/web/src/api/request.ts +++ b/web/src/api/request.ts @@ -18,11 +18,6 @@ class ApiError extends Error { } } -interface RequestConfig extends AxiosRequestConfig { - skipErrorHandler?: boolean - skipAuth?: boolean -} - const http: AxiosInstance = axios.create({ baseURL: import.meta.env.DEV ? '/dev' : '', timeout: 30000, @@ -51,11 +46,7 @@ http.interceptors.response.use( return response }, (error: AxiosError) => { - const { response, config } = error - - if ((config as RequestConfig)?.skipErrorHandler) { - return Promise.reject(error) - } + const { response } = error const status = response?.status || 0 const message = response?.data?.message || error.message || '请求失败' @@ -93,7 +84,7 @@ function isApiEnvelope(value: any): value is { data: any } { ) } -export default async function request(config: RequestConfig): Promise { +export default async function request(config: AxiosRequestConfig): Promise { const response = await http.request>({ ...config, }) diff --git a/web/src/components/base/BaseSwitch.vue b/web/src/components/base/BaseSwitch.vue deleted file mode 100644 index dbfb669b3..000000000 --- a/web/src/components/base/BaseSwitch.vue +++ /dev/null @@ -1,52 +0,0 @@ - - - diff --git a/web/src/components/base/BaseToast/index.ts b/web/src/components/base/BaseToast/index.ts index 698c0dca5..08b646b4b 100644 --- a/web/src/components/base/BaseToast/index.ts +++ b/web/src/components/base/BaseToast/index.ts @@ -1,3 +1,2 @@ export { default as BaseToast } from './BaseToast.vue' export { toast } from './toast' -export type { ToastMessage, ToastOptions, ToastType } from './toast' diff --git a/web/src/views/accounts/components/AccountCreateModal.vue b/web/src/views/accounts/components/AccountCreateModal.vue index 55c476fa9..1a46581ad 100644 --- a/web/src/views/accounts/components/AccountCreateModal.vue +++ b/web/src/views/accounts/components/AccountCreateModal.vue @@ -34,7 +34,7 @@ const fileError = ref('') const modeOptions = [ { label: 'OAuth 授权', value: 'oauth' }, { label: 'RT 导入', value: 'rt' }, - { label: 'CPR', value: 'native' }, + { label: 'CPR', value: 'cpr' }, { label: 'Sub2API', value: 'sub2api' }, ] diff --git a/web/src/views/accounts/composables/useAccountConnectionTest.ts b/web/src/views/accounts/composables/useAccountConnectionTest.ts index bfa6ef6ba..3f9734071 100644 --- a/web/src/views/accounts/composables/useAccountConnectionTest.ts +++ b/web/src/views/accounts/composables/useAccountConnectionTest.ts @@ -1,7 +1,7 @@ import { computed, onBeforeUnmount, ref, watch } from 'vue' import { CheckCircle2, Clock3, Wifi, XCircle } from '@lucide/vue' -import { getAccountTestModels, testAccountConnectionStream } from '@/api' +import { getAccountModels, testAccountConnectionStream } from '@/api' import { withMinimumDuration } from '@/utils/async' import { formatDateTime, formatTime } from '@/utils/date' @@ -182,7 +182,7 @@ export function useAccountConnectionTest() { loadingConnectionTestModels.value = true connectionTestError.value = '' try { - const result = await getAccountTestModels({ id: account.id }) + const result = await getAccountModels({ id: account.id }) connectionTestModelOptions.value = (result.models || []).map((model: any) => ({ label: model.label || model.id, value: model.id, diff --git a/web/src/views/accounts/composables/useAccountMutations.ts b/web/src/views/accounts/composables/useAccountMutations.ts index fcc7cf18a..265ecfa9a 100644 --- a/web/src/views/accounts/composables/useAccountMutations.ts +++ b/web/src/views/accounts/composables/useAccountMutations.ts @@ -146,7 +146,7 @@ export function useAccountMutations(options: { .filter(Boolean) .map((refreshToken) => ({ refreshToken })) - return accounts.length ? { sourceFormat: 'native', accounts } : null + return accounts.length ? { sourceFormat: 'cpr', accounts } : null } const text = createForm.value.importText.trim() @@ -159,7 +159,7 @@ export function useAccountMutations(options: { throw new Error('JSON 格式不正确') } - const sourceFormat = createForm.value.mode === 'sub2api' ? 'sub2api' : 'native' + const sourceFormat = createForm.value.mode === 'sub2api' ? 'sub2api' : 'cpr' if (Array.isArray(parsed)) { return { sourceFormat, accounts: parsed } } diff --git a/web/src/views/dashboard/components/EventLogCard.vue b/web/src/views/dashboard/components/EventLogCard.vue index 687745b72..6901fc65d 100644 --- a/web/src/views/dashboard/components/EventLogCard.vue +++ b/web/src/views/dashboard/components/EventLogCard.vue @@ -45,7 +45,6 @@ function rowTone(row: object) { } const filteredRows = computed(() => { - if (activeFilter.value === 'warn') return props.rows.filter((row) => row.tone === 'warning') if (activeFilter.value === 'error') return props.rows.filter((row) => row.tone === 'danger') return props.rows })