mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
docs: audit deleted feature parity
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Generated
-2
@@ -343,8 +343,6 @@ dependencies = [
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-stream",
|
||||
"tokio-tungstenite",
|
||||
"tower",
|
||||
"tower-http",
|
||||
|
||||
+1
-3
@@ -28,13 +28,10 @@ rustls-pki-types = "1.14.1"
|
||||
secrecy = "0.10.3"
|
||||
serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = "1.0.150"
|
||||
serde_yml = "0.0.13"
|
||||
sha2 = "0.11.0"
|
||||
sqlx = { version = "0.9.0", features = ["runtime-tokio", "tls-rustls", "sqlite", "chrono", "uuid"] }
|
||||
thiserror = "2.0.18"
|
||||
tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "signal", "fs", "time", "sync", "net", "test-util"] }
|
||||
tokio-rustls = { version = "0.26.4", default-features = false, features = ["ring", "tls12"] }
|
||||
tokio-stream = "0.1.17"
|
||||
tokio-tungstenite = { version = "0.28.0", features = ["proxy", "rustls-tls-native-roots"] }
|
||||
tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout"] }
|
||||
tracing = "0.1.44"
|
||||
@@ -44,6 +41,7 @@ tungstenite = { version = "0.27.0", features = ["deflate", "proxy"] }
|
||||
uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
|
||||
|
||||
[dev-dependencies]
|
||||
serde_yml = "0.0.13"
|
||||
tempfile = "3.27.0"
|
||||
tower = { version = "0.5.3", features = ["util"] }
|
||||
wiremock = "0.6.5"
|
||||
|
||||
@@ -39,7 +39,7 @@ Codex Proxy RS 关注这些细节:
|
||||
|
||||
**Codex 账号池**
|
||||
|
||||
- 支持 ChatGPT/Codex OAuth、device code、refresh token 账号接入
|
||||
- 支持 ChatGPT/Codex OAuth、CPR JSON、Sub2API JSON 和 refresh token 账号接入
|
||||
- 支持账号轮转、并发控制、请求间隔和会话亲和性
|
||||
- 支持账号状态、配额、Cookie、模型目录和连接测试管理
|
||||
|
||||
@@ -52,7 +52,7 @@ Codex Proxy RS 关注这些细节:
|
||||
**可观测与管理**
|
||||
|
||||
- 管理控制台内置账号、密钥、事件日志、用量统计和系统设置
|
||||
- SQLite 持久化,敏感密钥加密存储
|
||||
- SQLite 持久化,账号、Cookie 与访问密钥采用统一的本地存储策略
|
||||
- 结构化事件日志,便于定位请求失败、上游状态和调度路径
|
||||
|
||||
## 快速开始
|
||||
@@ -69,7 +69,7 @@ admin:
|
||||
default_password: admin
|
||||
```
|
||||
|
||||
长期运行前请修改默认密码,并确认数据库、密钥和日志路径符合你的部署环境。
|
||||
长期运行前请修改默认密码,并确认数据库和日志路径符合你的部署环境。
|
||||
|
||||
运行时数据默认写入 `.runtime`:
|
||||
|
||||
@@ -142,7 +142,7 @@ flowchart LR
|
||||
proxy --> codex[codex<br/>Codex 协议、模型和上游传输]
|
||||
http --> admin[admin<br/>管理端 API]
|
||||
admin --> runtime[runtime<br/>启动装配、状态和后台任务]
|
||||
runtime --> infra[infra<br/>SQLite、加密、路径和日志]
|
||||
runtime --> infra[infra<br/>SQLite、路径和日志]
|
||||
proxy --> telemetry[telemetry<br/>事件日志和用量记录]
|
||||
admin --> telemetry
|
||||
http --> web[web<br/>管理控制台静态资源]
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
# 删除代码功能等价审计
|
||||
|
||||
审计日期:2026-06-28
|
||||
|
||||
对照基准:
|
||||
|
||||
- 当前 Rust 工程:`/home/zyy/Codes/codex-proxy-rs`
|
||||
- TS 参考工程:`/home/zyy/桌面/Codes/codex-proxy`
|
||||
|
||||
本文档只审计近期删除或收敛的功能是否造成产品能力缺失。结论不以文件名为准,而以当前路由、服务实现和 TS 版已存在能力为准。
|
||||
|
||||
## 总体结论
|
||||
|
||||
当前删除并不全是问题。部分删除符合当前产品方向:后台设置减少人工配置、日志策略固定、接口去掉路径 id、账号导入增强为 CPR/Sub2API 双格式并支持重复账号更新,这些属于合理收敛。
|
||||
|
||||
但有几类能力被删后没有等价替代,会影响管理闭环:
|
||||
|
||||
1. 账号导出缺失,导入和备份迁移不闭环。
|
||||
2. 批量账号状态更新缺失,账号池批量维护效率下降。
|
||||
3. 批量健康检查缺失,目前只剩单账号模型测试。
|
||||
4. 手动 Cookie 维护缺失,Cloudflare 人工旁路能力不完整。
|
||||
5. 独立用量统计接口缺失,当前只被 Dashboard 和账号列表部分吸收。
|
||||
6. 手动模型刷新缺失,后台只剩启动后和周期任务刷新。
|
||||
7. 客户端 API Key 导出缺失,密钥备份/迁移能力被削弱。
|
||||
|
||||
## 当前 Rust 管理端路由事实
|
||||
|
||||
当前 `src/admin/router.rs` 保留的管理端接口:
|
||||
|
||||
- 登录会话:`/api/admin/login`、`/api/admin/auth/status`、`/api/admin/logout`
|
||||
- 设置:`/api/admin/settings`
|
||||
- 概览:`/api/admin/dashboard/summary`、`/api/admin/dashboard/trend`
|
||||
- 账号:`/api/admin/accounts`、`/api/admin/accounts/import`、`/api/admin/accounts/oauth/authorize`、`/api/admin/accounts/oauth/exchange`、`/api/admin/accounts/test`、`/api/admin/accounts/models`、`/api/admin/accounts/delete`、`/api/admin/accounts/update`、`/api/admin/accounts/refresh`、`/api/admin/accounts/quota`
|
||||
- 日志:`/api/admin/logs`、`/api/admin/logs/delete`、`/api/admin/logs/detail`
|
||||
- 客户端 API Key:`/api/admin/keys`、`/api/admin/keys/delete`、`/api/admin/keys/update`
|
||||
|
||||
已从路由层移除的旧 Rust 管理端接口:
|
||||
|
||||
- `/api/admin/models/refresh`
|
||||
- `/api/admin/usage`
|
||||
- `/api/admin/usage/summary`
|
||||
- `/api/admin/accounts/export`
|
||||
- `/api/admin/accounts/quota-warnings`
|
||||
- `/api/admin/accounts/health-check`
|
||||
- `/api/admin/accounts/reset-usage`
|
||||
- `/api/admin/accounts/cookies`
|
||||
- `/api/admin/logs/state`
|
||||
- `/api/admin/keys/export`
|
||||
|
||||
## 逐项审计
|
||||
|
||||
| 能力 | Rust 当前状态 | TS 参考状态 | 结论 | 建议 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 账号导入 | 已保留并增强。`src/admin/accounts/service/importing.rs` 支持 CPR/Sub2API,重复账号会按 ChatGPT identity 更新,缺少字段时会补充 account id、user id、email、quota。 | `src/routes/accounts.ts` 支持 `/auth/accounts/import`。 | 不缺失。当前 Rust 实现比旧导入更贴近当前产品。 | 保留当前实现。 |
|
||||
| 账号导出 | 当前路由已删除 `/api/admin/accounts/export`,旧 `export_accounts` 和 `export_with_tokens` 不再暴露。 | `src/routes/accounts.ts` 有 `/auth/accounts/export`,支持按 ids 和 format 导出。 | 功能缺失。导入存在但导出消失,备份、迁移、批量排障不闭环。 | P1 恢复,格式只保留 CPR,避免旧 native 命名和兼容分支。 |
|
||||
| 批量删除账号 | 当前保留 `/api/admin/accounts/delete`。 | TS 有 `/auth/accounts/batch-delete`。 | 不缺失。 | 保留。 |
|
||||
| 批量状态更新 | 当前 `/api/admin/accounts/update` 只解析单账号 `id` 更新,旧 `batch_update_status` 已删除。 | TS 有 `/auth/accounts/batch-status`。 | 功能缺失。批量禁用、恢复账号池时需要逐个操作。 | P1 恢复为 `/api/admin/accounts/update` 的批量 data 形式,保持无路径 id。 |
|
||||
| 单账号刷新/额度查询 | 当前保留 `/api/admin/accounts/refresh` 和 `/api/admin/accounts/quota`,id 放在 body 或 query。 | TS 有 `/auth/accounts/:id/refresh`、`/auth/accounts/:id/quota`。 | 不缺失,且当前 Rust 路由形态更符合项目规范。 | 保留当前形态。 |
|
||||
| 重置账号用量 | 当前删除 `/api/admin/accounts/reset-usage`,未发现等价管理端入口。 | TS 有 `/auth/accounts/:id/reset-usage`。 | 功能缺失但不是核心链路。主要影响调试、演示和清理异常统计。 | P3 低优先级恢复,或明确产品不提供手动重置。 |
|
||||
| 批量健康检查 | 当前删除 `/api/admin/accounts/health-check`,只保留 `/api/admin/accounts/test?id=...` 单账号 SSE 模型测试。 | TS 有 `/auth/accounts/health-check`,支持 ids、stagger、concurrency。 | 功能缺失。单账号测试不能替代账号池健康巡检。 | P1 恢复批量健康检查,复用现有单账号探测和 loading 流程。 |
|
||||
| 测试模型列表 | 当前保留 `/api/admin/accounts/models?id=...`,从上游返回模型。 | TS 模型列表能力来自模型 catalog 和上游刷新。 | 不缺失。 | 保留当前接口名 `models`。 |
|
||||
| OAuth 添加账号 | 当前新增 `/api/admin/accounts/oauth/authorize`、`/api/admin/accounts/oauth/exchange`。 | TS 有 `/auth/accounts/login` 和 OAuth 链路。 | 不缺失。Rust 当前已按后台弹窗流程实现。 | 保留。 |
|
||||
| 手动 Cookie 维护 | 当前删除 `/api/admin/accounts/cookies`。Cookie 存储仍存在,Cloudflare 逻辑会删除账号 cookie,但没有管理端查看/写入入口。 | TS 有 `/auth/accounts/:id/cookies` GET/POST/DELETE,并在 quota 失败时提示写入 cf cookie。 | 功能缺失。遇到 Cloudflare 挑战时缺少人工修复入口。 | P2 恢复,接口保持无路径 id,例如 `GET/POST /api/admin/accounts/cookies?id=...`,删除可并入 POST action 或独立 delete。 |
|
||||
| 配额告警列表 | 当前删除 `/api/admin/accounts/quota-warnings`。账号列表 summary 仍有 high_usage/attention,但没有独立 warning 列表。 | TS 有 `/auth/quota/warnings`。 | 部分缺失。列表页摘要不能替代可点击、可定位、可清理的告警集合。 | P2 视 UI 是否需要恢复;如果不做告警中心,应在产品上明确删除。 |
|
||||
| 独立账号用量分页 | 当前删除 `/api/admin/usage`,账号列表会内嵌账号用量,Dashboard 会读取 usage summary 和 buckets。 | TS 有 `/admin/usage-stats/summary` 与 `/admin/usage-stats/history`。 | 部分缺失。概览和账号列表替代了展示,但没有独立用量页/API。 | P2 如果后续有用量分析页,应恢复为 `/api/admin/usage` 和 `/api/admin/usage/summary`;否则可接受删除。 |
|
||||
| 用量历史 | 当前 Dashboard 趋势读 `usage_time_buckets`,但没有通用 history API。 | TS 有 `/admin/usage-stats/history`,支持 granularity 和 hours。 | 部分缺失。当前只服务 Dashboard,不适合复用到独立分析页面。 | P2 如要开放分析视图,新增 `/api/admin/usage/history`,直接读聚合事实表。 |
|
||||
| 手动模型刷新 | 当前删除 `/api/admin/models/refresh`,但 `ModelRefreshTask` 会启动后和每小时刷新,并有 `refresh_once()`。 | TS 有 `/admin/refresh-models` 手动触发。 | 功能缺失但已有后台替代。问题在于管理端无法立即刷新模型目录。 | P2 恢复轻量接口,直接调用现有 `refresh_once()` 或后台任务句柄,避免再引入 `admin_models` 旧模块。 |
|
||||
| OpenAI 模型列表 | 当前 `/v1/models`、`/v1/models/catalog`、`/v1/models/{model}`、`/v1/models/{model}/info` 仍存在。 | TS 有同类模型路由。 | 不缺失。 | 保留。 |
|
||||
| 日志列表/详情/清空 | 当前保留 `/api/admin/logs`、`/api/admin/logs/detail`、`/api/admin/logs/delete`。 | TS 有 `/admin/logs`、`/admin/logs/:id`、`/admin/logs/clear`。 | 不缺失。Rust 当前筛选字段更丰富。 | 保留当前形态。 |
|
||||
| 日志开关状态 | 当前删除 `/api/admin/logs/state`。 | TS 有 `/admin/logs/state` GET/POST。 | 不算缺失。此前产品已决定日志采集和容量走默认策略,不在系统设置暴露。 | 不恢复。 |
|
||||
| 设置项大而全 | 当前 `/api/admin/settings` 只保留 modelAliases、modelAccountRoutes、refreshMarginSeconds、refreshConcurrency、maxConcurrentPerAccount、requestIntervalMs、rotationStrategy。 | TS `general-settings`/`quota-settings` 暴露端口、TLS、默认模型、日志容量、更新策略等大量字段。 | 不算缺失。当前产品已收敛为只保留需要人工决策的运行配置。 | 不恢复 TS 大而全设置页。 |
|
||||
| 客户端 API Key 列表/创建/更新/删除 | 当前保留 `/api/admin/keys`、`/api/admin/keys/update`、`/api/admin/keys/delete`,并保存完整 `sk_` key 供复制。 | TS 的 `/auth/api-keys` 是第三方 provider key pool,领域不同。 | 当前核心能力不缺失。 | 保留当前“客户端访问密钥”定位,不混入 provider key pool。 |
|
||||
| 客户端 API Key 导出 | 当前删除 `/api/admin/keys/export`。旧 Rust 实现曾导出 `rustLocalClientApiKeys` 并标记 `rotation_required`。 | TS provider key pool 有 `/auth/api-keys/export`。 | 功能缺失但需谨慎。当前项目已允许保存完整 key,导出可以服务备份;但也会扩大密钥泄露面。 | P3 如恢复,只做管理员显式导出,并使用 CPR 格式;不恢复旧 `rustLocalClientApiKeys` 命名。 |
|
||||
|
||||
## 不建议恢复的删除项
|
||||
|
||||
以下删除符合当前项目方向,不建议回滚:
|
||||
|
||||
1. `/api/admin/logs/state`:日志策略已固定,系统设置不再暴露日志采集和容量。
|
||||
2. TS 版 `general-settings` 中的默认模型、默认推理强度、默认 service tier:这些应由客户端请求决定,网关不兜底。
|
||||
3. TS 版模型 aliases 的旧配置形态:当前 Rust 已拆成 `modelAliases` 和 `modelAccountRoutes`,边界更清楚。
|
||||
4. 路径中携带 id 的管理端接口形态:当前 Rust 统一使用 query/body 传 id,更符合近期接口规范。
|
||||
5. 第三方 provider API Key pool:当前 Rust 的 API Key 是客户端访问密钥,不应混入 TS 版 provider key pool 概念。
|
||||
|
||||
## 建议恢复优先级
|
||||
|
||||
P1:应该优先补齐,直接影响管理闭环。
|
||||
|
||||
- 账号导出:`GET /api/admin/accounts/export`
|
||||
- 批量账号状态更新:复用 `/api/admin/accounts/update`
|
||||
- 批量健康检查:`POST /api/admin/accounts/health-check`
|
||||
|
||||
P2:影响运维体验和可诊断性,建议排第二批。
|
||||
|
||||
- 手动 Cookie 维护:`GET/POST /api/admin/accounts/cookies`
|
||||
- 配额告警列表:`GET /api/admin/accounts/quota-warnings`
|
||||
- 独立用量统计和历史:`GET /api/admin/usage`、`GET /api/admin/usage/summary`、可选 `GET /api/admin/usage/history`
|
||||
- 手动模型刷新:`POST /api/admin/models/refresh`
|
||||
|
||||
P3:低频能力,可先产品确认。
|
||||
|
||||
- 重置账号用量:`POST /api/admin/accounts/reset-usage`
|
||||
- 客户端 API Key 导出:`GET /api/admin/keys/export`
|
||||
|
||||
## 恢复原则
|
||||
|
||||
如果恢复上述能力,不应该把旧代码原样搬回来。需要按当前项目规范重写:
|
||||
|
||||
1. 不在路径中放 id,id 统一放 query 或 body。
|
||||
2. 不恢复兼容命名,例如 `native`、`rustLocalClientApiKeys`。
|
||||
3. 不恢复大而全的 settings/config 写回逻辑。
|
||||
4. 不为了兼容旧数据增加迁移分支;已有开发库需要重建。
|
||||
5. 账号导出格式只使用 CPR,和当前导入格式闭环。
|
||||
6. 批量操作必须返回 `updated/deleted/notFound` 这类可被 UI 准确呈现的数据。
|
||||
7. 涉及真实网络请求的操作必须有后端测试覆盖,包括成功、部分失败、空 ids、未登录。
|
||||
|
||||
## 需要补的测试
|
||||
|
||||
当前被删除的测试里包含 `tests/admin/models/*`、`tests/admin/keys/import_export.rs`、旧账号 import/export 测试等。若恢复能力,至少补回以下测试面:
|
||||
|
||||
- 账号导出:全部导出、按 ids 导出、无效格式、未登录。
|
||||
- 批量状态:部分 id 不存在、空 ids、非法 status、状态写入后账号列表可见。
|
||||
- 批量健康检查:成功、失败、跳过、并发参数校验。
|
||||
- Cookie 维护:写入、读取、清空、账号不存在。
|
||||
- 用量接口:分页、汇总、空数据库。
|
||||
- 模型刷新:无账号、部分计划失败、成功刷新并写入模型快照。
|
||||
- API Key 导出:未登录、空列表、导出字段包含完整 key。
|
||||
+114
-654
File diff suppressed because it is too large
Load Diff
@@ -1,4 +1,4 @@
|
||||
use super::{types::*, AdminAccountService};
|
||||
use super::AdminAccountService;
|
||||
|
||||
impl AdminAccountService {
|
||||
pub(crate) async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
|
||||
@@ -8,70 +8,4 @@ impl AdminAccountService {
|
||||
.ok()
|
||||
.flatten()
|
||||
}
|
||||
|
||||
pub async fn cookies(&self, account_id: &str) -> Result<Option<String>, AdminAccountError> {
|
||||
self.ensure_cookie_account_exists(account_id).await?;
|
||||
self.cookies
|
||||
.cookie_header(account_id, "chatgpt.com")
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::LoadCookies)
|
||||
}
|
||||
pub async fn set_cookies(
|
||||
&self,
|
||||
account_id: &str,
|
||||
cookies: serde_json::Value,
|
||||
) -> Result<Option<String>, AdminAccountError> {
|
||||
let cookie_header = match cookies {
|
||||
serde_json::Value::String(ref s) => s.trim().to_string(),
|
||||
serde_json::Value::Object(ref obj) => {
|
||||
let pairs: Vec<String> = obj
|
||||
.iter()
|
||||
.filter_map(|(name, val)| {
|
||||
let v = val.as_str()?.trim();
|
||||
if name.trim().is_empty() || v.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(format!("{}={}", name.trim(), v))
|
||||
})
|
||||
.collect();
|
||||
if pairs.is_empty() {
|
||||
return Err(AdminAccountError::NoValidCookies);
|
||||
}
|
||||
pairs.join("; ")
|
||||
}
|
||||
_ => return Err(AdminAccountError::NoValidCookies),
|
||||
};
|
||||
self.ensure_cookie_account_exists(account_id).await?;
|
||||
match self
|
||||
.cookies
|
||||
.set_cookie_header(account_id, &cookie_header)
|
||||
.await
|
||||
{
|
||||
Ok(0) => Err(AdminAccountError::NoValidCookies),
|
||||
Ok(_) => self
|
||||
.cookies
|
||||
.cookie_header(account_id, "chatgpt.com")
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::LoadCookies),
|
||||
Err(_) => Err(AdminAccountError::StoreCookies),
|
||||
}
|
||||
}
|
||||
pub async fn delete_cookies(&self, account_id: &str) -> Result<(), AdminAccountError> {
|
||||
self.ensure_cookie_account_exists(account_id).await?;
|
||||
self.cookies
|
||||
.delete_account_cookies(account_id)
|
||||
.await
|
||||
.map(|_| ())
|
||||
.map_err(|_| AdminAccountError::DeleteCookies)
|
||||
}
|
||||
async fn ensure_cookie_account_exists(
|
||||
&self,
|
||||
account_id: &str,
|
||||
) -> Result<(), AdminAccountError> {
|
||||
match self.cookies.account_exists(account_id).await {
|
||||
Ok(true) => Ok(()),
|
||||
Ok(false) => Err(AdminAccountError::NotFound),
|
||||
Err(_) => Err(AdminAccountError::Inspect),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,74 +1,21 @@
|
||||
use chrono::Utc;
|
||||
use secrecy::SecretString;
|
||||
|
||||
use super::{types::*, AdminAccountService};
|
||||
use super::{
|
||||
types::{
|
||||
import_quota_plan_type, import_status_from_usage_error, import_usage_plan_type,
|
||||
import_usage_string, AdminAccountError, ImportSupplementalAccountInfo,
|
||||
ImportSupplementalNeeds, ImportedAccountState, ImportedAccounts, ResolvedImportTokens,
|
||||
},
|
||||
AdminAccountService,
|
||||
};
|
||||
|
||||
impl AdminAccountService {
|
||||
pub async fn export(
|
||||
&self,
|
||||
ids: Vec<String>,
|
||||
) -> Result<Vec<AdminAccountMetadata>, AdminAccountError> {
|
||||
if ids.is_empty() {
|
||||
let mut all = Vec::new();
|
||||
let mut cursor = None;
|
||||
loop {
|
||||
let page = self
|
||||
.store
|
||||
.list_metadata(cursor, 200)
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::Export)?;
|
||||
all.extend(page.items.into_iter().map(AdminAccountMetadata::from));
|
||||
if page.next_cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
cursor = page.next_cursor;
|
||||
}
|
||||
Ok(all)
|
||||
} else {
|
||||
let mut accounts = Vec::with_capacity(ids.len());
|
||||
for id in ids {
|
||||
if let Ok(Some(stored)) = self.store.get(&id).await {
|
||||
accounts.push(stored_to_admin_metadata(stored));
|
||||
}
|
||||
}
|
||||
Ok(accounts)
|
||||
}
|
||||
}
|
||||
pub async fn export_with_tokens(
|
||||
&self,
|
||||
ids: Vec<String>,
|
||||
) -> Result<Vec<crate::upstream::accounts::store::StoredAccount>, AdminAccountError> {
|
||||
if ids.is_empty() {
|
||||
let mut all = Vec::new();
|
||||
let mut cursor = None;
|
||||
loop {
|
||||
let page = self
|
||||
.store
|
||||
.list(cursor, 200)
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::Export)?;
|
||||
all.extend(page.items);
|
||||
if page.next_cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
cursor = page.next_cursor;
|
||||
}
|
||||
Ok(all)
|
||||
} else {
|
||||
let mut accounts = Vec::with_capacity(ids.len());
|
||||
for id in ids {
|
||||
if let Ok(Some(stored)) = self.store.get(&id).await {
|
||||
accounts.push(stored);
|
||||
}
|
||||
}
|
||||
Ok(accounts)
|
||||
}
|
||||
}
|
||||
pub async fn import(
|
||||
&self,
|
||||
data: serde_json::Value,
|
||||
) -> Result<ImportedAccounts, AdminAccountError> {
|
||||
let parsed = crate::upstream::accounts::import_export::parse_account_import_payload(&data)
|
||||
let parsed = crate::upstream::accounts::importing::parse_account_import_payload(&data)
|
||||
.map_err(|_| AdminAccountError::NoImportableAccounts)?;
|
||||
let source_format = parsed.source.as_str();
|
||||
let entries = parsed.entries;
|
||||
@@ -96,8 +43,8 @@ impl AdminAccountService {
|
||||
}
|
||||
async fn import_entry(
|
||||
&self,
|
||||
entry: crate::upstream::accounts::import_export::AccountImportEntry,
|
||||
source: crate::upstream::accounts::import_export::AccountImportSource,
|
||||
entry: crate::upstream::accounts::importing::AccountImportEntry,
|
||||
source: crate::upstream::accounts::importing::AccountImportSource,
|
||||
) -> Result<ImportedAccountState, AdminAccountError> {
|
||||
let Some(resolved_tokens) = self
|
||||
.resolve_import_tokens(entry.token, entry.refresh_token)
|
||||
@@ -105,7 +52,7 @@ impl AdminAccountService {
|
||||
else {
|
||||
return Ok(ImportedAccountState::Skipped);
|
||||
};
|
||||
let label = crate::upstream::accounts::import_export::normalize_label(entry.label);
|
||||
let label = crate::upstream::accounts::importing::normalize_label(entry.label);
|
||||
if label
|
||||
.as_ref()
|
||||
.is_some_and(|label| label.chars().count() > 64)
|
||||
@@ -116,13 +63,13 @@ impl AdminAccountService {
|
||||
let access_token_expires_at = entry
|
||||
.access_token_expires_at
|
||||
.as_deref()
|
||||
.map(crate::upstream::accounts::import_export::parse_account_import_datetime)
|
||||
.map(crate::upstream::accounts::importing::parse_account_import_datetime)
|
||||
.transpose()
|
||||
.map_err(|_| AdminAccountError::InvalidAccessTokenExpiresAt)?;
|
||||
let quota_fetched_at = entry
|
||||
.quota_fetched_at
|
||||
.as_deref()
|
||||
.map(crate::upstream::accounts::import_export::parse_account_import_datetime)
|
||||
.map(crate::upstream::accounts::importing::parse_account_import_datetime)
|
||||
.transpose()
|
||||
.map_err(|_| AdminAccountError::InvalidAccessTokenExpiresAt)?;
|
||||
let mut quota_json = entry
|
||||
@@ -131,16 +78,15 @@ impl AdminAccountService {
|
||||
.map(serde_json::Value::to_string);
|
||||
let mut quota_fetched_at = quota_fetched_at;
|
||||
let quota_verify_required = entry.quota_verify_required.unwrap_or(false);
|
||||
let parsed_status = crate::upstream::accounts::import_export::parse_account_import_status(
|
||||
let parsed_status = crate::upstream::accounts::importing::parse_account_import_status(
|
||||
entry.status.as_deref(),
|
||||
)
|
||||
.map_err(|error| AdminAccountError::InvalidStatus(error.to_string()))?;
|
||||
let mut status =
|
||||
crate::upstream::accounts::import_export::normalized_imported_account_status(
|
||||
parsed_status,
|
||||
source,
|
||||
&resolved_tokens.access_token,
|
||||
);
|
||||
let mut status = crate::upstream::accounts::importing::normalized_imported_account_status(
|
||||
parsed_status,
|
||||
source,
|
||||
&resolved_tokens.access_token,
|
||||
);
|
||||
let access_token_expires_at = resolved_tokens
|
||||
.claims
|
||||
.as_ref()
|
||||
@@ -151,26 +97,24 @@ impl AdminAccountService {
|
||||
let claims = resolved_tokens.claims.as_ref();
|
||||
let mut plan_type = claims
|
||||
.and_then(|claims| claims.plan_type.clone())
|
||||
.or_else(|| {
|
||||
crate::upstream::accounts::import_export::normalize_nonempty(entry.plan_type)
|
||||
});
|
||||
.or_else(|| crate::upstream::accounts::importing::normalize_nonempty(entry.plan_type));
|
||||
if plan_type.is_none() {
|
||||
plan_type = entry.cached_quota.as_ref().and_then(import_quota_plan_type);
|
||||
}
|
||||
let email = claims.and_then(|claims| claims.email.clone()).or_else(|| {
|
||||
crate::upstream::accounts::import_export::normalize_nonempty(entry.email.clone())
|
||||
crate::upstream::accounts::importing::normalize_nonempty(entry.email.clone())
|
||||
});
|
||||
let chatgpt_account_id = claims
|
||||
.and_then(|claims| claims.account_id.as_deref())
|
||||
.or_else(|| {
|
||||
crate::upstream::accounts::import_export::normalize_nonempty_str(
|
||||
crate::upstream::accounts::importing::normalize_nonempty_str(
|
||||
entry.account_id.as_deref(),
|
||||
)
|
||||
});
|
||||
let chatgpt_user_id = claims
|
||||
.and_then(|claims| claims.user_id.as_deref())
|
||||
.or_else(|| {
|
||||
crate::upstream::accounts::import_export::normalize_nonempty_str(
|
||||
crate::upstream::accounts::importing::normalize_nonempty_str(
|
||||
entry.user_id.as_deref(),
|
||||
)
|
||||
});
|
||||
@@ -333,9 +277,11 @@ impl AdminAccountService {
|
||||
refresh_token: Option<String>,
|
||||
) -> Result<Option<ResolvedImportTokens>, AdminAccountError> {
|
||||
let mut refresh_token =
|
||||
crate::upstream::accounts::import_export::normalize_nonempty(refresh_token);
|
||||
let Some(access_token) = crate::upstream::accounts::import_export::normalize_nonempty(
|
||||
token.map(crate::upstream::accounts::import_export::normalize_bearer_token),
|
||||
crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
|
||||
let Some(access_token) = crate::upstream::accounts::importing::normalize_nonempty(
|
||||
token
|
||||
.as_deref()
|
||||
.map(crate::upstream::accounts::importing::normalize_bearer_token),
|
||||
) else {
|
||||
let Some(existing_refresh_token) = refresh_token else {
|
||||
return Ok(None);
|
||||
@@ -345,9 +291,9 @@ impl AdminAccountService {
|
||||
.refresh(&existing_refresh_token)
|
||||
.await
|
||||
.map_err(AdminAccountError::RefreshTokenExchange)?;
|
||||
let access_token = crate::upstream::accounts::import_export::normalize_nonempty(Some(
|
||||
crate::upstream::accounts::import_export::normalize_bearer_token(
|
||||
refreshed.access_token,
|
||||
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
|
||||
crate::upstream::accounts::importing::normalize_bearer_token(
|
||||
&refreshed.access_token,
|
||||
),
|
||||
))
|
||||
.ok_or(AdminAccountError::TokenRequired)?;
|
||||
@@ -387,10 +333,8 @@ impl AdminAccountService {
|
||||
.refresh(&existing_refresh_token)
|
||||
.await
|
||||
.map_err(AdminAccountError::RefreshTokenExchange)?;
|
||||
let access_token = crate::upstream::accounts::import_export::normalize_nonempty(Some(
|
||||
crate::upstream::accounts::import_export::normalize_bearer_token(
|
||||
refreshed.access_token,
|
||||
),
|
||||
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
|
||||
crate::upstream::accounts::importing::normalize_bearer_token(&refreshed.access_token),
|
||||
))
|
||||
.ok_or(AdminAccountError::TokenRequired)?;
|
||||
refresh_token = refreshed.refresh_token.or(Some(existing_refresh_token));
|
||||
@@ -411,7 +355,7 @@ impl AdminAccountService {
|
||||
account_id: Option<&str>,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<String, AdminAccountError> {
|
||||
let provided_id = crate::upstream::accounts::import_export::normalize_nonempty_str(id)
|
||||
let provided_id = crate::upstream::accounts::importing::normalize_nonempty_str(id)
|
||||
.map(ToString::to_string);
|
||||
if let Some(id) = provided_id.as_deref() {
|
||||
match self.store.get(id).await {
|
||||
@@ -422,9 +366,8 @@ impl AdminAccountService {
|
||||
}
|
||||
|
||||
let chatgpt_account_id =
|
||||
crate::upstream::accounts::import_export::normalize_nonempty_str(account_id);
|
||||
let chatgpt_user_id =
|
||||
crate::upstream::accounts::import_export::normalize_nonempty_str(user_id);
|
||||
crate::upstream::accounts::importing::normalize_nonempty_str(account_id);
|
||||
let chatgpt_user_id = crate::upstream::accounts::importing::normalize_nonempty_str(user_id);
|
||||
if let Some(chatgpt_account_id) = chatgpt_account_id {
|
||||
if let Some(existing) = self
|
||||
.store
|
||||
@@ -436,8 +379,7 @@ impl AdminAccountService {
|
||||
}
|
||||
}
|
||||
|
||||
Ok(provided_id.unwrap_or_else(|| {
|
||||
crate::upstream::accounts::import_export::normalized_account_id(None)
|
||||
}))
|
||||
Ok(provided_id
|
||||
.unwrap_or_else(|| crate::upstream::accounts::importing::normalized_account_id(None)))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,13 @@ use secrecy::{ExposeSecret, SecretString};
|
||||
|
||||
use crate::infra::json::{NumberedPage, Page};
|
||||
|
||||
use super::{types::*, AdminAccountService};
|
||||
use super::{
|
||||
types::{
|
||||
parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountMetadata,
|
||||
AdminAccountMetadataUpdate, BatchDeleteAccounts, ManualCreateTokens, UpdatedAccountStatus,
|
||||
},
|
||||
AdminAccountService,
|
||||
};
|
||||
|
||||
impl AdminAccountService {
|
||||
pub async fn list(
|
||||
@@ -65,10 +71,12 @@ impl AdminAccountService {
|
||||
refresh_token: Option<String>,
|
||||
) -> Result<AdminAccountMetadata, AdminAccountError> {
|
||||
let provided_refresh_token =
|
||||
crate::upstream::accounts::import_export::normalize_nonempty(refresh_token);
|
||||
crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
|
||||
let tokens = if let Some(access_token) =
|
||||
crate::upstream::accounts::import_export::normalize_nonempty(
|
||||
token.map(crate::upstream::accounts::import_export::normalize_bearer_token),
|
||||
crate::upstream::accounts::importing::normalize_nonempty(
|
||||
token
|
||||
.as_deref()
|
||||
.map(crate::upstream::accounts::importing::normalize_bearer_token),
|
||||
) {
|
||||
ManualCreateTokens {
|
||||
access_token,
|
||||
@@ -81,9 +89,9 @@ impl AdminAccountService {
|
||||
.refresh(&refresh_token)
|
||||
.await
|
||||
.map_err(AdminAccountError::RefreshTokenExchange)?;
|
||||
let access_token = crate::upstream::accounts::import_export::normalize_nonempty(Some(
|
||||
crate::upstream::accounts::import_export::normalize_bearer_token(
|
||||
token_pair.access_token,
|
||||
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
|
||||
crate::upstream::accounts::importing::normalize_bearer_token(
|
||||
&token_pair.access_token,
|
||||
),
|
||||
))
|
||||
.ok_or(AdminAccountError::TokenRequired)?;
|
||||
@@ -141,7 +149,7 @@ impl AdminAccountService {
|
||||
}
|
||||
existing.id
|
||||
} else {
|
||||
let id = crate::upstream::accounts::import_export::normalized_account_id(None);
|
||||
let id = crate::upstream::accounts::importing::normalized_account_id(None);
|
||||
self.store
|
||||
.insert(crate::upstream::accounts::store::NewAccount {
|
||||
id: id.clone(),
|
||||
@@ -300,13 +308,12 @@ impl AdminAccountService {
|
||||
.await
|
||||
{
|
||||
Ok(tokens) => {
|
||||
let access_token =
|
||||
crate::upstream::accounts::import_export::normalize_nonempty(Some(
|
||||
crate::upstream::accounts::import_export::normalize_bearer_token(
|
||||
tokens.access_token,
|
||||
),
|
||||
))
|
||||
.ok_or(AdminAccountError::TokenRequired)?;
|
||||
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
|
||||
crate::upstream::accounts::importing::normalize_bearer_token(
|
||||
&tokens.access_token,
|
||||
),
|
||||
))
|
||||
.ok_or(AdminAccountError::TokenRequired)?;
|
||||
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
|
||||
&access_token,
|
||||
Utc::now(),
|
||||
@@ -351,7 +358,7 @@ impl AdminAccountService {
|
||||
if !updated {
|
||||
return Err(AdminAccountError::NotFound);
|
||||
}
|
||||
if crate::upstream::accounts::import_export::refresh_failure_status_clears_next_refresh_at(
|
||||
if crate::upstream::accounts::importing::refresh_failure_status_clears_next_refresh_at(
|
||||
status,
|
||||
) {
|
||||
let cleared = self
|
||||
@@ -377,47 +384,6 @@ impl AdminAccountService {
|
||||
.map(stored_to_admin_metadata)
|
||||
.ok_or(AdminAccountError::NotFound)
|
||||
}
|
||||
pub async fn reset_usage(
|
||||
&self,
|
||||
account_id: &str,
|
||||
) -> Result<AdminAccountMetadata, AdminAccountError> {
|
||||
self.store
|
||||
.reset_usage(account_id)
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::ResetUsage)?;
|
||||
self.sync_account_pool(account_id).await?;
|
||||
self.store
|
||||
.get(account_id)
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::NotFound)?
|
||||
.map(stored_to_admin_metadata)
|
||||
.ok_or(AdminAccountError::NotFound)
|
||||
}
|
||||
pub async fn batch_update_status(
|
||||
&self,
|
||||
ids: Vec<String>,
|
||||
status: &str,
|
||||
) -> Result<BatchUpdateAccountStatus, AdminAccountError> {
|
||||
if ids.is_empty() {
|
||||
return Err(AdminAccountError::EmptyIds);
|
||||
}
|
||||
let status = parse_batch_account_status(status)?;
|
||||
let mut updated = 0u32;
|
||||
let mut not_found = Vec::new();
|
||||
for id in ids {
|
||||
match self.store.set_status(&id, status).await {
|
||||
Ok(true) => {
|
||||
updated += 1;
|
||||
self.sync_account_pool_best_effort(&id, "account batch status update")
|
||||
.await;
|
||||
self.evict_account_websocket_pool(&id).await;
|
||||
}
|
||||
Ok(false) => not_found.push(id),
|
||||
Err(_) => return Err(AdminAccountError::UpdateStatus),
|
||||
}
|
||||
}
|
||||
Ok(BatchUpdateAccountStatus { updated, not_found })
|
||||
}
|
||||
}
|
||||
|
||||
fn manual_refresh_failure_status(
|
||||
|
||||
@@ -13,7 +13,6 @@ use std::sync::Arc as StdArc;
|
||||
use chrono::{DateTime, Duration};
|
||||
|
||||
use crate::{
|
||||
config::types::QuotaWarningThresholds,
|
||||
upstream::accounts::{
|
||||
cookies::SqliteCookieStore, pool::RuntimeAccountPoolService, store::SqliteAccountStore,
|
||||
token_refresh::TokenRefresher,
|
||||
@@ -23,15 +22,13 @@ use crate::{
|
||||
|
||||
pub use types::{
|
||||
AdminAccountError, AdminAccountMetadata, AdminAccountMetadataUpdate, BatchDeleteAccounts,
|
||||
BatchUpdateAccountStatus, ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput,
|
||||
UpdatedAccountStatus,
|
||||
ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, UpdatedAccountStatus,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AdminAccountService {
|
||||
pub store: SqliteAccountStore,
|
||||
pub(crate) cookies: SqliteCookieStore,
|
||||
pub(crate) quota_thresholds: QuotaWarningThresholds,
|
||||
pub(crate) codex: StdArc<CodexBackendClient>,
|
||||
pub(crate) account_pool: StdArc<RuntimeAccountPoolService>,
|
||||
pub(crate) token_refresher: StdArc<dyn TokenRefresher>,
|
||||
@@ -40,32 +37,28 @@ pub struct AdminAccountService {
|
||||
pub(crate) installation_id: Option<String>,
|
||||
}
|
||||
|
||||
pub struct AdminAccountServiceParts {
|
||||
pub store: SqliteAccountStore,
|
||||
pub cookies: SqliteCookieStore,
|
||||
pub codex: StdArc<CodexBackendClient>,
|
||||
pub account_pool: StdArc<RuntimeAccountPoolService>,
|
||||
pub token_refresher: StdArc<dyn TokenRefresher>,
|
||||
pub oauth: oauth::AccountOAuthService,
|
||||
pub refresh_margin_seconds: u64,
|
||||
pub installation_id: Option<String>,
|
||||
}
|
||||
|
||||
impl AdminAccountService {
|
||||
#[expect(
|
||||
clippy::too_many_arguments,
|
||||
reason = "service constructor wires independent stores and runtime collaborators"
|
||||
)]
|
||||
pub fn new(
|
||||
store: SqliteAccountStore,
|
||||
cookies: SqliteCookieStore,
|
||||
quota_thresholds: QuotaWarningThresholds,
|
||||
codex: StdArc<CodexBackendClient>,
|
||||
account_pool: StdArc<RuntimeAccountPoolService>,
|
||||
token_refresher: StdArc<dyn TokenRefresher>,
|
||||
oauth: oauth::AccountOAuthService,
|
||||
refresh_margin_seconds: u64,
|
||||
installation_id: Option<String>,
|
||||
) -> Self {
|
||||
pub fn new(parts: AdminAccountServiceParts) -> Self {
|
||||
Self {
|
||||
store,
|
||||
cookies,
|
||||
quota_thresholds,
|
||||
codex,
|
||||
account_pool,
|
||||
token_refresher,
|
||||
oauth,
|
||||
refresh_margin_seconds,
|
||||
installation_id,
|
||||
store: parts.store,
|
||||
cookies: parts.cookies,
|
||||
codex: parts.codex,
|
||||
account_pool: parts.account_pool,
|
||||
token_refresher: parts.token_refresher,
|
||||
oauth: parts.oauth,
|
||||
refresh_margin_seconds: parts.refresh_margin_seconds,
|
||||
installation_id: parts.installation_id,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
use std::{collections::HashMap, fmt::Write as _, sync::Arc};
|
||||
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
@@ -147,7 +147,7 @@ impl AdminAccountService {
|
||||
) -> Result<super::types::ImportedAccounts, AdminAccountError> {
|
||||
let tokens = self.oauth.exchange(input).await?;
|
||||
self.import(serde_json::json!({
|
||||
"sourceFormat": "native",
|
||||
"sourceFormat": "cpr",
|
||||
"accounts": [{
|
||||
"token": tokens.access_token,
|
||||
"refreshToken": tokens.refresh_token,
|
||||
@@ -218,7 +218,11 @@ fn random_url_token(size: usize) -> String {
|
||||
fn random_hex_token(size: usize) -> String {
|
||||
let mut bytes = vec![0u8; size];
|
||||
rand::rng().fill_bytes(&mut bytes);
|
||||
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
|
||||
let mut token = String::with_capacity(size * 2);
|
||||
for byte in bytes {
|
||||
let _ = write!(token, "{byte:02x}");
|
||||
}
|
||||
token
|
||||
}
|
||||
|
||||
fn exchange_error_text(status: StatusCode, body: &str) -> String {
|
||||
|
||||
@@ -1,8 +1,5 @@
|
||||
use chrono::Utc;
|
||||
use secrecy::ExposeSecret;
|
||||
|
||||
use crate::infra::time::china_rfc3339;
|
||||
|
||||
use super::{types::AdminAccountError, AdminAccountService};
|
||||
|
||||
impl AdminAccountService {
|
||||
@@ -52,191 +49,4 @@ impl AdminAccountService {
|
||||
}
|
||||
Ok(serde_json::json!({ "quota": normalized, "raw": raw }))
|
||||
}
|
||||
pub async fn quota_warnings(&self) -> Result<serde_json::Value, AdminAccountError> {
|
||||
let snapshots = self
|
||||
.store
|
||||
.list_quota_snapshots()
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::QuotaWarnings)?;
|
||||
|
||||
let mut warnings = Vec::new();
|
||||
for snap in &snapshots {
|
||||
let quota: serde_json::Value =
|
||||
serde_json::from_str(&snap.quota_json).unwrap_or(serde_json::Value::Null);
|
||||
let used = crate::upstream::accounts::quota::quota_snapshot_limit_reached("a);
|
||||
if used {
|
||||
warnings.push(serde_json::json!({
|
||||
"accountId": snap.account_id,
|
||||
"email": snap.email,
|
||||
"level": "exhausted"
|
||||
}));
|
||||
} else {
|
||||
for (quota_key, used_percent, role) in quota_warning_windows("a) {
|
||||
let thresholds = match role {
|
||||
"secondary" => &self.quota_thresholds.secondary,
|
||||
_ => &self.quota_thresholds.primary,
|
||||
};
|
||||
for threshold in thresholds {
|
||||
if used_percent >= f64::from(*threshold) {
|
||||
warnings.push(serde_json::json!({
|
||||
"accountId": snap.account_id,
|
||||
"email": snap.email,
|
||||
"level": "warning",
|
||||
"threshold": threshold,
|
||||
"usedPercent": used_percent,
|
||||
"quotaKey": quota_key,
|
||||
}));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(serde_json::json!({
|
||||
"warnings": warnings,
|
||||
"updatedAt": china_rfc3339(&Utc::now())
|
||||
}))
|
||||
}
|
||||
pub async fn health_check_accounts(
|
||||
&self,
|
||||
req: serde_json::Value,
|
||||
) -> Result<serde_json::Value, AdminAccountError> {
|
||||
use crate::upstream::accounts::store::StoredAccount;
|
||||
|
||||
let ids = req
|
||||
.get("ids")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(ToString::to_string))
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let mut results = Vec::new();
|
||||
let accounts = if ids.is_empty() {
|
||||
let mut cursor = None;
|
||||
let mut all: Vec<StoredAccount> = Vec::new();
|
||||
loop {
|
||||
let page = self
|
||||
.store
|
||||
.list(cursor, 200)
|
||||
.await
|
||||
.map_err(|_| AdminAccountError::HealthCheck)?;
|
||||
all.extend(page.items);
|
||||
if page.next_cursor.is_none() {
|
||||
break;
|
||||
}
|
||||
cursor = page.next_cursor;
|
||||
}
|
||||
all
|
||||
} else {
|
||||
let mut list = Vec::with_capacity(ids.len());
|
||||
for id in ids {
|
||||
if let Ok(Some(acct)) = self.store.get(&id).await {
|
||||
list.push(acct);
|
||||
}
|
||||
}
|
||||
list
|
||||
};
|
||||
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
for account in &accounts {
|
||||
let token = account.access_token.expose_secret().to_string();
|
||||
let cookie_header = self.usage_cookie_header(&account.id).await;
|
||||
let start = std::time::Instant::now();
|
||||
let context = crate::upstream::transport::CodexRequestContext {
|
||||
access_token: &token,
|
||||
account_id: account.account_id.as_deref(),
|
||||
request_id: &request_id,
|
||||
turn_state: None,
|
||||
turn_metadata: None,
|
||||
beta_features: None,
|
||||
include_timing_metrics: None,
|
||||
version: None,
|
||||
codex_window_id: None,
|
||||
parent_thread_id: None,
|
||||
cookie_header: cookie_header.as_deref(),
|
||||
installation_id: self.installation_id.as_deref(),
|
||||
session_id: None,
|
||||
};
|
||||
match self.codex.probe_connectivity(context).await {
|
||||
Ok(probe) => {
|
||||
let duration = start.elapsed().as_millis();
|
||||
results.push(serde_json::json!({
|
||||
"id": account.id,
|
||||
"email": account.email,
|
||||
"result": "alive",
|
||||
"endpoint": probe.endpoint,
|
||||
"status": probe.status.as_u16(),
|
||||
"durationMs": duration
|
||||
}));
|
||||
}
|
||||
Err(e) => {
|
||||
let duration = start.elapsed().as_millis();
|
||||
results.push(serde_json::json!({
|
||||
"id": account.id,
|
||||
"email": account.email,
|
||||
"result": "dead",
|
||||
"error": e.to_string(),
|
||||
"durationMs": duration
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let total = results.len();
|
||||
let alive = results
|
||||
.iter()
|
||||
.filter(|r| r.get("result") == Some(&serde_json::json!("alive")))
|
||||
.count();
|
||||
let dead = total - alive;
|
||||
|
||||
Ok(serde_json::json!({
|
||||
"summary": { "total": total, "alive": alive, "dead": dead, "skipped": 0 },
|
||||
"results": results
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_warning_windows(quota: &serde_json::Value) -> Vec<(String, f64, &'static str)> {
|
||||
let mut windows = Vec::new();
|
||||
if let Some(monthly) = quota.get("monthly_limit").filter(|value| !value.is_null()) {
|
||||
if let Some(used_percent) = used_percent(monthly) {
|
||||
windows.push(("monthly_limit".to_string(), used_percent, "primary"));
|
||||
}
|
||||
}
|
||||
if let Some(snapshots) = quota.get("snapshots").and_then(serde_json::Value::as_array) {
|
||||
for snapshot in snapshots {
|
||||
let source = snapshot
|
||||
.get("source")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or("quota");
|
||||
let limit_name = snapshot
|
||||
.get("limit_name")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.unwrap_or(source);
|
||||
for role in ["primary", "secondary"] {
|
||||
let Some(window) = snapshot.get(role).filter(|value| !value.is_null()) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(used_percent) = used_percent(window) {
|
||||
windows.push((format!("{source}:{limit_name}:{role}"), used_percent, role));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
windows
|
||||
}
|
||||
|
||||
fn used_percent(value: &serde_json::Value) -> Option<f64> {
|
||||
value
|
||||
.get("used_percent")
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_f64()
|
||||
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().ok()))
|
||||
})
|
||||
.filter(|value| value.is_finite())
|
||||
}
|
||||
|
||||
@@ -16,21 +16,21 @@ use crate::upstream::{
|
||||
|
||||
use super::{types::AdminAccountError, AdminAccountService};
|
||||
|
||||
pub type AccountTestStream = Pin<Box<dyn Stream<Item = Result<Bytes, Infallible>> + Send>>;
|
||||
pub(super) type AccountTestStream = Pin<Box<dyn Stream<Item = Result<Bytes, Infallible>> + Send>>;
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AccountTestModel {
|
||||
pub struct AccountModelOption {
|
||||
pub id: String,
|
||||
pub label: String,
|
||||
}
|
||||
|
||||
impl AdminAccountService {
|
||||
pub async fn test_connection_models(
|
||||
pub async fn account_models(
|
||||
&self,
|
||||
account_id: &str,
|
||||
request_id: &str,
|
||||
) -> Result<Vec<AccountTestModel>, AdminAccountError> {
|
||||
) -> Result<Vec<AccountModelOption>, AdminAccountError> {
|
||||
let account = self
|
||||
.store
|
||||
.get(account_id)
|
||||
@@ -50,8 +50,8 @@ impl AdminAccountService {
|
||||
.fetch_models(&request)
|
||||
.await
|
||||
.map_err(|error| AdminAccountError::FetchModels(error.to_string()))?
|
||||
.into_iter()
|
||||
.filter_map(account_test_model)
|
||||
.iter()
|
||||
.filter_map(account_model_option)
|
||||
.collect::<Vec<_>>();
|
||||
if models.is_empty() {
|
||||
return Err(AdminAccountError::NoModels);
|
||||
@@ -145,9 +145,9 @@ impl AdminAccountService {
|
||||
}
|
||||
}
|
||||
|
||||
fn account_test_model(
|
||||
entry: crate::upstream::models::BackendModelEntry,
|
||||
) -> Option<AccountTestModel> {
|
||||
fn account_model_option(
|
||||
entry: &crate::upstream::models::BackendModelEntry,
|
||||
) -> Option<AccountModelOption> {
|
||||
let id = first_non_empty([
|
||||
entry.slug.as_deref(),
|
||||
entry.id.as_deref(),
|
||||
@@ -162,7 +162,7 @@ fn account_test_model(
|
||||
])
|
||||
.unwrap_or(&id)
|
||||
.to_string();
|
||||
Some(AccountTestModel { id, label })
|
||||
Some(AccountModelOption { id, label })
|
||||
}
|
||||
|
||||
fn first_non_empty<'a>(values: impl IntoIterator<Item = Option<&'a str>>) -> Option<&'a str> {
|
||||
|
||||
@@ -54,8 +54,6 @@ impl From<StoredAccountMetadata> for AdminAccountMetadata {
|
||||
pub enum AdminAccountError {
|
||||
#[error("failed to list accounts")]
|
||||
List,
|
||||
#[error("failed to export accounts")]
|
||||
Export,
|
||||
#[error("failed to import accounts")]
|
||||
Import,
|
||||
#[error("failed to inspect account")]
|
||||
@@ -70,26 +68,14 @@ pub enum AdminAccountError {
|
||||
UpdateStatus,
|
||||
#[error("failed to delete account")]
|
||||
Delete,
|
||||
#[error("failed to load cookies")]
|
||||
LoadCookies,
|
||||
#[error("failed to store cookies")]
|
||||
StoreCookies,
|
||||
#[error("failed to delete cookies")]
|
||||
DeleteCookies,
|
||||
#[error("failed to update claims")]
|
||||
UpdateClaims,
|
||||
#[error("failed to reset usage")]
|
||||
ResetUsage,
|
||||
#[error("failed to sync account pool")]
|
||||
SyncAccountPool,
|
||||
#[error("failed to get quota warnings")]
|
||||
QuotaWarnings,
|
||||
#[error("failed to store quota")]
|
||||
StoreQuota,
|
||||
#[error("failed to fetch quota: {0}")]
|
||||
FetchQuota(String),
|
||||
#[error("health check failed")]
|
||||
HealthCheck,
|
||||
#[error("failed to fetch models: {0}")]
|
||||
FetchModels(String),
|
||||
#[error("no available models returned")]
|
||||
@@ -157,12 +143,6 @@ pub struct BatchDeleteAccounts {
|
||||
pub not_found: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BatchUpdateAccountStatus {
|
||||
pub updated: u32,
|
||||
pub not_found: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ImportedAccounts {
|
||||
pub imported: u32,
|
||||
@@ -259,9 +239,7 @@ pub(super) fn import_usage_string(usage: &serde_json::Value, key: &str) -> Optio
|
||||
usage
|
||||
.get(key)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.and_then(|value| {
|
||||
crate::upstream::accounts::import_export::normalize_nonempty_str(Some(value))
|
||||
})
|
||||
.and_then(|value| crate::upstream::accounts::importing::normalize_nonempty_str(Some(value)))
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
@@ -286,11 +264,6 @@ pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option
|
||||
}
|
||||
|
||||
pub(super) fn parse_account_status(status: &str) -> Result<AcctStatus, AdminAccountError> {
|
||||
crate::upstream::accounts::import_export::parse_account_status(status)
|
||||
.map_err(|_| AdminAccountError::InvalidStatus(status.trim().to_ascii_lowercase()))
|
||||
}
|
||||
|
||||
pub(super) fn parse_batch_account_status(status: &str) -> Result<AcctStatus, AdminAccountError> {
|
||||
crate::upstream::accounts::import_export::parse_batch_account_status(status)
|
||||
crate::upstream::accounts::importing::parse_account_status(status)
|
||||
.map_err(|_| AdminAccountError::InvalidStatus(status.trim().to_ascii_lowercase()))
|
||||
}
|
||||
|
||||
@@ -7,51 +7,6 @@ use uuid::Uuid;
|
||||
|
||||
use crate::infra::identity::{hash_admin_password, verify_admin_password};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 领域类型
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// 管理员会话。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AdminSession {
|
||||
/// 会话 ID。
|
||||
pub id: String,
|
||||
/// 过期时间。
|
||||
pub expires_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
impl AdminSession {
|
||||
/// 判断会话在给定时间是否已经过期。
|
||||
pub fn is_expired_at(&self, now: DateTime<Utc>) -> bool {
|
||||
now >= self.expires_at
|
||||
}
|
||||
}
|
||||
|
||||
/// 根据 TTL 计算会话过期时间。
|
||||
pub fn session_expiry(now: DateTime<Utc>, ttl_minutes: i64) -> DateTime<Utc> {
|
||||
now + Duration::minutes(ttl_minutes.max(0))
|
||||
}
|
||||
|
||||
/// 管理员认证服务。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AdminAuthService {
|
||||
default_username: String,
|
||||
}
|
||||
|
||||
impl AdminAuthService {
|
||||
/// 构造管理员认证服务。
|
||||
pub fn new(default_username: impl Into<String>) -> Self {
|
||||
Self {
|
||||
default_username: default_username.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 判断用户名是否匹配配置中的管理员用户名。
|
||||
pub fn username_matches(&self, username: &str) -> bool {
|
||||
self.default_username == username
|
||||
}
|
||||
}
|
||||
|
||||
/// 管理员登录成功后的会话。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AdminLoginSession {
|
||||
@@ -91,7 +46,6 @@ pub enum AdminSessionError {
|
||||
#[derive(Clone)]
|
||||
pub struct AdminSessionService {
|
||||
store: SqliteAdminSessionStore,
|
||||
auth: AdminAuthService,
|
||||
default_username: String,
|
||||
session_ttl_minutes: u64,
|
||||
}
|
||||
@@ -105,7 +59,6 @@ impl AdminSessionService {
|
||||
) -> Self {
|
||||
Self {
|
||||
store,
|
||||
auth: AdminAuthService::new(default_username.clone()),
|
||||
default_username,
|
||||
session_ttl_minutes,
|
||||
}
|
||||
@@ -139,7 +92,7 @@ impl AdminSessionService {
|
||||
password: &str,
|
||||
) -> Result<Option<AdminLoginSession>, AdminSessionError> {
|
||||
let username = username.unwrap_or(&self.default_username);
|
||||
if !self.auth.username_matches(username) {
|
||||
if username != self.default_username.as_str() {
|
||||
return Ok(None);
|
||||
}
|
||||
let Some(admin) = self
|
||||
@@ -209,11 +162,6 @@ impl SqliteAdminSessionStore {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// 返回底层连接池。
|
||||
pub fn pool(&self) -> &SqlitePool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
/// 如果还没有管理员用户,则创建默认管理员。
|
||||
pub async fn ensure_default_admin(
|
||||
&self,
|
||||
|
||||
+20
-61
@@ -4,57 +4,42 @@ use axum::{
|
||||
extract::State,
|
||||
http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
Extension, Json,
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{
|
||||
admin::response::{AdminEnvelope, AdminError, AdminResponse},
|
||||
http::middleware::request_id::RequestId,
|
||||
infra::time::china_rfc3339,
|
||||
runtime::state::AppState,
|
||||
};
|
||||
|
||||
const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session";
|
||||
|
||||
/// 管理员登录请求。
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminLoginRequest {
|
||||
/// 管理员用户名;缺省时使用配置中的默认管理员。
|
||||
pub username: Option<String>,
|
||||
/// 管理员密码。
|
||||
pub password: String,
|
||||
pub(crate) struct AdminLoginRequest {
|
||||
username: Option<String>,
|
||||
password: String,
|
||||
}
|
||||
|
||||
/// 管理员登录响应。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminLoginData {
|
||||
/// 会话过期时间。
|
||||
pub expires_at: String,
|
||||
struct AdminLoginData {
|
||||
expires_at: String,
|
||||
}
|
||||
|
||||
/// 管理员会话状态响应。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminSessionStatusData {
|
||||
/// 当前请求是否携带有效管理端会话。
|
||||
pub authenticated: bool,
|
||||
}
|
||||
|
||||
/// 会话登录是否成功。
|
||||
pub fn session_login_allowed() -> bool {
|
||||
true
|
||||
struct AdminSessionStatusData {
|
||||
authenticated: bool,
|
||||
}
|
||||
|
||||
/// `POST /api/admin/login`
|
||||
pub async fn login(
|
||||
pub(crate) async fn login(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
Json(payload): Json<AdminLoginRequest>,
|
||||
) -> Result<Response, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
let session = state
|
||||
.services
|
||||
.admin_sessions
|
||||
@@ -65,26 +50,17 @@ pub async fn login(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
"Failed to create admin session",
|
||||
request_id.clone(),
|
||||
)
|
||||
})?
|
||||
.ok_or_else(|| {
|
||||
AdminError::new(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
40102,
|
||||
"Invalid admin credentials",
|
||||
request_id.clone(),
|
||||
)
|
||||
AdminError::new(StatusCode::UNAUTHORIZED, 40102, "Invalid admin credentials")
|
||||
})?;
|
||||
|
||||
let mut response = AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(
|
||||
AdminLoginData {
|
||||
expires_at: china_rfc3339(&session.expires_at),
|
||||
},
|
||||
request_id.clone(),
|
||||
),
|
||||
AdminEnvelope::ok(AdminLoginData {
|
||||
expires_at: china_rfc3339(&session.expires_at),
|
||||
}),
|
||||
)
|
||||
.into_response();
|
||||
let cookie = format!(
|
||||
@@ -98,7 +74,6 @@ pub async fn login(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
"Failed to create admin session cookie",
|
||||
request_id.clone(),
|
||||
)
|
||||
})?,
|
||||
);
|
||||
@@ -106,12 +81,10 @@ pub async fn login(
|
||||
}
|
||||
|
||||
/// `GET /api/admin/auth/status`
|
||||
pub async fn session_status(
|
||||
pub(crate) async fn session_status(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
let authenticated = state
|
||||
.services
|
||||
.admin_sessions
|
||||
@@ -122,25 +95,20 @@ pub async fn session_status(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
"Failed to validate admin session",
|
||||
request_id.clone(),
|
||||
)
|
||||
})?;
|
||||
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminSessionStatusData { authenticated }, request_id),
|
||||
AdminEnvelope::ok(AdminSessionStatusData { authenticated }),
|
||||
))
|
||||
}
|
||||
|
||||
/// `POST /api/admin/logout`
|
||||
pub async fn logout(
|
||||
pub(crate) async fn logout(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Response, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
|
||||
// 获取并删除服务器端 session
|
||||
if let Some(session_id) = admin_session_cookie(&headers) {
|
||||
let _ = state
|
||||
.services
|
||||
@@ -149,19 +117,14 @@ pub async fn logout(
|
||||
.await;
|
||||
}
|
||||
|
||||
// 返回 JSON 响应并清除 cookie
|
||||
let mut response = AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(
|
||||
serde_json::json!({
|
||||
"message": "Logged out successfully"
|
||||
}),
|
||||
request_id.clone(),
|
||||
),
|
||||
AdminEnvelope::ok(serde_json::json!({
|
||||
"message": "Logged out successfully"
|
||||
})),
|
||||
)
|
||||
.into_response();
|
||||
|
||||
// 清除 cookie
|
||||
let cookie = format!("{ADMIN_SESSION_COOKIE}=; Path=/; HttpOnly; SameSite=Lax; Max-Age=0");
|
||||
response.headers_mut().insert(
|
||||
SET_COOKIE,
|
||||
@@ -170,7 +133,6 @@ pub async fn logout(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
"Failed to clear admin session cookie",
|
||||
request_id.clone(),
|
||||
)
|
||||
})?,
|
||||
);
|
||||
@@ -179,10 +141,9 @@ pub async fn logout(
|
||||
}
|
||||
|
||||
/// 要求请求携带有效管理员会话。
|
||||
pub async fn require_admin_session(
|
||||
pub(crate) async fn require_admin_session(
|
||||
state: &AppState,
|
||||
headers: &HeaderMap,
|
||||
request_id: &str,
|
||||
) -> Result<(), AdminError> {
|
||||
match state
|
||||
.services
|
||||
@@ -195,13 +156,11 @@ pub async fn require_admin_session(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
40101,
|
||||
"Admin session required",
|
||||
request_id,
|
||||
)),
|
||||
Err(_) => Err(AdminError::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
"Failed to validate admin session",
|
||||
request_id,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
+74
-200
@@ -1,10 +1,10 @@
|
||||
//! 管理端 v1 接口访问 Key 处理器(列表、创建、生命周期、导出)。
|
||||
//! 管理端 v1 接口访问 Key 处理器(列表、创建、生命周期)。
|
||||
|
||||
use axum::{
|
||||
extract::{Query, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
Extension, Json,
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
@@ -15,7 +15,6 @@ use crate::{
|
||||
AdminClientKeyError, AdminCreatedClientApiKey, AdminStoredClientApiKey,
|
||||
},
|
||||
admin::response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse},
|
||||
http::middleware::request_id::RequestId,
|
||||
infra::{
|
||||
json::{clamp_limit, Page},
|
||||
time::{china_relative_time_str, china_rfc3339_str},
|
||||
@@ -23,94 +22,62 @@ use crate::{
|
||||
runtime::state::AppState,
|
||||
};
|
||||
|
||||
// ---- Query types ----
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ApiKeysQuery {
|
||||
pub cursor: Option<String>,
|
||||
pub limit: Option<u32>,
|
||||
pub(crate) struct ApiKeysQuery {
|
||||
cursor: Option<String>,
|
||||
limit: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ApiKeyExportQuery {
|
||||
pub ids: Option<String>,
|
||||
pub(crate) struct CreateApiKeyRequest {
|
||||
name: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct CreateApiKeyRequest {
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BatchDeleteClientApiKeysRequest {
|
||||
pub ids: Vec<String>,
|
||||
}
|
||||
|
||||
// ---- Response types ----
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ClientApiKeyData {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub prefix: String,
|
||||
pub key: String,
|
||||
pub enabled: bool,
|
||||
pub created_at: String,
|
||||
pub created_at_display: String,
|
||||
pub last_used_at: Option<String>,
|
||||
pub last_used_at_display: String,
|
||||
pub(crate) struct BatchDeleteClientApiKeysRequest {
|
||||
ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct CreatedClientApiKeyData {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub prefix: String,
|
||||
pub key: String,
|
||||
pub enabled: bool,
|
||||
pub created_at: String,
|
||||
pub created_at_display: String,
|
||||
pub last_used_at: Option<String>,
|
||||
pub last_used_at_display: String,
|
||||
struct ClientApiKeyData {
|
||||
id: String,
|
||||
name: String,
|
||||
label: Option<String>,
|
||||
prefix: String,
|
||||
key: String,
|
||||
enabled: bool,
|
||||
created_at: String,
|
||||
created_at_display: String,
|
||||
last_used_at: Option<String>,
|
||||
last_used_at_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BatchDeleteClientApiKeysData {
|
||||
pub deleted: u32,
|
||||
pub not_found: Vec<String>,
|
||||
struct CreatedClientApiKeyData {
|
||||
id: String,
|
||||
name: String,
|
||||
label: Option<String>,
|
||||
prefix: String,
|
||||
key: String,
|
||||
enabled: bool,
|
||||
created_at: String,
|
||||
created_at_display: String,
|
||||
last_used_at: Option<String>,
|
||||
last_used_at_display: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ClientApiKeyExportData {
|
||||
pub source_format: &'static str,
|
||||
pub rotation_required: bool,
|
||||
pub api_keys: Vec<ClientApiKeyExportEntry>,
|
||||
struct BatchDeleteClientApiKeysData {
|
||||
deleted: u32,
|
||||
not_found: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ClientApiKeyExportEntry {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub label: Option<String>,
|
||||
pub prefix: String,
|
||||
pub enabled: bool,
|
||||
pub created_at: String,
|
||||
pub last_used_at: Option<String>,
|
||||
}
|
||||
|
||||
// ---- Conversions ----
|
||||
|
||||
impl From<AdminStoredClientApiKey> for ClientApiKeyData {
|
||||
fn from(k: AdminStoredClientApiKey) -> Self {
|
||||
Self {
|
||||
@@ -128,20 +95,6 @@ impl From<AdminStoredClientApiKey> for ClientApiKeyData {
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AdminStoredClientApiKey> for ClientApiKeyExportEntry {
|
||||
fn from(k: AdminStoredClientApiKey) -> Self {
|
||||
Self {
|
||||
id: k.id,
|
||||
name: k.name,
|
||||
label: k.label,
|
||||
prefix: k.prefix,
|
||||
enabled: k.enabled,
|
||||
created_at: china_rfc3339_str(&k.created_at),
|
||||
last_used_at: k.last_used_at.as_deref().map(china_rfc3339_str),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AdminCreatedClientApiKey> for CreatedClientApiKeyData {
|
||||
fn from(k: AdminCreatedClientApiKey) -> Self {
|
||||
Self {
|
||||
@@ -159,33 +112,20 @@ impl From<AdminCreatedClientApiKey> for CreatedClientApiKeyData {
|
||||
}
|
||||
}
|
||||
|
||||
fn client_key_error(error: AdminClientKeyError, request_id: String) -> AdminError {
|
||||
fn client_key_error(error: &AdminClientKeyError) -> AdminError {
|
||||
match error {
|
||||
AdminClientKeyError::InvalidStatus(_)
|
||||
| AdminClientKeyError::EmptyName
|
||||
| AdminClientKeyError::EmptyIds
|
||||
| AdminClientKeyError::LabelTooLong => AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
error.to_string(),
|
||||
request_id,
|
||||
),
|
||||
_ => AdminError::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
error.to_string(),
|
||||
request_id,
|
||||
),
|
||||
| AdminClientKeyError::LabelTooLong => {
|
||||
AdminError::new(StatusCode::BAD_REQUEST, 40001, error.to_string())
|
||||
}
|
||||
_ => AdminError::new(StatusCode::INTERNAL_SERVER_ERROR, 50001, error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn client_key_not_found(request_id: String) -> AdminError {
|
||||
AdminError::new(
|
||||
StatusCode::NOT_FOUND,
|
||||
40401,
|
||||
"Client API key not found",
|
||||
request_id,
|
||||
)
|
||||
fn client_key_not_found() -> AdminError {
|
||||
AdminError::new(StatusCode::NOT_FOUND, 40401, "Client API key not found")
|
||||
}
|
||||
|
||||
struct ParsedClientApiKeyUpdate {
|
||||
@@ -194,33 +134,28 @@ struct ParsedClientApiKeyUpdate {
|
||||
status: Option<String>,
|
||||
}
|
||||
|
||||
fn parse_client_api_key_update(
|
||||
payload: Value,
|
||||
request_id: &str,
|
||||
) -> Result<ParsedClientApiKeyUpdate, AdminError> {
|
||||
fn parse_client_api_key_update(payload: &Value) -> Result<ParsedClientApiKeyUpdate, AdminError> {
|
||||
let object = payload.as_object().ok_or_else(|| {
|
||||
AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
"Client API key update request must be an object",
|
||||
request_id,
|
||||
)
|
||||
})?;
|
||||
let id = required_string_field(object, "id", request_id)?;
|
||||
let id = required_string_field(object, "id")?;
|
||||
let label = object
|
||||
.get("label")
|
||||
.map(|value| optional_string_field(value, "label", request_id))
|
||||
.map(|value| optional_string_field(value, "label"))
|
||||
.transpose()?;
|
||||
let status = object
|
||||
.get("status")
|
||||
.map(|value| required_string_value(value, "status", request_id))
|
||||
.map(|value| required_string_value(value, "status"))
|
||||
.transpose()?;
|
||||
if label.is_none() && status.is_none() {
|
||||
return Err(AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
"Client API key update request must include label or status",
|
||||
request_id,
|
||||
));
|
||||
}
|
||||
Ok(ParsedClientApiKeyUpdate { id, label, status })
|
||||
@@ -229,24 +164,18 @@ fn parse_client_api_key_update(
|
||||
fn required_string_field(
|
||||
object: &serde_json::Map<String, Value>,
|
||||
field: &'static str,
|
||||
request_id: &str,
|
||||
) -> Result<String, AdminError> {
|
||||
let Some(value) = object.get(field) else {
|
||||
return Err(AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
format!("{field} is required"),
|
||||
request_id,
|
||||
));
|
||||
};
|
||||
required_string_value(value, field, request_id)
|
||||
required_string_value(value, field)
|
||||
}
|
||||
|
||||
fn required_string_value(
|
||||
value: &Value,
|
||||
field: &'static str,
|
||||
request_id: &str,
|
||||
) -> Result<String, AdminError> {
|
||||
fn required_string_value(value: &Value, field: &'static str) -> Result<String, AdminError> {
|
||||
value
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
@@ -257,16 +186,11 @@ fn required_string_value(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
format!("{field} must be a non-empty string"),
|
||||
request_id,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn optional_string_field(
|
||||
value: &Value,
|
||||
field: &'static str,
|
||||
request_id: &str,
|
||||
) -> Result<Option<String>, AdminError> {
|
||||
fn optional_string_field(value: &Value, field: &'static str) -> Result<Option<String>, AdminError> {
|
||||
if value.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
@@ -279,22 +203,17 @@ fn optional_string_field(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
format!("{field} must be a string or null"),
|
||||
request_id,
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
// ---- Handlers ----
|
||||
|
||||
/// `GET /api/admin/keys`
|
||||
pub async fn api_keys(
|
||||
pub(crate) async fn api_keys(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<ApiKeysQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
let limit = clamp_limit(query.limit.unwrap_or(50));
|
||||
match state
|
||||
.services
|
||||
@@ -309,42 +228,38 @@ pub async fn api_keys(
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminPageEnvelope::ok(page, limit, request_id),
|
||||
AdminPageEnvelope::ok(page, limit),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(client_key_error(error, request_id)),
|
||||
Err(error) => Err(client_key_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /api/admin/keys`
|
||||
pub async fn create_api_key(
|
||||
pub(crate) async fn create_api_key(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Json(payload): Json<CreateApiKeyRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
match state.services.admin_client_keys.create(&payload.name).await {
|
||||
Ok(created) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(CreatedClientApiKeyData::from(created), request_id),
|
||||
AdminEnvelope::ok(CreatedClientApiKeyData::from(created)),
|
||||
)),
|
||||
Err(error) => Err(client_key_error(error, request_id)),
|
||||
Err(error) => Err(client_key_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /api/admin/keys/update`
|
||||
pub async fn update_api_key(
|
||||
pub(crate) async fn update_api_key(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Json(payload): Json<Value>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
|
||||
let update = parse_client_api_key_update(payload, &request_id)?;
|
||||
let update = parse_client_api_key_update(&payload)?;
|
||||
if let Some(label) = update.label {
|
||||
match state
|
||||
.services
|
||||
@@ -353,8 +268,8 @@ pub async fn update_api_key(
|
||||
.await
|
||||
{
|
||||
Ok(Some(_)) => {}
|
||||
Ok(None) => return Err(client_key_not_found(request_id)),
|
||||
Err(error) => return Err(client_key_error(error, request_id)),
|
||||
Ok(None) => return Err(client_key_not_found()),
|
||||
Err(error) => return Err(client_key_error(&error)),
|
||||
}
|
||||
}
|
||||
if let Some(status) = update.status {
|
||||
@@ -364,31 +279,29 @@ pub async fn update_api_key(
|
||||
.update_status(&update.id, &status)
|
||||
.await
|
||||
{
|
||||
Ok(Some(_)) => {}
|
||||
Ok(None) => return Err(client_key_not_found(request_id)),
|
||||
Err(error) => return Err(client_key_error(error, request_id)),
|
||||
Ok(true) => {}
|
||||
Ok(false) => return Err(client_key_not_found()),
|
||||
Err(error) => return Err(client_key_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
match state.services.admin_client_keys.get(&update.id).await {
|
||||
Ok(Some(key)) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(ClientApiKeyData::from(key), request_id),
|
||||
AdminEnvelope::ok(ClientApiKeyData::from(key)),
|
||||
)),
|
||||
Ok(None) => Err(client_key_not_found(request_id)),
|
||||
Err(error) => Err(client_key_error(error, request_id)),
|
||||
Ok(None) => Err(client_key_not_found()),
|
||||
Err(error) => Err(client_key_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /api/admin/keys/delete`
|
||||
pub async fn batch_delete_api_keys(
|
||||
pub(crate) async fn batch_delete_api_keys(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Json(payload): Json<BatchDeleteClientApiKeysRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
match state
|
||||
.services
|
||||
.admin_client_keys
|
||||
@@ -397,50 +310,11 @@ pub async fn batch_delete_api_keys(
|
||||
{
|
||||
Ok(deleted) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(
|
||||
BatchDeleteClientApiKeysData {
|
||||
deleted: deleted.deleted,
|
||||
not_found: deleted.not_found,
|
||||
},
|
||||
request_id,
|
||||
),
|
||||
AdminEnvelope::ok(BatchDeleteClientApiKeysData {
|
||||
deleted: deleted.deleted,
|
||||
not_found: deleted.not_found,
|
||||
}),
|
||||
)),
|
||||
Err(error) => Err(client_key_error(error, request_id)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `GET /api/admin/keys/export`
|
||||
pub async fn export_api_keys(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<ApiKeyExportQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
let ids = query
|
||||
.ids
|
||||
.unwrap_or_default()
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|id| !id.is_empty())
|
||||
.map(ToString::to_string)
|
||||
.collect();
|
||||
match state.services.admin_client_keys.export(ids).await {
|
||||
Ok(keys) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(
|
||||
ClientApiKeyExportData {
|
||||
source_format: "rustLocalClientApiKeys",
|
||||
rotation_required: true,
|
||||
api_keys: keys
|
||||
.into_iter()
|
||||
.map(ClientApiKeyExportEntry::from)
|
||||
.collect(),
|
||||
},
|
||||
request_id,
|
||||
),
|
||||
)),
|
||||
Err(error) => Err(client_key_error(error, request_id)),
|
||||
Err(error) => Err(client_key_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
+13
-61
@@ -149,16 +149,12 @@ impl AdminClientKeyService {
|
||||
&self,
|
||||
key_id: &str,
|
||||
status: &str,
|
||||
) -> Result<Option<UpdatedClientApiKeyStatus>, AdminClientKeyError> {
|
||||
) -> Result<bool, AdminClientKeyError> {
|
||||
let enabled = parse_client_key_status(status)?;
|
||||
match self.store.set_enabled(key_id, enabled).await {
|
||||
Ok(true) => Ok(Some(UpdatedClientApiKeyStatus {
|
||||
id: key_id.to_string(),
|
||||
enabled,
|
||||
})),
|
||||
Ok(false) => Ok(None),
|
||||
Err(_) => Err(AdminClientKeyError::UpdateStatus),
|
||||
}
|
||||
self.store
|
||||
.set_enabled(key_id, enabled)
|
||||
.await
|
||||
.map_err(|_| AdminClientKeyError::UpdateStatus)
|
||||
}
|
||||
|
||||
/// 批量删除客户端 API Key。
|
||||
@@ -180,38 +176,6 @@ impl AdminClientKeyService {
|
||||
}
|
||||
Ok(BatchDeleteClientApiKeys { deleted, not_found })
|
||||
}
|
||||
|
||||
/// 导出客户端 API Key 元数据。
|
||||
pub async fn export(
|
||||
&self,
|
||||
ids: Vec<String>,
|
||||
) -> Result<Vec<AdminStoredClientApiKey>, AdminClientKeyError> {
|
||||
if ids.is_empty() {
|
||||
let mut all_keys = Vec::new();
|
||||
let mut cursor = None;
|
||||
loop {
|
||||
let page = self
|
||||
.store
|
||||
.list(cursor, 200)
|
||||
.await
|
||||
.map_err(|_| AdminClientKeyError::Export)?;
|
||||
all_keys.extend(page.items.into_iter().map(AdminStoredClientApiKey::from));
|
||||
if page.next_cursor.is_none() {
|
||||
return Ok(all_keys);
|
||||
}
|
||||
cursor = page.next_cursor;
|
||||
}
|
||||
}
|
||||
let mut keys = Vec::with_capacity(ids.len());
|
||||
for id in ids {
|
||||
match self.store.get(&id).await {
|
||||
Ok(Some(key)) => keys.push(AdminStoredClientApiKey::from(key)),
|
||||
Ok(None) => {}
|
||||
Err(_) => return Err(AdminClientKeyError::Export),
|
||||
}
|
||||
}
|
||||
Ok(keys)
|
||||
}
|
||||
}
|
||||
|
||||
/// 管理端可见的客户端 API Key 元数据。
|
||||
@@ -240,13 +204,6 @@ pub struct AdminCreatedClientApiKey {
|
||||
pub last_used_at: Option<String>,
|
||||
}
|
||||
|
||||
/// 客户端 API Key 状态更新结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UpdatedClientApiKeyStatus {
|
||||
pub id: String,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
/// 客户端 API Key 批量删除结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct BatchDeleteClientApiKeys {
|
||||
@@ -259,8 +216,6 @@ pub struct BatchDeleteClientApiKeys {
|
||||
pub enum AdminClientKeyError {
|
||||
#[error("failed to list client API keys")]
|
||||
List,
|
||||
#[error("failed to export client API keys")]
|
||||
Export,
|
||||
#[error("failed to create client API key")]
|
||||
Create,
|
||||
#[error("failed to delete client API key")]
|
||||
@@ -455,8 +410,8 @@ impl SqliteClientKeyStore {
|
||||
let items = rows
|
||||
.into_iter()
|
||||
.take(take_count)
|
||||
.map(|row| self.key_from_row(&row))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
.map(|row| Self::key_from_row(&row))
|
||||
.collect::<Vec<_>>();
|
||||
let next_cursor = if has_next {
|
||||
items
|
||||
.last()
|
||||
@@ -480,7 +435,7 @@ impl SqliteClientKeyStore {
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
row.map(|row| self.key_from_row(&row)).transpose()
|
||||
Ok(row.map(|row| Self::key_from_row(&row)))
|
||||
}
|
||||
|
||||
/// 更新客户端 API Key 启用状态。
|
||||
@@ -526,11 +481,8 @@ impl SqliteClientKeyStore {
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
fn key_from_row(
|
||||
&self,
|
||||
row: &sqlx::sqlite::SqliteRow,
|
||||
) -> Result<StoredClientApiKey, SqliteClientKeyStoreError> {
|
||||
Ok(StoredClientApiKey {
|
||||
fn key_from_row(row: &sqlx::sqlite::SqliteRow) -> StoredClientApiKey {
|
||||
StoredClientApiKey {
|
||||
id: row.get("id"),
|
||||
name: row.get("name"),
|
||||
label: row.get("label"),
|
||||
@@ -539,7 +491,7 @@ impl SqliteClientKeyStore {
|
||||
enabled: row.get::<i64, _>("enabled") != 0,
|
||||
created_at: row.get("created_at"),
|
||||
last_used_at: row.get("last_used_at"),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -548,7 +500,7 @@ impl ClientKeyStore for SqliteClientKeyStore {
|
||||
async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult<bool> {
|
||||
verify_and_touch(&self.pool, plaintext)
|
||||
.await
|
||||
.map_err(map_client_key_store_error)
|
||||
.map_err(|error| map_client_key_store_error(&error))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -574,7 +526,7 @@ async fn verify_and_touch(
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
fn map_client_key_store_error(error: SqliteClientKeyStoreError) -> ClientKeyStoreError {
|
||||
fn map_client_key_store_error(error: &SqliteClientKeyStoreError) -> ClientKeyStoreError {
|
||||
ClientKeyStoreError::OperationFailed {
|
||||
message: error.to_string(),
|
||||
}
|
||||
|
||||
@@ -6,8 +6,6 @@ pub mod accounts;
|
||||
pub mod auth;
|
||||
/// v1 接口访问 Key 处理器。
|
||||
pub mod keys;
|
||||
/// 模型处理器。
|
||||
pub mod models;
|
||||
/// 管理面日志、用量与诊断统计。
|
||||
pub mod monitoring;
|
||||
/// 响应封装。
|
||||
|
||||
@@ -1,3 +0,0 @@
|
||||
//! 管理端模型 HTTP 模块。
|
||||
|
||||
pub mod routes;
|
||||
@@ -1,92 +0,0 @@
|
||||
//! 管理端模型处理器。
|
||||
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
Extension,
|
||||
};
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::{
|
||||
admin::auth::session::require_admin_session,
|
||||
admin::response::{AdminEnvelope, AdminError, AdminResponse},
|
||||
http::middleware::request_id::RequestId,
|
||||
runtime::state::AppState,
|
||||
upstream::models::{AdminModelError, ModelRefreshResult},
|
||||
};
|
||||
|
||||
/// 模型刷新响应。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RefreshModelsData {
|
||||
pub refreshed_plans: usize,
|
||||
pub model_count: usize,
|
||||
pub failed_plans: usize,
|
||||
}
|
||||
|
||||
/// `POST /api/admin/models/refresh`
|
||||
pub async fn refresh_models(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
|
||||
match state
|
||||
.services
|
||||
.admin_models
|
||||
.refresh_backend_models(&request_id)
|
||||
.await
|
||||
{
|
||||
Ok(result) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(RefreshModelsData::from(result), request_id),
|
||||
)),
|
||||
Err(error) => model_error(error, request_id),
|
||||
}
|
||||
}
|
||||
|
||||
fn model_error(
|
||||
error: AdminModelError,
|
||||
request_id: String,
|
||||
) -> Result<AdminResponse<AdminEnvelope<RefreshModelsData>>, AdminError> {
|
||||
match error {
|
||||
AdminModelError::NoAccounts => Err(AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
error.to_string(),
|
||||
request_id,
|
||||
)),
|
||||
AdminModelError::AllPlansFailed(result) => Ok(AdminResponse::new(
|
||||
StatusCode::BAD_GATEWAY,
|
||||
AdminEnvelope::new(
|
||||
50201,
|
||||
"Failed to refresh backend models",
|
||||
RefreshModelsData::from(result),
|
||||
request_id,
|
||||
),
|
||||
)),
|
||||
AdminModelError::ListAccounts
|
||||
| AdminModelError::SnapshotStoreUnavailable
|
||||
| AdminModelError::UpstreamClientUnavailable
|
||||
| AdminModelError::StoreSnapshot
|
||||
| AdminModelError::LoadSnapshots => Err(AdminError::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
error.to_string(),
|
||||
request_id,
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ModelRefreshResult> for RefreshModelsData {
|
||||
fn from(result: ModelRefreshResult) -> Self {
|
||||
Self {
|
||||
refreshed_plans: result.refreshed_plans,
|
||||
model_count: result.model_count,
|
||||
failed_plans: result.failed_plans,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -151,7 +151,7 @@ const MODEL_PRICING_RULES: &[ModelPricingRule] = &[
|
||||
];
|
||||
|
||||
/// 按 codex2api 的 billing 引擎口径计算美元计费。
|
||||
pub(crate) fn calculate_cost(
|
||||
pub(in crate::admin) fn calculate_cost(
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
@@ -240,71 +240,62 @@ fn normalize_billing_model_name(model: &str) -> String {
|
||||
|
||||
fn normalize_codex_billing_model(model: &str) -> Option<&'static str> {
|
||||
let compact = model.replace([' ', '_'], "-").to_ascii_lowercase();
|
||||
match () {
|
||||
_ if compact.contains("gpt-5.5-pro")
|
||||
|| compact.contains("gpt5-5-pro")
|
||||
|| compact.contains("gpt5.5-pro") =>
|
||||
{
|
||||
Some("gpt-5.5-pro")
|
||||
}
|
||||
_ if compact.contains("gpt-5.5")
|
||||
|| compact.contains("gpt5-5")
|
||||
|| compact.contains("gpt5.5") =>
|
||||
{
|
||||
Some("gpt-5.5")
|
||||
}
|
||||
_ if compact.contains("gpt-5.4-mini")
|
||||
|| compact.contains("gpt5-4-mini")
|
||||
|| compact.contains("gpt5.4-mini") =>
|
||||
{
|
||||
Some("gpt-5.4-mini")
|
||||
}
|
||||
_ if compact.contains("gpt-5.4-nano")
|
||||
|| compact.contains("gpt5-4-nano")
|
||||
|| compact.contains("gpt5.4-nano") =>
|
||||
{
|
||||
Some("gpt-5.4-nano")
|
||||
}
|
||||
_ if compact.contains("gpt-5.4-pro")
|
||||
|| compact.contains("gpt5-4-pro")
|
||||
|| compact.contains("gpt5.4-pro") =>
|
||||
{
|
||||
Some("gpt-5.4-pro")
|
||||
}
|
||||
_ if compact.contains("gpt-5.4")
|
||||
|| compact.contains("gpt5-4")
|
||||
|| compact.contains("gpt5.4") =>
|
||||
{
|
||||
Some("gpt-5.4")
|
||||
}
|
||||
_ if compact.contains("gpt-5.2")
|
||||
|| compact.contains("gpt5-2")
|
||||
|| compact.contains("gpt5.2") =>
|
||||
{
|
||||
Some("gpt-5.2")
|
||||
}
|
||||
_ if compact.contains("gpt-5.3-codex-spark")
|
||||
|| compact.contains("gpt5-3-codex-spark")
|
||||
|| compact.contains("gpt5.3-codex-spark") =>
|
||||
{
|
||||
Some("gpt-5.3-codex-spark")
|
||||
}
|
||||
_ if compact.contains("gpt-5.3-codex")
|
||||
|| compact.contains("gpt5-3-codex")
|
||||
|| compact.contains("gpt5.3-codex") =>
|
||||
{
|
||||
Some("gpt-5.3-codex")
|
||||
}
|
||||
_ if compact.contains("gpt-5.3")
|
||||
|| compact.contains("gpt5-3")
|
||||
|| compact.contains("gpt5.3") =>
|
||||
{
|
||||
Some("gpt-5.3-codex")
|
||||
}
|
||||
_ if compact.contains("codex-auto-review") => Some("gpt-5.4"),
|
||||
_ if compact.contains("codex") => Some("gpt-5.3-codex"),
|
||||
_ if compact.contains("gpt-5") || compact.contains("gpt5") => Some("gpt-5.4"),
|
||||
_ => None,
|
||||
if compact.contains("gpt-5.5-pro")
|
||||
|| compact.contains("gpt5-5-pro")
|
||||
|| compact.contains("gpt5.5-pro")
|
||||
{
|
||||
Some("gpt-5.5-pro")
|
||||
} else if compact.contains("gpt-5.5")
|
||||
|| compact.contains("gpt5-5")
|
||||
|| compact.contains("gpt5.5")
|
||||
{
|
||||
Some("gpt-5.5")
|
||||
} else if compact.contains("gpt-5.4-mini")
|
||||
|| compact.contains("gpt5-4-mini")
|
||||
|| compact.contains("gpt5.4-mini")
|
||||
{
|
||||
Some("gpt-5.4-mini")
|
||||
} else if compact.contains("gpt-5.4-nano")
|
||||
|| compact.contains("gpt5-4-nano")
|
||||
|| compact.contains("gpt5.4-nano")
|
||||
{
|
||||
Some("gpt-5.4-nano")
|
||||
} else if compact.contains("gpt-5.4-pro")
|
||||
|| compact.contains("gpt5-4-pro")
|
||||
|| compact.contains("gpt5.4-pro")
|
||||
{
|
||||
Some("gpt-5.4-pro")
|
||||
} else if compact.contains("gpt-5.4")
|
||||
|| compact.contains("gpt5-4")
|
||||
|| compact.contains("gpt5.4")
|
||||
{
|
||||
Some("gpt-5.4")
|
||||
} else if compact.contains("gpt-5.2")
|
||||
|| compact.contains("gpt5-2")
|
||||
|| compact.contains("gpt5.2")
|
||||
{
|
||||
Some("gpt-5.2")
|
||||
} else if compact.contains("gpt-5.3-codex-spark")
|
||||
|| compact.contains("gpt5-3-codex-spark")
|
||||
|| compact.contains("gpt5.3-codex-spark")
|
||||
{
|
||||
Some("gpt-5.3-codex-spark")
|
||||
} else if compact.contains("gpt-5.3-codex")
|
||||
|| compact.contains("gpt5-3-codex")
|
||||
|| compact.contains("gpt5.3-codex")
|
||||
|| compact.contains("gpt-5.3")
|
||||
|| compact.contains("gpt5-3")
|
||||
|| compact.contains("gpt5.3")
|
||||
{
|
||||
Some("gpt-5.3-codex")
|
||||
} else if compact.contains("codex-auto-review") {
|
||||
Some("gpt-5.4")
|
||||
} else if compact.contains("codex") {
|
||||
Some("gpt-5.3-codex")
|
||||
} else if compact.contains("gpt-5") || compact.contains("gpt5") {
|
||||
Some("gpt-5.4")
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
@@ -327,30 +318,30 @@ fn model_matches_rule(model: &str, rule: &str) -> bool {
|
||||
}
|
||||
|
||||
fn claude_family_pricing(model: &str) -> Option<ModelPricing> {
|
||||
match () {
|
||||
_ if model.contains("opus") => {
|
||||
if model.contains("4.7")
|
||||
|| model.contains("4-7")
|
||||
|| model.contains("4.6")
|
||||
|| model.contains("4-6")
|
||||
|| model.contains("4.5")
|
||||
|| model.contains("4-5")
|
||||
{
|
||||
Some(ModelPricing::new(5.0, 25.0))
|
||||
} else {
|
||||
Some(ModelPricing::new(15.0, 75.0))
|
||||
}
|
||||
if model.contains("opus") {
|
||||
if model.contains("4.7")
|
||||
|| model.contains("4-7")
|
||||
|| model.contains("4.6")
|
||||
|| model.contains("4-6")
|
||||
|| model.contains("4.5")
|
||||
|| model.contains("4-5")
|
||||
{
|
||||
Some(ModelPricing::new(5.0, 25.0))
|
||||
} else {
|
||||
Some(ModelPricing::new(15.0, 75.0))
|
||||
}
|
||||
_ if model.contains("sonnet") => Some(ModelPricing::new(3.0, 15.0)),
|
||||
_ if model.contains("haiku") => {
|
||||
if model.contains("3-5") || model.contains("3.5") {
|
||||
Some(ModelPricing::new(1.0, 5.0))
|
||||
} else {
|
||||
Some(ModelPricing::new(0.25, 1.25))
|
||||
}
|
||||
} else if model.contains("sonnet") {
|
||||
Some(ModelPricing::new(3.0, 15.0))
|
||||
} else if model.contains("haiku") {
|
||||
if model.contains("3-5") || model.contains("3.5") {
|
||||
Some(ModelPricing::new(1.0, 5.0))
|
||||
} else {
|
||||
Some(ModelPricing::new(0.25, 1.25))
|
||||
}
|
||||
_ if model.contains("claude") => Some(ModelPricing::new(3.0, 15.0)),
|
||||
_ => None,
|
||||
} else if model.contains("claude") {
|
||||
Some(ModelPricing::new(3.0, 15.0))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+111
-124
@@ -6,7 +6,6 @@ use axum::{
|
||||
extract::{Query, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
Extension,
|
||||
};
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -26,7 +25,6 @@ use crate::{
|
||||
},
|
||||
response::{AdminEnvelope, AdminError, AdminResponse},
|
||||
},
|
||||
http::middleware::request_id::RequestId,
|
||||
infra::time::{
|
||||
china_datetime, china_datetime_rfc3339_str, china_day_start, china_hour, china_hour_start,
|
||||
china_quarter_hour_start, china_relative_time, china_time,
|
||||
@@ -41,13 +39,13 @@ const HEALTH_TIMELINE_SLOTS: i64 = 7 * 24 * 4;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardTrendQuery {
|
||||
pub kind: Option<DashboardTrendKind>,
|
||||
pub(crate) struct DashboardTrendQuery {
|
||||
kind: Option<DashboardTrendKind>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum DashboardTrendKind {
|
||||
enum DashboardTrendKind {
|
||||
#[default]
|
||||
Usage,
|
||||
Latency,
|
||||
@@ -56,143 +54,143 @@ pub enum DashboardTrendKind {
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardSummaryData {
|
||||
pub cards: DashboardCardsData,
|
||||
pub trend: DashboardTrendData,
|
||||
pub health_timeline: DashboardHealthTimelineData,
|
||||
pub account_usage: Vec<DashboardAccountUsageData>,
|
||||
pub service_statuses: Vec<DashboardServiceStatusData>,
|
||||
pub event_logs: Vec<DashboardEventLogData>,
|
||||
pub pool_summary: AccountPoolDiagnostics,
|
||||
pub capacity_info: AccountCapacityDiagnostics,
|
||||
pub rotation_strategy: Option<String>,
|
||||
struct DashboardSummaryData {
|
||||
cards: DashboardCardsData,
|
||||
trend: DashboardTrendData,
|
||||
health_timeline: DashboardHealthTimelineData,
|
||||
account_usage: Vec<DashboardAccountUsageData>,
|
||||
service_statuses: Vec<DashboardServiceStatusData>,
|
||||
event_logs: Vec<DashboardEventLogData>,
|
||||
pool_summary: AccountPoolDiagnostics,
|
||||
capacity_info: AccountCapacityDiagnostics,
|
||||
rotation_strategy: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardHealthTimelineData {
|
||||
pub title: String,
|
||||
pub description: String,
|
||||
pub range_display: String,
|
||||
pub reliability_display: String,
|
||||
pub oldest_label: String,
|
||||
pub newest_label: String,
|
||||
pub points: String,
|
||||
struct DashboardHealthTimelineData {
|
||||
title: String,
|
||||
description: String,
|
||||
range_display: String,
|
||||
reliability_display: String,
|
||||
oldest_label: String,
|
||||
newest_label: String,
|
||||
points: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardCardsData {
|
||||
pub accounts: DashboardAccountsCardData,
|
||||
pub traffic: DashboardTrafficCardData,
|
||||
pub tokens: DashboardTokenCardData,
|
||||
pub cache: DashboardCacheCardData,
|
||||
struct DashboardCardsData {
|
||||
accounts: DashboardAccountsCardData,
|
||||
traffic: DashboardTrafficCardData,
|
||||
tokens: DashboardTokenCardData,
|
||||
cache: DashboardCacheCardData,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardAccountsCardData {
|
||||
pub total: u64,
|
||||
pub enabled: u64,
|
||||
pub abnormal: u64,
|
||||
struct DashboardAccountsCardData {
|
||||
total: u64,
|
||||
enabled: u64,
|
||||
abnormal: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardTrafficCardData {
|
||||
pub today_requests: u64,
|
||||
pub yesterday_requests: u64,
|
||||
pub total_requests: u64,
|
||||
pub rpm: u64,
|
||||
pub tpm: u64,
|
||||
struct DashboardTrafficCardData {
|
||||
today_requests: u64,
|
||||
yesterday_requests: u64,
|
||||
total_requests: u64,
|
||||
rpm: u64,
|
||||
tpm: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardTokenCardData {
|
||||
pub today_tokens: u64,
|
||||
pub yesterday_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub today_cost_usd: Option<f64>,
|
||||
pub total_cost_usd: Option<f64>,
|
||||
struct DashboardTokenCardData {
|
||||
today_tokens: u64,
|
||||
yesterday_tokens: u64,
|
||||
total_tokens: u64,
|
||||
today_cost_usd: Option<f64>,
|
||||
total_cost_usd: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardCacheCardData {
|
||||
pub today_hit_rate: Option<f64>,
|
||||
pub yesterday_hit_rate: Option<f64>,
|
||||
pub total_hit_rate: Option<f64>,
|
||||
pub total_cached_tokens: u64,
|
||||
pub first_token_latency_ms: Option<u64>,
|
||||
pub completion_latency_ms: Option<u64>,
|
||||
struct DashboardCacheCardData {
|
||||
today_hit_rate: Option<f64>,
|
||||
yesterday_hit_rate: Option<f64>,
|
||||
total_hit_rate: Option<f64>,
|
||||
total_cached_tokens: u64,
|
||||
first_token_latency_ms: Option<u64>,
|
||||
completion_latency_ms: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardTrendData {
|
||||
pub kind: DashboardTrendKind,
|
||||
pub points: Vec<DashboardTrendPointData>,
|
||||
pub summary: Vec<DashboardTrendSummaryData>,
|
||||
struct DashboardTrendData {
|
||||
kind: DashboardTrendKind,
|
||||
points: Vec<DashboardTrendPointData>,
|
||||
summary: Vec<DashboardTrendSummaryData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardTrendPointData {
|
||||
pub time: String,
|
||||
pub requests: u64,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cached_tokens: u64,
|
||||
pub tokens: u64,
|
||||
pub errors: u64,
|
||||
pub latency: u64,
|
||||
pub max_latency: u64,
|
||||
pub min_latency: u64,
|
||||
pub success_rate: f64,
|
||||
struct DashboardTrendPointData {
|
||||
time: String,
|
||||
requests: u64,
|
||||
input_tokens: u64,
|
||||
output_tokens: u64,
|
||||
cached_tokens: u64,
|
||||
tokens: u64,
|
||||
errors: u64,
|
||||
latency: u64,
|
||||
max_latency: u64,
|
||||
min_latency: u64,
|
||||
success_rate: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardTrendSummaryData {
|
||||
pub label: String,
|
||||
pub value: u64,
|
||||
pub ratio: Option<f64>,
|
||||
struct DashboardTrendSummaryData {
|
||||
label: String,
|
||||
value: u64,
|
||||
ratio: Option<f64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardAccountUsageData {
|
||||
pub name: String,
|
||||
pub email: String,
|
||||
pub plan: String,
|
||||
pub requests: u64,
|
||||
pub tokens: u64,
|
||||
pub quota_used_percent: Option<f64>,
|
||||
pub last_used: String,
|
||||
pub status: String,
|
||||
struct DashboardAccountUsageData {
|
||||
name: String,
|
||||
email: String,
|
||||
plan: String,
|
||||
requests: u64,
|
||||
tokens: u64,
|
||||
quota_used_percent: Option<f64>,
|
||||
last_used: String,
|
||||
status: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardEventLogData {
|
||||
pub id: String,
|
||||
pub time: String,
|
||||
pub level: EventLevel,
|
||||
pub request_id: Option<String>,
|
||||
pub route: Option<String>,
|
||||
pub model: Option<String>,
|
||||
pub status_code: Option<i64>,
|
||||
pub latency_ms: Option<i64>,
|
||||
struct DashboardEventLogData {
|
||||
id: String,
|
||||
time: String,
|
||||
level: EventLevel,
|
||||
request_id: Option<String>,
|
||||
route: Option<String>,
|
||||
model: Option<String>,
|
||||
status_code: Option<i64>,
|
||||
latency_ms: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DashboardServiceStatusData {
|
||||
pub label: String,
|
||||
pub value: String,
|
||||
pub detail: String,
|
||||
pub tone: String,
|
||||
struct DashboardServiceStatusData {
|
||||
label: String,
|
||||
value: String,
|
||||
detail: String,
|
||||
tone: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Default)]
|
||||
@@ -234,13 +232,11 @@ impl UsageWindow {
|
||||
}
|
||||
|
||||
/// `GET /api/admin/dashboard/summary`
|
||||
pub async fn dashboard_summary(
|
||||
pub(crate) async fn dashboard_summary(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
|
||||
let accounts = state
|
||||
.services
|
||||
@@ -280,44 +276,35 @@ pub async fn dashboard_summary(
|
||||
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(
|
||||
DashboardSummaryData {
|
||||
cards: dashboard_cards(&accounts, &summary, &time_buckets),
|
||||
trend,
|
||||
health_timeline: dashboard_health_timeline_data(&time_buckets),
|
||||
account_usage: account_usage_data(
|
||||
&accounts,
|
||||
&usage_records,
|
||||
"a_used_by_account,
|
||||
),
|
||||
service_statuses: service_status_data(&state),
|
||||
event_logs: dashboard_logs,
|
||||
pool_summary,
|
||||
capacity_info: AccountCapacityDiagnostics::from(capacity),
|
||||
rotation_strategy: Some(settings.auth.rotation_strategy.clone()),
|
||||
},
|
||||
request_id,
|
||||
),
|
||||
AdminEnvelope::ok(DashboardSummaryData {
|
||||
cards: dashboard_cards(&accounts, &summary, &time_buckets),
|
||||
trend,
|
||||
health_timeline: dashboard_health_timeline_data(&time_buckets),
|
||||
account_usage: account_usage_data(&accounts, &usage_records, "a_used_by_account),
|
||||
service_statuses: service_status_data(&state),
|
||||
event_logs: dashboard_logs,
|
||||
pool_summary,
|
||||
capacity_info: AccountCapacityDiagnostics::from(capacity),
|
||||
rotation_strategy: Some(settings.auth.rotation_strategy.clone()),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
/// `GET /api/admin/dashboard/trend?kind=usage|latency|errors`
|
||||
pub async fn dashboard_trend(
|
||||
pub(crate) async fn dashboard_trend(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<DashboardTrendQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
let time_buckets = dashboard_time_buckets(&state, Utc::now()).await;
|
||||
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(
|
||||
dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default()),
|
||||
request_id,
|
||||
),
|
||||
AdminEnvelope::ok(dashboard_trend_data(
|
||||
&time_buckets,
|
||||
query.kind.unwrap_or_default(),
|
||||
)),
|
||||
))
|
||||
}
|
||||
|
||||
|
||||
@@ -1,19 +1,14 @@
|
||||
//! 事件日志存储实现(SQLite)。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde_json::Value;
|
||||
use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::admin::monitoring::events::{
|
||||
EventLevel, EventLog, EventLogStore, EventLogStoreError, EventLogStoreResult,
|
||||
};
|
||||
use crate::admin::monitoring::events::{EventLevel, EventLog};
|
||||
use crate::infra::json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page};
|
||||
use crate::infra::time::china_quarter_hour_start;
|
||||
|
||||
/// 默认启用事件日志,便于管理端诊断请求链路。
|
||||
pub const DEFAULT_EVENT_LOG_ENABLED: bool = true;
|
||||
/// 默认保留最新事件日志数量。
|
||||
pub const DEFAULT_EVENT_LOG_CAPACITY: u32 = 2_000;
|
||||
/// 默认不保存请求/响应体,避免把用户内容写入诊断日志。
|
||||
@@ -175,14 +170,6 @@ impl SqliteEventLogStore {
|
||||
row.map(|row| event_from_row(&row)).transpose()
|
||||
}
|
||||
|
||||
/// 统计事件日志数量。
|
||||
pub async fn count(&self) -> SqliteEventLogStoreResult<u64> {
|
||||
let count: (i64,) = sqlx::query_as("select count(*) from event_logs")
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
Ok(count.0.max(0) as u64)
|
||||
}
|
||||
|
||||
/// 清空事件日志。
|
||||
pub async fn clear(&self) -> SqliteEventLogStoreResult<u64> {
|
||||
let result = sqlx::query("delete from event_logs")
|
||||
@@ -192,18 +179,6 @@ impl SqliteEventLogStore {
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl EventLogStore for SqliteEventLogStore {
|
||||
async fn append(&self, event: &EventLog) -> EventLogStoreResult<()> {
|
||||
append_event(&self.pool, event).await.map_err(|error| {
|
||||
EventLogStoreError::OperationFailed {
|
||||
message: error.to_string(),
|
||||
}
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct EventLogSummaryFields {
|
||||
transport: Option<String>,
|
||||
@@ -444,7 +419,7 @@ async fn count_events(
|
||||
let mut builder = QueryBuilder::<Sqlite>::new("select count(*) from event_logs");
|
||||
push_filter(&mut builder, filter, None)?;
|
||||
let (total,): (i64,) = builder.build_query_as().fetch_one(pool).await?;
|
||||
Ok(total.max(0) as u64)
|
||||
Ok(total.max(0).cast_unsigned())
|
||||
}
|
||||
|
||||
fn push_filter(
|
||||
@@ -631,30 +606,6 @@ impl From<AdminLogFilter> for EventLogFilter {
|
||||
}
|
||||
}
|
||||
|
||||
/// 日志状态。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AdminLogState {
|
||||
/// 是否启用。
|
||||
pub enabled: bool,
|
||||
/// 内存容量。
|
||||
pub capacity: u32,
|
||||
/// 是否捕获请求体。
|
||||
pub capture_body: bool,
|
||||
/// 已存储数量。
|
||||
pub stored_count: u64,
|
||||
}
|
||||
|
||||
/// 日志状态更新。
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct AdminLogStateUpdate {
|
||||
/// 是否启用。
|
||||
pub enabled: Option<bool>,
|
||||
/// 日志容量。
|
||||
pub capacity: Option<u32>,
|
||||
/// 是否捕获请求体。
|
||||
pub capture_body: Option<bool>,
|
||||
}
|
||||
|
||||
/// 清空日志结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AdminClearLogs {
|
||||
@@ -671,9 +622,6 @@ pub enum AdminLogError {
|
||||
/// 读取失败。
|
||||
#[error("failed to get event log")]
|
||||
Get,
|
||||
/// 计数失败。
|
||||
#[error("failed to count event logs")]
|
||||
Count,
|
||||
/// 清空失败。
|
||||
#[error("failed to clear event logs")]
|
||||
Clear,
|
||||
@@ -683,9 +631,6 @@ pub enum AdminLogError {
|
||||
/// 裁剪失败。
|
||||
#[error("failed to trim event logs")]
|
||||
Trim,
|
||||
/// 日志容量非法。
|
||||
#[error("log capacity must be greater than zero")]
|
||||
InvalidCapacity,
|
||||
}
|
||||
|
||||
/// 管理端日志服务。
|
||||
@@ -751,50 +696,6 @@ impl AdminLogService {
|
||||
self.store.get(id).await.map_err(|_| AdminLogError::Get)
|
||||
}
|
||||
|
||||
/// 读取日志状态。
|
||||
pub async fn state(&self) -> Result<AdminLogState, AdminLogError> {
|
||||
let settings = *self.settings.read().await;
|
||||
Ok(AdminLogState {
|
||||
enabled: settings.enabled,
|
||||
capacity: settings.capacity,
|
||||
capture_body: settings.capture_body,
|
||||
stored_count: self.store.count().await.map_err(|_| AdminLogError::Count)?,
|
||||
})
|
||||
}
|
||||
|
||||
/// 更新日志状态。
|
||||
pub async fn update_state(
|
||||
&self,
|
||||
update: AdminLogStateUpdate,
|
||||
) -> Result<AdminLogState, AdminLogError> {
|
||||
if matches!(update.capacity, Some(0)) {
|
||||
return Err(AdminLogError::InvalidCapacity);
|
||||
}
|
||||
|
||||
let trim_capacity = {
|
||||
let mut settings = self.settings.write().await;
|
||||
if let Some(enabled) = update.enabled {
|
||||
settings.enabled = enabled;
|
||||
}
|
||||
if let Some(capacity) = update.capacity {
|
||||
settings.capacity = capacity;
|
||||
}
|
||||
if let Some(capture_body) = update.capture_body {
|
||||
settings.capture_body = capture_body;
|
||||
}
|
||||
update.capacity
|
||||
};
|
||||
|
||||
if let Some(capacity) = trim_capacity {
|
||||
self.store
|
||||
.trim_to_capacity(capacity)
|
||||
.await
|
||||
.map_err(|_| AdminLogError::Trim)?;
|
||||
}
|
||||
|
||||
self.state().await
|
||||
}
|
||||
|
||||
/// 清空日志。
|
||||
pub async fn clear(&self) -> Result<AdminClearLogs, AdminLogError> {
|
||||
self.store
|
||||
@@ -807,8 +708,7 @@ impl AdminLogService {
|
||||
/// 记录事件日志。
|
||||
pub async fn record(&self, mut event: EventLog) -> Result<(), AdminLogError> {
|
||||
let settings = *self.settings.read().await;
|
||||
let policy = crate::admin::monitoring::events::EventLogService::new(settings.enabled);
|
||||
if !policy.should_record(&event) {
|
||||
if !settings.enabled && event.level != EventLevel::Error {
|
||||
return Ok(());
|
||||
}
|
||||
apply_capture_body_policy(&mut event, settings.capture_body);
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
//! 事件日志模型、端口与策略服务。
|
||||
//! 事件日志模型。
|
||||
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use thiserror::Error;
|
||||
use serde_json::{Map, Value};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// 事件等级。
|
||||
@@ -84,51 +82,12 @@ impl EventLog {
|
||||
upstream_request_id: None,
|
||||
latency_ms: None,
|
||||
message: message.into(),
|
||||
metadata: Value::Object(Default::default()),
|
||||
metadata: Value::Object(Map::default()),
|
||||
created_at: Utc::now(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 事件日志存储错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum EventLogStoreError {
|
||||
/// 底层存储失败。
|
||||
#[error("event log store operation failed: {message}")]
|
||||
OperationFailed {
|
||||
/// 错误说明。
|
||||
message: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// 事件日志存储结果。
|
||||
pub type EventLogStoreResult<T> = Result<T, EventLogStoreError>;
|
||||
|
||||
/// 事件日志存储端口。
|
||||
#[async_trait]
|
||||
pub trait EventLogStore: Send + Sync + 'static {
|
||||
/// 写入事件日志。
|
||||
async fn append(&self, event: &EventLog) -> EventLogStoreResult<()>;
|
||||
}
|
||||
|
||||
/// 事件日志服务。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EventLogService {
|
||||
enabled: bool,
|
||||
}
|
||||
|
||||
impl EventLogService {
|
||||
/// 构造事件日志服务。
|
||||
pub fn new(enabled: bool) -> Self {
|
||||
Self { enabled }
|
||||
}
|
||||
|
||||
/// 判断事件是否应该记录。
|
||||
pub fn should_record(&self, event: &EventLog) -> bool {
|
||||
self.enabled || event.level == EventLevel::Error
|
||||
}
|
||||
}
|
||||
|
||||
/// 响应事件记录(供 dispatch 模块使用)。
|
||||
pub struct ResponseEventRecord<'a> {
|
||||
pub logs: &'a crate::admin::monitoring::event_store::AdminLogService,
|
||||
|
||||
+51
-182
@@ -4,20 +4,18 @@ use axum::{
|
||||
extract::{Query, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
Extension, Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{
|
||||
admin::monitoring::{
|
||||
event_store::{AdminLogError, AdminLogFilter, AdminLogState, AdminLogStateUpdate},
|
||||
event_store::{AdminLogError, AdminLogFilter},
|
||||
events::{EventLevel, EventLog},
|
||||
},
|
||||
admin::{
|
||||
auth::session::require_admin_session,
|
||||
response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse},
|
||||
},
|
||||
http::middleware::request_id::RequestId,
|
||||
infra::{
|
||||
json::{clamp_limit, clamp_page, NumberedPage, Page},
|
||||
time::china_datetime,
|
||||
@@ -25,115 +23,61 @@ use crate::{
|
||||
runtime::state::AppState,
|
||||
};
|
||||
|
||||
/// 日志查询参数。
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct LogsQuery {
|
||||
/// 分页游标。
|
||||
pub cursor: Option<String>,
|
||||
/// 分页大小。
|
||||
pub limit: Option<u32>,
|
||||
/// 页码。
|
||||
pub page: Option<u32>,
|
||||
/// 每页条目数。
|
||||
pub page_size: Option<u32>,
|
||||
/// 事件类别。
|
||||
pub kind: Option<String>,
|
||||
/// 事件等级。
|
||||
pub level: Option<String>,
|
||||
/// 请求 ID。
|
||||
pub request_id: Option<String>,
|
||||
/// 账号 ID。
|
||||
pub account_id: Option<String>,
|
||||
/// 路由。
|
||||
pub route: Option<String>,
|
||||
/// 模型。
|
||||
pub model: Option<String>,
|
||||
/// HTTP 状态码。
|
||||
pub status_code: Option<i64>,
|
||||
/// 上游传输方式。
|
||||
pub transport: Option<String>,
|
||||
/// 同一请求内的上游尝试序号。
|
||||
pub attempt_index: Option<i64>,
|
||||
/// 上游 HTTP 状态码。
|
||||
pub upstream_status_code: Option<i64>,
|
||||
/// 失败分类。
|
||||
pub failure_class: Option<String>,
|
||||
/// 上游响应 ID。
|
||||
pub response_id: Option<String>,
|
||||
/// 上游请求 ID。
|
||||
pub upstream_request_id: Option<String>,
|
||||
/// 搜索关键词。
|
||||
pub search: Option<String>,
|
||||
pub(crate) struct LogsQuery {
|
||||
cursor: Option<String>,
|
||||
limit: Option<u32>,
|
||||
page: Option<u32>,
|
||||
page_size: Option<u32>,
|
||||
kind: Option<String>,
|
||||
level: Option<String>,
|
||||
request_id: Option<String>,
|
||||
account_id: Option<String>,
|
||||
route: Option<String>,
|
||||
model: Option<String>,
|
||||
status_code: Option<i64>,
|
||||
transport: Option<String>,
|
||||
attempt_index: Option<i64>,
|
||||
upstream_status_code: Option<i64>,
|
||||
failure_class: Option<String>,
|
||||
response_id: Option<String>,
|
||||
upstream_request_id: Option<String>,
|
||||
search: Option<String>,
|
||||
}
|
||||
|
||||
/// 日志详情查询参数。
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct LogDetailQuery {
|
||||
/// 日志事件 ID。
|
||||
pub id: String,
|
||||
pub(crate) struct LogDetailQuery {
|
||||
id: String,
|
||||
}
|
||||
|
||||
/// 管理端日志状态响应。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct LogStateData {
|
||||
/// 是否启用。
|
||||
pub enabled: bool,
|
||||
/// 内存容量。
|
||||
pub capacity: u32,
|
||||
/// 是否捕获请求体。
|
||||
pub capture_body: bool,
|
||||
/// 已存储数量。
|
||||
pub stored_count: u64,
|
||||
struct ClearLogsData {
|
||||
cleared: u64,
|
||||
}
|
||||
|
||||
/// 清空日志响应。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ClearLogsData {
|
||||
/// 清理数量。
|
||||
pub cleared: u64,
|
||||
}
|
||||
|
||||
/// 管理端日志响应。
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct LogData {
|
||||
/// 原始日志字段。
|
||||
struct LogData {
|
||||
#[serde(flatten)]
|
||||
pub log: EventLog,
|
||||
/// 中国时区展示时间。
|
||||
pub created_at_display: String,
|
||||
}
|
||||
|
||||
/// 更新日志状态请求。
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UpdateLogStateRequest {
|
||||
/// 是否启用。
|
||||
pub enabled: Option<bool>,
|
||||
/// 日志容量。
|
||||
pub capacity: Option<u32>,
|
||||
/// 是否捕获请求体。
|
||||
pub capture_body: Option<bool>,
|
||||
log: EventLog,
|
||||
created_at_display: String,
|
||||
}
|
||||
|
||||
/// `GET /api/admin/logs`
|
||||
pub async fn logs(
|
||||
pub(crate) async fn logs(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<LogsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
let limit = clamp_limit(query.page_size.or(query.limit).unwrap_or(50));
|
||||
let page = query.page;
|
||||
let use_numbered_page = page.is_some() || query.page_size.is_some();
|
||||
let cursor = query.cursor.clone();
|
||||
let filter = filter_from_query(query, &request_id)?;
|
||||
let filter = filter_from_query(query)?;
|
||||
|
||||
if use_numbered_page {
|
||||
return match state
|
||||
@@ -151,10 +95,10 @@ pub async fn logs(
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminPageEnvelope::numbered(page, request_id),
|
||||
AdminPageEnvelope::numbered(page),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(log_error(error, request_id)),
|
||||
Err(error) => Err(log_error(&error)),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -166,100 +110,56 @@ pub async fn logs(
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminPageEnvelope::ok(page, limit, request_id),
|
||||
AdminPageEnvelope::ok(page, limit),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(log_error(error, request_id)),
|
||||
Err(error) => Err(log_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `GET /api/admin/logs/detail`
|
||||
pub async fn log_detail(
|
||||
pub(crate) async fn log_detail(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<LogDetailQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
match state.services.logs.get(&query.id).await {
|
||||
Ok(Some(log)) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(LogData::from(log), request_id),
|
||||
AdminEnvelope::ok(LogData::from(log)),
|
||||
)),
|
||||
Ok(None) => Err(AdminError::new(
|
||||
StatusCode::NOT_FOUND,
|
||||
40401,
|
||||
"Log event not found",
|
||||
request_id,
|
||||
)),
|
||||
Err(error) => Err(log_error(error, request_id)),
|
||||
Err(error) => Err(log_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /api/admin/logs/delete`
|
||||
pub async fn clear_logs(
|
||||
pub(crate) async fn clear_logs(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
match state.services.logs.clear().await {
|
||||
Ok(cleared) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(
|
||||
ClearLogsData {
|
||||
cleared: cleared.cleared,
|
||||
},
|
||||
request_id,
|
||||
),
|
||||
AdminEnvelope::ok(ClearLogsData {
|
||||
cleared: cleared.cleared,
|
||||
}),
|
||||
)),
|
||||
Err(error) => Err(log_error(error, request_id)),
|
||||
Err(error) => Err(log_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `GET /api/admin/logs/state`
|
||||
pub async fn logs_state(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
match state.services.logs.state().await {
|
||||
Ok(state) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(LogStateData::from(state), request_id),
|
||||
)),
|
||||
Err(error) => Err(log_error(error, request_id)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /api/admin/logs/state`
|
||||
pub async fn update_logs_state(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Json(payload): Json<UpdateLogStateRequest>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
match state.services.logs.update_state(payload.into()).await {
|
||||
Ok(state) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(LogStateData::from(state), request_id),
|
||||
)),
|
||||
Err(error) => Err(log_error(error, request_id)),
|
||||
}
|
||||
}
|
||||
|
||||
fn filter_from_query(query: LogsQuery, request_id: &str) -> Result<AdminLogFilter, AdminError> {
|
||||
fn filter_from_query(query: LogsQuery) -> Result<AdminLogFilter, AdminError> {
|
||||
Ok(AdminLogFilter {
|
||||
kind: non_empty(query.kind),
|
||||
level: level_from_query(query.level).map_err(|message| {
|
||||
AdminError::new(StatusCode::BAD_REQUEST, 40001, message, request_id)
|
||||
})?,
|
||||
level: level_from_query(query.level)
|
||||
.map_err(|message| AdminError::new(StatusCode::BAD_REQUEST, 40001, message))?,
|
||||
request_id: non_empty(query.request_id),
|
||||
account_id: non_empty(query.account_id),
|
||||
route: non_empty(query.route),
|
||||
@@ -275,25 +175,15 @@ fn filter_from_query(query: LogsQuery, request_id: &str) -> Result<AdminLogFilte
|
||||
})
|
||||
}
|
||||
|
||||
fn log_error(error: AdminLogError, request_id: String) -> AdminError {
|
||||
fn log_error(error: &AdminLogError) -> AdminError {
|
||||
match error {
|
||||
AdminLogError::List
|
||||
| AdminLogError::Get
|
||||
| AdminLogError::Count
|
||||
| AdminLogError::Clear
|
||||
| AdminLogError::Append
|
||||
| AdminLogError::Trim => AdminError::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
error.to_string(),
|
||||
request_id,
|
||||
),
|
||||
AdminLogError::InvalidCapacity => AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
error.to_string(),
|
||||
request_id,
|
||||
),
|
||||
| AdminLogError::Trim => {
|
||||
AdminError::new(StatusCode::INTERNAL_SERVER_ERROR, 50001, error.to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -324,24 +214,3 @@ impl From<EventLog> for LogData {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AdminLogState> for LogStateData {
|
||||
fn from(state: AdminLogState) -> Self {
|
||||
Self {
|
||||
enabled: state.enabled,
|
||||
capacity: state.capacity,
|
||||
capture_body: state.capture_body,
|
||||
stored_count: state.stored_count,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<UpdateLogStateRequest> for AdminLogStateUpdate {
|
||||
fn from(request: UpdateLogStateRequest) -> Self {
|
||||
Self {
|
||||
enabled: request.enabled,
|
||||
capacity: request.capacity,
|
||||
capture_body: request.capture_body,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,5 +7,4 @@ pub mod event_store;
|
||||
pub mod events;
|
||||
pub mod logs;
|
||||
pub mod service;
|
||||
pub mod usage;
|
||||
pub mod usage_store;
|
||||
|
||||
@@ -1,103 +1,11 @@
|
||||
//! 用量聚合模型、端口与策略服务。
|
||||
use async_trait::async_trait;
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::admin::monitoring::usage_store::{
|
||||
SqliteUsageStore, UsageListRecord, UsageSummary, UsageTimeBucketRecord,
|
||||
};
|
||||
use crate::infra::json::{NumberedPage, Page};
|
||||
use crate::upstream::accounts::model::AccountUsageDelta;
|
||||
use crate::upstream::protocol::events::TokenUsage;
|
||||
|
||||
/// API 响应用量记录。
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct UsageRecord {
|
||||
pub request_id: String,
|
||||
pub account_id: String,
|
||||
pub model: String,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cached_tokens: u64,
|
||||
pub reasoning_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
}
|
||||
|
||||
/// 用量聚合窗口。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsageWindow {
|
||||
/// 窗口开始时间。
|
||||
pub started_at: DateTime<Utc>,
|
||||
/// 请求数。
|
||||
pub request_count: u64,
|
||||
/// 输入 token 数。
|
||||
pub input_tokens: u64,
|
||||
/// 输出 token 数。
|
||||
pub output_tokens: u64,
|
||||
}
|
||||
|
||||
/// 用量快照。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsageSnapshot {
|
||||
/// 账号 ID。
|
||||
pub account_id: String,
|
||||
/// 当前窗口。
|
||||
pub window: UsageWindow,
|
||||
}
|
||||
|
||||
/// 用量存储错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum UsageStoreError {
|
||||
/// 底层存储失败。
|
||||
#[error("usage store operation failed: {message}")]
|
||||
OperationFailed {
|
||||
/// 错误说明。
|
||||
message: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// 用量存储结果。
|
||||
pub type UsageStoreResult<T> = Result<T, UsageStoreError>;
|
||||
|
||||
/// 用量存储端口。
|
||||
#[async_trait]
|
||||
pub trait UsageStore: Send + Sync + 'static {
|
||||
/// 写入用量快照。
|
||||
async fn record_snapshot(&self, snapshot: &UsageSnapshot) -> UsageStoreResult<()>;
|
||||
}
|
||||
|
||||
/// 用量聚合服务。
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct UsageService;
|
||||
|
||||
impl UsageService {
|
||||
/// 将 token 增量累加到窗口。
|
||||
pub fn add_tokens(window: &mut UsageWindow, input_tokens: u64, output_tokens: u64) {
|
||||
window.request_count += 1;
|
||||
window.input_tokens += input_tokens;
|
||||
window.output_tokens += output_tokens;
|
||||
}
|
||||
|
||||
/// 将标准化 token 用量转换为账号持久化用量增量。
|
||||
pub fn account_delta_from_token_usage(usage: TokenUsage) -> AccountUsageDelta {
|
||||
AccountUsageDelta {
|
||||
requests: 0,
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
cached_tokens: usage.cached_tokens,
|
||||
reasoning_tokens: 0,
|
||||
total_tokens: usage.input_tokens + usage.output_tokens + usage.cached_tokens,
|
||||
empty_responses: 0,
|
||||
image_input_tokens: 0,
|
||||
image_output_tokens: 0,
|
||||
image_requests: 0,
|
||||
image_request_failures: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
use crate::infra::json::Page;
|
||||
|
||||
/// 管理端用量统计服务。
|
||||
#[derive(Clone)]
|
||||
@@ -128,24 +36,6 @@ impl AdminUsageService {
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn list_page(
|
||||
&self,
|
||||
page: u32,
|
||||
page_size: u32,
|
||||
) -> Result<NumberedPage<AdminUsageRecord>, AdminUsageError> {
|
||||
let page = self
|
||||
.store
|
||||
.list_usage_page(page, page_size)
|
||||
.await
|
||||
.map_err(|_| AdminUsageError::List)?;
|
||||
Ok(NumberedPage {
|
||||
items: page.items.into_iter().map(AdminUsageRecord::from).collect(),
|
||||
total: page.total,
|
||||
page: page.page,
|
||||
page_size: page.page_size,
|
||||
})
|
||||
}
|
||||
|
||||
/// 汇总账号用量。
|
||||
pub async fn summary(&self) -> Result<AdminUsageSummary, AdminUsageError> {
|
||||
self.store
|
||||
|
||||
@@ -1,201 +0,0 @@
|
||||
//! 管理端用量处理器。
|
||||
|
||||
use axum::{
|
||||
extract::{Query, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
Extension,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{
|
||||
admin::auth::session::require_admin_session,
|
||||
admin::monitoring::service::{AdminUsageError, AdminUsageRecord, AdminUsageSummary},
|
||||
admin::response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse},
|
||||
http::middleware::request_id::RequestId,
|
||||
infra::{
|
||||
json::{clamp_limit, clamp_page, Page},
|
||||
time::{china_relative_time, china_rfc3339},
|
||||
},
|
||||
runtime::state::AppState,
|
||||
};
|
||||
|
||||
/// 用量统计查询参数。
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsageStatsQuery {
|
||||
pub cursor: Option<String>,
|
||||
pub limit: Option<u32>,
|
||||
pub page: Option<u32>,
|
||||
pub page_size: Option<u32>,
|
||||
}
|
||||
|
||||
/// 管理端账号用量响应。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminUsageStatsData {
|
||||
pub account_id: String,
|
||||
pub email: Option<String>,
|
||||
pub label: Option<String>,
|
||||
pub plan_type: Option<String>,
|
||||
pub request_count: i64,
|
||||
pub empty_response_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
pub reasoning_tokens: i64,
|
||||
pub total_tokens: i64,
|
||||
pub image_input_tokens: i64,
|
||||
pub image_output_tokens: i64,
|
||||
pub image_request_count: i64,
|
||||
pub image_request_failed_count: i64,
|
||||
pub last_used_at: Option<String>,
|
||||
pub last_used_at_display: String,
|
||||
}
|
||||
|
||||
/// 管理端账号用量汇总响应。
|
||||
#[derive(Debug, Clone, Copy, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminUsageStatsSummaryData {
|
||||
pub account_count: i64,
|
||||
pub request_count: i64,
|
||||
pub empty_response_count: i64,
|
||||
pub input_tokens: i64,
|
||||
pub output_tokens: i64,
|
||||
pub cached_tokens: i64,
|
||||
pub reasoning_tokens: i64,
|
||||
pub total_tokens: i64,
|
||||
pub image_input_tokens: i64,
|
||||
pub image_output_tokens: i64,
|
||||
pub image_request_count: i64,
|
||||
pub image_request_failed_count: i64,
|
||||
}
|
||||
|
||||
/// `GET /api/admin/usage`
|
||||
pub async fn usage_stats(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<UsageStatsQuery>,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
|
||||
let limit = clamp_limit(query.page_size.or(query.limit).unwrap_or(50));
|
||||
let use_numbered_page = query.page.is_some() || query.page_size.is_some();
|
||||
|
||||
if use_numbered_page {
|
||||
return match state
|
||||
.services
|
||||
.usage
|
||||
.list_page(clamp_page(query.page.unwrap_or(1)), limit)
|
||||
.await
|
||||
{
|
||||
Ok(page) => {
|
||||
let page = crate::infra::json::NumberedPage {
|
||||
items: page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(AdminUsageStatsData::from)
|
||||
.collect(),
|
||||
total: page.total,
|
||||
page: page.page,
|
||||
page_size: page.page_size,
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminPageEnvelope::numbered(page, request_id),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(usage_error(error, request_id)),
|
||||
};
|
||||
}
|
||||
|
||||
match state.services.usage.list(query.cursor, limit).await {
|
||||
Ok(page) => {
|
||||
let page = Page {
|
||||
items: page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(AdminUsageStatsData::from)
|
||||
.collect(),
|
||||
next_cursor: page.next_cursor,
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminPageEnvelope::ok(page, limit, request_id),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(usage_error(error, request_id)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `GET /api/admin/usage/summary`
|
||||
pub async fn usage_stats_summary(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
|
||||
match state.services.usage.summary().await {
|
||||
Ok(summary) => Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminUsageStatsSummaryData::from(summary), request_id),
|
||||
)),
|
||||
Err(error) => Err(usage_error(error, request_id)),
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_error(error: AdminUsageError, request_id: String) -> AdminError {
|
||||
AdminError::new(
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50001,
|
||||
error.to_string(),
|
||||
request_id,
|
||||
)
|
||||
}
|
||||
|
||||
impl From<AdminUsageRecord> for AdminUsageStatsData {
|
||||
fn from(usage: AdminUsageRecord) -> Self {
|
||||
Self {
|
||||
account_id: usage.account_id,
|
||||
email: usage.email,
|
||||
label: usage.label,
|
||||
plan_type: usage.plan_type,
|
||||
request_count: usage.request_count,
|
||||
empty_response_count: usage.empty_response_count,
|
||||
input_tokens: usage.input_tokens,
|
||||
output_tokens: usage.output_tokens,
|
||||
cached_tokens: usage.cached_tokens,
|
||||
reasoning_tokens: usage.reasoning_tokens,
|
||||
total_tokens: usage.total_tokens,
|
||||
image_input_tokens: usage.image_input_tokens,
|
||||
image_output_tokens: usage.image_output_tokens,
|
||||
image_request_count: usage.image_request_count,
|
||||
image_request_failed_count: usage.image_request_failed_count,
|
||||
last_used_at_display: china_relative_time(usage.last_used_at, chrono::Utc::now()),
|
||||
last_used_at: usage.last_used_at.map(|value| china_rfc3339(&value)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<AdminUsageSummary> for AdminUsageStatsSummaryData {
|
||||
fn from(summary: AdminUsageSummary) -> Self {
|
||||
Self {
|
||||
account_count: summary.account_count,
|
||||
request_count: summary.request_count,
|
||||
empty_response_count: summary.empty_response_count,
|
||||
input_tokens: summary.input_tokens,
|
||||
output_tokens: summary.output_tokens,
|
||||
cached_tokens: summary.cached_tokens,
|
||||
reasoning_tokens: summary.reasoning_tokens,
|
||||
total_tokens: summary.total_tokens,
|
||||
image_input_tokens: summary.image_input_tokens,
|
||||
image_output_tokens: summary.image_output_tokens,
|
||||
image_request_count: summary.image_request_count,
|
||||
image_request_failed_count: summary.image_request_failed_count,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,7 @@ use chrono::{DateTime, Utc};
|
||||
use sqlx::{sqlite::SqliteRow, Row, SqlitePool};
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::infra::json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page};
|
||||
use crate::infra::json::{decode_cursor, encode_cursor, Page};
|
||||
|
||||
const LIST_USAGE_AFTER_CURSOR_SQL: &str = r"
|
||||
select
|
||||
@@ -54,31 +54,6 @@ left join accounts a on a.id = au.account_id
|
||||
order by au.last_used_at desc, au.account_id desc
|
||||
limit ?";
|
||||
|
||||
const LIST_USAGE_PAGE_SQL: &str = r"
|
||||
select
|
||||
au.account_id,
|
||||
a.email,
|
||||
a.label,
|
||||
a.plan_type,
|
||||
au.request_count,
|
||||
au.empty_response_count,
|
||||
au.input_tokens,
|
||||
au.output_tokens,
|
||||
au.cached_tokens,
|
||||
au.reasoning_tokens,
|
||||
au.total_tokens,
|
||||
au.image_input_tokens,
|
||||
au.image_output_tokens,
|
||||
au.image_request_count,
|
||||
au.image_request_failed_count,
|
||||
au.last_used_at
|
||||
from account_usage au
|
||||
left join accounts a on a.id = au.account_id
|
||||
order by au.last_used_at desc, au.account_id desc
|
||||
limit ? offset ?";
|
||||
|
||||
const COUNT_USAGE_SQL: &str = "select count(*) from account_usage";
|
||||
|
||||
const USAGE_SUMMARY_SQL: &str = r"
|
||||
select
|
||||
count(*) as account_count,
|
||||
@@ -244,11 +219,6 @@ impl SqliteUsageStore {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// 返回底层连接池。
|
||||
pub fn pool(&self) -> &SqlitePool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
/// 分页列出账号用量。
|
||||
pub async fn list_usage(
|
||||
&self,
|
||||
@@ -266,45 +236,16 @@ impl SqliteUsageStore {
|
||||
.bind(limit + 1)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
Ok(to_page(rows, limit))
|
||||
Ok(to_page(&rows, limit))
|
||||
} else {
|
||||
let rows = sqlx::query(LIST_USAGE_SQL)
|
||||
.bind(limit + 1)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
Ok(to_page(rows, limit))
|
||||
Ok(to_page(&rows, limit))
|
||||
}
|
||||
}
|
||||
|
||||
/// 按页码列出账号用量。
|
||||
pub async fn list_usage_page(
|
||||
&self,
|
||||
page: u32,
|
||||
page_size: u32,
|
||||
) -> SqliteUsageStoreResult<NumberedPage<UsageListRecord>> {
|
||||
let page_size = page_size.clamp(1, 200);
|
||||
let offset = page_offset(page, page_size);
|
||||
let rows = sqlx::query(LIST_USAGE_PAGE_SQL)
|
||||
.bind(i64::from(page_size))
|
||||
.bind(offset.min(i64::MAX as u64) as i64)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
let items = rows
|
||||
.iter()
|
||||
.map(usage_list_from_row)
|
||||
.collect::<SqliteUsageStoreResult<Vec<_>>>()?;
|
||||
let (total,): (i64,) = sqlx::query_as(COUNT_USAGE_SQL)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(NumberedPage {
|
||||
items,
|
||||
total: total.max(0) as u64,
|
||||
page: page.max(1),
|
||||
page_size,
|
||||
})
|
||||
}
|
||||
|
||||
/// 汇总账号用量。
|
||||
pub async fn usage_summary(&self) -> SqliteUsageStoreResult<UsageSummary> {
|
||||
let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?;
|
||||
@@ -356,7 +297,9 @@ fn usage_list_from_row(row: &SqliteRow) -> SqliteUsageStoreResult<UsageListRecor
|
||||
image_output_tokens: row.get("image_output_tokens"),
|
||||
image_request_count: row.get("image_request_count"),
|
||||
image_request_failed_count: row.get("image_request_failed_count"),
|
||||
last_used_at: parse_optional_rfc3339(row.get::<Option<String>, _>("last_used_at"))?,
|
||||
last_used_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("last_used_at").as_deref(),
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -381,14 +324,13 @@ fn usage_time_bucket_from_row(row: &SqliteRow) -> SqliteUsageStoreResult<UsageTi
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_optional_rfc3339(value: Option<String>) -> SqliteUsageStoreResult<Option<DateTime<Utc>>> {
|
||||
fn parse_optional_rfc3339(value: Option<&str>) -> SqliteUsageStoreResult<Option<DateTime<Utc>>> {
|
||||
value
|
||||
.as_deref()
|
||||
.map(|value| Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc)))
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn to_page(rows: Vec<SqliteRow>, limit: u32) -> Page<UsageListRecord> {
|
||||
fn to_page(rows: &[SqliteRow], limit: u32) -> Page<UsageListRecord> {
|
||||
let has_more = rows.len() > limit as usize;
|
||||
let mut items = Vec::with_capacity(limit as usize);
|
||||
let mut last_row: Option<&SqliteRow> = None;
|
||||
|
||||
+29
-39
@@ -10,18 +10,13 @@ use serde::Serialize;
|
||||
|
||||
/// 管理端错误。
|
||||
pub struct AdminError {
|
||||
pub status: StatusCode,
|
||||
pub code: u32,
|
||||
pub message: String,
|
||||
status: StatusCode,
|
||||
code: u32,
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl AdminError {
|
||||
pub fn new(
|
||||
status: StatusCode,
|
||||
code: u32,
|
||||
message: impl Into<String>,
|
||||
_request_id: impl Into<String>,
|
||||
) -> Self {
|
||||
pub fn new(status: StatusCode, code: u32, message: impl Into<String>) -> Self {
|
||||
Self {
|
||||
status,
|
||||
code,
|
||||
@@ -48,52 +43,47 @@ impl IntoResponse for AdminError {
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminEnvelope<T> {
|
||||
pub code: u32,
|
||||
pub message: String,
|
||||
pub data: T,
|
||||
code: u32,
|
||||
message: String,
|
||||
data: T,
|
||||
}
|
||||
|
||||
impl<T> AdminEnvelope<T> {
|
||||
pub fn new(
|
||||
code: u32,
|
||||
message: impl Into<String>,
|
||||
data: T,
|
||||
_request_id: impl Into<String>,
|
||||
) -> Self {
|
||||
pub fn new(code: u32, message: impl Into<String>, data: T) -> Self {
|
||||
Self {
|
||||
code,
|
||||
message: message.into(),
|
||||
data,
|
||||
}
|
||||
}
|
||||
pub fn ok(data: T, request_id: impl Into<String>) -> Self {
|
||||
Self::new(200, "OK", data, request_id)
|
||||
pub fn ok(data: T) -> Self {
|
||||
Self::new(200, "OK", data)
|
||||
}
|
||||
}
|
||||
|
||||
/// 分页元数据。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct CursorPageMeta {
|
||||
pub limit: u32,
|
||||
pub(crate) struct CursorPageMeta {
|
||||
pub(crate) limit: u32,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub next_cursor: Option<String>,
|
||||
pub(crate) next_cursor: Option<String>,
|
||||
}
|
||||
|
||||
/// 页码分页元数据。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct NumberedPageMeta {
|
||||
pub page: u32,
|
||||
pub page_size: u32,
|
||||
pub total: u64,
|
||||
pub total_pages: u32,
|
||||
pub(crate) struct NumberedPageMeta {
|
||||
pub(crate) page: u32,
|
||||
pub(crate) page_size: u32,
|
||||
pub(crate) total: u64,
|
||||
pub(crate) total_pages: u32,
|
||||
}
|
||||
|
||||
/// 分页元数据。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum PageMeta {
|
||||
pub(crate) enum PageMeta {
|
||||
Cursor(CursorPageMeta),
|
||||
Numbered(NumberedPageMeta),
|
||||
}
|
||||
@@ -102,21 +92,21 @@ pub enum PageMeta {
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminPageEnvelope<T> {
|
||||
pub code: u32,
|
||||
pub message: String,
|
||||
pub data: PageData<T>,
|
||||
code: u32,
|
||||
message: String,
|
||||
data: PageData<T>,
|
||||
}
|
||||
|
||||
/// 分页响应数据。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PageData<T> {
|
||||
pub items: Vec<T>,
|
||||
pub page: PageMeta,
|
||||
struct PageData<T> {
|
||||
items: Vec<T>,
|
||||
page: PageMeta,
|
||||
}
|
||||
|
||||
impl<T> AdminPageEnvelope<T> {
|
||||
pub fn ok(page: Page<T>, limit: u32, _request_id: impl Into<String>) -> Self {
|
||||
pub fn ok(page: Page<T>, limit: u32) -> Self {
|
||||
let Page { items, next_cursor } = page;
|
||||
Self {
|
||||
code: 200,
|
||||
@@ -128,7 +118,7 @@ impl<T> AdminPageEnvelope<T> {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn numbered(page: NumberedPage<T>, _request_id: impl Into<String>) -> Self {
|
||||
pub fn numbered(page: NumberedPage<T>) -> Self {
|
||||
let NumberedPage {
|
||||
items,
|
||||
total,
|
||||
@@ -153,8 +143,8 @@ impl<T> AdminPageEnvelope<T> {
|
||||
|
||||
/// 管理端响应。
|
||||
pub struct AdminResponse<T: Serialize> {
|
||||
pub status: StatusCode,
|
||||
pub body: T,
|
||||
status: StatusCode,
|
||||
body: T,
|
||||
}
|
||||
|
||||
impl<T: Serialize> AdminResponse<T> {
|
||||
|
||||
+6
-30
@@ -9,20 +9,15 @@ use crate::runtime::state::AppState;
|
||||
|
||||
use super::{
|
||||
accounts::routes::{
|
||||
account_quota, account_test_models, accounts, batch_delete_accounts, create_account,
|
||||
export_accounts, get_account_cookies, health_check_accounts, import_accounts,
|
||||
oauth_authorize_account, oauth_exchange_account, quota_warnings, refresh_account,
|
||||
reset_account_usage, set_account_cookies, test_account_connection, update_account,
|
||||
account_models, account_quota, accounts, batch_delete_accounts, create_account,
|
||||
import_accounts, oauth_authorize_account, oauth_exchange_account, refresh_account,
|
||||
test_account_connection, update_account,
|
||||
},
|
||||
auth::session::{login, logout, session_status},
|
||||
keys::routes::{
|
||||
api_keys, batch_delete_api_keys, create_api_key, export_api_keys, update_api_key,
|
||||
},
|
||||
models::routes::refresh_models,
|
||||
keys::routes::{api_keys, batch_delete_api_keys, create_api_key, update_api_key},
|
||||
monitoring::{
|
||||
dashboard::{dashboard_summary, dashboard_trend},
|
||||
logs::{clear_logs, log_detail, logs, logs_state, update_logs_state},
|
||||
usage::{usage_stats, usage_stats_summary},
|
||||
logs::{clear_logs, log_detail, logs},
|
||||
},
|
||||
settings::routes::{settings, update_settings},
|
||||
};
|
||||
@@ -36,11 +31,7 @@ pub fn router() -> Router<AppState> {
|
||||
.route("/api/admin/settings", get(settings).post(update_settings))
|
||||
.route("/api/admin/dashboard/summary", get(dashboard_summary))
|
||||
.route("/api/admin/dashboard/trend", get(dashboard_trend))
|
||||
.route("/api/admin/models/refresh", post(refresh_models))
|
||||
.route("/api/admin/usage", get(usage_stats))
|
||||
.route("/api/admin/usage/summary", get(usage_stats_summary))
|
||||
.route("/api/admin/accounts", get(accounts).post(create_account))
|
||||
.route("/api/admin/accounts/export", get(export_accounts))
|
||||
.route("/api/admin/accounts/import", post(import_accounts))
|
||||
.route(
|
||||
"/api/admin/accounts/oauth/authorize",
|
||||
@@ -50,31 +41,16 @@ pub fn router() -> Router<AppState> {
|
||||
"/api/admin/accounts/oauth/exchange",
|
||||
post(oauth_exchange_account),
|
||||
)
|
||||
.route("/api/admin/accounts/quota-warnings", get(quota_warnings))
|
||||
.route(
|
||||
"/api/admin/accounts/health-check",
|
||||
post(health_check_accounts),
|
||||
)
|
||||
.route("/api/admin/accounts/test", post(test_account_connection))
|
||||
.route("/api/admin/accounts/models", get(account_test_models))
|
||||
.route("/api/admin/accounts/models", get(account_models))
|
||||
.route("/api/admin/accounts/delete", post(batch_delete_accounts))
|
||||
.route("/api/admin/accounts/update", post(update_account))
|
||||
.route("/api/admin/accounts/refresh", post(refresh_account))
|
||||
.route("/api/admin/accounts/reset-usage", post(reset_account_usage))
|
||||
.route(
|
||||
"/api/admin/accounts/cookies",
|
||||
get(get_account_cookies).post(set_account_cookies),
|
||||
)
|
||||
.route("/api/admin/accounts/quota", get(account_quota))
|
||||
.route("/api/admin/logs", get(logs))
|
||||
.route("/api/admin/logs/delete", post(clear_logs))
|
||||
.route(
|
||||
"/api/admin/logs/state",
|
||||
get(logs_state).post(update_logs_state),
|
||||
)
|
||||
.route("/api/admin/logs/detail", get(log_detail))
|
||||
.route("/api/admin/keys", get(api_keys).post(create_api_key))
|
||||
.route("/api/admin/keys/export", get(export_api_keys))
|
||||
.route("/api/admin/keys/delete", post(batch_delete_api_keys))
|
||||
.route("/api/admin/keys/update", post(update_api_key))
|
||||
}
|
||||
|
||||
@@ -5,7 +5,6 @@ use axum::{
|
||||
extract::State,
|
||||
http::{HeaderMap, StatusCode},
|
||||
response::IntoResponse,
|
||||
Extension,
|
||||
};
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
@@ -18,7 +17,6 @@ use crate::{
|
||||
settings::{AdminSettingsPatch, RuntimeSettingsError, SettingsServiceError},
|
||||
types::AppConfig,
|
||||
},
|
||||
http::middleware::request_id::RequestId,
|
||||
runtime::state::AppState,
|
||||
};
|
||||
|
||||
@@ -32,17 +30,16 @@ const ALLOWED_SETTINGS_KEYS: [&str; 7] = [
|
||||
"rotationStrategy",
|
||||
];
|
||||
|
||||
/// 管理端设置响应。
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AdminSettingsData {
|
||||
pub model_aliases: BTreeMap<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,
|
||||
struct AdminSettingsData {
|
||||
model_aliases: BTreeMap<String, String>,
|
||||
model_account_routes: BTreeMap<String, Vec<String>>,
|
||||
refresh_margin_seconds: u64,
|
||||
refresh_concurrency: u32,
|
||||
max_concurrent_per_account: usize,
|
||||
request_interval_ms: u64,
|
||||
rotation_strategy: String,
|
||||
}
|
||||
|
||||
impl AdminSettingsData {
|
||||
@@ -60,59 +57,48 @@ impl AdminSettingsData {
|
||||
}
|
||||
|
||||
/// `GET /api/admin/settings`
|
||||
pub async fn settings(
|
||||
pub(crate) async fn settings(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
let config = state.services.settings.current();
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminSettingsData::from_config(&config), request_id),
|
||||
AdminEnvelope::ok(AdminSettingsData::from_config(&config)),
|
||||
))
|
||||
}
|
||||
|
||||
/// `POST /api/admin/settings`
|
||||
pub async fn update_settings(
|
||||
pub(crate) async fn update_settings(
|
||||
State(state): State<AppState>,
|
||||
Extension(request_id): Extension<RequestId>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Result<impl IntoResponse, AdminError> {
|
||||
let request_id = request_id.as_str().to_string();
|
||||
require_admin_session(&state, &headers, &request_id).await?;
|
||||
require_admin_session(&state, &headers).await?;
|
||||
|
||||
let patch = parse_settings_patch(&body, &request_id)?;
|
||||
let patch = parse_settings_patch(&body)?;
|
||||
let config = state
|
||||
.services
|
||||
.settings
|
||||
.update(patch)
|
||||
.await
|
||||
.map_err(|error| settings_service_error(error, &request_id))?;
|
||||
.map_err(settings_service_error)?;
|
||||
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminSettingsData::from_config(&config), request_id),
|
||||
AdminEnvelope::ok(AdminSettingsData::from_config(&config)),
|
||||
))
|
||||
}
|
||||
|
||||
fn parse_settings_patch(body: &[u8], request_id: &str) -> Result<AdminSettingsPatch, AdminError> {
|
||||
let value = serde_json::from_slice::<Value>(body).map_err(|_| {
|
||||
AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40000,
|
||||
"Malformed JSON body",
|
||||
request_id,
|
||||
)
|
||||
})?;
|
||||
fn parse_settings_patch(body: &[u8]) -> Result<AdminSettingsPatch, AdminError> {
|
||||
let value = serde_json::from_slice::<Value>(body)
|
||||
.map_err(|_| AdminError::new(StatusCode::BAD_REQUEST, 40000, "Malformed JSON body"))?;
|
||||
let object = value.as_object().ok_or_else(|| {
|
||||
AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
"Settings payload must be an object",
|
||||
request_id,
|
||||
)
|
||||
})?;
|
||||
|
||||
@@ -124,21 +110,14 @@ fn parse_settings_patch(body: &[u8], request_id: &str) -> Result<AdminSettingsPa
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
format!("Unsupported settings field: {key}"),
|
||||
request_id,
|
||||
));
|
||||
}
|
||||
|
||||
serde_json::from_value(value).map_err(|_| {
|
||||
AdminError::new(
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
"Invalid settings payload",
|
||||
request_id,
|
||||
)
|
||||
})
|
||||
serde_json::from_value(value)
|
||||
.map_err(|_| AdminError::new(StatusCode::BAD_REQUEST, 40001, "Invalid settings payload"))
|
||||
}
|
||||
|
||||
fn settings_service_error(error: RuntimeSettingsError, request_id: &str) -> AdminError {
|
||||
fn settings_service_error(error: RuntimeSettingsError) -> AdminError {
|
||||
match error {
|
||||
RuntimeSettingsError::InvalidField(SettingsServiceError::InvalidField {
|
||||
field,
|
||||
@@ -147,7 +126,6 @@ fn settings_service_error(error: RuntimeSettingsError, request_id: &str) -> Admi
|
||||
StatusCode::BAD_REQUEST,
|
||||
40001,
|
||||
format!("Invalid setting {field}: {message}"),
|
||||
request_id,
|
||||
),
|
||||
RuntimeSettingsError::Database(_)
|
||||
| RuntimeSettingsError::Json(_)
|
||||
@@ -155,7 +133,6 @@ fn settings_service_error(error: RuntimeSettingsError, request_id: &str) -> Admi
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
50000,
|
||||
"Failed to persist settings",
|
||||
request_id,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,15 +1,15 @@
|
||||
use std::path::Path;
|
||||
|
||||
use super::{types::AppConfig, ConfigResult};
|
||||
use super::types::AppConfig;
|
||||
|
||||
impl AppConfig {
|
||||
/// 从当前目录加载配置文件。
|
||||
pub fn load() -> ConfigResult<Self> {
|
||||
pub fn load() -> Result<Self, ::config::ConfigError> {
|
||||
Self::load_from_dir(".")
|
||||
}
|
||||
|
||||
/// 从指定目录加载配置文件。
|
||||
pub fn load_from_dir(config_dir: impl AsRef<Path>) -> ConfigResult<Self> {
|
||||
pub fn load_from_dir(config_dir: impl AsRef<Path>) -> Result<Self, ::config::ConfigError> {
|
||||
let config_dir = config_dir.as_ref();
|
||||
::config::Config::builder()
|
||||
.add_source(::config::File::from(config_dir.join("config.yaml")).required(true))
|
||||
|
||||
@@ -8,6 +8,3 @@ pub mod loader;
|
||||
pub mod settings;
|
||||
/// 配置类型定义。
|
||||
pub mod types;
|
||||
|
||||
/// 配置加载结果类型。
|
||||
pub type ConfigResult<T> = Result<T, ::config::ConfigError>;
|
||||
|
||||
+8
-23
@@ -102,7 +102,7 @@ impl SettingsService {
|
||||
current.request_interval_ms = request_interval_ms;
|
||||
}
|
||||
if let Some(rotation_strategy) = patch.rotation_strategy {
|
||||
current.rotation_strategy = validate_rotation_strategy(rotation_strategy)?;
|
||||
current.rotation_strategy = validate_rotation_strategy(&rotation_strategy)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -128,12 +128,6 @@ impl SettingsServiceError {
|
||||
Self::InvalidField { field, .. } => field,
|
||||
}
|
||||
}
|
||||
/// 返回字段错误说明。
|
||||
pub fn message(&self) -> &str {
|
||||
match self {
|
||||
Self::InvalidField { message, .. } => message,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_model_aliases(
|
||||
@@ -141,8 +135,8 @@ fn validate_model_aliases(
|
||||
) -> Result<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)?;
|
||||
let alias = non_empty_string("modelAliases", &alias)?;
|
||||
let target = non_empty_string("modelAliases", &target)?;
|
||||
if alias == target {
|
||||
return Err(invalid_field(
|
||||
"modelAliases",
|
||||
@@ -159,10 +153,10 @@ fn validate_model_account_routes(
|
||||
) -> Result<BTreeMap<String, Vec<String>>, SettingsServiceError> {
|
||||
let mut normalized = BTreeMap::new();
|
||||
for (model, account_ids) in routes {
|
||||
let model = non_empty_string("modelAccountRoutes", model)?;
|
||||
let model = non_empty_string("modelAccountRoutes", &model)?;
|
||||
let mut normalized_account_ids = Vec::new();
|
||||
for account_id in account_ids {
|
||||
let account_id = non_empty_string("modelAccountRoutes", account_id)?;
|
||||
let account_id = non_empty_string("modelAccountRoutes", &account_id)?;
|
||||
if normalized_account_ids.contains(&account_id) {
|
||||
return Err(invalid_field(
|
||||
"modelAccountRoutes",
|
||||
@@ -182,7 +176,7 @@ fn validate_model_account_routes(
|
||||
Ok(normalized)
|
||||
}
|
||||
|
||||
fn validate_rotation_strategy(strategy: String) -> Result<String, SettingsServiceError> {
|
||||
fn validate_rotation_strategy(strategy: &str) -> Result<String, SettingsServiceError> {
|
||||
let strategy = non_empty_string("rotationStrategy", strategy)?;
|
||||
if ROTATION_STRATEGIES.contains(&strategy.as_str()) {
|
||||
Ok(strategy)
|
||||
@@ -194,7 +188,7 @@ fn validate_rotation_strategy(strategy: String) -> Result<String, SettingsServic
|
||||
}
|
||||
}
|
||||
|
||||
fn non_empty_string(field: &str, value: String) -> Result<String, SettingsServiceError> {
|
||||
fn non_empty_string(field: &str, value: &str) -> Result<String, SettingsServiceError> {
|
||||
let value = value.trim().to_string();
|
||||
if value.is_empty() {
|
||||
Err(invalid_field(field, "must not be empty"))
|
||||
@@ -243,15 +237,6 @@ pub struct RuntimeSettingsService {
|
||||
}
|
||||
|
||||
impl RuntimeSettingsService {
|
||||
/// 构造运行时设置服务。
|
||||
pub fn new(config: AppConfig, pool: SqlitePool) -> Self {
|
||||
Self {
|
||||
current: Arc::new(StdRwLock::new(Arc::new(config))),
|
||||
pool,
|
||||
account_pool: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 构造带账号池同步的运行时设置服务。
|
||||
pub fn with_account_pool(
|
||||
config: AppConfig,
|
||||
@@ -494,7 +479,7 @@ fn runtime_settings_from_row(row: &SqliteRow) -> Result<AdminSettings, RuntimeSe
|
||||
"requestIntervalMs",
|
||||
row.get("request_interval_ms"),
|
||||
)?,
|
||||
rotation_strategy: validate_rotation_strategy(row.get("rotation_strategy"))?,
|
||||
rotation_strategy: validate_rotation_strategy(&row.get::<String, _>("rotation_strategy"))?,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -101,33 +101,12 @@ impl Default for AuthConfig {
|
||||
pub struct QuotaConfig {
|
||||
/// 配额刷新周期。
|
||||
pub refresh_interval_minutes: u64,
|
||||
/// 预警阈值。
|
||||
pub warning_thresholds: QuotaWarningThresholds,
|
||||
}
|
||||
|
||||
impl Default for QuotaConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
refresh_interval_minutes: 5,
|
||||
warning_thresholds: QuotaWarningThresholds::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 配额预警阈值集合。
|
||||
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
|
||||
pub struct QuotaWarningThresholds {
|
||||
/// 主阈值列表。
|
||||
pub primary: Vec<u8>,
|
||||
/// 次阈值列表。
|
||||
pub secondary: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Default for QuotaWarningThresholds {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
primary: vec![80, 90],
|
||||
secondary: vec![80, 90],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,9 +14,6 @@ use tracing::{info, warn, Span};
|
||||
|
||||
use crate::http::middleware::request_id::RequestId;
|
||||
|
||||
/// 默认 tracing 事件名。
|
||||
pub const TRACE_LAYER_NAME: &str = "http";
|
||||
|
||||
/// HTTP tracing layer type.
|
||||
pub type HttpTraceLayer = TraceLayer<
|
||||
SharedClassifier<ServerErrorsAsFailures>,
|
||||
@@ -65,6 +62,10 @@ fn on_http_response(response: &Response<Body>, latency: Duration, _span: &Span)
|
||||
);
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::needless_pass_by_value,
|
||||
reason = "tower-http OnFailure callbacks receive the failure class by value"
|
||||
)]
|
||||
fn on_http_failure(failure_class: ServerErrorsFailureClass, latency: Duration, _span: &Span) {
|
||||
warn!(
|
||||
failure = %failure_class,
|
||||
|
||||
@@ -2,8 +2,7 @@
|
||||
|
||||
use std::{fs, path::Path, str::FromStr, time::Duration};
|
||||
|
||||
use sqlx::sqlite::SqliteConnectOptions;
|
||||
pub use sqlx::SqlitePool;
|
||||
use sqlx::{sqlite::SqliteConnectOptions, SqlitePool};
|
||||
|
||||
/// 连接 SQLite 并初始化 schema。
|
||||
pub async fn connect_sqlite(database_url: &str) -> Result<SqlitePool, sqlx::Error> {
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 分页
|
||||
@@ -71,25 +70,3 @@ pub fn total_pages(total: u64, page_size: u32) -> u32 {
|
||||
let pages = total.div_ceil(page_size);
|
||||
pages.min(u64::from(u32::MAX)) as u32
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// JSON 值提取
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// 从多个路径中提取第一个非空字符串。
|
||||
pub fn first_string(value: &Value, paths: &[&[&str]]) -> Option<String> {
|
||||
paths.iter().find_map(|path| string_at(value, path))
|
||||
}
|
||||
|
||||
/// 从指定路径提取字符串。
|
||||
pub fn string_at(value: &Value, path: &[&str]) -> Option<String> {
|
||||
let mut current = value;
|
||||
for segment in path {
|
||||
current = current.get(*segment)?;
|
||||
}
|
||||
current
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
@@ -73,8 +73,8 @@ pub fn build_file_appender(config: &RotationConfig) -> Result<ChinaDailyFileAppe
|
||||
}
|
||||
|
||||
/// 初始化 tracing。
|
||||
pub fn init_tracing(config: RotationConfig) -> Result<LogGuard, LogError> {
|
||||
let appender = build_file_appender(&config)?;
|
||||
pub fn init_tracing(config: &RotationConfig) -> Result<LogGuard, LogError> {
|
||||
let appender = build_file_appender(config)?;
|
||||
let (writer, guard) = tracing_appender::non_blocking(appender);
|
||||
let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
|
||||
|
||||
|
||||
+4
-10
@@ -9,8 +9,7 @@ use uuid::Uuid;
|
||||
|
||||
const INSTALLATION_ID_FILE_NAME: &str = "installation_id";
|
||||
|
||||
/// 返回本地数据目录。
|
||||
pub fn data_dir() -> PathBuf {
|
||||
fn data_dir() -> PathBuf {
|
||||
dirs::data_local_dir()
|
||||
.unwrap_or_else(|| PathBuf::from("."))
|
||||
.join("codex-proxy-rs")
|
||||
@@ -24,13 +23,11 @@ pub fn ensure_data_dir() -> std::io::Result<PathBuf> {
|
||||
Ok(dir)
|
||||
}
|
||||
|
||||
/// 返回真实 Codex Desktop installation ID 文件路径。
|
||||
pub fn codex_desktop_installation_id_path() -> Option<PathBuf> {
|
||||
fn codex_desktop_installation_id_path() -> Option<PathBuf> {
|
||||
dirs::home_dir().map(|home| home.join(".codex").join(INSTALLATION_ID_FILE_NAME))
|
||||
}
|
||||
|
||||
/// 返回本应用数据目录中的 installation ID 文件路径。
|
||||
pub fn data_installation_id_path(data_dir: &Path) -> PathBuf {
|
||||
fn data_installation_id_path(data_dir: &Path) -> PathBuf {
|
||||
data_dir.join(INSTALLATION_ID_FILE_NAME)
|
||||
}
|
||||
|
||||
@@ -44,10 +41,7 @@ pub fn load_or_create_installation_id(data_dir: Option<&Path>) -> io::Result<Str
|
||||
load_or_create_installation_id_from_paths(codex_path.as_deref(), data_path.as_deref())
|
||||
}
|
||||
|
||||
/// 从显式文件路径按读取优先级读取或生成 installation ID。
|
||||
///
|
||||
/// 此函数用于测试和调用方已经完成路径解析的场景。
|
||||
pub fn load_or_create_installation_id_from_paths(
|
||||
fn load_or_create_installation_id_from_paths(
|
||||
codex_desktop_path: Option<&Path>,
|
||||
data_path: Option<&Path>,
|
||||
) -> io::Result<String> {
|
||||
|
||||
+13
-10
@@ -19,21 +19,23 @@ pub fn china_rfc3339_millis(value: &DateTime<Utc>) -> String {
|
||||
|
||||
/// 将 RFC3339 字符串输出为中国时区 RFC3339 字符串。
|
||||
pub fn china_rfc3339_str(value: &str) -> String {
|
||||
DateTime::parse_from_rfc3339(value)
|
||||
.map(|datetime| datetime.with_timezone(&china_offset()).to_rfc3339())
|
||||
.unwrap_or_else(|_| value.to_string())
|
||||
DateTime::parse_from_rfc3339(value).map_or_else(
|
||||
|_| value.to_string(),
|
||||
|datetime| datetime.with_timezone(&china_offset()).to_rfc3339(),
|
||||
)
|
||||
}
|
||||
|
||||
/// 将 RFC3339 字符串输出为中国时区日期时间。
|
||||
pub fn china_datetime_rfc3339_str(value: &str) -> String {
|
||||
DateTime::parse_from_rfc3339(value)
|
||||
.map(|datetime| {
|
||||
DateTime::parse_from_rfc3339(value).map_or_else(
|
||||
|_| value.to_string(),
|
||||
|datetime| {
|
||||
datetime
|
||||
.with_timezone(&china_offset())
|
||||
.format("%Y-%m-%d %H:%M:%S")
|
||||
.to_string()
|
||||
})
|
||||
.unwrap_or_else(|_| value.to_string())
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// 将 UTC 时间输出为中国时区日期时间。
|
||||
@@ -149,9 +151,10 @@ pub fn china_relative_time_str(value: Option<&str>) -> String {
|
||||
let Some(value) = value else {
|
||||
return "-".to_string();
|
||||
};
|
||||
DateTime::parse_from_rfc3339(value)
|
||||
.map(|datetime| china_relative_time(Some(datetime.with_timezone(&Utc)), Utc::now()))
|
||||
.unwrap_or_else(|_| value.to_string())
|
||||
DateTime::parse_from_rfc3339(value).map_or_else(
|
||||
|_| value.to_string(),
|
||||
|datetime| china_relative_time(Some(datetime.with_timezone(&Utc)), Utc::now()),
|
||||
)
|
||||
}
|
||||
|
||||
/// Serde 序列化 UTC 时间为中国时区 RFC3339 字符串。
|
||||
|
||||
@@ -26,7 +26,8 @@ use crate::{
|
||||
},
|
||||
upstream::{
|
||||
create_response_with_account, verify_acquired_quota_if_required,
|
||||
QuotaVerificationDecision, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
|
||||
QuotaVerificationContext, QuotaVerificationDecision,
|
||||
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
|
||||
},
|
||||
},
|
||||
proxy::openai::chat::ChatStreamTranslationError,
|
||||
@@ -149,7 +150,7 @@ impl ChatDispatchService {
|
||||
let (account_id, response) = loop {
|
||||
let acquire_request = AccountAcquireRequest::new(&request.model, Utc::now())
|
||||
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
|
||||
let acquired = match self.account_pool.acquire_with(acquire_request).await {
|
||||
let acquired = match self.account_pool.acquire_with(&acquire_request).await {
|
||||
Some(acquired) => acquired,
|
||||
None if quota_exhausted_count > 0 => {
|
||||
return_dispatch_error!(ChatDispatchError::QuotaExhausted {
|
||||
@@ -203,14 +204,16 @@ impl ChatDispatchService {
|
||||
None => return_dispatch_error!(ChatDispatchError::NoActiveAccount),
|
||||
};
|
||||
let acquired = match verify_acquired_quota_if_required(
|
||||
self.account_pool.as_ref(),
|
||||
self.codex.as_ref(),
|
||||
&self.cloudflare,
|
||||
self.installation_id.as_deref(),
|
||||
request_id,
|
||||
QuotaVerificationContext {
|
||||
account_pool: self.account_pool.as_ref(),
|
||||
codex: self.codex.as_ref(),
|
||||
cloudflare: &self.cloudflare,
|
||||
installation_id: self.installation_id.as_deref(),
|
||||
request_id,
|
||||
excluded_account_ids: &mut excluded_account_ids,
|
||||
verify_attempts: &mut quota_verify_attempts,
|
||||
},
|
||||
acquired,
|
||||
&mut excluded_account_ids,
|
||||
&mut quota_verify_attempts,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -25,8 +25,8 @@ pub struct CloudflareRecovery {
|
||||
impl CloudflareRecovery {
|
||||
pub fn new(cookie_store: SqliteCookieStore) -> Self {
|
||||
Self {
|
||||
path_block_tracker: Default::default(),
|
||||
challenge_tracker: Default::default(),
|
||||
path_block_tracker: CloudflarePathBlockTracker::default(),
|
||||
challenge_tracker: CloudflareChallengeCooldownTracker::default(),
|
||||
cookie_store,
|
||||
}
|
||||
}
|
||||
@@ -104,7 +104,7 @@ impl CloudflareRecovery {
|
||||
fn clear_challenge_cookies_after_cooldown(&self, account_id: &str, delay_seconds: i64) {
|
||||
let cookie_store = self.cookie_store.clone();
|
||||
let account_id = account_id.to_string();
|
||||
let delay = StdDuration::from_secs(delay_seconds.max(0) as u64);
|
||||
let delay = StdDuration::from_secs(delay_seconds.max(0).cast_unsigned());
|
||||
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(delay).await;
|
||||
|
||||
+210
-195
@@ -49,7 +49,8 @@ use crate::{
|
||||
create_compact_response_with_account_retrying_5xx,
|
||||
create_response_stream_with_account_retrying_5xx,
|
||||
create_response_with_account_retrying_5xx, verify_acquired_quota_if_required,
|
||||
QuotaVerificationDecision, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
|
||||
QuotaVerificationContext, QuotaVerificationDecision,
|
||||
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
|
||||
},
|
||||
},
|
||||
upstream::accounts::{
|
||||
@@ -430,89 +431,89 @@ impl ResponseDispatchService {
|
||||
.clone()
|
||||
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
|
||||
attempt_acquire_request.now = Utc::now();
|
||||
let acquired = match self
|
||||
let Some(acquired) = self
|
||||
.account_pool
|
||||
.acquire_with(attempt_acquire_request)
|
||||
.acquire_with(&attempt_acquire_request)
|
||||
.await
|
||||
{
|
||||
Some(acquired) => acquired,
|
||||
None => {
|
||||
let error = match last_exhausted_account_class {
|
||||
Some(ExhaustedAccountClass::QuotaExhausted) => {
|
||||
ResponseDispatchError::QuotaExhausted {
|
||||
count: quota_exhausted_count,
|
||||
upstream_error: last_quota_error.unwrap_or_default(),
|
||||
}
|
||||
else {
|
||||
let error = match last_exhausted_account_class {
|
||||
Some(ExhaustedAccountClass::QuotaExhausted) => {
|
||||
ResponseDispatchError::QuotaExhausted {
|
||||
count: quota_exhausted_count,
|
||||
upstream_error: last_quota_error.unwrap_or_default(),
|
||||
}
|
||||
Some(ExhaustedAccountClass::RateLimited) => {
|
||||
ResponseDispatchError::RateLimited {
|
||||
count: rate_limited_count,
|
||||
upstream_error: last_rate_limit_error.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
Some(ExhaustedAccountClass::RateLimited) => {
|
||||
ResponseDispatchError::RateLimited {
|
||||
count: rate_limited_count,
|
||||
upstream_error: last_rate_limit_error.unwrap_or_default(),
|
||||
}
|
||||
Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired {
|
||||
count: expired_count,
|
||||
upstream_error: last_auth_error.unwrap_or_default(),
|
||||
},
|
||||
Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled {
|
||||
count: disabled_count,
|
||||
upstream_error: last_disabled_auth_error.unwrap_or_default(),
|
||||
},
|
||||
Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned {
|
||||
count: banned_count,
|
||||
upstream_error: last_banned_auth_error.unwrap_or_default(),
|
||||
status_code: last_banned_status_code.unwrap_or(403),
|
||||
},
|
||||
Some(ExhaustedAccountClass::CloudflareChallenge) => {
|
||||
ResponseDispatchError::CloudflareChallenge {
|
||||
count: cloudflare_challenge_count,
|
||||
upstream_error: last_cloudflare_challenge_error.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired {
|
||||
count: expired_count,
|
||||
upstream_error: last_auth_error.unwrap_or_default(),
|
||||
},
|
||||
Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled {
|
||||
count: disabled_count,
|
||||
upstream_error: last_disabled_auth_error.unwrap_or_default(),
|
||||
},
|
||||
Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned {
|
||||
count: banned_count,
|
||||
upstream_error: last_banned_auth_error.unwrap_or_default(),
|
||||
status_code: last_banned_status_code.unwrap_or(403),
|
||||
},
|
||||
Some(ExhaustedAccountClass::CloudflareChallenge) => {
|
||||
ResponseDispatchError::CloudflareChallenge {
|
||||
count: cloudflare_challenge_count,
|
||||
upstream_error: last_cloudflare_challenge_error.unwrap_or_default(),
|
||||
}
|
||||
Some(ExhaustedAccountClass::CloudflarePathBlocked) => {
|
||||
ResponseDispatchError::CloudflarePathBlocked {
|
||||
count: cloudflare_path_block_count,
|
||||
upstream_error: last_cloudflare_path_block_error
|
||||
.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
Some(ExhaustedAccountClass::CloudflarePathBlocked) => {
|
||||
ResponseDispatchError::CloudflarePathBlocked {
|
||||
count: cloudflare_path_block_count,
|
||||
upstream_error: last_cloudflare_path_block_error.unwrap_or_default(),
|
||||
}
|
||||
Some(ExhaustedAccountClass::ModelUnsupported) => {
|
||||
ResponseDispatchError::ModelUnsupported {
|
||||
count: model_unsupported_count,
|
||||
upstream_error: last_model_unsupported_error.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
Some(ExhaustedAccountClass::ModelUnsupported) => {
|
||||
ResponseDispatchError::ModelUnsupported {
|
||||
count: model_unsupported_count,
|
||||
upstream_error: last_model_unsupported_error.unwrap_or_default(),
|
||||
}
|
||||
None => ResponseDispatchError::NoActiveAccount,
|
||||
};
|
||||
self.record_response_dispatch_error(
|
||||
request_id,
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
last_attempted_account_id.as_deref(),
|
||||
false,
|
||||
false,
|
||||
Some(backend_transport_name(
|
||||
}
|
||||
None => ResponseDispatchError::NoActiveAccount,
|
||||
};
|
||||
self.record_response_dispatch_error(
|
||||
request_id,
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: last_attempted_account_id.as_deref(),
|
||||
stream: false,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(
|
||||
backend_transport_for_response_request(&request),
|
||||
)),
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
return Err(error);
|
||||
}
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
return Err(error);
|
||||
};
|
||||
let acquired_account_id = acquired.account.id.clone();
|
||||
|
||||
// 配额验证
|
||||
let acquired = match verify_acquired_quota_if_required(
|
||||
self.account_pool.as_ref(),
|
||||
self.codex.as_ref(),
|
||||
&self.cloudflare,
|
||||
self.installation_id.as_deref(),
|
||||
request_id,
|
||||
QuotaVerificationContext {
|
||||
account_pool: self.account_pool.as_ref(),
|
||||
codex: self.codex.as_ref(),
|
||||
cloudflare: &self.cloudflare,
|
||||
installation_id: self.installation_id.as_deref(),
|
||||
request_id,
|
||||
excluded_account_ids: &mut excluded_account_ids,
|
||||
verify_attempts: &mut quota_verify_attempts,
|
||||
},
|
||||
acquired,
|
||||
&mut excluded_account_ids,
|
||||
&mut quota_verify_attempts,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -533,12 +534,14 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
Some(&acquired_account_id),
|
||||
false,
|
||||
false,
|
||||
Some(backend_transport_name(
|
||||
backend_transport_for_response_request(&request),
|
||||
)),
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: Some(&acquired_account_id),
|
||||
stream: false,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(
|
||||
backend_transport_for_response_request(&request),
|
||||
)),
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -595,10 +598,12 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
Some(&release_account_id),
|
||||
false,
|
||||
false,
|
||||
Some(backend_transport_name(response.transport)),
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: Some(&release_account_id),
|
||||
stream: false,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(response.transport)),
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -641,10 +646,12 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
Some(&release_account_id),
|
||||
false,
|
||||
false,
|
||||
Some(backend_transport_name(response.transport)),
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: Some(&release_account_id),
|
||||
stream: false,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(response.transport)),
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -792,12 +799,14 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
Some(&release_account_id),
|
||||
false,
|
||||
false,
|
||||
Some(backend_transport_name(
|
||||
backend_transport_for_response_request(&request),
|
||||
)),
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: Some(&release_account_id),
|
||||
stream: false,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(
|
||||
backend_transport_for_response_request(&request),
|
||||
)),
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -890,10 +899,12 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
Some(&account.id),
|
||||
false,
|
||||
false,
|
||||
Some(backend_transport_name(response.transport)),
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: Some(&account.id),
|
||||
stream: false,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(response.transport)),
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -906,10 +917,12 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
Some(&account.id),
|
||||
false,
|
||||
false,
|
||||
Some(backend_transport_name(response.transport)),
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: Some(&account.id),
|
||||
stream: false,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(response.transport)),
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -922,10 +935,12 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
Some(&account.id),
|
||||
false,
|
||||
false,
|
||||
Some(backend_transport_name(response.transport)),
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: Some(&account.id),
|
||||
stream: false,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(response.transport)),
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -934,32 +949,25 @@ impl ResponseDispatchService {
|
||||
}
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::too_many_arguments,
|
||||
reason = "keeps response dispatch logging call sites explicit at branch exits"
|
||||
)]
|
||||
async fn record_response_dispatch_error(
|
||||
&self,
|
||||
request_id: &str,
|
||||
route: &str,
|
||||
requested_model: &str,
|
||||
started_at: Instant,
|
||||
account_id: Option<&str>,
|
||||
stream: bool,
|
||||
compact: bool,
|
||||
transport: Option<&str>,
|
||||
details: ResponseDispatchErrorDetails<'_>,
|
||||
error: &ResponseDispatchError,
|
||||
) {
|
||||
record_response_dispatch_error_event(ResponseDispatchErrorEventRecord {
|
||||
logs: &self.logs,
|
||||
request_id,
|
||||
account_id,
|
||||
account_id: details.account_id,
|
||||
route,
|
||||
model: requested_model,
|
||||
started_at,
|
||||
stream,
|
||||
compact,
|
||||
transport,
|
||||
stream: details.stream,
|
||||
compact: details.compact,
|
||||
transport: details.transport,
|
||||
error,
|
||||
})
|
||||
.await;
|
||||
@@ -1018,12 +1026,14 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
last_attempted_account_id.as_deref(),
|
||||
true,
|
||||
false,
|
||||
Some(backend_transport_name(
|
||||
backend_transport_for_response_request(&request),
|
||||
)),
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: last_attempted_account_id.as_deref(),
|
||||
stream: true,
|
||||
compact: false,
|
||||
transport: Some(backend_transport_name(
|
||||
backend_transport_for_response_request(&request),
|
||||
)),
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -1036,10 +1046,12 @@ impl ResponseDispatchService {
|
||||
route,
|
||||
requested_model,
|
||||
started_at,
|
||||
$account_id,
|
||||
true,
|
||||
false,
|
||||
$transport,
|
||||
ResponseDispatchErrorDetails {
|
||||
account_id: $account_id,
|
||||
stream: true,
|
||||
compact: false,
|
||||
transport: $transport,
|
||||
},
|
||||
&error,
|
||||
)
|
||||
.await;
|
||||
@@ -1051,73 +1063,71 @@ impl ResponseDispatchService {
|
||||
.clone()
|
||||
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
|
||||
attempt_acquire_request.now = Utc::now();
|
||||
let acquired = match self
|
||||
let Some(acquired) = self
|
||||
.account_pool
|
||||
.acquire_with(attempt_acquire_request)
|
||||
.acquire_with(&attempt_acquire_request)
|
||||
.await
|
||||
{
|
||||
Some(acquired) => acquired,
|
||||
None => {
|
||||
let error = match last_exhausted_account_class {
|
||||
Some(ExhaustedAccountClass::QuotaExhausted) => {
|
||||
ResponseDispatchError::QuotaExhausted {
|
||||
count: quota_exhausted_count,
|
||||
upstream_error: last_quota_error.unwrap_or_default(),
|
||||
}
|
||||
else {
|
||||
let error = match last_exhausted_account_class {
|
||||
Some(ExhaustedAccountClass::QuotaExhausted) => {
|
||||
ResponseDispatchError::QuotaExhausted {
|
||||
count: quota_exhausted_count,
|
||||
upstream_error: last_quota_error.unwrap_or_default(),
|
||||
}
|
||||
Some(ExhaustedAccountClass::RateLimited) => {
|
||||
ResponseDispatchError::RateLimited {
|
||||
count: rate_limited_count,
|
||||
upstream_error: last_rate_limit_error.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
Some(ExhaustedAccountClass::RateLimited) => {
|
||||
ResponseDispatchError::RateLimited {
|
||||
count: rate_limited_count,
|
||||
upstream_error: last_rate_limit_error.unwrap_or_default(),
|
||||
}
|
||||
Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired {
|
||||
count: expired_count,
|
||||
upstream_error: last_auth_error.unwrap_or_default(),
|
||||
},
|
||||
Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled {
|
||||
count: disabled_count,
|
||||
upstream_error: last_disabled_auth_error.unwrap_or_default(),
|
||||
},
|
||||
Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned {
|
||||
count: banned_count,
|
||||
upstream_error: last_banned_auth_error.unwrap_or_default(),
|
||||
status_code: last_banned_status_code.unwrap_or(403),
|
||||
},
|
||||
Some(ExhaustedAccountClass::CloudflareChallenge) => {
|
||||
ResponseDispatchError::CloudflareChallenge {
|
||||
count: cloudflare_challenge_count,
|
||||
upstream_error: last_cloudflare_challenge_error.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired {
|
||||
count: expired_count,
|
||||
upstream_error: last_auth_error.unwrap_or_default(),
|
||||
},
|
||||
Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled {
|
||||
count: disabled_count,
|
||||
upstream_error: last_disabled_auth_error.unwrap_or_default(),
|
||||
},
|
||||
Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned {
|
||||
count: banned_count,
|
||||
upstream_error: last_banned_auth_error.unwrap_or_default(),
|
||||
status_code: last_banned_status_code.unwrap_or(403),
|
||||
},
|
||||
Some(ExhaustedAccountClass::CloudflareChallenge) => {
|
||||
ResponseDispatchError::CloudflareChallenge {
|
||||
count: cloudflare_challenge_count,
|
||||
upstream_error: last_cloudflare_challenge_error.unwrap_or_default(),
|
||||
}
|
||||
Some(ExhaustedAccountClass::CloudflarePathBlocked) => {
|
||||
ResponseDispatchError::CloudflarePathBlocked {
|
||||
count: cloudflare_path_block_count,
|
||||
upstream_error: last_cloudflare_path_block_error
|
||||
.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
Some(ExhaustedAccountClass::CloudflarePathBlocked) => {
|
||||
ResponseDispatchError::CloudflarePathBlocked {
|
||||
count: cloudflare_path_block_count,
|
||||
upstream_error: last_cloudflare_path_block_error.unwrap_or_default(),
|
||||
}
|
||||
Some(ExhaustedAccountClass::ModelUnsupported) => {
|
||||
ResponseDispatchError::ModelUnsupported {
|
||||
count: model_unsupported_count,
|
||||
upstream_error: last_model_unsupported_error.unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
Some(ExhaustedAccountClass::ModelUnsupported) => {
|
||||
ResponseDispatchError::ModelUnsupported {
|
||||
count: model_unsupported_count,
|
||||
upstream_error: last_model_unsupported_error.unwrap_or_default(),
|
||||
}
|
||||
None => ResponseDispatchError::NoActiveAccount,
|
||||
};
|
||||
return_stream_dispatch_error!(error);
|
||||
}
|
||||
}
|
||||
None => ResponseDispatchError::NoActiveAccount,
|
||||
};
|
||||
return_stream_dispatch_error!(error);
|
||||
};
|
||||
let acquired_account_id = acquired.account.id.clone();
|
||||
let acquired = match verify_acquired_quota_if_required(
|
||||
self.account_pool.as_ref(),
|
||||
self.codex.as_ref(),
|
||||
&self.cloudflare,
|
||||
self.installation_id.as_deref(),
|
||||
request_id,
|
||||
QuotaVerificationContext {
|
||||
account_pool: self.account_pool.as_ref(),
|
||||
codex: self.codex.as_ref(),
|
||||
cloudflare: &self.cloudflare,
|
||||
installation_id: self.installation_id.as_deref(),
|
||||
request_id,
|
||||
excluded_account_ids: &mut excluded_account_ids,
|
||||
verify_attempts: &mut quota_verify_attempts,
|
||||
},
|
||||
acquired,
|
||||
&mut excluded_account_ids,
|
||||
&mut quota_verify_attempts,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -1525,7 +1535,7 @@ impl ResponseDispatchService {
|
||||
loop {
|
||||
let acquire_request = AccountAcquireRequest::new(&request.model, Utc::now())
|
||||
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
|
||||
let acquired = match self.account_pool.acquire_with(acquire_request).await {
|
||||
let acquired = match self.account_pool.acquire_with(&acquire_request).await {
|
||||
Some(acquired) => acquired,
|
||||
None if quota_exhausted_count > 0 => {
|
||||
let error = ResponseDispatchError::QuotaExhausted {
|
||||
@@ -1663,14 +1673,16 @@ impl ResponseDispatchService {
|
||||
};
|
||||
last_attempted_account_id = Some(acquired.account.id.clone());
|
||||
let acquired = match verify_acquired_quota_if_required(
|
||||
self.account_pool.as_ref(),
|
||||
self.codex.as_ref(),
|
||||
&self.cloudflare,
|
||||
self.installation_id.as_deref(),
|
||||
request_id,
|
||||
QuotaVerificationContext {
|
||||
account_pool: self.account_pool.as_ref(),
|
||||
codex: self.codex.as_ref(),
|
||||
cloudflare: &self.cloudflare,
|
||||
installation_id: self.installation_id.as_deref(),
|
||||
request_id,
|
||||
excluded_account_ids: &mut excluded_account_ids,
|
||||
verify_attempts: &mut quota_verify_attempts,
|
||||
},
|
||||
acquired,
|
||||
&mut excluded_account_ids,
|
||||
&mut quota_verify_attempts,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -2261,7 +2273,8 @@ async fn send_live_response_stream_tail(
|
||||
return None;
|
||||
}
|
||||
}
|
||||
let failure = premature_close_failed_event(latest_response_id(&body_text), failure_detail);
|
||||
let failure =
|
||||
premature_close_failed_event(latest_response_id(&body_text).as_deref(), failure_detail);
|
||||
body_text.push_str(&failure);
|
||||
body_bytes.extend_from_slice(failure.as_bytes());
|
||||
if sender.send(Ok(Bytes::from(failure))).await.is_err() {
|
||||
@@ -2324,17 +2337,12 @@ fn latest_response_id(body: &str) -> Option<String> {
|
||||
})
|
||||
}
|
||||
|
||||
fn premature_close_failed_event(response_id: Option<String>, detail: Option<&str>) -> String {
|
||||
fn premature_close_failed_event(response_id: Option<&str>, detail: Option<&str>) -> String {
|
||||
let message = match detail.filter(|value| !value.trim().is_empty()) {
|
||||
Some(detail) => format!("Upstream stream closed before response.completed: {detail}"),
|
||||
None => "Upstream stream closed before response.completed".to_string(),
|
||||
};
|
||||
response_failed_sse_event_with_id(
|
||||
response_id.as_deref(),
|
||||
"server_error",
|
||||
"stream_disconnected",
|
||||
&message,
|
||||
)
|
||||
response_failed_sse_event_with_id(response_id, "server_error", "stream_disconnected", &message)
|
||||
}
|
||||
|
||||
async fn finalize_live_response_stream(context: LiveResponseStreamContext, body: String) {
|
||||
@@ -2500,6 +2508,13 @@ struct ResponseDispatchErrorEventRecord<'a> {
|
||||
error: &'a ResponseDispatchError,
|
||||
}
|
||||
|
||||
struct ResponseDispatchErrorDetails<'a> {
|
||||
account_id: Option<&'a str>,
|
||||
stream: bool,
|
||||
compact: bool,
|
||||
transport: Option<&'a str>,
|
||||
}
|
||||
|
||||
async fn record_response_dispatch_error_event(record: ResponseDispatchErrorEventRecord<'_>) {
|
||||
let mut metadata = dispatch_error_metadata(
|
||||
record.error,
|
||||
|
||||
@@ -15,12 +15,6 @@ use crate::upstream::protocol::responses::CodexResponsesRequest;
|
||||
// 亲和性核心类型
|
||||
// ====================================================================
|
||||
|
||||
/// 会话亲和性键。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SessionAffinityKey {
|
||||
pub conversation_id: String,
|
||||
}
|
||||
|
||||
/// 会话亲和性条目。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SessionAffinityEntry {
|
||||
@@ -401,10 +395,6 @@ impl SqliteSessionAffinityStore {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &SqlitePool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub async fn upsert(
|
||||
&self,
|
||||
response_id: &str,
|
||||
|
||||
@@ -35,33 +35,35 @@ pub(crate) enum QuotaVerificationDecision {
|
||||
MaxAttemptsReached,
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::too_many_arguments,
|
||||
reason = "quota verification is the boundary where account, transport, and retry state meet"
|
||||
)]
|
||||
pub(crate) struct QuotaVerificationContext<'a> {
|
||||
pub account_pool: &'a RuntimeAccountPoolService,
|
||||
pub codex: &'a CodexBackendClient,
|
||||
pub cloudflare: &'a CloudflareRecovery,
|
||||
pub installation_id: Option<&'a str>,
|
||||
pub request_id: &'a str,
|
||||
pub excluded_account_ids: &'a mut Vec<String>,
|
||||
pub verify_attempts: &'a mut usize,
|
||||
}
|
||||
|
||||
pub(crate) async fn verify_acquired_quota_if_required(
|
||||
account_pool: &RuntimeAccountPoolService,
|
||||
codex: &CodexBackendClient,
|
||||
cloudflare: &CloudflareRecovery,
|
||||
installation_id: Option<&str>,
|
||||
request_id: &str,
|
||||
context: QuotaVerificationContext<'_>,
|
||||
acquired: AcquiredAccount,
|
||||
excluded_account_ids: &mut Vec<String>,
|
||||
verify_attempts: &mut usize,
|
||||
) -> QuotaVerificationDecision {
|
||||
if !acquired.account.quota_verify_required {
|
||||
return QuotaVerificationDecision::Ready(Box::new(acquired));
|
||||
}
|
||||
|
||||
let account_id = acquired.account.id.clone();
|
||||
let cookie_header = cloudflare
|
||||
let cookie_header = context
|
||||
.cloudflare
|
||||
.cookie_header_for_request(&account_id, "/codex/usage")
|
||||
.await;
|
||||
let usage = codex
|
||||
let usage = context
|
||||
.codex
|
||||
.fetch_usage(CodexRequestContext {
|
||||
access_token: &acquired.account.access_token,
|
||||
account_id: acquired.account.account_id.as_deref(),
|
||||
request_id,
|
||||
request_id: context.request_id,
|
||||
turn_state: None,
|
||||
turn_metadata: None,
|
||||
beta_features: None,
|
||||
@@ -70,7 +72,7 @@ pub(crate) async fn verify_acquired_quota_if_required(
|
||||
codex_window_id: None,
|
||||
parent_thread_id: None,
|
||||
cookie_header: cookie_header.as_deref(),
|
||||
installation_id,
|
||||
installation_id: context.installation_id,
|
||||
session_id: None,
|
||||
})
|
||||
.await;
|
||||
@@ -88,12 +90,15 @@ pub(crate) async fn verify_acquired_quota_if_required(
|
||||
};
|
||||
|
||||
let quota = quota_from_usage(&raw);
|
||||
account_pool.apply_quota_snapshot(&account_id, "a).await;
|
||||
context
|
||||
.account_pool
|
||||
.apply_quota_snapshot(&account_id, "a)
|
||||
.await;
|
||||
if quota_snapshot_limit_reached("a) {
|
||||
account_pool.release(&account_id).await;
|
||||
excluded_account_ids.push(account_id);
|
||||
*verify_attempts += 1;
|
||||
if *verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS {
|
||||
context.account_pool.release(&account_id).await;
|
||||
context.excluded_account_ids.push(account_id);
|
||||
*context.verify_attempts += 1;
|
||||
if *context.verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS {
|
||||
return QuotaVerificationDecision::MaxAttemptsReached;
|
||||
}
|
||||
return QuotaVerificationDecision::RetryWithAnotherAccount;
|
||||
|
||||
+30
-36
@@ -442,8 +442,7 @@ impl ChatCompletionStreamTranslator {
|
||||
let call_id = self
|
||||
.function_call_items
|
||||
.get(event_id)
|
||||
.map(|info| info.call_id.as_str())
|
||||
.unwrap_or(event_id)
|
||||
.map_or(event_id, |info| info.call_id.as_str())
|
||||
.to_string();
|
||||
Some((call_id, delta.to_string()))
|
||||
}
|
||||
@@ -456,8 +455,7 @@ impl ChatCompletionStreamTranslator {
|
||||
let call_id = self
|
||||
.function_call_items
|
||||
.get(event_id)
|
||||
.map(|info| info.call_id.as_str())
|
||||
.unwrap_or(event_id)
|
||||
.map_or(event_id, |info| info.call_id.as_str())
|
||||
.to_string();
|
||||
let arguments = value
|
||||
.get("arguments")
|
||||
@@ -840,11 +838,9 @@ pub fn chat_completion_from_codex_sse(
|
||||
return Err(ChatStreamTranslationError::Upstream(failure));
|
||||
}
|
||||
|
||||
if completed_response.is_none() {
|
||||
let Some(response) = completed_response else {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let response = completed_response.unwrap();
|
||||
};
|
||||
let output_text_value = if output_text.is_empty() {
|
||||
output_text_from_response(&response)
|
||||
} else {
|
||||
@@ -913,11 +909,8 @@ pub async fn chat_completions(
|
||||
let parsed_model = catalog.parse_model_name(&model);
|
||||
let display_model = ModelCatalog::build_display_model_name(&parsed_model);
|
||||
let stream = chat_request.stream;
|
||||
let mut codex_request = match translate_chat_to_codex(chat_request) {
|
||||
Ok(request) => request,
|
||||
Err(_) => {
|
||||
return invalid_chat_completion_request_response().into_response();
|
||||
}
|
||||
let Ok(mut codex_request) = translate_chat_to_codex(chat_request) else {
|
||||
return invalid_chat_completion_request_response().into_response();
|
||||
};
|
||||
apply_response_model_options(&mut codex_request, &parsed_model);
|
||||
let include_reasoning = codex_request
|
||||
@@ -946,7 +939,7 @@ pub async fn chat_completions(
|
||||
include_reasoning,
|
||||
tuple_schema,
|
||||
),
|
||||
Err(error) => response_dispatch_chat_stream_error_response(error),
|
||||
Err(error) => response_dispatch_chat_stream_error_response(&error),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -957,7 +950,7 @@ pub async fn chat_completions(
|
||||
.await
|
||||
{
|
||||
Ok(body) => (StatusCode::OK, Json(body)).into_response(),
|
||||
Err(error) => chat_dispatch_error_response(error),
|
||||
Err(error) => chat_dispatch_error_response(&error),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -991,8 +984,8 @@ fn live_chat_event_stream_response(
|
||||
)
|
||||
}
|
||||
|
||||
fn response_dispatch_chat_stream_error_response(error: ResponseDispatchError) -> Response {
|
||||
let message = chat_stream_dispatch_error_message(&error);
|
||||
fn response_dispatch_chat_stream_error_response(error: &ResponseDispatchError) -> Response {
|
||||
let message = chat_stream_dispatch_error_message(error);
|
||||
chat_stream_error_response(&message)
|
||||
}
|
||||
|
||||
@@ -1024,20 +1017,22 @@ fn output_text_from_response(response: &Value) -> String {
|
||||
response
|
||||
.get("output_text")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToString::to_string)
|
||||
.unwrap_or_else(|| {
|
||||
response
|
||||
.get("output")
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(output_text_from_item)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n")
|
||||
})
|
||||
.unwrap_or_default()
|
||||
})
|
||||
.map_or_else(
|
||||
|| {
|
||||
response
|
||||
.get("output")
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.filter_map(output_text_from_item)
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n")
|
||||
})
|
||||
.unwrap_or_default()
|
||||
},
|
||||
ToString::to_string,
|
||||
)
|
||||
}
|
||||
|
||||
fn output_text_from_item(item: &Value) -> Option<String> {
|
||||
@@ -1077,8 +1072,7 @@ fn tool_call_from_done_event(
|
||||
.and_then(Value::as_str)?;
|
||||
let info = function_call_items.get(event_id);
|
||||
let call_id = info
|
||||
.map(|info| info.call_id.as_str())
|
||||
.unwrap_or(event_id)
|
||||
.map_or(event_id, |info| info.call_id.as_str())
|
||||
.to_string();
|
||||
if !finished_call_ids.insert(call_id.clone()) {
|
||||
return None;
|
||||
@@ -1093,7 +1087,7 @@ fn tool_call_from_done_event(
|
||||
.get("arguments")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
Some(openai_tool_call(call_id, name, arguments))
|
||||
Some(openai_tool_call(&call_id, name, arguments))
|
||||
}
|
||||
|
||||
fn tool_call_from_output_item(
|
||||
@@ -1125,10 +1119,10 @@ fn tool_call_from_output_item(
|
||||
.get("arguments")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
Some(openai_tool_call(call_id.to_string(), name, arguments))
|
||||
Some(openai_tool_call(call_id, name, arguments))
|
||||
}
|
||||
|
||||
fn openai_tool_call(call_id: String, name: &str, arguments: &str) -> Value {
|
||||
fn openai_tool_call(call_id: &str, name: &str, arguments: &str) -> Value {
|
||||
json!({
|
||||
"id": call_id,
|
||||
"type": "function",
|
||||
|
||||
@@ -62,8 +62,8 @@ pub fn openai_error_response(
|
||||
)
|
||||
}
|
||||
|
||||
pub fn chat_dispatch_error_response(error: ChatDispatchError) -> Response {
|
||||
let error = chat_dispatch_openai_error(&error);
|
||||
pub fn chat_dispatch_error_response(error: &ChatDispatchError) -> Response {
|
||||
let error = chat_dispatch_openai_error(error);
|
||||
openai_error_response(error.status, &error.message, error.error_type, error.code)
|
||||
.into_response()
|
||||
}
|
||||
@@ -82,7 +82,7 @@ pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Respon
|
||||
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
|
||||
responses_no_available_accounts_response().into_response()
|
||||
}
|
||||
error => response_dispatch_openai_error_response(error),
|
||||
error => response_dispatch_openai_error_response(&error),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -140,8 +140,8 @@ fn responses_error_response(status: StatusCode, message: &str) -> (StatusCode, J
|
||||
)
|
||||
}
|
||||
|
||||
fn response_dispatch_openai_error_response(error: ResponseDispatchError) -> Response {
|
||||
let error = response_dispatch_openai_error(&error);
|
||||
fn response_dispatch_openai_error_response(error: &ResponseDispatchError) -> Response {
|
||||
let error = response_dispatch_openai_error(error);
|
||||
openai_error_response(error.status, &error.message, error.error_type, error.code)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
@@ -6,7 +6,6 @@ use axum::{
|
||||
response::IntoResponse,
|
||||
Json,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::{proxy::auth::authorize_client_api_key, runtime::state::AppState};
|
||||
@@ -15,26 +14,6 @@ use super::errors::{missing_client_api_key_response, model_not_found_response};
|
||||
|
||||
const MODEL_CREATED_TIMESTAMP: i64 = 1_700_000_000;
|
||||
|
||||
// ====================================================================
|
||||
// OpenAI 模型类型
|
||||
// ====================================================================
|
||||
|
||||
/// OpenAI 模型对象。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct OpenAiModel {
|
||||
pub id: String,
|
||||
pub object: String,
|
||||
pub created: i64,
|
||||
pub owned_by: String,
|
||||
}
|
||||
|
||||
/// OpenAI 模型列表。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct OpenAiModelList {
|
||||
pub object: String,
|
||||
pub data: Vec<OpenAiModel>,
|
||||
}
|
||||
|
||||
// ====================================================================
|
||||
// HTTP 处理器
|
||||
// ====================================================================
|
||||
|
||||
@@ -40,14 +40,6 @@ const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent";
|
||||
// 协议类型
|
||||
// ====================================================================
|
||||
|
||||
/// OpenAI 请求声明的响应格式类别。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ResponseFormat {
|
||||
Text,
|
||||
JsonObject,
|
||||
JsonSchema,
|
||||
}
|
||||
|
||||
/// Responses API 请求体。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct OpenAiResponsesRequest {
|
||||
@@ -242,7 +234,7 @@ fn sanitize_responses_input(input: Value) -> Vec<Value> {
|
||||
match input {
|
||||
Value::Array(items) => sanitize_codex_input_items(items),
|
||||
Value::Null => Vec::new(),
|
||||
Value::String(text) => vec![responses_input_text_message(text)],
|
||||
Value::String(text) => vec![responses_input_text_message(&text)],
|
||||
value => vec![value],
|
||||
}
|
||||
}
|
||||
@@ -252,7 +244,7 @@ fn sanitize_codex_input_items(input: Vec<Value>) -> Vec<Value> {
|
||||
.into_iter()
|
||||
.filter_map(|item| {
|
||||
if let Value::String(text) = item {
|
||||
return Some(responses_input_text_message(text));
|
||||
return Some(responses_input_text_message(&text));
|
||||
}
|
||||
let Value::Object(object) = item else {
|
||||
return Some(item);
|
||||
@@ -266,7 +258,7 @@ fn sanitize_codex_input_items(input: Vec<Value>) -> Vec<Value> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn responses_input_text_message(text: String) -> Value {
|
||||
fn responses_input_text_message(text: &str) -> Value {
|
||||
json!({
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
@@ -476,8 +468,10 @@ pub fn response_failed_sse_event_with_id(
|
||||
});
|
||||
let response_id = response_id
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.map(ToString::to_string)
|
||||
.unwrap_or_else(|| format!("resp_proxy_{}", uuid::Uuid::new_v4().simple()));
|
||||
.map_or_else(
|
||||
|| format!("resp_proxy_{}", uuid::Uuid::new_v4().simple()),
|
||||
ToString::to_string,
|
||||
);
|
||||
let data = json!({
|
||||
"type": "response.failed",
|
||||
"response": {
|
||||
@@ -576,7 +570,7 @@ async fn handle_responses(
|
||||
.await
|
||||
{
|
||||
Ok(stream) => live_event_stream_response(stream),
|
||||
Err(error) => response_dispatch_stream_error_response(error),
|
||||
Err(error) => response_dispatch_stream_error_response(&error),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -724,8 +718,8 @@ fn live_event_stream_response(stream: ResponseDispatchStream) -> Response {
|
||||
sse_event_stream_response(Body::from_stream(stream.body), SseResponseOptions::BASIC)
|
||||
}
|
||||
|
||||
fn response_dispatch_stream_error_response(error: ResponseDispatchError) -> Response {
|
||||
event_stream_response(responses_stream_dispatch_failed_sse_event(&error))
|
||||
fn response_dispatch_stream_error_response(error: &ResponseDispatchError) -> Response {
|
||||
event_stream_response(responses_stream_dispatch_failed_sse_event(error))
|
||||
}
|
||||
|
||||
fn ensure_done_sse_frame(body: &mut String) {
|
||||
|
||||
@@ -105,5 +105,5 @@ fn init_logging(config: &AppConfig) -> Result<Option<LogGuard>, crate::infra::lo
|
||||
}
|
||||
|
||||
let rotation = RotationConfig::new(&config.logging.directory, config.logging.retention_days);
|
||||
Ok(Some(init_tracing(rotation)?))
|
||||
Ok(Some(init_tracing(&rotation)?))
|
||||
}
|
||||
|
||||
+48
-40
@@ -6,13 +6,13 @@ use crate::{
|
||||
admin::monitoring::{
|
||||
event_store::{
|
||||
AdminLogService, SqliteEventLogStore, DEFAULT_EVENT_LOG_CAPACITY,
|
||||
DEFAULT_EVENT_LOG_CAPTURE_BODY, DEFAULT_EVENT_LOG_ENABLED,
|
||||
DEFAULT_EVENT_LOG_CAPTURE_BODY,
|
||||
},
|
||||
service::AdminUsageService,
|
||||
usage_store::SqliteUsageStore,
|
||||
},
|
||||
admin::{
|
||||
accounts::service::AdminAccountService,
|
||||
accounts::service::{AdminAccountService, AdminAccountServiceParts},
|
||||
auth::service::{AdminSessionService, SqliteAdminSessionStore},
|
||||
keys::service::{AdminClientKeyService, ClientKeyService, SqliteClientKeyStore},
|
||||
},
|
||||
@@ -33,10 +33,10 @@ use crate::{
|
||||
},
|
||||
upstream::{
|
||||
fingerprint::{Fingerprint, FingerprintRepository},
|
||||
models::{AdminModelService, ModelService, ModelSnapshotStore},
|
||||
models::{ModelService, ModelSnapshotStore},
|
||||
token_client::{default_openai_token_client, TokenClientConfig},
|
||||
transport::{
|
||||
build_reqwest_client, CodexBackendClient, CodexModelCatalogClient, CustomCaError,
|
||||
build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient,
|
||||
},
|
||||
},
|
||||
};
|
||||
@@ -66,7 +66,6 @@ pub struct BackgroundTaskStores {
|
||||
#[derive(Clone)]
|
||||
pub struct Services {
|
||||
pub models: StdArc<ModelService>,
|
||||
pub admin_models: StdArc<AdminModelService>,
|
||||
pub accounts: StdArc<dyn AccountStoreTrait>,
|
||||
pub client_keys: StdArc<ClientKeyService>,
|
||||
pub admin_client_keys: StdArc<AdminClientKeyService>,
|
||||
@@ -85,6 +84,24 @@ pub struct Services {
|
||||
pub background_tasks: BackgroundTaskStores,
|
||||
}
|
||||
|
||||
/// 事件日志运行选项。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct EventLogOptions {
|
||||
pub enabled: bool,
|
||||
pub capacity: u32,
|
||||
pub capture_body: bool,
|
||||
}
|
||||
|
||||
impl EventLogOptions {
|
||||
pub fn from_config(config: &AppConfig) -> Self {
|
||||
Self {
|
||||
enabled: config.logging.enabled,
|
||||
capacity: DEFAULT_EVENT_LOG_CAPACITY,
|
||||
capture_body: DEFAULT_EVENT_LOG_CAPTURE_BODY,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Services {
|
||||
pub fn new(config: &AppConfig, stores: BackgroundTaskStores, fingerprint: Fingerprint) -> Self {
|
||||
Self::try_new(config, stores, fingerprint)
|
||||
@@ -99,21 +116,27 @@ impl Services {
|
||||
Self::try_with_installation_id(config, stores, fingerprint, None)
|
||||
}
|
||||
|
||||
pub fn with_installation_id(
|
||||
pub(crate) fn try_with_installation_id(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: Fingerprint,
|
||||
installation_id: Option<String>,
|
||||
) -> Self {
|
||||
Self::try_with_installation_id(config, stores, fingerprint, installation_id)
|
||||
.expect("failed to build runtime services with configured TLS transport")
|
||||
) -> Result<Self, CustomCaError> {
|
||||
Self::try_with_installation_id_and_log_options(
|
||||
config,
|
||||
stores,
|
||||
fingerprint,
|
||||
installation_id,
|
||||
EventLogOptions::from_config(config),
|
||||
)
|
||||
}
|
||||
|
||||
pub fn try_with_installation_id(
|
||||
pub fn try_with_installation_id_and_log_options(
|
||||
config: &AppConfig,
|
||||
stores: BackgroundTaskStores,
|
||||
fingerprint: Fingerprint,
|
||||
installation_id: Option<String>,
|
||||
log_options: EventLogOptions,
|
||||
) -> Result<Self, CustomCaError> {
|
||||
let installation_id = installation_id.filter(|id| !id.trim().is_empty());
|
||||
let account_store_trait =
|
||||
@@ -159,26 +182,25 @@ impl Services {
|
||||
stores.client_keys.clone(),
|
||||
)));
|
||||
let token_client = StdArc::new(default_openai_token_client(token_client_config(config)));
|
||||
let admin_accounts = StdArc::new(AdminAccountService::new(
|
||||
stores.accounts.clone(),
|
||||
stores.cookies.clone(),
|
||||
config.quota.warning_thresholds.clone(),
|
||||
codex.clone(),
|
||||
account_pool.clone(),
|
||||
token_client,
|
||||
crate::admin::accounts::service::oauth::AccountOAuthService::new(
|
||||
let admin_accounts = StdArc::new(AdminAccountService::new(AdminAccountServiceParts {
|
||||
store: stores.accounts.clone(),
|
||||
cookies: stores.cookies.clone(),
|
||||
codex: codex.clone(),
|
||||
account_pool: account_pool.clone(),
|
||||
token_refresher: token_client,
|
||||
oauth: crate::admin::accounts::service::oauth::AccountOAuthService::new(
|
||||
reqwest::Client::new(),
|
||||
config.auth.oauth_client_id.clone(),
|
||||
config.auth.oauth_token_endpoint.clone(),
|
||||
),
|
||||
config.auth.refresh_margin_seconds,
|
||||
installation_id.clone(),
|
||||
));
|
||||
refresh_margin_seconds: config.auth.refresh_margin_seconds,
|
||||
installation_id: installation_id.clone(),
|
||||
}));
|
||||
let logs = StdArc::new(AdminLogService::new(
|
||||
stores.event_logs.clone(),
|
||||
DEFAULT_EVENT_LOG_ENABLED,
|
||||
DEFAULT_EVENT_LOG_CAPACITY,
|
||||
DEFAULT_EVENT_LOG_CAPTURE_BODY,
|
||||
log_options.enabled,
|
||||
log_options.capacity,
|
||||
log_options.capture_body,
|
||||
));
|
||||
let usage_store = SqliteUsageStore::new(stores.accounts.pool().clone());
|
||||
let usage = StdArc::new(AdminUsageService::new(usage_store));
|
||||
@@ -194,21 +216,8 @@ impl Services {
|
||||
crate::upstream::models::ModelConfig {
|
||||
model_aliases: config.model_aliases.clone(),
|
||||
},
|
||||
Some(snapshot_store.clone()),
|
||||
Some(upstream_client.clone()),
|
||||
None,
|
||||
));
|
||||
let admin_models = StdArc::new(AdminModelService::new(
|
||||
StdArc::new(ModelService::new(
|
||||
crate::upstream::models::ModelConfig {
|
||||
model_aliases: config.model_aliases.clone(),
|
||||
},
|
||||
Some(snapshot_store),
|
||||
Some(upstream_client),
|
||||
None,
|
||||
)),
|
||||
account_store_trait.clone(),
|
||||
installation_id.clone(),
|
||||
Some(snapshot_store),
|
||||
Some(upstream_client),
|
||||
));
|
||||
let cloudflare_recovery =
|
||||
crate::proxy::dispatch::cloudflare::CloudflareRecovery::new(stores.cookies.clone());
|
||||
@@ -232,7 +241,6 @@ impl Services {
|
||||
|
||||
Ok(Self {
|
||||
models,
|
||||
admin_models,
|
||||
accounts: account_store_trait,
|
||||
client_keys,
|
||||
admin_client_keys,
|
||||
|
||||
@@ -22,7 +22,7 @@ pub async fn shutdown_signal() {
|
||||
let terminate = std::future::pending::<()>();
|
||||
|
||||
tokio::select! {
|
||||
_ = ctrl_c => {},
|
||||
_ = terminate => {},
|
||||
() = ctrl_c => {},
|
||||
() = terminate => {},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -27,14 +27,6 @@ impl CookieCleanupTask {
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用自定义间隔构造 Cookie 清理任务。
|
||||
pub fn with_interval_secs(store: SqliteCookieStore, interval_secs: u64) -> Self {
|
||||
Self {
|
||||
store,
|
||||
interval_secs,
|
||||
}
|
||||
}
|
||||
|
||||
/// 启动后台清理任务。
|
||||
pub fn start(self) -> SchedulerHandle {
|
||||
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
|
||||
|
||||
@@ -30,17 +30,12 @@ pub struct TaskCoordinator {
|
||||
|
||||
impl TaskCoordinator {
|
||||
/// 注册一个后台任务句柄。
|
||||
pub fn push(&mut self, name: &'static str, handle: SchedulerHandle) {
|
||||
pub(crate) fn push(&mut self, name: &'static str, handle: SchedulerHandle) {
|
||||
self.handles.push((name, handle));
|
||||
}
|
||||
|
||||
/// 返回已注册任务名称。
|
||||
pub fn task_names(&self) -> Vec<&'static str> {
|
||||
self.handles.iter().map(|(name, _)| *name).collect()
|
||||
}
|
||||
|
||||
/// 启动所有后台任务。
|
||||
pub fn start(
|
||||
pub(crate) fn start(
|
||||
config: &crate::runtime::state::RuntimeConfig,
|
||||
services: &crate::runtime::services::Services,
|
||||
) -> Self {
|
||||
@@ -69,9 +64,13 @@ impl TaskCoordinator {
|
||||
);
|
||||
coordinator.push(
|
||||
"model_refresh",
|
||||
ModelRefreshTask::new(services.models.clone(), services.accounts.clone())
|
||||
.with_installation_id(services.installation_id.clone())
|
||||
.start(),
|
||||
ModelRefreshTask::new(
|
||||
services.models.clone(),
|
||||
services.accounts.clone(),
|
||||
services.account_pool.clone(),
|
||||
)
|
||||
.with_installation_id(services.installation_id.clone())
|
||||
.start(),
|
||||
);
|
||||
if config.auth.refresh_enabled {
|
||||
coordinator.push(
|
||||
@@ -146,12 +145,12 @@ pub enum SchedulerHandle {
|
||||
|
||||
impl SchedulerHandle {
|
||||
/// 使用关闭 channel 构造句柄。
|
||||
pub fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self {
|
||||
pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self {
|
||||
Self::Channel(shutdown_tx)
|
||||
}
|
||||
|
||||
/// 使用 `JoinHandle` 构造句柄。
|
||||
pub fn from_join_handle(handle: JoinHandle<()>) -> Self {
|
||||
pub(crate) fn from_join_handle(handle: JoinHandle<()>) -> Self {
|
||||
Self::JoinHandle(handle)
|
||||
}
|
||||
|
||||
|
||||
@@ -8,5 +8,3 @@ pub mod quota_refresh;
|
||||
pub mod session_affinity_cleanup;
|
||||
pub mod session_cleanup;
|
||||
pub mod token_refresh;
|
||||
|
||||
pub use coordinator::{SchedulerHandle, TaskCoordinator};
|
||||
|
||||
@@ -5,8 +5,9 @@ use std::{sync::Arc, time::Duration};
|
||||
use tokio::time::{interval, sleep};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::upstream::accounts::model::Account;
|
||||
use crate::upstream::accounts::store::AccountStore;
|
||||
use crate::upstream::accounts::{
|
||||
model::Account, pool::RuntimeAccountPoolService, store::AccountStore,
|
||||
};
|
||||
use crate::upstream::models::{ModelRefreshResult, ModelService, ModelServiceError};
|
||||
|
||||
use super::coordinator::SchedulerHandle;
|
||||
@@ -15,6 +16,7 @@ use super::coordinator::SchedulerHandle;
|
||||
pub struct ModelRefreshTask {
|
||||
model_service: Arc<ModelService>,
|
||||
account_store: Arc<dyn AccountStore>,
|
||||
account_pool: Arc<RuntimeAccountPoolService>,
|
||||
refresh_interval_secs: u64,
|
||||
initial_delay_ms: u64,
|
||||
retry_delay_ms: u64,
|
||||
@@ -29,10 +31,15 @@ const MAX_RETRIES: u32 = 12;
|
||||
|
||||
impl ModelRefreshTask {
|
||||
/// 构造默认任务。
|
||||
pub fn new(model_service: Arc<ModelService>, account_store: Arc<dyn AccountStore>) -> Self {
|
||||
pub fn new(
|
||||
model_service: Arc<ModelService>,
|
||||
account_store: Arc<dyn AccountStore>,
|
||||
account_pool: Arc<RuntimeAccountPoolService>,
|
||||
) -> Self {
|
||||
Self {
|
||||
model_service,
|
||||
account_store,
|
||||
account_pool,
|
||||
refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS,
|
||||
initial_delay_ms: INITIAL_DELAY_MS,
|
||||
retry_delay_ms: RETRY_DELAY_MS,
|
||||
@@ -65,7 +72,7 @@ impl ModelRefreshTask {
|
||||
info!("模型刷新任务在首次拉取前关闭");
|
||||
return;
|
||||
}
|
||||
_ = async {} => {}
|
||||
() = async {} => {}
|
||||
}
|
||||
|
||||
match self.refresh_once().await {
|
||||
@@ -98,7 +105,7 @@ impl ModelRefreshTask {
|
||||
info!("模型刷新任务在重试等待期间关闭");
|
||||
return;
|
||||
}
|
||||
_ = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
|
||||
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -130,16 +137,23 @@ impl ModelRefreshTask {
|
||||
SchedulerHandle::new(shutdown_tx)
|
||||
}
|
||||
|
||||
async fn refresh_once(&self) -> Result<ModelRefreshResult, ModelServiceError> {
|
||||
/// 执行一次模型刷新并同步账号池模型计划约束。
|
||||
pub async fn refresh_once(&self) -> Result<ModelRefreshResult, ModelServiceError> {
|
||||
let request_id = uuid::Uuid::new_v4().to_string();
|
||||
let accounts = self.load_accounts().await?;
|
||||
self.model_service
|
||||
let result = self
|
||||
.model_service
|
||||
.refresh_backend_models_with_installation_id(
|
||||
&accounts,
|
||||
&request_id,
|
||||
self.installation_id.as_deref(),
|
||||
)
|
||||
.await
|
||||
.await?;
|
||||
let allowlist = self.model_service.model_plan_allowlist().await?;
|
||||
self.account_pool
|
||||
.apply_model_plan_allowlist(allowlist)
|
||||
.await;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn load_accounts(&self) -> Result<Vec<Account>, ModelServiceError> {
|
||||
|
||||
@@ -16,8 +16,6 @@ use crate::{
|
||||
|
||||
use super::coordinator::SchedulerHandle;
|
||||
|
||||
const DEFAULT_INTERVAL_SECS: u64 = 15 * 60;
|
||||
|
||||
/// 主动配额刷新后台任务。
|
||||
pub struct QuotaRefreshTask {
|
||||
service: RuntimeQuotaRefreshService,
|
||||
@@ -25,14 +23,6 @@ pub struct QuotaRefreshTask {
|
||||
}
|
||||
|
||||
impl QuotaRefreshTask {
|
||||
/// 构造默认配额刷新后台任务。
|
||||
pub fn new(store: SqliteAccountStore, codex: Arc<CodexBackendClient>) -> Self {
|
||||
Self {
|
||||
service: RuntimeQuotaRefreshService::new(store, codex),
|
||||
interval_secs: DEFAULT_INTERVAL_SECS,
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用自定义间隔构造配额刷新后台任务。
|
||||
pub fn with_intervals(
|
||||
store: SqliteAccountStore,
|
||||
|
||||
@@ -35,19 +35,6 @@ where
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用自定义间隔构造后台任务。
|
||||
pub fn with_interval_secs(
|
||||
store: SqliteAccountStore,
|
||||
policy: RefreshPolicy,
|
||||
client: C,
|
||||
interval_secs: u64,
|
||||
) -> Self {
|
||||
Self {
|
||||
service: RuntimeTokenRefreshService::new(store, policy, client),
|
||||
interval_secs,
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用刷新租约存储保护账号刷新。
|
||||
pub fn with_refresh_lease_store(mut self, refresh_leases: RefreshLeaseStore) -> Self {
|
||||
self.service = self.service.with_refresh_lease_store(refresh_leases);
|
||||
|
||||
@@ -54,20 +54,6 @@ impl SqliteCookieStore {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// 返回底层连接池。
|
||||
pub fn pool(&self) -> &SqlitePool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
/// 检查账号是否存在。
|
||||
pub async fn account_exists(&self, account_id: &str) -> SqliteCookieStoreResult<bool> {
|
||||
let row = sqlx::query("select 1 from accounts where id = ?")
|
||||
.bind(account_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
Ok(row.is_some())
|
||||
}
|
||||
|
||||
/// 捕获上游 `Set-Cookie` 响应头中允许持久化的 Cookie。
|
||||
pub async fn capture_set_cookie(
|
||||
&self,
|
||||
@@ -83,20 +69,6 @@ impl SqliteCookieStore {
|
||||
self.upsert_cookie(account_id, parsed).await
|
||||
}
|
||||
|
||||
/// 将 Cookie 请求头写入账号 Cookie 存储。
|
||||
pub async fn set_cookie_header(
|
||||
&self,
|
||||
account_id: &str,
|
||||
raw: &str,
|
||||
) -> SqliteCookieStoreResult<usize> {
|
||||
let parsed = parse_cookie_header(raw);
|
||||
let count = parsed.len();
|
||||
for cookie in parsed {
|
||||
self.upsert_cookie(account_id, cookie).await?;
|
||||
}
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// 为请求域名读取账号 Cookie 请求头。
|
||||
pub async fn cookie_header(
|
||||
&self,
|
||||
@@ -267,27 +239,6 @@ fn max_age_expires_at(seconds: i64) -> String {
|
||||
(now + Duration::seconds(seconds.min(i32::MAX as i64))).to_rfc3339()
|
||||
}
|
||||
|
||||
fn parse_cookie_header(raw: &str) -> Vec<ParsedCookie> {
|
||||
raw.split(';')
|
||||
.map(str::trim)
|
||||
.filter_map(|part| {
|
||||
let (name, value) = part.split_once('=')?;
|
||||
let name = name.trim();
|
||||
let value = value.trim();
|
||||
if name.is_empty() || value.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some(ParsedCookie {
|
||||
domain: DEFAULT_COOKIE_DOMAIN.to_string(),
|
||||
name: name.to_string(),
|
||||
value: value.to_string(),
|
||||
path: "/".to_string(),
|
||||
expires_at: None,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn domain_matches(request_domain: &str, cookie_domain: &str) -> bool {
|
||||
request_domain == cookie_domain
|
||||
|| request_domain
|
||||
@@ -386,8 +337,7 @@ impl CloudflarePathBlockTracker {
|
||||
let count = counts
|
||||
.get(account_id)
|
||||
.filter(|state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER)
|
||||
.map(|state| state.count.saturating_add(1))
|
||||
.unwrap_or(1);
|
||||
.map_or(1, |state| state.count.saturating_add(1));
|
||||
counts.insert(
|
||||
account_id.to_string(),
|
||||
PathBlockState {
|
||||
@@ -442,8 +392,7 @@ impl CloudflareChallengeCooldownTracker {
|
||||
let challenge_count = states
|
||||
.get(account_id)
|
||||
.filter(|state| now.signed_duration_since(state.updated_at) <= CHALLENGE_STALE_AFTER)
|
||||
.map(|state| state.challenge_count.saturating_add(1))
|
||||
.unwrap_or(1);
|
||||
.map_or(1, |state| state.challenge_count.saturating_add(1));
|
||||
states.insert(
|
||||
account_id.to_string(),
|
||||
ChallengeCooldownState {
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//! 账号导入导出逻辑。
|
||||
//! 账号导入逻辑。
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde_json::Value;
|
||||
@@ -29,7 +29,7 @@ pub struct AccountImportEntry {
|
||||
pub access_token_expires_at: Option<String>,
|
||||
/// 导入时的原始状态。
|
||||
pub status: Option<String>,
|
||||
/// Sub2api 导入时的缓存配额 JSON。
|
||||
/// Sub2API 导入时的缓存配额 JSON。
|
||||
pub cached_quota: Option<Value>,
|
||||
/// 配额抓取时间。
|
||||
pub quota_fetched_at: Option<String>,
|
||||
@@ -40,18 +40,18 @@ pub struct AccountImportEntry {
|
||||
/// 账号导入来源。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AccountImportSource {
|
||||
/// 原生格式。
|
||||
Native,
|
||||
/// Sub2api 格式。
|
||||
Sub2api,
|
||||
/// CPR 格式。
|
||||
Cpr,
|
||||
/// Sub2API 格式。
|
||||
Sub2Api,
|
||||
}
|
||||
|
||||
impl AccountImportSource {
|
||||
/// 返回来源的字符串标识。
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Native => "native",
|
||||
Self::Sub2api => "sub2api",
|
||||
Self::Cpr => "cpr",
|
||||
Self::Sub2Api => "sub2api",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -179,8 +179,8 @@ fn account_import_entry_from_value(
|
||||
return Ok(None);
|
||||
};
|
||||
let allowed_keys = match source {
|
||||
AccountImportSource::Native => ACCOUNT_IMPORT_ACCOUNT_KEYS,
|
||||
AccountImportSource::Sub2api => SUB2API_ACCOUNT_IMPORT_KEYS,
|
||||
AccountImportSource::Cpr => ACCOUNT_IMPORT_ACCOUNT_KEYS,
|
||||
AccountImportSource::Sub2Api => SUB2API_ACCOUNT_IMPORT_KEYS,
|
||||
};
|
||||
if account
|
||||
.keys()
|
||||
@@ -209,13 +209,13 @@ fn account_import_entry_from_value(
|
||||
&["accessTokenExpiresAt", "access_token_expires_at"],
|
||||
),
|
||||
status: first_string(value, &["status"]),
|
||||
cached_quota: (source == AccountImportSource::Sub2api)
|
||||
cached_quota: (source == AccountImportSource::Sub2Api)
|
||||
.then(|| first_value(value, &["cachedQuota", "cached_quota"]))
|
||||
.flatten(),
|
||||
quota_fetched_at: (source == AccountImportSource::Sub2api)
|
||||
quota_fetched_at: (source == AccountImportSource::Sub2Api)
|
||||
.then(|| first_string(value, &["quotaFetchedAt", "quota_fetched_at"]))
|
||||
.flatten(),
|
||||
quota_verify_required: (source == AccountImportSource::Sub2api)
|
||||
quota_verify_required: (source == AccountImportSource::Sub2Api)
|
||||
.then(|| first_bool(value, &["quotaVerifyRequired", "quota_verify_required"]))
|
||||
.flatten(),
|
||||
}))
|
||||
@@ -241,16 +241,16 @@ fn account_import_source(
|
||||
) -> Result<AccountImportSource, &'static str> {
|
||||
if let Some(source_format) = first_string(value, &["sourceFormat", "source_format"]) {
|
||||
return match source_format.trim().to_ascii_lowercase().as_str() {
|
||||
"native" => Ok(AccountImportSource::Native),
|
||||
"sub2api" => Ok(AccountImportSource::Sub2api),
|
||||
"cpr" => Ok(AccountImportSource::Cpr),
|
||||
"sub2api" => Ok(AccountImportSource::Sub2Api),
|
||||
_ => Err("no importable accounts"),
|
||||
};
|
||||
}
|
||||
|
||||
if accounts.iter().any(account_import_entry_looks_sub2api) {
|
||||
Ok(AccountImportSource::Sub2api)
|
||||
Ok(AccountImportSource::Sub2Api)
|
||||
} else {
|
||||
Ok(AccountImportSource::Native)
|
||||
Ok(AccountImportSource::Cpr)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -284,7 +284,7 @@ pub fn normalized_imported_account_status(
|
||||
source: AccountImportSource,
|
||||
access_token: &str,
|
||||
) -> AccountStatus {
|
||||
if source == AccountImportSource::Sub2api
|
||||
if source == AccountImportSource::Sub2Api
|
||||
&& status == AccountStatus::Active
|
||||
&& jwt_expiry(access_token, Utc::now()) != JwtExpiry::Valid
|
||||
{
|
||||
@@ -314,46 +314,13 @@ pub fn parse_account_status(status: &str) -> Result<AccountStatus, &'static str>
|
||||
}
|
||||
}
|
||||
|
||||
/// 对管理端修改用的状态做严格校验。
|
||||
pub fn parse_batch_account_status(status: &str) -> Result<AccountStatus, &'static str> {
|
||||
match status.trim().to_ascii_lowercase().as_str() {
|
||||
"active" => Ok(AccountStatus::Active),
|
||||
"disabled" => Ok(AccountStatus::Disabled),
|
||||
_ => Err("invalid account status (only active/disabled allowed)"),
|
||||
}
|
||||
}
|
||||
|
||||
/// 将 RefreshFailure 映射为账号状态。
|
||||
pub fn refresh_failure_status(
|
||||
failure: &crate::upstream::accounts::token_refresh::RefreshFailure,
|
||||
) -> AccountStatus {
|
||||
match failure {
|
||||
crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant => {
|
||||
AccountStatus::Expired
|
||||
}
|
||||
crate::upstream::accounts::token_refresh::RefreshFailure::QuotaExhausted => {
|
||||
AccountStatus::QuotaExhausted
|
||||
}
|
||||
crate::upstream::accounts::token_refresh::RefreshFailure::Banned => AccountStatus::Banned,
|
||||
crate::upstream::accounts::token_refresh::RefreshFailure::Disabled => {
|
||||
AccountStatus::Disabled
|
||||
}
|
||||
crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport => {
|
||||
AccountStatus::Active
|
||||
}
|
||||
crate::upstream::accounts::token_refresh::RefreshFailure::Transport => {
|
||||
AccountStatus::Active
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 判断是否需要清除 next_refresh_at。
|
||||
pub fn refresh_failure_status_clears_next_refresh_at(status: AccountStatus) -> bool {
|
||||
!matches!(status, AccountStatus::Active)
|
||||
}
|
||||
|
||||
/// 标准化 Bearer token(去除 Bearer 前缀,去除空白)。
|
||||
pub fn normalize_bearer_token(value: String) -> String {
|
||||
pub fn normalize_bearer_token(value: &str) -> String {
|
||||
value
|
||||
.trim()
|
||||
.strip_prefix("Bearer ")
|
||||
@@ -385,7 +352,7 @@ pub fn normalize_nonempty_str(value: Option<&str>) -> Option<&str> {
|
||||
value.map(str::trim).filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
pub fn first_string(value: &Value, keys: &[&str]) -> Option<String> {
|
||||
fn first_string(value: &Value, keys: &[&str]) -> Option<String> {
|
||||
keys.iter()
|
||||
.find_map(|key| value.get(key).and_then(Value::as_str))
|
||||
.map(str::trim)
|
||||
@@ -1,9 +1,9 @@
|
||||
//! 账号领域模块 —— 模型、仓库、池调度、Cookie、配额、刷新与导入导出。
|
||||
//! 账号领域模块 —— 模型、仓库、池调度、Cookie、配额、刷新与导入。
|
||||
|
||||
/// Cookie 捕获与重放策略。
|
||||
pub mod cookies;
|
||||
/// 账号导入导出逻辑。
|
||||
pub mod import_export;
|
||||
/// 账号导入逻辑。
|
||||
pub mod importing;
|
||||
/// 账号领域模型。
|
||||
pub mod model;
|
||||
/// 账号池调度策略。
|
||||
|
||||
@@ -106,7 +106,7 @@ pub struct Account {
|
||||
pub last_used_at: Option<String>,
|
||||
}
|
||||
|
||||
/// 账号用量增量(来自 usage.rs)。
|
||||
/// 账号用量增量。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct AccountUsageDelta {
|
||||
/// 请求数增量。
|
||||
@@ -133,25 +133,6 @@ pub struct AccountUsageDelta {
|
||||
pub image_request_failures: u64,
|
||||
}
|
||||
|
||||
impl AccountUsageDelta {
|
||||
/// 合并两个用量增量。
|
||||
pub fn merged(self, other: Self) -> Self {
|
||||
Self {
|
||||
requests: self.requests + other.requests,
|
||||
input_tokens: self.input_tokens + other.input_tokens,
|
||||
output_tokens: self.output_tokens + other.output_tokens,
|
||||
cached_tokens: self.cached_tokens + other.cached_tokens,
|
||||
reasoning_tokens: self.reasoning_tokens + other.reasoning_tokens,
|
||||
total_tokens: self.total_tokens + other.total_tokens,
|
||||
empty_responses: self.empty_responses + other.empty_responses,
|
||||
image_input_tokens: self.image_input_tokens + other.image_input_tokens,
|
||||
image_output_tokens: self.image_output_tokens + other.image_output_tokens,
|
||||
image_requests: self.image_requests + other.image_requests,
|
||||
image_request_failures: self.image_request_failures + other.image_request_failures,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号模型维度用量增量。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct AccountModelUsageDelta {
|
||||
|
||||
+24
-145
@@ -164,27 +164,6 @@ pub struct AccountWindowUsageDelta {
|
||||
pub image_request_failed: bool,
|
||||
}
|
||||
|
||||
/// 账号池状态摘要。
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct AccountPoolStatusSummary {
|
||||
/// 账号总数。
|
||||
pub total: usize,
|
||||
/// 可调度账号数。
|
||||
pub active: usize,
|
||||
/// token 过期账号数。
|
||||
pub expired: usize,
|
||||
/// 配额耗尽账号数。
|
||||
pub quota_exhausted: usize,
|
||||
/// 处于限流冷却的账号数。
|
||||
pub rate_limited: usize,
|
||||
/// 正在刷新账号数。
|
||||
pub refreshing: usize,
|
||||
/// 禁用账号数。
|
||||
pub disabled: usize,
|
||||
/// 封禁账号数。
|
||||
pub banned: usize,
|
||||
}
|
||||
|
||||
/// 纯内存账号池,负责账号调度和运行时状态维护。
|
||||
#[derive(Debug)]
|
||||
pub struct AccountPool {
|
||||
@@ -238,30 +217,11 @@ impl AccountPool {
|
||||
self.least_used_cursor = 0;
|
||||
}
|
||||
|
||||
/// 替换模型计划 allowlist。
|
||||
pub fn set_model_plan_allowlist(&mut self, allowlist: BTreeMap<String, Vec<String>>) {
|
||||
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 {
|
||||
return false;
|
||||
};
|
||||
account.label = label;
|
||||
true
|
||||
}
|
||||
|
||||
/// 更新账号状态。
|
||||
pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool {
|
||||
let Some(account) = self.accounts.get_mut(account_id) else {
|
||||
@@ -274,30 +234,6 @@ impl AccountPool {
|
||||
true
|
||||
}
|
||||
|
||||
/// 重置账号累计和窗口用量。
|
||||
pub fn reset_usage(&mut self, account_id: &str) -> bool {
|
||||
let Some(account) = self.accounts.get_mut(account_id) else {
|
||||
return false;
|
||||
};
|
||||
account.last_used_at = None;
|
||||
account.request_count = 0;
|
||||
account.empty_response_count = 0;
|
||||
account.image_input_tokens = 0;
|
||||
account.image_output_tokens = 0;
|
||||
account.image_request_count = 0;
|
||||
account.image_request_failed_count = 0;
|
||||
account.window_request_count = 0;
|
||||
account.window_input_tokens = 0;
|
||||
account.window_output_tokens = 0;
|
||||
account.window_cached_tokens = 0;
|
||||
account.window_image_input_tokens = 0;
|
||||
account.window_image_output_tokens = 0;
|
||||
account.window_image_request_count = 0;
|
||||
account.window_image_request_failed_count = 0;
|
||||
account.window_started_at = account.window_reset_at.map(|_| Utc::now());
|
||||
true
|
||||
}
|
||||
|
||||
/// 标记账号配额限流直到指定时间。
|
||||
pub fn mark_quota_limited_until(
|
||||
&mut self,
|
||||
@@ -327,15 +263,6 @@ impl AccountPool {
|
||||
true
|
||||
}
|
||||
|
||||
/// 设置账号是否需要额外配额校验。
|
||||
pub fn set_quota_verify_required(&mut self, account_id: &str, required: bool) -> bool {
|
||||
let Some(account) = self.accounts.get_mut(account_id) else {
|
||||
return false;
|
||||
};
|
||||
account.quota_verify_required = required;
|
||||
true
|
||||
}
|
||||
|
||||
/// 应用已经验证过的配额状态。
|
||||
pub fn apply_quota_state(
|
||||
&mut self,
|
||||
@@ -433,24 +360,18 @@ impl AccountPool {
|
||||
true
|
||||
}
|
||||
|
||||
/// 使用当前时间为指定模型获取账号。
|
||||
pub fn acquire(&mut self, model: &str) -> Option<Account> {
|
||||
self.acquire_with(AccountAcquireRequest::new(model, Utc::now()))
|
||||
.map(|acquired| acquired.account)
|
||||
}
|
||||
|
||||
/// 使用完整调度请求获取账号。
|
||||
pub fn acquire_with(&mut self, request: AccountAcquireRequest) -> Option<AcquiredAccount> {
|
||||
pub fn acquire_with(&mut self, request: &AccountAcquireRequest) -> Option<AcquiredAccount> {
|
||||
self.acquire_with_status_refresh(request).acquired
|
||||
}
|
||||
|
||||
fn acquire_with_status_refresh(
|
||||
&mut self,
|
||||
request: AccountAcquireRequest,
|
||||
request: &AccountAcquireRequest,
|
||||
) -> AccountAcquireWithStatusRefresh {
|
||||
self.cleanup_stale_slots(request.now);
|
||||
let expired_account_ids = self.refresh_account_statuses(request.now);
|
||||
let candidates = self.candidates(&request);
|
||||
let candidates = self.candidates(request);
|
||||
let selected = if let Some(preferred_account_id) = &request.preferred_account_id {
|
||||
let preferred = candidates
|
||||
.iter()
|
||||
@@ -473,7 +394,7 @@ impl AccountPool {
|
||||
(account, previous)
|
||||
}
|
||||
RotationStrategy::Sticky => {
|
||||
let account = self.select_sticky(&candidates)?;
|
||||
let account = Self::select_sticky(&candidates)?;
|
||||
let previous = self.previous_slot_at(&account.id);
|
||||
(account, previous)
|
||||
}
|
||||
@@ -562,47 +483,6 @@ impl AccountPool {
|
||||
}
|
||||
}
|
||||
|
||||
/// 计算账号池状态摘要。
|
||||
pub fn status_summary(&mut self, now: DateTime<Utc>) -> AccountPoolStatusSummary {
|
||||
self.cleanup_stale_slots(now);
|
||||
self.refresh_account_statuses(now);
|
||||
let mut summary = AccountPoolStatusSummary {
|
||||
total: self.accounts.len(),
|
||||
..AccountPoolStatusSummary::default()
|
||||
};
|
||||
for account in self.accounts.values() {
|
||||
match account.status {
|
||||
AccountStatus::Active
|
||||
if !AccountService::quota_available_at(
|
||||
account,
|
||||
now,
|
||||
self.options.skip_quota_limited,
|
||||
) =>
|
||||
{
|
||||
summary.rate_limited += 1;
|
||||
}
|
||||
AccountStatus::Active => summary.active += 1,
|
||||
AccountStatus::Expired => summary.expired += 1,
|
||||
AccountStatus::QuotaExhausted => summary.quota_exhausted += 1,
|
||||
AccountStatus::Refreshing => summary.refreshing += 1,
|
||||
AccountStatus::Disabled => summary.disabled += 1,
|
||||
AccountStatus::Banned => summary.banned += 1,
|
||||
}
|
||||
}
|
||||
summary
|
||||
}
|
||||
|
||||
/// 获取所有处于配额锁定状态的账号 ID。
|
||||
pub fn list_quota_locked_accounts(&self) -> Vec<String> {
|
||||
self.accounts
|
||||
.values()
|
||||
.filter(|account| {
|
||||
account.status == AccountStatus::Active && account.quota_limit_reached
|
||||
})
|
||||
.map(|account| account.id.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn select_least_used(&mut self, candidates: &[Account]) -> Option<Account> {
|
||||
let best_key = candidates
|
||||
.iter()
|
||||
@@ -618,7 +498,7 @@ impl AccountPool {
|
||||
Some((*tied[index]).clone())
|
||||
}
|
||||
|
||||
fn select_sticky(&self, candidates: &[Account]) -> Option<Account> {
|
||||
fn select_sticky(candidates: &[Account]) -> Option<Account> {
|
||||
candidates
|
||||
.iter()
|
||||
.max_by_key(|account| account.last_used_at.clone())
|
||||
@@ -793,8 +673,10 @@ impl RuntimeAccountPoolService {
|
||||
}
|
||||
|
||||
/// 更新账号池运行参数。
|
||||
pub async fn apply_options(&self, options: AccountPoolOptions, request_interval_ms: u64) {
|
||||
self.pool.lock().await.set_options(options);
|
||||
pub async fn apply_options(&self, mut options: AccountPoolOptions, request_interval_ms: u64) {
|
||||
let mut pool = self.pool.lock().await;
|
||||
options.model_plan_allowlist = pool.options.model_plan_allowlist.clone();
|
||||
pool.set_options(options);
|
||||
*self
|
||||
.request_interval
|
||||
.write()
|
||||
@@ -802,6 +684,14 @@ impl RuntimeAccountPoolService {
|
||||
StdDuration::from_millis(request_interval_ms);
|
||||
}
|
||||
|
||||
/// 更新模型到可用账号计划的调度约束。
|
||||
pub(crate) async fn apply_model_plan_allowlist(
|
||||
&self,
|
||||
model_plan_allowlist: BTreeMap<String, Vec<String>>,
|
||||
) {
|
||||
self.pool.lock().await.options.model_plan_allowlist = model_plan_allowlist;
|
||||
}
|
||||
|
||||
/// 从持久化账号恢复运行时账号池。
|
||||
pub async fn restore_from_repository(&self) -> Result<usize, RuntimeAccountPoolError> {
|
||||
let accounts = self
|
||||
@@ -818,11 +708,6 @@ impl RuntimeAccountPoolService {
|
||||
Ok(count)
|
||||
}
|
||||
|
||||
/// 清空运行时账号池。
|
||||
pub async fn clear(&self) {
|
||||
self.pool.lock().await.clear();
|
||||
}
|
||||
|
||||
/// 读取账号池容量摘要。
|
||||
pub async fn capacity_summary(&self, now: DateTime<Utc>) -> AccountCapacitySummary {
|
||||
self.pool.lock().await.capacity_summary(now)
|
||||
@@ -877,7 +762,7 @@ impl RuntimeAccountPoolService {
|
||||
}
|
||||
|
||||
/// 按请求上下文获取可用账号。
|
||||
pub async fn acquire_with(&self, request: AccountAcquireRequest) -> Option<AcquiredAccount> {
|
||||
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)
|
||||
@@ -1162,11 +1047,6 @@ impl RuntimeAccountPoolService {
|
||||
self.pool.lock().await.remove(account_id)
|
||||
}
|
||||
|
||||
/// 从运行时账号池删除账号。
|
||||
pub async fn delete_account(&self, account_id: &str) -> bool {
|
||||
self.remove_account(account_id).await
|
||||
}
|
||||
|
||||
/// 从仓储同步单个账号到运行时账号池。
|
||||
pub async fn sync_account_from_repository(
|
||||
&self,
|
||||
@@ -1185,11 +1065,6 @@ impl RuntimeAccountPoolService {
|
||||
Ok(pool.remove(account_id))
|
||||
}
|
||||
|
||||
/// 重置运行时账号用量。
|
||||
pub async fn reset_usage(&self, account_id: &str) -> bool {
|
||||
self.pool.lock().await.reset_usage(account_id)
|
||||
}
|
||||
|
||||
/// 设置 Cloudflare 冷却状态。
|
||||
pub async fn set_cloudflare_cooldown_until(
|
||||
&self,
|
||||
@@ -1298,7 +1173,11 @@ fn next_window_reset_at(
|
||||
return None;
|
||||
}
|
||||
|
||||
let elapsed_seconds = now.signed_duration_since(reset_at).num_seconds().max(0) as u64;
|
||||
let elapsed_seconds = now
|
||||
.signed_duration_since(reset_at)
|
||||
.num_seconds()
|
||||
.max(0)
|
||||
.cast_unsigned();
|
||||
let windows_to_advance = elapsed_seconds / window_seconds + 1;
|
||||
let advance_seconds = window_seconds
|
||||
.saturating_mul(windows_to_advance)
|
||||
@@ -1342,7 +1221,7 @@ fn compare_least_used_group(a: &LeastUsedGroupKey, b: &LeastUsedGroupKey) -> Ord
|
||||
fn compare_window_reset(a: Option<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> Ordering {
|
||||
match (a, b) {
|
||||
(Some(a), Some(b)) => a.cmp(&b),
|
||||
(Some(_), None) | (None, Some(_)) | (None, None) => Ordering::Equal,
|
||||
(Some(_) | None, None) | (None, Some(_)) => Ordering::Equal,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1,16 +1,13 @@
|
||||
//! 配额检查策略。
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
const MONTH_WINDOW_MINUTES: u64 = 43_200;
|
||||
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
|
||||
|
||||
/// 判断配额百分比是否达到阈值。
|
||||
pub fn quota_reached(used_percent: f64, threshold: f64) -> bool {
|
||||
used_percent >= threshold
|
||||
}
|
||||
|
||||
/// 判断配额快照是否已经触顶。
|
||||
pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
|
||||
core_snapshot_blocked(quota)
|
||||
@@ -125,7 +122,8 @@ fn window_from_rate_limit(window: Option<&Value>) -> Value {
|
||||
.get("limit_window_seconds")
|
||||
.and_then(number_value)
|
||||
.filter(|seconds| *seconds > 0.0)
|
||||
.map(|seconds| (seconds / 60.0).round() as u64);
|
||||
.and_then(|seconds| Duration::try_from_secs_f64((seconds / 60.0).round()).ok())
|
||||
.map(|duration| duration.as_secs());
|
||||
json!({
|
||||
"used_percent": used_percent,
|
||||
"remaining_percent": remaining_percent(used_percent),
|
||||
|
||||
@@ -45,8 +45,7 @@ impl RuntimeQuotaRefreshService {
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回连续 quota 请求之间的默认错峰间隔。
|
||||
pub fn default_request_spacing() -> Duration {
|
||||
fn default_request_spacing() -> Duration {
|
||||
Duration::from_secs(DEFAULT_REQUEST_SPACING_SECS)
|
||||
}
|
||||
|
||||
@@ -84,14 +83,6 @@ impl RuntimeQuotaRefreshService {
|
||||
self
|
||||
}
|
||||
|
||||
/// 执行一次配额锁定或待验证账号刷新。
|
||||
pub async fn refresh_locked_accounts_once(
|
||||
&self,
|
||||
) -> QuotaRefreshServiceResult<QuotaRefreshSummary> {
|
||||
let mut last_refreshed = HashMap::new();
|
||||
self.refresh_locked_accounts(&mut last_refreshed).await
|
||||
}
|
||||
|
||||
/// 执行一次配额锁定或待验证账号刷新,复用调用方传入的最近刷新记录。
|
||||
pub async fn refresh_locked_accounts(
|
||||
&self,
|
||||
|
||||
@@ -2,40 +2,13 @@
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
|
||||
use crate::upstream::accounts::model::{Account, AccountStatus};
|
||||
|
||||
/// 根据 token 过期时间推导账号运行时状态。
|
||||
pub fn effective_status(account: &Account, now: DateTime<Utc>) -> AccountStatus {
|
||||
if account.status != AccountStatus::Active {
|
||||
return account.status;
|
||||
}
|
||||
if account
|
||||
.access_token_expires_at
|
||||
.is_some_and(|expires_at| expires_at <= now)
|
||||
{
|
||||
AccountStatus::Expired
|
||||
} else {
|
||||
AccountStatus::Active
|
||||
}
|
||||
}
|
||||
use crate::upstream::accounts::model::Account;
|
||||
|
||||
/// 账号服务。
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct AccountService;
|
||||
|
||||
impl AccountService {
|
||||
/// 判断账号是否可参与请求调度。
|
||||
pub fn can_serve(account: &Account) -> bool {
|
||||
account.status == AccountStatus::Active && !account.quota_limit_reached
|
||||
}
|
||||
|
||||
/// 判断账号在指定时间是否处于可服务的基础状态。
|
||||
pub fn can_serve_at(account: &Account, now: DateTime<Utc>, skip_quota_limited: bool) -> bool {
|
||||
account.status == AccountStatus::Active
|
||||
&& Self::quota_available_at(account, now, skip_quota_limited)
|
||||
&& Self::cloudflare_available_at(account, now)
|
||||
}
|
||||
|
||||
/// 判断账号配额锁定是否已经释放。
|
||||
pub fn quota_available_at(
|
||||
account: &Account,
|
||||
|
||||
+53
-222
@@ -350,61 +350,6 @@ select
|
||||
from account_model_usage
|
||||
order by account_id asc, request_count desc, last_used_at desc, model asc";
|
||||
|
||||
const GET_USAGE_SQL: &str = r"
|
||||
select
|
||||
account_id,
|
||||
request_count,
|
||||
empty_response_count,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
reasoning_tokens,
|
||||
total_tokens,
|
||||
image_input_tokens,
|
||||
image_output_tokens,
|
||||
image_request_count,
|
||||
image_request_failed_count,
|
||||
window_request_count,
|
||||
window_input_tokens,
|
||||
window_output_tokens,
|
||||
window_cached_tokens,
|
||||
window_image_input_tokens,
|
||||
window_image_output_tokens,
|
||||
window_image_request_count,
|
||||
window_image_request_failed_count,
|
||||
window_started_at,
|
||||
window_reset_at,
|
||||
limit_window_seconds,
|
||||
last_used_at
|
||||
from account_usage
|
||||
where account_id = ?";
|
||||
|
||||
const RESET_USAGE_SQL: &str = r"
|
||||
update account_usage
|
||||
set
|
||||
request_count = 0,
|
||||
empty_response_count = 0,
|
||||
input_tokens = 0,
|
||||
output_tokens = 0,
|
||||
cached_tokens = 0,
|
||||
reasoning_tokens = 0,
|
||||
total_tokens = 0,
|
||||
image_input_tokens = 0,
|
||||
image_output_tokens = 0,
|
||||
image_request_count = 0,
|
||||
image_request_failed_count = 0,
|
||||
window_request_count = 0,
|
||||
window_input_tokens = 0,
|
||||
window_output_tokens = 0,
|
||||
window_cached_tokens = 0,
|
||||
window_image_input_tokens = 0,
|
||||
window_image_output_tokens = 0,
|
||||
window_image_request_count = 0,
|
||||
window_image_request_failed_count = 0,
|
||||
window_started_at = null,
|
||||
last_used_at = null
|
||||
where account_id = ?";
|
||||
|
||||
const LIST_QUOTA_SNAPSHOTS_SQL: &str = r"
|
||||
select
|
||||
id,
|
||||
@@ -585,8 +530,6 @@ set
|
||||
updated_at = ?
|
||||
where id = ?";
|
||||
|
||||
const DELETE_ALL_ACCOUNTS_SQL: &str = "delete from accounts";
|
||||
|
||||
const DELETE_ACCOUNT_SQL: &str = "delete from accounts where id = ?";
|
||||
|
||||
// ============================================================================
|
||||
@@ -757,59 +700,6 @@ impl AccountMetadataUpdate {
|
||||
}
|
||||
}
|
||||
|
||||
/// 账号用量记录。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccountUsageRecord {
|
||||
/// 账号 ID。
|
||||
pub account_id: String,
|
||||
/// 历史请求总数。
|
||||
pub request_count: i64,
|
||||
/// 历史空响应次数。
|
||||
pub empty_response_count: i64,
|
||||
/// 累计输入 token。
|
||||
pub input_tokens: i64,
|
||||
/// 累计输出 token。
|
||||
pub output_tokens: i64,
|
||||
/// 累计缓存 token。
|
||||
pub cached_tokens: i64,
|
||||
/// 累计 reasoning token。
|
||||
pub reasoning_tokens: i64,
|
||||
/// 累计总 token。
|
||||
pub total_tokens: i64,
|
||||
/// 累计图片输入 token。
|
||||
pub image_input_tokens: i64,
|
||||
/// 累计图片输出 token。
|
||||
pub image_output_tokens: i64,
|
||||
/// 累计图片请求成功次数。
|
||||
pub image_request_count: i64,
|
||||
/// 累计图片请求失败次数。
|
||||
pub image_request_failed_count: i64,
|
||||
/// 当前窗口请求数。
|
||||
pub window_request_count: i64,
|
||||
/// 当前窗口输入 token。
|
||||
pub window_input_tokens: i64,
|
||||
/// 当前窗口输出 token。
|
||||
pub window_output_tokens: i64,
|
||||
/// 当前窗口缓存 token。
|
||||
pub window_cached_tokens: i64,
|
||||
/// 当前窗口图片输入 token。
|
||||
pub window_image_input_tokens: i64,
|
||||
/// 当前窗口图片输出 token。
|
||||
pub window_image_output_tokens: i64,
|
||||
/// 当前窗口图片请求数。
|
||||
pub window_image_request_count: i64,
|
||||
/// 当前窗口图片请求失败数。
|
||||
pub window_image_request_failed_count: i64,
|
||||
/// 当前窗口起始时间。
|
||||
pub window_started_at: Option<DateTime<Utc>>,
|
||||
/// 当前窗口重置时间。
|
||||
pub window_reset_at: Option<DateTime<Utc>>,
|
||||
/// 限流窗口大小。
|
||||
pub limit_window_seconds: Option<u64>,
|
||||
/// 最近使用时间。
|
||||
pub last_used_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
/// 账号模型维度用量记录。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AccountModelUsageRecord {
|
||||
@@ -1010,8 +900,7 @@ impl SqliteAccountStore {
|
||||
.bind(
|
||||
account
|
||||
.added_at
|
||||
.map(|dt| dt.to_rfc3339())
|
||||
.unwrap_or_else(|| Utc::now().to_rfc3339()),
|
||||
.map_or_else(|| Utc::now().to_rfc3339(), |dt| dt.to_rfc3339()),
|
||||
)
|
||||
.bind(&now)
|
||||
.execute(&self.pool)
|
||||
@@ -1061,7 +950,7 @@ impl SqliteAccountStore {
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
Ok(to_page(
|
||||
rows,
|
||||
&rows,
|
||||
limit,
|
||||
stored_account_from_row,
|
||||
("added_at", "id"),
|
||||
@@ -1072,7 +961,7 @@ impl SqliteAccountStore {
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
Ok(to_page(
|
||||
rows,
|
||||
&rows,
|
||||
limit,
|
||||
stored_account_from_row,
|
||||
("added_at", "id"),
|
||||
@@ -1097,13 +986,13 @@ impl SqliteAccountStore {
|
||||
.bind(limit + 1)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
Ok(to_page(rows, limit, metadata_from_row, ("added_at", "id")))
|
||||
Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id")))
|
||||
} else {
|
||||
let rows = sqlx::query(LIST_ACCOUNT_METADATA_SQL)
|
||||
.bind(limit + 1)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
Ok(to_page(rows, limit, metadata_from_row, ("added_at", "id")))
|
||||
Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id")))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1426,27 +1315,6 @@ impl SqliteAccountStore {
|
||||
rows.iter().map(model_usage_from_row).collect()
|
||||
}
|
||||
|
||||
/// 获取用量记录。
|
||||
pub async fn get_usage(
|
||||
&self,
|
||||
account_id: &str,
|
||||
) -> SqliteAccountStoreResult<Option<AccountUsageRecord>> {
|
||||
let row = sqlx::query(GET_USAGE_SQL)
|
||||
.bind(account_id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(usage_from_row).transpose()
|
||||
}
|
||||
|
||||
/// 重置用量。
|
||||
pub async fn reset_usage(&self, account_id: &str) -> SqliteAccountStoreResult<bool> {
|
||||
let result = sqlx::query(RESET_USAGE_SQL)
|
||||
.bind(account_id)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
/// 读取配额快照列表。
|
||||
pub async fn list_quota_snapshots(
|
||||
&self,
|
||||
@@ -1551,10 +1419,11 @@ impl SqliteAccountStore {
|
||||
let should_reset = existing
|
||||
.as_ref()
|
||||
.map(|row| {
|
||||
let existing_reset_at =
|
||||
parse_optional_rfc3339(row.get::<Option<String>, _>("window_reset_at"))
|
||||
.ok()
|
||||
.flatten();
|
||||
let existing_reset_at = parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("window_reset_at").as_deref(),
|
||||
)
|
||||
.ok()
|
||||
.flatten();
|
||||
let existing_limit_window_seconds =
|
||||
optional_positive_i64_to_u64(row.get::<Option<i64>, _>("limit_window_seconds"));
|
||||
should_reset_usage_window(
|
||||
@@ -1587,14 +1456,6 @@ impl SqliteAccountStore {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 删除全部账号。
|
||||
pub async fn delete_all(&self) -> SqliteAccountStoreResult<u64> {
|
||||
let result = sqlx::query(DELETE_ALL_ACCOUNTS_SQL)
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
|
||||
/// 删除单个账号。
|
||||
pub async fn delete(&self, account_id: &str) -> SqliteAccountStoreResult<bool> {
|
||||
let result = sqlx::query(DELETE_ACCOUNT_SQL)
|
||||
@@ -1622,16 +1483,6 @@ impl SqliteAccountStore {
|
||||
.map_err(SqliteAccountStoreError::Database)?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
/// 按 ID 删除账号。
|
||||
pub async fn delete_account(&self, account_id: &str) -> Result<bool, SqliteAccountStoreError> {
|
||||
let result = sqlx::query("DELETE FROM accounts WHERE id = ?1")
|
||||
.bind(account_id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_err(SqliteAccountStoreError::Database)?;
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -1639,13 +1490,13 @@ impl AccountStore for SqliteAccountStore {
|
||||
async fn list_pool_accounts(&self) -> AccountStoreResult<Vec<Account>> {
|
||||
list_pool_accounts(self)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult<Option<Account>> {
|
||||
get_pool_account(self, account_id)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn mark_quota_limited_until(
|
||||
@@ -1655,7 +1506,7 @@ impl AccountStore for SqliteAccountStore {
|
||||
) -> AccountStoreResult<bool> {
|
||||
SqliteAccountStore::mark_quota_limited_until(self, account_id, cooldown_until)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn set_cloudflare_cooldown_until(
|
||||
@@ -1665,7 +1516,7 @@ impl AccountStore for SqliteAccountStore {
|
||||
) -> AccountStoreResult<bool> {
|
||||
SqliteAccountStore::set_cloudflare_cooldown_until(self, account_id, cooldown_until)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn set_status(
|
||||
@@ -1675,7 +1526,7 @@ impl AccountStore for SqliteAccountStore {
|
||||
) -> AccountStoreResult<bool> {
|
||||
SqliteAccountStore::set_status(self, account_id, status)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn record_usage_delta(
|
||||
@@ -1685,7 +1536,7 @@ impl AccountStore for SqliteAccountStore {
|
||||
) -> AccountStoreResult<()> {
|
||||
self.record_usage(account_id, sqlite_usage_delta(usage))
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn record_model_usage_delta(
|
||||
@@ -1696,13 +1547,13 @@ impl AccountStore for SqliteAccountStore {
|
||||
) -> AccountStoreResult<()> {
|
||||
self.record_model_usage(account_id, model, usage)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult<Option<String>> {
|
||||
SqliteAccountStore::get_quota_json(self, account_id)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn update_quota_json(
|
||||
@@ -1712,7 +1563,7 @@ impl AccountStore for SqliteAccountStore {
|
||||
) -> AccountStoreResult<bool> {
|
||||
SqliteAccountStore::update_quota_json(self, account_id, quota_json)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn apply_quota_snapshot(
|
||||
@@ -1730,7 +1581,7 @@ impl AccountStore for SqliteAccountStore {
|
||||
cooldown_until,
|
||||
)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
|
||||
async fn sync_rate_limit_window(
|
||||
@@ -1741,7 +1592,7 @@ impl AccountStore for SqliteAccountStore {
|
||||
) -> AccountStoreResult<()> {
|
||||
SqliteAccountStore::sync_rate_limit_window(self, account_id, reset_at, limit_window_seconds)
|
||||
.await
|
||||
.map_err(map_account_store_error)
|
||||
.map_err(|error| map_account_store_error(&error))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1784,17 +1635,22 @@ fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult<Account> {
|
||||
access_token: row.get("access_token"),
|
||||
refresh_token: row.get("refresh_token"),
|
||||
access_token_expires_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("access_token_expires_at"),
|
||||
row.get::<Option<String>, _>("access_token_expires_at")
|
||||
.as_deref(),
|
||||
)?,
|
||||
next_refresh_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("next_refresh_at").as_deref(),
|
||||
)?,
|
||||
next_refresh_at: parse_optional_rfc3339(row.get::<Option<String>, _>("next_refresh_at"))?,
|
||||
status: status_from_db(&row.get::<String, _>("status"))?,
|
||||
quota_limit_reached: row.get::<i64, _>("quota_limit_reached") != 0,
|
||||
quota_verify_required: row.get::<i64, _>("quota_verify_required") != 0,
|
||||
quota_cooldown_until: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("quota_cooldown_until"),
|
||||
row.get::<Option<String>, _>("quota_cooldown_until")
|
||||
.as_deref(),
|
||||
)?,
|
||||
cloudflare_cooldown_until: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("cloudflare_cooldown_until"),
|
||||
row.get::<Option<String>, _>("cloudflare_cooldown_until")
|
||||
.as_deref(),
|
||||
)?,
|
||||
request_count: nonnegative_i64_to_u64(row.get::<i64, _>("usage_request_count")),
|
||||
empty_response_count: nonnegative_i64_to_u64(
|
||||
@@ -1829,10 +1685,12 @@ fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult<Account> {
|
||||
row.get::<i64, _>("usage_window_image_request_failed_count"),
|
||||
),
|
||||
window_started_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("usage_window_started_at"),
|
||||
row.get::<Option<String>, _>("usage_window_started_at")
|
||||
.as_deref(),
|
||||
)?,
|
||||
window_reset_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("usage_window_reset_at"),
|
||||
row.get::<Option<String>, _>("usage_window_reset_at")
|
||||
.as_deref(),
|
||||
)?,
|
||||
limit_window_seconds: optional_positive_i64_to_u64(
|
||||
row.get::<Option<i64>, _>("usage_limit_window_seconds"),
|
||||
@@ -1855,9 +1713,12 @@ fn stored_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult<StoredAc
|
||||
.get::<Option<String>, _>("refresh_token")
|
||||
.map(|token| SecretString::new(token.into())),
|
||||
access_token_expires_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("access_token_expires_at"),
|
||||
row.get::<Option<String>, _>("access_token_expires_at")
|
||||
.as_deref(),
|
||||
)?,
|
||||
next_refresh_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("next_refresh_at").as_deref(),
|
||||
)?,
|
||||
next_refresh_at: parse_optional_rfc3339(row.get::<Option<String>, _>("next_refresh_at"))?,
|
||||
status: status_from_db(&row.get::<String, _>("status"))?,
|
||||
added_at: row.get("added_at"),
|
||||
updated_at: row.get("updated_at"),
|
||||
@@ -1873,7 +1734,8 @@ fn metadata_from_row(row: &SqliteRow) -> SqliteAccountStoreResult<StoredAccountM
|
||||
label: row.get("label"),
|
||||
plan_type: row.get("plan_type"),
|
||||
access_token_expires_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("access_token_expires_at"),
|
||||
row.get::<Option<String>, _>("access_token_expires_at")
|
||||
.as_deref(),
|
||||
)?,
|
||||
status: status_from_db(&row.get::<String, _>("status"))?,
|
||||
added_at: row.get("added_at"),
|
||||
@@ -1881,7 +1743,7 @@ fn metadata_from_row(row: &SqliteRow) -> SqliteAccountStoreResult<StoredAccountM
|
||||
})
|
||||
}
|
||||
|
||||
fn map_account_store_error(error: impl ToString) -> AccountStoreError {
|
||||
fn map_account_store_error(error: &impl ToString) -> AccountStoreError {
|
||||
AccountStoreError::OperationFailed {
|
||||
message: error.to_string(),
|
||||
}
|
||||
@@ -1958,39 +1820,6 @@ fn quota_plan_type(quota_json: &str) -> Option<String> {
|
||||
.map(ToString::to_string)
|
||||
}
|
||||
|
||||
fn usage_from_row(row: &sqlx::sqlite::SqliteRow) -> SqliteAccountStoreResult<AccountUsageRecord> {
|
||||
Ok(AccountUsageRecord {
|
||||
account_id: row.get("account_id"),
|
||||
request_count: row.get("request_count"),
|
||||
empty_response_count: row.get("empty_response_count"),
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
reasoning_tokens: row.get("reasoning_tokens"),
|
||||
total_tokens: row.get("total_tokens"),
|
||||
image_input_tokens: row.get("image_input_tokens"),
|
||||
image_output_tokens: row.get("image_output_tokens"),
|
||||
image_request_count: row.get("image_request_count"),
|
||||
image_request_failed_count: row.get("image_request_failed_count"),
|
||||
window_request_count: row.get("window_request_count"),
|
||||
window_input_tokens: row.get("window_input_tokens"),
|
||||
window_output_tokens: row.get("window_output_tokens"),
|
||||
window_cached_tokens: row.get("window_cached_tokens"),
|
||||
window_image_input_tokens: row.get("window_image_input_tokens"),
|
||||
window_image_output_tokens: row.get("window_image_output_tokens"),
|
||||
window_image_request_count: row.get("window_image_request_count"),
|
||||
window_image_request_failed_count: row.get("window_image_request_failed_count"),
|
||||
window_started_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("window_started_at"),
|
||||
)?,
|
||||
window_reset_at: parse_optional_rfc3339(row.get::<Option<String>, _>("window_reset_at"))?,
|
||||
limit_window_seconds: optional_positive_i64_to_u64(
|
||||
row.get::<Option<i64>, _>("limit_window_seconds"),
|
||||
),
|
||||
last_used_at: parse_optional_rfc3339(row.get::<Option<String>, _>("last_used_at"))?,
|
||||
})
|
||||
}
|
||||
|
||||
fn model_usage_from_row(
|
||||
row: &sqlx::sqlite::SqliteRow,
|
||||
) -> SqliteAccountStoreResult<AccountModelUsageRecord> {
|
||||
@@ -2002,7 +1831,9 @@ fn model_usage_from_row(
|
||||
input_tokens: row.get("input_tokens"),
|
||||
output_tokens: row.get("output_tokens"),
|
||||
cached_tokens: row.get("cached_tokens"),
|
||||
last_used_at: parse_optional_rfc3339(row.get::<Option<String>, _>("last_used_at"))?,
|
||||
last_used_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("last_used_at").as_deref(),
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2013,7 +1844,9 @@ fn quota_snapshot_from_row(
|
||||
account_id: row.get("id"),
|
||||
email: row.get("email"),
|
||||
quota_json: row.get("quota_json"),
|
||||
quota_fetched_at: parse_optional_rfc3339(row.get::<Option<String>, _>("quota_fetched_at"))?,
|
||||
quota_fetched_at: parse_optional_rfc3339(
|
||||
row.get::<Option<String>, _>("quota_fetched_at").as_deref(),
|
||||
)?,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2044,10 +1877,8 @@ fn status_from_db(value: &str) -> SqliteAccountStoreResult<AccountStatus> {
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_optional_rfc3339(
|
||||
value: Option<String>,
|
||||
) -> SqliteAccountStoreResult<Option<DateTime<Utc>>> {
|
||||
value.as_deref().map(parse_rfc3339).transpose()
|
||||
fn parse_optional_rfc3339(value: Option<&str>) -> SqliteAccountStoreResult<Option<DateTime<Utc>>> {
|
||||
value.map(parse_rfc3339).transpose()
|
||||
}
|
||||
|
||||
fn parse_rfc3339(value: &str) -> SqliteAccountStoreResult<DateTime<Utc>> {
|
||||
@@ -2055,7 +1886,7 @@ fn parse_rfc3339(value: &str) -> SqliteAccountStoreResult<DateTime<Utc>> {
|
||||
}
|
||||
|
||||
fn nonnegative_i64_to_u64(value: i64) -> u64 {
|
||||
value.max(0) as u64
|
||||
value.max(0).cast_unsigned()
|
||||
}
|
||||
|
||||
fn u64_to_i64_saturating(value: u64) -> i64 {
|
||||
@@ -2075,7 +1906,7 @@ async fn count_account_metadata(
|
||||
let mut builder = QueryBuilder::<Sqlite>::new("select count(*) from accounts");
|
||||
push_account_metadata_search(&mut builder, search);
|
||||
let (total,): (i64,) = builder.build_query_as().fetch_one(pool).await?;
|
||||
Ok(total.max(0) as u64)
|
||||
Ok(total.max(0).cast_unsigned())
|
||||
}
|
||||
|
||||
fn push_account_metadata_search(builder: &mut QueryBuilder<Sqlite>, search: Option<&str>) {
|
||||
@@ -2097,7 +1928,7 @@ fn push_account_metadata_search(builder: &mut QueryBuilder<Sqlite>, search: Opti
|
||||
}
|
||||
|
||||
fn to_page<T>(
|
||||
rows: Vec<SqliteRow>,
|
||||
rows: &[SqliteRow],
|
||||
limit: u32,
|
||||
mapper: impl Fn(&SqliteRow) -> SqliteAccountStoreResult<T>,
|
||||
cursor_fields: (&str, &str),
|
||||
|
||||
@@ -364,11 +364,6 @@ impl RefreshLeaseStore {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// 返回底层连接池。
|
||||
pub fn pool(&self) -> &SqlitePool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
/// 尝试获取账号刷新租约。
|
||||
pub async fn try_acquire(
|
||||
&self,
|
||||
@@ -378,7 +373,7 @@ impl RefreshLeaseStore {
|
||||
now: DateTime<Utc>,
|
||||
) -> RefreshLeaseStoreResult<bool> {
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
r"
|
||||
insert into account_refresh_leases (account_id, owner, expires_at, updated_at)
|
||||
values (?, ?, ?, ?)
|
||||
on conflict(account_id) do update set
|
||||
@@ -387,7 +382,7 @@ on conflict(account_id) do update set
|
||||
updated_at = excluded.updated_at
|
||||
where account_refresh_leases.expires_at <= ?
|
||||
or account_refresh_leases.owner = ?
|
||||
"#,
|
||||
",
|
||||
)
|
||||
.bind(account_id)
|
||||
.bind(owner)
|
||||
|
||||
@@ -113,13 +113,6 @@ where
|
||||
self
|
||||
}
|
||||
|
||||
/// 执行一次到期刷新扫描。
|
||||
pub async fn refresh_due_accounts_once(
|
||||
&self,
|
||||
) -> TokenRefreshServiceResult<TokenRefreshSummary> {
|
||||
self.refresh_due_accounts_once_at(Utc::now()).await
|
||||
}
|
||||
|
||||
/// 执行一次账号级刷新定时器调度。
|
||||
pub async fn schedule_account_timers_once(
|
||||
&self,
|
||||
@@ -530,11 +523,8 @@ where
|
||||
}
|
||||
Ok(_) => return Ok(TokenRefreshOutcome::Skipped),
|
||||
Err(RefreshError::RetryableTransport)
|
||||
if self.sleep_before_retry(attempt_index).await =>
|
||||
{
|
||||
continue;
|
||||
}
|
||||
Err(RefreshError::RetryableTransport) | Err(RefreshError::Transport) => {
|
||||
if self.sleep_before_retry(attempt_index).await => {}
|
||||
Err(RefreshError::RetryableTransport | RefreshError::Transport) => {
|
||||
persist_status(&self.store, &attempt.id, AccountStatus::Active).await?;
|
||||
let next_refresh_at = now + chrono::Duration::seconds(RECOVERY_DELAY_SECONDS);
|
||||
persist_next_refresh_at(&self.store, &attempt.id, Some(next_refresh_at))
|
||||
|
||||
+5
-192
@@ -11,15 +11,13 @@ use tokio::sync::Mutex;
|
||||
use tracing::{info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
/// 指纹历史记录的来源标识。
|
||||
pub const CODEX_DESKTOP_UPDATE_SOURCE: &str = "codex_desktop_update_source";
|
||||
/// 运行时当前指纹槽位。
|
||||
pub const CURRENT_FINGERPRINT_ID: &str = "current";
|
||||
|
||||
const AUTO_UPDATE_SOURCE: &str = "auto_update";
|
||||
const CONFIG_SEED_SOURCE: &str = "config_seed";
|
||||
const APPCAST_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const POLL_INTERVAL: Duration = Duration::from_secs(3 * 24 * 60 * 60);
|
||||
const POLL_INTERVAL: Duration = Duration::from_hours(72);
|
||||
|
||||
type AppcastFields = (Option<String>, Option<String>, Option<String>);
|
||||
|
||||
@@ -74,64 +72,6 @@ impl Fingerprint {
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回默认的 Codex Desktop 指纹。
|
||||
pub fn default_codex_desktop() -> Self {
|
||||
Self {
|
||||
originator: "Codex Desktop".to_string(),
|
||||
app_version: "26.519.81530".to_string(),
|
||||
build_number: "3178".to_string(),
|
||||
platform: "darwin".to_string(),
|
||||
arch: "arm64".to_string(),
|
||||
chromium_version: "146".to_string(),
|
||||
user_agent_template: "Codex Desktop/{version} ({platform}; {arch})".to_string(),
|
||||
default_headers: Self::default_headers(),
|
||||
header_order: Self::default_header_order(),
|
||||
updated_at: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回默认请求头集合。
|
||||
pub fn default_headers() -> IndexMap<String, String> {
|
||||
let mut headers = IndexMap::new();
|
||||
headers.insert(
|
||||
"Accept-Encoding".to_string(),
|
||||
"gzip, deflate, br, zstd".to_string(),
|
||||
);
|
||||
headers.insert("Accept-Language".to_string(), "en-US,en;q=0.9".to_string());
|
||||
headers.insert("sec-ch-ua-mobile".to_string(), "?0".to_string());
|
||||
headers.insert("sec-ch-ua-platform".to_string(), "\"macOS\"".to_string());
|
||||
headers.insert("sec-fetch-site".to_string(), "same-origin".to_string());
|
||||
headers.insert("sec-fetch-mode".to_string(), "cors".to_string());
|
||||
headers.insert("sec-fetch-dest".to_string(), "empty".to_string());
|
||||
headers
|
||||
}
|
||||
|
||||
/// 返回默认请求头顺序。
|
||||
pub fn default_header_order() -> Vec<String> {
|
||||
vec![
|
||||
"authorization".to_string(),
|
||||
"chatgpt-account-id".to_string(),
|
||||
"originator".to_string(),
|
||||
"x-openai-internal-codex-residency".to_string(),
|
||||
"x-client-request-id".to_string(),
|
||||
"x-codex-installation-id".to_string(),
|
||||
"x-codex-turn-state".to_string(),
|
||||
"openai-beta".to_string(),
|
||||
"user-agent".to_string(),
|
||||
"sec-ch-ua".to_string(),
|
||||
"sec-ch-ua-mobile".to_string(),
|
||||
"sec-ch-ua-platform".to_string(),
|
||||
"accept-encoding".to_string(),
|
||||
"accept-language".to_string(),
|
||||
"sec-fetch-site".to_string(),
|
||||
"sec-fetch-mode".to_string(),
|
||||
"sec-fetch-dest".to_string(),
|
||||
"content-type".to_string(),
|
||||
"accept".to_string(),
|
||||
"cookie".to_string(),
|
||||
]
|
||||
}
|
||||
|
||||
/// 根据模板展开最终 User-Agent。
|
||||
pub fn user_agent(&self) -> String {
|
||||
self.user_agent_template
|
||||
@@ -149,63 +89,10 @@ impl Fingerprint {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 更新清单解析
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// 更新清单解析结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct FingerprintUpdate {
|
||||
/// 应用版本。
|
||||
pub app_version: String,
|
||||
/// 构建号。
|
||||
pub build_number: String,
|
||||
}
|
||||
|
||||
/// 指纹更新错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum FingerprintError {
|
||||
/// 更新清单 JSON 无效。
|
||||
#[error("invalid update manifest: {0}")]
|
||||
InvalidManifest(#[from] serde_json::Error),
|
||||
/// HTTP 请求失败。
|
||||
#[error("failed to fetch update manifest: {0}")]
|
||||
Http(#[from] reqwest::Error),
|
||||
/// 数据库存储失败。
|
||||
#[error("failed to persist fingerprint update: {0}")]
|
||||
Database(#[from] sqlx::Error),
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Manifest {
|
||||
version: String,
|
||||
build_number: String,
|
||||
}
|
||||
|
||||
/// 从更新清单 JSON 中提取版本信息。
|
||||
pub fn parse_update_manifest(input: &str) -> Result<FingerprintUpdate, FingerprintError> {
|
||||
let manifest: Manifest = serde_json::from_str(input)?;
|
||||
Ok(FingerprintUpdate {
|
||||
app_version: manifest.version,
|
||||
build_number: manifest.build_number,
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 指纹历史记录与仓储
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// 指纹历史记录。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct StoredFingerprint {
|
||||
/// 应用版本。
|
||||
pub app_version: String,
|
||||
/// 构建号。
|
||||
pub build_number: String,
|
||||
/// 来源。
|
||||
pub source: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
struct StoredHeader {
|
||||
name: String,
|
||||
@@ -266,22 +153,6 @@ impl FingerprintRepository {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
/// 插入一条指纹更新历史。
|
||||
pub async fn insert_update(&self, update: &FingerprintUpdate) -> Result<(), sqlx::Error> {
|
||||
sqlx::query(
|
||||
"insert into fingerprint_update_history (id, current_fingerprint_id, app_version, build_number, source, created_at) values (?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(CURRENT_FINGERPRINT_ID)
|
||||
.bind(&update.app_version)
|
||||
.bind(&update.build_number)
|
||||
.bind(CODEX_DESKTOP_UPDATE_SOURCE)
|
||||
.bind(Utc::now().to_rfc3339())
|
||||
.execute(&self.pool)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 写入当前指纹默认值;如果当前槽位已存在,直接读取数据库值。
|
||||
pub async fn ensure_current_seed(
|
||||
&self,
|
||||
@@ -293,7 +164,7 @@ impl FingerprintRepository {
|
||||
|
||||
let now = Utc::now().to_rfc3339();
|
||||
sqlx::query(
|
||||
r#"
|
||||
r"
|
||||
insert into fingerprints (
|
||||
id,
|
||||
originator,
|
||||
@@ -309,7 +180,7 @@ impl FingerprintRepository {
|
||||
created_at,
|
||||
updated_at
|
||||
) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
",
|
||||
)
|
||||
.bind(CURRENT_FINGERPRINT_ID)
|
||||
.bind(&default_fingerprint.originator)
|
||||
@@ -371,7 +242,7 @@ impl FingerprintRepository {
|
||||
/// 读取当前运行时指纹。
|
||||
pub async fn load_current(&self) -> Result<Option<Fingerprint>, sqlx::Error> {
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
r"
|
||||
select
|
||||
originator,
|
||||
app_version,
|
||||
@@ -385,7 +256,7 @@ impl FingerprintRepository {
|
||||
updated_at
|
||||
from fingerprints
|
||||
where id = ?
|
||||
"#,
|
||||
",
|
||||
)
|
||||
.bind(CURRENT_FINGERPRINT_ID)
|
||||
.fetch_optional(&self.pool)
|
||||
@@ -417,21 +288,6 @@ impl FingerprintRepository {
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 读取最新一条指纹历史记录。
|
||||
pub async fn latest(&self) -> Result<Option<StoredFingerprint>, sqlx::Error> {
|
||||
let row = sqlx::query(
|
||||
"select app_version, build_number, source from fingerprint_update_history order by created_at desc, id desc limit 1",
|
||||
)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(row.map(|row| StoredFingerprint {
|
||||
app_version: row.get("app_version"),
|
||||
build_number: row.get("build_number"),
|
||||
source: row.get("source"),
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
fn fingerprint_from_row(row: &sqlx::sqlite::SqliteRow) -> Result<Fingerprint, sqlx::Error> {
|
||||
@@ -482,49 +338,6 @@ fn json_error(error: serde_json::Error) -> sqlx::Error {
|
||||
sqlx::Error::Decode(Box::new(error))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 指纹更新器(HTTP 拉取 + 持久化)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// 通过 HTTP 拉取更新清单并写入历史记录。
|
||||
#[derive(Clone)]
|
||||
pub struct FingerprintUpdater {
|
||||
client: reqwest::Client,
|
||||
repository: FingerprintRepository,
|
||||
update_url: String,
|
||||
}
|
||||
|
||||
impl FingerprintUpdater {
|
||||
/// 构造更新器。
|
||||
pub fn new(
|
||||
client: reqwest::Client,
|
||||
repository: FingerprintRepository,
|
||||
update_url: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
client,
|
||||
repository,
|
||||
update_url: update_url.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// 执行一次更新清单拉取并持久化历史记录。
|
||||
pub async fn poll_once(&self) -> Result<FingerprintUpdate, FingerprintError> {
|
||||
let manifest = self
|
||||
.client
|
||||
.get(&self.update_url)
|
||||
.timeout(APPCAST_TIMEOUT)
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?
|
||||
.text()
|
||||
.await?;
|
||||
let update = parse_update_manifest(&manifest)?;
|
||||
self.repository.insert_update(&update).await?;
|
||||
Ok(update)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Appcast 指纹更新检查器
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
+2
-108
@@ -7,10 +7,7 @@ use serde::{Deserialize, Serialize};
|
||||
use sqlx::Row;
|
||||
use thiserror::Error;
|
||||
|
||||
use crate::upstream::accounts::{
|
||||
model::{Account, AccountStatus},
|
||||
store::AccountStore,
|
||||
};
|
||||
use crate::upstream::accounts::model::{Account, AccountStatus};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 数据类型
|
||||
@@ -195,85 +192,6 @@ pub trait ModelSnapshotStore: Send + Sync + 'static {
|
||||
async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult<Vec<ModelPlanSnapshot>>;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// 管理端模型刷新服务
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// 管理端模型服务。
|
||||
#[derive(Clone)]
|
||||
pub struct AdminModelService {
|
||||
models: Arc<ModelService>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
installation_id: Option<String>,
|
||||
}
|
||||
|
||||
impl AdminModelService {
|
||||
/// 构造服务。
|
||||
pub fn new(
|
||||
models: Arc<ModelService>,
|
||||
accounts: Arc<dyn AccountStore>,
|
||||
installation_id: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
models,
|
||||
accounts,
|
||||
installation_id,
|
||||
}
|
||||
}
|
||||
|
||||
/// 使用账号池中的账号刷新上游模型目录。
|
||||
pub async fn refresh_backend_models(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<ModelRefreshResult, AdminModelError> {
|
||||
let accounts = self
|
||||
.accounts
|
||||
.list_pool_accounts()
|
||||
.await
|
||||
.map_err(|_| AdminModelError::ListAccounts)?;
|
||||
self.models
|
||||
.refresh_backend_models_with_installation_id(
|
||||
&accounts,
|
||||
request_id,
|
||||
self.installation_id.as_deref(),
|
||||
)
|
||||
.await
|
||||
.map_err(AdminModelError::from)
|
||||
}
|
||||
}
|
||||
|
||||
/// 管理端模型刷新错误。
|
||||
#[derive(Debug, Error)]
|
||||
pub enum AdminModelError {
|
||||
#[error("failed to list accounts")]
|
||||
ListAccounts,
|
||||
#[error("no active accounts available for model refresh")]
|
||||
NoAccounts,
|
||||
#[error("model snapshot store is unavailable")]
|
||||
SnapshotStoreUnavailable,
|
||||
#[error("model upstream client is unavailable")]
|
||||
UpstreamClientUnavailable,
|
||||
#[error("failed to store model snapshot")]
|
||||
StoreSnapshot,
|
||||
#[error("failed to load model snapshots")]
|
||||
LoadSnapshots,
|
||||
#[error("all model refresh plans failed")]
|
||||
AllPlansFailed(ModelRefreshResult),
|
||||
}
|
||||
|
||||
impl From<ModelServiceError> for AdminModelError {
|
||||
fn from(error: ModelServiceError) -> Self {
|
||||
match error {
|
||||
ModelServiceError::SnapshotStoreUnavailable => Self::SnapshotStoreUnavailable,
|
||||
ModelServiceError::UpstreamClientUnavailable => Self::UpstreamClientUnavailable,
|
||||
ModelServiceError::NoAccounts => Self::NoAccounts,
|
||||
ModelServiceError::StoreSnapshot => Self::StoreSnapshot,
|
||||
ModelServiceError::LoadSnapshots => Self::LoadSnapshots,
|
||||
ModelServiceError::AllPlansFailed(result) => Self::AllPlansFailed(result),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// SQLite 模型快照存储。
|
||||
#[derive(Clone)]
|
||||
pub struct SqliteModelSnapshotStore {
|
||||
@@ -658,16 +576,12 @@ pub enum ModelServiceError {
|
||||
/// 模型到可用计划的映射。
|
||||
pub type ModelPlanAllowlist = BTreeMap<String, Vec<String>>;
|
||||
|
||||
/// 可共享更新的模型计划映射缓存。
|
||||
pub type SharedModelPlanAllowlist = Arc<tokio::sync::Mutex<ModelPlanAllowlist>>;
|
||||
|
||||
/// 模型目录服务。
|
||||
#[derive(Clone)]
|
||||
pub struct ModelService {
|
||||
config: ModelConfig,
|
||||
snapshot_store: Option<Arc<dyn ModelSnapshotStore>>,
|
||||
upstream_client: Option<Arc<dyn CodexModelCatalogClient>>,
|
||||
model_plan_allowlist: Option<SharedModelPlanAllowlist>,
|
||||
}
|
||||
|
||||
impl ModelService {
|
||||
@@ -676,21 +590,14 @@ impl ModelService {
|
||||
config: ModelConfig,
|
||||
snapshot_store: Option<Arc<dyn ModelSnapshotStore>>,
|
||||
upstream_client: Option<Arc<dyn CodexModelCatalogClient>>,
|
||||
model_plan_allowlist: Option<SharedModelPlanAllowlist>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
snapshot_store,
|
||||
upstream_client,
|
||||
model_plan_allowlist,
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回模型目录的静态配置。
|
||||
pub fn config(&self) -> &ModelConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// 构造当前对外暴露的模型目录。
|
||||
pub async fn catalog(&self) -> ModelCatalog {
|
||||
let Some(snapshot_store) = self.snapshot_store.as_ref() else {
|
||||
@@ -709,16 +616,6 @@ impl ModelService {
|
||||
}
|
||||
}
|
||||
|
||||
/// 刷新活跃账号对应的后端模型目录。
|
||||
pub async fn refresh_backend_models(
|
||||
&self,
|
||||
accounts: &[Account],
|
||||
request_id: &str,
|
||||
) -> Result<ModelRefreshResult, ModelServiceError> {
|
||||
self.refresh_backend_models_with_installation_id(accounts, request_id, None)
|
||||
.await
|
||||
}
|
||||
|
||||
/// 使用运行时 installation id 刷新活跃账号对应的后端模型目录。
|
||||
pub async fn refresh_backend_models_with_installation_id(
|
||||
&self,
|
||||
@@ -781,10 +678,7 @@ impl ModelService {
|
||||
return Err(ModelServiceError::AllPlansFailed(result));
|
||||
}
|
||||
|
||||
let allowlist = self.model_plan_allowlist_from_store().await?;
|
||||
if let Some(shared_allowlist) = &self.model_plan_allowlist {
|
||||
*shared_allowlist.lock().await = allowlist;
|
||||
}
|
||||
self.model_plan_allowlist_from_store().await?;
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
@@ -1,9 +1,5 @@
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
time::{SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
use std::{collections::BTreeMap, time::Duration};
|
||||
|
||||
use chrono::{DateTime, Duration, Utc};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
@@ -29,34 +25,6 @@ pub struct TokenUsage {
|
||||
pub total_tokens: u64,
|
||||
}
|
||||
|
||||
impl TokenUsage {
|
||||
/// 构造仅包含主模型 token 的用量快照。
|
||||
pub fn new(input_tokens: u64, output_tokens: u64, cached_tokens: u64) -> Self {
|
||||
Self {
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cached_tokens,
|
||||
reasoning_tokens: 0,
|
||||
image_input_tokens: 0,
|
||||
image_output_tokens: 0,
|
||||
total_tokens: input_tokens + output_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
/// 将两份用量按字段累加。
|
||||
pub fn merged(self, other: Self) -> Self {
|
||||
Self {
|
||||
input_tokens: self.input_tokens + other.input_tokens,
|
||||
output_tokens: self.output_tokens + other.output_tokens,
|
||||
cached_tokens: self.cached_tokens + other.cached_tokens,
|
||||
reasoning_tokens: self.reasoning_tokens + other.reasoning_tokens,
|
||||
image_input_tokens: self.image_input_tokens + other.image_input_tokens,
|
||||
image_output_tokens: self.image_output_tokens + other.image_output_tokens,
|
||||
total_tokens: self.total_tokens + other.total_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 从单个 JSON 响应体中提取用量。
|
||||
///
|
||||
/// 该函数同时支持 Codex usage 结构和 OpenAI usage 结构。
|
||||
@@ -198,20 +166,6 @@ pub struct RateLimitWindow {
|
||||
pub reset_at: Option<i64>,
|
||||
}
|
||||
|
||||
impl RateLimitWindow {
|
||||
/// 以秒返回限流窗口大小。
|
||||
pub fn limit_window_seconds(self) -> Option<u64> {
|
||||
self.window_minutes
|
||||
.and_then(|minutes| minutes.checked_mul(60))
|
||||
}
|
||||
|
||||
/// 将重置时间戳转换为 UTC 时间。
|
||||
pub fn reset_at_datetime(self) -> Option<DateTime<Utc>> {
|
||||
let reset_at = self.reset_at?;
|
||||
DateTime::<Utc>::from_timestamp(reset_at, 0)
|
||||
}
|
||||
}
|
||||
|
||||
/// 复合限流信息,通常用于 code review 配额。
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub struct RateLimitDetails {
|
||||
@@ -236,24 +190,6 @@ pub struct ParsedRateLimits {
|
||||
pub code_review: Option<RateLimitDetails>,
|
||||
}
|
||||
|
||||
impl ParsedRateLimits {
|
||||
/// 返回主窗口的 UTC 重置时间。
|
||||
pub fn primary_reset_at(self) -> Option<DateTime<Utc>> {
|
||||
self.primary.and_then(RateLimitWindow::reset_at_datetime)
|
||||
}
|
||||
|
||||
/// 返回主窗口的秒级窗口长度。
|
||||
pub fn primary_limit_window_seconds(self) -> Option<u64> {
|
||||
self.primary.and_then(RateLimitWindow::limit_window_seconds)
|
||||
}
|
||||
|
||||
/// 返回主窗口是否已经触顶。
|
||||
pub fn primary_limit_reached(self) -> bool {
|
||||
self.primary
|
||||
.is_some_and(|window| window.used_percent >= 100.0)
|
||||
}
|
||||
}
|
||||
|
||||
/// 从响应头对中解析限流信息。
|
||||
pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option<ParsedRateLimits> {
|
||||
let mut normalized = BTreeMap::new();
|
||||
@@ -412,22 +348,6 @@ pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(Strin
|
||||
headers
|
||||
}
|
||||
|
||||
/// 为退避秒数增加抖动,避免同一时间窗口内集中重试。
|
||||
pub fn cooldown_with_jitter(seconds: u64, factor_bps: u16) -> Duration {
|
||||
let factor_bps = u64::from(factor_bps);
|
||||
let variance = seconds.saturating_mul(factor_bps) / 10_000;
|
||||
if variance == 0 {
|
||||
return Duration::seconds(seconds.min(i64::MAX as u64) as i64);
|
||||
}
|
||||
let lower = seconds.saturating_sub(variance);
|
||||
let span = variance.saturating_mul(2).saturating_add(1);
|
||||
let offset = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| u64::from(duration.subsec_nanos()) % span)
|
||||
.unwrap_or(variance);
|
||||
Duration::seconds(lower.saturating_add(offset).min(i64::MAX as u64) as i64)
|
||||
}
|
||||
|
||||
fn number_field(value: &Value, field: &str) -> Option<u64> {
|
||||
value.get(field)?.as_u64()
|
||||
}
|
||||
@@ -510,10 +430,12 @@ fn unit_token(input: &str) -> Option<&str> {
|
||||
}
|
||||
|
||||
fn positive_seconds_ceil(seconds: f64) -> Option<u64> {
|
||||
if !seconds.is_finite() || seconds <= 0.0 || seconds > u64::MAX as f64 {
|
||||
if !seconds.is_finite() || seconds <= 0.0 {
|
||||
return None;
|
||||
}
|
||||
Some(seconds.ceil() as u64)
|
||||
Duration::try_from_secs_f64(seconds.ceil())
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
}
|
||||
|
||||
fn push_window_headers(
|
||||
|
||||
@@ -8,30 +8,6 @@ use crate::upstream::protocol::sse::encode_sse_event;
|
||||
|
||||
const REDACTED_PAYLOAD_VALUE: &str = "<redacted>";
|
||||
|
||||
/// WebSocket 消息包装。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct WsMessage<T> {
|
||||
/// 消息类型。
|
||||
#[serde(rename = "type")]
|
||||
pub message_type: String,
|
||||
/// 消息负载。
|
||||
pub payload: T,
|
||||
}
|
||||
|
||||
/// Codex WebSocket 聚合响应。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
|
||||
pub struct CodexWebSocketResponse {
|
||||
/// 聚合后的事件。
|
||||
pub events: Vec<Value>,
|
||||
}
|
||||
|
||||
/// Codex WebSocket 流响应。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
|
||||
pub struct CodexWebSocketStreamResponse {
|
||||
/// 聚合响应。
|
||||
pub response: CodexWebSocketResponse,
|
||||
}
|
||||
|
||||
/// WebSocket 握手审计快照。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct OpeningAuditSnapshot {
|
||||
@@ -61,13 +37,6 @@ pub struct PayloadAuditSnapshot {
|
||||
pub body: Value,
|
||||
}
|
||||
|
||||
/// WebSocket 一致性差异。
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct WsParityDiff {
|
||||
/// 差异说明。
|
||||
pub differences: Vec<String>,
|
||||
}
|
||||
|
||||
/// WebSocket audit artifact.
|
||||
#[derive(Debug, Clone, Serialize, PartialEq)]
|
||||
pub struct WebSocketAuditArtifact {
|
||||
@@ -92,24 +61,6 @@ pub struct WebSocketAuditErrorSnapshot {
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
/// Structured parity diff between two WebSocket audit artifacts.
|
||||
#[derive(Debug, Clone, Serialize, PartialEq)]
|
||||
pub struct WebSocketParityDiff {
|
||||
/// 逐字段差异。
|
||||
pub differences: Vec<WebSocketParityDifference>,
|
||||
}
|
||||
|
||||
/// Single WebSocket parity difference.
|
||||
#[derive(Debug, Clone, Serialize, PartialEq)]
|
||||
pub struct WebSocketParityDifference {
|
||||
/// 字段路径。
|
||||
pub path: String,
|
||||
/// 当前值。
|
||||
pub current: Value,
|
||||
/// 参考值。
|
||||
pub reference: Value,
|
||||
}
|
||||
|
||||
/// WebSocket 错误帧分类模式。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum WebSocketErrorClassificationProfile {
|
||||
@@ -186,20 +137,6 @@ struct ResponsesStreamEventShape {
|
||||
_content_index: Option<i64>,
|
||||
}
|
||||
|
||||
/// 将单条 WebSocket 消息编码为文本帧。
|
||||
pub fn encode_ws_message<T: Serialize>(
|
||||
message: &WsMessage<T>,
|
||||
) -> Result<String, serde_json::Error> {
|
||||
serde_json::to_string(message)
|
||||
}
|
||||
|
||||
/// 从文本帧解码单条 WebSocket 消息。
|
||||
pub fn decode_ws_message<T: for<'de> Deserialize<'de>>(
|
||||
frame: &str,
|
||||
) -> Result<WsMessage<T>, serde_json::Error> {
|
||||
serde_json::from_str(frame)
|
||||
}
|
||||
|
||||
/// 将一条公开 WebSocket JSON 事件编码为 SSE 帧。
|
||||
pub fn websocket_event_to_sse_frame(raw: &str) -> Option<String> {
|
||||
let event = websocket_event_type(raw)?;
|
||||
@@ -1081,51 +1018,6 @@ pub fn websocket_audit_artifact_from_attempt(
|
||||
}
|
||||
}
|
||||
|
||||
/// 按 parity 检查字段比较两个 WebSocket audit artifact。
|
||||
pub fn websocket_parity_diff(
|
||||
current: &WebSocketAuditArtifact,
|
||||
reference: &WebSocketAuditArtifact,
|
||||
) -> WebSocketParityDiff {
|
||||
let mut differences = Vec::new();
|
||||
push_websocket_parity_difference(
|
||||
&mut differences,
|
||||
"transport_mode",
|
||||
serde_json::json!(current.transport_mode),
|
||||
serde_json::json!(reference.transport_mode),
|
||||
);
|
||||
push_websocket_parity_difference(
|
||||
&mut differences,
|
||||
"fallback_allowed",
|
||||
serde_json::json!(current.fallback_allowed),
|
||||
serde_json::json!(reference.fallback_allowed),
|
||||
);
|
||||
push_websocket_parity_difference(
|
||||
&mut differences,
|
||||
"opening.request_line",
|
||||
opening_request_line(¤t.opening),
|
||||
opening_request_line(&reference.opening),
|
||||
);
|
||||
push_websocket_parity_difference(
|
||||
&mut differences,
|
||||
"opening.header_order",
|
||||
opening_header_order(¤t.opening),
|
||||
opening_header_order(&reference.opening),
|
||||
);
|
||||
push_websocket_parity_difference(
|
||||
&mut differences,
|
||||
"opening.sec_websocket_extensions",
|
||||
opening_header_value(¤t.opening, "Sec-WebSocket-Extensions"),
|
||||
opening_header_value(&reference.opening, "Sec-WebSocket-Extensions"),
|
||||
);
|
||||
push_websocket_parity_difference(
|
||||
&mut differences,
|
||||
"payload.top_level_keys",
|
||||
payload_top_level_keys(¤t.payload),
|
||||
payload_top_level_keys(&reference.payload),
|
||||
);
|
||||
WebSocketParityDiff { differences }
|
||||
}
|
||||
|
||||
fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static str {
|
||||
match transport_for_request(request) {
|
||||
CodexTransport::HttpSse => "http_sse",
|
||||
@@ -1134,69 +1026,6 @@ fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static st
|
||||
}
|
||||
}
|
||||
|
||||
fn push_websocket_parity_difference(
|
||||
differences: &mut Vec<WebSocketParityDifference>,
|
||||
path: &'static str,
|
||||
current: Value,
|
||||
reference: Value,
|
||||
) {
|
||||
if current != reference {
|
||||
differences.push(WebSocketParityDifference {
|
||||
path: path.to_string(),
|
||||
current,
|
||||
reference,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn opening_request_line(opening: &Option<OpeningAuditSnapshot>) -> Value {
|
||||
opening
|
||||
.as_ref()
|
||||
.map(|o| Value::String(o.request_line.clone()))
|
||||
.unwrap_or(Value::Null)
|
||||
}
|
||||
|
||||
fn opening_header_order(opening: &Option<OpeningAuditSnapshot>) -> Value {
|
||||
opening
|
||||
.as_ref()
|
||||
.map(|o| {
|
||||
Value::Array(
|
||||
o.header_order
|
||||
.iter()
|
||||
.map(|h| Value::String(h.clone()))
|
||||
.collect(),
|
||||
)
|
||||
})
|
||||
.unwrap_or(Value::Null)
|
||||
}
|
||||
|
||||
fn opening_header_value(opening: &Option<OpeningAuditSnapshot>, name: &str) -> Value {
|
||||
opening
|
||||
.as_ref()
|
||||
.map(|o| {
|
||||
o.header_order
|
||||
.iter()
|
||||
.find(|h| h.starts_with(&format!("{name}:")))
|
||||
.map(|h| Value::String(h.clone()))
|
||||
.unwrap_or(Value::Null)
|
||||
})
|
||||
.unwrap_or(Value::Null)
|
||||
}
|
||||
|
||||
fn payload_top_level_keys(payload: &Option<PayloadAuditSnapshot>) -> Value {
|
||||
payload
|
||||
.as_ref()
|
||||
.map(|p| {
|
||||
Value::Array(
|
||||
p.top_level_keys
|
||||
.iter()
|
||||
.map(|k| Value::String(k.clone()))
|
||||
.collect(),
|
||||
)
|
||||
})
|
||||
.unwrap_or(Value::Null)
|
||||
}
|
||||
|
||||
/// 生成 Responses WebSocket `response.create` payload。
|
||||
pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Value {
|
||||
let mut body = Map::new();
|
||||
|
||||
@@ -62,7 +62,7 @@ impl TokenRefresher for OpenAiTokenClient {
|
||||
if !status.is_success() {
|
||||
return Err(classify_refresh_failure(status, &body));
|
||||
}
|
||||
parse_token_pair(&body).map_err(|_| RefreshFailure::Transport)
|
||||
parse_token_pair(&body).map_err(|()| RefreshFailure::Transport)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -65,7 +65,7 @@ pub fn build_reqwest_client(force_http11: bool) -> Result<Client, CustomCaError>
|
||||
let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new()));
|
||||
let mut clients = cache
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
|
||||
if let Some(client) = clients.get(&cache_key) {
|
||||
return Ok(client.clone());
|
||||
@@ -302,15 +302,6 @@ pub struct CodexCompactResponse {
|
||||
pub rate_limit_headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
/// 上游模型端点探测结果。
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CodexConnectivityProbe {
|
||||
/// 请求的完整端点。
|
||||
pub endpoint: String,
|
||||
/// 返回状态码。
|
||||
pub status: StatusCode,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// CodexBackendClient
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -650,32 +641,6 @@ impl CodexBackendClient {
|
||||
})
|
||||
}
|
||||
|
||||
/// 探测模型端点连通性。
|
||||
pub async fn probe_connectivity(
|
||||
&self,
|
||||
context: CodexRequestContext<'_>,
|
||||
) -> CodexClientResult<CodexConnectivityProbe> {
|
||||
let headers = self.auxiliary_request_headers(context)?;
|
||||
let endpoint = format!(
|
||||
"{}/codex/models?client_version={}",
|
||||
self.base_url, self.fingerprint.app_version
|
||||
);
|
||||
let response = self.client.get(&endpoint).headers(headers).send().await?;
|
||||
let status = response.status();
|
||||
let retry_after_seconds = retry_after_seconds(response.headers(), None);
|
||||
if !status.is_success() {
|
||||
let body = read_capped_error_body(response).await?;
|
||||
return Err(CodexClientError::Upstream {
|
||||
status,
|
||||
retry_after_seconds: retry_after_seconds
|
||||
.or_else(|| retry_after_seconds_from_body(&body)),
|
||||
body,
|
||||
set_cookie_headers: Vec::new(),
|
||||
});
|
||||
}
|
||||
Ok(CodexConnectivityProbe { endpoint, status })
|
||||
}
|
||||
|
||||
/// 获取后端模型目录条目。
|
||||
async fn fetch_models_with_context(
|
||||
&self,
|
||||
|
||||
@@ -42,15 +42,6 @@ pub fn usage_endpoint_urls(base_url: &str) -> Vec<String> {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// 返回 usage 主请求路径。
|
||||
pub fn primary_usage_request_path(base_url: &str) -> String {
|
||||
let endpoint_path = usage_endpoint_paths(base_url)
|
||||
.into_iter()
|
||||
.next()
|
||||
.unwrap_or(CODEX_USAGE_API_PATH);
|
||||
endpoint_request_path(base_url, endpoint_path)
|
||||
}
|
||||
|
||||
fn usage_endpoint_paths(base_url: &str) -> Vec<&'static str> {
|
||||
if has_backend_api_base_path(base_url) {
|
||||
vec![WHAM_USAGE_PATH, CODEX_USAGE_PATH]
|
||||
|
||||
@@ -13,23 +13,17 @@ pub use self::{
|
||||
backend_transport_for_response_request, build_reqwest_client, is_banned_auth_signal,
|
||||
is_banned_upstream_error, CodexBackendClient, CodexBackendResponse, CodexBackendSseStream,
|
||||
CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexClientResult,
|
||||
CodexCompactResponse, CodexConnectivityProbe, CodexModelCatalogClient,
|
||||
CodexModelCatalogClientError, CodexModelCatalogRequest, CodexRateLimitHeaderUpdates,
|
||||
CodexRequestContext, CodexTurnStateUpdate,
|
||||
CodexCompactResponse, CodexModelCatalogClient, CodexModelCatalogClientError,
|
||||
CodexModelCatalogRequest, CodexRateLimitHeaderUpdates, CodexRequestContext,
|
||||
CodexTurnStateUpdate,
|
||||
},
|
||||
endpoints::{
|
||||
endpoint_request_path, endpoint_url, primary_usage_request_path, usage_endpoint_urls,
|
||||
CODEX_RESPONSES_COMPACT_PATH, CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, CODEX_USAGE_PATH,
|
||||
WHAM_USAGE_PATH,
|
||||
endpoint_request_path, endpoint_url, usage_endpoint_urls, CODEX_RESPONSES_COMPACT_PATH,
|
||||
CODEX_RESPONSES_PATH, CODEX_USAGE_API_PATH, CODEX_USAGE_PATH, WHAM_USAGE_PATH,
|
||||
},
|
||||
headers::{
|
||||
build_codex_base_headers, build_codex_headers, build_ordered_codex_base_headers,
|
||||
build_ordered_codex_headers, order_headers,
|
||||
},
|
||||
tls::{
|
||||
build_reqwest_client_with_custom_ca, custom_ca_env_cache_key,
|
||||
maybe_build_rustls_client_config_with_custom_ca, CustomCaError, CustomCaResult,
|
||||
CODEX_CA_CERT_ENV, SSL_CERT_FILE_ENV,
|
||||
},
|
||||
websocket_pool::{CodexWebSocketPool, CodexWebSocketPoolConfig, CodexWebSocketPoolKey},
|
||||
};
|
||||
|
||||
@@ -46,9 +46,10 @@ impl CodexBackendClient {
|
||||
Err(CodexClientError::Upstream {
|
||||
status: StatusCode::BAD_GATEWAY,
|
||||
retry_after_seconds: None,
|
||||
body: last_invalid_body
|
||||
.map(|body| format!("invalid usage response: {}", truncate_for_error(&body)))
|
||||
.unwrap_or_else(|| "usage endpoint is unavailable".to_string()),
|
||||
body: last_invalid_body.map_or_else(
|
||||
|| "usage endpoint is unavailable".to_string(),
|
||||
|body| format!("invalid usage response: {}", truncate_for_error(&body)),
|
||||
),
|
||||
set_cookie_headers: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -61,14 +61,6 @@ pub struct CodexWebSocketConnection {
|
||||
headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
/// 构造握手审计快照。
|
||||
pub fn opening_audit_snapshot(header_order: Vec<String>) -> OpeningAuditSnapshot {
|
||||
OpeningAuditSnapshot {
|
||||
header_order,
|
||||
..OpeningAuditSnapshot::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// 显式写入 WebSocket audit artifact。
|
||||
pub async fn write_websocket_audit_artifact_for_dir(
|
||||
dir: Option<&Path>,
|
||||
@@ -206,14 +198,6 @@ pub enum CodexWebSocketExchangeError {
|
||||
}
|
||||
|
||||
impl CodexWebSocketRequest {
|
||||
/// 构造 WebSocket 请求。
|
||||
pub fn new(connection: CodexWebSocketConnection, payload_text: String) -> Self {
|
||||
Self {
|
||||
connection,
|
||||
payload_text,
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回连接描述。
|
||||
pub fn connection(&self) -> &CodexWebSocketConnection {
|
||||
&self.connection
|
||||
@@ -282,11 +266,6 @@ impl CodexWebSocketConnection {
|
||||
&self.headers
|
||||
}
|
||||
|
||||
/// 生成真实 WebSocket opening 请求文本,用于 capture/audit parity。
|
||||
pub fn opening_request_text(&self) -> String {
|
||||
String::from_utf8_lossy(&opening_request_bytes(self)).into_owned()
|
||||
}
|
||||
|
||||
/// 生成打开握手审计快照。
|
||||
pub fn opening_audit_snapshot(&self) -> OpeningAuditSnapshot {
|
||||
OpeningAuditSnapshot {
|
||||
@@ -337,12 +316,12 @@ async fn connect_websocket(
|
||||
connection: &CodexWebSocketConnection,
|
||||
) -> Result<(CodexWsStream, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
|
||||
let request = websocket_handshake_request(connection)?;
|
||||
let connector = super::maybe_build_rustls_client_config_with_custom_ca()
|
||||
let connector = super::tls::maybe_build_rustls_client_config_with_custom_ca()
|
||||
.map_err(|error| tungstenite::Error::Io(std::io::Error::other(error)))?
|
||||
.map(Connector::Rustls);
|
||||
match connect_async_tls_with_config(request, Some(websocket_config()), false, connector).await {
|
||||
Ok((websocket, response)) => Ok((websocket, response)),
|
||||
Err(tungstenite::Error::Http(response)) => Err(websocket_opening_error(*response)),
|
||||
Err(tungstenite::Error::Http(response)) => Err(websocket_opening_error(response.as_ref())),
|
||||
Err(error) => Err(error.into()),
|
||||
}
|
||||
}
|
||||
@@ -371,7 +350,7 @@ fn websocket_config() -> WebSocketConfig {
|
||||
config
|
||||
}
|
||||
|
||||
fn websocket_opening_error(response: WsResponse<Option<Vec<u8>>>) -> CodexWebSocketExchangeError {
|
||||
fn websocket_opening_error(response: &WsResponse<Option<Vec<u8>>>) -> CodexWebSocketExchangeError {
|
||||
let status_code = response.status().as_u16();
|
||||
let body = response
|
||||
.body()
|
||||
@@ -393,20 +372,6 @@ fn websocket_opening_error(response: WsResponse<Option<Vec<u8>>>) -> CodexWebSoc
|
||||
}
|
||||
}
|
||||
|
||||
fn opening_request_bytes(connection: &CodexWebSocketConnection) -> Vec<u8> {
|
||||
let mut bytes = Vec::new();
|
||||
bytes.extend_from_slice(request_line_for_endpoint(connection.endpoint()).as_bytes());
|
||||
bytes.extend_from_slice(b"\r\n");
|
||||
for (name, value) in connection.headers() {
|
||||
bytes.extend_from_slice(name.as_bytes());
|
||||
bytes.extend_from_slice(b": ");
|
||||
bytes.extend_from_slice(value.as_bytes());
|
||||
bytes.extend_from_slice(b"\r\n");
|
||||
}
|
||||
bytes.extend_from_slice(b"\r\n");
|
||||
bytes
|
||||
}
|
||||
|
||||
fn request_line_for_endpoint(endpoint: &str) -> String {
|
||||
let path = reqwest::Url::parse(endpoint)
|
||||
.ok()
|
||||
@@ -673,7 +638,7 @@ async fn first_reused_stream_text(
|
||||
},
|
||||
);
|
||||
}
|
||||
Err(error) => return Err(reused_connection_died_before_first_frame(error)),
|
||||
Err(error) => return Err(reused_connection_died_before_first_frame(&error)),
|
||||
};
|
||||
match message {
|
||||
Message::Text(text) => return Ok(text.to_string()),
|
||||
@@ -685,7 +650,7 @@ async fn first_reused_stream_text(
|
||||
);
|
||||
}
|
||||
Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent),
|
||||
_ => continue,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -724,7 +689,7 @@ async fn collect_websocket_response(
|
||||
let message = match next_websocket_message(&mut websocket).await {
|
||||
Ok(message) => message,
|
||||
Err(error) if reused_connection && !saw_response_frame => {
|
||||
return Err(reused_connection_died_before_first_frame(error));
|
||||
return Err(reused_connection_died_before_first_frame(&error));
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
@@ -808,7 +773,7 @@ async fn collect_websocket_response(
|
||||
}
|
||||
|
||||
fn reused_connection_died_before_first_frame(
|
||||
error: CodexWebSocketExchangeError,
|
||||
error: &CodexWebSocketExchangeError,
|
||||
) -> CodexWebSocketExchangeError {
|
||||
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
|
||||
message: error.to_string(),
|
||||
|
||||
@@ -11,7 +11,7 @@ use tokio::{net::TcpStream, sync::Mutex, time::timeout};
|
||||
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
|
||||
|
||||
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
|
||||
const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60);
|
||||
const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55);
|
||||
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
|
||||
const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25);
|
||||
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5);
|
||||
@@ -113,46 +113,6 @@ impl CodexWebSocketPool {
|
||||
pool
|
||||
}
|
||||
|
||||
/// 使用默认配置构造连接池。
|
||||
pub fn with_default_max_age() -> Self {
|
||||
Self::with_config(CodexWebSocketPoolConfig::default())
|
||||
}
|
||||
|
||||
/// 使用最大存活时间与账号容量限制构造连接池。
|
||||
pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(WebSocketPoolState::default())),
|
||||
config: CodexWebSocketPoolConfig {
|
||||
max_age,
|
||||
max_per_account: max_per_account.max(1),
|
||||
maintenance_interval: None,
|
||||
ping_interval: None,
|
||||
liveness_timeout: None,
|
||||
..CodexWebSocketPoolConfig::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// 返回单账号最大连接数。
|
||||
pub fn max_per_account(&self) -> usize {
|
||||
self.config.max_per_account
|
||||
}
|
||||
|
||||
/// 返回连接最大存活时长。
|
||||
pub fn max_age(&self) -> Duration {
|
||||
self.config.max_age
|
||||
}
|
||||
|
||||
/// 判断当前账号是否还能打开新连接。
|
||||
pub fn permits_new_connection(&self, current_connections: usize) -> bool {
|
||||
current_connections < self.config.max_per_account
|
||||
}
|
||||
|
||||
/// 判断连接是否已达到回收年龄。
|
||||
pub fn should_recycle(&self, age: Duration) -> bool {
|
||||
age >= self.config.max_age
|
||||
}
|
||||
|
||||
pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
|
||||
let mut expired_connection = None;
|
||||
let acquire = {
|
||||
@@ -270,11 +230,6 @@ impl CodexWebSocketPool {
|
||||
close_pooled_connections(idle_connections).await;
|
||||
}
|
||||
|
||||
/// 维护 idle 连接:关闭过期/失活连接,并按需 ping 探活。
|
||||
pub async fn gc_sweep(&self) {
|
||||
self.maintain_idle_connections().await;
|
||||
}
|
||||
|
||||
/// 维护 idle 连接:关闭过期/失活连接,并按需 ping 探活。
|
||||
pub async fn maintain_idle_connections(&self) {
|
||||
let maintenance = self.take_idle_connections_for_maintenance().await;
|
||||
@@ -520,7 +475,7 @@ async fn probe_idle_connection(
|
||||
return Some(connection);
|
||||
}
|
||||
match timeout(ping_timeout, connection.websocket.next()).await {
|
||||
Ok(Some(Ok(Message::Close(_)))) | Ok(None) | Ok(Some(Err(_))) | Err(_) => {
|
||||
Ok(Some(Ok(Message::Close(_)) | Err(_)) | None) | Err(_) => {
|
||||
close_pooled_connection(connection).await;
|
||||
None
|
||||
}
|
||||
|
||||
@@ -1,78 +1,12 @@
|
||||
use super::*;
|
||||
use crate::support::jwt::unsigned_jwt;
|
||||
use wiremock::{
|
||||
matchers::{method, path},
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_export_should_return_native_account_tokens() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-accounts-export.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
seed_account(
|
||||
&pool,
|
||||
NewAccount {
|
||||
id: "acct_export".to_string(),
|
||||
email: Some("export@example.com".to_string()),
|
||||
account_id: Some("chatgpt_export".to_string()),
|
||||
user_id: Some("user_export".to_string()),
|
||||
label: Some("primary".to_string()),
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new("access-export".to_string().into()),
|
||||
refresh_token: Some(SecretString::new("refresh-export".to_string().into())),
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let config = test_config(url);
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/accounts/export?ids=acct_export")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_accounts_export")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["sourceFormat"], "native");
|
||||
assert_eq!(body["data"]["accounts"][0]["id"], "acct_export");
|
||||
assert_eq!(body["data"]["accounts"][0]["token"], "access-export");
|
||||
assert_eq!(
|
||||
body["data"]["accounts"][0]["refreshToken"],
|
||||
"refresh-export"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_import_should_store_native_account_tokens() {
|
||||
async fn admin_accounts_import_should_store_cpr_account_tokens() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-accounts-import.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
@@ -138,7 +72,7 @@ async fn admin_accounts_import_should_store_native_account_tokens() {
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(body["data"]["imported"], 1);
|
||||
assert_eq!(body["data"]["skipped"], 0);
|
||||
assert_eq!(body["data"]["sourceFormat"], "native");
|
||||
assert_eq!(body["data"]["sourceFormat"], "cpr");
|
||||
assert_eq!(stored.access_token.expose_secret(), access_token);
|
||||
assert_eq!(
|
||||
stored.refresh_token.unwrap().expose_secret(),
|
||||
@@ -192,7 +126,6 @@ async fn admin_accounts_import_should_fetch_wham_usage_for_current_openai_token_
|
||||
format!("{}/backend-api", server.uri()),
|
||||
)
|
||||
.await;
|
||||
let header = json!({"alg": "none", "typ": "JWT"});
|
||||
let payload = json!({
|
||||
"exp": 4_102_444_800i64,
|
||||
"https://api.openai.com/auth": {
|
||||
@@ -204,7 +137,7 @@ async fn admin_accounts_import_should_fetch_wham_usage_for_current_openai_token_
|
||||
"email_verified": true,
|
||||
},
|
||||
});
|
||||
let access_token = format!("{}.{}.", jwt_part(&header), jwt_part(&payload));
|
||||
let access_token = unsigned_jwt(&payload);
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
@@ -689,44 +622,3 @@ async fn admin_accounts_import_should_store_plain_tokens_and_list_sanitized_acco
|
||||
);
|
||||
assert!(list_body["data"]["items"][0].get("token").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_accounts_export_should_return_native_accounts_with_tokens_and_filter_ids() {
|
||||
let (app, _state, pool, _dir) =
|
||||
admin_accounts_test_app("admin-accounts-export-filter.sqlite", 117).await;
|
||||
seed_account(
|
||||
&pool,
|
||||
NewAccount {
|
||||
id: "acct_export_a".to_string(),
|
||||
email: Some("export-a@example.com".to_string()),
|
||||
account_id: Some("chatgpt-export-a".to_string()),
|
||||
user_id: Some("user-export-a".to_string()),
|
||||
label: None,
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new("access-acct_export_a".to_string().into()),
|
||||
refresh_token: Some(SecretString::new(
|
||||
"refresh-acct_export_a".to_string().into(),
|
||||
)),
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/accounts/export?ids=acct_export_a")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["accounts"].as_array().unwrap().len(), 1);
|
||||
assert_eq!(body["data"]["accounts"][0]["id"], "acct_export_a");
|
||||
}
|
||||
@@ -359,159 +359,6 @@ async fn admin_account_refresh_should_not_mark_valid_account_banned_when_refresh
|
||||
assert_eq!(stored.status, AccountStatus::Active);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_update_should_accept_batch_status_payload() {
|
||||
let (app, _state, pool, _dir) =
|
||||
admin_accounts_test_app("admin-account-batch-status-update.sqlite", 122).await;
|
||||
sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)")
|
||||
.bind("acct_batch_status_1").bind("batch-1@example.com").bind("cipher").bind("active")
|
||||
.bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z").execute(&pool).await.unwrap();
|
||||
sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)")
|
||||
.bind("acct_batch_status_2").bind("batch-2@example.com").bind("cipher").bind("active")
|
||||
.bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z").execute(&pool).await.unwrap();
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/accounts/update")
|
||||
.header("content-type", "application/json")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"ids": ["acct_batch_status_1", "acct_batch_status_2", "acct_missing"],
|
||||
"status": "disabled"
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["updated"], 2);
|
||||
assert_eq!(body["data"]["notFound"], json!(["acct_missing"]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_batch_status_route_should_not_exist() {
|
||||
let (_app, state, _pool, _dir) =
|
||||
admin_accounts_test_app("admin-account-batch-status-route.sqlite", 123).await;
|
||||
let admin_router = codex_proxy_rs::admin::router::router().with_state(state);
|
||||
|
||||
let response = admin_router
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/accounts/batch-status")
|
||||
.header("content-type", "application/json")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::from(
|
||||
json!({"ids": ["acct_batch_status_1"], "status": "disabled"}).to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_cookies_should_store_plain_cookie_header() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-account-cookies.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
seed_account(
|
||||
&pool,
|
||||
NewAccount {
|
||||
id: "acct_cookies".to_string(),
|
||||
email: Some("cookies@example.com".to_string()),
|
||||
account_id: None,
|
||||
user_id: None,
|
||||
label: None,
|
||||
plan_type: None,
|
||||
access_token: SecretString::new("access-cookies".to_string().into()),
|
||||
refresh_token: None,
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
let config = test_config(url);
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let set = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/accounts/cookies")
|
||||
.header("content-type", "application/json")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::from(
|
||||
json!({
|
||||
"id": "acct_cookies",
|
||||
"cookies":"cf_clearance=clear-secret; __cf_bm=bm-secret"
|
||||
})
|
||||
.to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(set.status(), StatusCode::OK);
|
||||
|
||||
let get = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/accounts/cookies?id=acct_cookies")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(get.status(), StatusCode::OK);
|
||||
let get_body = response_json(get).await;
|
||||
let cookies = get_body["data"]["cookies"].as_str().unwrap();
|
||||
assert!(cookies.contains("cf_clearance=clear-secret"));
|
||||
assert!(cookies.contains("__cf_bm=bm-secret"));
|
||||
|
||||
let stored_values: Vec<String> =
|
||||
sqlx::query_scalar("select value from account_cookies where account_id = ?1")
|
||||
.bind("acct_cookies")
|
||||
.fetch_all(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!stored_values.is_empty());
|
||||
assert!(stored_values.iter().any(|value| value == "clear-secret"));
|
||||
assert!(stored_values.iter().any(|value| value == "bm-secret"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_create_should_derive_claims_and_store_plain_tokens() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
@@ -59,16 +59,18 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
|
||||
admin_accounts_test_app("admin-accounts-stats.sqlite", 68).await;
|
||||
seed_usage_account(
|
||||
&pool,
|
||||
"acct_stats",
|
||||
"stats@example.com",
|
||||
"stats",
|
||||
"free",
|
||||
41,
|
||||
5,
|
||||
3_500_000,
|
||||
13_900,
|
||||
3_400_000,
|
||||
"2026-06-23T08:50:13Z",
|
||||
UsageAccountSeed {
|
||||
id: "acct_stats",
|
||||
email: "stats@example.com",
|
||||
label: "stats",
|
||||
plan_type: "free",
|
||||
request_count: 41,
|
||||
empty_response_count: 5,
|
||||
input_tokens: 3_500_000,
|
||||
output_tokens: 13_900,
|
||||
cached_tokens: 3_400_000,
|
||||
last_used_at: "2026-06-23T08:50:13Z",
|
||||
},
|
||||
)
|
||||
.await;
|
||||
sqlx::query("update accounts set quota_json = ?, quota_fetched_at = ? where id = ?")
|
||||
|
||||
+24
-106
@@ -1,19 +1,13 @@
|
||||
use std::{collections::BTreeMap, sync::Arc};
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use codex_proxy_rs::{
|
||||
admin::auth::service::SqliteAdminSessionStore,
|
||||
admin::keys::service::SqliteClientKeyStore,
|
||||
admin::monitoring::event_store::SqliteEventLogStore,
|
||||
config::types::AppConfig,
|
||||
config::types::{
|
||||
AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig,
|
||||
QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig,
|
||||
},
|
||||
infra::database::connect_sqlite,
|
||||
proxy::dispatch::session_affinity::SqliteSessionAffinityStore,
|
||||
runtime::services::{BackgroundTaskStores, Services},
|
||||
@@ -31,67 +25,45 @@ use serde_json::{json, Value};
|
||||
use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
mod import_export;
|
||||
use crate::support::jwt::unsigned_jwt;
|
||||
use crate::support::{admin::seed_admin_session, config::test_config, http::response_json};
|
||||
|
||||
mod importing;
|
||||
mod lifecycle;
|
||||
mod list;
|
||||
mod oauth;
|
||||
mod quota;
|
||||
mod testing;
|
||||
|
||||
async fn seed_admin_session(pool: &SqlitePool, session_id: &str) {
|
||||
sqlx::query(
|
||||
"insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind("admin_1")
|
||||
.bind("hash")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind("admin_1")
|
||||
.bind("2999-01-01T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::too_many_arguments,
|
||||
reason = "test fixture keeps usage rows explicit"
|
||||
)]
|
||||
async fn seed_usage_account(
|
||||
pool: &SqlitePool,
|
||||
id: &str,
|
||||
email: &str,
|
||||
label: &str,
|
||||
plan_type: &str,
|
||||
struct UsageAccountSeed<'a> {
|
||||
id: &'a str,
|
||||
email: &'a str,
|
||||
label: &'a str,
|
||||
plan_type: &'a str,
|
||||
request_count: i64,
|
||||
empty_response_count: i64,
|
||||
input_tokens: i64,
|
||||
output_tokens: i64,
|
||||
cached_tokens: i64,
|
||||
last_used_at: &str,
|
||||
) {
|
||||
last_used_at: &'a str,
|
||||
}
|
||||
|
||||
async fn seed_usage_account(pool: &SqlitePool, seed: UsageAccountSeed<'_>) {
|
||||
sqlx::query("insert into accounts (id, email, label, plan_type, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, 'active', ?, ?)")
|
||||
.bind(id).bind(email).bind(label).bind(plan_type).bind("access-token")
|
||||
.bind(seed.id).bind(seed.email).bind(seed.label).bind(seed.plan_type).bind("access-token")
|
||||
.bind("2026-06-11T00:00:00Z").bind("2026-06-11T00:00:00Z")
|
||||
.execute(pool).await.unwrap();
|
||||
sqlx::query("insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)")
|
||||
.bind(id).bind(request_count).bind(empty_response_count).bind(input_tokens).bind(output_tokens).bind(cached_tokens).bind(last_used_at)
|
||||
.bind(seed.id)
|
||||
.bind(seed.request_count)
|
||||
.bind(seed.empty_response_count)
|
||||
.bind(seed.input_tokens)
|
||||
.bind(seed.output_tokens)
|
||||
.bind(seed.cached_tokens)
|
||||
.bind(seed.last_used_at)
|
||||
.execute(pool).await.unwrap();
|
||||
}
|
||||
|
||||
async fn response_json(response: axum::response::Response) -> Value {
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
|
||||
async fn post_admin_account(app: &axum::Router, payload: Value) -> axum::response::Response {
|
||||
app.clone()
|
||||
.oneshot(
|
||||
@@ -215,7 +187,6 @@ fn test_jwt_with_exp(
|
||||
plan_type: Option<&str>,
|
||||
exp: i64,
|
||||
) -> String {
|
||||
let header = json!({"alg": "none", "typ": "JWT"});
|
||||
let payload = json!({
|
||||
"exp": exp,
|
||||
"https://api.openai.com/auth": {
|
||||
@@ -225,58 +196,5 @@ fn test_jwt_with_exp(
|
||||
},
|
||||
"https://api.openai.com/profile": { "email": email }
|
||||
});
|
||||
format!("{}.{}.", jwt_part(&header), jwt_part(&payload))
|
||||
}
|
||||
|
||||
fn jwt_part(value: &Value) -> String {
|
||||
URL_SAFE_NO_PAD.encode(serde_json::to_vec(value).unwrap())
|
||||
}
|
||||
|
||||
fn test_config(database_url: String) -> AppConfig {
|
||||
AppConfig {
|
||||
server: ServerConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 0,
|
||||
},
|
||||
api: ApiConfig {
|
||||
base_url: "https://chatgpt.com/backend-api".to_string(),
|
||||
},
|
||||
model_aliases: BTreeMap::new(),
|
||||
model_account_routes: BTreeMap::new(),
|
||||
auth: AuthConfig {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_enabled: true,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 3,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
tier_priority: Vec::new(),
|
||||
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
|
||||
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
|
||||
},
|
||||
quota: QuotaConfig {
|
||||
refresh_interval_minutes: 5,
|
||||
warning_thresholds: QuotaWarningThresholds {
|
||||
primary: vec![80, 90],
|
||||
secondary: vec![80, 90],
|
||||
},
|
||||
},
|
||||
database: DatabaseConfig { url: database_url },
|
||||
tls: TlsConfig {
|
||||
force_http11: false,
|
||||
},
|
||||
ws_pool: WebSocketPoolConfig::default(),
|
||||
fingerprint: Default::default(),
|
||||
admin: AdminConfig {
|
||||
session_ttl_minutes: 1440,
|
||||
session_cleanup_interval_secs: 3600,
|
||||
default_username: "admin".to_string(),
|
||||
default_password: "admin".to_string(),
|
||||
},
|
||||
logging: LoggingConfig {
|
||||
directory: "logs".to_string(),
|
||||
retention_days: 14,
|
||||
enabled: false,
|
||||
},
|
||||
}
|
||||
unsigned_jwt(&payload)
|
||||
}
|
||||
|
||||
@@ -4,206 +4,6 @@ use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_stats_should_return_page_and_summary() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-usage.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
sqlx::query("insert into accounts (id, email, label, plan_type, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)")
|
||||
.bind("acct_usage").bind("usage@example.com").bind("primary").bind("plus")
|
||||
.bind("cipher").bind("active").bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z")
|
||||
.execute(&pool).await.unwrap();
|
||||
sqlx::query("insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, 3, 1, 21, 8, 5, ?)")
|
||||
.bind("acct_usage").bind("2026-06-18T00:10:00Z").execute(&pool).await.unwrap();
|
||||
let config = test_config(url);
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let page = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/usage?page=1&pageSize=10")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_usage_page")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.status(), StatusCode::OK);
|
||||
let page_body = response_json(page).await;
|
||||
assert_eq!(page_body["data"]["items"][0]["accountId"], "acct_usage");
|
||||
assert_eq!(page_body["data"]["items"][0]["requestCount"], 3);
|
||||
assert_eq!(page_body["data"]["page"]["page"], 1);
|
||||
assert_eq!(page_body["data"]["page"]["pageSize"], 10);
|
||||
assert_eq!(page_body["data"]["page"]["total"], 1);
|
||||
assert_eq!(page_body["data"]["page"]["totalPages"], 1);
|
||||
|
||||
let summary = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/usage/summary")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_usage_summary")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(summary.status(), StatusCode::OK);
|
||||
assert_eq!(response_json(summary).await["data"]["accountCount"], 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_stats_should_require_admin_session_cookie() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-usage-auth.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
let config = test_config(url);
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/usage")
|
||||
.header("x-request-id", "req_usage_auth")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
assert_eq!(response_json(response).await["code"], 40101);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_usage_stats_should_cursor_page_account_usage() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-usage-cursor.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
seed_usage_account(
|
||||
&pool,
|
||||
"acct_a",
|
||||
"a@example.com",
|
||||
"primary",
|
||||
"plus",
|
||||
3,
|
||||
0,
|
||||
12,
|
||||
5,
|
||||
1,
|
||||
"2026-06-11T00:00:00Z",
|
||||
)
|
||||
.await;
|
||||
seed_usage_account(
|
||||
&pool,
|
||||
"acct_b",
|
||||
"b@example.com",
|
||||
"backup",
|
||||
"free",
|
||||
2,
|
||||
1,
|
||||
7,
|
||||
3,
|
||||
2,
|
||||
"2026-06-11T00:01:00Z",
|
||||
)
|
||||
.await;
|
||||
let config = test_config(url);
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let first = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/usage?limit=1")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_usage_cursor")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let first_body = response_json(first).await;
|
||||
assert_eq!(first_body["code"], 200);
|
||||
assert_eq!(first_body["data"]["items"].as_array().unwrap().len(), 1);
|
||||
let cursor = first_body["data"]["page"]["nextCursor"].as_str().unwrap();
|
||||
|
||||
let second = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri(format!("/api/admin/usage?limit=1&cursor={cursor}"))
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(second.status(), StatusCode::OK);
|
||||
assert_eq!(
|
||||
response_json(second).await["data"]["items"][0]["accountId"],
|
||||
"acct_a"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_quota_should_send_usage_cookie() {
|
||||
let server = MockServer::start().await;
|
||||
@@ -247,7 +47,10 @@ async fn admin_account_quota_should_send_usage_cookie() {
|
||||
)
|
||||
.await;
|
||||
SqliteCookieStore::new(pool.clone())
|
||||
.set_cookie_header("acct_quota_cookie", "cf_clearance=admin-quota")
|
||||
.capture_set_cookie(
|
||||
"acct_quota_cookie",
|
||||
"cf_clearance=admin-quota; Domain=.chatgpt.com; Path=/",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
@@ -275,168 +78,3 @@ async fn admin_account_quota_should_send_usage_cookie() {
|
||||
(StatusCode::OK, Some("cf_clearance=admin-quota"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_health_check_should_send_probe_cookie() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/backend-api/codex/models"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"models": []
|
||||
})))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url(
|
||||
"admin-account-health-cookie.sqlite",
|
||||
89,
|
||||
format!("{}/backend-api", server.uri()),
|
||||
)
|
||||
.await;
|
||||
seed_account(
|
||||
&pool,
|
||||
NewAccount {
|
||||
id: "acct_health_cookie".to_string(),
|
||||
email: Some("health-cookie@example.com".to_string()),
|
||||
account_id: Some("chatgpt-health-cookie".to_string()),
|
||||
user_id: None,
|
||||
label: None,
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new("access-health-cookie".to_string().into()),
|
||||
refresh_token: None,
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
SqliteCookieStore::new(pool.clone())
|
||||
.set_cookie_header("acct_health_cookie", "cf_clearance=admin-health")
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/accounts/health-check")
|
||||
.header("content-type", "application/json")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_admin_health_cookie")
|
||||
.body(Body::from(
|
||||
json!({ "ids": ["acct_health_cookie"] }).to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let requests = server.received_requests().await.unwrap();
|
||||
let cookie_header = requests
|
||||
.iter()
|
||||
.find(|request| request.url.path() == "/backend-api/codex/models")
|
||||
.and_then(|request| request.headers.get("cookie"))
|
||||
.and_then(|value| value.to_str().ok());
|
||||
|
||||
assert_eq!(
|
||||
(response.status(), cookie_header),
|
||||
(StatusCode::OK, Some("cf_clearance=admin-health"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_account_quota_warnings_should_return_threshold_matches_from_cached_quota() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-account-quota-warnings.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
seed_account(
|
||||
&pool,
|
||||
NewAccount {
|
||||
id: "acct_warn".to_string(),
|
||||
email: Some("warn@example.com".to_string()),
|
||||
account_id: None,
|
||||
user_id: None,
|
||||
label: None,
|
||||
plan_type: None,
|
||||
access_token: SecretString::new("access-warn".to_string().into()),
|
||||
refresh_token: None,
|
||||
access_token_expires_at: None,
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
sqlx::query(
|
||||
"update accounts set quota_json = ?, quota_fetched_at = ?, updated_at = ? where id = ?",
|
||||
)
|
||||
.bind(
|
||||
json!({
|
||||
"snapshots": [{
|
||||
"source": "core",
|
||||
"limit_name": null,
|
||||
"metered_feature": null,
|
||||
"allowed": true,
|
||||
"limit_reached": false,
|
||||
"blocked": false,
|
||||
"primary": {
|
||||
"used_percent": 85,
|
||||
"remaining_percent": 15,
|
||||
"reset_at": 1770000100,
|
||||
"window_minutes": 300,
|
||||
"limit_reached": false
|
||||
},
|
||||
"secondary": {
|
||||
"used_percent": 91,
|
||||
"remaining_percent": 9,
|
||||
"reset_at": 1770000200,
|
||||
"window_minutes": 10080,
|
||||
"limit_reached": false
|
||||
}
|
||||
}],
|
||||
"monthly_limit": null,
|
||||
"credits": null,
|
||||
"spend_control": null
|
||||
})
|
||||
.to_string(),
|
||||
)
|
||||
.bind("2026-06-13T00:00:00Z")
|
||||
.bind("2026-06-13T00:00:00Z")
|
||||
.bind("acct_warn")
|
||||
.execute(&pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let config = test_config(url);
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/accounts/quota-warnings")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
let warnings = body["data"]["warnings"].as_array().unwrap();
|
||||
assert_eq!(warnings.len(), 2);
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use super::*;
|
||||
use axum::body::to_bytes;
|
||||
use codex_proxy_rs::upstream::protocol::sse::{encode_sse_event, parse_sse_events};
|
||||
use wiremock::{
|
||||
matchers::{method, path},
|
||||
@@ -34,7 +35,7 @@ fn test_events(body: &str) -> Vec<Value> {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn account_test_models_should_return_upstream_models_only() {
|
||||
async fn account_models_should_return_upstream_models_only() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/backend-api/codex/models"))
|
||||
@@ -53,7 +54,7 @@ async fn account_test_models_should_return_upstream_models_only() {
|
||||
.mount(&server)
|
||||
.await;
|
||||
let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url(
|
||||
"admin-account-test-models.sqlite",
|
||||
"admin-account-models.sqlite",
|
||||
91,
|
||||
format!("{}/backend-api", server.uri()),
|
||||
)
|
||||
@@ -66,7 +67,7 @@ async fn account_test_models_should_return_upstream_models_only() {
|
||||
.method("GET")
|
||||
.uri("/api/admin/accounts/models?id=acct_test")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_account_test_models")
|
||||
.header("x-request-id", "req_account_models")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
pub mod password;
|
||||
pub mod session;
|
||||
pub mod session_routes;
|
||||
pub mod session_store;
|
||||
pub(crate) mod password;
|
||||
pub(crate) mod session;
|
||||
pub(crate) mod session_store;
|
||||
|
||||
+17
-13
@@ -1,5 +1,5 @@
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use chrono::Utc;
|
||||
@@ -7,7 +7,7 @@ use codex_proxy_rs::infra::{database::connect_sqlite, identity::hash_admin_passw
|
||||
use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::config::test_config;
|
||||
use crate::support::{config::test_config, http::response_json};
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_login_should_issue_http_only_session_cookie() {
|
||||
@@ -58,9 +58,7 @@ async fn admin_login_should_issue_http_only_session_cookie() {
|
||||
.uri("/api/admin/login")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-request-id", "req_login")
|
||||
.body(Body::from(
|
||||
r#"{"username":"admin","password":"correct-password"}"#,
|
||||
))
|
||||
.body(Body::from(r#"{"password":"correct-password"}"#))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
@@ -80,6 +78,19 @@ async fn admin_login_should_issue_http_only_session_cookie() {
|
||||
assert!(cookie.contains("SameSite=Lax"));
|
||||
assert_eq!(body["code"], 200);
|
||||
assert!(body["data"]["expiresAt"].is_string());
|
||||
|
||||
let logs_response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/logs")
|
||||
.header("cookie", cookie.split(';').next().unwrap())
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(logs_response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -132,9 +143,7 @@ async fn admin_login_should_reject_client_api_key_as_password_or_authorization()
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer cpr_not_an_admin_session")
|
||||
.header("x-request-id", "req_login_bad")
|
||||
.body(Body::from(
|
||||
r#"{"username":"admin","password":"cpr_not_an_admin_password"}"#,
|
||||
))
|
||||
.body(Body::from(r#"{"password":"cpr_not_an_admin_password"}"#))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
@@ -162,8 +171,3 @@ async fn seed_admin_user(pool: &SqlitePool, password: &str) {
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn response_json(response: axum::response::Response) -> serde_json::Value {
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
|
||||
@@ -1,208 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use chrono::Utc;
|
||||
use codex_proxy_rs::{
|
||||
admin::auth::service::SqliteAdminSessionStore,
|
||||
admin::keys::service::SqliteClientKeyStore,
|
||||
admin::monitoring::event_store::SqliteEventLogStore,
|
||||
config::types::AppConfig,
|
||||
config::types::{
|
||||
AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig,
|
||||
QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig,
|
||||
},
|
||||
infra::{database::connect_sqlite, identity::hash_admin_password},
|
||||
proxy::dispatch::session_affinity::SqliteSessionAffinityStore,
|
||||
runtime::services::{BackgroundTaskStores, Services},
|
||||
runtime::state::AppState,
|
||||
upstream::accounts::token_refresh::RefreshLeaseStore,
|
||||
upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore},
|
||||
upstream::fingerprint::FingerprintRepository,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_login_should_issue_http_only_session_cookie() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-login.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_user(&pool, "correct-password").await;
|
||||
let config = test_config(url);
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/login")
|
||||
.header("content-type", "application/json")
|
||||
.header("x-request-id", "req_login")
|
||||
.body(Body::from(r#"{"password":"correct-password"}"#))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let cookie = response
|
||||
.headers()
|
||||
.get("set-cookie")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(ToString::to_string);
|
||||
let body = response_json(response).await;
|
||||
let cookie = cookie.expect("login should set admin session cookie");
|
||||
assert!(cookie.starts_with("cpr_admin_session="));
|
||||
assert!(cookie.contains("HttpOnly"));
|
||||
assert_eq!(body["code"], 200);
|
||||
|
||||
let logs_response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri("/api/admin/logs")
|
||||
.header("cookie", cookie.split(';').next().unwrap())
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(logs_response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_login_should_reject_client_api_key_as_password_or_authorization() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-login-rejects-client-key.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_user(&pool, "correct-password").await;
|
||||
let config = test_config(url);
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/login")
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer cpr_not_an_admin_session")
|
||||
.header("x-request-id", "req_login_bad")
|
||||
.body(Body::from(r#"{"password":"cpr_not_an_admin_password"}"#))
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
assert!(response.headers().get("set-cookie").is_none());
|
||||
assert_eq!(response_json(response).await["code"], 40102);
|
||||
}
|
||||
|
||||
async fn seed_admin_user(pool: &SqlitePool, password: &str) {
|
||||
let now = Utc::now().to_rfc3339();
|
||||
let hash = hash_admin_password(password).unwrap();
|
||||
sqlx::query(
|
||||
"insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind("admin_1")
|
||||
.bind(hash)
|
||||
.bind(&now)
|
||||
.bind(&now)
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn response_json(response: axum::response::Response) -> Value {
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
|
||||
fn test_config(database_url: String) -> AppConfig {
|
||||
let mut aliases = BTreeMap::new();
|
||||
aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string());
|
||||
AppConfig {
|
||||
server: ServerConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 0,
|
||||
},
|
||||
api: ApiConfig {
|
||||
base_url: "https://chatgpt.com/backend-api".to_string(),
|
||||
},
|
||||
model_aliases: aliases,
|
||||
model_account_routes: BTreeMap::new(),
|
||||
auth: AuthConfig {
|
||||
refresh_margin_seconds: 240,
|
||||
refresh_enabled: true,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 4,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
tier_priority: vec!["team".to_string(), "plus".to_string()],
|
||||
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
|
||||
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
|
||||
},
|
||||
quota: QuotaConfig {
|
||||
refresh_interval_minutes: 5,
|
||||
warning_thresholds: QuotaWarningThresholds {
|
||||
primary: vec![80, 90],
|
||||
secondary: vec![70, 95],
|
||||
},
|
||||
},
|
||||
database: DatabaseConfig { url: database_url },
|
||||
tls: TlsConfig {
|
||||
force_http11: false,
|
||||
},
|
||||
ws_pool: WebSocketPoolConfig::default(),
|
||||
fingerprint: Default::default(),
|
||||
admin: AdminConfig {
|
||||
session_ttl_minutes: 60,
|
||||
session_cleanup_interval_secs: 3600,
|
||||
default_username: "admin".to_string(),
|
||||
default_password: "admin".to_string(),
|
||||
},
|
||||
logging: LoggingConfig {
|
||||
directory: "logs".to_string(),
|
||||
retention_days: 14,
|
||||
enabled: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,5 @@
|
||||
use super::*;
|
||||
|
||||
// Ignored: duplicate of admin::keys::service::authorization test that passes
|
||||
#[tokio::test]
|
||||
async fn admin_client_keys_route_should_create_list_and_authorize_v1_requests() {
|
||||
let (app, _dir) = admin_client_key_test_app("admin-client-key-create.sqlite").await;
|
||||
@@ -71,7 +70,6 @@ async fn admin_client_keys_route_should_create_list_and_authorize_v1_requests()
|
||||
assert_eq!(models_response.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
// Ignored: duplicate of admin::keys::service::authorization test that passes
|
||||
#[tokio::test]
|
||||
async fn admin_client_keys_status_should_disable_and_enable_authorization() {
|
||||
let (app, _dir) = admin_client_key_test_app("admin-client-key-status.sqlite").await;
|
||||
@@ -141,7 +139,6 @@ async fn admin_client_keys_status_should_disable_and_enable_authorization() {
|
||||
assert_eq!(accepted.status(), StatusCode::OK);
|
||||
}
|
||||
|
||||
// Ignored: duplicate of admin::keys::service::authorization test that passes
|
||||
#[tokio::test]
|
||||
async fn admin_client_keys_delete_should_remove_authorization() {
|
||||
let (app, _dir) = admin_client_key_test_app("admin-client-key-delete.sqlite").await;
|
||||
|
||||
@@ -1,25 +0,0 @@
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_client_keys_export_should_return_metadata_without_secret_material() {
|
||||
let (app, _dir) = admin_client_key_test_app("admin-client-key-export.sqlite").await;
|
||||
let (key_id, _key) = create_admin_client_key(&app, "export-key").await;
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("GET")
|
||||
.uri(format!("/api/admin/keys/export?ids={key_id}"))
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_api_key_export")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["apiKeys"][0]["id"], key_id);
|
||||
assert_eq!(body["data"]["apiKeys"][0]["name"], "export-key");
|
||||
assert!(body["data"]["apiKeys"][0].get("key").is_none());
|
||||
}
|
||||
@@ -57,7 +57,6 @@ async fn admin_client_keys_label_should_update_clear_and_validate_label() {
|
||||
assert_eq!(missing.status(), StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
// Ignored: duplicate of admin::keys::service::lifecycle test that passes
|
||||
#[tokio::test]
|
||||
async fn admin_client_keys_batch_delete_should_remove_found_keys_and_report_missing_ids() {
|
||||
let (app, _dir) = admin_client_key_test_app("admin-client-key-delete.sqlite").await;
|
||||
|
||||
+4
-88
@@ -1,18 +1,11 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use codex_proxy_rs::{
|
||||
admin::auth::service::SqliteAdminSessionStore,
|
||||
admin::keys::service::SqliteClientKeyStore,
|
||||
admin::monitoring::event_store::SqliteEventLogStore,
|
||||
config::types::AppConfig,
|
||||
config::types::{
|
||||
AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig,
|
||||
QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig,
|
||||
},
|
||||
infra::database::connect_sqlite,
|
||||
proxy::dispatch::session_affinity::SqliteSessionAffinityStore,
|
||||
runtime::services::{BackgroundTaskStores, Services},
|
||||
@@ -21,12 +14,12 @@ use codex_proxy_rs::{
|
||||
upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore},
|
||||
upstream::fingerprint::FingerprintRepository,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use sqlx::SqlitePool;
|
||||
use serde_json::json;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{admin::seed_admin_session, config::test_config, http::response_json};
|
||||
|
||||
mod authorization;
|
||||
mod import_export;
|
||||
mod lifecycle;
|
||||
mod store;
|
||||
|
||||
@@ -79,80 +72,3 @@ async fn create_admin_client_key(app: &axum::Router, name: &str) -> (String, Str
|
||||
body["data"]["key"].as_str().unwrap().to_string(),
|
||||
)
|
||||
}
|
||||
|
||||
async fn seed_admin_session(pool: &SqlitePool, session_id: &str) {
|
||||
sqlx::query(
|
||||
"insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind("admin_1")
|
||||
.bind("hash")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind("admin_1")
|
||||
.bind("2999-01-01T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn response_json(response: axum::response::Response) -> Value {
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
|
||||
fn test_config(database_url: String) -> AppConfig {
|
||||
AppConfig {
|
||||
server: ServerConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 0,
|
||||
},
|
||||
api: ApiConfig {
|
||||
base_url: "https://chatgpt.com/backend-api".to_string(),
|
||||
},
|
||||
model_aliases: BTreeMap::new(),
|
||||
model_account_routes: BTreeMap::new(),
|
||||
auth: AuthConfig {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_enabled: true,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 3,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
tier_priority: Vec::new(),
|
||||
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
|
||||
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
|
||||
},
|
||||
quota: QuotaConfig {
|
||||
refresh_interval_minutes: 5,
|
||||
warning_thresholds: QuotaWarningThresholds {
|
||||
primary: vec![80, 90],
|
||||
secondary: vec![80, 90],
|
||||
},
|
||||
},
|
||||
database: DatabaseConfig { url: database_url },
|
||||
tls: TlsConfig {
|
||||
force_http11: false,
|
||||
},
|
||||
ws_pool: WebSocketPoolConfig::default(),
|
||||
fingerprint: Default::default(),
|
||||
admin: AdminConfig {
|
||||
session_ttl_minutes: 1440,
|
||||
session_cleanup_interval_secs: 3600,
|
||||
default_username: "admin".to_string(),
|
||||
default_password: "admin".to_string(),
|
||||
},
|
||||
logging: LoggingConfig {
|
||||
directory: "logs".to_string(),
|
||||
retention_days: 14,
|
||||
enabled: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
+6
-7
@@ -1,7 +1,6 @@
|
||||
pub mod accounts;
|
||||
pub mod auth;
|
||||
pub mod keys;
|
||||
pub mod models;
|
||||
pub mod monitoring;
|
||||
pub mod response;
|
||||
pub mod settings;
|
||||
pub(crate) mod accounts;
|
||||
pub(crate) mod auth;
|
||||
pub(crate) mod keys;
|
||||
pub(crate) mod monitoring;
|
||||
pub(crate) mod response;
|
||||
pub(crate) mod settings;
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
pub mod routes;
|
||||
@@ -1,247 +0,0 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use chrono::{Duration, Utc};
|
||||
use codex_proxy_rs::{
|
||||
admin::auth::service::SqliteAdminSessionStore,
|
||||
admin::keys::service::SqliteClientKeyStore,
|
||||
admin::monitoring::event_store::SqliteEventLogStore,
|
||||
config::types::AppConfig,
|
||||
config::types::{
|
||||
AdminConfig, ApiConfig, AuthConfig, DatabaseConfig, LoggingConfig, QuotaConfig,
|
||||
QuotaWarningThresholds, ServerConfig, TlsConfig, WebSocketPoolConfig,
|
||||
},
|
||||
infra::database::connect_sqlite,
|
||||
proxy::dispatch::session_affinity::SqliteSessionAffinityStore,
|
||||
runtime::services::{BackgroundTaskStores, Services},
|
||||
runtime::state::AppState,
|
||||
upstream::accounts::{
|
||||
cookies::SqliteCookieStore,
|
||||
model::AccountStatus,
|
||||
store::{NewAccount, SqliteAccountStore},
|
||||
token_refresh::RefreshLeaseStore,
|
||||
},
|
||||
upstream::fingerprint::FingerprintRepository,
|
||||
};
|
||||
use secrecy::SecretString;
|
||||
use serde_json::{json, Value};
|
||||
use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
use wiremock::{
|
||||
matchers::{method, path},
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn admin_refresh_models_should_require_admin_session() {
|
||||
let upstream = MockServer::start().await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-models-auth.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
let config = test_config(url, upstream.uri());
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: SqliteAccountStore::new(pool.clone()),
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state);
|
||||
|
||||
let response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/models/refresh")
|
||||
.header("x-request-id", "req_models_auth")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
|
||||
assert_eq!(response_json(response).await["code"], 40101);
|
||||
}
|
||||
|
||||
// Ignored: requires ModelService with snapshot_store and upstream_client configured
|
||||
#[tokio::test]
|
||||
|
||||
async fn admin_refresh_models_should_store_snapshots_for_distinct_active_plans() {
|
||||
let upstream = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/codex/models"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
|
||||
"models": [{
|
||||
"slug": "gpt-refresh", "display_name": "GPT Refresh",
|
||||
"description": "Backend refreshed model",
|
||||
"default_reasoning_level": "minimal",
|
||||
"supported_reasoning_levels": [{"effort": "minimal", "description": "Minimal"}],
|
||||
"contextWindow": 200000
|
||||
}]
|
||||
})))
|
||||
.expect(2)
|
||||
.mount(&upstream)
|
||||
.await;
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let db = dir.path().join("admin-models-refresh.sqlite");
|
||||
let url = format!("sqlite://{}", db.display());
|
||||
let pool = connect_sqlite(&url).await.unwrap();
|
||||
seed_admin_session(&pool, "session_1").await;
|
||||
let config = test_config(url, upstream.uri());
|
||||
let store = SqliteAccountStore::new(pool.clone());
|
||||
store
|
||||
.insert(NewAccount {
|
||||
id: "acct_plus".to_string(),
|
||||
email: Some("plus@test.com".to_string()),
|
||||
account_id: Some("chatgpt-plus".to_string()),
|
||||
user_id: Some("user_plus".to_string()),
|
||||
label: None,
|
||||
plan_type: Some("plus".to_string()),
|
||||
access_token: SecretString::new("token-plus".to_string().into()),
|
||||
refresh_token: Some(SecretString::new("refresh-plus".to_string().into())),
|
||||
access_token_expires_at: Some(Utc::now() + Duration::hours(1)),
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
store
|
||||
.insert(NewAccount {
|
||||
id: "acct_team".to_string(),
|
||||
email: Some("team@test.com".to_string()),
|
||||
account_id: Some("chatgpt-team".to_string()),
|
||||
user_id: Some("user_team".to_string()),
|
||||
label: None,
|
||||
plan_type: Some("team".to_string()),
|
||||
access_token: SecretString::new("token-team".to_string().into()),
|
||||
refresh_token: Some(SecretString::new("refresh-team".to_string().into())),
|
||||
access_token_expires_at: Some(Utc::now() + Duration::hours(1)),
|
||||
status: AccountStatus::Active,
|
||||
added_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let stores = BackgroundTaskStores {
|
||||
accounts: store,
|
||||
admin_sessions: SqliteAdminSessionStore::new(pool.clone()),
|
||||
cookies: SqliteCookieStore::new(pool.clone()),
|
||||
fingerprints: FingerprintRepository::new(pool.clone()),
|
||||
session_affinity: SqliteSessionAffinityStore::new(pool.clone()),
|
||||
refresh_leases: RefreshLeaseStore::new(pool.clone()),
|
||||
client_keys: SqliteClientKeyStore::new(pool.clone()),
|
||||
event_logs: SqliteEventLogStore::new(pool.clone()),
|
||||
};
|
||||
let fingerprint = crate::support::fingerprint::test_fingerprint();
|
||||
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
|
||||
let state = AppState {
|
||||
config,
|
||||
services: (*services).clone(),
|
||||
};
|
||||
let app = codex_proxy_rs::http::router::router().with_state(state.clone());
|
||||
|
||||
let response = app
|
||||
.clone()
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.method("POST")
|
||||
.uri("/api/admin/models/refresh")
|
||||
.header("cookie", "cpr_admin_session=session_1")
|
||||
.header("x-request-id", "req_models_refresh")
|
||||
.body(Body::empty())
|
||||
.unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let body = response_json(response).await;
|
||||
assert_eq!(body["data"]["refreshedPlans"], 2);
|
||||
assert_eq!(body["data"]["modelCount"], 2);
|
||||
}
|
||||
|
||||
async fn seed_admin_session(pool: &SqlitePool, session_id: &str) {
|
||||
sqlx::query(
|
||||
"insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind("admin_1")
|
||||
.bind("hash")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind("admin_1")
|
||||
.bind("2999-01-01T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn response_json(response: axum::response::Response) -> Value {
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
|
||||
fn test_config(database_url: String, base_url: String) -> AppConfig {
|
||||
AppConfig {
|
||||
server: ServerConfig {
|
||||
host: "127.0.0.1".to_string(),
|
||||
port: 0,
|
||||
},
|
||||
api: ApiConfig { base_url },
|
||||
model_aliases: BTreeMap::new(),
|
||||
model_account_routes: BTreeMap::new(),
|
||||
auth: AuthConfig {
|
||||
refresh_margin_seconds: 300,
|
||||
refresh_enabled: true,
|
||||
refresh_concurrency: 2,
|
||||
max_concurrent_per_account: 3,
|
||||
request_interval_ms: 50,
|
||||
rotation_strategy: "least_used".to_string(),
|
||||
tier_priority: Vec::new(),
|
||||
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
|
||||
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
|
||||
},
|
||||
quota: QuotaConfig {
|
||||
refresh_interval_minutes: 5,
|
||||
warning_thresholds: QuotaWarningThresholds {
|
||||
primary: vec![80, 90],
|
||||
secondary: vec![80, 90],
|
||||
},
|
||||
},
|
||||
database: DatabaseConfig { url: database_url },
|
||||
tls: TlsConfig {
|
||||
force_http11: false,
|
||||
},
|
||||
ws_pool: WebSocketPoolConfig::default(),
|
||||
fingerprint: Default::default(),
|
||||
admin: AdminConfig {
|
||||
session_ttl_minutes: 1440,
|
||||
session_cleanup_interval_secs: 3600,
|
||||
default_username: "admin".to_string(),
|
||||
default_password: "admin".to_string(),
|
||||
},
|
||||
logging: LoggingConfig {
|
||||
directory: "logs".to_string(),
|
||||
retention_days: 14,
|
||||
enabled: false,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,5 @@
|
||||
use axum::{
|
||||
body::{to_bytes, Body},
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
};
|
||||
use chrono::{Duration, Utc};
|
||||
@@ -29,6 +29,8 @@ use serde_json::{json, Value};
|
||||
use sqlx::SqlitePool;
|
||||
use tower::util::ServiceExt;
|
||||
|
||||
use crate::support::{admin::seed_admin_session, http::response_json};
|
||||
|
||||
#[tokio::test]
|
||||
async fn dashboard_summary_should_render_fingerprint_updated_at_as_beijing_time() {
|
||||
let (app, _store, _pool, _dir) = dashboard_test_app(
|
||||
@@ -346,29 +348,6 @@ async fn dashboard_summary(app: axum::Router) -> Value {
|
||||
response_json(response).await
|
||||
}
|
||||
|
||||
async fn seed_admin_session(pool: &SqlitePool, session_id: &str) {
|
||||
sqlx::query(
|
||||
"insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind("admin_1")
|
||||
.bind("hash")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
sqlx::query(
|
||||
"insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)",
|
||||
)
|
||||
.bind(session_id)
|
||||
.bind("admin_1")
|
||||
.bind("2999-01-01T00:00:00Z")
|
||||
.bind("2026-06-18T00:00:00Z")
|
||||
.execute(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn seed_usage_summary(pool: &SqlitePool, input_tokens: u64, output_tokens: u64) {
|
||||
seed_usage_summary_with_last_used(pool, input_tokens, output_tokens, "2026-06-18T00:00:00Z")
|
||||
.await;
|
||||
@@ -405,11 +384,6 @@ async fn seed_usage_summary_with_last_used(
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
async fn response_json(response: axum::response::Response) -> Value {
|
||||
let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
|
||||
serde_json::from_slice(&bytes).unwrap()
|
||||
}
|
||||
|
||||
fn service_status_value<'a>(body: &'a Value, label: &str) -> &'a str {
|
||||
body["data"]["serviceStatuses"]
|
||||
.as_array()
|
||||
|
||||
@@ -92,7 +92,7 @@ async fn event_log_store_should_promote_diagnostic_metadata_fields() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn event_log_store_should_get_count_and_clear_events() {
|
||||
async fn event_log_store_should_get_and_clear_events() {
|
||||
let (store, _dir) = event_log_store("event-logs-clear.sqlite").await;
|
||||
let mut event = EventLog::new("request", EventLevel::Warn, "detail");
|
||||
event.id = "log_detail".to_string();
|
||||
@@ -100,9 +100,12 @@ async fn event_log_store_should_get_count_and_clear_events() {
|
||||
|
||||
let loaded = store.get("log_detail").await.unwrap().unwrap();
|
||||
assert_eq!(loaded.id, "log_detail");
|
||||
assert_eq!(store.count().await.unwrap(), 1);
|
||||
assert_eq!(store.clear().await.unwrap(), 1);
|
||||
assert_eq!(store.count().await.unwrap(), 0);
|
||||
let page = store
|
||||
.list(EventLogFilter::default(), None, 10)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(page.items.is_empty());
|
||||
}
|
||||
|
||||
async fn event_log_store(db_name: &str) -> (SqliteEventLogStore, tempfile::TempDir) {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user