mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 16:06:08 +00:00
refactor: converge database design
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Generated
+1
-104
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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> {
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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,4 +1,4 @@
|
||||
//! SQLite 数据库连接与 Schema 迁移。
|
||||
//! SQLite 数据库连接与 Schema 初始化。
|
||||
|
||||
use std::{fs, path::Path, str::FromStr, time::Duration};
|
||||
|
||||
|
||||
+1
-3
@@ -1,9 +1,7 @@
|
||||
#![deny(missing_docs)]
|
||||
|
||||
//! 平台基础设施层,承载加密、身份、存储、日志和 JSON 原语。
|
||||
//! 平台基础设施层,承载身份、存储、日志和 JSON 原语。
|
||||
|
||||
/// 加密原语。
|
||||
pub mod crypto;
|
||||
/// 数据库连接与 SQLite 管理。
|
||||
pub mod database;
|
||||
/// 身份验证原语(管理员密码、API Key 哈希)。
|
||||
|
||||
+97
-10
@@ -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);
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,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
@@ -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(),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 [
|
||||
(
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()),
|
||||
|
||||
@@ -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()),
|
||||
|
||||
@@ -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(),
|
||||
),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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()),
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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");
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,4 +1,3 @@
|
||||
pub mod crypto;
|
||||
pub mod log_rotation;
|
||||
pub mod storage_schema;
|
||||
pub mod time;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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,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")
|
||||
|
||||
@@ -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>
|
||||
|
||||
Reference in New Issue
Block a user