refactor: converge database design

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-27 23:58:07 +08:00
co-authored by Codex
parent 159eb49e3d
commit a7a72079f5
61 changed files with 2791 additions and 1013 deletions
Generated
+1 -104
View File
@@ -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"
+2 -3
View File
@@ -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"
+1 -1
View File
@@ -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
+440
View File
@@ -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 仍读取聚合趋势。
## 最终建议
当前数据库设计已经完成第一轮“无历史包袱”收敛。当前边界已经更清晰:启动配置归配置文件,后台运行状态归数据库,长期统计归统计表,事件日志归诊断窗口。
+43 -104
View File
@@ -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<AdminUsageRecord> {
}
}
async fn list_all_event_logs(state: &AppState) -> Vec<EventLog> {
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<DateTime<Utc>>,
}
fn models_by_account(logs: &[EventLog]) -> HashMap<String, Vec<AdminAccountModelUsageData>> {
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::<String, Vec<ModelUsageAggregate>>::new();
for ((account_id, _), aggregate) in by_account_model {
by_account.entry(account_id).or_default().push(aggregate);
fn models_by_account(
records: Vec<AccountModelUsageRecord>,
) -> HashMap<String, Vec<AdminAccountModelUsageData>> {
let mut by_account = HashMap::<String, Vec<AccountModelUsageRecord>>::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<f64> {
value
.as_f64()
+123 -143
View File
@@ -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<impl IntoResponse, AdminError> {
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<Utc>,
@@ -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<Utc>,
) -> Vec<AdminUsageTimeBucketRecord> {
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<Utc>,
) -> DashboardTrendData {
@@ -444,16 +458,16 @@ fn dashboard_trend_data_at(
})
.collect::<Vec<_>>();
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<Utc>,
) -> DashboardHealthTimelineData {
let current_slot = china_quarter_hour_start(now);
@@ -508,16 +524,16 @@ fn dashboard_health_timeline_data_at(
})
.collect::<Vec<_>>();
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<Utc>, end: DateTime<Utc>) -> UsageWindow {
fn usage_window(
records: &[AdminUsageTimeBucketRecord],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> 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<Utc>,
end: DateTime<Utc>,
context: BillingContext<'_>,
) -> Option<f64> {
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<f64> {
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<f64> {
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<f64> {
))
}
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<u64> {
[
"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<Utc> {
+126
View File
@@ -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<i64> {
[
"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<u64> {
let result = sqlx::query(
"delete from event_logs where id in (
+63 -1
View File
@@ -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<Utc>,
end: DateTime<Utc>,
) -> Result<Vec<AdminUsageTimeBucketRecord>, 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<Utc>,
pub model: String,
pub service_tier: Option<String>,
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<UsageSummary> for AdminUsageSummary {
fn from(s: UsageSummary) -> Self {
Self {
@@ -220,6 +261,27 @@ impl From<UsageSummary> for AdminUsageSummary {
}
}
impl From<UsageTimeBucketRecord> 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<UsageListRecord> for AdminUsageRecord {
fn from(usage: UsageListRecord) -> Self {
Self {
+89
View File
@@ -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<Utc>,
/// 模型。
pub model: String,
/// 服务层级。
pub service_tier: Option<String>,
/// 请求数。
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<Utc>,
end: DateTime<Utc>,
) -> SqliteUsageStoreResult<Vec<UsageTimeBucketRecord>> {
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<UsageListRecord> {
@@ -292,6 +360,27 @@ fn usage_list_from_row(row: &SqliteRow) -> SqliteUsageStoreResult<UsageListRecor
})
}
fn usage_time_bucket_from_row(row: &SqliteRow) -> SqliteUsageStoreResult<UsageTimeBucketRecord> {
let service_tier = row.get::<String, _>("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<String>) -> SqliteUsageStoreResult<Option<DateTime<Utc>>> {
value
.as_deref()
+17 -2
View File
@@ -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<String, String>,
pub model_account_routes: BTreeMap<String, Vec<String>>,
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",
-21
View File
@@ -8,27 +8,6 @@ pub mod loader;
pub mod settings;
/// 配置类型定义。
pub mod types;
/// 配置写回。
pub mod writeback;
use thiserror::Error;
/// 配置加载结果类型。
pub type ConfigResult<T> = Result<T, ::config::ConfigError>;
/// 配置写入结果类型。
pub type ConfigWriteResult<T> = Result<T, ConfigWriteError>;
/// 配置写入错误。
#[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),
}
+414 -37
View File
@@ -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<String>,
/// 模型别名映射。
pub model_aliases: BTreeMap<String, String>,
/// 模型到账号 ID 的显式路由。
pub model_account_routes: BTreeMap<String, Vec<String>>,
/// 是否启用访问令牌刷新。
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<String>,
/// 模型别名映射。
pub model_aliases: Option<BTreeMap<String, String>>,
/// 模型到账号 ID 的显式路由。
pub model_account_routes: Option<BTreeMap<String, Vec<String>>>,
/// 访问令牌过期前多少秒开始刷新。
pub refresh_margin_seconds: Option<u64>,
/// 访问令牌刷新并发数。
pub refresh_concurrency: Option<u32>,
/// 单账号最大并发请求数。
pub max_concurrent_per_account: Option<usize>,
/// 同账号请求间隔毫秒数。
@@ -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<String, String>,
) -> Result<BTreeMap<String, String>, 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<String, Vec<String>>,
) -> Result<BTreeMap<String, Vec<String>>, 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<String, SettingsServiceError> {
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<String, SettingsServic
}
}
fn positive_u64(field: &str, value: u64) -> Result<u64, SettingsServiceError> {
if value == 0 {
Err(invalid_field(field, "must be greater than 0"))
} else {
Ok(value)
}
}
fn positive_u32(field: &str, value: u32) -> Result<u32, SettingsServiceError> {
if value == 0 {
Err(invalid_field(field, "must be greater than 0"))
} else {
Ok(value)
}
}
fn positive_usize(field: &str, value: usize) -> Result<usize, SettingsServiceError> {
if value == 0 {
Err(invalid_field(field, "must be greater than 0"))
@@ -210,20 +298,30 @@ fn invalid_field(field: &str, message: impl Into<String>) -> SettingsServiceErro
#[derive(Clone)]
pub struct RuntimeSettingsService {
current: Arc<StdRwLock<Arc<AppConfig>>>,
config_path: Arc<PathBuf>,
pool: SqlitePool,
account_pool: Option<Arc<RuntimeAccountPoolService>>,
}
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<PathBuf>) -> 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<RuntimeAccountPoolService>,
) -> 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<PathBuf> {
self.config_path.clone()
/// 初始化缺省运行设置,并返回数据库中的运行配置快照。
pub async fn load_or_initialize_config(
mut config: AppConfig,
pool: &SqlitePool,
) -> Result<AppConfig, RuntimeSettingsError> {
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<AdminSettings, RuntimeSettingsError> {
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<AdminSettings, RuntimeSettingsError> {
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<String, Vec<String>>,
) -> 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<BTreeMap<String, Vec<String>>, 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::<String, Vec<String>>::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<u64, RuntimeSettingsError> {
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<u32, RuntimeSettingsError> {
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<usize, RuntimeSettingsError> {
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<u64, RuntimeSettingsError> {
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<String>,
message: impl Into<String>,
) -> RuntimeSettingsError {
RuntimeSettingsError::StoredField {
field: field.into(),
message: message.into(),
}
}
+2 -2
View File
@@ -62,6 +62,8 @@ pub struct ModelConfig {
pub service_tier: Option<String>,
/// 模型别名映射。
pub aliases: BTreeMap<String, String>,
/// 模型到账号 ID 的显式路由。
pub account_routes: BTreeMap<String, Vec<String>>,
}
/// 认证、轮换与 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,
}
-27
View File
@@ -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<Path>,
) -> 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)
}
}
-160
View File
@@ -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<Sha256>;
/// 计算 HMAC-SHA256 并以 URL-safe base64 输出。
pub fn hmac_sha256_base64(key: &[u8], payload: &[u8]) -> Option<String> {
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<T> = Result<T, CryptoError>;
/// 用于加密和解密敏感字符串的密钥盒。
#[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<Path>) -> CryptoResult<Self> {
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<String> {
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<SecretString> {
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)
}
+1 -1
View File
@@ -1,4 +1,4 @@
//! SQLite 数据库连接与 Schema 迁移。
//! SQLite 数据库连接与 Schema 初始化。
use std::{fs, path::Path, str::FromStr, time::Duration};
+1 -3
View File
@@ -1,9 +1,7 @@
#![deny(missing_docs)]
//! 平台基础设施层,承载加密、身份、存储、日志和 JSON 原语。
//! 平台基础设施层,承载身份、存储、日志和 JSON 原语。
/// 加密原语。
pub mod crypto;
/// 数据库连接与 SQLite 管理。
pub mod database;
/// 身份验证原语(管理员密码、API Key 哈希)。
+97 -10
View File
@@ -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);
+1 -1
View File
@@ -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 {
+9 -3
View File
@@ -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)
}
+3 -5
View File
@@ -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<dyn Error + Send + Sync>> {
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(),
+13 -9
View File
@@ -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(),
+6 -18
View File
@@ -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<T> = Result<T, SqliteCookieStoreError>;
#[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<Option<String>> {
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::<String, _>("name");
let value_cipher = row.get::<String, _>("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::<String, _>("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)
+15
View File
@@ -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,
}
+110 -8
View File
@@ -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<String>,
/// 模型到允许订阅计划的映射。
pub model_plan_allowlist: BTreeMap<String, Vec<String>>,
/// 模型到指定账号 ID 的显式路由。
pub model_account_routes: BTreeMap<String, Vec<String>>,
}
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<String, Vec<String>>) {
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<String>) -> 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<String> {
self.options.tier_priority.iter().find_map(|tier| {
candidates
@@ -750,7 +773,7 @@ fn access_token_expired(account: &Account, now: DateTime<Utc>) -> bool {
pub struct RuntimeAccountPoolService {
pool: StdArc<tokio::sync::Mutex<AccountPool>>,
store: StdArc<dyn AccountStore>,
request_interval: StdDuration,
request_interval: StdArc<std::sync::RwLock<StdDuration>>,
}
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<usize, RuntimeAccountPoolError> {
let accounts = self
@@ -843,6 +878,7 @@ impl RuntimeAccountPoolService {
/// 按请求上下文获取可用账号。
pub async fn acquire_with(&self, request: AccountAcquireRequest) -> Option<AcquiredAccount> {
let model = request.model.clone();
let refresh = self.pool.lock().await.acquire_with_status_refresh(request);
self.persist_expired_statuses(refresh.expired_account_ids)
.await;
@@ -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,
+124 -1
View File
@@ -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<DateTime<Utc>>,
}
/// 账号模型维度用量记录。
#[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<DateTime<Utc>>,
}
/// 账号配额快照。
#[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<Option<String>>;
@@ -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<Vec<AccountModelUsageRecord>> {
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<Option<String>> {
SqliteAccountStore::get_quota_json(self, account_id)
.await
@@ -1883,6 +1991,21 @@ fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> SqliteAccountStoreResult<Acc
})
}
fn model_usage_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountModelUsageRecord> {
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::<Option<String>, _>("last_used_at"))?,
})
}
fn quota_snapshot_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountQuotaSnapshot> {
+9 -13
View File
@@ -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,
+119 -17
View File
@@ -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<String> =
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,
+29 -42
View File
@@ -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 [
(
+9 -40
View File
@@ -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<String>,
) -> (
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 {
+8 -12
View File
@@ -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()),
+2 -2
View File
@@ -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()),
+2 -11
View File
@@ -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(),
),
+3 -6
View File
@@ -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 {
+3 -4
View File
@@ -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 {
+4 -6
View File
@@ -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 {
+31 -6
View File
@@ -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::<u64>(),
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()),
+3 -4
View File
@@ -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 {
+60
View File
@@ -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");
}
+99 -37
View File
@@ -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<String, String> = 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 {
+2 -3
View File
@@ -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
+3 -7
View File
@@ -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 {
-31
View File
@@ -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");
}
-1
View File
@@ -1,4 +1,3 @@
pub mod crypto;
pub mod log_rotation;
pub mod storage_schema;
pub mod time;
+316 -3
View File
@@ -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::<Vec<_>>(),
[
"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::<Vec<_>>(),
[
"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::<Vec<_>>(),
[
"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::<Vec<_>>(),
[
"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::<Vec<_>>(),
[
"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::<Vec<_>>(),
[
"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();
@@ -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
+3 -18
View File
@@ -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 {
@@ -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",
@@ -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
@@ -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(),
);
@@ -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]
+4 -6
View File
@@ -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 {
+3 -4
View File
@@ -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 {
+2 -3
View File
@@ -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(
+1 -2
View File
@@ -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(
+11 -3
View File
@@ -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<Option<String>> {
Ok(None)
}
@@ -173,7 +181,7 @@ async fn insert_cookie(
expires_at: chrono::DateTime<Utc>,
) {
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)
+1 -1
View File
@@ -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 {
@@ -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();
+7 -7
View File
@@ -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) {
+1 -3
View File
@@ -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")
+298 -47
View File
@@ -1,51 +1,151 @@
<script setup lang="ts">
import { computed, onMounted, ref } from 'vue'
import { Cpu, Gauge, RefreshCw, Save, Timer } from '@lucide/vue'
import { computed, onMounted, reactive, ref } from 'vue'
import { Cpu, Gauge, GitBranch, Plus, RefreshCw, Save, Timer, Trash2, Zap } from '@lucide/vue'
import { getSettings, updateSettings } from '@/api'
import BaseButton from '@/components/base/BaseButton.vue'
import BaseCard from '@/components/base/BaseCard.vue'
import BaseInput from '@/components/base/BaseInput.vue'
import BaseSelect from '@/components/base/BaseSelect.vue'
import { toast } from '@/components/base/BaseToast'
import { withMinimumDuration } from '@/utils/async'
import { getSettings, updateSettings } from '@/api'
import { toast } from '@/components/base/BaseToast'
type RotationStrategy = 'least_used' | 'round_robin' | 'sticky'
interface AliasRow {
alias: string
target: string
}
interface SettingsForm {
defaultModel: string
modelAliases: Record<string, string>
refreshMarginSeconds: number
refreshConcurrency: number
maxConcurrentPerAccount: number
requestIntervalMs: number
rotationStrategy: RotationStrategy
}
const loading = ref(true)
const refreshing = ref(false)
const saving = ref(false)
const aliasError = ref('')
const form = ref({
const form = reactive<SettingsForm>({
defaultModel: '',
modelAliases: {},
refreshMarginSeconds: 300,
refreshConcurrency: 2,
maxConcurrentPerAccount: 3,
requestIntervalMs: 0,
requestIntervalMs: 50,
rotationStrategy: 'least_used',
})
const rotationOptions = [
{ label: '最少使用优先', value: 'least_used' },
{ label: '轮询', value: 'round_robin' },
{ label: '会话粘性', value: 'sticky' },
const aliasRows = ref<AliasRow[]>([])
const rotationOptions: Array<{
label: string
value: RotationStrategy
description: string
}> = [
{
label: '智能分配',
value: 'least_used',
description: '优先使用压力更低的账号,适合长期稳定运行。',
},
{
label: '轮询',
value: 'round_robin',
description: '按顺序分配请求,让账号池保持均匀节奏。',
},
{
label: '粘性',
value: 'sticky',
description: '持续使用同一账号,直到遇到限速或不可用。',
},
]
function numericModel(key: 'maxConcurrentPerAccount' | 'requestIntervalMs') {
function numericModel(key: keyof Pick<
SettingsForm,
| 'refreshMarginSeconds'
| 'refreshConcurrency'
| 'maxConcurrentPerAccount'
| 'requestIntervalMs'
>) {
return computed({
get: () => String(form.value[key] ?? 0),
get: () => String(form[key] ?? 0),
set: (value: string) => {
const parsed = Number(value)
form.value[key] = Number.isFinite(parsed) ? parsed : 0
form[key] = Number.isFinite(parsed) ? parsed : 0
},
})
}
const refreshMarginSecondsValue = numericModel('refreshMarginSeconds')
const refreshConcurrencyValue = numericModel('refreshConcurrency')
const maxConcurrentPerAccountValue = numericModel('maxConcurrentPerAccount')
const requestIntervalMsValue = numericModel('requestIntervalMs')
function aliasesToRows(aliases: Record<string, string> = {}) {
return Object.entries(aliases)
.sort(([left], [right]) => left.localeCompare(right))
.map(([alias, target]) => ({ alias, target }))
}
function rowsToAliases(rows: AliasRow[]) {
const aliases: Record<string, string> = {}
for (const row of rows) {
const alias = row.alias.trim()
const target = row.target.trim()
if (!alias && !target) continue
if (!alias || !target) {
return { aliases: {}, error: '别名和目标模型需要同时填写' }
}
if (alias === target) {
return { aliases: {}, error: '别名不能指向自身' }
}
if (aliases[alias] !== undefined) {
return { aliases: {}, error: `别名重复:${alias}` }
}
aliases[alias] = target
}
return { aliases, error: '' }
}
function applySettings(data: any) {
form.defaultModel = data.defaultModel || ''
form.modelAliases = data.modelAliases || {}
form.refreshMarginSeconds = Number(data.refreshMarginSeconds ?? 300)
form.refreshConcurrency = Number(data.refreshConcurrency ?? 2)
form.maxConcurrentPerAccount = Number(data.maxConcurrentPerAccount ?? 3)
form.requestIntervalMs = Number(data.requestIntervalMs ?? 50)
form.rotationStrategy = (data.rotationStrategy || 'least_used') as RotationStrategy
aliasRows.value = aliasesToRows(form.modelAliases)
aliasError.value = ''
}
function addAliasRow() {
aliasRows.value = [...aliasRows.value, { alias: '', target: '' }]
aliasError.value = ''
}
function updateAliasRow(index: number, key: keyof AliasRow, value: string) {
aliasRows.value = aliasRows.value.map((row, rowIndex) =>
rowIndex === index ? { ...row, [key]: value } : row,
)
aliasError.value = ''
}
function removeAliasRow(index: number) {
aliasRows.value = aliasRows.value.filter((_row, rowIndex) => rowIndex !== index)
aliasError.value = ''
}
async function loadSettings() {
try {
loading.value = true
const data = await getSettings()
form.value = data
applySettings(await getSettings())
} catch (error: any) {
toast.error(error.message || '加载失败')
} finally {
@@ -65,15 +165,25 @@ async function refreshSettings() {
async function handleSave() {
if (saving.value || loading.value || refreshing.value) return
const aliasResult = rowsToAliases(aliasRows.value)
if (aliasResult.error) {
aliasError.value = aliasResult.error
return
}
try {
saving.value = true
const patch = {
defaultModel: form.value.defaultModel,
maxConcurrentPerAccount: form.value.maxConcurrentPerAccount,
requestIntervalMs: form.value.requestIntervalMs,
rotationStrategy: form.value.rotationStrategy,
}
await updateSettings(patch)
const data = await updateSettings({
defaultModel: form.defaultModel,
modelAliases: aliasResult.aliases,
refreshMarginSeconds: form.refreshMarginSeconds,
refreshConcurrency: form.refreshConcurrency,
maxConcurrentPerAccount: form.maxConcurrentPerAccount,
requestIntervalMs: form.requestIntervalMs,
rotationStrategy: form.rotationStrategy,
})
applySettings(data)
toast.success('设置已保存')
} catch (error: any) {
toast.error(error.message || '保存失败')
@@ -82,29 +192,27 @@ async function handleSave() {
}
}
onMounted(() => {
loadSettings()
})
onMounted(loadSettings)
</script>
<template>
<div class="w-full">
<header class="flex min-h-17 items-start justify-between gap-4">
<div>
<h1 class="mt-0 text-[34px] leading-[1.15] font-extrabold mb-0 text-(--cp-text-primary)">
<h1 class="m-0 text-[34px] leading-[1.15] font-extrabold text-(--cp-text-primary)">
系统设置
</h1>
<p class="mt-2.5 text-[15px] leading-[1.15] font-semibold mb-0 text-(--cp-text-secondary)">
配置默认模型和账号调度参数。
<p class="mt-2.5 mb-0 text-[15px] leading-[1.15] font-semibold text-(--cp-text-secondary)">
让模型入口、账号选择和 Token 刷新保持可控。
</p>
</div>
<div class="mt-0.5 flex shrink-0 items-center gap-2">
<BaseButton variant="ghost" :disabled="loading || refreshing" @click="refreshSettings">
<BaseButton variant="ghost" :disabled="loading || refreshing || saving" @click="refreshSettings">
<template #icon>
<RefreshCw class="size-4" :class="refreshing ? 'animate-spin' : undefined" />
</template>
重置
刷新
</BaseButton>
<BaseButton
variant="primary"
@@ -115,21 +223,28 @@ onMounted(() => {
<template #icon>
<Save class="size-4" />
</template>
{{ saving ? '保存中...' : '保存设置' }}
{{ saving ? '保存中...' : '保存' }}
</BaseButton>
</div>
</header>
<BaseCard v-loading="loading" :padded="false" class="mt-5 max-w-5xl" body-class="px-5 py-5">
<template #body>
<div class="grid gap-5 md:grid-cols-2">
<div class="mt-5 grid max-w-6xl gap-5">
<BaseCard
v-loading="loading"
:padded="false"
title="运行参数"
description="默认模型、请求节奏和 Token 刷新。"
header-class="px-5 pt-4"
body-class="px-5 py-5"
>
<div class="grid gap-4 lg:grid-cols-2">
<div class="grid gap-2">
<span class="text-xs leading-[1.15] font-bold text-(--cp-text-secondary)"
>默认模型</span
>
<span class="text-xs leading-[1.15] font-bold text-(--cp-text-secondary)">
默认模型
</span>
<BaseInput v-model="form.defaultModel">
<template #prefix>
<Cpu class="size-4 text-(--cp-text-muted)" />
<Cpu class="size-4" />
</template>
</BaseInput>
</div>
@@ -140,30 +255,166 @@ onMounted(() => {
</span>
<BaseInput v-model="maxConcurrentPerAccountValue" type="number">
<template #prefix>
<Gauge class="size-4 text-(--cp-text-muted)" />
<Gauge class="size-4" />
</template>
</BaseInput>
</div>
<div class="grid gap-2">
<span class="text-xs leading-[1.15] font-bold text-(--cp-text-secondary)"
>轮换策略</span
>
<BaseSelect v-model="form.rotationStrategy" :options="rotationOptions" />
<span class="text-xs leading-[1.15] font-bold text-(--cp-text-secondary)">
提前刷新秒数
</span>
<BaseInput v-model="refreshMarginSecondsValue" type="number">
<template #prefix>
<Timer class="size-4" />
</template>
</BaseInput>
<span class="text-xs leading-[1.15] font-[650] text-(--cp-text-muted)">
Token 过期前多少秒触发刷新。
</span>
</div>
<div class="grid gap-2">
<span class="text-xs leading-[1.15] font-bold text-(--cp-text-secondary)">
刷新并发数
</span>
<BaseInput v-model="refreshConcurrencyValue" type="number">
<template #prefix>
<Zap class="size-4" />
</template>
</BaseInput>
<span class="text-xs leading-[1.15] font-[650] text-(--cp-text-muted)">
同时刷新 Token 的最大请求数,减小可避免限流。
</span>
</div>
<div class="grid gap-2 lg:col-span-2">
<span class="text-xs leading-[1.15] font-bold text-(--cp-text-secondary)">
请求间隔 ms
</span>
<BaseInput v-model="requestIntervalMsValue" type="number">
<template #prefix>
<Timer class="size-4 text-(--cp-text-muted)" />
<Timer class="size-4" />
</template>
</BaseInput>
</div>
</div>
</template>
</BaseCard>
</BaseCard>
<BaseCard
v-loading="loading"
:padded="false"
title="模型映射"
description="把客户端可见名称指向真实上游模型。"
header-class="px-5 pt-4"
body-class="px-5 py-5"
>
<div class="grid gap-3">
<div
class="hidden grid-cols-[minmax(0,1fr)_minmax(0,1fr)_2.5rem] gap-2 px-0.75 text-xs leading-none font-bold text-(--cp-text-secondary) sm:grid"
>
<span>别名</span>
<span>目标模型</span>
<span />
</div>
<div
v-if="aliasRows.length === 0"
class="rounded-(--cp-input-radius-base) bg-(--cp-bg-subtle) px-4 py-3 text-[13px] font-[650] text-(--cp-text-muted)"
>
还没有模型映射。
</div>
<div
v-for="(row, index) in aliasRows"
:key="index"
class="grid items-start gap-2 sm:grid-cols-[minmax(0,1fr)_minmax(0,1fr)_2.5rem]"
>
<BaseInput
:model-value="row.alias"
placeholder="gpt-5.2"
@update:model-value="updateAliasRow(index, 'alias', $event)"
>
<template #prefix>
<GitBranch class="size-4" />
</template>
</BaseInput>
<BaseInput
:model-value="row.target"
placeholder="gpt-5.5"
@update:model-value="updateAliasRow(index, 'target', $event)"
/>
<BaseButton
variant="ghost"
size="default"
icon-only
label="删除映射"
:disabled="saving || loading"
@click="removeAliasRow(index)"
>
<Trash2 class="size-4" />
</BaseButton>
</div>
<div class="flex flex-wrap items-center gap-3 pt-1">
<BaseButton variant="default" :disabled="saving || loading" @click="addAliasRow">
<template #icon>
<Plus class="size-4" />
</template>
添加映射
</BaseButton>
<span v-if="aliasError" class="text-xs font-[650] text-(--cp-danger-text)">
{{ aliasError }}
</span>
</div>
</div>
</BaseCard>
<BaseCard
v-loading="loading"
:padded="false"
title="账号选择"
description="决定每次请求如何使用账号池。"
header-class="px-5 pt-4"
body-class="px-5 py-5"
>
<div class="grid gap-3 lg:grid-cols-3">
<button
v-for="option in rotationOptions"
:key="option.value"
type="button"
class="min-h-25 cursor-pointer rounded-(--cp-input-radius-base) border-0 px-4 py-3.5 text-left shadow-(--cp-shadow-input) outline-none transition-[background-color,box-shadow,color] duration-160 focus-visible:ring-2 focus-visible:ring-(--cp-info-border)"
:class="
form.rotationStrategy === option.value
? 'bg-(--cp-info-bg) text-(--cp-info-text) shadow-(--cp-shadow-control)'
: 'bg-(--cp-input-current-bg,var(--cp-input-context-bg)) text-(--cp-text-primary) hover:bg-(--cp-input-current-bg-hover,var(--cp-input-context-bg-hover)) hover:shadow-(--cp-shadow-input-hover)'
"
:aria-pressed="form.rotationStrategy === option.value"
@click="form.rotationStrategy = option.value"
>
<span class="flex items-center gap-2">
<span
class="inline-flex size-4 shrink-0 items-center justify-center rounded-full bg-(--cp-bg-surface) shadow-[inset_0_0_0_1px_var(--cp-default-border-hover)]"
>
<span
class="size-2 rounded-full transition-opacity duration-150"
:class="
form.rotationStrategy === option.value
? 'bg-(--cp-info) opacity-100'
: 'opacity-0'
"
/>
</span>
<span class="text-[14px] leading-[1.15] font-[760]">{{ option.label }}</span>
</span>
<span
class="mt-2 block text-[13px] leading-[1.5] font-[650] text-(--cp-text-secondary)"
>
{{ option.description }}
</span>
</button>
</div>
</BaseCard>
</div>
</div>
</template>