diff --git a/Cargo.lock b/Cargo.lock index 25a93e4b7..81f5105c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,41 +8,6 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" -[[package]] -name = "aead" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" -dependencies = [ - "crypto-common 0.1.6", - "generic-array", -] - -[[package]] -name = "aes" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" -dependencies = [ - "cfg-if", - "cipher", - "cpufeatures 0.2.17", -] - -[[package]] -name = "aes-gcm" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" -dependencies = [ - "aead", - "aes", - "cipher", - "ctr", - "ghash", - "subtle", -] - [[package]] name = "aho-corasick" version = "1.1.4" @@ -348,16 +313,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "cipher" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" -dependencies = [ - "crypto-common 0.1.6", - "inout", -] - [[package]] name = "cmov" version = "0.5.4" @@ -368,7 +323,6 @@ checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" name = "codex-proxy-rs" version = "0.1.0" dependencies = [ - "aes-gcm", "argon2", "async-trait", "axum", @@ -382,6 +336,7 @@ dependencies = [ "hmac", "indexmap", "rand 0.10.1", + "rand_core 0.6.4", "reqwest", "rustls", "rustls-native-certs", @@ -406,7 +361,6 @@ dependencies = [ "tungstenite", "uuid", "wiremock", - "zeroize", ] [[package]] @@ -617,7 +571,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" dependencies = [ "generic-array", - "rand_core 0.6.4", "typenum", ] @@ -630,15 +583,6 @@ dependencies = [ "hybrid-array", ] -[[package]] -name = "ctr" -version = "0.9.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" -dependencies = [ - "cipher", -] - [[package]] name = "ctutils" version = "0.4.2" @@ -1037,16 +981,6 @@ dependencies = [ "wasip3", ] -[[package]] -name = "ghash" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" -dependencies = [ - "opaque-debug", - "polyval", -] - [[package]] name = "h2" version = "0.4.15" @@ -1427,15 +1361,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "inout" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" -dependencies = [ - "generic-array", -] - [[package]] name = "ipnet" version = "2.12.0" @@ -1679,12 +1604,6 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" -[[package]] -name = "opaque-debug" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" - [[package]] name = "openssl-probe" version = "0.2.1" @@ -1814,18 +1733,6 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" -[[package]] -name = "polyval" -version = "0.6.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" -dependencies = [ - "cfg-if", - "cpufeatures 0.2.17", - "opaque-debug", - "universal-hash", -] - [[package]] name = "potential_utf" version = "0.1.5" @@ -3183,16 +3090,6 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" -[[package]] -name = "universal-hash" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" -dependencies = [ - "crypto-common 0.1.6", - "subtle", -] - [[package]] name = "untrusted" version = "0.9.0" diff --git a/Cargo.toml b/Cargo.toml index a3002367c..dcbd0d00e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,7 +8,6 @@ rust-version = "1.95" description = "Codex/OpenAI proxy backed by ChatGPT accounts." [dependencies] -aes-gcm = "0.10.3" argon2 = "0.5.3" async-trait = "0.1.89" axum = "0.8.9" @@ -22,6 +21,7 @@ hex = "0.4.3" hmac = "0.13.0" indexmap = { version = "2.0", features = ["serde"] } rand = "0.10.1" +rand_core = { version = "0.6.4", features = ["getrandom"] } reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] } rustls = { version = "=0.23.36", default-features = false, features = ["ring", "std", "tls12"] } rustls-native-certs = "0.8.4" @@ -31,7 +31,7 @@ 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", "migrate"] } +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"] } @@ -43,7 +43,6 @@ tracing-appender = "0.2.5" tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] } tungstenite = { version = "0.27.0", features = ["deflate", "proxy"] } uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] } -zeroize = "1.8.2" [dev-dependencies] tempfile = "3.27.0" diff --git a/config.yaml b/config.yaml index 89df33f4c..94743e73b 100644 --- a/config.yaml +++ b/config.yaml @@ -8,6 +8,7 @@ model: default_reasoning_effort: null default_service_tier: null aliases: {} + account_routes: {} auth: refresh_margin_seconds: 300 refresh_enabled: true @@ -34,7 +35,6 @@ usage_stats: database: url: 'sqlite://.runtime/data/codex-proxy-rs.sqlite' security: - master_key_file: .runtime/data/master.key api_key_pepper_file: .runtime/data/api-key-pepper.key tls: force_http11: false diff --git a/docs/database-design-audit.md b/docs/database-design-audit.md new file mode 100644 index 000000000..3f4b22d01 --- /dev/null +++ b/docs/database-design-audit.md @@ -0,0 +1,440 @@ +# 数据库设计审计 + +审计日期：2026-06-27 +执行状态：P0/P1/P2/P3/P4/P5 已落地，当前源码 schema 已移除空壳迁移表，新增运行设置表、模型用量统计表、时间桶聚合表与模型账号路由表，统一本地敏感字段策略，并补齐关键查询与诊断筛选索引。 + +本文档审计当前 SQLite schema、真实查询路径、系统设置持久化方式与测试覆盖。结论以“新工程、无历史包袱”为前提：不保留半套迁移、不保留误导性表、不为了兼容旧数据维持不合理结构。 + +## 总体结论 + +当前数据库已经覆盖了轻量模型网关的核心运行域：管理员、客户端 API Key、账号池、账号用量、时间桶趋势聚合、账号 Cookie、设备指纹、事件日志、模型快照、会话亲和。表边界整体可理解，外键、布尔检查、状态检查、非负计数约束也有基础质量。 + +当前已完成第一轮数据库边界收敛： + +1. 已删除源码 schema 中的 `schema_migrations` 空壳表。 +2. 已新增 `runtime_settings`，后台系统设置改为数据库持久化。 +3. 已补齐账号、用量、API Key、指纹历史的关键分页索引。 +4. 已删除设置写回 `config.yaml` 的代码路径。 + +当前数据库审计文档中的结构性收敛项已经落地。趋势统计不再依赖事件日志窗口，而是写入独立聚合事实表；事件日志继续只承担诊断窗口。 + +## 当前事实 + +### Schema 初始化 + +当前初始化入口在 `src/infra/database.rs`： + +- 启动时创建 SQLite 父目录。 +- 使用 WAL。 +- 开启 foreign keys。 +- 设置 busy timeout。 +- 执行整段 `src/infra/schema.sql`。 + +当前选择是“当前态 schema 初始化”，不是版本化迁移。`schema.sql` 使用 `create table if not exists` 只负责创建新库；已有开发库如果曾经创建过旧结构，直接重建数据库。项目不保留迁移动作代码，也不保留兼容旧字段的分支。 + +### 当前表 + +`src/infra/schema.sql` 当前定义： + +- `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` + +源码 schema 不再包含 `schema_migrations`。已有开发库如果曾经创建过旧表，需要直接重建数据库；项目不保留迁移或兼容清理代码。 + +### 系统设置现状 + +后台系统设置当前已经进入数据库： + +- `runtime_settings` 保存后台可改的运行配置。 +- 启动时如果缺少设置行，会用启动配置初始化缺省行。 +- 启动时如果已有设置行，会用数据库运行设置覆盖运行时配置。 +- `POST /api/admin/settings` 只写 `runtime_settings`，不再写回 `config.yaml`。 + +`config.yaml` 只保留启动级配置，不再承担后台设置持久化。 + +## 高优先级结论 + +### 1. 空壳迁移表必须处理 + +状态：已处理。 + +`schema_migrations` 曾经存在，但没有任何代码读取、写入、校验版本。它会制造错误预期：看起来有迁移体系，实际只是 `create table if not exists`。 + +当前已经按新工程策略处理：直接删除空壳迁移表。schema 以当前文件为唯一事实；已有旧开发库直接重建，不提供迁移动作代码。 + +### 2. 系统设置应该入库 + +状态：已处理。 + +页面可改的运行配置不应该写回整份 YAML。`config.yaml` 应保留为启动级配置，例如： + +- 服务监听地址和端口。 +- 数据库 URL。 +- API Key pepper 文件路径。 +- 上游基础地址。 +- OAuth endpoint。 +- 日志目录。 +- 首次启动管理员默认值。 + +后台设置页改动的是运行态业务配置，更适合入库，例如： + +- 默认模型。 +- 模型别名。 +- 提前刷新秒数。 +- 刷新并发数。 +- 单账号并发数。 +- 请求间隔。 +- 账号调度策略。 +- 模型到账户的路由或绑定。 + +当前使用单行强类型表，而不是散乱 key-value： + +```sql +create table runtime_settings ( + id integer primary key check (id = 1), + default_model text not null, + model_aliases_json text not null default '{}', + refresh_margin_seconds integer not null check (refresh_margin_seconds > 0), + refresh_concurrency integer not null check (refresh_concurrency > 0), + max_concurrent_per_account integer not null check (max_concurrent_per_account > 0), + request_interval_ms integer not null check (request_interval_ms >= 0), + rotation_strategy text not null check (rotation_strategy in ('least_used', 'round_robin', 'sticky')), + updated_at text not null +); +``` + +模型映射绑定账号不塞进 `model_aliases_json`。这属于调度关系，不是配置字符串，当前使用独立关系表： + +```sql +create table model_account_routes ( + model text not null, + account_id text not null references accounts(id) on delete cascade, + priority integer not null default 0, + enabled integer not null default 1 check (enabled in (0, 1)), + created_at text not null, + updated_at text not null, + primary key (model, account_id) +); + +create index idx_model_account_routes_account on model_account_routes(account_id); +create index idx_model_account_routes_enabled_model on model_account_routes(enabled, model, priority asc, account_id); +``` + +这样页面配置、调度逻辑、账号删除行为都有明确数据库约束。 + +设置接口通过 `modelAccountRoutes` 读写该表，运行时账号池按请求模型过滤显式账号列表；账号删除时由外键级联清理相关路由。 + +### 3. 关键查询缺复合索引 + +状态：核心索引已处理。 + +已补齐： + +```sql +create index idx_accounts_added_id on accounts(added_at desc, id desc); +create index idx_account_usage_last_used_account on account_usage(last_used_at desc, account_id desc); +create index idx_client_api_keys_created_id on client_api_keys(created_at desc, id desc); +create index idx_fingerprint_update_history_created_id on fingerprint_update_history(created_at desc, id desc); +``` + +运行清理和诊断筛选索引已补齐： + +```sql +create index idx_account_cookies_expires on account_cookies(expires_at) where expires_at is not null; +create index idx_session_affinities_active_order on session_affinities(expires_at, created_at, response_id); +create index idx_event_logs_level_created on event_logs(level, created_at desc); +create index idx_event_logs_route_created on event_logs(route, created_at desc) where route is not null; +create index idx_event_logs_model_created on event_logs(model, created_at desc) where model is not null; +create index idx_event_logs_status_created on event_logs(status_code, created_at desc) where status_code is not null; +create index idx_event_logs_upstream_status_created on event_logs(upstream_status_code, created_at desc) where upstream_status_code is not null; +``` + +事件日志继续保持诊断窗口定位，不引入 FTS。`message like '%...%'` 和 `metadata_json like '%...%'` 不适合普通 BTree；当前不把全文搜索定义为核心能力。 + +### 4. 长期统计不应该依赖事件日志 + +状态：已处理。 + +账号模型用量已经从事件日志分页回放中剥离。请求调度成功会记录模型请求数，请求完成会记录模型维度 token 用量，空响应会记录模型维度错误数。账号列表的模型统计现在读取 `account_model_usage`，事件日志容量裁剪不会影响该统计。 + +概览卡片、24 小时趋势与 7 天健康时间线不再从 `event_logs` 反推。事件写入时会同步写入 `usage_time_buckets`，Dashboard 读取聚合事实表；清空或裁剪事件日志不会影响趋势统计。 + +当前使用独立统计表： + +```sql +create table account_model_usage ( + account_id text not null references accounts(id) on delete cascade, + model text not null, + request_count integer not null default 0 check (request_count >= 0), + error_count integer not null default 0 check (error_count >= 0), + input_tokens integer not null default 0 check (input_tokens >= 0), + output_tokens integer not null default 0 check (output_tokens >= 0), + cached_tokens integer not null default 0 check (cached_tokens >= 0), + last_used_at text, + primary key (account_id, model) +); + +create index idx_account_model_usage_last_used on account_model_usage(last_used_at desc, account_id, model); +``` + +当前使用 15 分钟时间桶，兼顾 24 小时趋势、7 天健康时间线和按模型成本计算： + +```sql +create table usage_time_buckets ( + bucket_start text not null, + account_id text not null default '', + model text not null default '', + service_tier text not null default '', + request_count integer not null default 0 check (request_count >= 0), + error_count integer not null default 0 check (error_count >= 0), + input_tokens integer not null default 0 check (input_tokens >= 0), + output_tokens integer not null default 0 check (output_tokens >= 0), + cached_tokens integer not null default 0 check (cached_tokens >= 0), + first_token_latency_sum integer not null default 0 check (first_token_latency_sum >= 0), + first_token_latency_count integer not null default 0 check (first_token_latency_count >= 0), + latency_sum integer not null default 0 check (latency_sum >= 0), + latency_count integer not null default 0 check (latency_count >= 0), + max_latency_ms integer not null default 0 check (max_latency_ms >= 0), + min_latency_ms integer not null default 0 check (min_latency_ms >= 0), + updated_at text not null, + primary key (bucket_start, account_id, model, service_tier) +); +``` + +## 表级审计 + +### `admin_users` + +设计合理。当前规模极小，按 `created_at asc, id asc` 找首个管理员不需要额外优化。若未来支持多管理员列表页，再补 `created_at` 索引。 + +### `admin_sessions` + +设计合理。`user_id` 外键级联删除、`expires_at` 索引用于会话校验和清理。保持。 + +### `client_api_keys` + +整体合理。`prefix` 部分索引用于启用状态下的认证查找，方向正确。 + +已补 `created_at desc, id desc` 列表索引。当前明文 `key` 与 `key_hash` 同存是产品取舍：如果管理端需要复制完整 key，这个结构一致；但接口输出要严格控制，避免列表接口无意泄漏。 + +### `accounts` + +核心字段基本合理，已经没有旧的 token 加密字段。`chatgpt_account_id` 与 `chatgpt_user_id` 的唯一表达式索引是正确方向。 + +已补 `accounts(added_at desc, id desc)`，用于账号列表和账号池恢复排序。如果状态筛选加排序成为常用查询，再补 `(status, added_at desc, id desc)`，当前不提前扩展。 + +`quota_json` 当前可以接受，因为上游 quota 形态不稳定；但如果页面要按余额、额度、冷却时间筛选排序，应把关键字段提升为列，不要用 JSON 做业务查询。 + +### `account_refresh_leases` + +设计合理。以 `account_id` 为主键，`expires_at` 索引用于过期释放。保持。 + +### `account_usage` + +作为账号级累计聚合表是合理的。非负 check 覆盖较完整。 + +它承担账号列表用量展示和用量分页，当前已经补齐 `last_used_at desc, account_id desc` 索引。边界是它只保存账号总量，不保存模型维度和时间趋势。 + +### `account_model_usage` + +用于账号 + 模型维度的长期聚合。它补齐了 `account_usage` 缺少的模型维度，也避免从事件日志回放业务统计。 + +当前记录： + +- 请求数。 +- 错误数。 +- input tokens。 +- output tokens。 +- cached tokens。 +- 最近使用时间。 + +这张表是长期事实表，事件日志是诊断窗口，两者不再混用。趋势图不扩展这张表为时序明细，而是由 `usage_time_buckets` 承担。 + +### `usage_time_buckets` + +用于保存 15 分钟粒度的请求趋势聚合。它不是事件日志副本，不保存请求内容、响应内容或长尾元数据，只保存 Dashboard 和长期统计需要的数值： + +- 请求数与错误数。 +- input/output/cached tokens。 +- 首 token 延迟与完成延迟的 sum/count/max/min。 +- model 与 service tier，用于成本计算。 + +写入路径在事件日志落库时同步更新聚合表。读取路径在 Dashboard 使用该表生成卡片、趋势和健康时间线；`event_logs` 清空或容量裁剪不影响这些趋势数据。 + +### `model_account_routes` + +用于保存模型到账号 ID 的显式调度关系。它是运行配置的一部分，但不能压进 `runtime_settings.model_aliases_json`： + +- `model_aliases_json` 只表达客户端模型名到真实上游模型名的别名。 +- `model_account_routes` 表达真实调度约束，受账号生命周期影响。 +- `account_id` 外键级联删除，避免账号删除后留下无效路由。 +- `priority` 保存账号选择顺序，`enabled` 用于保留可控开关。 + +运行时账号池会先按模型账号路由过滤候选账号，再执行计划 allowlist、套餐优先级、并发槽位与轮换策略。 + +### `account_cookies` + +按账号、域名、名称、路径唯一是合理的。当前已经统一为本地明文可运维策略，字段为 `value`。 + +已删除 `value_cipher`、Cookie 加解密路径、`master_key_file` 配置和不再使用的加密模块。Cookie、token、API Key 的本地持久化策略不再混用。`expires_at` 清理索引已补齐。 + +### `fingerprints` 与 `fingerprint_update_history` + +当前把默认 headers 和 header order 存 JSON 是合理的，因为它们是配置形态，不适合拆成多张表。 + +已补 `fingerprint_update_history(created_at desc, id desc)`，匹配读取最新历史记录的排序。 + +### `event_logs` + +这是当前 schema 里设计最完整的一块：热字段结构化，长尾信息放 `metadata_json`，基础索引覆盖最新列表、kind、request、account、transport、failure、response、upstream request。 + +边界要明确：事件日志是诊断窗口，不是长期统计账本。容量裁剪后，任何从日志反推的统计都只是窗口统计。 + +事件日志页面的结构化筛选索引已补齐。全文搜索暂不引入 FTS，避免把日志系统做成半个分析数据库。Dashboard 不再依赖事件日志计算长期趋势。 + +### `model_plan_snapshots` + +按 plan 存模型列表 JSON 是合理的。模型目录来自上游，结构变化概率高，当前不需要拆表。 + +按模型控制账号路由已经由 `model_account_routes` 承担，不改造这个快照表承担调度职责。 + +### `session_affinities` + +设计基本合理。它是运行时恢复和会话亲和状态，不是长期数据。`conversation_id` 和 `expires_at` 索引覆盖主要查询。 + +已补 `expires_at, created_at, response_id` 复合索引，匹配活跃亲和恢复和过期清理的排序路径。 + +## 目标数据库边界 + +建议把数据分成三类。 + +启动配置留在 `config.yaml`： + +- server。 +- database。 +- API Key pepper 文件路径。 +- API base URL。 +- OAuth endpoint。 +- TLS。 +- 日志目录。 +- 首次管理员默认值。 + +运行配置进入数据库： + +- default model。 +- model aliases。 +- model-account routes。 +- refresh margin。 +- refresh concurrency。 +- per-account concurrency。 +- request interval。 +- rotation strategy。 +后台默认策略不进入运行配置表，避免把不需要人工决策的开关暴露成持久配置。 + +运行事实进入数据库： + +- accounts。 +- account usage。 +- account model usage。 +- usage time buckets。 +- sessions。 +- API keys。 +- event logs。 +- cookies。 +- fingerprints。 +- model snapshots。 +- session affinities。 + +这个边界更符合网关项目：配置文件负责启动，数据库负责后台可管理状态。 + +## 推荐落地顺序 + +### P0：数据库基础收敛 + +1. 已删除空壳 `schema_migrations`。 +2. 已新增 `runtime_settings` 表，把后台系统设置从 YAML 写回改为数据库持久化。 +3. 已让 `config.yaml` 只保留启动级配置。 +4. 已将后台设置接口改为读写数据库，启动时从 DB 合成运行时配置。 +5. 已补齐账号、用量、API Key、指纹历史的复合索引。 +6. 已新增 `model_account_routes`，模型账号绑定进入关系表并参与调度。 +7. 已更新 `tests/infra/storage_schema`，覆盖新表、索引、约束。 + +### P1：统计模型收敛 + +状态：已完成。 + +1. 已新增 `account_model_usage`。 +2. 已在账号调度成功时记录模型请求数。 +3. 已在 Responses、Chat、Compact 路径完成时记录模型 token 用量。 +4. 已在空响应路径记录模型错误数。 +5. 账号列表模型统计已改为读取 `account_model_usage`。 +6. 已明确事件日志只作为诊断窗口。 +7. 已新增 `usage_time_buckets`，Dashboard 趋势不再依赖事件日志容量。 + +### P2：日志与搜索能力决策 + +状态：已完成。 + +1. 已保留轻量事件日志表，不引入 FTS。 +2. 已补事件日志结构化筛选索引。 +3. 已补 Cookie 过期清理索引。 +4. 已补活跃 session affinity 排序索引。 + +### P3：敏感字段策略统一 + +状态：已完成。 + +1. 已明确本地数据库不单独加密 Cookie。 +2. 已删除 `value_cipher` 和 Cookie 加解密路径。 +3. 已删除 `master_key_file` 启动配置与不再使用的加密模块。 +4. 已补 schema 与账号 Cookie 行为测试，固定 Cookie 明文持久化策略。 + +### P4：模型账号路由收敛 + +状态：已完成。 + +1. 已新增 `model_account_routes`。 +2. 已让后台设置接口读写 `modelAccountRoutes`。 +3. 已让运行时账号池按模型账号路由过滤候选账号。 +4. 已补账号删除级联测试，账号删除后模型路由同步清理。 + +### P5：趋势聚合收敛 + +状态：已完成。 + +1. 已新增 `usage_time_buckets`。 +2. 已在事件日志写入时同步更新时间桶聚合。 +3. 已将 Dashboard 卡片、趋势和健康时间线切换到聚合表。 +4. 已补测试验证清空事件日志后趋势统计仍然保留。 + +## 测试覆盖 + +本轮数据库收敛对应的测试覆盖： + +1. schema 创建测试：所有目标表存在，空壳表不存在。 +2. 索引测试：关键列表查询需要的索引存在。 +3. 约束测试：布尔、状态、非负计数、单行设置表 id 约束。 +4. 设置持久化测试：后台设置写入 DB，重启服务后仍生效。 +5. YAML 边界测试：后台设置更新不再重写整份 `config.yaml`。 +6. 已补统计测试：模型用量来自 `account_model_usage`，不依赖事件日志容量。 +7. 已补删除级联测试：账号删除后相关 usage、model usage、cookies、leases、session affinities 正确清理。事件日志按诊断历史保留，不参与账号外键级联。 +8. 已补模型账号路由测试：schema、索引、约束、设置持久化、重启恢复、运行时调度过滤和账号删除级联。 +9. 已补时间桶趋势测试：schema、索引、约束、日志写入同步聚合，以及清空事件日志后 Dashboard 仍读取聚合趋势。 + +## 最终建议 + +当前数据库设计已经完成第一轮“无历史包袱”收敛。当前边界已经更清晰：启动配置归配置文件，后台运行状态归数据库，长期统计归统计表，事件日志归诊断窗口。 diff --git a/src/admin/accounts/routes.rs b/src/admin/accounts/routes.rs index f83d6004a..fab314c3e 100644 --- a/src/admin/accounts/routes.rs +++ b/src/admin/accounts/routes.rs @@ -24,12 +24,7 @@ use crate::{ }, admin::{ accounts::service::{AdminAccountError, AdminAccountMetadata, AdminAccountMetadataUpdate}, - monitoring::{ - billing, - event_store::AdminLogFilter, - events::{EventLevel, EventLog}, - service::AdminUsageRecord, - }, + monitoring::{billing, service::AdminUsageRecord}, }, http::middleware::request_id::RequestId, infra::{ @@ -38,7 +33,7 @@ use crate::{ }, runtime::state::AppState, upstream::accounts::model::AccountStatus, - upstream::accounts::store::StoredAccount, + upstream::accounts::store::{AccountModelUsageRecord, StoredAccount}, }; const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200; @@ -1022,7 +1017,13 @@ async fn account_list_stats(state: &AppState) -> AccountListStats { .list_quota_snapshots() .await .unwrap_or_default(); - let logs = list_all_event_logs(state).await; + let model_usage_records = state + .services + .background_tasks + .accounts + .list_model_usage() + .await + .unwrap_or_default(); AccountListStats { usage_by_account: usage_records @@ -1038,7 +1039,7 @@ async fn account_list_stats(state: &AppState) -> AccountListStats { ) }) .collect(), - models_by_account: models_by_account(&logs), + models_by_account: models_by_account(model_usage_records), } } @@ -1129,120 +1130,67 @@ async fn list_all_usage_records(state: &AppState) -> Vec { } } -async fn list_all_event_logs(state: &AppState) -> Vec { - let mut cursor = None; - let mut logs = Vec::new(); - loop { - let Ok(page) = state - .services - .logs - .list(cursor, ACCOUNT_STATS_PAGE_LIMIT, AdminLogFilter::default()) - .await - else { - return Vec::new(); - }; - logs.extend(page.items); - let Some(next_cursor) = page.next_cursor else { - return logs; - }; - cursor = Some(next_cursor); - } -} - -#[derive(Debug, Clone, Default)] -struct ModelUsageAggregate { - model: String, - request_count: u64, - error_count: u64, - input_tokens: u64, - output_tokens: u64, - cached_tokens: u64, - last_used_at: Option>, -} - -fn models_by_account(logs: &[EventLog]) -> HashMap> { - let mut by_account_model = HashMap::<(String, String), ModelUsageAggregate>::new(); - for log in logs { - let Some(account_id) = log - .account_id - .as_ref() - .filter(|value| !value.trim().is_empty()) - else { - continue; - }; - let Some(model) = log.model.as_ref().filter(|value| !value.trim().is_empty()) else { - continue; - }; - let aggregate = by_account_model - .entry((account_id.clone(), model.clone())) - .or_insert_with(|| ModelUsageAggregate { - model: model.clone(), - ..ModelUsageAggregate::default() - }); - aggregate.request_count += 1; - if log.level == EventLevel::Error || log.status_code.is_some_and(|status| status >= 400) { - aggregate.error_count += 1; - } - aggregate.input_tokens += metadata_usage_number(&log.metadata, "inputTokens"); - aggregate.output_tokens += metadata_usage_number(&log.metadata, "outputTokens"); - aggregate.cached_tokens += metadata_usage_number(&log.metadata, "cachedTokens"); - aggregate.last_used_at = Some( - aggregate - .last_used_at - .map_or(log.created_at, |last| last.max(log.created_at)), - ); - } - - let mut by_account = HashMap::>::new(); - for ((account_id, _), aggregate) in by_account_model { - by_account.entry(account_id).or_default().push(aggregate); +fn models_by_account( + records: Vec, +) -> HashMap> { + let mut by_account = HashMap::>::new(); + for record in records { + by_account + .entry(record.account_id.clone()) + .or_default() + .push(record); } by_account .into_iter() - .map(|(account_id, mut aggregates)| { - aggregates.sort_by(|a, b| { + .map(|(account_id, mut records)| { + records.sort_by(|a, b| { b.request_count .cmp(&a.request_count) .then_with(|| b.last_used_at.cmp(&a.last_used_at)) + .then_with(|| a.model.cmp(&b.model)) }); ( account_id, - aggregates.into_iter().map(model_usage_data).collect(), + records.into_iter().map(model_usage_data).collect(), ) }) .collect() } -fn model_usage_data(usage: ModelUsageAggregate) -> AdminAccountModelUsageData { - let total_tokens = usage.input_tokens + usage.output_tokens; - let success_rate = if usage.request_count > 0 { - ((usage.request_count - usage.error_count) as f64 / usage.request_count as f64 * 1000.0) - .round() +fn model_usage_data(usage: AccountModelUsageRecord) -> AdminAccountModelUsageData { + let request_count = nonnegative_i64_to_u64(usage.request_count); + let error_count = nonnegative_i64_to_u64(usage.error_count); + let input_tokens = nonnegative_i64_to_u64(usage.input_tokens); + let output_tokens = nonnegative_i64_to_u64(usage.output_tokens); + let cached_tokens = nonnegative_i64_to_u64(usage.cached_tokens); + let total_tokens = input_tokens + output_tokens; + let success_rate = if request_count > 0 { + ((request_count.saturating_sub(error_count)) as f64 / request_count as f64 * 1000.0).round() / 10.0 } else { 0.0 }; let total_cost_usd = billing::calculate_cost( - usage.input_tokens, - usage.output_tokens, - usage.cached_tokens, + input_tokens, + output_tokens, + cached_tokens, &usage.model, None, ); AdminAccountModelUsageData { model: usage.model, - request_count: usage.request_count, - request_count_display: format_count(usage.request_count), + request_count, + request_count_display: format_count(request_count), success_rate, success_rate_display: format_percent(success_rate), - input_tokens: usage.input_tokens, - input_tokens_display: format_tokens(usage.input_tokens), - output_tokens: usage.output_tokens, - output_tokens_display: format_tokens(usage.output_tokens), - cached_tokens: usage.cached_tokens, - cached_tokens_display: format_tokens(usage.cached_tokens), + input_tokens, + input_tokens_display: format_tokens(input_tokens), + output_tokens, + output_tokens_display: format_tokens(output_tokens), + cached_tokens, + cached_tokens_display: format_tokens(cached_tokens), total_tokens, total_tokens_display: format_tokens(total_tokens), total_cost_usd, @@ -1552,15 +1500,6 @@ fn format_duration_days(seconds: u64) -> String { } } -fn metadata_usage_number(metadata: &Value, field: &str) -> u64 { - metadata - .get("usage") - .and_then(|usage| usage.get(field)) - .or_else(|| metadata.get(field)) - .and_then(Value::as_u64) - .unwrap_or(0) -} - fn number_value(value: &Value) -> Option { value .as_f64() diff --git a/src/admin/monitoring/dashboard.rs b/src/admin/monitoring/dashboard.rs index 1c1afcbbc..6eee91ab6 100644 --- a/src/admin/monitoring/dashboard.rs +++ b/src/admin/monitoring/dashboard.rs @@ -22,7 +22,7 @@ use crate::{ }, event_store::AdminLogFilter, events::{EventLevel, EventLog}, - service::{AdminUsageRecord, AdminUsageSummary}, + service::{AdminUsageRecord, AdminUsageSummary, AdminUsageTimeBucketRecord}, }, response::{AdminEnvelope, AdminError, AdminResponse}, }, @@ -276,9 +276,11 @@ pub async fn dashboard_summary( .map(|page| page.items) .unwrap_or_default(); + let now = Utc::now(); + let time_buckets = dashboard_time_buckets(&state, now).await; let pool_summary = account_pool_summary(&accounts); let dashboard_logs = logs.iter().take(10).map(dashboard_event_log_data).collect(); - let trend = dashboard_trend_data(&logs, DashboardTrendKind::Usage); + let trend = dashboard_trend_data(&time_buckets, DashboardTrendKind::Usage); let quota_used_by_account = account_quota_used_percent_by_id(&state, &usage_records).await; let settings = state.services.settings.current(); @@ -289,12 +291,12 @@ pub async fn dashboard_summary( cards: dashboard_cards( &accounts, &summary, - &logs, + &time_buckets, &settings.model.default_model, settings.model.service_tier.as_deref(), ), trend, - health_timeline: dashboard_health_timeline_data(&logs), + health_timeline: dashboard_health_timeline_data(&time_buckets), account_usage: account_usage_data( &accounts, &usage_records, @@ -320,18 +322,12 @@ pub async fn dashboard_trend( ) -> Result { let request_id = request_id.as_str().to_string(); require_admin_session(&state, &headers, &request_id).await?; - let logs = state - .services - .logs - .list(None, DASHBOARD_LOG_LIMIT, AdminLogFilter::default()) - .await - .map(|page| page.items) - .unwrap_or_default(); + let time_buckets = dashboard_time_buckets(&state, Utc::now()).await; Ok(AdminResponse::new( StatusCode::OK, AdminEnvelope::ok( - dashboard_trend_data(&logs, query.kind.unwrap_or_default()), + dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default()), request_id, ), )) @@ -340,14 +336,14 @@ pub async fn dashboard_trend( fn dashboard_cards( accounts: &[Account], summary: &AdminUsageSummary, - logs: &[EventLog], + buckets: &[AdminUsageTimeBucketRecord], default_model: &str, default_service_tier: Option<&str>, ) -> DashboardCardsData { dashboard_cards_at( accounts, summary, - logs, + buckets, default_model, default_service_tier, Utc::now(), @@ -357,7 +353,7 @@ fn dashboard_cards( fn dashboard_cards_at( accounts: &[Account], summary: &AdminUsageSummary, - logs: &[EventLog], + buckets: &[AdminUsageTimeBucketRecord], default_model: &str, default_service_tier: Option<&str>, now: DateTime, @@ -369,10 +365,10 @@ fn dashboard_cards_at( default_service_tier, }; - let today = usage_window(logs, today_start, now); - let yesterday = usage_window(logs, yesterday_start, today_start); - let today_cost = cost_window(logs, today_start, now, billing_context).unwrap_or(0.0); - let total_cost = total_cost(logs, summary, billing_context); + let today = usage_window(buckets, today_start, now); + let yesterday = usage_window(buckets, yesterday_start, today_start); + let today_cost = cost_window(buckets, today_start, now, billing_context).unwrap_or(0.0); + let total_cost = total_cost(buckets, summary, billing_context); let total_input = nonnegative_i64_to_u64(summary.input_tokens); let total_cached = nonnegative_i64_to_u64(summary.cached_tokens); @@ -426,12 +422,30 @@ fn dashboard_cards_at( } } -fn dashboard_trend_data(logs: &[EventLog], kind: DashboardTrendKind) -> DashboardTrendData { - dashboard_trend_data_at(logs, kind, Utc::now()) +async fn dashboard_time_buckets( + state: &AppState, + now: DateTime, +) -> Vec { + let current_slot = china_quarter_hour_start(now); + let start = current_slot + - Duration::minutes(HEALTH_TIMELINE_SLOT_MINUTES * (HEALTH_TIMELINE_SLOTS - 1)); + state + .services + .usage + .time_buckets(start, now) + .await + .unwrap_or_default() +} + +fn dashboard_trend_data( + buckets: &[AdminUsageTimeBucketRecord], + kind: DashboardTrendKind, +) -> DashboardTrendData { + dashboard_trend_data_at(buckets, kind, Utc::now()) } fn dashboard_trend_data_at( - logs: &[EventLog], + records: &[AdminUsageTimeBucketRecord], kind: DashboardTrendKind, now: DateTime, ) -> DashboardTrendData { @@ -444,16 +458,16 @@ fn dashboard_trend_data_at( }) .collect::>(); - for log in logs { - if log.created_at < start || log.created_at > now { + for record in records { + if record.bucket_start < start || record.bucket_start > now { continue; } - let log_hour = china_hour_start(log.created_at); + let log_hour = china_hour_start(record.bucket_start); if let Some((_, bucket)) = buckets .iter_mut() .find(|(bucket_start, _)| *bucket_start == log_hour) { - apply_log(bucket, log); + apply_bucket(bucket, record); } } @@ -490,12 +504,14 @@ fn dashboard_trend_data_at( } } -fn dashboard_health_timeline_data(logs: &[EventLog]) -> DashboardHealthTimelineData { - dashboard_health_timeline_data_at(logs, Utc::now()) +fn dashboard_health_timeline_data( + buckets: &[AdminUsageTimeBucketRecord], +) -> DashboardHealthTimelineData { + dashboard_health_timeline_data_at(buckets, Utc::now()) } fn dashboard_health_timeline_data_at( - logs: &[EventLog], + records: &[AdminUsageTimeBucketRecord], now: DateTime, ) -> DashboardHealthTimelineData { let current_slot = china_quarter_hour_start(now); @@ -508,16 +524,16 @@ fn dashboard_health_timeline_data_at( }) .collect::>(); - for log in logs { - if log.created_at < start || log.created_at > now { + for record in records { + if record.bucket_start < start || record.bucket_start > now { continue; } - let log_slot = china_quarter_hour_start(log.created_at); + let log_slot = china_quarter_hour_start(record.bucket_start); if let Some((_, bucket)) = buckets .iter_mut() .find(|(bucket_start, _)| *bucket_start == log_slot) { - apply_log(bucket, log); + apply_bucket(bucket, record); } } @@ -659,27 +675,31 @@ fn trend_summary( } } -fn usage_window(logs: &[EventLog], start: DateTime, end: DateTime) -> UsageWindow { +fn usage_window( + records: &[AdminUsageTimeBucketRecord], + start: DateTime, + end: DateTime, +) -> UsageWindow { let mut window = UsageWindow::default(); - for log in logs { - if log.created_at >= start && log.created_at < end { - apply_log(&mut window, log); + for record in records { + if record.bucket_start >= start && record.bucket_start < end { + apply_bucket(&mut window, record); } } window } fn cost_window( - logs: &[EventLog], + records: &[AdminUsageTimeBucketRecord], start: DateTime, end: DateTime, context: BillingContext<'_>, ) -> Option { let mut total = 0.0; let mut has_usage = false; - for log in logs { - if log.created_at >= start && log.created_at < end { - if let Some(cost) = log_cost(log, context) { + for record in records { + if record.bucket_start >= start && record.bucket_start < end { + if let Some(cost) = bucket_cost(record, context) { total += cost; has_usage = true; } @@ -688,11 +708,15 @@ fn cost_window( has_usage.then_some(total) } -fn total_cost(logs: &[EventLog], summary: &AdminUsageSummary, context: BillingContext<'_>) -> f64 { +fn total_cost( + records: &[AdminUsageTimeBucketRecord], + summary: &AdminUsageSummary, + context: BillingContext<'_>, +) -> f64 { let mut total = 0.0; let mut has_usage = false; - for log in logs { - if let Some(cost) = log_cost(log, context) { + for record in records { + if let Some(cost) = bucket_cost(record, context) { total += cost; has_usage = true; } @@ -713,21 +737,26 @@ fn total_cost(logs: &[EventLog], summary: &AdminUsageSummary, context: BillingCo ) } -fn log_cost(log: &EventLog, context: BillingContext<'_>) -> Option { - let input_tokens = metadata_usage_number(&log.metadata, "inputTokens"); - let output_tokens = metadata_usage_number(&log.metadata, "outputTokens"); - let cached_tokens = metadata_usage_number(&log.metadata, "cachedTokens"); +fn bucket_cost(record: &AdminUsageTimeBucketRecord, context: BillingContext<'_>) -> Option { + let input_tokens = nonnegative_i64_to_u64(record.input_tokens); + let output_tokens = nonnegative_i64_to_u64(record.output_tokens); + let cached_tokens = nonnegative_i64_to_u64(record.cached_tokens); if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 { return None; } - let model = log - .model + let model = { + let value = record.model.trim(); + if value.is_empty() { + context.default_model + } else { + value + } + }; + let service_tier = record + .service_tier .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .unwrap_or(context.default_model); - let service_tier = metadata_service_tier(&log.metadata).or(context.default_service_tier); + .or(context.default_service_tier); Some(billing::calculate_cost( input_tokens, @@ -738,84 +767,30 @@ fn log_cost(log: &EventLog, context: BillingContext<'_>) -> Option { )) } -fn apply_log(window: &mut UsageWindow, log: &EventLog) { - window.requests += 1; - window.input_tokens += metadata_usage_number(&log.metadata, "inputTokens"); - window.output_tokens += metadata_usage_number(&log.metadata, "outputTokens"); - window.cached_tokens += metadata_usage_number(&log.metadata, "cachedTokens"); - if let Some(first_token_latency) = metadata_first_token_latency(&log.metadata) { - window.first_token_latency_sum += first_token_latency; - window.first_token_latency_count += 1; +fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) { + window.requests += nonnegative_i64_to_u64(record.request_count); + window.input_tokens += nonnegative_i64_to_u64(record.input_tokens); + window.output_tokens += nonnegative_i64_to_u64(record.output_tokens); + window.cached_tokens += nonnegative_i64_to_u64(record.cached_tokens); + window.errors += nonnegative_i64_to_u64(record.error_count); + window.first_token_latency_sum += nonnegative_i64_to_u64(record.first_token_latency_sum); + window.first_token_latency_count += nonnegative_i64_to_u64(record.first_token_latency_count); + window.latency_sum += nonnegative_i64_to_u64(record.latency_sum); + window.latency_count += nonnegative_i64_to_u64(record.latency_count); + let max_latency = nonnegative_i64_to_u64(record.max_latency_ms); + if max_latency > 0 { + window.max_latency = window.max_latency.max(max_latency); } - if is_error_log(log) { - window.errors += 1; - } - if let Some(latency) = log - .latency_ms - .filter(|value| *value > 0) - .map(nonnegative_i64_to_u64) - { - window.latency_sum += latency; - window.latency_count += 1; - window.max_latency = window.max_latency.max(latency); + let min_latency = nonnegative_i64_to_u64(record.min_latency_ms); + if min_latency > 0 { window.min_latency = if window.min_latency == 0 { - latency + min_latency } else { - window.min_latency.min(latency) + window.min_latency.min(min_latency) }; } } -fn metadata_first_token_latency(metadata: &Value) -> Option { - [ - "firstTokenMs", - "first_token_ms", - "firstTokenLatencyMs", - "first_token_latency_ms", - ] - .into_iter() - .find_map(|field| { - metadata - .get(field) - .or_else(|| metadata.get("usage").and_then(|usage| usage.get(field))) - .and_then(Value::as_u64) - .filter(|value| *value > 0) - }) -} - -fn metadata_usage_number(metadata: &Value, field: &str) -> u64 { - metadata - .get("usage") - .and_then(|usage| usage.get(field)) - .or_else(|| metadata.get(field)) - .and_then(Value::as_u64) - .unwrap_or(0) -} - -fn metadata_service_tier(metadata: &Value) -> Option<&str> { - let direct = metadata_service_tier_value(metadata); - direct.or_else(|| metadata.get("usage").and_then(metadata_service_tier_value)) -} - -fn metadata_service_tier_value(value: &Value) -> Option<&str> { - [ - "billingServiceTier", - "billing_service_tier", - "actualServiceTier", - "actual_service_tier", - "serviceTier", - "service_tier", - ] - .into_iter() - .find_map(|field| value.get(field).and_then(Value::as_str)) - .map(str::trim) - .filter(|value| !value.is_empty()) -} - -fn is_error_log(log: &EventLog) -> bool { - log.level == EventLevel::Error || log.status_code.is_some_and(|status| status >= 400) -} - fn account_pool_summary(accounts: &[Account]) -> AccountPoolDiagnostics { let mut summary = AccountPoolDiagnostics { total: accounts.len(), @@ -1046,12 +1021,12 @@ mod tests { #[test] fn dashboard_cards_should_use_china_day_boundary() { let now = utc("2026-06-24T18:00:00Z"); - let logs = vec![ - usage_log_at("2026-06-24T16:30:00Z", 10), - usage_log_at("2026-06-24T15:30:00Z", 20), + let buckets = vec![ + usage_bucket_at("2026-06-24T16:30:00Z", 10), + usage_bucket_at("2026-06-24T15:30:00Z", 20), ]; - let cards = dashboard_cards_at(&[], &empty_usage_summary(), &logs, "gpt-5.5", None, now); + let cards = dashboard_cards_at(&[], &empty_usage_summary(), &buckets, "gpt-5.5", None, now); assert_eq!(cards.traffic.today_requests, 1); assert_eq!(cards.traffic.yesterday_requests, 1); @@ -1062,27 +1037,32 @@ mod tests { #[test] fn dashboard_trend_should_bucket_and_label_by_china_hour() { let now = utc("2026-06-24T18:37:00Z"); - let logs = vec![usage_log_at("2026-06-24T18:15:00Z", 10)]; + let buckets = vec![usage_bucket_at("2026-06-24T18:15:00Z", 10)]; - let trend = dashboard_trend_data_at(&logs, DashboardTrendKind::Usage, now); + let trend = dashboard_trend_data_at(&buckets, DashboardTrendKind::Usage, now); let last = trend.points.last().expect("trend should contain points"); assert_eq!(last.time, "02"); assert_eq!(last.requests, 1); } - fn usage_log_at(created_at: &str, input_tokens: u64) -> EventLog { - let mut log = EventLog::new("v1.response", EventLevel::Info, "completed"); - log.created_at = utc(created_at); - log.model = Some("gpt-5.5".to_string()); - log.metadata = serde_json::json!({ - "usage": { - "inputTokens": input_tokens, - "outputTokens": 0u64, - "cachedTokens": 0u64 - } - }); - log + fn usage_bucket_at(created_at: &str, input_tokens: i64) -> AdminUsageTimeBucketRecord { + AdminUsageTimeBucketRecord { + bucket_start: utc(created_at), + model: "gpt-5.5".to_string(), + service_tier: None, + request_count: 1, + error_count: 0, + input_tokens, + output_tokens: 0, + cached_tokens: 0, + first_token_latency_sum: 0, + first_token_latency_count: 0, + latency_sum: 0, + latency_count: 0, + max_latency_ms: 0, + min_latency_ms: 0, + } } fn utc(value: &str) -> DateTime { diff --git a/src/admin/monitoring/event_store.rs b/src/admin/monitoring/event_store.rs index 85dbd8af5..3613533a4 100644 --- a/src/admin/monitoring/event_store.rs +++ b/src/admin/monitoring/event_store.rs @@ -10,6 +10,7 @@ use crate::admin::monitoring::events::{ EventLevel, EventLog, EventLogStore, EventLogStoreError, EventLogStoreResult, }; use crate::infra::json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page}; +use crate::infra::time::china_quarter_hour_start; /// SQLite 事件日志错误。 #[derive(Debug, Error)] @@ -287,9 +288,134 @@ async fn append_event(pool: &SqlitePool, event: &EventLog) -> SqliteEventLogStor .bind(event.created_at.to_rfc3339()) .execute(pool) .await?; + append_usage_time_bucket(pool, event).await?; Ok(()) } +async fn append_usage_time_bucket( + pool: &SqlitePool, + event: &EventLog, +) -> SqliteEventLogStoreResult<()> { + let bucket_start = china_quarter_hour_start(event.created_at); + let error_count = i64::from(is_error_event(event)); + let input_tokens = metadata_usage_i64(&event.metadata, "inputTokens"); + let output_tokens = metadata_usage_i64(&event.metadata, "outputTokens"); + let cached_tokens = metadata_usage_i64(&event.metadata, "cachedTokens"); + let first_token_latency = metadata_first_token_latency(&event.metadata).unwrap_or(0); + let first_token_latency_count = i64::from(first_token_latency > 0); + let latency = event.latency_ms.filter(|value| *value > 0).unwrap_or(0); + let latency_count = i64::from(latency > 0); + let now = Utc::now().to_rfc3339(); + + sqlx::query( + r" +insert into usage_time_buckets ( + bucket_start, + account_id, + model, + service_tier, + request_count, + error_count, + input_tokens, + output_tokens, + cached_tokens, + first_token_latency_sum, + first_token_latency_count, + latency_sum, + latency_count, + max_latency_ms, + min_latency_ms, + updated_at +) values (?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +on conflict(bucket_start, account_id, model, service_tier) do update set + request_count = request_count + excluded.request_count, + error_count = error_count + excluded.error_count, + input_tokens = input_tokens + excluded.input_tokens, + output_tokens = output_tokens + excluded.output_tokens, + cached_tokens = cached_tokens + excluded.cached_tokens, + first_token_latency_sum = first_token_latency_sum + excluded.first_token_latency_sum, + first_token_latency_count = first_token_latency_count + excluded.first_token_latency_count, + latency_sum = latency_sum + excluded.latency_sum, + latency_count = latency_count + excluded.latency_count, + max_latency_ms = max(max_latency_ms, excluded.max_latency_ms), + min_latency_ms = case + when min_latency_ms = 0 then excluded.min_latency_ms + when excluded.min_latency_ms = 0 then min_latency_ms + else min(min_latency_ms, excluded.min_latency_ms) + end, + updated_at = excluded.updated_at", + ) + .bind(bucket_start.to_rfc3339()) + .bind(event.account_id.as_deref().unwrap_or_default()) + .bind(event.model.as_deref().unwrap_or_default()) + .bind(metadata_service_tier(&event.metadata).unwrap_or_default()) + .bind(error_count) + .bind(input_tokens) + .bind(output_tokens) + .bind(cached_tokens) + .bind(first_token_latency) + .bind(first_token_latency_count) + .bind(latency) + .bind(latency_count) + .bind(latency) + .bind(latency) + .bind(now) + .execute(pool) + .await?; + Ok(()) +} + +fn is_error_event(event: &EventLog) -> bool { + event.level == EventLevel::Error || event.status_code.is_some_and(|status| status >= 400) +} + +fn metadata_usage_i64(metadata: &Value, field: &str) -> i64 { + metadata + .get("usage") + .and_then(|usage| usage.get(field)) + .or_else(|| metadata.get(field)) + .and_then(Value::as_i64) + .filter(|value| *value > 0) + .unwrap_or(0) +} + +fn metadata_first_token_latency(metadata: &Value) -> Option { + [ + "firstTokenMs", + "first_token_ms", + "firstTokenLatencyMs", + "first_token_latency_ms", + ] + .into_iter() + .find_map(|field| { + metadata + .get(field) + .or_else(|| metadata.get("usage").and_then(|usage| usage.get(field))) + .and_then(Value::as_i64) + .filter(|value| *value > 0) + }) +} + +fn metadata_service_tier(metadata: &Value) -> Option<&str> { + metadata_service_tier_value(metadata) + .or_else(|| metadata.get("usage").and_then(metadata_service_tier_value)) +} + +fn metadata_service_tier_value(value: &Value) -> Option<&str> { + [ + "billingServiceTier", + "billing_service_tier", + "actualServiceTier", + "actual_service_tier", + "serviceTier", + "service_tier", + ] + .into_iter() + .find_map(|field| value.get(field).and_then(Value::as_str)) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + async fn trim_to_capacity(pool: &SqlitePool, capacity: u32) -> SqliteEventLogStoreResult { let result = sqlx::query( "delete from event_logs where id in ( diff --git a/src/admin/monitoring/service.rs b/src/admin/monitoring/service.rs index eaab12f76..2fc18a021 100644 --- a/src/admin/monitoring/service.rs +++ b/src/admin/monitoring/service.rs @@ -4,7 +4,9 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use thiserror::Error; -use crate::admin::monitoring::usage_store::{SqliteUsageStore, UsageListRecord, UsageSummary}; +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; @@ -152,6 +154,24 @@ impl AdminUsageService { .map(AdminUsageSummary::from) .map_err(|_| AdminUsageError::Summary) } + + /// 列出指定时间范围内的聚合时间桶。 + pub async fn time_buckets( + &self, + start: DateTime, + end: DateTime, + ) -> Result, AdminUsageError> { + self.store + .list_time_buckets(start, end) + .await + .map(|records| { + records + .into_iter() + .map(AdminUsageTimeBucketRecord::from) + .collect() + }) + .map_err(|_| AdminUsageError::TimeBuckets) + } } /// 管理端用量统计错误。 @@ -161,6 +181,8 @@ pub enum AdminUsageError { List, #[error("failed to summarize account usage")] Summary, + #[error("failed to list usage time buckets")] + TimeBuckets, } /// 管理端账号用量记录。 @@ -201,6 +223,25 @@ pub struct AdminUsageSummary { pub image_request_failed_count: i64, } +/// 管理端时间桶用量记录。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AdminUsageTimeBucketRecord { + pub bucket_start: DateTime, + pub model: String, + pub service_tier: Option, + pub request_count: i64, + pub error_count: i64, + pub input_tokens: i64, + pub output_tokens: i64, + pub cached_tokens: i64, + pub first_token_latency_sum: i64, + pub first_token_latency_count: i64, + pub latency_sum: i64, + pub latency_count: i64, + pub max_latency_ms: i64, + pub min_latency_ms: i64, +} + impl From for AdminUsageSummary { fn from(s: UsageSummary) -> Self { Self { @@ -220,6 +261,27 @@ impl From for AdminUsageSummary { } } +impl From for AdminUsageTimeBucketRecord { + fn from(record: UsageTimeBucketRecord) -> Self { + Self { + bucket_start: record.bucket_start, + model: record.model, + service_tier: record.service_tier, + request_count: record.request_count, + error_count: record.error_count, + input_tokens: record.input_tokens, + output_tokens: record.output_tokens, + cached_tokens: record.cached_tokens, + first_token_latency_sum: record.first_token_latency_sum, + first_token_latency_count: record.first_token_latency_count, + latency_sum: record.latency_sum, + latency_count: record.latency_count, + max_latency_ms: record.max_latency_ms, + min_latency_ms: record.min_latency_ms, + } + } +} + impl From for AdminUsageRecord { fn from(usage: UsageListRecord) -> Self { Self { diff --git a/src/admin/monitoring/usage_store.rs b/src/admin/monitoring/usage_store.rs index e43bb46d0..6a59f2dc7 100644 --- a/src/admin/monitoring/usage_store.rs +++ b/src/admin/monitoring/usage_store.rs @@ -95,6 +95,27 @@ select coalesce(sum(image_request_failed_count), 0) as image_request_failed_count from account_usage"; +const LIST_USAGE_TIME_BUCKETS_SQL: &str = r" +select + bucket_start, + model, + service_tier, + coalesce(sum(request_count), 0) as request_count, + coalesce(sum(error_count), 0) as error_count, + coalesce(sum(input_tokens), 0) as input_tokens, + coalesce(sum(output_tokens), 0) as output_tokens, + coalesce(sum(cached_tokens), 0) as cached_tokens, + coalesce(sum(first_token_latency_sum), 0) as first_token_latency_sum, + coalesce(sum(first_token_latency_count), 0) as first_token_latency_count, + coalesce(sum(latency_sum), 0) as latency_sum, + coalesce(sum(latency_count), 0) as latency_count, + coalesce(max(max_latency_ms), 0) as max_latency_ms, + coalesce(min(nullif(min_latency_ms, 0)), 0) as min_latency_ms +from usage_time_buckets +where bucket_start >= ? and bucket_start <= ? +group by bucket_start, model, service_tier +order by bucket_start asc, model asc, service_tier asc"; + /// SQLite 用量存储错误。 #[derive(Debug, Error)] pub enum SqliteUsageStoreError { @@ -178,6 +199,39 @@ pub struct UsageSummary { pub image_request_failed_count: i64, } +/// 时间桶聚合用量记录。 +#[derive(Debug, Clone)] +pub struct UsageTimeBucketRecord { + /// 桶开始时间。 + pub bucket_start: DateTime, + /// 模型。 + pub model: String, + /// 服务层级。 + pub service_tier: Option, + /// 请求数。 + pub request_count: i64, + /// 错误数。 + pub error_count: i64, + /// 输入 token。 + pub input_tokens: i64, + /// 输出 token。 + pub output_tokens: i64, + /// 缓存 token。 + pub cached_tokens: i64, + /// 首 token 延迟总和。 + pub first_token_latency_sum: i64, + /// 首 token 延迟样本数。 + pub first_token_latency_count: i64, + /// 完成延迟总和。 + pub latency_sum: i64, + /// 完成延迟样本数。 + pub latency_count: i64, + /// 最大完成延迟。 + pub max_latency_ms: i64, + /// 最小完成延迟。 + pub min_latency_ms: i64, +} + /// SQLite 用量存储。 #[derive(Clone)] pub struct SqliteUsageStore { @@ -269,6 +323,20 @@ impl SqliteUsageStore { image_request_failed_count: row.get("image_request_failed_count"), }) } + + /// 列出指定时间范围内的时间桶聚合用量。 + pub async fn list_time_buckets( + &self, + start: DateTime, + end: DateTime, + ) -> SqliteUsageStoreResult> { + let rows = sqlx::query(LIST_USAGE_TIME_BUCKETS_SQL) + .bind(start.to_rfc3339()) + .bind(end.to_rfc3339()) + .fetch_all(&self.pool) + .await?; + rows.iter().map(usage_time_bucket_from_row).collect() + } } fn usage_list_from_row(row: &SqliteRow) -> SqliteUsageStoreResult { @@ -292,6 +360,27 @@ fn usage_list_from_row(row: &SqliteRow) -> SqliteUsageStoreResult SqliteUsageStoreResult { + let service_tier = row.get::("service_tier").trim().to_string(); + Ok(UsageTimeBucketRecord { + bucket_start: DateTime::parse_from_rfc3339(row.get::<&str, _>("bucket_start"))? + .with_timezone(&Utc), + model: row.get("model"), + service_tier: (!service_tier.is_empty()).then_some(service_tier), + request_count: row.get("request_count"), + error_count: row.get("error_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + first_token_latency_sum: row.get("first_token_latency_sum"), + first_token_latency_count: row.get("first_token_latency_count"), + latency_sum: row.get("latency_sum"), + latency_count: row.get("latency_count"), + max_latency_ms: row.get("max_latency_ms"), + min_latency_ms: row.get("min_latency_ms"), + }) +} + fn parse_optional_rfc3339(value: Option) -> SqliteUsageStoreResult>> { value .as_deref() diff --git a/src/admin/settings/routes.rs b/src/admin/settings/routes.rs index b9316822c..4e35138a5 100644 --- a/src/admin/settings/routes.rs +++ b/src/admin/settings/routes.rs @@ -9,6 +9,7 @@ use axum::{ }; use serde::Serialize; use serde_json::Value; +use std::collections::BTreeMap; use crate::{ admin::auth::session::require_admin_session, @@ -21,8 +22,12 @@ use crate::{ runtime::state::AppState, }; -const ALLOWED_SETTINGS_KEYS: [&str; 4] = [ +const ALLOWED_SETTINGS_KEYS: [&str; 8] = [ "defaultModel", + "modelAliases", + "modelAccountRoutes", + "refreshMarginSeconds", + "refreshConcurrency", "maxConcurrentPerAccount", "requestIntervalMs", "rotationStrategy", @@ -33,6 +38,10 @@ const ALLOWED_SETTINGS_KEYS: [&str; 4] = [ #[serde(rename_all = "camelCase")] pub struct AdminSettingsData { pub default_model: String, + 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, @@ -42,6 +51,10 @@ impl AdminSettingsData { fn from_config(config: &AppConfig) -> Self { Self { default_model: config.model.default_model.clone(), + model_aliases: config.model.aliases.clone(), + model_account_routes: config.model.account_routes.clone(), + refresh_margin_seconds: config.auth.refresh_margin_seconds, + refresh_concurrency: config.auth.refresh_concurrency, max_concurrent_per_account: config.auth.max_concurrent_per_account, request_interval_ms: config.auth.request_interval_ms, rotation_strategy: config.auth.rotation_strategy.clone(), @@ -139,7 +152,9 @@ fn settings_service_error(error: RuntimeSettingsError, request_id: &str) -> Admi format!("Invalid setting {field}: {message}"), request_id, ), - RuntimeSettingsError::Persist(_) => AdminError::new( + RuntimeSettingsError::Database(_) + | RuntimeSettingsError::Json(_) + | RuntimeSettingsError::StoredField { .. } => AdminError::new( StatusCode::INTERNAL_SERVER_ERROR, 50000, "Failed to persist settings", diff --git a/src/config/mod.rs b/src/config/mod.rs index c7516bee5..724c9cd3f 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -8,27 +8,6 @@ pub mod loader; pub mod settings; /// 配置类型定义。 pub mod types; -/// 配置写回。 -pub mod writeback; - -use thiserror::Error; /// 配置加载结果类型。 pub type ConfigResult = Result; - -/// 配置写入结果类型。 -pub type ConfigWriteResult = Result; - -/// 配置写入错误。 -#[derive(Debug, Error)] -pub enum ConfigWriteError { - /// 序列化配置文件失败。 - #[error("serialize config file")] - Serialize(#[source] serde_yml::Error), - /// 创建配置目录失败。 - #[error("create config directory")] - CreateDirectory(#[source] std::io::Error), - /// 写入配置文件失败。 - #[error("write config file")] - Write(#[source] std::io::Error), -} diff --git a/src/config/settings.rs b/src/config/settings.rs index 4c5d81ff0..7718f43d2 100644 --- a/src/config/settings.rs +++ b/src/config/settings.rs @@ -1,20 +1,22 @@ -//! 配置设置领域逻辑与运行时写回服务。 +//! 配置设置领域逻辑与运行时数据库服务。 use std::{ collections::BTreeMap, - path::PathBuf, sync::{Arc, RwLock as StdRwLock}, }; +use chrono::Utc; use serde::{Deserialize, Serialize}; +use sqlx::{sqlite::SqliteRow, Row, SqlitePool}; use thiserror::Error; -use crate::config::{ - types::{AppConfig, QuotaWarningThresholds}, - ConfigWriteError, +use crate::config::types::AppConfig; +use crate::upstream::accounts::pool::{ + AccountPoolOptions, RotationStrategy, RuntimeAccountPoolService, }; const ROTATION_STRATEGIES: [&str; 3] = ["least_used", "round_robin", "sticky"]; +const RUNTIME_SETTINGS_ID: i64 = 1; /// 管理端可变设置。 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -27,6 +29,8 @@ pub struct AdminSettings { pub service_tier: Option, /// 模型别名映射。 pub model_aliases: BTreeMap, + /// 模型到账号 ID 的显式路由。 + pub model_account_routes: BTreeMap>, /// 是否启用访问令牌刷新。 pub refresh_enabled: bool, /// 访问令牌过期前多少秒开始刷新。 @@ -64,6 +68,7 @@ impl Default for AdminSettings { default_reasoning_effort: None, service_tier: None, model_aliases: BTreeMap::new(), + model_account_routes: BTreeMap::new(), refresh_enabled: true, refresh_margin_seconds: 300, refresh_concurrency: 2, @@ -107,6 +112,14 @@ impl Default for AdminQuotaWarningThresholds { pub struct AdminSettingsPatch { /// 默认模型 ID。 pub default_model: Option, + /// 模型别名映射。 + pub model_aliases: Option>, + /// 模型到账号 ID 的显式路由。 + pub model_account_routes: Option>>, + /// 访问令牌过期前多少秒开始刷新。 + pub refresh_margin_seconds: Option, + /// 访问令牌刷新并发数。 + pub refresh_concurrency: Option, /// 单账号最大并发请求数。 pub max_concurrent_per_account: Option, /// 同账号请求间隔毫秒数。 @@ -128,6 +141,19 @@ impl SettingsService { if let Some(default_model) = patch.default_model { current.default_model = non_empty_string("defaultModel", default_model)?; } + if let Some(model_aliases) = patch.model_aliases { + current.model_aliases = validate_model_aliases(model_aliases)?; + } + if let Some(model_account_routes) = patch.model_account_routes { + current.model_account_routes = validate_model_account_routes(model_account_routes)?; + } + if let Some(refresh_margin_seconds) = patch.refresh_margin_seconds { + current.refresh_margin_seconds = + positive_u64("refreshMarginSeconds", refresh_margin_seconds)?; + } + if let Some(refresh_concurrency) = patch.refresh_concurrency { + current.refresh_concurrency = positive_u32("refreshConcurrency", refresh_concurrency)?; + } if let Some(max_concurrent_per_account) = patch.max_concurrent_per_account { current.max_concurrent_per_account = positive_usize("maxConcurrentPerAccount", max_concurrent_per_account)?; @@ -170,6 +196,52 @@ impl SettingsServiceError { } } +fn validate_model_aliases( + aliases: BTreeMap, +) -> 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)?; + if alias == target { + return Err(invalid_field( + "modelAliases", + "alias and target must differ", + )); + } + normalized.insert(alias, target); + } + Ok(normalized) +} + +fn validate_model_account_routes( + routes: BTreeMap>, +) -> Result>, SettingsServiceError> { + let mut normalized = BTreeMap::new(); + for (model, account_ids) in routes { + 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)?; + if normalized_account_ids.contains(&account_id) { + return Err(invalid_field( + "modelAccountRoutes", + format!("duplicate account id `{account_id}`"), + )); + } + normalized_account_ids.push(account_id); + } + if normalized_account_ids.is_empty() { + return Err(invalid_field( + "modelAccountRoutes", + "account list must not be empty", + )); + } + normalized.insert(model, normalized_account_ids); + } + Ok(normalized) +} + fn validate_rotation_strategy(strategy: String) -> Result { let strategy = non_empty_string("rotationStrategy", strategy)?; if ROTATION_STRATEGIES.contains(&strategy.as_str()) { @@ -191,6 +263,22 @@ fn non_empty_string(field: &str, value: String) -> Result Result { + if value == 0 { + Err(invalid_field(field, "must be greater than 0")) + } else { + Ok(value) + } +} + +fn positive_u32(field: &str, value: u32) -> Result { + if value == 0 { + Err(invalid_field(field, "must be greater than 0")) + } else { + Ok(value) + } +} + fn positive_usize(field: &str, value: usize) -> Result { if value == 0 { Err(invalid_field(field, "must be greater than 0")) @@ -210,20 +298,30 @@ fn invalid_field(field: &str, message: impl Into) -> SettingsServiceErro #[derive(Clone)] pub struct RuntimeSettingsService { current: Arc>>, - config_path: Arc, + pool: SqlitePool, + account_pool: Option>, } impl RuntimeSettingsService { - /// 使用默认 `config.yaml` 路径构造运行时设置服务。 - pub fn new(config: AppConfig) -> Self { - Self::with_config_path(config, "config.yaml") - } - - /// 使用指定配置路径构造运行时设置服务。 - pub fn with_config_path(config: AppConfig, config_path: impl Into) -> Self { + /// 构造运行时设置服务。 + pub fn new(config: AppConfig, pool: SqlitePool) -> Self { Self { current: Arc::new(StdRwLock::new(Arc::new(config))), - config_path: Arc::new(config_path.into()), + pool, + account_pool: None, + } + } + + /// 构造带账号池同步的运行时设置服务。 + pub fn with_account_pool( + config: AppConfig, + pool: SqlitePool, + account_pool: Arc, + ) -> Self { + Self { + current: Arc::new(StdRwLock::new(Arc::new(config))), + pool, + account_pool: Some(account_pool), } } @@ -235,12 +333,19 @@ impl RuntimeSettingsService { .clone() } - /// 返回配置写回路径。 - pub fn config_path(&self) -> Arc { - self.config_path.clone() + /// 初始化缺省运行设置，并返回数据库中的运行配置快照。 + pub async fn load_or_initialize_config( + mut config: AppConfig, + pool: &SqlitePool, + ) -> Result { + let settings = admin_settings_from_config(&config); + insert_runtime_settings_if_missing(pool, &settings).await?; + let settings = load_runtime_settings(pool).await?; + apply_admin_settings_to_config(&mut config, settings); + Ok(config) } - /// 应用设置补丁、写回配置文件并更新运行时配置快照。 + /// 应用设置补丁、写入数据库并更新运行时配置快照。 pub async fn update( &self, patch: AdminSettingsPatch, @@ -248,9 +353,16 @@ impl RuntimeSettingsService { let mut next = (*self.current()).clone(); let mut settings = admin_settings_from_config(&next); SettingsService::apply_patch(&mut settings, patch)?; + save_runtime_settings(&self.pool, &settings).await?; apply_admin_settings_to_config(&mut next, settings); - next.write_settings_config(self.config_path.as_ref()) - .await?; + if let Some(account_pool) = &self.account_pool { + account_pool + .apply_options( + account_pool_options_from_config(&next), + next.auth.request_interval_ms, + ) + .await; + } let next = Arc::new(next); *self .current @@ -266,9 +378,20 @@ pub enum RuntimeSettingsError { /// 设置字段校验失败。 #[error(transparent)] InvalidField(#[from] SettingsServiceError), - /// 配置持久化失败。 - #[error(transparent)] - Persist(#[from] ConfigWriteError), + /// 数据库操作失败。 + #[error("runtime settings database error: {0}")] + Database(#[from] sqlx::Error), + /// JSON 编解码失败。 + #[error("runtime settings json error: {0}")] + Json(#[from] serde_json::Error), + /// 数据库存储值非法。 + #[error("invalid stored setting `{field}`: {message}")] + StoredField { + /// 字段名。 + field: String, + /// 错误说明。 + message: String, + }, } fn admin_settings_from_config(config: &AppConfig) -> AdminSettings { @@ -278,6 +401,7 @@ fn admin_settings_from_config(config: &AppConfig) -> AdminSettings { default_reasoning_effort: config.model.default_reasoning_effort.clone(), service_tier: config.model.service_tier.clone(), model_aliases: config.model.aliases.clone(), + model_account_routes: config.model.account_routes.clone(), refresh_enabled: defaults.refresh_enabled, refresh_margin_seconds: config.auth.refresh_margin_seconds, refresh_concurrency: config.auth.refresh_concurrency, @@ -300,24 +424,277 @@ fn admin_settings_from_config(config: &AppConfig) -> AdminSettings { fn apply_admin_settings_to_config(config: &mut AppConfig, settings: AdminSettings) { config.model.default_model = settings.default_model; - config.model.default_reasoning_effort = settings.default_reasoning_effort; - config.model.service_tier = settings.service_tier; config.model.aliases = settings.model_aliases; - config.auth.refresh_enabled = settings.refresh_enabled; + config.model.account_routes = settings.model_account_routes; config.auth.refresh_margin_seconds = settings.refresh_margin_seconds; config.auth.refresh_concurrency = settings.refresh_concurrency; config.auth.max_concurrent_per_account = settings.max_concurrent_per_account; config.auth.request_interval_ms = settings.request_interval_ms; config.auth.rotation_strategy = settings.rotation_strategy; - config.auth.tier_priority = settings.tier_priority; - config.quota.refresh_interval_minutes = settings.quota_refresh_interval_minutes; - config.quota.warning_thresholds = QuotaWarningThresholds { - primary: settings.quota_warning_thresholds.primary, - secondary: settings.quota_warning_thresholds.secondary, - }; - config.quota.skip_exhausted = settings.quota_skip_exhausted; - config.logging.enabled = settings.logs_enabled; - config.logging.capacity = settings.logs_capacity; - config.logging.capture_body = settings.logs_capture_body; - config.usage_stats.history_retention_days = settings.usage_history_retention_days; +} + +async fn insert_runtime_settings_if_missing( + pool: &SqlitePool, + settings: &AdminSettings, +) -> Result<(), RuntimeSettingsError> { + let model_aliases_json = serde_json::to_string(&settings.model_aliases)?; + let mut tx = pool.begin().await?; + let result = sqlx::query( + r" +insert or ignore into runtime_settings ( + id, + default_model, + model_aliases_json, + refresh_margin_seconds, + refresh_concurrency, + max_concurrent_per_account, + request_interval_ms, + rotation_strategy, + updated_at +) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(RUNTIME_SETTINGS_ID) + .bind(&settings.default_model) + .bind(model_aliases_json) + .bind( + i64::try_from(settings.refresh_margin_seconds) + .map_err(|_| stored_field_error("refreshMarginSeconds", "out of range"))?, + ) + .bind(i64::from(settings.refresh_concurrency)) + .bind( + i64::try_from(settings.max_concurrent_per_account) + .map_err(|_| stored_field_error("maxConcurrentPerAccount", "out of range"))?, + ) + .bind( + i64::try_from(settings.request_interval_ms) + .map_err(|_| stored_field_error("requestIntervalMs", "out of range"))?, + ) + .bind(&settings.rotation_strategy) + .bind(Utc::now().to_rfc3339()) + .execute(&mut *tx) + .await?; + if result.rows_affected() > 0 { + replace_model_account_routes(&mut tx, &settings.model_account_routes).await?; + } + tx.commit().await?; + Ok(()) +} + +async fn save_runtime_settings( + pool: &SqlitePool, + settings: &AdminSettings, +) -> Result<(), RuntimeSettingsError> { + let model_aliases_json = serde_json::to_string(&settings.model_aliases)?; + let mut tx = pool.begin().await?; + sqlx::query( + r" +insert into runtime_settings ( + id, + default_model, + model_aliases_json, + refresh_margin_seconds, + refresh_concurrency, + max_concurrent_per_account, + request_interval_ms, + rotation_strategy, + updated_at +) values (?, ?, ?, ?, ?, ?, ?, ?, ?) +on conflict(id) do update set + default_model = excluded.default_model, + model_aliases_json = excluded.model_aliases_json, + refresh_margin_seconds = excluded.refresh_margin_seconds, + refresh_concurrency = excluded.refresh_concurrency, + max_concurrent_per_account = excluded.max_concurrent_per_account, + request_interval_ms = excluded.request_interval_ms, + rotation_strategy = excluded.rotation_strategy, + updated_at = excluded.updated_at", + ) + .bind(RUNTIME_SETTINGS_ID) + .bind(&settings.default_model) + .bind(model_aliases_json) + .bind( + i64::try_from(settings.refresh_margin_seconds) + .map_err(|_| stored_field_error("refreshMarginSeconds", "out of range"))?, + ) + .bind(i64::from(settings.refresh_concurrency)) + .bind( + i64::try_from(settings.max_concurrent_per_account) + .map_err(|_| stored_field_error("maxConcurrentPerAccount", "out of range"))?, + ) + .bind( + i64::try_from(settings.request_interval_ms) + .map_err(|_| stored_field_error("requestIntervalMs", "out of range"))?, + ) + .bind(&settings.rotation_strategy) + .bind(Utc::now().to_rfc3339()) + .execute(&mut *tx) + .await?; + replace_model_account_routes(&mut tx, &settings.model_account_routes).await?; + tx.commit().await?; + Ok(()) +} + +async fn load_runtime_settings(pool: &SqlitePool) -> Result { + let row = sqlx::query( + r" +select + default_model, + model_aliases_json, + refresh_margin_seconds, + refresh_concurrency, + max_concurrent_per_account, + request_interval_ms, + rotation_strategy +from runtime_settings +where id = ?", + ) + .bind(RUNTIME_SETTINGS_ID) + .fetch_one(pool) + .await?; + let mut settings = runtime_settings_from_row(&row)?; + settings.model_account_routes = load_model_account_routes(pool).await?; + Ok(settings) +} + +fn runtime_settings_from_row(row: &SqliteRow) -> Result { + let model_aliases_json: String = row.get("model_aliases_json"); + Ok(AdminSettings { + default_model: row.get("default_model"), + model_aliases: serde_json::from_str(&model_aliases_json)?, + model_account_routes: BTreeMap::new(), + refresh_margin_seconds: positive_i64_to_u64( + "refreshMarginSeconds", + row.get("refresh_margin_seconds"), + )?, + refresh_concurrency: positive_i64_to_u32( + "refreshConcurrency", + row.get("refresh_concurrency"), + )?, + max_concurrent_per_account: positive_i64_to_usize( + "maxConcurrentPerAccount", + row.get("max_concurrent_per_account"), + )?, + request_interval_ms: non_negative_i64_to_u64( + "requestIntervalMs", + row.get("request_interval_ms"), + )?, + rotation_strategy: validate_rotation_strategy(row.get("rotation_strategy"))?, + ..AdminSettings::default() + }) +} + +async fn replace_model_account_routes( + tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + routes: &BTreeMap>, +) -> Result<(), RuntimeSettingsError> { + sqlx::query("delete from model_account_routes") + .execute(&mut **tx) + .await?; + let now = Utc::now().to_rfc3339(); + for (model, account_ids) in routes { + for (priority, account_id) in account_ids.iter().enumerate() { + sqlx::query( + r" +insert into model_account_routes ( + model, + account_id, + priority, + enabled, + created_at, + updated_at +) values (?, ?, ?, 1, ?, ?)", + ) + .bind(model) + .bind(account_id) + .bind( + i64::try_from(priority) + .map_err(|_| stored_field_error("modelAccountRoutes", "out of range"))?, + ) + .bind(&now) + .bind(&now) + .execute(&mut **tx) + .await?; + } + } + Ok(()) +} + +async fn load_model_account_routes( + pool: &SqlitePool, +) -> Result>, RuntimeSettingsError> { + let rows: Vec<(String, String)> = sqlx::query_as( + r" +select model, account_id +from model_account_routes +where enabled = 1 +order by model asc, priority asc, account_id asc", + ) + .fetch_all(pool) + .await?; + let mut routes = BTreeMap::>::new(); + for (model, account_id) in rows { + routes.entry(model).or_default().push(account_id); + } + Ok(routes) +} + +/// 从当前运行配置生成账号池调度参数。 +pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions { + AccountPoolOptions { + max_concurrent_per_account: config.auth.max_concurrent_per_account, + rotation_strategy: rotation_strategy_from_config(&config.auth.rotation_strategy), + skip_quota_limited: config.quota.skip_exhausted, + tier_priority: config.auth.tier_priority.clone(), + model_account_routes: config.model.account_routes.clone(), + ..AccountPoolOptions::default() + } +} + +fn rotation_strategy_from_config(strategy: &str) -> RotationStrategy { + match strategy { + "round_robin" => RotationStrategy::RoundRobin, + "sticky" => RotationStrategy::Sticky, + _ => RotationStrategy::LeastUsed, + } +} + +fn positive_i64_to_u64(field: &'static str, value: i64) -> Result { + if value <= 0 { + return Err(stored_field_error(field, "must be greater than 0")); + } + u64::try_from(value).map_err(|_| stored_field_error(field, "out of range")) +} + +fn positive_i64_to_u32(field: &'static str, value: i64) -> Result { + if value <= 0 { + return Err(stored_field_error(field, "must be greater than 0")); + } + u32::try_from(value).map_err(|_| stored_field_error(field, "out of range")) +} + +fn positive_i64_to_usize(field: &'static str, value: i64) -> Result { + if value <= 0 { + return Err(stored_field_error(field, "must be greater than 0")); + } + usize::try_from(value).map_err(|_| stored_field_error(field, "out of range")) +} + +fn non_negative_i64_to_u64(field: &'static str, value: i64) -> Result { + if value < 0 { + return Err(stored_field_error( + field, + "must be greater than or equal to 0", + )); + } + u64::try_from(value).map_err(|_| stored_field_error(field, "out of range")) +} + +fn stored_field_error( + field: impl Into, + message: impl Into, +) -> RuntimeSettingsError { + RuntimeSettingsError::StoredField { + field: field.into(), + message: message.into(), + } } diff --git a/src/config/types.rs b/src/config/types.rs index 1bd3b35f1..4efecc613 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -62,6 +62,8 @@ pub struct ModelConfig { pub service_tier: Option, /// 模型别名映射。 pub aliases: BTreeMap, + /// 模型到账号 ID 的显式路由。 + pub account_routes: BTreeMap>, } /// 认证、轮换与 token 续期配置。 @@ -126,8 +128,6 @@ pub struct DatabaseConfig { /// 本地安全文件路径配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub struct SecurityConfig { - /// 主密钥文件路径。 - pub master_key_file: String, /// API Key pepper 文件路径。 pub api_key_pepper_file: String, } diff --git a/src/config/writeback.rs b/src/config/writeback.rs deleted file mode 100644 index 3b9ec19ae..000000000 --- a/src/config/writeback.rs +++ /dev/null @@ -1,27 +0,0 @@ -use std::path::Path; - -use super::types::AppConfig; -use crate::config::{ConfigWriteError, ConfigWriteResult}; - -impl AppConfig { - /// 写入管理端设置到主配置文件。 - pub async fn write_settings_config( - &self, - config_path: impl AsRef, - ) -> ConfigWriteResult<()> { - let config_path = config_path.as_ref(); - if let Some(parent) = config_path - .parent() - .filter(|path| !path.as_os_str().is_empty()) - { - tokio::fs::create_dir_all(parent) - .await - .map_err(ConfigWriteError::CreateDirectory)?; - } - - let content = serde_yml::to_string(self).map_err(ConfigWriteError::Serialize)?; - tokio::fs::write(config_path, content) - .await - .map_err(ConfigWriteError::Write) - } -} diff --git a/src/infra/crypto.rs b/src/infra/crypto.rs deleted file mode 100644 index 997c2a53f..000000000 --- a/src/infra/crypto.rs +++ /dev/null @@ -1,160 +0,0 @@ -//! 加密原语 —— HMAC、常数时间比较、AES-256-GCM 密钥盒。 - -use std::{fs, path::Path}; - -use aes_gcm::{ - aead::{rand_core::RngCore, Aead, OsRng}, - Aes256Gcm, KeyInit, Nonce, -}; -use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; -use hmac::KeyInit as _; -use hmac::{Hmac, Mac}; -use secrecy::{ExposeSecret, SecretString}; -use sha2::Sha256; -use thiserror::Error; - -// --------------------------------------------------------------------------- -// HMAC-SHA256 -// --------------------------------------------------------------------------- - -/// HMAC-SHA256 类型别名。 -pub type HmacSha256 = Hmac; - -/// 计算 HMAC-SHA256 并以 URL-safe base64 输出。 -pub fn hmac_sha256_base64(key: &[u8], payload: &[u8]) -> Option { - let mut mac = HmacSha256::new_from_slice(key).ok()?; - mac.update(payload); - Some(URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes())) -} - -/// 常数时间比较两个字节切片。 -pub fn constant_time_eq(left: &[u8], right: &[u8]) -> bool { - if left.len() != right.len() { - return false; - } - - left.iter() - .zip(right) - .fold(0u8, |acc, (left, right)| acc | (left ^ right)) - == 0 -} - -// --------------------------------------------------------------------------- -// SecretBox (AES-256-GCM) -// --------------------------------------------------------------------------- - -/// 密钥盒相关错误。 -#[derive(Debug, Error)] -pub enum CryptoError { - /// 密钥长度不正确。 - #[error("invalid secret key length")] - InvalidKeyLength, - /// 加密失败。 - #[error("secret encryption failed")] - Encrypt, - /// 解密失败。 - #[error("secret decryption failed")] - Decrypt, - /// nonce 长度不正确。 - #[error("invalid nonce length")] - InvalidNonceLength, - /// Base64 编码/解码失败。 - #[error("invalid secret encoding: {0}")] - Decode(#[from] base64::DecodeError), - /// 密文版本不支持。 - #[error("unsupported secret version")] - UnsupportedVersion, - /// 解密结果不是 UTF-8。 - #[error("secret is not valid utf-8: {0}")] - Utf8(#[from] std::string::FromUtf8Error), - /// 读写密钥文件失败。 - #[error("secret key file io error: {0}")] - Io(#[from] std::io::Error), - /// 持久化密钥长度不正确。 - #[error("stored secret key must decode to 32 bytes")] - InvalidStoredKeyLength, -} - -/// 密钥盒操作结果。 -pub type CryptoResult = Result; - -/// 用于加密和解密敏感字符串的密钥盒。 -#[derive(Clone)] -pub struct SecretBox { - key: [u8; 32], -} - -impl SecretBox { - /// 使用 32 字节密钥创建密钥盒。 - pub fn new(key: [u8; 32]) -> Self { - Self { key } - } - - /// 从文件加载或生成密钥。 - pub fn load_or_create(path: impl AsRef) -> CryptoResult { - let path = path.as_ref(); - if path.exists() { - let encoded = fs::read_to_string(path)?; - let key = decode_key(encoded.trim())?; - return Ok(Self::new(key)); - } - - if let Some(parent) = path - .parent() - .filter(|parent| !parent.as_os_str().is_empty()) - { - fs::create_dir_all(parent)?; - } - let mut key = [0u8; 32]; - OsRng.fill_bytes(&mut key); - fs::write(path, URL_SAFE_NO_PAD.encode(key))?; - Ok(Self::new(key)) - } - - /// 加密敏感明文。 - pub fn encrypt(&self, plaintext: &SecretString) -> CryptoResult { - let cipher = - Aes256Gcm::new_from_slice(&self.key).map_err(|_| CryptoError::InvalidKeyLength)?; - let mut nonce_bytes = [0u8; 12]; - OsRng.fill_bytes(&mut nonce_bytes); - let nonce = Nonce::from(nonce_bytes); - let ciphertext = cipher - .encrypt(&nonce, plaintext.expose_secret().as_bytes()) - .map_err(|_| CryptoError::Encrypt)?; - Ok(format!( - "v1:{}:{}", - URL_SAFE_NO_PAD.encode(nonce_bytes), - URL_SAFE_NO_PAD.encode(ciphertext) - )) - } - - /// 解密密文。 - pub fn decrypt(&self, encoded: &str) -> CryptoResult { - let mut parts = encoded.split(':'); - let version = parts.next().unwrap_or_default(); - let nonce = parts.next().unwrap_or_default(); - let ciphertext = parts.next().unwrap_or_default(); - if version != "v1" { - return Err(CryptoError::UnsupportedVersion); - } - let nonce: [u8; 12] = URL_SAFE_NO_PAD - .decode(nonce)? - .try_into() - .map_err(|_| CryptoError::InvalidNonceLength)?; - let ciphertext = URL_SAFE_NO_PAD.decode(ciphertext)?; - let cipher = - Aes256Gcm::new_from_slice(&self.key).map_err(|_| CryptoError::InvalidKeyLength)?; - let nonce = Nonce::from(nonce); - let plaintext = cipher - .decrypt(&nonce, ciphertext.as_ref()) - .map_err(|_| CryptoError::Decrypt)?; - Ok(SecretString::new(String::from_utf8(plaintext)?.into())) - } -} - -fn decode_key(encoded: &str) -> CryptoResult<[u8; 32]> { - URL_SAFE_NO_PAD - .decode(encoded)? - .try_into() - .map_err(|_| CryptoError::InvalidStoredKeyLength) -} diff --git a/src/infra/database.rs b/src/infra/database.rs index 766f6a056..89b246875 100644 --- a/src/infra/database.rs +++ b/src/infra/database.rs @@ -1,4 +1,4 @@ -//! SQLite 数据库连接与 Schema 迁移。 +//! SQLite 数据库连接与 Schema 初始化。 use std::{fs, path::Path, str::FromStr, time::Duration}; diff --git a/src/infra/mod.rs b/src/infra/mod.rs index 9423daef3..e5a482dbf 100644 --- a/src/infra/mod.rs +++ b/src/infra/mod.rs @@ -1,9 +1,7 @@ #![deny(missing_docs)] -//! 平台基础设施层，承载加密、身份、存储、日志和 JSON 原语。 +//! 平台基础设施层，承载身份、存储、日志和 JSON 原语。 -/// 加密原语。 -pub mod crypto; /// 数据库连接与 SQLite 管理。 pub mod database; /// 身份验证原语（管理员密码、API Key 哈希）。 diff --git a/src/infra/schema.sql b/src/infra/schema.sql index 3ef0dd33b..4937c2c6d 100644 --- a/src/infra/schema.sql +++ b/src/infra/schema.sql @@ -1,14 +1,5 @@ pragma foreign_keys = on; --- ============================================ --- Schema Migrations --- ============================================ -create table if not exists schema_migrations ( - version integer primary key, - name text not null, - applied_at text not null -); - -- ============================================ -- Admin Users -- ============================================ @@ -47,6 +38,22 @@ create table if not exists client_api_keys ( ); create index if not exists idx_client_api_keys_prefix on client_api_keys(prefix) where enabled = 1; +create index if not exists idx_client_api_keys_created_id on client_api_keys(created_at desc, id desc); + +-- ============================================ +-- Runtime Settings +-- ============================================ +create table if not exists runtime_settings ( + id integer primary key check (id = 1), + default_model text not null, + model_aliases_json text not null default '{}', + refresh_margin_seconds integer not null check (refresh_margin_seconds > 0), + refresh_concurrency integer not null check (refresh_concurrency > 0), + max_concurrent_per_account integer not null check (max_concurrent_per_account > 0), + request_interval_ms integer not null check (request_interval_ms >= 0), + rotation_strategy text not null check (rotation_strategy in ('least_used', 'round_robin', 'sticky')), + updated_at text not null +); -- ============================================ -- Accounts @@ -74,6 +81,7 @@ create table if not exists accounts ( ); create index if not exists idx_accounts_status on accounts(status); +create index if not exists idx_accounts_added_id on accounts(added_at desc, id desc); create unique index if not exists ux_accounts_chatgpt_identity on accounts(chatgpt_account_id, coalesce(chatgpt_user_id, '')) where chatgpt_account_id is not null; @@ -120,6 +128,73 @@ create table if not exists account_usage ( last_used_at text ); +create index if not exists idx_account_usage_last_used_account +on account_usage(last_used_at desc, account_id desc); + +-- ============================================ +-- Account Model Usage Statistics +-- ============================================ +create table if not exists account_model_usage ( + account_id text not null references accounts(id) on delete cascade, + model text not null, + request_count integer not null default 0 check (request_count >= 0), + error_count integer not null default 0 check (error_count >= 0), + input_tokens integer not null default 0 check (input_tokens >= 0), + output_tokens integer not null default 0 check (output_tokens >= 0), + cached_tokens integer not null default 0 check (cached_tokens >= 0), + last_used_at text, + primary key (account_id, model) +); + +create index if not exists idx_account_model_usage_last_used +on account_model_usage(last_used_at desc, account_id, model); + +-- ============================================ +-- Usage Time Buckets +-- ============================================ +create table if not exists usage_time_buckets ( + bucket_start text not null, + account_id text not null default '', + model text not null default '', + service_tier text not null default '', + request_count integer not null default 0 check (request_count >= 0), + error_count integer not null default 0 check (error_count >= 0), + input_tokens integer not null default 0 check (input_tokens >= 0), + output_tokens integer not null default 0 check (output_tokens >= 0), + cached_tokens integer not null default 0 check (cached_tokens >= 0), + first_token_latency_sum integer not null default 0 check (first_token_latency_sum >= 0), + first_token_latency_count integer not null default 0 check (first_token_latency_count >= 0), + latency_sum integer not null default 0 check (latency_sum >= 0), + latency_count integer not null default 0 check (latency_count >= 0), + max_latency_ms integer not null default 0 check (max_latency_ms >= 0), + min_latency_ms integer not null default 0 check (min_latency_ms >= 0), + updated_at text not null, + primary key (bucket_start, account_id, model, service_tier) +); + +create index if not exists idx_usage_time_buckets_bucket +on usage_time_buckets(bucket_start); +create index if not exists idx_usage_time_buckets_model_bucket +on usage_time_buckets(model, bucket_start); + +-- ============================================ +-- Model Account Routes +-- ============================================ +create table if not exists model_account_routes ( + model text not null, + account_id text not null references accounts(id) on delete cascade, + priority integer not null default 0 check (priority >= 0), + enabled integer not null default 1 check (enabled in (0, 1)), + created_at text not null, + updated_at text not null, + primary key (model, account_id) +); + +create index if not exists idx_model_account_routes_account +on model_account_routes(account_id); +create index if not exists idx_model_account_routes_enabled_model +on model_account_routes(enabled, model, priority asc, account_id); + -- ============================================ -- Account Cookies -- ============================================ @@ -128,7 +203,7 @@ create table if not exists account_cookies ( account_id text not null references accounts(id) on delete cascade, domain text not null, name text not null, - value_cipher text not null, + value text not null, path text not null default '/', expires_at text, updated_at text not null, @@ -137,6 +212,9 @@ create table if not exists account_cookies ( create index if not exists idx_account_cookies_account on account_cookies(account_id); create index if not exists idx_account_cookies_account_domain on account_cookies(account_id, domain); +create index if not exists idx_account_cookies_expires +on account_cookies(expires_at) +where expires_at is not null; -- ============================================ -- Device Fingerprints @@ -168,6 +246,9 @@ create table if not exists fingerprint_update_history ( created_at text not null ); +create index if not exists idx_fingerprint_update_history_created_id +on fingerprint_update_history(created_at desc, id desc); + -- ============================================ -- Event Logs -- ============================================ @@ -200,6 +281,11 @@ create index if not exists idx_event_logs_transport on event_logs(transport, cre create index if not exists idx_event_logs_failure_class on event_logs(failure_class, created_at desc) where failure_class is not null; create index if not exists idx_event_logs_response_id on event_logs(response_id) where response_id is not null; create index if not exists idx_event_logs_upstream_request_id on event_logs(upstream_request_id) where upstream_request_id is not null; +create index if not exists idx_event_logs_level_created on event_logs(level, created_at desc); +create index if not exists idx_event_logs_route_created on event_logs(route, created_at desc) where route is not null; +create index if not exists idx_event_logs_model_created on event_logs(model, created_at desc) where model is not null; +create index if not exists idx_event_logs_status_created on event_logs(status_code, created_at desc) where status_code is not null; +create index if not exists idx_event_logs_upstream_status_created on event_logs(upstream_status_code, created_at desc) where upstream_status_code is not null; -- ============================================ -- Model Plan Snapshots @@ -228,3 +314,4 @@ create table if not exists session_affinities ( create index if not exists idx_session_affinities_conversation on session_affinities(conversation_id, created_at desc); create index if not exists idx_session_affinities_expires on session_affinities(expires_at); +create index if not exists idx_session_affinities_active_order on session_affinities(expires_at, created_at, response_id); diff --git a/src/proxy/dispatch/chat.rs b/src/proxy/dispatch/chat.rs index 96b64f6ec..b75669aea 100644 --- a/src/proxy/dispatch/chat.rs +++ b/src/proxy/dispatch/chat.rs @@ -373,7 +373,7 @@ impl ChatDispatchService { self.cloudflare.reset_account_recovery(&account_id).await; if let Some(ref usage) = response.usage { self.account_pool - .record_token_usage(&account_id, usage) + .record_token_usage(&account_id, &request.model, usage) .await; } record_response_event(ResponseEventRecord { diff --git a/src/proxy/dispatch/responses.rs b/src/proxy/dispatch/responses.rs index 8cfda9ed7..3bce6753d 100644 --- a/src/proxy/dispatch/responses.rs +++ b/src/proxy/dispatch/responses.rs @@ -609,6 +609,7 @@ impl ResponseDispatchService { self.account_pool .record_empty_response_attempt( &release_account_id, + &request.model, image_generation_requested, ) .await; @@ -845,7 +846,12 @@ impl ResponseDispatchService { .await; if let Some(usage) = response.usage { self.account_pool - .record_response_usage(&account.id, usage, image_generation_requested) + .record_response_usage( + &account.id, + &request.model, + usage, + image_generation_requested, + ) .await; } self.record_response_affinity( @@ -1723,7 +1729,7 @@ impl ResponseDispatchService { let usage = extract_usage(&response.body); if let Some(usage) = usage { self.account_pool - .record_token_usage(&account.id, &usage) + .record_token_usage(&account.id, &request.model, &usage) .await; } record_response_event(ResponseEventRecord { @@ -2338,7 +2344,7 @@ async fn finalize_live_response_stream(context: LiveResponseStreamContext, body: Ok(Some(usage)) => { context .account_pool - .record_token_usage(&context.account_id, &usage) + .record_token_usage(&context.account_id, &context.model, &usage) .await; Some(usage) } diff --git a/src/runtime/bootstrap.rs b/src/runtime/bootstrap.rs index 7d392bf27..75868c613 100644 --- a/src/runtime/bootstrap.rs +++ b/src/runtime/bootstrap.rs @@ -3,6 +3,7 @@ use std::error::Error; use axum::Router; use chrono::Utc; +use crate::config::settings::RuntimeSettingsService; use crate::config::types::AppConfig; use crate::infra::database::connect_sqlite; use crate::infra::logging::{init_tracing, LogGuard, RotationConfig}; @@ -35,7 +36,7 @@ async fn build_application( config: AppConfig, ) -> Result<(Router, TaskCoordinator), Box> { let pool = connect_sqlite(&config.database.url).await?; - let secret = crate::infra::crypto::SecretBox::load_or_create(&config.security.master_key_file)?; + let config = RuntimeSettingsService::load_or_initialize_config(config, &pool).await?; let hasher = crate::infra::identity::ApiKeyHasher::load_or_create(&config.security.api_key_pepper_file)?; @@ -43,10 +44,7 @@ async fn build_application( let stores = BackgroundTaskStores { accounts: crate::upstream::accounts::store::SqliteAccountStore::new(pool.clone()), admin_sessions: crate::admin::auth::service::SqliteAdminSessionStore::new(pool.clone()), - cookies: crate::upstream::accounts::cookies::SqliteCookieStore::new( - pool.clone(), - secret.clone(), - ), + cookies: crate::upstream::accounts::cookies::SqliteCookieStore::new(pool.clone()), fingerprints: fingerprint_repository.clone(), session_affinity: crate::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new( pool.clone(), diff --git a/src/runtime/services.rs b/src/runtime/services.rs index 2330d5b62..d97431c31 100644 --- a/src/runtime/services.rs +++ b/src/runtime/services.rs @@ -13,7 +13,10 @@ use crate::{ auth::service::{AdminSessionService, SqliteAdminSessionStore}, keys::service::{AdminClientKeyService, ClientKeyService, SqliteClientKeyStore}, }, - config::{settings::RuntimeSettingsService, types::AppConfig}, + config::{ + settings::{account_pool_options_from_config, RuntimeSettingsService}, + types::AppConfig, + }, proxy::dispatch::{ chat::ChatDispatchService, responses::ResponseDispatchService, @@ -21,7 +24,7 @@ use crate::{ }, upstream::accounts::{ cookies::SqliteCookieStore, - pool::{AccountPoolOptions, RuntimeAccountPoolService}, + pool::RuntimeAccountPoolService, store::{AccountStore as AccountStoreTrait, SqliteAccountStore}, token_refresh::RefreshLeaseStore as SqliteRefreshLeaseStore, }, @@ -133,9 +136,15 @@ impl Services { StdArc::new(client) } }; - let settings = StdArc::new(RuntimeSettingsService::with_config_path( + let account_pool = StdArc::new(RuntimeAccountPoolService::new( + account_store_trait.clone(), + account_pool_options_from_config(config), + config.auth.request_interval_ms, + )); + let settings = StdArc::new(RuntimeSettingsService::with_account_pool( config.clone(), - "config.yaml", + stores.accounts.pool().clone(), + account_pool.clone(), )); let admin_sessions = StdArc::new(AdminSessionService::new( stores.admin_sessions.clone(), @@ -146,11 +155,6 @@ impl Services { let client_keys = StdArc::new(ClientKeyService::new(StdArc::new( stores.client_keys.clone(), ))); - let account_pool = StdArc::new(RuntimeAccountPoolService::new( - account_store_trait.clone(), - AccountPoolOptions::default(), - config.auth.request_interval_ms, - )); let token_client = StdArc::new(default_openai_token_client(token_client_config(config))); let admin_accounts = StdArc::new(AdminAccountService::new( stores.accounts.clone(), diff --git a/src/upstream/accounts/cookies.rs b/src/upstream/accounts/cookies.rs index 676cc8a35..bcb9d8439 100644 --- a/src/upstream/accounts/cookies.rs +++ b/src/upstream/accounts/cookies.rs @@ -3,14 +3,11 @@ use std::{collections::HashMap, sync::Arc}; use chrono::{DateTime, Duration, Utc}; -use secrecy::{ExposeSecret, SecretString}; use sqlx::{Row, SqlitePool}; use thiserror::Error; use tokio::sync::RwLock; use uuid::Uuid; -use crate::infra::crypto::{CryptoError, SecretBox}; - /// 单个 Cookie 条目。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct CookieEntry { @@ -40,9 +37,6 @@ pub enum SqliteCookieStoreError { /// 数据库错误。 #[error("sqlite cookie store database error: {0}")] Database(#[from] sqlx::Error), - /// 加解密错误。 - #[error("sqlite cookie store crypto error: {0}")] - Crypto(#[from] CryptoError), } /// SQLite Cookie 存储结果。 @@ -52,13 +46,12 @@ pub type SqliteCookieStoreResult = Result; #[derive(Clone)] pub struct SqliteCookieStore { pool: SqlitePool, - secret: SecretBox, } impl SqliteCookieStore { /// 构造 Cookie 存储。 - pub fn new(pool: SqlitePool, secret: SecretBox) -> Self { - Self { pool, secret } + pub fn new(pool: SqlitePool) -> Self { + Self { pool } } /// 返回底层连接池。 @@ -122,7 +115,7 @@ impl SqliteCookieStore { request_path: &str, ) -> SqliteCookieStoreResult> { let rows = sqlx::query( - "select domain, name, value_cipher, path, expires_at from account_cookies where account_id = ?", + "select domain, name, value, path, expires_at from account_cookies where account_id = ?", ) .bind(account_id) .fetch_all(&self.pool) @@ -143,12 +136,10 @@ impl SqliteCookieStore { continue; } let name = row.get::("name"); - let value_cipher = row.get::("value_cipher"); - let value = self.secret.decrypt(&value_cipher)?; pairs.push(CookieHeaderPair { path_len: path.len(), name: name.clone(), - value: format!("{name}={}", value.expose_secret()), + value: format!("{name}={}", row.get::("value")), }); } if pairs.is_empty() { @@ -196,17 +187,14 @@ impl SqliteCookieStore { parsed: ParsedCookie, ) -> SqliteCookieStoreResult<()> { let now = Utc::now().to_rfc3339(); - let value_cipher = self - .secret - .encrypt(&SecretString::new(parsed.value.into()))?; sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?) on conflict(account_id, domain, name, path) do update set value_cipher = excluded.value_cipher, expires_at = excluded.expires_at, updated_at = excluded.updated_at", + "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?) on conflict(account_id, domain, name, path) do update set value = excluded.value, expires_at = excluded.expires_at, updated_at = excluded.updated_at", ) .bind(Uuid::new_v4().to_string()) .bind(account_id) .bind(parsed.domain) .bind(parsed.name) - .bind(value_cipher) + .bind(parsed.value) .bind(parsed.path) .bind(parsed.expires_at) .bind(now) diff --git a/src/upstream/accounts/model.rs b/src/upstream/accounts/model.rs index 7ea8bf32d..50ea971fc 100644 --- a/src/upstream/accounts/model.rs +++ b/src/upstream/accounts/model.rs @@ -151,3 +151,18 @@ impl AccountUsageDelta { } } } + +/// 账号模型维度用量增量。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct AccountModelUsageDelta { + /// 请求数增量。 + pub requests: u64, + /// 错误数增量。 + pub errors: u64, + /// 输入 token 增量。 + pub input_tokens: u64, + /// 输出 token 增量。 + pub output_tokens: u64, + /// 缓存 token 增量。 + pub cached_tokens: u64, +} diff --git a/src/upstream/accounts/pool.rs b/src/upstream/accounts/pool.rs index 09ced8a89..a0c222d7d 100644 --- a/src/upstream/accounts/pool.rs +++ b/src/upstream/accounts/pool.rs @@ -14,7 +14,7 @@ use crate::upstream::accounts::quota::{quota_snapshot_limit_reached, quota_snaps use crate::upstream::accounts::store::AccountStore; use crate::upstream::accounts::token_refresh::{jwt_expiry, JwtExpiry}; use crate::upstream::accounts::{ - model::{Account, AccountStatus, AccountUsageDelta}, + model::{Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta}, service::AccountService, }; use crate::upstream::protocol::events::{ @@ -59,6 +59,8 @@ pub struct AccountPoolOptions { pub tier_priority: Vec, /// 模型到允许订阅计划的映射。 pub model_plan_allowlist: BTreeMap>, + /// 模型到指定账号 ID 的显式路由。 + pub model_account_routes: BTreeMap>, } impl Default for AccountPoolOptions { @@ -70,6 +72,7 @@ impl Default for AccountPoolOptions { skip_quota_limited: true, tier_priority: Vec::new(), model_plan_allowlist: BTreeMap::new(), + model_account_routes: BTreeMap::new(), } } } @@ -240,6 +243,16 @@ impl AccountPool { 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 { @@ -642,6 +655,7 @@ impl AccountPool { request.now, self.options.skip_quota_limited, ) + && self.is_model_account_allowed(account, &request.model) && self.is_model_allowed(account, &request.model) && !request .exclude_account_ids @@ -720,6 +734,15 @@ impl AccountPool { .any(|plan| account.plan_type.as_deref() == Some(plan.as_str())) } + fn is_model_account_allowed(&self, account: &Account, model: &str) -> bool { + let Some(account_ids) = self.options.model_account_routes.get(model) else { + return true; + }; + account_ids + .iter() + .any(|account_id| account_id == &account.id) + } + fn best_available_tier(&self, candidates: &[Account]) -> Option { self.options.tier_priority.iter().find_map(|tier| { candidates @@ -750,7 +773,7 @@ fn access_token_expired(account: &Account, now: DateTime) -> bool { pub struct RuntimeAccountPoolService { pool: StdArc>, store: StdArc, - request_interval: StdDuration, + request_interval: StdArc>, } impl RuntimeAccountPoolService { @@ -763,10 +786,22 @@ impl RuntimeAccountPoolService { Self { pool: StdArc::new(tokio::sync::Mutex::new(AccountPool::with_options(options))), store, - request_interval: StdDuration::from_millis(request_interval_ms), + request_interval: StdArc::new(std::sync::RwLock::new(StdDuration::from_millis( + request_interval_ms, + ))), } } + /// 更新账号池运行参数。 + pub async fn apply_options(&self, options: AccountPoolOptions, request_interval_ms: u64) { + self.pool.lock().await.set_options(options); + *self + .request_interval + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + StdDuration::from_millis(request_interval_ms); + } + /// 从持久化账号恢复运行时账号池。 pub async fn restore_from_repository(&self) -> Result { let accounts = self @@ -843,6 +878,7 @@ impl RuntimeAccountPoolService { /// 按请求上下文获取可用账号。 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) .await; @@ -854,6 +890,25 @@ impl RuntimeAccountPoolService { "failed to persist account request usage" ); } + if let Err(error) = self + .store + .record_model_usage_delta( + &acquired.account.id, + &model, + AccountModelUsageDelta { + requests: 1, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id = acquired.account.id, + model, + error = %error, + "failed to persist account model request usage" + ); + } Some(acquired) } @@ -875,7 +930,11 @@ impl RuntimeAccountPoolService { /// 按账号上一个在途槽位控制请求间隔。 pub async fn wait_for_request_interval(&self, acquired: &AcquiredAccount) { - if self.request_interval.is_zero() { + let request_interval = *self + .request_interval + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if request_interval.is_zero() { return; } let Some(previous_slot_at) = acquired.previous_slot_at else { @@ -885,20 +944,22 @@ impl RuntimeAccountPoolService { .signed_duration_since(previous_slot_at) .to_std() .unwrap_or_default(); - if elapsed < self.request_interval { - tokio::time::sleep(self.request_interval - elapsed).await; + if elapsed < request_interval { + tokio::time::sleep(request_interval - elapsed).await; } } /// 记录上游 token 用量。 - pub async fn record_token_usage(&self, account_id: &str, usage: &CodexTokenUsage) { - self.record_response_usage(account_id, *usage, false).await; + pub async fn record_token_usage(&self, account_id: &str, model: &str, usage: &CodexTokenUsage) { + self.record_response_usage(account_id, model, *usage, false) + .await; } /// 记录 Responses 请求用量。 pub async fn record_response_usage( &self, account_id: &str, + model: &str, usage: CodexTokenUsage, image_generation_requested: bool, ) { @@ -927,6 +988,27 @@ impl RuntimeAccountPoolService { "failed to persist account token usage" ); } + if let Err(error) = self + .store + .record_model_usage_delta( + account_id, + model, + AccountModelUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id, + model, + error = %error, + "failed to persist account model token usage" + ); + } self.pool.lock().await.record_window_token_usage( account_id, AccountWindowUsageDelta { @@ -945,6 +1027,7 @@ impl RuntimeAccountPoolService { pub async fn record_empty_response_attempt( &self, account_id: &str, + model: &str, image_generation_requested: bool, ) { let usage = AccountUsageDelta { @@ -959,6 +1042,25 @@ impl RuntimeAccountPoolService { "failed to persist empty response usage" ); } + if let Err(error) = self + .store + .record_model_usage_delta( + account_id, + model, + AccountModelUsageDelta { + errors: 1, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id, + model, + error = %error, + "failed to persist account model empty response usage" + ); + } if image_generation_requested { self.pool.lock().await.record_window_token_usage( account_id, diff --git a/src/upstream/accounts/store.rs b/src/upstream/accounts/store.rs index 6c69d0d03..552576c8d 100644 --- a/src/upstream/accounts/store.rs +++ b/src/upstream/accounts/store.rs @@ -8,7 +8,9 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool}; use thiserror::Error; use crate::infra::json::{decode_cursor, page_offset, NumberedPage, Page}; -use crate::upstream::accounts::model::{Account, AccountStatus, AccountUsageDelta}; +use crate::upstream::accounts::model::{ + Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta, +}; // ============================================================================ // SQL 常量 @@ -314,8 +316,40 @@ on conflict(account_id) do update set then excluded.last_used_at else account_usage.window_started_at end, + last_used_at = excluded.last_used_at"; + +const RECORD_MODEL_USAGE_SQL: &str = r" +insert into account_model_usage ( + account_id, + model, + request_count, + error_count, + input_tokens, + output_tokens, + cached_tokens, + last_used_at +) values (?, ?, ?, ?, ?, ?, ?, ?) +on conflict(account_id, model) do update set + request_count = request_count + excluded.request_count, + error_count = error_count + excluded.error_count, + input_tokens = input_tokens + excluded.input_tokens, + output_tokens = output_tokens + excluded.output_tokens, + cached_tokens = cached_tokens + excluded.cached_tokens, last_used_at = excluded.last_used_at"; +const LIST_MODEL_USAGE_SQL: &str = r" +select + account_id, + model, + request_count, + error_count, + input_tokens, + output_tokens, + cached_tokens, + last_used_at +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, @@ -776,6 +810,27 @@ pub struct AccountUsageRecord { pub last_used_at: Option>, } +/// 账号模型维度用量记录。 +#[derive(Debug, Clone)] +pub struct AccountModelUsageRecord { + /// 账号 ID。 + pub account_id: String, + /// 模型 ID。 + pub model: String, + /// 历史请求总数。 + pub request_count: i64, + /// 历史错误数。 + pub error_count: i64, + /// 累计输入 token。 + pub input_tokens: i64, + /// 累计输出 token。 + pub output_tokens: i64, + /// 累计缓存 token。 + pub cached_tokens: i64, + /// 最近使用时间。 + pub last_used_at: Option>, +} + /// 账号配额快照。 #[derive(Debug, Clone)] pub struct AccountQuotaSnapshot { @@ -870,6 +925,14 @@ pub trait AccountStore: Send + Sync + 'static { usage: AccountUsageDelta, ) -> AccountStoreResult<()>; + /// 记录账号模型维度用量增量。 + async fn record_model_usage_delta( + &self, + account_id: &str, + model: &str, + usage: AccountModelUsageDelta, + ) -> AccountStoreResult<()>; + /// 读取账号当前配额 JSON。 async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult>; @@ -1329,6 +1392,40 @@ impl SqliteAccountStore { Ok(()) } + /// 记录模型维度用量。 + pub async fn record_model_usage( + &self, + account_id: &str, + model: &str, + delta: AccountModelUsageDelta, + ) -> SqliteAccountStoreResult<()> { + let model = model.trim(); + if model.is_empty() { + return Ok(()); + } + let last_used_at = Utc::now().to_rfc3339(); + sqlx::query(RECORD_MODEL_USAGE_SQL) + .bind(account_id) + .bind(model) + .bind(u64_to_i64_saturating(delta.requests)) + .bind(u64_to_i64_saturating(delta.errors)) + .bind(u64_to_i64_saturating(delta.input_tokens)) + .bind(u64_to_i64_saturating(delta.output_tokens)) + .bind(u64_to_i64_saturating(delta.cached_tokens)) + .bind(&last_used_at) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 列出模型维度用量。 + pub async fn list_model_usage(&self) -> SqliteAccountStoreResult> { + let rows = sqlx::query(LIST_MODEL_USAGE_SQL) + .fetch_all(&self.pool) + .await?; + rows.iter().map(model_usage_from_row).collect() + } + /// 获取用量记录。 pub async fn get_usage( &self, @@ -1591,6 +1688,17 @@ impl AccountStore for SqliteAccountStore { .map_err(map_account_store_error) } + async fn record_model_usage_delta( + &self, + account_id: &str, + model: &str, + usage: AccountModelUsageDelta, + ) -> AccountStoreResult<()> { + self.record_model_usage(account_id, model, usage) + .await + .map_err(map_account_store_error) + } + async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult> { SqliteAccountStore::get_quota_json(self, account_id) .await @@ -1883,6 +1991,21 @@ fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> SqliteAccountStoreResult SqliteAccountStoreResult { + Ok(AccountModelUsageRecord { + account_id: row.get("account_id"), + model: row.get("model"), + request_count: row.get("request_count"), + error_count: row.get("error_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + last_used_at: parse_optional_rfc3339(row.get::, _>("last_used_at"))?, + }) +} + fn quota_snapshot_from_row( row: &sqlx::sqlite::SqliteRow, ) -> SqliteAccountStoreResult { diff --git a/tests/admin/accounts/import_export.rs b/tests/admin/accounts/import_export.rs index 7db3ab723..4bc9d2f11 100644 --- a/tests/admin/accounts/import_export.rs +++ b/tests/admin/accounts/import_export.rs @@ -11,7 +11,6 @@ async fn admin_accounts_export_should_return_native_account_tokens() { let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([77u8; 32]); seed_account( &pool, NewAccount { @@ -34,7 +33,7 @@ async fn admin_accounts_export_should_return_native_account_tokens() { let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box.clone()), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -80,7 +79,6 @@ async fn admin_accounts_import_should_store_native_account_tokens() { let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([79u8; 32]); let access_token = test_jwt( "chatgpt_import", Some("user_import"), @@ -92,7 +90,7 @@ async fn admin_accounts_import_should_store_native_account_tokens() { let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box.clone()), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -190,7 +188,7 @@ async fn admin_accounts_import_should_fetch_wham_usage_for_current_openai_token_ }))) .mount(&server) .await; - let (app, _state, pool, _dir, _) = admin_accounts_test_app_with_api_base_url( + let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( "admin-accounts-import-current-openai-token.sqlite", 118, format!("{}/backend-api", server.uri()), @@ -295,7 +293,7 @@ async fn admin_accounts_import_should_complete_chatgpt_account_id_from_refresh_t }))) .mount(&server) .await; - let (app, _state, pool, _dir, _) = + let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url_and_oauth_token_endpoint( "admin-accounts-import-rt-complete.sqlite", 125, @@ -382,14 +380,13 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([120u8; 32]); let mut config = test_config(url); config.api.base_url = format!("{}/backend-api", server.uri()); let hasher = ApiKeyHasher::new([121u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box.clone()), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -499,7 +496,7 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q #[tokio::test] async fn admin_accounts_import_should_update_existing_sub2api_account_and_clear_stale_quota_lock() { - let (app, _state, pool, _dir, _) = + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-accounts-import-update-existing.sqlite", 116).await; let new_access_token = test_jwt( "chatgpt-import-update", @@ -603,12 +600,11 @@ async fn admin_accounts_import_should_require_admin_session_cookie() { let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); let config = test_config(url); - let secret_box = SecretBox::new([111u8; 32]); let hasher = ApiKeyHasher::new([112u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box.clone()), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -641,7 +637,7 @@ async fn admin_accounts_import_should_require_admin_session_cookie() { #[tokio::test] async fn admin_accounts_import_should_store_plain_tokens_and_list_sanitized_accounts() { - let (app, _state, pool, _dir, _secret_box) = + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-accounts-import-sanitized.sqlite", 113).await; let access_token = test_jwt( "chatgpt-account", @@ -700,7 +696,7 @@ async fn admin_accounts_import_should_store_plain_tokens_and_list_sanitized_acco #[tokio::test] async fn admin_accounts_export_should_return_native_accounts_with_tokens_and_filter_ids() { - let (app, _state, pool, _dir, _) = + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-accounts-export-filter.sqlite", 117).await; seed_account( &pool, diff --git a/tests/admin/accounts/lifecycle.rs b/tests/admin/accounts/lifecycle.rs index 517e94de5..47c73b51c 100644 --- a/tests/admin/accounts/lifecycle.rs +++ b/tests/admin/accounts/lifecycle.rs @@ -10,13 +10,13 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() { sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)") .bind("acct_lifecycle").bind("life@example.com").bind("cipher").bind("active") .bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z").execute(&pool).await.unwrap(); + seed_account_related_rows(&pool, "acct_lifecycle").await; let config = test_config(url); - let secret_box = SecretBox::new([75u8; 32]); let hasher = ApiKeyHasher::new([76u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -114,11 +114,116 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() { .unwrap(); assert_eq!(deleted.status(), StatusCode::OK); assert_eq!(response_json(deleted).await["data"]["deleted"], 1); + assert_account_related_rows_deleted(&pool, "acct_lifecycle").await; +} + +async fn seed_account_related_rows(pool: &SqlitePool, account_id: &str) { + sqlx::query("insert into account_usage (account_id, request_count) values (?, 1)") + .bind(account_id) + .execute(pool) + .await + .unwrap(); + sqlx::query( + "insert into account_model_usage (account_id, model, request_count) values (?, ?, 1)", + ) + .bind(account_id) + .bind("gpt-5.5") + .execute(pool) + .await + .unwrap(); + sqlx::query( + "insert into model_account_routes (model, account_id, priority, created_at, updated_at) values (?, ?, 0, ?, ?)", + ) + .bind("gpt-5.5") + .bind(account_id) + .bind("2026-06-18T00:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); + sqlx::query("insert into account_cookies (id, account_id, domain, name, value, updated_at) values (?, ?, ?, ?, ?, ?)") + .bind("cookie_lifecycle") + .bind(account_id) + .bind("chatgpt.com") + .bind("cf_clearance") + .bind("cipher") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); + sqlx::query("insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?)") + .bind(account_id) + .bind("test") + .bind("2026-06-18T00:10:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); + sqlx::query("insert into session_affinities (response_id, account_id, conversation_id, expires_at, created_at) values (?, ?, ?, ?, ?)") + .bind("resp_lifecycle") + .bind(account_id) + .bind("conv_lifecycle") + .bind("2026-06-18T01:00:00Z") + .bind("2026-06-18T00:00:00Z") + .execute(pool) + .await + .unwrap(); +} + +async fn assert_account_related_rows_deleted(pool: &SqlitePool, account_id: &str) { + let accounts: i64 = sqlx::query_scalar("select count(*) from accounts where id = ?") + .bind(account_id) + .fetch_one(pool) + .await + .unwrap(); + let usage: i64 = sqlx::query_scalar("select count(*) from account_usage where account_id = ?") + .bind(account_id) + .fetch_one(pool) + .await + .unwrap(); + let model_usage: i64 = + sqlx::query_scalar("select count(*) from account_model_usage where account_id = ?") + .bind(account_id) + .fetch_one(pool) + .await + .unwrap(); + let cookies: i64 = + sqlx::query_scalar("select count(*) from account_cookies where account_id = ?") + .bind(account_id) + .fetch_one(pool) + .await + .unwrap(); + let model_routes: i64 = + sqlx::query_scalar("select count(*) from model_account_routes where account_id = ?") + .bind(account_id) + .fetch_one(pool) + .await + .unwrap(); + let leases: i64 = + sqlx::query_scalar("select count(*) from account_refresh_leases where account_id = ?") + .bind(account_id) + .fetch_one(pool) + .await + .unwrap(); + let affinities: i64 = + sqlx::query_scalar("select count(*) from session_affinities where account_id = ?") + .bind(account_id) + .fetch_one(pool) + .await + .unwrap(); + + assert_eq!(accounts, 0); + assert_eq!(usage, 0); + assert_eq!(model_usage, 0); + assert_eq!(model_routes, 0); + assert_eq!(cookies, 0); + assert_eq!(leases, 0); + assert_eq!(affinities, 0); } #[tokio::test] async fn admin_account_status_update_should_update_proxy_account_pool() { - let (app, state, pool, _dir, _) = + let (app, state, pool, _dir) = admin_accounts_test_app("admin-account-status-runtime-pool.sqlite", 121).await; seed_account( &pool, @@ -199,7 +304,7 @@ async fn admin_account_refresh_should_not_mark_valid_account_banned_when_refresh }))) .mount(&server) .await; - let (app, _state, pool, _dir, _) = admin_accounts_test_app_with_oauth_token_endpoint( + let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( "admin-account-refresh-rt-reused.sqlite", 124, format!("{}/oauth/token", server.uri()), @@ -257,7 +362,7 @@ async fn admin_account_refresh_should_not_mark_valid_account_banned_when_refresh #[tokio::test] async fn admin_account_update_should_accept_batch_status_payload() { - let (app, _state, pool, _dir, _secret_box) = + 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") @@ -293,7 +398,7 @@ async fn admin_account_update_should_accept_batch_status_payload() { #[tokio::test] async fn admin_account_batch_status_route_should_not_exist() { - let (_app, state, _pool, _dir, _secret_box) = + 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); @@ -316,13 +421,12 @@ async fn admin_account_batch_status_route_should_not_exist() { } #[tokio::test] -async fn admin_account_cookies_should_store_encrypted_cookie_header() { +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; - let secret_box = SecretBox::new([85u8; 32]); seed_account( &pool, NewAccount { @@ -345,7 +449,7 @@ async fn admin_account_cookies_should_store_encrypted_cookie_header() { let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -400,14 +504,14 @@ async fn admin_account_cookies_should_store_encrypted_cookie_header() { assert!(cookies.contains("__cf_bm=bm-secret")); let stored_values: Vec = - sqlx::query_scalar("select value_cipher from account_cookies where account_id = ?1") + 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")); + assert!(stored_values.iter().any(|value| value == "clear-secret")); + assert!(stored_values.iter().any(|value| value == "bm-secret")); } #[tokio::test] @@ -417,7 +521,6 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() { let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([91u8; 32]); let token = test_jwt( "jwt-account", Some("jwt-user"), @@ -429,7 +532,7 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() { let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box.clone()), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -470,12 +573,11 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; let config = test_config(url); - let secret_box = SecretBox::new([93u8; 32]); let hasher = ApiKeyHasher::new([94u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -507,7 +609,7 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un #[tokio::test] async fn admin_account_manual_create_should_accept_current_openai_token_without_chatgpt_account_id() { - let (app, _state, pool, _dir, _) = + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-account-create-current-openai-token.sqlite", 119).await; let token = test_jwt_with_exp( None, diff --git a/tests/admin/accounts/list.rs b/tests/admin/accounts/list.rs index 99ca04724..d6e1c89d8 100644 --- a/tests/admin/accounts/list.rs +++ b/tests/admin/accounts/list.rs @@ -1,5 +1,4 @@ use super::*; -use codex_proxy_rs::admin::monitoring::events::{EventLevel, EventLog}; #[tokio::test] async fn admin_accounts_list_should_not_expose_account_tokens() { @@ -13,12 +12,11 @@ async fn admin_accounts_list_should_not_expose_account_tokens() { .bind("active").bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z") .execute(&pool).await.unwrap(); let config = test_config(url); - let secret_box = SecretBox::new([63u8; 32]); let hasher = ApiKeyHasher::new([64u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -58,7 +56,7 @@ async fn admin_accounts_list_should_not_expose_account_tokens() { #[tokio::test] async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { - let (app, _state, pool, _dir, _secret_box) = + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-accounts-stats.sqlite", 68).await; seed_usage_account( &pool, @@ -125,43 +123,32 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { .execute(&pool) .await .unwrap(); - let store = SqliteEventLogStore::new(pool.clone()); - let mut log = EventLog::new("response", EventLevel::Info, "ok"); - log.account_id = Some("acct_stats".to_string()); - log.model = Some("gpt-5.5".to_string()); - log.created_at = "2026-06-23T08:50:13Z".parse().unwrap(); - log.metadata = json!({ - "usage": { - "inputTokens": 3_500_000u64, - "outputTokens": 13_900u64, - "cachedTokens": 3_400_000u64 - } - }); - store.append(&log).await.unwrap(); - let mut second_gpt55_log = EventLog::new("response", EventLevel::Info, "ok"); - second_gpt55_log.account_id = Some("acct_stats".to_string()); - second_gpt55_log.model = Some("gpt-5.5".to_string()); - second_gpt55_log.created_at = "2026-06-23T08:51:13Z".parse().unwrap(); - second_gpt55_log.metadata = json!({ - "usage": { - "inputTokens": 100u64, - "outputTokens": 20u64, - "cachedTokens": 10u64 - } - }); - store.append(&second_gpt55_log).await.unwrap(); - let mut gpt5_log = EventLog::new("response", EventLevel::Info, "ok"); - gpt5_log.account_id = Some("acct_stats".to_string()); - gpt5_log.model = Some("gpt-5".to_string()); - gpt5_log.created_at = "2026-06-23T08:52:13Z".parse().unwrap(); - gpt5_log.metadata = json!({ - "usage": { - "inputTokens": 200u64, - "outputTokens": 30u64, - "cachedTokens": 20u64 - } - }); - store.append(&gpt5_log).await.unwrap(); + sqlx::query( + "insert into account_model_usage (account_id, model, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", + ) + .bind("acct_stats") + .bind("gpt-5.5") + .bind(2) + .bind(3_500_100) + .bind(13_920) + .bind(3_400_010) + .bind("2026-06-23T08:51:13Z") + .execute(&pool) + .await + .unwrap(); + sqlx::query( + "insert into account_model_usage (account_id, model, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", + ) + .bind("acct_stats") + .bind("gpt-5") + .bind(1) + .bind(200) + .bind(30) + .bind(20) + .bind("2026-06-23T08:52:13Z") + .execute(&pool) + .await + .unwrap(); let response = app .oneshot( @@ -199,7 +186,7 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { #[tokio::test] async fn admin_accounts_list_should_return_numbered_page_with_search_total() { - let (app, _state, pool, _dir, _secret_box) = + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-accounts-numbered.sqlite", 67).await; for (id, email, label, added_at) in [ ( diff --git a/tests/admin/accounts/mod.rs b/tests/admin/accounts/mod.rs index 9f617d07f..ccc057229 100644 --- a/tests/admin/accounts/mod.rs +++ b/tests/admin/accounts/mod.rs @@ -15,7 +15,7 @@ use codex_proxy_rs::{ QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, WebSocketPoolConfig, }, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, runtime::state::AppState, @@ -110,13 +110,7 @@ async fn post_admin_account(app: &axum::Router, payload: Value) -> axum::respons async fn admin_accounts_test_app( db_name: &str, key_byte: u8, -) -> ( - axum::Router, - AppState, - SqlitePool, - tempfile::TempDir, - SecretBox, -) { +) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { admin_accounts_test_app_with_api_base_url( db_name, key_byte, @@ -129,13 +123,7 @@ async fn admin_accounts_test_app_with_api_base_url( db_name: &str, key_byte: u8, api_base_url: String, -) -> ( - axum::Router, - AppState, - SqlitePool, - tempfile::TempDir, - SecretBox, -) { +) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { admin_accounts_test_app_with_overrides(db_name, key_byte, api_base_url, None).await } @@ -143,13 +131,7 @@ async fn admin_accounts_test_app_with_oauth_token_endpoint( db_name: &str, key_byte: u8, oauth_token_endpoint: String, -) -> ( - axum::Router, - AppState, - SqlitePool, - tempfile::TempDir, - SecretBox, -) { +) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { admin_accounts_test_app_with_overrides( db_name, key_byte, @@ -164,13 +146,7 @@ async fn admin_accounts_test_app_with_api_base_url_and_oauth_token_endpoint( key_byte: u8, api_base_url: String, oauth_token_endpoint: String, -) -> ( - axum::Router, - AppState, - SqlitePool, - tempfile::TempDir, - SecretBox, -) { +) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { admin_accounts_test_app_with_overrides( db_name, key_byte, @@ -185,19 +161,12 @@ async fn admin_accounts_test_app_with_overrides( key_byte: u8, api_base_url: String, oauth_token_endpoint: Option, -) -> ( - axum::Router, - AppState, - SqlitePool, - tempfile::TempDir, - SecretBox, -) { +) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); let db = dir.path().join(db_name); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([key_byte; 32]); let hasher = ApiKeyHasher::new([key_byte; 32]); let mut config = test_config(url); config.api.base_url = api_base_url; @@ -207,7 +176,7 @@ async fn admin_accounts_test_app_with_overrides( let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box.clone()), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -221,7 +190,7 @@ async fn admin_accounts_test_app_with_overrides( services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state.clone()); - (app, state, pool, dir, secret_box) + (app, state, pool, dir) } async fn seed_account(pool: &SqlitePool, account: NewAccount) { @@ -278,6 +247,7 @@ fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: BTreeMap::new(), + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 300, @@ -304,7 +274,6 @@ fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/admin/accounts/quota.rs b/tests/admin/accounts/quota.rs index 9e03f0b9a..f44435164 100644 --- a/tests/admin/accounts/quota.rs +++ b/tests/admin/accounts/quota.rs @@ -18,12 +18,11 @@ async fn admin_usage_stats_should_return_page_and_summary() { 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 secret_box = SecretBox::new([73u8; 32]); let hasher = ApiKeyHasher::new([74u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -83,12 +82,11 @@ async fn admin_usage_stats_should_require_admin_session_cookie() { let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); let config = test_config(url); - let secret_box = SecretBox::new([125u8; 32]); let hasher = ApiKeyHasher::new([126u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -154,12 +152,11 @@ async fn admin_usage_stats_should_cursor_page_account_usage() { ) .await; let config = test_config(url); - let secret_box = SecretBox::new([127u8; 32]); let hasher = ApiKeyHasher::new([128u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -229,7 +226,7 @@ async fn admin_account_quota_should_send_usage_cookie() { }))) .mount(&server) .await; - let (app, _state, pool, _dir, secret_box) = admin_accounts_test_app_with_api_base_url( + let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( "admin-account-quota-cookie.sqlite", 88, format!("{}/backend-api", server.uri()), @@ -252,7 +249,7 @@ async fn admin_account_quota_should_send_usage_cookie() { }, ) .await; - SqliteCookieStore::new(pool.clone(), secret_box) + SqliteCookieStore::new(pool.clone()) .set_cookie_header("acct_quota_cookie", "cf_clearance=admin-quota") .await .unwrap(); @@ -292,7 +289,7 @@ async fn admin_account_health_check_should_send_probe_cookie() { }))) .mount(&server) .await; - let (app, _state, pool, _dir, secret_box) = admin_accounts_test_app_with_api_base_url( + let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( "admin-account-health-cookie.sqlite", 89, format!("{}/backend-api", server.uri()), @@ -315,7 +312,7 @@ async fn admin_account_health_check_should_send_probe_cookie() { }, ) .await; - SqliteCookieStore::new(pool.clone(), secret_box) + SqliteCookieStore::new(pool.clone()) .set_cookie_header("acct_health_cookie", "cf_clearance=admin-health") .await .unwrap(); @@ -355,7 +352,6 @@ async fn admin_account_quota_warnings_should_return_threshold_matches_from_cache let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([89u8; 32]); seed_account( &pool, NewAccount { @@ -417,7 +413,7 @@ async fn admin_account_quota_warnings_should_return_threshold_matches_from_cache let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), diff --git a/tests/admin/accounts/testing.rs b/tests/admin/accounts/testing.rs index 7419480e2..1cf20ca23 100644 --- a/tests/admin/accounts/testing.rs +++ b/tests/admin/accounts/testing.rs @@ -52,7 +52,7 @@ async fn account_test_models_should_return_upstream_models_only() { }))) .mount(&server) .await; - let (app, _state, pool, _dir, _secret_box) = admin_accounts_test_app_with_api_base_url( + let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( "admin-account-test-models.sqlite", 91, format!("{}/backend-api", server.uri()), @@ -124,7 +124,7 @@ async fn account_test_stream_should_translate_upstream_responses_sse() { ) .mount(&server) .await; - let (app, _state, pool, _dir, _secret_box) = admin_accounts_test_app_with_api_base_url( + let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( "admin-account-test-stream.sqlite", 92, format!("{}/backend-api", server.uri()), diff --git a/tests/admin/auth/session.rs b/tests/admin/auth/session.rs index c06c1d9d9..21a4b3c67 100644 --- a/tests/admin/auth/session.rs +++ b/tests/admin/auth/session.rs @@ -4,7 +4,6 @@ use axum::{ }; use chrono::Utc; use codex_proxy_rs::infra::{ - crypto::SecretBox, database::connect_sqlite, identity::{hash_admin_password, ApiKeyHasher}, }; @@ -21,17 +20,13 @@ async fn admin_login_should_issue_http_only_session_cookie() { let pool = connect_sqlite(&url).await.unwrap(); seed_admin_user(&pool, "correct-password").await; let config = test_config(url); - let secret_box = SecretBox::new([121u8; 32]); let hasher = ApiKeyHasher::new([122u8; 32]); let stores = codex_proxy_rs::runtime::services::BackgroundTaskStores { accounts: codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()), admin_sessions: codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore::new( pool.clone(), ), - cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new( - pool.clone(), - secret_box.clone(), - ), + cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new(pool.clone()), fingerprints: codex_proxy_rs::upstream::fingerprint::FingerprintRepository::new( pool.clone(), ), @@ -102,17 +97,13 @@ async fn admin_login_should_reject_client_api_key_as_password_or_authorization() let pool = connect_sqlite(&url).await.unwrap(); seed_admin_user(&pool, "correct-password").await; let config = test_config(url); - let secret_box = SecretBox::new([123u8; 32]); let hasher = ApiKeyHasher::new([124u8; 32]); let stores = codex_proxy_rs::runtime::services::BackgroundTaskStores { accounts: codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()), admin_sessions: codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore::new( pool.clone(), ), - cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new( - pool.clone(), - secret_box.clone(), - ), + cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new(pool.clone()), fingerprints: codex_proxy_rs::upstream::fingerprint::FingerprintRepository::new( pool.clone(), ), diff --git a/tests/admin/auth/session_routes.rs b/tests/admin/auth/session_routes.rs index 70bad880d..44cf250a3 100644 --- a/tests/admin/auth/session_routes.rs +++ b/tests/admin/auth/session_routes.rs @@ -16,7 +16,6 @@ use codex_proxy_rs::{ UsageStatsConfig, WebSocketPoolConfig, }, infra::{ - crypto::SecretBox, database::connect_sqlite, identity::{hash_admin_password, ApiKeyHasher}, }, @@ -39,12 +38,11 @@ async fn admin_login_should_issue_http_only_session_cookie() { let pool = connect_sqlite(&url).await.unwrap(); seed_admin_user(&pool, "correct-password").await; let config = test_config(url); - let secret_box = SecretBox::new([121u8; 32]); let hasher = ApiKeyHasher::new([122u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -106,12 +104,11 @@ async fn admin_login_should_reject_client_api_key_as_password_or_authorization() let pool = connect_sqlite(&url).await.unwrap(); seed_admin_user(&pool, "correct-password").await; let config = test_config(url); - let secret_box = SecretBox::new([123u8; 32]); let hasher = ApiKeyHasher::new([124u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -181,6 +178,7 @@ fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: Some("high".to_string()), service_tier: Some("flex".to_string()), aliases, + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 240, @@ -207,7 +205,6 @@ fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/admin/keys/mod.rs b/tests/admin/keys/mod.rs index 3b73a21a6..bdca04f5c 100644 --- a/tests/admin/keys/mod.rs +++ b/tests/admin/keys/mod.rs @@ -14,7 +14,7 @@ use codex_proxy_rs::{ QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, WebSocketPoolConfig, }, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, runtime::state::AppState, @@ -38,12 +38,11 @@ async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, tempfile::Te let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; let config = test_config(url); - let secret_box = SecretBox::new([43u8; 32]); let hasher = ApiKeyHasher::new([44u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -125,6 +124,7 @@ fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: BTreeMap::new(), + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 300, @@ -151,7 +151,6 @@ fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/admin/models/routes.rs b/tests/admin/models/routes.rs index d32c28452..25ce4128c 100644 --- a/tests/admin/models/routes.rs +++ b/tests/admin/models/routes.rs @@ -15,7 +15,7 @@ use codex_proxy_rs::{ QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, WebSocketPoolConfig, }, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, runtime::state::AppState, @@ -44,12 +44,11 @@ async fn admin_refresh_models_should_require_admin_session() { let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); let config = test_config(url, upstream.uri()); - let secret_box = SecretBox::new([81u8; 32]); let hasher = ApiKeyHasher::new([82u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -103,7 +102,6 @@ async fn admin_refresh_models_should_store_snapshots_for_distinct_active_plans() let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([83u8; 32]); let config = test_config(url, upstream.uri()); let hasher = ApiKeyHasher::new([84u8; 32]); let store = SqliteAccountStore::new(pool.clone()); @@ -142,7 +140,7 @@ async fn admin_refresh_models_should_store_snapshots_for_distinct_active_plans() let stores = BackgroundTaskStores { accounts: store, admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box.clone()), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -216,6 +214,7 @@ fn test_config(database_url: String, base_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: BTreeMap::new(), + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 300, @@ -242,7 +241,6 @@ fn test_config(database_url: String, base_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/admin/monitoring/dashboard.rs b/tests/admin/monitoring/dashboard.rs index 4d6ac775f..8838d51ab 100644 --- a/tests/admin/monitoring/dashboard.rs +++ b/tests/admin/monitoring/dashboard.rs @@ -12,9 +12,7 @@ use codex_proxy_rs::{ events::{EventLevel, EventLog}, }, }, - infra::{ - crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher, time::china_day_start, - }, + infra::{database::connect_sqlite, identity::ApiKeyHasher, time::china_day_start}, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::{ services::{BackgroundTaskStores, Services}, @@ -63,7 +61,7 @@ async fn dashboard_summary_should_render_dash_when_fingerprint_updated_at_is_mis } #[tokio::test] -async fn dashboard_summary_should_calculate_today_traffic_and_latency_from_event_logs() { +async fn dashboard_summary_should_calculate_today_traffic_and_latency_from_time_buckets() { let (app, store, _pool, _dir) = dashboard_test_app( "dashboard-traffic-latency.sqlite", crate::support::fingerprint::test_fingerprint(), @@ -113,6 +111,34 @@ async fn dashboard_summary_should_calculate_today_traffic_and_latency_from_event ); } +#[tokio::test] +async fn dashboard_summary_should_keep_trend_after_event_logs_are_cleared() { + let (app, store, _pool, _dir) = dashboard_test_app( + "dashboard-time-buckets-survive-log-clear.sqlite", + crate::support::fingerprint::test_fingerprint(), + ) + .await; + let mut log = usage_log_with_tokens(Utc::now(), 12); + log.latency_ms = Some(900); + store.append(&log).await.unwrap(); + store.clear().await.unwrap(); + + let body = dashboard_summary(app).await; + + assert_eq!(body["data"]["eventLogs"].as_array().unwrap().len(), 0); + assert_eq!(body["data"]["cards"]["traffic"]["todayRequests"], 1); + assert_eq!(body["data"]["cards"]["tokens"]["todayTokens"], 12); + assert_eq!( + body["data"]["trend"]["points"] + .as_array() + .unwrap() + .iter() + .map(|point| point["requests"].as_u64().unwrap()) + .sum::(), + 1 + ); +} + #[tokio::test] async fn dashboard_summary_should_calculate_priority_cost_from_event_service_tier() { let cost = dashboard_summary_total_cost_for_log( @@ -280,12 +306,11 @@ async fn dashboard_test_app( let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; let config = crate::support::config::test_config(url); - let secret_box = SecretBox::new([73u8; 32]); let hasher = ApiKeyHasher::new([74u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), diff --git a/tests/admin/monitoring/logs_routes.rs b/tests/admin/monitoring/logs_routes.rs index f9689e409..b3296a00c 100644 --- a/tests/admin/monitoring/logs_routes.rs +++ b/tests/admin/monitoring/logs_routes.rs @@ -18,7 +18,7 @@ use codex_proxy_rs::{ QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, WebSocketPoolConfig, }, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, runtime::state::AppState, @@ -258,12 +258,11 @@ async fn admin_logs_test_app( let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; let config = test_config(url); - let secret_box = SecretBox::new([73u8; 32]); let hasher = ApiKeyHasher::new([74u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -325,6 +324,7 @@ fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: BTreeMap::new(), + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 300, @@ -351,7 +351,6 @@ fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/admin/settings/mod.rs b/tests/admin/settings/mod.rs index bf1af306e..0a7e45175 100644 --- a/tests/admin/settings/mod.rs +++ b/tests/admin/settings/mod.rs @@ -17,6 +17,22 @@ fn settings_service_should_apply_retained_settings_patch() { &mut current, AdminSettingsPatch { default_model: Some(" gpt-6 ".to_string()), + model_aliases: Some( + [ + (" claude-sonnet ".to_string(), " gpt-5.5 ".to_string()), + ("openai-fast".to_string(), "openai:gpt-4o".to_string()), + ] + .into(), + ), + model_account_routes: Some( + [( + " gpt-5.5 ".to_string(), + vec![" acct_a ".to_string(), "acct_b".to_string()], + )] + .into(), + ), + refresh_margin_seconds: Some(600), + refresh_concurrency: Some(3), max_concurrent_per_account: Some(5), request_interval_ms: Some(125), rotation_strategy: Some("round_robin".to_string()), @@ -25,12 +41,41 @@ fn settings_service_should_apply_retained_settings_patch() { .expect("settings patch should be valid"); assert_eq!(current.default_model, "gpt-6"); + assert_eq!(current.model_aliases["claude-sonnet"], "gpt-5.5"); + assert_eq!(current.model_aliases["openai-fast"], "openai:gpt-4o"); + assert_eq!( + current.model_account_routes["gpt-5.5"], + vec!["acct_a".to_string(), "acct_b".to_string()] + ); + assert_eq!(current.refresh_margin_seconds, 600); + assert_eq!(current.refresh_concurrency, 3); assert_eq!(current.rotation_strategy, "round_robin"); assert_eq!(current.max_concurrent_per_account, 5); assert_eq!(current.request_interval_ms, 125); assert_eq!(current.default_reasoning_effort.as_deref(), Some("high")); } +#[test] +fn settings_service_should_reject_invalid_model_account_routes() { + let mut current = AdminSettings::default(); + let error = SettingsService::apply_patch( + &mut current, + AdminSettingsPatch { + model_account_routes: Some( + [( + "gpt-5.5".to_string(), + vec!["acct_a".to_string(), "acct_a".to_string()], + )] + .into(), + ), + ..AdminSettingsPatch::default() + }, + ) + .expect_err("duplicate model route account should be rejected"); + + assert_eq!(error.field(), "modelAccountRoutes"); +} + #[test] fn settings_service_should_reject_invalid_patch_values() { let mut current = AdminSettings::default(); @@ -45,3 +90,18 @@ fn settings_service_should_reject_invalid_patch_values() { assert_eq!(error.field(), "rotationStrategy"); } + +#[test] +fn settings_service_should_reject_invalid_model_aliases() { + let mut current = AdminSettings::default(); + let error = SettingsService::apply_patch( + &mut current, + AdminSettingsPatch { + model_aliases: Some([("gpt-5.5".to_string(), "gpt-5.5".to_string())].into()), + ..AdminSettingsPatch::default() + }, + ) + .expect_err("self-referencing model alias should be rejected"); + + assert_eq!(error.field(), "modelAliases"); +} diff --git a/tests/admin/settings/routes.rs b/tests/admin/settings/routes.rs index aa884108d..1793a17d1 100644 --- a/tests/admin/settings/routes.rs +++ b/tests/admin/settings/routes.rs @@ -1,4 +1,4 @@ -use std::{collections::BTreeMap, fs, path::Path}; +use std::collections::BTreeMap; use axum::{ body::{to_bytes, Body}, @@ -8,13 +8,13 @@ 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, ModelConfig, QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, WebSocketPoolConfig, }, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + config::{settings::RuntimeSettingsService, types::AppConfig}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, runtime::state::AppState, @@ -61,6 +61,11 @@ async fn admin_settings_should_return_runtime_fields() { .unwrap(); let body = response_json(response).await; assert_eq!(body["data"]["defaultModel"], "gpt-5.5"); + assert_eq!(body["data"]["modelAliases"]["codex-fast"], "gpt-5.5"); + assert_eq!(body["data"]["modelAccountRoutes"], json!({})); + assert_eq!(body["data"]["refreshMarginSeconds"], 240); + assert_eq!(body["data"]["refreshConcurrency"], 2); + assert_eq!(body["data"]["rotationStrategy"], "least_used"); } #[tokio::test] @@ -82,20 +87,20 @@ async fn admin_settings_update_should_require_admin_session_cookie() { } #[tokio::test] -async fn admin_settings_update_should_persist_retained_fields_to_config_yaml() { +async fn admin_settings_update_should_persist_runtime_settings_to_database() { let dir = tempfile::tempdir().unwrap(); let db = dir.path().join("admin-settings-update.sqlite"); let url = format!("sqlite://{}", db.display()); let config = test_config(url); - write_config_yaml(dir.path(), &config); let pool = connect_sqlite(&config.database.url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([53u8; 32]); + seed_account(&pool, "acct_route_a").await; + seed_account(&pool, "acct_route_b").await; let hasher = ApiKeyHasher::new([54u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -103,13 +108,7 @@ async fn admin_settings_update_should_persist_retained_fields_to_config_yaml() { event_logs: SqliteEventLogStore::new(pool.clone()), }; let fingerprint = crate::support::fingerprint::test_fingerprint(); - let mut services = Services::new(&config, stores, fingerprint); - services.settings = std::sync::Arc::new( - codex_proxy_rs::config::settings::RuntimeSettingsService::with_config_path( - config.clone(), - dir.path().join("config.yaml"), - ), - ); + let services = Services::new(&config, stores, fingerprint); let services = std::sync::Arc::new(services); let state = AppState { config: config.clone(), @@ -129,6 +128,15 @@ async fn admin_settings_update_should_persist_retained_fields_to_config_yaml() { .body(Body::from( json!({ "defaultModel": "gpt-6", "rotationStrategy": "round_robin", + "modelAliases": { + "gpt-5.2": "gpt-5.5", + "claude-sonnet": "gpt-5.5" + }, + "modelAccountRoutes": { + "gpt-5.5": ["acct_route_a", "acct_route_b"] + }, + "refreshMarginSeconds": 900, + "refreshConcurrency": 4, "maxConcurrentPerAccount": 7, "requestIntervalMs": 80 }) @@ -141,6 +149,13 @@ async fn admin_settings_update_should_persist_retained_fields_to_config_yaml() { assert_eq!(response.status(), StatusCode::OK); let body = response_json(response).await; assert_eq!(body["data"]["defaultModel"], "gpt-6"); + assert_eq!(body["data"]["modelAliases"]["gpt-5.2"], "gpt-5.5"); + assert_eq!( + body["data"]["modelAccountRoutes"]["gpt-5.5"][0], + "acct_route_a" + ); + assert_eq!(body["data"]["refreshMarginSeconds"], 900); + assert_eq!(body["data"]["refreshConcurrency"], 4); let get_response = app .oneshot( @@ -158,17 +173,55 @@ async fn admin_settings_update_should_persist_retained_fields_to_config_yaml() { "gpt-6" ); - let persisted = AppConfig::load_from_dir(dir.path()).unwrap(); - assert_eq!(persisted.model.default_model, "gpt-6"); - assert_eq!(persisted.auth.rotation_strategy, "round_robin"); - assert_eq!(persisted.auth.max_concurrent_per_account, 7); - assert_eq!(persisted.auth.request_interval_ms, 80); - assert!(persisted.auth.refresh_enabled); - assert_eq!(persisted.database.url, config.database.url); + let row: (String, String, i64, i64, i64, i64, String) = sqlx::query_as( + "select default_model, model_aliases_json, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy from runtime_settings where id = 1", + ) + .fetch_one(&pool) + .await + .unwrap(); + let aliases: BTreeMap = serde_json::from_str(&row.1).unwrap(); + assert_eq!(row.0, "gpt-6"); + assert_eq!(aliases["gpt-5.2"], "gpt-5.5"); + assert_eq!(aliases["claude-sonnet"], "gpt-5.5"); + assert_eq!(row.2, 900); + assert_eq!(row.3, 4); + assert_eq!(row.4, 7); + assert_eq!(row.5, 80); + assert_eq!(row.6, "round_robin"); + let route_rows: Vec<(String, String, i64)> = sqlx::query_as( + "select model, account_id, priority from model_account_routes order by model, priority", + ) + .fetch_all(&pool) + .await + .unwrap(); assert_eq!( - persisted.security.master_key_file, - config.security.master_key_file + route_rows, + vec![ + ("gpt-5.5".to_string(), "acct_route_a".to_string(), 0), + ("gpt-5.5".to_string(), "acct_route_b".to_string(), 1), + ] ); + assert!(!dir.path().join("config.yaml").exists()); + + let restarted_config = RuntimeSettingsService::load_or_initialize_config(config.clone(), &pool) + .await + .unwrap(); + assert_eq!(restarted_config.model.default_model, "gpt-6"); + assert_eq!(restarted_config.model.aliases["gpt-5.2"], "gpt-5.5"); + assert_eq!( + restarted_config.model.account_routes["gpt-5.5"], + vec!["acct_route_a".to_string(), "acct_route_b".to_string()] + ); + assert_eq!(restarted_config.auth.refresh_margin_seconds, 900); + assert_eq!(restarted_config.auth.refresh_concurrency, 4); + assert_eq!(restarted_config.auth.rotation_strategy, "round_robin"); + assert_eq!(restarted_config.auth.max_concurrent_per_account, 7); + assert_eq!(restarted_config.auth.request_interval_ms, 80); + assert_eq!( + restarted_config.model.default_reasoning_effort, + config.model.default_reasoning_effort + ); + assert_eq!(restarted_config.database.url, config.database.url); } #[tokio::test] @@ -177,15 +230,13 @@ async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() { let db = dir.path().join("admin-settings-update-invalid.sqlite"); let url = format!("sqlite://{}", db.display()); let config = test_config(url); - write_config_yaml(dir.path(), &config); let pool = connect_sqlite(&config.database.url).await.unwrap(); seed_admin_session(&pool, "session_1").await; - let secret_box = SecretBox::new([53u8; 32]); let hasher = ApiKeyHasher::new([54u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -225,12 +276,11 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp let pool = connect_sqlite(&url).await.unwrap(); seed_admin_session(&pool, "session_1").await; let config = test_config(url); - let secret_box = SecretBox::new([53u8; 32]); let hasher = ApiKeyHasher::new([54u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -272,19 +322,31 @@ async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { .unwrap(); } +async fn seed_account(pool: &SqlitePool, account_id: &str) { + sqlx::query( + r" +insert into accounts ( + id, + access_token, + status, + added_at, + updated_at +) values (?, ?, 'active', ?, ?)", + ) + .bind(account_id) + .bind("access-token") + .bind("2026-06-18T00: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 write_config_yaml(path: &Path, config: &AppConfig) { - fs::write( - path.join("config.yaml"), - serde_yml::to_string(config).unwrap(), - ) - .unwrap(); -} - fn test_config(database_url: String) -> AppConfig { let mut aliases = BTreeMap::new(); aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string()); @@ -301,6 +363,7 @@ fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: Some("high".to_string()), service_tier: Some("flex".to_string()), aliases, + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 240, @@ -327,7 +390,6 @@ fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/config/mod.rs b/tests/config/mod.rs index 02c2a04a6..b821c14a6 100644 --- a/tests/config/mod.rs +++ b/tests/config/mod.rs @@ -13,6 +13,7 @@ model: default_reasoning_effort: null default_service_tier: null aliases: {} + account_routes: {} auth: refresh_margin_seconds: 300 refresh_enabled: true @@ -39,7 +40,6 @@ usage_stats: database: url: sqlite://.runtime/data/codex-proxy-rs.sqlite security: - master_key_file: .runtime/data/master.key api_key_pepper_file: .runtime/data/api-key-pepper.key tls: force_http11: false @@ -132,7 +132,6 @@ fn default_config_keeps_runtime_artifacts_under_runtime_directory() { cfg.database.url, "sqlite://.runtime/data/codex-proxy-rs.sqlite" ); - assert_eq!(cfg.security.master_key_file, ".runtime/data/master.key"); assert_eq!( cfg.security.api_key_pepper_file, ".runtime/data/api-key-pepper.key" @@ -156,6 +155,7 @@ model: default_reasoning_effort: null default_service_tier: null aliases: {} + account_routes: {} auth: refresh_margin_seconds: 300 refresh_enabled: true @@ -182,7 +182,6 @@ usage_stats: database: url: sqlite://.runtime/data/codex-proxy-rs.sqlite security: - master_key_file: .runtime/data/master.key api_key_pepper_file: .runtime/data/api-key-pepper.key tls: force_http11: false diff --git a/tests/http/web_assets/mod.rs b/tests/http/web_assets/mod.rs index afeaeab36..cd4a633f9 100644 --- a/tests/http/web_assets/mod.rs +++ b/tests/http/web_assets/mod.rs @@ -8,7 +8,7 @@ use codex_proxy_rs::{ QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, WebSocketPoolConfig, }, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, }; use serde_json::Value; use tower::ServiceExt; @@ -41,17 +41,13 @@ async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes let database_url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&database_url).await.expect("sqlite pool"); let config = test_config(database_url); - let secret_box = SecretBox::new([61u8; 32]); let hasher = ApiKeyHasher::new([62u8; 32]); let stores = codex_proxy_rs::runtime::services::BackgroundTaskStores { accounts: codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()), admin_sessions: codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore::new( pool.clone(), ), - cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new( - pool.clone(), - secret_box.clone(), - ), + cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new(pool.clone()), fingerprints: codex_proxy_rs::upstream::fingerprint::FingerprintRepository::new( pool.clone(), ), @@ -147,6 +143,7 @@ fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: Default::default(), + account_routes: Default::default(), }, auth: AuthConfig { refresh_margin_seconds: 240, @@ -173,7 +170,6 @@ fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.file".to_string(), }, tls: TlsConfig { diff --git a/tests/infra/crypto/mod.rs b/tests/infra/crypto/mod.rs deleted file mode 100644 index 300f27b06..000000000 --- a/tests/infra/crypto/mod.rs +++ /dev/null @@ -1,31 +0,0 @@ -use codex_proxy_rs::infra::crypto::SecretBox; -use secrecy::{ExposeSecret, SecretString}; - -#[test] -fn secret_box_encrypts_and_decrypts_without_storing_plaintext() { - let secret_box = SecretBox::new([7u8; 32]); - let plaintext = SecretString::new("rt_example_refresh_token".to_string().into()); - let ciphertext = secret_box.encrypt(&plaintext).unwrap(); - - assert!(ciphertext.starts_with("v1:")); - assert!(!ciphertext.contains("rt_example_refresh_token")); - assert_eq!( - secret_box.decrypt(&ciphertext).unwrap().expose_secret(), - "rt_example_refresh_token" - ); -} - -#[test] -fn secret_box_load_or_create_should_reuse_existing_master_key_file() { - let dir = tempfile::tempdir().unwrap(); - let key_path = dir.path().join("master.key"); - let first_box = SecretBox::load_or_create(&key_path).unwrap(); - let ciphertext = first_box - .encrypt(&SecretString::new("persisted-secret".to_string().into())) - .unwrap(); - - let second_box = SecretBox::load_or_create(&key_path).unwrap(); - let decrypted = second_box.decrypt(&ciphertext).unwrap(); - - assert_eq!(decrypted.expose_secret(), "persisted-secret"); -} diff --git a/tests/infra/mod.rs b/tests/infra/mod.rs index 1a12e5984..03bb9c9dd 100644 --- a/tests/infra/mod.rs +++ b/tests/infra/mod.rs @@ -1,4 +1,3 @@ -pub mod crypto; pub mod log_rotation; pub mod storage_schema; pub mod time; diff --git a/tests/infra/storage_schema/mod.rs b/tests/infra/storage_schema/mod.rs index f16de1844..b66cf400f 100644 --- a/tests/infra/storage_schema/mod.rs +++ b/tests/infra/storage_schema/mod.rs @@ -1,19 +1,232 @@ use codex_proxy_rs::infra::database::connect_sqlite; #[tokio::test] -async fn sqlite_schema_creates_accounts_and_event_tables() { +async fn sqlite_schema_creates_current_runtime_tables_without_migration_table() { 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', 'event_logs', 'model_plan_snapshots')", + "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')", ) .fetch_one(&pool) .await .unwrap(); - assert_eq!(row.0, 4); + 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); +} + +#[tokio::test] +async fn sqlite_schema_should_persist_runtime_settings_columns() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("settings.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let rows: Vec<(String,)> = sqlx::query_as( + "select name from pragma_table_info('runtime_settings') where name in ('id', 'default_model', 'model_aliases_json', 'refresh_margin_seconds', 'refresh_concurrency', 'max_concurrent_per_account', 'request_interval_ms', 'rotation_strategy', 'updated_at') order by name", + ) + .fetch_all(&pool) + .await + .unwrap(); + + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "default_model", + "id", + "max_concurrent_per_account", + "model_aliases_json", + "refresh_concurrency", + "refresh_margin_seconds", + "request_interval_ms", + "rotation_strategy", + "updated_at", + ] + ); +} + +#[tokio::test] +async fn sqlite_schema_should_persist_runtime_query_indexes() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("indexes.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let rows: Vec<(String,)> = sqlx::query_as( + "select name from sqlite_master where type = 'index' and name in ('idx_accounts_added_id', 'idx_account_usage_last_used_account', 'idx_account_model_usage_last_used', 'idx_client_api_keys_created_id', 'idx_fingerprint_update_history_created_id') order by name", + ) + .fetch_all(&pool) + .await + .unwrap(); + + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "idx_account_model_usage_last_used", + "idx_account_usage_last_used_account", + "idx_accounts_added_id", + "idx_client_api_keys_created_id", + "idx_fingerprint_update_history_created_id", + ] + ); +} + +#[tokio::test] +async fn sqlite_schema_should_persist_diagnostic_filter_indexes() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("diagnostic-indexes.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let rows: Vec<(String,)> = sqlx::query_as( + "select name from sqlite_master where type = 'index' and name in ('idx_account_cookies_expires', 'idx_event_logs_level_created', 'idx_event_logs_model_created', 'idx_event_logs_route_created', 'idx_event_logs_status_created', 'idx_event_logs_upstream_status_created', 'idx_session_affinities_active_order', 'idx_model_account_routes_account', 'idx_model_account_routes_enabled_model', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name", + ) + .fetch_all(&pool) + .await + .unwrap(); + + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "idx_account_cookies_expires", + "idx_event_logs_level_created", + "idx_event_logs_model_created", + "idx_event_logs_route_created", + "idx_event_logs_status_created", + "idx_event_logs_upstream_status_created", + "idx_model_account_routes_account", + "idx_model_account_routes_enabled_model", + "idx_session_affinities_active_order", + "idx_usage_time_buckets_bucket", + "idx_usage_time_buckets_model_bucket", + ] + ); +} + +#[tokio::test] +async fn sqlite_schema_should_persist_account_model_usage_columns() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("model-usage.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let rows: Vec<(String,)> = sqlx::query_as( + "select name from pragma_table_info('account_model_usage') where name in ('account_id', 'model', 'request_count', 'error_count', 'input_tokens', 'output_tokens', 'cached_tokens', 'last_used_at') order by name", + ) + .fetch_all(&pool) + .await + .unwrap(); + + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "account_id", + "cached_tokens", + "error_count", + "input_tokens", + "last_used_at", + "model", + "output_tokens", + "request_count", + ] + ); +} + +#[tokio::test] +async fn sqlite_schema_should_persist_usage_time_bucket_columns() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage-time-buckets.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let rows: Vec<(String,)> = sqlx::query_as( + "select name from pragma_table_info('usage_time_buckets') where name in ('bucket_start', 'account_id', 'model', 'service_tier', 'request_count', 'error_count', 'input_tokens', 'output_tokens', 'cached_tokens', 'first_token_latency_sum', 'first_token_latency_count', 'latency_sum', 'latency_count', 'max_latency_ms', 'min_latency_ms', 'updated_at') order by name", + ) + .fetch_all(&pool) + .await + .unwrap(); + + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "account_id", + "bucket_start", + "cached_tokens", + "error_count", + "first_token_latency_count", + "first_token_latency_sum", + "input_tokens", + "latency_count", + "latency_sum", + "max_latency_ms", + "min_latency_ms", + "model", + "output_tokens", + "request_count", + "service_tier", + "updated_at", + ] + ); +} + +#[tokio::test] +async fn sqlite_schema_should_persist_model_account_routes_columns() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("model-account-routes.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let rows: Vec<(String,)> = sqlx::query_as( + "select name from pragma_table_info('model_account_routes') where name in ('model', 'account_id', 'priority', 'enabled', 'created_at', 'updated_at') order by name", + ) + .fetch_all(&pool) + .await + .unwrap(); + + assert_eq!( + rows.into_iter().map(|row| row.0).collect::>(), + [ + "account_id", + "created_at", + "enabled", + "model", + "priority", + "updated_at", + ] + ); +} + +#[tokio::test] +async fn sqlite_schema_should_persist_plain_account_cookie_value_column() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("cookies.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let value_column: Option<(String,)> = sqlx::query_as( + "select name from pragma_table_info('account_cookies') where name = 'value'", + ) + .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] @@ -279,10 +492,88 @@ async fn sqlite_schema_should_reject_non_boolean_flags() { ) .execute(&pool) .await; + sqlx::query( + "insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_route', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await + .unwrap(); + let model_route_result = sqlx::query( + "insert into model_account_routes (model, account_id, enabled, created_at, updated_at) values ('gpt-5.5', 'acct_route', 2, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; assert!(api_key_result.is_err()); assert!(account_result.is_err()); assert!(quota_verify_result.is_err()); + assert!(model_route_result.is_err()); +} + +#[tokio::test] +async fn sqlite_schema_should_reject_invalid_runtime_settings() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("settings-check.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let invalid_id_result = sqlx::query( + "insert into runtime_settings (id, default_model, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (2, 'gpt-5.5', 300, 2, 3, 50, 'least_used', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + let invalid_refresh_result = sqlx::query( + "insert into runtime_settings (id, default_model, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (1, 'gpt-5.5', 0, 2, 3, 50, 'least_used', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + let invalid_strategy_result = sqlx::query( + "insert into runtime_settings (id, default_model, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (1, 'gpt-5.5', 300, 2, 3, 50, 'random', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + + assert!(invalid_id_result.is_err()); + assert!(invalid_refresh_result.is_err()); + assert!(invalid_strategy_result.is_err()); +} + +#[tokio::test] +async fn sqlite_schema_should_reject_invalid_model_account_route_priority() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("model-route-check.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + sqlx::query( + "insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_route', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await + .unwrap(); + let result = sqlx::query( + "insert into model_account_routes (model, account_id, priority, created_at, updated_at) values ('gpt-5.5', 'acct_route', -1, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + + assert!(result.is_err()); +} + +#[tokio::test] +async fn sqlite_schema_should_reject_negative_usage_time_bucket_counts() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("usage-time-bucket-check.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + let result = sqlx::query( + "insert into usage_time_buckets (bucket_start, request_count, updated_at) values ('2026-06-14T00:00:00Z', -1, '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await; + + assert!(result.is_err()); } #[tokio::test] @@ -306,6 +597,28 @@ async fn sqlite_schema_should_reject_negative_account_usage_counts() { assert!(result.is_err()); } +#[tokio::test] +async fn sqlite_schema_should_reject_negative_account_model_usage_counts() { + let dir = tempfile::tempdir().unwrap(); + let db = dir.path().join("model-usage-check.sqlite"); + let url = format!("sqlite://{}", db.display()); + let pool = connect_sqlite(&url).await.unwrap(); + + sqlx::query( + "insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_a', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + ) + .execute(&pool) + .await + .unwrap(); + let result = sqlx::query( + "insert into account_model_usage (account_id, model, request_count) values ('acct_a', 'gpt-5.5', -1)", + ) + .execute(&pool) + .await; + + assert!(result.is_err()); +} + #[tokio::test] async fn connect_sqlite_creates_parent_directory_for_file_database() { let dir = tempfile::tempdir().unwrap(); diff --git a/tests/proxy/dispatch/chat_upstream/chat_routes.rs b/tests/proxy/dispatch/chat_upstream/chat_routes.rs index 95f845b01..8b92b2226 100644 --- a/tests/proxy/dispatch/chat_upstream/chat_routes.rs +++ b/tests/proxy/dispatch/chat_upstream/chat_routes.rs @@ -77,6 +77,14 @@ async fn chat_completions_should_dispatch_to_codex_and_return_openai_response() .await .unwrap(); assert_eq!(usage, (1, 9, 3)); + let model_usage: (String, i64, i64, i64) = sqlx::query_as( + "select model, request_count, input_tokens, output_tokens from account_model_usage where account_id = ?", + ) + .bind("acct_chat") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(model_usage, ("gpt-5.5".to_string(), 1, 9, 3)); let event = latest_event_log(&pool, "v1.chat").await; let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); assert_eq!(event.request_id.as_deref(), Some("req_chat_nonstream_log")); @@ -743,7 +751,7 @@ async fn chat_completions_should_cool_down_cloudflare_403_and_fallback() { .await; let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store .set_cookie_header("acct_primary", "cf_clearance=old") .await diff --git a/tests/proxy/dispatch/chat_upstream/mod.rs b/tests/proxy/dispatch/chat_upstream/mod.rs index 479f29830..ca0b9169b 100644 --- a/tests/proxy/dispatch/chat_upstream/mod.rs +++ b/tests/proxy/dispatch/chat_upstream/mod.rs @@ -26,7 +26,7 @@ use codex_proxy_rs::{ UsageStatsConfig, WebSocketPoolConfig, }, http::router, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, proxy::dispatch::session_affinity::{ build_conversation_identity, SessionAffinityEntry, SqliteSessionAffinityStore, }, @@ -239,14 +239,13 @@ mod usage_logging; fn test_app_state_with_pool_secret_api_key_hasher_and_installation_id( config: AppConfig, pool: SqlitePool, - secret_box: SecretBox, hasher: ApiKeyHasher, installation_id: String, ) -> AppState { let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -274,14 +273,12 @@ async fn test_app_with_account_and_installation_id( let db = dir.path().join("openai-chat-upstream.sqlite"); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([83u8; 32]); let hasher = ApiKeyHasher::new([84u8; 32]); let api_key = insert_client_api_key(&pool, &hasher).await; insert_account(&pool).await; let state = test_app_state_with_pool_secret_api_key_hasher_and_installation_id( test_config(url, base_url), pool, - secret_box, hasher, installation_id, ); @@ -305,13 +302,11 @@ async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, t let db = dir.path().join("openai-responses-no-accounts.sqlite"); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([83u8; 32]); let hasher = ApiKeyHasher::new([84u8; 32]); let api_key = insert_client_api_key(&pool, &hasher).await; let state = test_app_state_with_pool_secret_api_key_hasher_and_installation_id( test_config(url, base_url), pool, - secret_box, hasher, TEST_INSTALLATION_ID.to_string(), ); @@ -351,7 +346,6 @@ async fn test_app_with_account_pool_config( let db = dir.path().join("openai-record-affinity.sqlite"); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([83u8; 32]); let hasher = ApiKeyHasher::new([84u8; 32]); let api_key = insert_client_api_key(&pool, &hasher).await; insert_account(&pool).await; @@ -360,7 +354,6 @@ async fn test_app_with_account_pool_config( let state = test_app_state_with_pool_secret_api_key_hasher_and_installation_id( config, pool.clone(), - secret_box, hasher, TEST_INSTALLATION_ID.to_string(), ); @@ -380,14 +373,12 @@ async fn test_app_with_restored_pool_then_disabled_account( let db = dir.path().join("openai-chat-restored-pool.sqlite"); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([83u8; 32]); let hasher = ApiKeyHasher::new([84u8; 32]); let api_key = insert_client_api_key(&pool, &hasher).await; insert_account(&pool).await; let state = test_app_state_with_pool_secret_api_key_hasher_and_installation_id( test_config(url, base_url), pool.clone(), - secret_box, hasher, TEST_INSTALLATION_ID.to_string(), ); @@ -412,7 +403,6 @@ async fn test_app_with_two_accounts_and_affinity( let db = dir.path().join("openai-responses-affinity.sqlite"); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([83u8; 32]); let hasher = ApiKeyHasher::new([84u8; 32]); let api_key = insert_client_api_key(&pool, &hasher).await; insert_named_account(&pool, "acct_a", "access-default", "chatgpt-default").await; @@ -438,7 +428,6 @@ async fn test_app_with_two_accounts_and_affinity( let state = test_app_state_with_pool_secret_api_key_hasher_and_installation_id( test_config(url, base_url), pool, - secret_box, hasher, TEST_INSTALLATION_ID.to_string(), ); @@ -467,7 +456,6 @@ async fn test_app_with_two_accounts_and_affinity_status( )); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([83u8; 32]); let hasher = ApiKeyHasher::new([84u8; 32]); let api_key = insert_client_api_key(&pool, &hasher).await; insert_named_account(&pool, "acct_primary", "access-primary", "chatgpt-primary").await; @@ -505,7 +493,6 @@ async fn test_app_with_two_accounts_and_affinity_status( let state = test_app_state_with_pool_secret_api_key_hasher_and_installation_id( test_config(url, base_url), pool.clone(), - secret_box, hasher, TEST_INSTALLATION_ID.to_string(), ); @@ -531,7 +518,6 @@ async fn test_app_with_two_accounts( let db = dir.path().join("openai-responses-fallback.sqlite"); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); - let secret_box = SecretBox::new([83u8; 32]); let hasher = ApiKeyHasher::new([84u8; 32]); let api_key = insert_client_api_key(&pool, &hasher).await; insert_named_account(&pool, "acct_primary", "access-primary", "chatgpt-primary").await; @@ -545,7 +531,6 @@ async fn test_app_with_two_accounts( let state = test_app_state_with_pool_secret_api_key_hasher_and_installation_id( test_config(url, base_url), pool.clone(), - secret_box, hasher, TEST_INSTALLATION_ID.to_string(), ); @@ -1601,6 +1586,7 @@ fn test_config(database_url: String, base_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: BTreeMap::new(), + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 300, @@ -1627,7 +1613,6 @@ fn test_config(database_url: String, base_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/proxy/dispatch/chat_upstream/responses_http.rs b/tests/proxy/dispatch/chat_upstream/responses_http.rs index 9c99eab3e..7a97e34a1 100644 --- a/tests/proxy/dispatch/chat_upstream/responses_http.rs +++ b/tests/proxy/dispatch/chat_upstream/responses_http.rs @@ -806,7 +806,7 @@ async fn responses_should_scope_upstream_cookie_by_codex_response_path() { .await; let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store .capture_set_cookie( "acct_chat", diff --git a/tests/proxy/dispatch/chat_upstream/responses_recovery.rs b/tests/proxy/dispatch/chat_upstream/responses_recovery.rs index 77d2f89dc..78b8c53ff 100644 --- a/tests/proxy/dispatch/chat_upstream/responses_recovery.rs +++ b/tests/proxy/dispatch/chat_upstream/responses_recovery.rs @@ -484,7 +484,7 @@ async fn responses_should_return_cloudflare_challenge_error_when_403_fallback_is .await; let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store .set_cookie_header("acct_chat", "cf_clearance=old") .await @@ -560,7 +560,7 @@ async fn responses_should_cool_down_cloudflare_403_and_fallback() { .await; let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store .set_cookie_header("acct_primary", "cf_clearance=old") .await @@ -646,7 +646,7 @@ async fn responses_should_clear_cookies_after_cloudflare_path_block_404_and_fall .await; let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store .set_cookie_header("acct_primary", "cf_clearance=old") .await @@ -715,7 +715,7 @@ async fn responses_should_disable_account_after_three_cloudflare_path_block_404s .await; let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store .set_cookie_header("acct_chat", "cf_clearance=old") .await @@ -1597,7 +1597,7 @@ async fn responses_stream_should_cool_down_cloudflare_403_and_fallback() { .await; let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store .set_cookie_header("acct_primary", "cf_clearance=old") .await @@ -1680,7 +1680,7 @@ async fn responses_stream_should_return_cloudflare_path_block_error_when_404_fal .await; let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])); + let cookie_store = SqliteCookieStore::new(pool.clone()); cookie_store .set_cookie_header("acct_chat", "cf_clearance=old") .await diff --git a/tests/proxy/dispatch/chat_upstream/responses_websocket.rs b/tests/proxy/dispatch/chat_upstream/responses_websocket.rs index eae203058..e41412bd5 100644 --- a/tests/proxy/dispatch/chat_upstream/responses_websocket.rs +++ b/tests/proxy/dispatch/chat_upstream/responses_websocket.rs @@ -897,7 +897,6 @@ async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_rest let restored_state = test_app_state_with_pool_secret_api_key_hasher_and_installation_id( test_config(format!("sqlite://{}", db.display()), server_base_url), pool.clone(), - SecretBox::new([83u8; 32]), ApiKeyHasher::new([84u8; 32]), TEST_INSTALLATION_ID.to_string(), ); diff --git a/tests/proxy/dispatch/chat_upstream/usage_logging.rs b/tests/proxy/dispatch/chat_upstream/usage_logging.rs index 31b22ca81..63e55a9c2 100644 --- a/tests/proxy/dispatch/chat_upstream/usage_logging.rs +++ b/tests/proxy/dispatch/chat_upstream/usage_logging.rs @@ -56,7 +56,16 @@ async fn responses_should_use_imported_account_record_usage_cookie_and_event_log .await .unwrap(); assert_eq!(usage, (1, 7, 4, 2)); - let cookie_header = SqliteCookieStore::new(pool.clone(), SecretBox::new([83u8; 32])) + let model_usage: (i64, i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens, cached_tokens from account_model_usage where account_id = ? and model = ?", + ) + .bind("acct_chat") + .bind("gpt-5.5") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(model_usage, (1, 7, 4, 2)); + let cookie_header = SqliteCookieStore::new(pool.clone()) .cookie_header("acct_chat", "chatgpt.com") .await .unwrap(); @@ -398,6 +407,14 @@ async fn responses_stream_should_proxy_sse_and_record_usage() { .fetch_one(&pool) .await .unwrap(); + let model_usage: (i64, i64, i64) = sqlx::query_as( + "select request_count, input_tokens, output_tokens from account_model_usage where account_id = ? and model = ?", + ) + .bind("acct_chat") + .bind("gpt-5.5") + .fetch_one(&pool) + .await + .unwrap(); assert_eq!(status, StatusCode::OK); assert!(content_type.starts_with("text/event-stream")); @@ -408,6 +425,7 @@ async fn responses_stream_should_proxy_sse_and_record_usage() { "stream responses should terminate clients, body was {body:?}" ); assert_eq!(usage, (1, 3, 5)); + assert_eq!(model_usage, (1, 3, 5)); } #[tokio::test] diff --git a/tests/proxy/openai/models_routes.rs b/tests/proxy/openai/models_routes.rs index 2a67b516a..7ee871c7b 100644 --- a/tests/proxy/openai/models_routes.rs +++ b/tests/proxy/openai/models_routes.rs @@ -14,7 +14,7 @@ use codex_proxy_rs::{ QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, WebSocketPoolConfig, }, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, runtime::state::AppState, @@ -32,12 +32,11 @@ async fn models_route_should_reject_unknown_client_api_key() { let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); let config = test_config(url); - let secret_box = SecretBox::new([91u8; 32]); let hasher = ApiKeyHasher::new([92u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -75,11 +74,10 @@ async fn models_route_should_accept_stored_client_api_key() { let hasher = ApiKeyHasher::new([93u8; 32]); let plaintext = insert_client_api_key(&pool, &hasher).await; let config = test_config(url); - let secret_box = SecretBox::new([94u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -130,6 +128,7 @@ fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: BTreeMap::new(), + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 300, @@ -156,7 +155,6 @@ fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/proxy/openai/responses_routes.rs b/tests/proxy/openai/responses_routes.rs index 0b2e55fb7..732e9295d 100644 --- a/tests/proxy/openai/responses_routes.rs +++ b/tests/proxy/openai/responses_routes.rs @@ -14,7 +14,7 @@ use codex_proxy_rs::{ QuotaConfig, QuotaWarningThresholds, SecurityConfig, ServerConfig, TlsConfig, UsageStatsConfig, WebSocketPoolConfig, }, - infra::{crypto::SecretBox, database::connect_sqlite, identity::ApiKeyHasher}, + infra::{database::connect_sqlite, identity::ApiKeyHasher}, proxy::dispatch::session_affinity::SqliteSessionAffinityStore, runtime::services::{BackgroundTaskStores, Services}, runtime::state::AppState, @@ -142,11 +142,10 @@ async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::Temp let hasher = ApiKeyHasher::new([95u8; 32]); let plaintext = insert_client_api_key(&pool, &hasher).await; let config = test_config(url); - let secret_box = SecretBox::new([96u8; 32]); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: SqliteSessionAffinityStore::new(pool.clone()), refresh_leases: RefreshLeaseStore::new(pool.clone()), @@ -193,6 +192,7 @@ fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: BTreeMap::new(), + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 300, @@ -219,7 +219,6 @@ fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/runtime/services.rs b/tests/runtime/services.rs index bc66e1d3a..71b34613e 100644 --- a/tests/runtime/services.rs +++ b/tests/runtime/services.rs @@ -5,7 +5,7 @@ use codex_proxy_rs::{ auth::service::SqliteAdminSessionStore, keys::service::SqliteClientKeyStore, monitoring::event_store::SqliteEventLogStore, }, - infra::{crypto::SecretBox, identity::ApiKeyHasher}, + infra::identity::ApiKeyHasher, runtime::services::{BackgroundTaskStores, Services}, upstream::{ accounts::{ @@ -28,13 +28,12 @@ async fn services_try_new_should_use_configured_tls_transport_builder() { "sqlite://{}", temp_dir.path().join("services-tls.sqlite").display() )); - let secret_box = SecretBox::load_or_create(temp_dir.path().join("master.key")).unwrap(); let hasher = ApiKeyHasher::load_or_create(temp_dir.path().join("api-key-pepper.key")) .expect("api key hasher"); let stores = BackgroundTaskStores { accounts: SqliteAccountStore::new(pool.clone()), admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone(), secret_box), + cookies: SqliteCookieStore::new(pool.clone()), fingerprints: FingerprintRepository::new(pool.clone()), session_affinity: codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new( diff --git a/tests/runtime/tasks/cleanup.rs b/tests/runtime/tasks/cleanup.rs index 1d96648cf..4f8978d6e 100644 --- a/tests/runtime/tasks/cleanup.rs +++ b/tests/runtime/tasks/cleanup.rs @@ -7,8 +7,7 @@ async fn cookie_cleanup_task_should_delete_only_expired_cookie_rows() { let pool = connect_sqlite(&format!("sqlite://{}", db.display())) .await .expect("sqlite pool"); - let secret_box = SecretBox::new([7u8; 32]); - let store = SqliteCookieStore::new(pool.clone(), secret_box); + let store = SqliteCookieStore::new(pool.clone()); insert_account(&pool, "acct-cookie").await; let now = Utc::now(); insert_cookie( diff --git a/tests/runtime/tasks/mod.rs b/tests/runtime/tasks/mod.rs index 82265df46..228af33c6 100644 --- a/tests/runtime/tasks/mod.rs +++ b/tests/runtime/tasks/mod.rs @@ -3,13 +3,12 @@ use std::{sync::Arc, time::Duration as StdDuration}; use async_trait::async_trait; use chrono::{Duration, Utc}; use codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore; -use codex_proxy_rs::infra::crypto::SecretBox; use codex_proxy_rs::infra::database::connect_sqlite; use codex_proxy_rs::proxy::dispatch::session_affinity::SessionAffinityEntry; use codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore; use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore; -use codex_proxy_rs::upstream::accounts::model::AccountUsageDelta; use codex_proxy_rs::upstream::accounts::model::{Account, AccountStatus}; +use codex_proxy_rs::upstream::accounts::model::{AccountModelUsageDelta, AccountUsageDelta}; use codex_proxy_rs::upstream::accounts::store::{AccountStore, AccountStoreResult}; use codex_proxy_rs::upstream::fingerprint::FingerprintRepository; use codex_proxy_rs::upstream::models::ModelConfig; @@ -67,6 +66,15 @@ impl AccountStore for FakeAccountStore { 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) } @@ -173,7 +181,7 @@ async fn insert_cookie( expires_at: chrono::DateTime, ) { sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, 'chatgpt.com', ?, 'cipher', '/', ?, ?)", + "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values (?, ?, 'chatgpt.com', ?, 'cipher', '/', ?, ?)", ) .bind(id) .bind(account_id) diff --git a/tests/support/config.rs b/tests/support/config.rs index 00011f800..343b92025 100644 --- a/tests/support/config.rs +++ b/tests/support/config.rs @@ -20,6 +20,7 @@ pub fn test_config(database_url: String) -> AppConfig { default_reasoning_effort: None, service_tier: None, aliases: BTreeMap::new(), + account_routes: BTreeMap::new(), }, auth: AuthConfig { refresh_margin_seconds: 300, @@ -46,7 +47,6 @@ pub fn test_config(database_url: String) -> AppConfig { }, database: DatabaseConfig { url: database_url }, security: SecurityConfig { - master_key_file: "data/master.key".to_string(), api_key_pepper_file: "data/api-key-pepper.key".to_string(), }, tls: TlsConfig { diff --git a/tests/upstream/accounts/account_pool/selection.rs b/tests/upstream/accounts/account_pool/selection.rs index 39d689437..f9b0e1a39 100644 --- a/tests/upstream/accounts/account_pool/selection.rs +++ b/tests/upstream/accounts/account_pool/selection.rs @@ -154,6 +154,28 @@ fn account_pool_should_filter_by_model_plan_allowlist() { assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "plus"); } +#[test] +fn account_pool_should_filter_by_model_account_routes() { + let mut model_routes = BTreeMap::new(); + model_routes.insert("gpt-5.5".to_string(), vec!["acct_b".to_string()]); + let mut pool = AccountPool::with_options(AccountPoolOptions { + max_concurrent_per_account: 1, + model_account_routes: model_routes, + ..AccountPoolOptions::default() + }); + pool.insert(crate::support::accounts::test_account( + "acct_a", + AccountStatus::Active, + )); + pool.insert(crate::support::accounts::test_account( + "acct_b", + AccountStatus::Active, + )); + + assert_eq!(pool.acquire("gpt-5.5").unwrap().id, "acct_b"); + assert!(pool.acquire("gpt-5.5").is_none()); +} + #[test] fn account_pool_should_exclude_requested_account_ids() { let mut pool = AccountPool::default(); diff --git a/tests/upstream/accounts/cookies.rs b/tests/upstream/accounts/cookies.rs index 4af3926a6..44ddff4d3 100644 --- a/tests/upstream/accounts/cookies.rs +++ b/tests/upstream/accounts/cookies.rs @@ -1,6 +1,6 @@ use chrono::{Duration, Utc}; use codex_proxy_rs::{ - infra::{crypto::SecretBox, database::connect_sqlite}, + infra::database::connect_sqlite, upstream::accounts::{ cookies::SqliteCookieStore, model::AccountStatus, @@ -12,16 +12,16 @@ use sqlx::SqlitePool; #[tokio::test] async fn cookie_store_should_cleanup_expired_cookies() { - let (pool, secret_box, _dir) = sqlite_cookie_store_parts("cookies-cleanup.sqlite").await; + let (pool, _dir) = sqlite_cookie_store_parts("cookies-cleanup.sqlite").await; seed_account(&pool, "acct_a").await; - let store = SqliteCookieStore::new(pool.clone(), secret_box); + let store = SqliteCookieStore::new(pool.clone()); let now = Utc::now(); let past = now - Duration::hours(2); let future = now + Duration::hours(2); sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", ) .bind("cookie_expired") .bind("acct_a") @@ -36,7 +36,7 @@ async fn cookie_store_should_cleanup_expired_cookies() { .unwrap(); sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", ) .bind("cookie_valid") .bind("acct_a") @@ -60,12 +60,12 @@ async fn cookie_store_should_cleanup_expired_cookies() { assert_eq!(remaining.0, 1); } -async fn sqlite_cookie_store_parts(name: &str) -> (SqlitePool, SecretBox, tempfile::TempDir) { +async fn sqlite_cookie_store_parts(name: &str) -> (SqlitePool, tempfile::TempDir) { let dir = tempfile::tempdir().unwrap(); let db = dir.path().join(name); let url = format!("sqlite://{}", db.display()); let pool = connect_sqlite(&url).await.unwrap(); - (pool, SecretBox::new([21u8; 32]), dir) + (pool, dir) } async fn seed_account(pool: &SqlitePool, id: &str) { diff --git a/tests/upstream/accounts/quota_refresh.rs b/tests/upstream/accounts/quota_refresh.rs index 21bf546ab..f0ea58ae7 100644 --- a/tests/upstream/accounts/quota_refresh.rs +++ b/tests/upstream/accounts/quota_refresh.rs @@ -1,7 +1,6 @@ use std::{sync::Arc, time::Duration as StdDuration}; use chrono::{Duration, Utc}; -use codex_proxy_rs::infra::crypto::SecretBox; use codex_proxy_rs::infra::database::connect_sqlite; use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore; use codex_proxy_rs::upstream::accounts::model::AccountStatus; @@ -204,9 +203,8 @@ async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_con let pool = connect_sqlite(&format!("sqlite://{}", db.display())) .await .expect("sqlite pool"); - let secret_box = SecretBox::new([14u8; 32]); let store = SqliteAccountStore::new(pool.clone()); - let cookies = SqliteCookieStore::new(pool.clone(), secret_box); + 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") diff --git a/web/src/views/settings/index.vue b/web/src/views/settings/index.vue index 582a293b2..e23e6f625 100644 --- a/web/src/views/settings/index.vue +++ b/web/src/views/settings/index.vue @@ -1,51 +1,151 @@