Compare commits

..
50 Commits
Author SHA1 Message Date
ZyyandCodex cbf9ce91a1 chore: bump version to 1.0.11
Co-authored-by: Codex <noreply@openai.com>
2026-07-09 15:44:19 +08:00
ZyyandCodex 71e51a868f fix: keep quota exhausted status during token refresh
Co-authored-by: Codex <noreply@openai.com>
2026-07-09 15:41:42 +08:00
ZyyandCodex 3c9ec50bf2 chore: bump version to 1.0.10
Co-authored-by: Codex <noreply@openai.com>
2026-07-09 13:08:08 +08:00
ZyyandCodex c1b81809a6 refactor: derive account refresh status from leases
Co-authored-by: Codex <noreply@openai.com>
2026-07-09 13:06:54 +08:00
Zyy 939f392438 docs: refresh architecture and database docs 2026-07-08 21:35:26 +08:00
ZyyandCodex ff0de4742d chore: bump version to 1.0.9
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 18:06:08 +08:00
ZyyandCodex 44f45101b8 test: assert account quota local usage
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 17:44:01 +08:00
ZyyandCodex a240668a47 fix: polish sidebar focus state
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 16:45:52 +08:00
ZyyandCodex 30f354965c fix: align account quota usage tracking
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 16:37:32 +08:00
ZyyandCodex 3b462f7ecf refactor: remove trusted proxy config
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 15:22:23 +08:00
ZyyandCodex 5682d4fbe9 test: simplify ops error log snapshot helper
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 14:55:05 +08:00
ZyyandCodex 146776c742 fix: separate ops errors from usage records
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 14:04:49 +08:00
ZyyandCodex e8c14d4324 fix: complete update log streaming
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 13:58:28 +08:00
zyycnandCodex 5483e37b00 refactor: extract upstream scheduler
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 13:56:27 +08:00
ZyyandCodex daad7c746c fix: strengthen response account affinity
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 18:31:50 +08:00
ZyyandCodex 6e24037c4d chore: bump version to 1.0.8
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 17:06:08 +08:00
ZyyandCodex 95daf7093e feat: add smart account scheduling
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 17:00:56 +08:00
ZyyandCodex 2e121c0670 fix: hot reload fingerprint updates
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 15:01:28 +08:00
ZyyandCodex 66916ab8b2 chore: bump version to 1.0.7
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 14:04:19 +08:00
ZyyandCodex 2945523140 fix: clean account routing and affinity fallback
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 14:01:35 +08:00
ZyyandCodex f375d2d77c fix: allow large OpenAI request bodies
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 14:00:04 +08:00
zyycnandCodex c6e1b347be chore: bump version to 1.0.6
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 01:54:27 +08:00
zyycnandCodex 1ac9803c34 fix: stop recording account state usage events
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 01:54:14 +08:00
zyycnandCodex 352b781340 chore: add engineering standards skill
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 01:10:02 +08:00
zyycnandCodex 62a6dfc489 chore: bump version to 1.0.5
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 00:16:26 +08:00
zyycnandCodex e91f91b542 test: move source tests to integration suite
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 00:15:30 +08:00
zyycnandCodex 380556369b fix: harden runtime diagnostics and quota refresh
Co-authored-by: Codex <noreply@openai.com>
2026-07-06 23:59:27 +08:00
ZyyandCodex 595cda541d style: apply cargo fmt
Co-authored-by: Codex <noreply@openai.com>
2026-07-06 17:55:54 +08:00
ZyyandCodex 90ed8bf6bc chore: bump version to 1.0.4
Co-authored-by: Codex <noreply@openai.com>
2026-07-06 17:49:48 +08:00
ZyyandCodex 1101eee628 fix: align quota window accounting
Co-authored-by: Codex <noreply@openai.com>
2026-07-06 17:49:21 +08:00
ZyyandCodex 0e952c9d83 fix: enforce websocket first token retries
Co-authored-by: Codex <noreply@openai.com>
2026-07-06 17:49:15 +08:00
ZyyandCodex 875d610812 fix: refine dashboard metrics
Co-authored-by: Codex <noreply@openai.com>
2026-07-06 16:53:31 +08:00
ZyyandCodex 9c37d262d1 fix: retry websocket first token stalls
Co-authored-by: Codex <noreply@openai.com>
2026-07-06 16:53:07 +08:00
ZyyandCodex bd95adadf9 fix: align dashboard daily metrics
Co-authored-by: Codex <noreply@openai.com>
2026-07-06 13:19:02 +08:00
zyycnandCodex 6e6627b0c9 docs: refresh project documentation
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 22:39:26 +08:00
zyycnandCodex 3579c28b75 chore: bump version to 1.0.3
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 22:27:00 +08:00
zyycnandCodex 2143de2791 feat: animate dashboard metrics
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 22:26:30 +08:00
zyycnandCodex 298357f97c feat: refine admin console ergonomics
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 22:15:52 +08:00
zyycnandCodex 07829ca385 chore: bump version to 1.0.2
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 19:39:35 +08:00
zyycnandCodex 617bd676f0 refactor: move websocket keepalive into pump
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 19:38:15 +08:00
zyycnandCodex 0c18b44246 style: refine quota usage display
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 16:36:24 +08:00
zyycnandCodex 73658238e5 fix: refine admin usage reporting
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 16:25:01 +08:00
zyycnandCodex be3d648fee style: format responses proxy code
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 05:11:27 +08:00
zyycnandCodex d123704172 chore: bump version to 1.0.1
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 05:05:55 +08:00
zyycnandCodex 7d41151b72 feat: align responses proxy and api keys
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 05:05:05 +08:00
zyycnandCodex 52dd6212ca refactor: remove chat completions path
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 23:04:23 +08:00
zyycnandCodex c0cd7ccbc1 fix: clean up quota dispatch audit issues
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 21:34:05 +08:00
zyycnandCodex 50d9b6eef8 fix: harden account scheduling and client context
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 19:56:32 +08:00
zyycnandCodex 94ef9784ba fix: serve vite public assets
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 16:45:25 +08:00
zyycnandCodex 38ab14fb19 fix: defer update event stream connection
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 16:00:25 +08:00
226 changed files with 14774 additions and 10284 deletions
+18 -17
View File
@@ -15,11 +15,11 @@
## 概览
Codex Proxy RS 是一个单进程服务:
Codex Proxy RS 是一个单进程网关:
- Rust/Axum 后端提供 OpenAI 兼容 `/v1/*` 代理、管理端 API、SQLite 持久化和静态前端托管。
- Vue 管理端用于账号导入、API Key、用量统计、请求记录、运行参数、模型映射和系统更新。
- 运行数据默认收敛到仓库 `.runtime/`,Docker 也使用同一套宿主机目录做 bind mount。
- Rust/Axum 后端提供 Responses 兼容代理、管理端 API、SQLite 持久化和静态前端托管。
- Vue 管理端负责账号导入、API Key、用量统计、请求记录、运行参数、模型映射和系统更新。
- 运行数据默认写入仓库 `.runtime/`;Docker 部署也使用这套宿主机目录挂载。
## 环境
@@ -55,7 +55,7 @@ mkdir -p .runtime/data .runtime/logs
CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.toml
```
服务默认监听 `http://0.0.0.0:8080`。
服务默认监听 `0.0.0.0:8080`。本机访问通常使用 `http://127.0.0.1:8080`。
## Docker 部署
@@ -106,9 +106,11 @@ docker compose -f deploy/docker-compose.yml up -d
| 端点 | 说明 |
| --- | --- |
| `POST /v1/responses` | OpenAI Responses 兼容接口 |
| `POST /v1/chat/completions` | Chat Completions 兼容接口 |
| `POST /v1/responses/review` | Review 模型请求入口 |
| `POST /v1/responses/compact` | Compact 请求入口 |
| `GET /v1/models` | 模型列表 |
| `GET /v1/models/{id}` | 模型详情 |
| `GET /v1/models/{id}/info` | 模型运行信息 |
| `GET /v1/models/catalog` | 管理端可见模型目录 |
| `/api/admin/*` | 管理端 API |
@@ -137,11 +139,11 @@ pnpm --dir frontend build
CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.toml
```
管理端主要功能:
管理端功能:
- 账号导入、OAuth 授权、连接测试、quota 刷新、token 刷新。
- 客户端 API Key 管理。
- 请求记录、Token 用量、模型用量和账号额度。
- Dashboard、请求记录、Token 用量、模型用量和账号额度。
- 运行参数、模型别名、账号选择策略。
- 版本检查、在线更新、更新日志、重启和回滚。
@@ -152,7 +154,7 @@ CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.tom
`.runtime` 约定:
```text
.runtime/config.yaml # 本地或 Docker 启动配置,按实际运行环境写路径
.runtime/config.yaml # 本地或 Docker 启动配置,路径按运行环境写
.runtime/data/codex-proxy-rs.sqlite # SQLite 数据库
.runtime/data/installation_id # 上游 installation id
.runtime/data/update-state.json # 在线更新状态
@@ -162,22 +164,21 @@ CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.tom
## 发布与更新
版本信息在 `release/version.yaml`,发布平台在 `release/platforms.yaml`。
当前版本记录在 `release/version.yaml`,发布平台记录在 `release/platforms.yaml`。
发布命令:
```bash
release/publish 0.1.5
release/publish 1.0.4
```
该命令会更新版本文件、生成版本提交、创建 `v<version>` tag 并推送,随后 GitHub Actions 构建 Release 归档和 GHCR 多平台镜像。
该命令会更新版本文件、生成版本提交、创建 `v<version>` tag 并推送。GitHub Actions 随后构建 Release 归档和 GHCR 多平台镜像。
在线更新由主服务处理,管理端调用 `/api/admin/system/*`:
- 检查 GitHub Release。
- 下载匹配当前平台的归档和 `checksums.txt`。
- 校验 checksum。
- 替换二进制和 `web/dist`。
- 查询 GitHub Release。
- 下载当前平台对应的归档和 `checksums.txt`。
- 校验 checksum 后替换二进制和 `web/dist`。
- 重启时,Docker 模式依赖 `restart: unless-stopped` 拉起新容器进程;非 Docker 模式会先安排新进程延迟启动,再关闭当前进程。
## 项目结构
@@ -186,7 +187,7 @@ release/publish 0.1.5
backend/ Rust/Axum 后端、SQLite migration、集成测试
frontend/ Vue 3 管理端
deploy/ Dockerfile、Compose、部署配置模板
docs/ 架构、审计和维护文档
docs/ 架构和维护文档
release/ 版本、平台和发布脚本
skills/ 项目本地 Codex skill
```
+28
View File
@@ -1112,6 +1112,12 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "http-range-header"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9171a2ea8a68358193d15dd5d70c1c10a2afc3e7e4c5bc92bc9f025cebd7359c"
[[package]]
name = "httparse"
version = "1.10.1"
@@ -1496,6 +1502,16 @@ version = "0.3.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
[[package]]
name = "mime_guess"
version = "2.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
dependencies = [
"mime",
"unicase",
]
[[package]]
name = "miniz_oxide"
version = "0.8.9"
@@ -2176,6 +2192,7 @@ version = "1.0.150"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
dependencies = [
"indexmap",
"itoa",
"memchr",
"serde",
@@ -2839,6 +2856,11 @@ dependencies = [
"http",
"http-body",
"http-body-util",
"http-range-header",
"httpdate",
"mime",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"tokio",
"tokio-util",
@@ -2994,6 +3016,12 @@ version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
[[package]]
name = "unicase"
version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
[[package]]
name = "unicode-bidi"
version = "0.3.18"
+3 -3
View File
@@ -29,7 +29,7 @@ rustls-pki-types = "1.14.1"
secrecy = "0.10.3"
semver = "1.0.27"
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.150"
serde_json = { version = "1.0.150", features = ["preserve_order"] }
sha2 = "0.10.9"
sqlx = { version = "0.9.0", features = ["runtime-tokio", "tls-rustls", "sqlite", "chrono", "uuid"] }
subtle = "2.6.1"
@@ -37,7 +37,8 @@ tar = "0.4.44"
thiserror = "2.0.18"
tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "signal", "fs", "time", "sync", "net", "test-util"] }
tokio-tungstenite = { version = "0.28.0", features = ["proxy", "rustls-tls-native-roots"] }
tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout"] }
tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout", "fs"] }
tracing = "0.1.44"
tracing-appender = "0.2.5"
tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] }
@@ -46,7 +47,6 @@ uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
[dev-dependencies]
tempfile = "3.27.0"
tower = { version = "0.5.3", features = ["util"] }
wiremock = "0.6.5"
[patch.crates-io]
+118 -4
View File
@@ -1,11 +1,13 @@
//! 管理端账号 quota 展示视图。
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use chrono::{DateTime, Duration, Utc};
use serde::Serialize;
use serde_json::Value;
use crate::infra::{
format::format_percent,
format::{format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64},
time::{china_datetime, china_relative_time},
};
@@ -29,8 +31,79 @@ struct AdminAccountQuotaWindowData {
label_display: String,
used_percent: Option<f64>,
used_percent_display: String,
#[serde(skip_serializing_if = "Option::is_none")]
local_usage: Option<AdminAccountQuotaWindowLocalUsageData>,
reset_at_display: String,
window_used_display: String,
#[serde(skip)]
reset_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountQuotaWindowLocalUsageData {
request_count: i64,
request_count_display: String,
input_tokens: i64,
input_tokens_display: String,
output_tokens: i64,
output_tokens_display: String,
cached_tokens: i64,
cached_tokens_display: String,
total_tokens: i64,
total_tokens_display: String,
}
#[derive(Debug, Clone)]
pub(crate) struct AdminAccountQuotaUsageWindow {
pub(crate) key: String,
pub(crate) start: DateTime<Utc>,
pub(crate) end: DateTime<Utc>,
pub(crate) window_seconds: u64,
}
impl AdminAccountQuotaUsageWindow {
pub(crate) fn duration_seconds(&self) -> u64 {
self.end
.signed_duration_since(self.start)
.num_seconds()
.max(0) as u64
}
}
#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct AdminAccountQuotaWindowLocalUsage {
pub(crate) request_count: i64,
pub(crate) input_tokens: i64,
pub(crate) output_tokens: i64,
pub(crate) cached_tokens: i64,
}
impl AdminAccountQuotaWindowLocalUsage {
pub(crate) fn add(&mut self, other: Self) {
self.request_count = self.request_count.saturating_add(other.request_count);
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens);
}
}
impl From<AdminAccountQuotaWindowLocalUsage> for AdminAccountQuotaWindowLocalUsageData {
fn from(usage: AdminAccountQuotaWindowLocalUsage) -> Self {
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
Self {
request_count: usage.request_count,
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
input_tokens: usage.input_tokens,
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
output_tokens: usage.output_tokens,
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
cached_tokens: usage.cached_tokens,
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
total_tokens,
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
}
}
}
impl Default for AdminAccountQuotaData {
@@ -48,6 +121,43 @@ impl AdminAccountQuotaData {
.iter()
.any(|window| window.used_percent.is_some_and(|percent| percent >= 80.0))
}
pub(crate) fn usage_windows(&self) -> Vec<AdminAccountQuotaUsageWindow> {
self.windows
.iter()
.filter_map(AdminAccountQuotaWindowData::usage_window)
.collect()
}
pub(crate) fn apply_local_usage(
&mut self,
usage_by_window: &HashMap<String, AdminAccountQuotaWindowLocalUsage>,
) {
for window in &mut self.windows {
window.local_usage = window.usage_window().map(|_| {
usage_by_window
.get(&window.key)
.copied()
.unwrap_or_default()
.into()
});
}
}
}
impl AdminAccountQuotaWindowData {
fn usage_window(&self) -> Option<AdminAccountQuotaUsageWindow> {
let end = self.reset_at?;
let window_seconds = self.window_seconds?;
let seconds = i64::try_from(window_seconds).ok()?;
let start = end.checked_sub_signed(Duration::seconds(seconds))?;
(start <= end).then(|| AdminAccountQuotaUsageWindow {
key: self.key.clone(),
start,
end,
window_seconds,
})
}
}
pub(crate) fn quota_data(
@@ -111,11 +221,13 @@ fn push_monthly_quota_window(
window_seconds,
label_display: "月限额".to_string(),
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent),
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
local_usage: None,
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
reset_at,
});
true
}
@@ -198,11 +310,13 @@ fn push_quota_window(
window_seconds,
label_display,
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent),
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
local_usage: None,
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
reset_at,
});
}
+334 -70
View File
@@ -10,7 +10,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
@@ -25,7 +25,10 @@ use crate::{
},
admin::update_payload::{parse_editable_update, EditableUpdateMessages},
admin::{
accounts::quota_view::{quota_data, AdminAccountQuotaData},
accounts::quota_view::{
quota_data, AdminAccountQuotaData, AdminAccountQuotaUsageWindow,
AdminAccountQuotaWindowLocalUsage,
},
accounts::service::{
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata,
AdminAccountRefreshOutcome, AdminAccountRefreshResult, AdminAccountUpdate,
@@ -38,7 +41,10 @@ use crate::{
format_cost, format_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
},
json::{clamp_limit, clamp_page, total_pages, Page},
time::{china_datetime, china_quarter_hour_start, china_relative_time, china_rfc3339},
time::{
china_datetime, china_quarter_hour_start, china_relative_time, china_rfc3339,
parse_rfc3339_utc,
},
},
runtime::state::AppState,
upstream::accounts::model::AccountStatus,
@@ -129,6 +135,8 @@ struct AdminAccountData {
plan_type: Option<String>,
has_refresh_token: bool,
status: String,
display_status: String,
token_refreshing: bool,
access_token_expires_at: Option<String>,
access_token_expires_at_display: Option<String>,
added_at: String,
@@ -141,7 +149,7 @@ struct AdminAccountData {
impl From<AdminAccountMetadata> for AdminAccountData {
fn from(a: AdminAccountMetadata) -> Self {
Self::from_parts(a, None, None, Vec::new())
Self::from_parts(a, None, None, Vec::new(), false)
}
}
@@ -151,10 +159,12 @@ impl AdminAccountData {
usage: Option<&AdminUsageRecord>,
quota: Option<AdminAccountQuotaData>,
models: Vec<AdminAccountModelUsageData>,
token_refreshing: bool,
) -> Self {
let access_token_expires_at = a.access_token_expires_at.as_ref().map(china_rfc3339);
let access_token_expires_at_display =
a.access_token_expires_at.as_ref().map(china_datetime);
let display_status = account_display_status(a.status, token_refreshing).to_string();
Self {
id: a.id,
email: a.email,
@@ -163,7 +173,9 @@ impl AdminAccountData {
label: a.label,
plan_type: a.plan_type,
has_refresh_token: a.has_refresh_token,
status: account_status_str(a.status).to_string(),
status: a.status.as_str().to_string(),
display_status,
token_refreshing,
access_token_expires_at,
access_token_expires_at_display,
added_at: china_rfc3339(&a.added_at),
@@ -283,6 +295,7 @@ struct AccountListStats {
usage_by_account: HashMap<String, AdminUsageRecord>,
quota_by_account: HashMap<String, AdminAccountQuotaData>,
models_by_account: HashMap<String, Vec<AdminAccountModelUsageData>>,
refreshing_account_ids: HashSet<String>,
}
struct AccountModelUsageRecord {
@@ -297,6 +310,12 @@ struct AccountModelUsageRecord {
last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone)]
struct AccountQuotaUsageWindow {
account_id: String,
window: AdminAccountQuotaUsageWindow,
}
impl AccountListStats {
fn data_for(&self, account: AdminAccountMetadata) -> AdminAccountData {
let account_id = account.id.clone();
@@ -308,6 +327,7 @@ impl AccountListStats {
.get(&account_id)
.cloned()
.unwrap_or_default(),
self.refreshing_account_ids.contains(&account_id),
)
}
}
@@ -396,7 +416,7 @@ impl From<AdminAccountRefreshResult> for AccountRefreshData {
Self {
id: result.id,
email: result.email,
previous_status: account_status_str(result.previous_status).to_string(),
previous_status: result.previous_status.as_str().to_string(),
result: account_refresh_outcome_str(result.outcome),
error: result.error,
duration_ms: result.duration_ms,
@@ -404,6 +424,34 @@ impl From<AdminAccountRefreshResult> for AccountRefreshData {
}
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountQuotaData {
quota: Value,
raw: Value,
quota_data: AdminAccountQuotaData,
plan_type: Option<String>,
account: AdminAccountData,
}
impl AccountQuotaData {
fn from_account(
quota: Value,
raw: Value,
quota_data: AdminAccountQuotaData,
account: AdminAccountData,
) -> Self {
let plan_type = account.plan_type.clone();
Self {
quota,
raw,
quota_data,
plan_type,
account,
}
}
}
#[derive(Debug, Clone, Deserialize, Default)]
#[serde(default)]
pub(crate) struct AccountHealthCheckRequest {
@@ -472,8 +520,8 @@ pub(crate) async fn accounts(
) -> Result<impl IntoResponse, AdminError> {
let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50));
let use_numbered_page = params.page.is_some() || params.page_size.is_some();
let stats = account_list_stats(&state).await;
let summary = account_summary_data(&state, &stats).await;
let quota_by_account = quota_snapshots_by_account(&state).await;
let summary = account_summary_data(&state, &quota_by_account).await;
if use_numbered_page {
return match state
@@ -483,6 +531,7 @@ pub(crate) async fn accounts(
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await?;
let page = crate::infra::json::NumberedPage {
items: page
.items
@@ -509,6 +558,7 @@ pub(crate) async fn accounts(
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await?;
let page = Page {
items: page
.items
@@ -640,25 +690,21 @@ pub(crate) async fn account_quota(
let quota = data.get("quota").cloned().unwrap_or(Value::Null);
let raw = data.get("raw").cloned().unwrap_or(Value::Null);
let quota_json = quota.to_string();
let quota_data = quota_data(&quota_json, Some(Utc::now()));
let plan_type = quota
.get("plan_type")
.and_then(Value::as_str)
.map(ToString::to_string);
let mut quota_data = quota_data(&quota_json, Some(Utc::now()));
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
let account =
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(serde_json::json!({
"quota": quota,
"raw": raw,
"quotaData": quota_data,
"planType": plan_type,
})),
AdminEnvelope::ok(AccountQuotaData::from_account(
quota, raw, quota_data, account,
)),
))
}
Err(AdminAccountError::NotFound) => Err(account_not_found()),
Err(AdminAccountError::Inactive(status)) => Err(AdminError::conflict(format!(
"Account is {}, cannot query quota",
account_status_str(status)
status.as_str()
))),
Err(AdminAccountError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!(
"Failed to fetch quota from Codex API: {msg}"
@@ -667,6 +713,26 @@ pub(crate) async fn account_quota(
}
}
async fn account_data_for_quota_refresh(
state: &AppState,
account_id: &str,
quota_data: AdminAccountQuotaData,
) -> Result<AdminAccountData, AdminError> {
let Some(account) = state
.services
.admin_accounts
.get(account_id)
.await
.map_err(|error| account_error(&error))?
else {
return Err(account_not_found());
};
let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]);
let stats =
account_list_stats(state, std::slice::from_ref(&account), &quota_by_account).await?;
Ok(stats.data_for(account))
}
/// `POST /api/admin/accounts/import`
pub(crate) async fn import_accounts(
State(state): State<AppState>,
@@ -841,38 +907,250 @@ pub(crate) async fn update_account(
}
}
async fn account_list_stats(state: &AppState) -> AccountListStats {
let usage_records = list_all_usage_records(state).await;
let quota_snapshots = state
async fn account_list_stats(
state: &AppState,
accounts: &[AdminAccountMetadata],
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
) -> Result<AccountListStats, AdminError> {
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let usage_records = state
.services
.background_tasks
.accounts
.list_quota_snapshots()
.usage
.list_by_account_ids(&account_ids)
.await
.unwrap_or_default();
let mut quota_by_account = quota_by_account.clone();
let all_quota_windows = quota_usage_windows_by_account(accounts, &quota_by_account);
let quota_window_usage_stats =
quota_window_local_usage_by_account(state, &all_quota_windows).await;
apply_quota_window_local_usage(accounts, &mut quota_by_account, &quota_window_usage_stats);
let selected_quota_windows = selected_quota_windows_by_account(accounts, &quota_by_account);
let usage_records = usage_records
.into_iter()
.map(|usage| {
apply_selected_quota_window_usage(
usage,
&selected_quota_windows,
&quota_window_usage_stats,
)
})
.collect::<Vec<_>>();
let models_by_account = list_current_window_model_usage(state, &usage_records).await;
let refreshing_account_ids = state
.services
.token_refresh
.refreshing_account_ids(&account_ids, Utc::now())
.await
.map_err(|error| AdminError::internal(error.to_string()))?;
AccountListStats {
Ok(AccountListStats {
usage_by_account: usage_records
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect(),
quota_by_account: quota_snapshots
.into_iter()
.map(|snapshot| {
(
snapshot.account_id,
quota_data(&snapshot.quota_json, snapshot.quota_fetched_at),
)
})
.collect(),
quota_by_account,
models_by_account,
refreshing_account_ids,
})
}
fn account_display_status(status: AccountStatus, token_refreshing: bool) -> &'static str {
if token_refreshing && matches!(status, AccountStatus::Active | AccountStatus::Expired) {
"refreshing"
} else {
status.as_str()
}
}
fn selected_quota_windows_by_account(
accounts: &[AdminAccountMetadata],
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
) -> HashMap<String, AdminAccountQuotaUsageWindow> {
accounts
.iter()
.filter_map(|account| {
quota_by_account
.get(&account.id)
.and_then(selected_quota_window)
.map(|window| (account.id.clone(), window))
})
.collect()
}
fn selected_quota_window(quota: &AdminAccountQuotaData) -> Option<AdminAccountQuotaUsageWindow> {
quota.usage_windows().into_iter().max_by(|left, right| {
left.duration_seconds()
.cmp(&right.duration_seconds())
.then_with(|| left.end.cmp(&right.end))
.then_with(|| left.key.cmp(&right.key))
})
}
fn quota_usage_windows_by_account(
accounts: &[AdminAccountMetadata],
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
) -> Vec<AccountQuotaUsageWindow> {
accounts
.iter()
.filter_map(|account| {
quota_by_account
.get(&account.id)
.map(|quota| (account.id.as_str(), quota))
})
.flat_map(|(account_id, quota)| {
quota
.usage_windows()
.into_iter()
.map(move |window| AccountQuotaUsageWindow {
account_id: account_id.to_string(),
window,
})
})
.collect()
}
async fn apply_account_quota_window_local_usage(
state: &AppState,
account_id: &str,
quota: &mut AdminAccountQuotaData,
) {
let windows = quota
.usage_windows()
.into_iter()
.map(|window| AccountQuotaUsageWindow {
account_id: account_id.to_string(),
window,
})
.collect::<Vec<_>>();
let usage_by_account = quota_window_local_usage_by_account(state, &windows).await;
let empty_usage = HashMap::new();
quota.apply_local_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage));
}
fn apply_quota_window_local_usage(
accounts: &[AdminAccountMetadata],
quota_by_account: &mut HashMap<String, AdminAccountQuotaData>,
usage_by_account: &HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>,
) {
for account in accounts {
if let Some(quota) = quota_by_account.get_mut(&account.id) {
let empty_usage = HashMap::new();
quota.apply_local_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage));
}
}
}
async fn quota_window_local_usage_by_account(
state: &AppState,
windows: &[AccountQuotaUsageWindow],
) -> HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>> {
let Some(min_start) = windows.iter().map(|window| window.window.start).min() else {
return HashMap::new();
};
let Some(max_end) = windows.iter().map(|window| window.window.end).max() else {
return HashMap::new();
};
let account_ids = windows
.iter()
.map(|window| window.account_id.as_str())
.collect::<HashSet<_>>();
let mut windows_by_account = HashMap::<&str, Vec<&AdminAccountQuotaUsageWindow>>::new();
for window in windows {
windows_by_account
.entry(window.account_id.as_str())
.or_default()
.push(&window.window);
}
let mut builder = QueryBuilder::<Sqlite>::new(
"select
account_id,
bucket_start,
coalesce(sum(request_count), 0) as request_count,
coalesce(sum(input_tokens), 0) as input_tokens,
coalesce(sum(output_tokens), 0) as output_tokens,
coalesce(sum(cached_tokens), 0) as cached_tokens
from usage_time_buckets
where account_id in (",
);
let mut separated = builder.separated(", ");
for account_id in &account_ids {
separated.push_bind(*account_id);
}
separated.push_unseparated(")");
builder.push(" and bucket_start >= ");
builder.push_bind(china_quarter_hour_start(min_start).to_rfc3339());
builder.push(" and bucket_start <= ");
builder.push_bind(max_end.to_rfc3339());
builder.push(" group by account_id, bucket_start");
let rows = builder
.build()
.fetch_all(state.services.background_tasks.accounts.pool())
.await
.unwrap_or_default();
let mut usage_by_account =
HashMap::<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>::new();
for row in rows {
let account_id: String = row.get("account_id");
let Ok(bucket_start) = parse_rfc3339_utc(row.get::<&str, _>("bucket_start")) else {
continue;
};
let bucket_usage = AdminAccountQuotaWindowLocalUsage {
request_count: row.get("request_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
};
let Some(account_windows) = windows_by_account.get(account_id.as_str()) else {
continue;
};
for window in account_windows.iter().filter(|window| {
bucket_start >= china_quarter_hour_start(window.start) && bucket_start <= window.end
}) {
usage_by_account
.entry(account_id.clone())
.or_default()
.entry(window.key.clone())
.or_default()
.add(bucket_usage);
}
}
usage_by_account
}
fn apply_selected_quota_window_usage(
mut usage: AdminUsageRecord,
selected_quota_windows: &HashMap<String, AdminAccountQuotaUsageWindow>,
quota_window_usage_stats: &HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>,
) -> AdminUsageRecord {
let Some(window) = selected_quota_windows.get(&usage.account_id) else {
return usage;
};
let stats = quota_window_usage_stats
.get(&usage.account_id)
.and_then(|usage_by_window| usage_by_window.get(&window.key))
.copied()
.unwrap_or_default();
usage.window_request_count = stats.request_count;
usage.window_input_tokens = stats.input_tokens;
usage.window_output_tokens = stats.output_tokens;
usage.window_cached_tokens = stats.cached_tokens;
usage.window_started_at = Some(window.start);
usage.window_reset_at = Some(window.end);
usage.limit_window_seconds = Some(window.window_seconds);
usage
}
async fn account_summary_data(
state: &AppState,
stats: &AccountListStats,
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
) -> AdminAccountSummaryData {
let accounts = list_all_account_metadata(state).await;
let total = accounts.len() as u64;
@@ -883,8 +1161,7 @@ async fn account_summary_data(
let high_usage = accounts
.iter()
.filter(|account| {
stats
.quota_by_account
quota_by_account
.get(&account.id)
.is_some_and(account_quota_has_high_usage)
})
@@ -902,6 +1179,24 @@ async fn account_summary_data(
}
}
async fn quota_snapshots_by_account(state: &AppState) -> HashMap<String, AdminAccountQuotaData> {
state
.services
.background_tasks
.accounts
.list_quota_snapshots()
.await
.unwrap_or_default()
.into_iter()
.map(|snapshot| {
(
snapshot.account_id,
quota_data(&snapshot.quota_json, snapshot.quota_fetched_at),
)
})
.collect()
}
async fn list_all_account_metadata(state: &AppState) -> Vec<AdminAccountMetadata> {
let mut page = 1;
let mut accounts = Vec::new();
@@ -934,26 +1229,6 @@ fn account_summary_needs_attention(status: AccountStatus) -> bool {
)
}
async fn list_all_usage_records(state: &AppState) -> Vec<AdminUsageRecord> {
let mut cursor = None;
let mut records = Vec::new();
loop {
let Ok(page) = state
.services
.usage
.list(cursor, ACCOUNT_STATS_PAGE_LIMIT)
.await
else {
return Vec::new();
};
records.extend(page.items);
let Some(next_cursor) = page.next_cursor else {
return records;
};
cursor = Some(next_cursor);
}
}
async fn list_current_window_model_usage(
state: &AppState,
usage_records: &[AdminUsageRecord],
@@ -1202,14 +1477,3 @@ fn account_export_ids(value: Option<&str>) -> Vec<String> {
.map(ToString::to_string)
.collect()
}
fn account_status_str(status: crate::upstream::accounts::model::AccountStatus) -> &'static str {
match status {
crate::upstream::accounts::model::AccountStatus::Active => "active",
crate::upstream::accounts::model::AccountStatus::Expired => "expired",
crate::upstream::accounts::model::AccountStatus::QuotaExhausted => "quota_exhausted",
crate::upstream::accounts::model::AccountStatus::Refreshing => "refreshing",
crate::upstream::accounts::model::AccountStatus::Disabled => "disabled",
crate::upstream::accounts::model::AccountStatus::Banned => "banned",
}
}
@@ -75,6 +75,8 @@ pub enum AdminAccountError {
Delete,
#[error("failed to update claims")]
UpdateClaims,
#[error("failed to coordinate token refresh lease")]
RefreshLease,
#[error("failed to sync account pool")]
SyncAccountPool,
#[error("failed to store quota")]
@@ -113,12 +115,6 @@ pub enum AdminAccountError {
Inactive(AcctStatus),
}
#[derive(Debug, Clone)]
pub struct UpdatedAccountStatus {
pub id: String,
pub status: AcctStatus,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdminAccountRefreshOutcome {
Alive,
@@ -233,6 +229,14 @@ pub(super) struct ManualCreateTokens {
pub(super) access_token: String,
pub(super) refresh_token_for_new: Option<String>,
pub(super) refresh_token_for_existing: Option<String>,
pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims,
}
#[derive(Debug, Clone)]
pub(super) struct RefreshedAccountTokens {
pub(super) access_token: String,
pub(super) refresh_token: Option<String>,
pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims,
}
#[derive(Debug, Clone)]
@@ -4,10 +4,7 @@ use super::{
contracts::{ExportedAccount, ExportedAccounts},
AdminAccountError, AdminAccountService,
};
use crate::{
infra::time::china_rfc3339_str,
upstream::accounts::{model::AccountStatus, store::StoredAccount},
};
use crate::{infra::time::china_rfc3339_str, upstream::accounts::store::StoredAccount};
impl AdminAccountService {
pub async fn export(&self, ids: Vec<String>) -> Result<ExportedAccounts, AdminAccountError> {
@@ -58,20 +55,9 @@ impl From<StoredAccount> for ExportedAccount {
access_token_expires_at: account
.access_token_expires_at
.map(|value| value.to_rfc3339()),
status: export_status_str(account.status).to_string(),
status: account.status.as_str().to_string(),
added_at: china_rfc3339_str(&account.added_at),
updated_at: china_rfc3339_str(&account.updated_at),
}
}
}
fn export_status_str(status: AccountStatus) -> &'static str {
match status {
AccountStatus::Active => "active",
AccountStatus::Expired => "expired",
AccountStatus::QuotaExhausted => "quota_exhausted",
AccountStatus::Refreshing => "refreshing",
AccountStatus::Disabled => "disabled",
AccountStatus::Banned => "banned",
}
}
+11 -38
View File
@@ -281,8 +281,7 @@ impl AdminAccountService {
token: Option<String>,
refresh_token: Option<String>,
) -> Result<Option<ResolvedImportTokens>, AdminAccountError> {
let mut refresh_token =
crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
let refresh_token = crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
let Some(access_token) = crate::upstream::accounts::importing::normalize_nonempty(
token
.as_deref()
@@ -292,26 +291,12 @@ impl AdminAccountService {
return Ok(None);
};
let refreshed = self
.token_refresher
.refresh(&existing_refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
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;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
.refresh_tokens_from_refresh_token(&existing_refresh_token)
.await?;
return Ok(Some(ResolvedImportTokens {
access_token,
refresh_token,
claims: Some(claims),
access_token: refreshed.access_token,
refresh_token: refreshed.refresh_token,
claims: Some(refreshed.claims),
}));
};
@@ -334,24 +319,12 @@ impl AdminAccountService {
}));
};
let refreshed = self
.token_refresher
.refresh(&existing_refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
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;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
.refresh_tokens_from_refresh_token(&existing_refresh_token)
.await?;
Ok(Some(ResolvedImportTokens {
access_token,
refresh_token,
claims: Some(claims),
access_token: refreshed.access_token,
refresh_token: refreshed.refresh_token,
claims: Some(refreshed.claims),
}))
}
async fn import_target_account_id(
+64 -92
View File
@@ -3,10 +3,11 @@ use std::{
hash::{Hash, Hasher},
};
use chrono::{DateTime, Utc};
use chrono::{DateTime, Duration, Utc};
use futures::{stream, StreamExt};
use secrecy::{ExposeSecret, SecretString};
use tokio::time::sleep;
use uuid::Uuid;
use crate::{
infra::{
@@ -20,11 +21,13 @@ use super::{
contracts::{
parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountHealthCheck,
AdminAccountMetadata, AdminAccountRefreshOutcome, AdminAccountRefreshResult,
AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens, UpdatedAccountStatus,
AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens,
},
AdminAccountService,
};
const ADMIN_REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60;
impl AdminAccountService {
pub async fn list(
&self,
@@ -92,37 +95,32 @@ impl AdminAccountService {
.as_deref()
.map(crate::upstream::accounts::importing::normalize_bearer_token),
) {
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
ManualCreateTokens {
access_token,
refresh_token_for_new: provided_refresh_token.clone(),
refresh_token_for_existing: provided_refresh_token,
claims,
}
} else if let Some(refresh_token) = provided_refresh_token {
let token_pair = self
.token_refresher
.refresh(&refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
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)?;
let refreshed = self
.refresh_tokens_from_refresh_token(&refresh_token)
.await?;
ManualCreateTokens {
access_token,
refresh_token_for_new: token_pair.refresh_token.clone(),
refresh_token_for_existing: token_pair.refresh_token,
access_token: refreshed.access_token,
refresh_token_for_new: refreshed.refresh_token.clone(),
refresh_token_for_existing: refreshed.refresh_token,
claims: refreshed.claims,
}
} else {
return Err(AdminAccountError::TokenRequired);
};
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&tokens.access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
let claims = tokens.claims;
let existing = if let Some(account_id) = claims.account_id.as_deref() {
self.store
.find_by_chatgpt_identity(account_id, claims.user_id.as_deref())
@@ -195,45 +193,6 @@ impl AdminAccountService {
.map(stored_to_admin_metadata)
.ok_or(AdminAccountError::NotFound)
}
pub async fn update_label(
&self,
account_id: &str,
label: Option<String>,
) -> Result<bool, AdminAccountError> {
if label.as_ref().is_some_and(|l| l.chars().count() > 64) {
return Err(AdminAccountError::LabelTooLong);
}
let updated = self
.store
.set_label(account_id, label)
.await
.map_err(|_| AdminAccountError::UpdateLabel)?;
if updated {
self.sync_account_pool_best_effort(account_id, "account label update")
.await;
}
Ok(updated)
}
pub async fn update_status(
&self,
account_id: &str,
status: &str,
) -> Result<Option<UpdatedAccountStatus>, AdminAccountError> {
let status = parse_account_status(status)?;
match self.store.set_status(account_id, status).await {
Ok(true) => {
self.sync_account_pool_best_effort(account_id, "account status update")
.await;
self.evict_account_websocket_pool(account_id).await;
Ok(Some(UpdatedAccountStatus {
id: account_id.to_string(),
status,
}))
}
Ok(false) => Ok(None),
Err(_) => Err(AdminAccountError::UpdateStatus),
}
}
pub async fn update_account(
&self,
account_id: &str,
@@ -325,50 +284,63 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
};
let Some(refresh_token) = account.refresh_token.as_ref() else {
let Some(refresh_token) = account
.refresh_token
.as_ref()
.map(|token| token.expose_secret().to_string())
else {
return Ok(skipped("no refresh token"));
};
if account.status == crate::upstream::accounts::model::AccountStatus::Disabled {
return Ok(skipped("manually disabled"));
}
if account.status == crate::upstream::accounts::model::AccountStatus::Refreshing {
let lease_now = Utc::now();
let lease_owner = format!(
"{}:{}",
self.refresh_lease_owner_prefix,
Uuid::new_v4().simple()
);
let lease_acquired = self
.refresh_leases
.try_acquire(
account_id,
&lease_owner,
lease_now + Duration::seconds(ADMIN_REFRESH_LEASE_TTL_SECONDS),
lease_now,
)
.await
.map_err(|_| AdminAccountError::RefreshLease)?;
if !lease_acquired {
return Ok(skipped("refresh already in progress"));
}
match self
.token_refresher
.refresh(refresh_token.expose_secret())
.await
{
Ok(tokens) => {
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(),
)
.map_err(AdminAccountError::InvalidToken)?;
let refresh_result = self.refresh_tokens_from_refresh_token(&refresh_token).await;
let release_result = self.refresh_leases.release(account_id, &lease_owner).await;
if release_result.is_err() {
return Err(AdminAccountError::RefreshLease);
}
match refresh_result {
Ok(refreshed) => {
let updated = self
.store
.update_from_claims(
account_id,
crate::upstream::accounts::store::AccountClaimsUpdate {
email: claims.email,
account_id: claims.account_id.or(account.account_id),
user_id: claims.user_id,
plan_type: claims.plan_type,
access_token: SecretString::new(access_token.into()),
refresh_token: tokens
email: refreshed.claims.email,
account_id: refreshed.claims.account_id.or(account.account_id),
user_id: refreshed.claims.user_id,
plan_type: refreshed.claims.plan_type,
access_token: SecretString::new(refreshed.access_token.into()),
refresh_token: refreshed
.refresh_token
.map(|token| SecretString::new(token.into())),
access_token_expires_at: Some(claims.expires_at),
next_refresh_at: Some(
self.next_refresh_at_for_expires_at(account_id, claims.expires_at),
),
access_token_expires_at: Some(refreshed.claims.expires_at),
next_refresh_at: Some(self.next_refresh_at_for_expires_at(
account_id,
refreshed.claims.expires_at,
)),
status: crate::upstream::accounts::model::AccountStatus::Active,
},
)
@@ -387,7 +359,7 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
})
}
Err(failure) => {
Err(AdminAccountError::RefreshTokenExchange(failure)) => {
if manual_refresh_failure_is_permanent(&failure) {
let status = crate::upstream::accounts::model::AccountStatus::Expired;
let updated = self
@@ -422,6 +394,7 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
})
}
Err(error) => Err(error),
}
}
@@ -494,8 +467,7 @@ fn manual_refresh_failure_is_permanent(
match failure {
crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant
| crate::upstream::accounts::token_refresh::RefreshFailure::Banned => true,
crate::upstream::accounts::token_refresh::RefreshFailure::QuotaExhausted
| crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport
crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport
| crate::upstream::accounts::token_refresh::RefreshFailure::Transport => false,
}
}
+38 -3
View File
@@ -11,14 +11,15 @@ mod testing;
use std::sync::Arc as StdArc;
use chrono::DateTime;
use chrono::{DateTime, Utc};
use uuid::Uuid;
use crate::{
upstream::accounts::{
cookies::SqliteCookieStore,
pool::RuntimeAccountPoolService,
store::SqliteAccountStore,
token_refresh::{RuntimeRefreshPolicy, TokenRefresher},
token_refresh::{RefreshLeaseStore, RuntimeRefreshPolicy, TokenRefresher},
},
upstream::models::service::ModelService,
upstream::transport::CodexBackendClient,
@@ -27,7 +28,7 @@ use crate::{
pub use contracts::{
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, AdminAccountRefreshOutcome,
AdminAccountRefreshResult, AdminAccountUpdate, BatchDeleteAccounts, ExportedAccounts,
ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, UpdatedAccountStatus,
ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput,
};
#[derive(Clone)]
@@ -38,6 +39,8 @@ pub struct AdminAccountService {
pub(crate) models: StdArc<ModelService>,
pub(crate) account_pool: StdArc<RuntimeAccountPoolService>,
pub(crate) token_refresher: StdArc<dyn TokenRefresher>,
pub(crate) refresh_leases: RefreshLeaseStore,
pub(crate) refresh_lease_owner_prefix: String,
pub(crate) oauth: oauth::AccountOAuthService,
pub(crate) refresh_policy: RuntimeRefreshPolicy,
pub(crate) installation_id: Option<String>,
@@ -50,6 +53,7 @@ pub struct AdminAccountServiceParts {
pub models: StdArc<ModelService>,
pub account_pool: StdArc<RuntimeAccountPoolService>,
pub token_refresher: StdArc<dyn TokenRefresher>,
pub refresh_leases: RefreshLeaseStore,
pub oauth: oauth::AccountOAuthService,
pub refresh_policy: RuntimeRefreshPolicy,
pub installation_id: Option<String>,
@@ -64,6 +68,11 @@ impl AdminAccountService {
models: parts.models,
account_pool: parts.account_pool,
token_refresher: parts.token_refresher,
refresh_leases: parts.refresh_leases,
refresh_lease_owner_prefix: format!(
"admin-account-refresh:{}",
Uuid::new_v4().simple()
),
oauth: parts.oauth,
refresh_policy: parts.refresh_policy,
installation_id: parts.installation_id,
@@ -82,6 +91,32 @@ impl AdminAccountService {
)
}
async fn refresh_tokens_from_refresh_token(
&self,
refresh_token: &str,
) -> Result<contracts::RefreshedAccountTokens, AdminAccountError> {
let token_pair = self
.token_refresher
.refresh(refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
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)?;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
Ok(contracts::RefreshedAccountTokens {
access_token,
refresh_token: token_pair.refresh_token,
claims,
})
}
pub(crate) async fn sync_account_pool(
&self,
account_id: &str,
+5 -9
View File
@@ -38,15 +38,11 @@ impl AdminAccountService {
.await
.map_err(|e| AdminAccountError::FetchQuota(e.to_string()))?;
let normalized = crate::upstream::accounts::quota::quota_from_usage(&raw);
if let Ok(json_str) = serde_json::to_string(&normalized) {
if matches!(
self.store.update_quota_json(account_id, &json_str).await,
Ok(true)
) {
self.sync_account_pool_best_effort(account_id, "account quota refresh")
.await;
}
}
self.account_pool
.apply_quota_snapshot(account_id, &normalized)
.await;
self.sync_account_pool_best_effort(account_id, "account quota refresh")
.await;
Ok(serde_json::json!({ "quota": normalized, "raw": raw }))
}
}
+6 -33
View File
@@ -220,11 +220,10 @@ fn stored_account_to_model_refresh_account(stored: StoredAccount) -> Account {
}
fn test_responses_request(model: String) -> CodexResponsesRequest {
CodexResponsesRequest {
let mut request = CodexResponsesRequest::new_http_sse(
model,
instructions: "You are checking whether this Codex account can answer. Reply with ok."
.to_string(),
input: vec![json!({
"You are checking whether this Codex account can answer. Reply with ok.",
vec![json!({
"role": "user",
"content": [
{
@@ -233,35 +232,9 @@ fn test_responses_request(model: String) -> CodexResponsesRequest {
}
]
})],
stream: true,
store: false,
reasoning: None,
tools: None,
tool_choice: None,
parallel_tool_calls: None,
text: None,
generate: None,
tuple_schema: None,
service_tier: None,
previous_response_id: None,
prompt_cache_key: None,
explicit_prompt_cache_key: false,
client_conversation_id: None,
variant_identity: None,
include: None,
client_metadata: None,
client_ip: None,
client_user_agent: None,
use_websocket: false,
force_http_sse: true,
turn_state: None,
turn_metadata: None,
beta_features: None,
version: None,
include_timing_metrics: None,
codex_window_id: None,
parent_thread_id: None,
}
);
request.force_http_sse = true;
request
}
async fn process_upstream_test_stream(
+11 -11
View File
@@ -4,7 +4,7 @@ use axum::{
extract::{FromRequestParts, State},
http::{header::SET_COOKIE, request::Parts, HeaderMap, HeaderValue, StatusCode},
response::{IntoResponse, Response},
Json,
Extension, Json,
};
use serde::{Deserialize, Serialize};
@@ -13,6 +13,7 @@ use crate::{
auth::service::AdminSessionError,
response::{AdminEnvelope, AdminError, AdminResponse},
},
http::middleware::request_id::ClientIp,
infra::time::china_rfc3339,
runtime::state::AppState,
};
@@ -57,10 +58,10 @@ struct AdminSessionStatusData {
/// `POST /api/admin/login`
pub(crate) async fn login(
State(state): State<AppState>,
headers: HeaderMap,
client_ip: Option<Extension<ClientIp>>,
Json(payload): Json<AdminLoginRequest>,
) -> Result<Response, AdminError> {
let source = admin_login_source(&headers);
let source = admin_login_source(optional_client_ip(client_ip));
let login_result = state
.services
.admin_sessions
@@ -191,12 +192,11 @@ fn admin_api_key_header(headers: &HeaderMap) -> Option<String> {
(!value.is_empty()).then(|| value.to_string())
}
fn admin_login_source(headers: &HeaderMap) -> String {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("unknown")
.to_string()
fn optional_client_ip(client_ip: Option<Extension<ClientIp>>) -> Option<String> {
let Extension(client_ip) = client_ip?;
Some(client_ip.as_str().to_string())
}
fn admin_login_source(fallback_client_ip: Option<String>) -> String {
fallback_client_ip.unwrap_or_else(|| "unknown".to_string())
}
@@ -37,6 +37,18 @@ impl AdminUsageService {
})
}
/// 按账号 ID 批量读取账号用量。
pub async fn list_by_account_ids(
&self,
account_ids: &[String],
) -> Result<Vec<AdminUsageRecord>, AdminUsageError> {
self.store
.list_usage_by_account_ids(account_ids)
.await
.map(|records| records.into_iter().map(AdminUsageRecord::from).collect())
.map_err(|_| AdminUsageError::List)
}
/// 汇总账号用量。
pub async fn summary(&self) -> Result<AdminUsageSummary, AdminUsageError> {
self.store
@@ -245,6 +245,27 @@ impl SqliteUsageStore {
Ok(to_page(&rows, limit))
}
/// 按账号 ID 批量读取账号用量。
pub async fn list_usage_by_account_ids(
&self,
account_ids: &[String],
) -> SqliteUsageStoreResult<Vec<UsageListRecord>> {
if account_ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(LIST_USAGE_SELECT_SQL);
builder.push("\nwhere au.account_id in (");
let mut separated = builder.separated(", ");
for account_id in account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await?;
rows.iter().map(usage_list_from_row).collect()
}
/// 汇总账号用量。
pub async fn usage_summary(&self) -> SqliteUsageStoreResult<UsageSummary> {
let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?;
+291 -192
View File
@@ -1,6 +1,6 @@
//! 管理端 Dashboard 聚合视图。
use std::{cmp::Reverse, collections::HashMap};
use std::collections::{HashMap, HashSet};
use axum::{
extract::{Query, State},
@@ -15,14 +15,19 @@ use crate::{
admin::{
auth::session::AdminAuth,
monitoring::{
account_usage_service::{
AdminUsageRecord, AdminUsageSummary, AdminUsageTimeBucketRecord,
},
account_usage_service::AdminUsageTimeBucketRecord,
diagnostics::{
fingerprint_diagnostics, AccountCapacityDiagnostics, AccountPoolDiagnostics,
},
usage_record_routes::{usage_record_items, UsageRecordData},
usage_record_model::{
metadata_nonnegative_i64, metadata_string, UsageRecord, UsageRecordLevel,
},
usage_record_routes::{
usage_cost_details, usage_record_models, usage_token_details,
UsageRecordCostDetailsData, UsageRecordTokenDetailsData,
},
usage_record_service::AdminUsageRecordFilter,
usage_record_store::{UsageRecordAccountUsage, UsageRecordSummary},
},
response::{AdminEnvelope, AdminError, AdminResponse},
},
@@ -40,11 +45,15 @@ use crate::{
upstream::accounts::model::{Account, AccountStatus},
};
const DASHBOARD_USAGE_RECORD_LIMIT: u32 = 1000;
const DASHBOARD_ACCOUNT_USAGE_LIMIT: u32 = 4;
const DASHBOARD_USAGE_RECORD_LIMIT: u32 = 10;
const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15;
const DASHBOARD_TIME_BUCKET_SLOTS: i64 = 7 * 24 * 4;
const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4;
const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3;
const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000;
const WEEK_WINDOW_SECONDS: u64 = 604_800;
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -69,7 +78,7 @@ struct DashboardSummaryData {
health_timeline: DashboardHealthTimelineData,
account_usage: Vec<DashboardAccountUsageData>,
service_statuses: Vec<DashboardServiceStatusData>,
usage_records: Vec<UsageRecordData>,
usage_records: Vec<DashboardUsageRecordData>,
pool_summary: AccountPoolDiagnostics,
capacity_info: AccountCapacityDiagnostics,
rotation_strategy: Option<String>,
@@ -80,10 +89,7 @@ struct DashboardSummaryData {
struct DashboardHealthTimelineData {
title: String,
description: String,
range_display: String,
reliability_display: String,
oldest_label: String,
newest_label: String,
points: String,
}
@@ -112,11 +118,8 @@ struct DashboardAccountsCardData {
struct DashboardTrafficCardData {
today_requests: String,
today_requests_value: u64,
yesterday_requests: String,
yesterday_requests_value: u64,
total_requests: String,
rpm: String,
tpm: String,
}
#[derive(Debug, Serialize)]
@@ -124,13 +127,9 @@ struct DashboardTrafficCardData {
struct DashboardTokenCardData {
today_tokens: String,
today_tokens_value: u64,
yesterday_tokens: String,
yesterday_tokens_value: u64,
total_tokens: String,
today_cost_usd: String,
today_cost_usd_value: f64,
total_cost_usd: String,
total_cost_usd_value: f64,
}
#[derive(Debug, Serialize)]
@@ -138,12 +137,10 @@ struct DashboardTokenCardData {
struct DashboardCacheCardData {
today_hit_rate: String,
today_hit_rate_value: Option<f64>,
yesterday_hit_rate: String,
yesterday_hit_rate_value: Option<f64>,
total_hit_rate: String,
total_cached_tokens: String,
first_token_latency_ms: String,
completion_latency_ms: String,
average_first_token_latency_ms: String,
}
#[derive(Debug, Serialize)]
@@ -166,7 +163,7 @@ struct DashboardTrendPointData {
output_tokens_value: u64,
cached_tokens: String,
cached_tokens_value: u64,
tokens: String,
cache_hit_rate_value: f64,
tokens_value: u64,
errors: String,
errors_value: u64,
@@ -191,14 +188,34 @@ struct DashboardTrendSummaryData {
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardAccountUsageData {
name: String,
id: String,
email: String,
plan: String,
requests: String,
plan_type: Option<String>,
tokens: String,
quota_used_percent: Option<f64>,
last_used: String,
status: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardUsageRecordData {
id: String,
level: UsageRecordLevel,
route: Option<String>,
model: Option<String>,
status_code: Option<i64>,
transport: Option<String>,
created_at_display: String,
account_email: Option<String>,
requested_model: Option<String>,
upstream_model: Option<String>,
client_ip: Option<String>,
user_agent: Option<String>,
reasoning_effort: Option<String>,
token_details: UsageRecordTokenDetailsData,
cost_details: Option<UsageRecordCostDetailsData>,
first_token_latency_ms_display: String,
latency_ms_display: String,
}
#[derive(Debug, Serialize)]
@@ -219,10 +236,16 @@ struct UsageWindow {
errors: u64,
first_token_latency_sum: u64,
first_token_latency_count: u64,
latency_sum: u64,
latency_count: u64,
max_latency: u64,
min_latency: u64,
max_first_token_bucket_latency: u64,
min_first_token_bucket_latency: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
enum QuotaWindowPriority {
FiveHour,
Weekly,
Monthly,
Other,
}
impl UsageWindow {
@@ -243,8 +266,12 @@ impl UsageWindow {
.checked_div(self.first_token_latency_count)
}
fn avg_latency(self) -> Option<u64> {
self.latency_sum.checked_div(self.latency_count)
fn max_first_token_bucket_latency(self) -> Option<u64> {
(self.max_first_token_bucket_latency > 0).then_some(self.max_first_token_bucket_latency)
}
fn min_first_token_bucket_latency(self) -> Option<u64> {
(self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency)
}
}
@@ -252,6 +279,7 @@ impl UsageWindow {
pub(crate) async fn dashboard_summary(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
let accounts = state
.services
@@ -260,53 +288,66 @@ pub(crate) async fn dashboard_summary(
.await
.unwrap_or_default();
let capacity = state.services.account_pool.capacity_summary_now().await;
let summary = state
let now = Utc::now();
let today_filter = today_usage_record_filter(now);
let usage_summary = state
.services
.usage
.summary()
.usage_records
.summary(AdminUsageRecordFilter::default())
.await
.unwrap_or_else(|_| empty_usage_summary());
let usage_records = state
.unwrap_or_else(|_| UsageRecordSummary::default());
let account_usage_records = state
.services
.usage
.list(None, 200)
.usage_records
.account_usage(today_filter.clone(), DASHBOARD_ACCOUNT_USAGE_LIMIT)
.await
.map(|page| page.items)
.unwrap_or_default();
let recent_events = state
.services
.usage_records
.list(
None,
DASHBOARD_USAGE_RECORD_LIMIT,
AdminUsageRecordFilter::default(),
)
.list(None, DASHBOARD_USAGE_RECORD_LIMIT, today_filter)
.await
.map(|page| page.items)
.unwrap_or_default();
let now = Utc::now();
let time_buckets = dashboard_time_buckets(&state, now).await;
let pool_summary = account_pool_summary(&accounts);
let trend = dashboard_trend_data(&time_buckets, DashboardTrendKind::Usage);
let quota_used_by_account = account_quota_used_percent_by_id(&state, &usage_records).await;
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let refreshing_account_ids = state
.services
.token_refresh
.refreshing_account_ids(&account_ids, now)
.await
.map_err(|error| AdminError::internal(error.to_string()))?;
let pool_summary = account_pool_summary(&accounts, &refreshing_account_ids);
let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default());
let quota_used_by_account =
account_quota_used_percent_by_id(&state, &account_usage_records).await;
let settings = state.services.settings.current();
let recent_usage_records = recent_events.into_iter().take(10).collect::<Vec<_>>();
let recent_usage_records = recent_events;
let account_emails = state
.services
.usage_records
.account_email_map(&recent_usage_records)
.await
.map_err(|error| AdminError::usage_record_accounts_failed(error.to_string()))?;
let dashboard_usage_records = usage_record_items(recent_usage_records, &account_emails);
let dashboard_usage_records =
dashboard_usage_record_items(recent_usage_records, &account_emails);
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(DashboardSummaryData {
cards: dashboard_cards(&accounts, &summary, &time_buckets),
cards: dashboard_cards(&accounts, &time_buckets, &usage_summary),
trend,
health_timeline: dashboard_health_timeline_data(&time_buckets),
account_usage: account_usage_data(&accounts, &usage_records, &quota_used_by_account),
account_usage: account_usage_data(
&accounts,
&account_usage_records,
&quota_used_by_account,
now,
),
service_statuses: service_status_data(&state),
usage_records: dashboard_usage_records,
pool_summary,
@@ -316,6 +357,14 @@ pub(crate) async fn dashboard_summary(
))
}
fn today_usage_record_filter(now: DateTime<Utc>) -> AdminUsageRecordFilter {
AdminUsageRecordFilter {
start_time: Some(china_day_start(now)),
end_time: Some(now),
..AdminUsageRecordFilter::default()
}
}
/// `GET /api/admin/dashboard/trend?kind=usage|latency|errors`
pub(crate) async fn dashboard_trend(
State(state): State<AppState>,
@@ -335,16 +384,16 @@ pub(crate) async fn dashboard_trend(
fn dashboard_cards(
accounts: &[Account],
summary: &AdminUsageSummary,
buckets: &[AdminUsageTimeBucketRecord],
usage_summary: &UsageRecordSummary,
) -> DashboardCardsData {
dashboard_cards_at(accounts, summary, buckets, Utc::now())
dashboard_cards_at(accounts, buckets, usage_summary, Utc::now())
}
fn dashboard_cards_at(
accounts: &[Account],
summary: &AdminUsageSummary,
buckets: &[AdminUsageTimeBucketRecord],
usage_summary: &UsageRecordSummary,
now: DateTime<Utc>,
) -> DashboardCardsData {
let today_start = china_day_start(now);
@@ -352,10 +401,11 @@ fn dashboard_cards_at(
let today = usage_window(buckets, today_start, now);
let yesterday = usage_window(buckets, yesterday_start, today_start);
let today_cost = cost_window(buckets, today_start, now).unwrap_or(0.0);
let total_cost = total_cost(buckets);
let total_requests = usage_summary.total_requests;
let total_tokens = usage_summary.total_tokens;
let total_cached_tokens = usage_summary.cached_tokens;
let total_hit_rate = summary_cache_hit_rate(usage_summary);
let total_input = nonnegative_i64_to_u64(summary.input_tokens);
let total_cached = nonnegative_i64_to_u64(summary.cached_tokens);
let total_accounts = accounts.len() as u64;
let enabled_accounts = accounts
.iter()
@@ -373,13 +423,6 @@ fn dashboard_cards_at(
)
})
.count() as u64;
let total_requests = nonnegative_i64_to_u64(summary.request_count);
let total_tokens = nonnegative_i64_to_u64(summary.input_tokens + summary.output_tokens);
let total_hit_rate = if total_input > 0 {
Some(total_cached as f64 / total_input as f64)
} else {
None
};
DashboardCardsData {
accounts: DashboardAccountsCardData {
@@ -393,38 +436,36 @@ fn dashboard_cards_at(
traffic: DashboardTrafficCardData {
today_requests: format_compact_number(today.requests),
today_requests_value: today.requests,
yesterday_requests: format_compact_number(yesterday.requests),
yesterday_requests_value: yesterday.requests,
total_requests: format_compact_number(total_requests),
rpm: format_compact_number(today.requests),
tpm: format_tokens(today.tokens()),
},
tokens: DashboardTokenCardData {
today_tokens: format_tokens(today.tokens()),
today_tokens_value: today.tokens(),
yesterday_tokens: format_tokens(yesterday.tokens()),
yesterday_tokens_value: yesterday.tokens(),
total_tokens: format_tokens(total_tokens),
today_cost_usd: format_optional_cost(Some(today_cost)),
today_cost_usd_value: today_cost,
total_cost_usd: format_optional_cost(Some(total_cost)),
total_cost_usd_value: total_cost,
today_cost_usd: format_cost(today_cost),
},
cache: DashboardCacheCardData {
today_hit_rate: format_rate(today.cache_hit_rate()),
today_hit_rate_value: today.cache_hit_rate(),
yesterday_hit_rate: format_rate(yesterday.cache_hit_rate()),
yesterday_hit_rate_value: yesterday.cache_hit_rate(),
total_hit_rate: format_rate(total_hit_rate),
total_cached_tokens: format_tokens(total_cached),
first_token_latency_ms: format_optional_duration_ms(today.avg_first_token_latency()),
completion_latency_ms: format_optional_duration_ms(today.avg_latency()),
total_cached_tokens: format_tokens(total_cached_tokens),
average_first_token_latency_ms: format_optional_duration_ms(
today.avg_first_token_latency(),
),
},
}
}
fn format_optional_cost(value: Option<f64>) -> String {
value.map_or_else(|| "-".to_string(), format_cost)
fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option<f64> {
let input_tokens = summary.input_tokens;
if input_tokens > 0 {
Some(summary.cached_tokens as f64 / input_tokens as f64)
} else {
None
}
}
fn format_optional_duration_ms(value: Option<u64>) -> String {
@@ -459,8 +500,9 @@ fn dashboard_trend_data_at(
now: DateTime<Utc>,
) -> DashboardTrendData {
let current_hour = china_hour_start(now);
let start = current_hour - Duration::hours(23);
let mut buckets = (0..24)
let start = china_day_start(now);
let elapsed_hours = current_hour.signed_duration_since(start).num_hours();
let mut buckets = (0..=elapsed_hours)
.map(|index| {
let bucket_start = start + Duration::hours(index);
(bucket_start, UsageWindow::default())
@@ -483,8 +525,11 @@ fn dashboard_trend_data_at(
let points = buckets
.iter()
.map(|(bucket_start, bucket)| {
let latency = bucket.avg_latency().unwrap_or(0);
let latency = bucket.avg_first_token_latency().unwrap_or(0);
let max_latency = bucket.max_first_token_bucket_latency().unwrap_or(0);
let min_latency = bucket.min_first_token_bucket_latency().unwrap_or(0);
let tokens = bucket.tokens();
let cache_hit_rate = bucket.cache_hit_rate().unwrap_or(0.0);
let success_rate = if bucket.requests > 0 {
((bucket.requests - bucket.errors) as f64 / bucket.requests as f64 * 1000.0).round()
/ 10.0
@@ -501,16 +546,16 @@ fn dashboard_trend_data_at(
output_tokens_value: bucket.output_tokens,
cached_tokens: format_tokens(bucket.cached_tokens),
cached_tokens_value: bucket.cached_tokens,
tokens: format_tokens(tokens),
cache_hit_rate_value: cache_hit_rate,
tokens_value: tokens,
errors: format_compact_number(bucket.errors),
errors_value: bucket.errors,
latency: format_optional_duration_ms(Some(latency)),
latency_value: latency,
max_latency: format_optional_duration_ms(Some(bucket.max_latency)),
max_latency_value: bucket.max_latency,
min_latency: format_optional_duration_ms(Some(bucket.min_latency)),
min_latency_value: bucket.min_latency,
max_latency: format_optional_duration_ms(Some(max_latency)),
max_latency_value: max_latency,
min_latency: format_optional_duration_ms(Some(min_latency)),
min_latency_value: min_latency,
success_rate: format_percent(success_rate),
success_rate_value: success_rate,
}
@@ -574,15 +619,12 @@ fn dashboard_health_timeline_data_at(
};
DashboardHealthTimelineData {
title: "请求健康时间线".to_string(),
description: "今日请求可靠性".to_string(),
range_display: format!("{} - {}", china_datetime(&start), china_datetime(&now)),
description: "请求可靠性".to_string(),
reliability_display: if requests > 0 {
format!("{reliability:.1}%")
} else {
"-".to_string()
},
oldest_label: "00:00".to_string(),
newest_label: "24:00".to_string(),
points: buckets
.into_iter()
.map(|(bucket_start, bucket)| {
@@ -748,22 +790,6 @@ fn cost_window(
has_usage.then_some(total)
}
fn total_cost(records: &[AdminUsageTimeBucketRecord]) -> f64 {
let mut total = 0.0;
let mut has_usage = false;
for record in records {
if let Some(cost) = record.cost_usd {
total += cost;
has_usage = true;
}
}
if has_usage {
return total;
}
0.0
}
fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) {
window.requests += nonnegative_i64_to_u64(record.request_count);
window.input_tokens += nonnegative_i64_to_u64(record.input_tokens);
@@ -772,33 +798,41 @@ fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) {
window.errors += nonnegative_i64_to_u64(record.error_count);
window.first_token_latency_sum += nonnegative_i64_to_u64(record.first_token_latency_sum);
window.first_token_latency_count += nonnegative_i64_to_u64(record.first_token_latency_count);
window.latency_sum += nonnegative_i64_to_u64(record.latency_sum);
window.latency_count += nonnegative_i64_to_u64(record.latency_count);
let max_latency = nonnegative_i64_to_u64(record.max_latency_ms);
if max_latency > 0 {
window.max_latency = window.max_latency.max(max_latency);
}
let min_latency = nonnegative_i64_to_u64(record.min_latency_ms);
if min_latency > 0 {
window.min_latency = if window.min_latency == 0 {
min_latency
let bucket_first_token_latency = first_token_bucket_latency(record);
if let Some(latency) = bucket_first_token_latency {
window.max_first_token_bucket_latency = window.max_first_token_bucket_latency.max(latency);
window.min_first_token_bucket_latency = if window.min_first_token_bucket_latency == 0 {
latency
} else {
window.min_latency.min(min_latency)
window.min_first_token_bucket_latency.min(latency)
};
}
}
fn account_pool_summary(accounts: &[Account]) -> AccountPoolDiagnostics {
fn first_token_bucket_latency(record: &AdminUsageTimeBucketRecord) -> Option<u64> {
let sum = nonnegative_i64_to_u64(record.first_token_latency_sum);
let count = nonnegative_i64_to_u64(record.first_token_latency_count);
sum.checked_div(count).filter(|latency| *latency > 0)
}
fn account_pool_summary(
accounts: &[Account],
refreshing_account_ids: &HashSet<String>,
) -> AccountPoolDiagnostics {
let mut summary = AccountPoolDiagnostics {
total: accounts.len(),
..AccountPoolDiagnostics::default()
};
for account in accounts {
match account.status {
AccountStatus::Active | AccountStatus::Expired
if refreshing_account_ids.contains(&account.id) =>
{
summary.refreshing += 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,
}
@@ -808,68 +842,51 @@ fn account_pool_summary(accounts: &[Account]) -> AccountPoolDiagnostics {
fn account_usage_data(
accounts: &[Account],
usage_records: &[AdminUsageRecord],
usage_records: &[UsageRecordAccountUsage],
quota_used_by_account: &HashMap<String, f64>,
now: DateTime<Utc>,
) -> Vec<DashboardAccountUsageData> {
let account_by_id = accounts
.iter()
.map(|account| (account.id.as_str(), account))
.collect::<HashMap<_, _>>();
let mut records = usage_records.iter().collect::<Vec<_>>();
records.sort_by_key(|record| Reverse((record.last_used_at, record.account_id.clone())));
let mut rows = records
usage_records
.iter()
.map(|usage| {
let account = account_by_id.get(usage.account_id.as_str()).copied();
let name = account
.and_then(|account| account.label.clone())
.or_else(|| usage.label.clone())
.or_else(|| {
usage
.email
.as_ref()
.and_then(|email| email.split('@').next().map(ToString::to_string))
})
.unwrap_or_else(|| usage.account_id.chars().take(8).collect());
let status = account.map_or_else(
|| "active".to_string(),
|account| account.status.to_string(),
);
let quota_used_percent = quota_used_by_account
.get(&usage.account_id)
.copied()
.or_else(|| (status == "quota_exhausted").then_some(100.0));
let requests = nonnegative_i64_to_u64(usage.request_count);
let tokens = nonnegative_i64_to_u64(usage.input_tokens + usage.output_tokens);
.or_else(|| {
matches!(
account.map(|account| account.status),
Some(AccountStatus::QuotaExhausted)
)
.then_some(100.0)
});
DashboardAccountUsageData {
name,
email: usage.email.clone().unwrap_or_else(|| "-".to_string()),
plan: usage
.plan_type
.clone()
.unwrap_or_else(|| "free".to_string()),
requests: format_compact_number(requests),
tokens: format_tokens(tokens),
id: usage.account_id.clone(),
email: account
.and_then(|account| account.email.as_ref())
.cloned()
.expect("dashboard account usage requires email"),
plan_type: account
.and_then(|account| account.plan_type.as_ref())
.cloned(),
tokens: format_tokens(usage.total_tokens),
quota_used_percent,
last_used: china_relative_time(usage.last_used_at, Utc::now()),
status,
last_used: china_relative_time(Some(usage.last_used_at), now),
}
})
.collect::<Vec<_>>();
rows.truncate(4);
rows
.collect()
}
async fn account_quota_used_percent_by_id(
state: &AppState,
usage_records: &[AdminUsageRecord],
usage_records: &[UsageRecordAccountUsage],
) -> HashMap<String, f64> {
let mut records = usage_records.iter().collect::<Vec<_>>();
records.sort_by_key(|record| Reverse((record.last_used_at, record.account_id.clone())));
records.truncate(4);
let mut quota_used_by_account = HashMap::with_capacity(records.len());
for usage in records {
let mut quota_used_by_account = HashMap::with_capacity(usage_records.len());
for usage in usage_records {
let Ok(Some(quota_json)) = state
.services
.accounts
@@ -887,27 +904,67 @@ async fn account_quota_used_percent_by_id(
fn quota_used_percent(quota_json: &str) -> Option<f64> {
let quota = serde_json::from_str::<Value>(quota_json).ok()?;
let mut values = Vec::new();
if let Some(used_percent) = quota
.pointer("/monthly_limit/used_percent")
.and_then(percent_value)
{
values.push(used_percent);
let mut selected = None;
if let Some(monthly_limit) = quota.get("monthly_limit") {
select_quota_used_percent(
&mut selected,
quota_window_priority(monthly_limit, QuotaWindowPriority::Monthly),
monthly_limit.get("used_percent").and_then(percent_value),
);
}
if let Some(snapshots) = quota.get("snapshots").and_then(Value::as_array) {
for snapshot in snapshots {
for role in ["primary", "secondary"] {
if let Some(used_percent) = snapshot
.get(role)
.and_then(|window| window.get("used_percent"))
.and_then(percent_value)
{
values.push(used_percent);
if let Some(window) = snapshot.get(role) {
select_quota_used_percent(
&mut selected,
quota_window_priority(window, QuotaWindowPriority::Other),
window.get("used_percent").and_then(percent_value),
);
}
}
}
}
values.into_iter().max_by(f64::total_cmp)
selected.map(|(_, used_percent)| used_percent)
}
fn select_quota_used_percent(
selected: &mut Option<(QuotaWindowPriority, f64)>,
priority: QuotaWindowPriority,
used_percent: Option<f64>,
) {
let Some(used_percent) = used_percent else {
return;
};
match selected {
Some((selected_priority, selected_percent))
if priority > *selected_priority
|| (priority == *selected_priority && used_percent <= *selected_percent) => {}
_ => *selected = Some((priority, used_percent)),
}
}
fn quota_window_priority(window: &Value, fallback: QuotaWindowPriority) -> QuotaWindowPriority {
let window_seconds = window
.get("window_minutes")
.and_then(Value::as_u64)
.and_then(|minutes| minutes.checked_mul(60));
match window_seconds {
Some(seconds) if quota_window_matches(seconds, FIVE_HOUR_WINDOW_SECONDS) => {
QuotaWindowPriority::FiveHour
}
Some(seconds) if quota_window_matches(seconds, WEEK_WINDOW_SECONDS) => {
QuotaWindowPriority::Weekly
}
Some(seconds) if quota_window_matches(seconds, MONTH_WINDOW_SECONDS) => {
QuotaWindowPriority::Monthly
}
_ => fallback,
}
}
fn quota_window_matches(actual: u64, expected: u64) -> bool {
actual > 0 && actual.abs_diff(expected) <= expected / 20
}
fn percent_value(value: &Value) -> Option<f64> {
@@ -917,8 +974,67 @@ fn percent_value(value: &Value) -> Option<f64> {
percent.is_finite().then_some(percent.clamp(0.0, 100.0))
}
fn dashboard_usage_record_items(
items: Vec<UsageRecord>,
account_emails: &HashMap<String, String>,
) -> Vec<DashboardUsageRecordData> {
items
.into_iter()
.map(|record| dashboard_usage_record_data(record, account_emails))
.collect()
}
fn dashboard_usage_record_data(
record: UsageRecord,
account_emails: &HashMap<String, String>,
) -> DashboardUsageRecordData {
let account_email = record
.account_id
.as_deref()
.and_then(|account_id| account_emails.get(account_id))
.cloned()
.or_else(|| metadata_string(&record.metadata, &["accountEmail", "account_email"]));
let (requested_model, upstream_model) = usage_record_models(&record);
let client_ip = metadata_string(&record.metadata, &["clientIp", "ipAddress", "ip_address"]);
let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]);
let reasoning_effort =
metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]);
let token_details = usage_token_details(&record.metadata);
let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details);
let first_token_latency_ms = metadata_nonnegative_i64(
&record.metadata,
&[
"firstTokenMs",
"first_token_ms",
"firstTokenLatencyMs",
"first_token_latency_ms",
],
);
DashboardUsageRecordData {
id: record.id,
level: record.level,
route: record.route,
model: record.model,
status_code: record.status_code,
transport: record.transport,
created_at_display: china_datetime(&record.created_at),
account_email,
requested_model,
upstream_model,
client_ip,
user_agent,
reasoning_effort,
token_details,
cost_details,
first_token_latency_ms_display: format_duration_ms(first_token_latency_ms),
latency_ms_display: format_duration_ms(record.latency_ms),
}
}
fn service_status_data(state: &AppState) -> Vec<DashboardServiceStatusData> {
let fingerprint = fingerprint_diagnostics(&state.services.fingerprint);
let current_fingerprint = state.services.fingerprint.snapshot();
let fingerprint = fingerprint_diagnostics(&current_fingerprint);
vec![
DashboardServiceStatusData {
label: "客户端版本".to_string(),
@@ -979,20 +1095,3 @@ fn format_fingerprint_updated_at(value: Option<String>) -> String {
}
china_datetime_rfc3339_str(value)
}
fn empty_usage_summary() -> AdminUsageSummary {
AdminUsageSummary {
account_count: 0,
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,
}
}
+3
View File
@@ -5,6 +5,9 @@ pub mod account_usage_store;
pub(crate) mod billing;
pub mod dashboard;
pub mod diagnostics;
pub mod ops_error_model;
pub mod ops_error_service;
pub mod ops_error_store;
pub mod usage_record_model;
pub mod usage_record_routes;
pub mod usage_record_service;
@@ -0,0 +1,75 @@
//! 运维错误事件模型。
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use uuid::Uuid;
/// 结构化运维错误事件。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct OpsErrorLog {
/// 事件 ID。
pub id: String,
/// 请求 ID。
pub request_id: Option<String>,
/// 事件类别。
pub kind: String,
/// 账号 ID。预账号调度失败时为空。
pub account_id: Option<String>,
/// HTTP 路由。
pub route: Option<String>,
/// 模型名。
pub model: Option<String>,
/// 事件状态码。
pub status_code: Option<i64>,
/// 客户端看到的状态码。
pub client_status_code: Option<i64>,
/// 上游 HTTP 状态码。
pub upstream_status_code: Option<i64>,
/// 上游传输方式,例如 http_sse 或 websocket。
pub transport: Option<String>,
/// 同一请求内的上游尝试序号。
pub attempt_index: Option<i64>,
/// 失败分类。
pub failure_class: Option<String>,
/// 上游响应 ID。
pub response_id: Option<String>,
/// 上游请求 ID。
pub upstream_request_id: Option<String>,
/// 请求耗时毫秒。
pub latency_ms: Option<i64>,
/// 展示消息。
pub message: String,
/// 结构化元数据。
pub metadata: Value,
/// 创建时间。
#[serde(serialize_with = "crate::infra::time::serialize_china_rfc3339")]
pub created_at: DateTime<Utc>,
}
impl OpsErrorLog {
/// 构造运维错误事件。
pub fn new(kind: impl Into<String>, message: impl Into<String>) -> Self {
Self {
id: format!("ops_error_{}", Uuid::new_v4().simple()),
request_id: None,
kind: kind.into(),
account_id: None,
route: None,
model: None,
status_code: None,
client_status_code: None,
upstream_status_code: None,
transport: None,
attempt_index: None,
failure_class: None,
response_id: None,
upstream_request_id: None,
latency_ms: None,
message: message.into(),
metadata: Value::Object(Map::default()),
created_at: Utc::now(),
}
}
}
@@ -0,0 +1,68 @@
//! 管理端运维错误事件服务。
use chrono::Utc;
use thiserror::Error;
use crate::admin::monitoring::{
ops_error_model::OpsErrorLog, ops_error_store::SqliteOpsErrorLogStore,
};
/// 管理端运维错误事件错误。
#[derive(Debug, Error)]
pub enum AdminOpsErrorLogError {
/// 写入失败。
#[error("failed to append ops error log")]
Append,
/// 保留期清理失败。
#[error("failed to trim expired ops error logs")]
Retention,
}
/// 管理端运维错误事件服务。
#[derive(Clone)]
pub struct AdminOpsErrorLogService {
store: SqliteOpsErrorLogStore,
capture_body: bool,
}
impl AdminOpsErrorLogService {
/// 构造管理端运维错误事件服务。
pub fn new(store: SqliteOpsErrorLogStore, capture_body: bool) -> Self {
Self {
store,
capture_body,
}
}
/// 记录运维错误事件。
pub async fn record(&self, mut event: OpsErrorLog) -> Result<(), AdminOpsErrorLogError> {
apply_capture_body_policy(&mut event, self.capture_body);
self.store
.append(&event)
.await
.map_err(|_| AdminOpsErrorLogError::Append)?;
self.store
.trim_to_retention(Utc::now())
.await
.map_err(|_| AdminOpsErrorLogError::Retention)?;
Ok(())
}
}
fn apply_capture_body_policy(event: &mut OpsErrorLog, capture_body: bool) {
if capture_body {
return;
}
let Some(metadata) = event.metadata.as_object_mut() else {
return;
};
for key in [
"body",
"rawBody",
"requestBody",
"responseBody",
"upstreamBody",
] {
metadata.remove(key);
}
}
@@ -0,0 +1,261 @@
//! 运维错误事件存储实现(SQLite)。
use chrono::{DateTime, Duration, Utc};
use serde_json::Value;
use sqlx::SqlitePool;
use thiserror::Error;
use crate::{
admin::monitoring::{
ops_error_model::OpsErrorLog,
usage_record_model::{metadata_i64, metadata_service_tier, metadata_string},
usage_record_store::USAGE_RECORD_RETENTION_DAYS,
},
infra::time::china_quarter_hour_start,
};
/// SQLite 运维错误事件错误。
#[derive(Debug, Error)]
pub enum SqliteOpsErrorLogStoreError {
/// 数据库错误。
#[error("sqlite ops error log database error: {0}")]
Database(#[from] sqlx::Error),
}
/// SQLite 运维错误事件结果。
pub type SqliteOpsErrorLogStoreResult<T> = Result<T, SqliteOpsErrorLogStoreError>;
/// SQLite 运维错误事件存储。
#[derive(Clone)]
pub struct SqliteOpsErrorLogStore {
pool: SqlitePool,
}
impl SqliteOpsErrorLogStore {
/// 构造存储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 写入运维错误事件。
pub async fn append(&self, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> {
append_event(&self.pool, event).await
}
/// 按保留期清理运维错误事件。
pub async fn trim_to_retention(&self, now: DateTime<Utc>) -> SqliteOpsErrorLogStoreResult<u64> {
trim_to_retention(&self.pool, now).await
}
/// 清空运维错误事件。
pub async fn clear(&self) -> SqliteOpsErrorLogStoreResult<u64> {
let result = sqlx::query("delete from ops_error_logs")
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
/// 读取错误事件数量。
pub async fn count(&self) -> SqliteOpsErrorLogStoreResult<u64> {
let (total,): (i64,) = sqlx::query_as("select count(*) from ops_error_logs")
.fetch_one(&self.pool)
.await?;
Ok(total.max(0) as u64)
}
}
#[derive(Debug, Clone)]
struct OpsErrorSummaryFields {
client_status_code: Option<i64>,
upstream_status_code: Option<i64>,
transport: Option<String>,
attempt_index: Option<i64>,
failure_class: Option<String>,
response_id: Option<String>,
upstream_request_id: Option<String>,
}
fn ops_error_summary_fields(event: &OpsErrorLog) -> OpsErrorSummaryFields {
OpsErrorSummaryFields {
client_status_code: event.client_status_code.or_else(|| {
metadata_i64(
&event.metadata,
&["clientStatusCode", "client_status_code", "clientStatus"],
)
}),
upstream_status_code: event.upstream_status_code.or_else(|| {
metadata_i64(
&event.metadata,
&[
"upstreamStatus",
"upstreamStatusCode",
"upstream_status_code",
],
)
}),
transport: event
.transport
.clone()
.or_else(|| metadata_string(&event.metadata, &["transport"])),
attempt_index: event
.attempt_index
.or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"])),
failure_class: event
.failure_class
.clone()
.or_else(|| metadata_string(&event.metadata, &["failureClass", "failure_class"])),
response_id: event
.response_id
.clone()
.or_else(|| metadata_string(&event.metadata, &["responseId", "response_id"])),
upstream_request_id: event.upstream_request_id.clone().or_else(|| {
metadata_string(
&event.metadata,
&[
"upstreamRequestId",
"upstream_request_id",
"openaiRequestId",
],
)
}),
}
}
async fn append_event(pool: &SqlitePool, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> {
let summary = ops_error_summary_fields(event);
sqlx::query(
"insert into ops_error_logs (id, request_id, kind, account_id, route, model, status_code, client_status_code, upstream_status_code, transport, attempt_index, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(&event.id)
.bind(&event.request_id)
.bind(&event.kind)
.bind(&event.account_id)
.bind(&event.route)
.bind(&event.model)
.bind(event.status_code)
.bind(summary.client_status_code)
.bind(summary.upstream_status_code)
.bind(summary.transport)
.bind(summary.attempt_index)
.bind(summary.failure_class)
.bind(summary.response_id)
.bind(summary.upstream_request_id)
.bind(event.latency_ms)
.bind(&event.message)
.bind(event.metadata.to_string())
.bind(event.created_at.to_rfc3339())
.execute(pool)
.await?;
append_error_time_bucket(pool, event).await?;
Ok(())
}
async fn append_error_time_bucket(
pool: &SqlitePool,
event: &OpsErrorLog,
) -> SqliteOpsErrorLogStoreResult<()> {
let bucket_start = china_quarter_hour_start(event.created_at);
let input_tokens = metadata_usage_i64(&event.metadata, "inputTokens");
let output_tokens = metadata_usage_i64(&event.metadata, "outputTokens");
let cached_tokens = metadata_usage_i64(&event.metadata, "cachedTokens");
let first_token_latency = metadata_first_token_latency(&event.metadata).unwrap_or(0);
let first_token_latency_count = i64::from(first_token_latency > 0);
let latency = event.latency_ms.filter(|value| *value > 0).unwrap_or(0);
let latency_count = i64::from(latency > 0);
let now = Utc::now().to_rfc3339();
sqlx::query(
r"
insert into usage_time_buckets (
bucket_start,
account_id,
model,
service_tier,
request_count,
error_count,
input_tokens,
output_tokens,
cached_tokens,
first_token_latency_sum,
first_token_latency_count,
latency_sum,
latency_count,
max_latency_ms,
min_latency_ms,
updated_at
) values (?, ?, ?, ?, 1, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
on conflict(bucket_start, account_id, model, service_tier) do update set
request_count = request_count + excluded.request_count,
error_count = error_count + excluded.error_count,
input_tokens = input_tokens + excluded.input_tokens,
output_tokens = output_tokens + excluded.output_tokens,
cached_tokens = cached_tokens + excluded.cached_tokens,
first_token_latency_sum = first_token_latency_sum + excluded.first_token_latency_sum,
first_token_latency_count = first_token_latency_count + excluded.first_token_latency_count,
latency_sum = latency_sum + excluded.latency_sum,
latency_count = latency_count + excluded.latency_count,
max_latency_ms = max(max_latency_ms, excluded.max_latency_ms),
min_latency_ms = case
when min_latency_ms = 0 then excluded.min_latency_ms
when excluded.min_latency_ms = 0 then min_latency_ms
else min(min_latency_ms, excluded.min_latency_ms)
end,
updated_at = excluded.updated_at",
)
.bind(bucket_start.to_rfc3339())
.bind(event.account_id.as_deref().unwrap_or_default())
.bind(event.model.as_deref().unwrap_or_default())
.bind(metadata_service_tier(&event.metadata).unwrap_or_default())
.bind(input_tokens)
.bind(output_tokens)
.bind(cached_tokens)
.bind(first_token_latency)
.bind(first_token_latency_count)
.bind(latency)
.bind(latency_count)
.bind(latency)
.bind(latency)
.bind(now)
.execute(pool)
.await?;
Ok(())
}
fn metadata_usage_i64(metadata: &Value, field: &str) -> i64 {
metadata
.get("usage")
.and_then(|usage| usage.get(field))
.or_else(|| metadata.get(field))
.and_then(Value::as_i64)
.filter(|value| *value > 0)
.unwrap_or(0)
}
fn metadata_first_token_latency(metadata: &Value) -> Option<i64> {
[
"firstTokenMs",
"first_token_ms",
"firstTokenLatencyMs",
"first_token_latency_ms",
]
.into_iter()
.find_map(|field| {
metadata
.get(field)
.or_else(|| metadata.get("usage").and_then(|usage| usage.get(field)))
.and_then(Value::as_i64)
.filter(|value| *value > 0)
})
}
async fn trim_to_retention(
pool: &SqlitePool,
now: DateTime<Utc>,
) -> SqliteOpsErrorLogStoreResult<u64> {
let cutoff = now - Duration::days(USAGE_RECORD_RETENTION_DAYS);
let result = sqlx::query("delete from ops_error_logs where created_at < ?")
.bind(cutoff.to_rfc3339())
.execute(pool)
.await?;
Ok(result.rows_affected())
}
@@ -29,8 +29,8 @@ use crate::{
},
infra::{
format::{
format_compact_number, format_duration_ms, format_duration_ms_f64, format_multiplier,
format_number, format_percent, format_precise_cost, format_rate, format_token_price,
format_compact_number, format_cost, format_duration_ms, format_duration_ms_f64,
format_multiplier, format_number, format_percent, format_rate, format_token_price,
format_tokens,
},
json::{clamp_limit, clamp_page, NumberedPage, Page},
@@ -125,7 +125,6 @@ pub(crate) struct UsageRecordCostDetailsData {
output_cost: f64,
cache_read_cost: f64,
total_cost: f64,
billed_cost: f64,
original_cost: f64,
input_price_per_mtoken: f64,
output_price_per_mtoken: f64,
@@ -137,7 +136,6 @@ pub(crate) struct UsageRecordCostDetailsData {
output_cost_display: String,
cache_read_cost_display: String,
total_cost_display: String,
billed_cost_display: String,
original_cost_display: String,
input_price_display: String,
output_price_display: String,
@@ -332,19 +330,7 @@ pub(crate) async fn usage_records_token_trend(
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
let filter = filter_from_query(query)?;
match state.services.usage_records.token_trend(filter).await {
Ok(points) => Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
points
.into_iter()
.map(UsageRecordTrendPointData::from)
.collect::<Vec<_>>(),
),
)),
Err(error) => Err(log_error(&error)),
}
usage_records_trend_response(state, query).await
}
/// `GET /api/admin/usage/records/insights/latency-trend`
@@ -353,8 +339,15 @@ pub(crate) async fn usage_records_latency_trend(
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
usage_records_trend_response(state, query).await
}
async fn usage_records_trend_response(
state: AppState,
query: UsageRecordsQuery,
) -> Result<AdminResponse<AdminEnvelope<Vec<UsageRecordTrendPointData>>>, AdminError> {
let filter = filter_from_query(query)?;
match state.services.usage_records.latency_trend(filter).await {
match state.services.usage_records.trend(filter).await {
Ok(points) => Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
@@ -470,7 +463,7 @@ fn log_error(error: &AdminUsageRecordError) -> AdminError {
| AdminUsageRecordError::Get
| AdminUsageRecordError::Clear
| AdminUsageRecordError::Append
| AdminUsageRecordError::Trim => AdminError::internal(error.to_string()),
| AdminUsageRecordError::Retention => AdminError::internal(error.to_string()),
AdminUsageRecordError::Accounts => {
AdminError::usage_record_accounts_failed(error.to_string())
}
@@ -552,7 +545,7 @@ fn usage_record_data(
}
}
fn usage_record_models(record: &UsageRecord) -> (Option<String>, Option<String>) {
pub(crate) fn usage_record_models(record: &UsageRecord) -> (Option<String>, Option<String>) {
let stored_model = record
.model
.as_deref()
@@ -567,7 +560,7 @@ fn usage_record_models(record: &UsageRecord) -> (Option<String>, Option<String>)
(requested_model, upstream_model)
}
fn usage_token_details(metadata: &Value) -> UsageRecordTokenDetailsData {
pub(crate) fn usage_token_details(metadata: &Value) -> UsageRecordTokenDetailsData {
let input_tokens = metadata_usage_u64(metadata, "inputTokens");
let output_tokens = metadata_usage_u64(metadata, "outputTokens");
let cached_tokens = metadata_usage_u64(metadata, "cachedTokens");
@@ -589,7 +582,7 @@ fn usage_token_details(metadata: &Value) -> UsageRecordTokenDetailsData {
}
}
fn usage_cost_details(
pub(crate) fn usage_cost_details(
record: &UsageRecord,
upstream_model: Option<&str>,
tokens: &UsageRecordTokenDetailsData,
@@ -607,7 +600,6 @@ fn usage_cost_details(
output_cost: cost.output_cost,
cache_read_cost: cost.cache_read_cost,
total_cost: cost.total_cost,
billed_cost: cost.billed_cost,
original_cost: cost.original_cost,
input_price_per_mtoken: cost.input_price_per_mtoken,
output_price_per_mtoken: cost.output_price_per_mtoken,
@@ -615,12 +607,11 @@ fn usage_cost_details(
service_tier: cost.service_tier,
service_tier_display: cost.service_tier_display,
multiplier: cost.multiplier,
input_cost_display: format_precise_cost(cost.input_cost),
output_cost_display: format_precise_cost(cost.output_cost),
cache_read_cost_display: format_precise_cost(cost.cache_read_cost),
total_cost_display: format_precise_cost(cost.total_cost),
billed_cost_display: format_precise_cost(cost.billed_cost),
original_cost_display: format_precise_cost(cost.original_cost),
input_cost_display: format_cost(cost.input_cost),
output_cost_display: format_cost(cost.output_cost),
cache_read_cost_display: format_cost(cost.cache_read_cost),
total_cost_display: format_cost(cost.total_cost),
original_cost_display: format_cost(cost.original_cost),
input_price_display: format_token_price(cost.input_price_per_mtoken),
output_price_display: format_token_price(cost.output_price_per_mtoken),
cache_read_price_display: format_token_price(cost.cache_read_price_per_mtoken),
@@ -682,9 +673,9 @@ impl UsageRecordBreakdownData {
total_tokens_value: item.total_tokens,
total_tokens_total: format_tokens(total_tokens_total),
total_tokens_total_value: total_tokens_total,
cost: format_precise_cost(item.cost),
actual_cost: format_precise_cost(item.actual_cost),
account_cost: format_precise_cost(item.account_cost),
cost: format_cost(item.cost),
actual_cost: format_cost(item.actual_cost),
account_cost: format_cost(item.account_cost),
average_latency_ms: format_duration_ms_f64(item.average_latency_ms),
}
}
@@ -726,8 +717,8 @@ impl From<UsageRecordTrendPoint> for UsageRecordTrendPointData {
total_tokens_value: point.total_tokens,
cache_hit_rate: format_percent(cache_hit_rate_value),
cache_hit_rate_value,
cost: format_precise_cost(point.cost),
actual_cost: format_precise_cost(point.actual_cost),
cost: format_cost(point.cost),
actual_cost: format_cost(point.actual_cost),
average_latency_ms: format_duration_ms_f64(point.average_latency_ms),
average_latency_ms_value: point.average_latency_ms,
}
@@ -11,8 +11,9 @@ use crate::{
usage_record_model::metadata_service_tier,
usage_record_model::{UsageRecord, UsageRecordLevel},
usage_record_store::{
SqliteUsageRecordStore, UsageRecordBreakdown, UsageRecordEndpointSource,
UsageRecordFilter, UsageRecordModelSource, UsageRecordSummary, UsageRecordTrendPoint,
SqliteUsageRecordStore, UsageRecordAccountUsage, UsageRecordBreakdown,
UsageRecordEndpointSource, UsageRecordFilter, UsageRecordModelSource,
UsageRecordSummary, UsageRecordTrendPoint,
},
},
infra::json::{NumberedPage, Page},
@@ -92,7 +93,6 @@ pub(crate) struct AdminUsageRecordCostDetails {
pub output_cost: f64,
pub cache_read_cost: f64,
pub total_cost: f64,
pub billed_cost: f64,
pub original_cost: f64,
pub input_price_per_mtoken: f64,
pub output_price_per_mtoken: f64,
@@ -117,9 +117,9 @@ pub enum AdminUsageRecordError {
/// 写入失败。
#[error("failed to append usage record")]
Append,
/// 裁剪失败。
#[error("failed to trim usage records")]
Trim,
/// 保留期清理失败。
#[error("failed to trim expired usage records")]
Retention,
/// 账号关联信息读取失败。
#[error("failed to load usage record accounts")]
Accounts,
@@ -135,23 +135,16 @@ pub struct AdminUsageRecordService {
#[derive(Debug, Clone, Copy)]
struct AdminUsageRecordSettings {
enabled: bool,
capacity: u32,
capture_body: bool,
}
impl AdminUsageRecordService {
/// 构造管理端使用记录服务。
pub fn new(
store: SqliteUsageRecordStore,
enabled: bool,
capacity: u32,
capture_body: bool,
) -> Self {
pub fn new(store: SqliteUsageRecordStore, enabled: bool, capture_body: bool) -> Self {
Self {
store,
settings: Arc::new(tokio::sync::RwLock::new(AdminUsageRecordSettings {
enabled,
capacity,
capture_body,
})),
}
@@ -213,6 +206,18 @@ impl AdminUsageRecordService {
.map_err(|_| AdminUsageRecordError::List)
}
/// 按账号聚合使用记录。
pub async fn account_usage(
&self,
filter: AdminUsageRecordFilter,
limit: u32,
) -> Result<Vec<UsageRecordAccountUsage>, AdminUsageRecordError> {
self.store
.account_usage(filter.into(), limit)
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 按模型来源聚合使用记录分布。
pub async fn model_distribution(
&self,
@@ -237,19 +242,8 @@ impl AdminUsageRecordService {
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合 Token 趋势。
pub async fn token_trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
self.store
.trend(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合延迟趋势。
pub async fn latency_trend(
/// 聚合使用记录趋势。
pub async fn trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
@@ -271,7 +265,19 @@ impl AdminUsageRecordService {
/// 记录使用记录。
pub async fn record(&self, event: UsageRecord) -> Result<(), AdminUsageRecordError> {
let settings = *self.settings.read().await;
if !settings.enabled && event.level != UsageRecordLevel::Error {
if !settings.enabled {
return Ok(());
}
if !is_usage_log_event(&event) {
tracing::warn!(
usage_record_id = %event.id,
request_id = event.request_id.as_deref().unwrap_or(""),
account_id = event.account_id.as_deref().unwrap_or(""),
model = event.model.as_deref().unwrap_or(""),
level = ?event.level,
status_code = ?event.status_code,
"skipped non-usage event in usage record service"
);
return Ok(());
}
self.append_with_settings(event, settings).await
@@ -288,13 +294,28 @@ impl AdminUsageRecordService {
.await
.map_err(|_| AdminUsageRecordError::Append)?;
self.store
.trim_to_capacity(settings.capacity)
.trim_to_retention(Utc::now())
.await
.map_err(|_| AdminUsageRecordError::Trim)?;
.map_err(|_| AdminUsageRecordError::Retention)?;
Ok(())
}
}
fn is_usage_log_event(event: &UsageRecord) -> bool {
event.level != UsageRecordLevel::Error
&& event.status_code.is_none_or(|status| status < 400)
&& event
.account_id
.as_deref()
.map(str::trim)
.is_some_and(|account_id| !account_id.is_empty())
&& event
.model
.as_deref()
.map(str::trim)
.is_some_and(|model| !model.is_empty())
}
fn apply_capture_body_policy(event: &mut UsageRecord, capture_body: bool) {
if capture_body {
return;
@@ -343,7 +364,6 @@ pub(crate) fn usage_record_cost_details(
output_cost: breakdown.output_cost,
cache_read_cost: breakdown.cache_read_cost,
total_cost: breakdown.total_cost,
billed_cost: breakdown.total_cost,
original_cost,
input_price_per_mtoken: breakdown.input_price_per_mtoken,
output_price_per_mtoken: breakdown.output_price_per_mtoken,
@@ -2,7 +2,7 @@
use std::collections::{BTreeMap, HashMap};
use chrono::{DateTime, Utc};
use chrono::{DateTime, Duration, Utc};
use serde_json::Value;
use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool};
use thiserror::Error;
@@ -19,8 +19,8 @@ use crate::infra::{
time::{china_quarter_hour_start, parse_rfc3339_utc as parse_rfc3339},
};
/// 默认保留最新使用记录数量。
pub const DEFAULT_USAGE_RECORD_CAPACITY: u32 = 2_000;
/// 使用记录保留天数。
pub const USAGE_RECORD_RETENTION_DAYS: i64 = 30;
/// 默认不保存请求/响应体,避免把用户内容写入诊断日志。
pub const DEFAULT_USAGE_RECORD_CAPTURE_BODY: bool = false;
@@ -100,14 +100,19 @@ impl SqliteUsageRecordStore {
Self { pool }
}
/// 暴露底层连接池,供运行时组合层复用。
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
/// 写入使用记录。
pub async fn append(&self, event: &UsageRecord) -> SqliteUsageRecordStoreResult<()> {
append_event(&self.pool, event).await
}
/// 保留最新的指定数量使用记录。
pub async fn trim_to_capacity(&self, capacity: u32) -> SqliteUsageRecordStoreResult<u64> {
trim_to_capacity(&self.pool, capacity).await
/// 按保留期清理使用记录。
pub async fn trim_to_retention(&self, now: DateTime<Utc>) -> SqliteUsageRecordStoreResult<u64> {
trim_to_retention(&self.pool, now).await
}
/// 分页查询使用记录。
@@ -198,6 +203,15 @@ impl SqliteUsageRecordStore {
usage_summary(&self.pool, &filter).await
}
/// 按账号聚合使用记录。
pub async fn account_usage(
&self,
filter: UsageRecordFilter,
limit: u32,
) -> SqliteUsageRecordStoreResult<Vec<UsageRecordAccountUsage>> {
usage_account_usage(&self.pool, &filter, limit).await
}
/// 按模型来源聚合使用记录分布。
pub async fn model_distribution(
&self,
@@ -277,7 +291,6 @@ fn usage_record_summary_fields(event: &UsageRecord) -> UsageRecordSummaryFields
"upstreamRequestId",
"upstream_request_id",
"openaiRequestId",
"cfRay",
],
)
}),
@@ -417,17 +430,15 @@ fn metadata_first_token_latency(metadata: &Value) -> Option<i64> {
})
}
async fn trim_to_capacity(pool: &SqlitePool, capacity: u32) -> SqliteUsageRecordStoreResult<u64> {
let result = sqlx::query(
"delete from usage_records where id in (
select id from usage_records
order by created_at desc, id desc
limit -1 offset ?
)",
)
.bind(i64::from(capacity))
.execute(pool)
.await?;
async fn trim_to_retention(
pool: &SqlitePool,
now: DateTime<Utc>,
) -> SqliteUsageRecordStoreResult<u64> {
let cutoff = now - Duration::days(USAGE_RECORD_RETENTION_DAYS);
let result = sqlx::query("delete from usage_records where created_at < ?")
.bind(cutoff.to_rfc3339())
.execute(pool)
.await?;
Ok(result.rows_affected())
}
@@ -477,6 +488,45 @@ async fn usage_summary(
})
}
async fn usage_account_usage(
pool: &SqlitePool,
filter: &UsageRecordFilter,
limit: u32,
) -> SqliteUsageRecordStoreResult<Vec<UsageRecordAccountUsage>> {
let mut builder = QueryBuilder::<Sqlite>::new(
"select
account_id,
sum(coalesce(cast(json_extract(metadata_json, '$.usage.inputTokens') as integer), cast(json_extract(metadata_json, '$.inputTokens') as integer), 0)) as input_tokens,
sum(coalesce(cast(json_extract(metadata_json, '$.usage.outputTokens') as integer), cast(json_extract(metadata_json, '$.outputTokens') as integer), 0)) as output_tokens,
max(created_at) as last_used_at
from usage_records",
);
push_filter(&mut builder, filter, None)?;
builder.push(" and account_id is not null and trim(account_id) <> ''");
builder.push(" group by account_id order by last_used_at desc, account_id asc limit ");
builder.push_bind(i64::from(limit.clamp(1, 50)));
builder
.build()
.fetch_all(pool)
.await?
.into_iter()
.map(|row| {
let input_tokens =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("input_tokens"));
let output_tokens =
optional_nonnegative_i64_to_u64(row.get::<Option<i64>, _>("output_tokens"));
Ok(UsageRecordAccountUsage {
account_id: row.get("account_id"),
input_tokens,
output_tokens,
total_tokens: input_tokens + output_tokens,
last_used_at: parse_rfc3339(&row.get::<String, _>("last_used_at"))?,
})
})
.collect()
}
async fn usage_record_account_email_map(
pool: &SqlitePool,
items: &[UsageRecord],
@@ -978,6 +1028,21 @@ pub struct UsageRecordSummary {
pub average_latency_ms: Option<f64>,
}
/// 账号使用记录聚合。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageRecordAccountUsage {
/// 账号 ID。
pub account_id: String,
/// 输入 Token。
pub input_tokens: u64,
/// 输出 Token。
pub output_tokens: u64,
/// 总 Token。
pub total_tokens: u64,
/// 最近使用时间。
pub last_used_at: DateTime<Utc>,
}
/// 模型分布来源。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UsageRecordModelSource {
+1 -4
View File
@@ -16,9 +16,8 @@ use crate::{
upstream::models::config::ModelConfig,
};
const ALLOWED_SETTINGS_KEYS: [&str; 7] = [
const ALLOWED_SETTINGS_KEYS: [&str; 6] = [
"modelAliases",
"modelAccountRoutes",
"refreshMarginSeconds",
"refreshConcurrency",
"maxConcurrentPerAccount",
@@ -30,7 +29,6 @@ const ALLOWED_SETTINGS_KEYS: [&str; 7] = [
#[serde(rename_all = "camelCase")]
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,
@@ -42,7 +40,6 @@ impl AdminSettingsData {
fn from_config(config: &AppConfig) -> Self {
Self {
model_aliases: config.model_aliases.clone(),
model_account_routes: config.model_account_routes.clone(),
refresh_margin_seconds: config.auth.refresh_margin_seconds,
refresh_concurrency: config.auth.refresh_concurrency,
max_concurrent_per_account: config.auth.max_concurrent_per_account,
+40 -9
View File
@@ -110,6 +110,8 @@ struct SystemUpdateEvent {
level: UpdateLogLevel,
step: Option<String>,
message: String,
#[serde(skip_serializing_if = "is_false")]
terminal: bool,
at: String,
}
@@ -157,18 +159,23 @@ pub(crate) async fn update_event_stream(
let receiver = update_event_sender().subscribe();
let shutdown = shutdown_subscription();
let stream = futures::stream::unfold(
(receiver, shutdown),
|(mut receiver, mut shutdown)| async move {
(receiver, shutdown, false),
|(mut receiver, mut shutdown, close_after_send)| async move {
if close_after_send {
return None;
}
loop {
tokio::select! {
_ = shutdown.recv() => return None,
received = receiver.recv() => {
match received {
Ok(message) => {
let close_after_send = message.terminal;
let id = message.id.clone();
let data = serde_json::to_string(&message).unwrap_or_else(|_| "{}".to_string());
let event = Event::default().event("update").id(id).data(data);
return Some((Ok(event), (receiver, shutdown)));
return Some((Ok(event), (receiver, shutdown, close_after_send)));
}
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(broadcast::error::RecvError::Closed) => return None,
@@ -192,7 +199,7 @@ pub(crate) async fn perform_update(
.map_err(|_| AdminError::conflict("System update already running"))?;
let config = SystemUpdateConfig::from_env();
if let Some(reason) = config.update_support_error() {
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Error,
None,
Some("preflight"),
@@ -217,7 +224,7 @@ pub(crate) async fn perform_update(
)
.await
.map_err(|error| {
emit_update_event(UpdateLogLevel::Error, None, Some("release"), error.clone());
emit_terminal_update_event(UpdateLogLevel::Error, None, Some("release"), error.clone());
AdminError::bad_gateway(error)
})?;
let info = update_info_from_release(&config, release.clone());
@@ -226,7 +233,7 @@ pub(crate) async fn perform_update(
"远端最新版本已变更为 v{},请重新检查并确认",
info.latest_version
);
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Warning,
None,
Some("release"),
@@ -235,7 +242,7 @@ pub(crate) async fn perform_update(
return Err(AdminError::conflict(message));
}
if !info.has_update {
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Warning,
None,
Some("release"),
@@ -261,14 +268,14 @@ pub(crate) async fn perform_update(
);
let result = perform_release_update(&config, &release, &target_version, &operation_id).await;
if let Err(error) = &result {
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Error,
Some(&operation_id),
Some("failed"),
error.to_string(),
);
} else {
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Success,
Some(&operation_id),
Some("done"),
@@ -732,6 +739,25 @@ fn emit_update_event(
operation_id: Option<&str>,
step: Option<&str>,
message: impl Into<String>,
) {
emit_update_event_with_terminal(level, operation_id, step, message, false);
}
fn emit_terminal_update_event(
level: UpdateLogLevel,
operation_id: Option<&str>,
step: Option<&str>,
message: impl Into<String>,
) {
emit_update_event_with_terminal(level, operation_id, step, message, true);
}
fn emit_update_event_with_terminal(
level: UpdateLogLevel,
operation_id: Option<&str>,
step: Option<&str>,
message: impl Into<String>,
terminal: bool,
) {
let now = Utc::now();
let event = SystemUpdateEvent {
@@ -744,11 +770,16 @@ fn emit_update_event(
level,
step: step.map(ToString::to_string),
message: message.into(),
terminal,
at: now.to_rfc3339(),
};
let _ = update_event_sender().send(event);
}
fn is_false(value: &bool) -> bool {
!*value
}
fn format_bytes(bytes: u64) -> String {
const KIB: u64 = 1024;
const MIB: u64 = KIB * 1024;
+13 -7
View File
@@ -24,9 +24,6 @@ pub struct AppConfig {
/// 运行时模型别名,由数据库设置加载,配置文件不承载该字段。
#[serde(skip)]
pub model_aliases: BTreeMap<String, String>,
/// 运行时模型账号路由,由 `model_account_routes` 表加载,配置文件不承载该字段。
#[serde(skip)]
pub model_account_routes: BTreeMap<String, Vec<String>>,
/// 认证与刷新配置,由固定默认值和数据库运行时设置承载,配置文件不承载该字段。
#[serde(skip)]
pub auth: AuthConfig,
@@ -60,9 +57,6 @@ pub struct ServerConfig {
pub host: String,
/// 监听端口。
pub port: u16,
/// 可信反向代理 IP 或 CIDR。只有这些 peer 的转发 IP 头会被采信。
#[serde(default)]
pub trusted_proxies: Vec<String>,
}
/// 上游 API 地址配置。
@@ -103,7 +97,7 @@ impl Default for AuthConfig {
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
rotation_strategy: "least_used".to_string(),
rotation_strategy: "smart".to_string(),
tier_priority: Vec::new(),
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
@@ -164,6 +158,17 @@ pub struct WebSocketPoolConfig {
pub max_age_ms: u64,
/// 单账号最大连接数。
pub max_per_account: usize,
/// 首个内容帧(首 token)到达前的绝对超时;`0` 表示禁用。
///
/// 从发出 `response.create` 起算,覆盖建连/发送/上游排队/思考的全程,直到收到
/// 首个真实内容帧(`response.created`/`response.in_progress` 等纯生命周期帧不计)。
/// 超时判定为连接落到病态上游后端,丢弃并换新连接重试。
#[serde(default = "default_ws_first_token_timeout_ms")]
pub first_token_timeout_ms: u64,
}
fn default_ws_first_token_timeout_ms() -> u64 {
20_000
}
impl Default for WebSocketPoolConfig {
@@ -172,6 +177,7 @@ impl Default for WebSocketPoolConfig {
enabled: true,
max_age_ms: 3_300_000,
max_per_account: 8,
first_token_timeout_ms: default_ws_first_token_timeout_ms(),
}
}
}
+10 -113
View File
@@ -18,7 +18,7 @@ use crate::upstream::accounts::{
token_refresh::{RefreshPolicy, RuntimeRefreshPolicy},
};
const ROTATION_STRATEGIES: [&str; 3] = ["least_used", "round_robin", "sticky"];
const ROTATION_STRATEGIES: [&str; 4] = ["smart", "quota_reset_priority", "round_robin", "sticky"];
const RUNTIME_SETTINGS_ID: i64 = 1;
/// 管理端可变设置。
@@ -26,8 +26,6 @@ const RUNTIME_SETTINGS_ID: i64 = 1;
pub struct AdminSettings {
/// 模型别名映射。
pub model_aliases: BTreeMap<String, String>,
/// 模型到账号 ID 的显式路由。
pub model_account_routes: BTreeMap<String, Vec<String>>,
/// 访问令牌过期前多少秒开始刷新。
pub refresh_margin_seconds: u64,
/// 访问令牌刷新并发数。
@@ -44,12 +42,11 @@ impl Default for AdminSettings {
fn default() -> Self {
Self {
model_aliases: BTreeMap::new(),
model_account_routes: BTreeMap::new(),
refresh_margin_seconds: 300,
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
rotation_strategy: "least_used".to_string(),
rotation_strategy: "smart".to_string(),
}
}
}
@@ -60,8 +57,6 @@ impl Default for AdminSettings {
pub struct AdminSettingsPatch {
/// 模型别名映射。
pub model_aliases: Option<BTreeMap<String, String>>,
/// 模型到账号 ID 的显式路由。
pub model_account_routes: Option<BTreeMap<String, Vec<String>>>,
/// 访问令牌过期前多少秒开始刷新。
pub refresh_margin_seconds: Option<u64>,
/// 访问令牌刷新并发数。
@@ -87,9 +82,6 @@ impl SettingsService {
if let Some(model_aliases) = patch.model_aliases {
current.model_aliases = validate_model_aliases(model_aliases)?;
}
if let Some(model_account_routes) = patch.model_account_routes {
current.model_account_routes = validate_model_account_routes(model_account_routes)?;
}
if let Some(refresh_margin_seconds) = patch.refresh_margin_seconds {
current.refresh_margin_seconds =
positive_u64("refreshMarginSeconds", refresh_margin_seconds)?;
@@ -151,34 +143,6 @@ fn validate_model_aliases(
Ok(normalized)
}
fn validate_model_account_routes(
routes: BTreeMap<String, Vec<String>>,
) -> Result<BTreeMap<String, Vec<String>>, SettingsServiceError> {
let mut normalized = BTreeMap::new();
for (model, account_ids) in routes {
let model = non_empty_string("modelAccountRoutes", &model)?;
let mut normalized_account_ids = Vec::new();
for account_id in account_ids {
let account_id = non_empty_string("modelAccountRoutes", &account_id)?;
if normalized_account_ids.contains(&account_id) {
return Err(invalid_field(
"modelAccountRoutes",
format!("duplicate account id `{account_id}`"),
));
}
normalized_account_ids.push(account_id);
}
if normalized_account_ids.is_empty() {
return Err(invalid_field(
"modelAccountRoutes",
"account list must not be empty",
));
}
normalized.insert(model, normalized_account_ids);
}
Ok(normalized)
}
fn validate_rotation_strategy(strategy: &str) -> Result<String, SettingsServiceError> {
let strategy = non_empty_string("rotationStrategy", strategy)?;
if ROTATION_STRATEGIES.contains(&strategy.as_str()) {
@@ -186,7 +150,7 @@ fn validate_rotation_strategy(strategy: &str) -> Result<String, SettingsServiceE
} else {
Err(invalid_field(
"rotationStrategy",
"must be one of least_used, round_robin, sticky",
"must be one of smart, quota_reset_priority, round_robin, sticky",
))
}
}
@@ -394,7 +358,6 @@ pub enum RuntimeSettingsError {
fn admin_settings_from_config(config: &AppConfig) -> AdminSettings {
AdminSettings {
model_aliases: config.model_aliases.clone(),
model_account_routes: config.model_account_routes.clone(),
refresh_margin_seconds: config.auth.refresh_margin_seconds,
refresh_concurrency: config.auth.refresh_concurrency,
max_concurrent_per_account: config.auth.max_concurrent_per_account,
@@ -405,7 +368,6 @@ fn admin_settings_from_config(config: &AppConfig) -> AdminSettings {
fn apply_admin_settings_to_config(config: &mut AppConfig, settings: AdminSettings) {
config.model_aliases = settings.model_aliases;
config.model_account_routes = settings.model_account_routes;
config.auth.refresh_margin_seconds = settings.refresh_margin_seconds;
config.auth.refresh_concurrency = settings.refresh_concurrency;
config.auth.max_concurrent_per_account = settings.max_concurrent_per_account;
@@ -418,8 +380,7 @@ async fn insert_runtime_settings_if_missing(
settings: &AdminSettings,
) -> Result<(), RuntimeSettingsError> {
let model_aliases_json = serde_json::to_string(&settings.model_aliases)?;
let mut tx = pool.begin().await?;
let result = sqlx::query(
sqlx::query(
r"
insert or ignore into runtime_settings (
id,
@@ -449,12 +410,8 @@ insert or ignore into runtime_settings (
)
.bind(&settings.rotation_strategy)
.bind(Utc::now().to_rfc3339())
.execute(&mut *tx)
.execute(pool)
.await?;
if result.rows_affected() > 0 {
replace_model_account_routes(&mut tx, &settings.model_account_routes).await?;
}
tx.commit().await?;
Ok(())
}
@@ -463,7 +420,6 @@ async fn save_runtime_settings(
settings: &AdminSettings,
) -> Result<(), RuntimeSettingsError> {
let model_aliases_json = serde_json::to_string(&settings.model_aliases)?;
let mut tx = pool.begin().await?;
sqlx::query(
r"
insert into runtime_settings (
@@ -502,10 +458,8 @@ on conflict(id) do update set
)
.bind(&settings.rotation_strategy)
.bind(Utc::now().to_rfc3339())
.execute(&mut *tx)
.execute(pool)
.await?;
replace_model_account_routes(&mut tx, &settings.model_account_routes).await?;
tx.commit().await?;
Ok(())
}
@@ -525,9 +479,7 @@ where id = ?",
.bind(RUNTIME_SETTINGS_ID)
.fetch_one(pool)
.await?;
let mut settings = runtime_settings_from_row(&row)?;
settings.model_account_routes = load_model_account_routes(pool).await?;
Ok(settings)
runtime_settings_from_row(&row)
}
async fn admin_api_key_status(
@@ -566,7 +518,6 @@ fn runtime_settings_from_row(row: &SqliteRow) -> Result<AdminSettings, RuntimeSe
let model_aliases_json: String = row.get("model_aliases_json");
Ok(AdminSettings {
model_aliases: serde_json::from_str(&model_aliases_json)?,
model_account_routes: BTreeMap::new(),
refresh_margin_seconds: positive_i64_to_u64(
"refreshMarginSeconds",
row.get("refresh_margin_seconds"),
@@ -587,61 +538,6 @@ fn runtime_settings_from_row(row: &SqliteRow) -> Result<AdminSettings, RuntimeSe
})
}
async fn replace_model_account_routes(
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
routes: &BTreeMap<String, Vec<String>>,
) -> Result<(), RuntimeSettingsError> {
sqlx::query("delete from model_account_routes")
.execute(&mut **tx)
.await?;
let now = Utc::now().to_rfc3339();
for (model, account_ids) in routes {
for (priority, account_id) in account_ids.iter().enumerate() {
sqlx::query(
r"
insert into model_account_routes (
model,
account_id,
priority,
enabled,
created_at,
updated_at
) values (?, ?, ?, 1, ?, ?)",
)
.bind(model)
.bind(account_id)
.bind(
i64::try_from(priority)
.map_err(|_| stored_field_error("modelAccountRoutes", "out of range"))?,
)
.bind(&now)
.bind(&now)
.execute(&mut **tx)
.await?;
}
}
Ok(())
}
async fn load_model_account_routes(
pool: &SqlitePool,
) -> Result<BTreeMap<String, Vec<String>>, RuntimeSettingsError> {
let rows: Vec<(String, String)> = sqlx::query_as(
r"
select model, account_id
from model_account_routes
where enabled = 1
order by model asc, priority asc, account_id asc",
)
.fetch_all(pool)
.await?;
let mut routes = BTreeMap::<String, Vec<String>>::new();
for (model, account_id) in rows {
routes.entry(model).or_default().push(account_id);
}
Ok(routes)
}
/// 从当前运行配置生成账号池调度参数。
pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions {
AccountPoolOptions {
@@ -649,16 +545,17 @@ pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOption
rotation_strategy: rotation_strategy_from_config(&config.auth.rotation_strategy),
skip_quota_limited: config.quota.skip_exhausted,
tier_priority: config.auth.tier_priority.clone(),
model_account_routes: config.model_account_routes.clone(),
..AccountPoolOptions::default()
}
}
fn rotation_strategy_from_config(strategy: &str) -> RotationStrategy {
match strategy {
"smart" => RotationStrategy::Smart,
"quota_reset_priority" => RotationStrategy::QuotaResetPriority,
"round_robin" => RotationStrategy::RoundRobin,
"sticky" => RotationStrategy::Sticky,
_ => RotationStrategy::LeastUsed,
_ => RotationStrategy::Smart,
}
}
+77 -105
View File
@@ -1,9 +1,10 @@
//! 请求上下文中间件。
use std::net::{IpAddr, SocketAddr};
use std::sync::LazyLock;
use axum::{
extract::{connect_info::ConnectInfo, Request, State},
extract::{connect_info::ConnectInfo, Request},
http::{HeaderMap, HeaderValue},
middleware::Next,
response::Response,
@@ -17,39 +18,24 @@ const MAX_REQUEST_ID_LEN: usize = 128;
const CF_CONNECTING_IP_HEADER: &str = "cf-connecting-ip";
const REAL_IP_HEADER: &str = "x-real-ip";
const X_FORWARDED_FOR_HEADER: &str = "x-forwarded-for";
const UNTRUSTED_CLIENT_IP_HEADERS: &[&str] = &[
CF_CONNECTING_IP_HEADER,
REAL_IP_HEADER,
X_FORWARDED_FOR_HEADER,
];
/// 可信反向代理配置。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct TrustedProxyConfig {
networks: Vec<IpNet>,
}
/// 私有 / 回环 / 唯一本地地址网段,用于从 `X-Forwarded-For` 中跳过内网跳板 IP。
static PRIVATE_NETWORKS: LazyLock<Vec<IpNet>> = LazyLock::new(|| {
[
"10.0.0.0/8",
"172.16.0.0/12",
"192.168.0.0/16",
"127.0.0.0/8",
"::1/128",
"fc00::/7",
]
.iter()
.filter_map(|entry| entry.parse().ok())
.collect()
});
impl TrustedProxyConfig {
/// 从 IP 或 CIDR 字符串构造可信代理配置。
pub fn from_entries(entries: &[String]) -> Result<Self, TrustedProxyConfigError> {
let networks = entries
.iter()
.map(|entry| parse_trusted_proxy_entry(entry))
.collect::<Result<Vec<_>, _>>()?;
Ok(Self { networks })
}
fn is_trusted(&self, ip: IpAddr) -> bool {
self.networks.iter().any(|network| network.contains(&ip))
}
}
/// 可信代理配置错误。
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum TrustedProxyConfigError {
/// 配置项不是 IP 地址或 CIDR。
#[error("invalid trusted proxy entry `{entry}`; expected an IP address or CIDR network")]
InvalidEntry { entry: String },
fn is_private_ip(ip: IpAddr) -> bool {
PRIVATE_NETWORKS.iter().any(|network| network.contains(&ip))
}
/// 请求 ID。
@@ -77,6 +63,17 @@ impl RequestId {
}
}
/// 由连接来源和转发头解析出的客户端 IP。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientIp(String);
impl ClientIp {
/// 返回 IP 字符串。
pub fn as_str(&self) -> &str {
&self.0
}
}
fn valid_request_id(value: &str) -> bool {
!value.is_empty()
&& value.len() <= MAX_REQUEST_ID_LEN
@@ -87,17 +84,7 @@ fn valid_request_id(value: &str) -> bool {
/// 为请求附加请求 ID,并在响应头中回写。
pub async fn attach_request_id(mut request: Request, next: Next) -> Response {
attach_request_context(&TrustedProxyConfig::default(), &mut request);
run_with_request_id(request, next).await
}
/// 为请求附加请求上下文,并按可信代理配置解析真实客户端 IP。
pub async fn attach_request_id_with_proxy_config(
State(trusted_proxies): State<TrustedProxyConfig>,
mut request: Request,
next: Next,
) -> Response {
attach_request_context(&trusted_proxies, &mut request);
attach_request_context(&mut request);
run_with_request_id(request, next).await
}
@@ -116,84 +103,69 @@ async fn run_with_request_id(mut request: Request, next: Next) -> Response {
response
}
fn attach_request_context(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
attach_real_ip_from_connection(trusted_proxies, request);
fn attach_request_context(request: &mut Request) {
attach_client_ip_from_connection(request);
}
fn attach_real_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
fn attach_client_ip_from_connection(request: &mut Request) {
let peer_ip = request
.extensions()
.get::<ConnectInfo<SocketAddr>>()
.map(|ConnectInfo(addr)| addr.ip());
let client_ip =
peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers(), trusted_proxies));
for header in UNTRUSTED_CLIENT_IP_HEADERS {
request.headers_mut().remove(*header);
}
let client_ip = peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers()));
let Some(client_ip) = client_ip else {
return;
};
if let Ok(value) = HeaderValue::from_str(&client_ip.to_string()) {
request.headers_mut().insert(REAL_IP_HEADER, value);
request.extensions_mut().insert(ClientIp(client_ip));
}
fn client_ip_from_peer(peer_ip: IpAddr, headers: &HeaderMap) -> String {
auto_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string())
}
/// 解析转发头中的真实客户端 IP。
///
/// `CF-Connecting-IP` → `X-Real-IP` → `X-Forwarded-For` 首个公网 IP。
fn auto_forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
if let Some(cf_ip) = header_string(headers, CF_CONNECTING_IP_HEADER) {
return Some(cf_ip);
}
}
fn client_ip_from_peer(
peer_ip: IpAddr,
headers: &HeaderMap,
trusted_proxies: &TrustedProxyConfig,
) -> IpAddr {
if trusted_proxies.is_trusted(peer_ip) {
trusted_forwarded_client_ip(headers).unwrap_or(peer_ip)
} else {
peer_ip
if let Some(real_ip) = header_string(headers, REAL_IP_HEADER) {
return Some(real_ip);
}
forwarded_for_public_ip(headers).or_else(|| first_forwarded_for_value(headers))
}
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<IpAddr> {
header_ip(headers, CF_CONNECTING_IP_HEADER)
.or_else(|| first_forwarded_for_ip(headers))
.or_else(|| header_ip(headers, REAL_IP_HEADER))
/// 返回 `X-Forwarded-For` 链中第一个公网 IP(跳过内网跳板)。
fn forwarded_for_public_ip(headers: &HeaderMap) -> Option<String> {
let value = header_string(headers, X_FORWARDED_FOR_HEADER)?;
value
.split(',')
.map(str::trim)
.filter(|entry| !entry.is_empty())
.find(|entry| match entry.parse::<IpAddr>() {
Ok(ip) => !is_private_ip(ip),
Err(_) => false,
})
.map(ToString::to_string)
}
fn header_ip(headers: &HeaderMap, name: &str) -> Option<IpAddr> {
parse_ip(headers.get(name)?.to_str().ok()?)
fn header_string(headers: &HeaderMap, name: &str) -> Option<String> {
headers
.get(name)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn first_forwarded_for_ip(headers: &HeaderMap) -> Option<IpAddr> {
let value = headers.get(X_FORWARDED_FOR_HEADER)?.to_str().ok()?;
parse_ip(value.split(',').next()?)
}
fn parse_ip(value: &str) -> Option<IpAddr> {
value.trim().parse().ok()
}
fn parse_trusted_proxy_entry(entry: &str) -> Result<IpNet, TrustedProxyConfigError> {
let entry = entry.trim();
if entry.is_empty() {
return Err(invalid_trusted_proxy_entry(entry));
}
if let Ok(network) = entry.parse::<IpNet>() {
return Ok(network);
}
let ip = entry
.parse::<IpAddr>()
.map_err(|_| invalid_trusted_proxy_entry(entry))?;
IpNet::new(
ip,
match ip {
IpAddr::V4(_) => 32,
IpAddr::V6(_) => 128,
},
)
.map_err(|_| invalid_trusted_proxy_entry(entry))
}
fn invalid_trusted_proxy_entry(entry: &str) -> TrustedProxyConfigError {
TrustedProxyConfigError::InvalidEntry {
entry: entry.to_string(),
}
fn first_forwarded_for_value(headers: &HeaderMap) -> Option<String> {
let value = header_string(headers, X_FORWARDED_FOR_HEADER)?;
value
.split(',')
.next()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
+3 -22
View File
@@ -5,10 +5,7 @@ use std::path::Path;
use axum::{extract::State, http::StatusCode, middleware, routing::get, Router};
use crate::admin;
use crate::http::middleware::{
request_id::{attach_request_id_with_proxy_config, TrustedProxyConfig},
trace::http_trace_layer,
};
use crate::http::middleware::{request_id::attach_request_id, trace::http_trace_layer};
use crate::proxy;
use crate::runtime::state::AppState;
use crate::web;
@@ -18,34 +15,18 @@ pub const DEFAULT_ASSET_DIST_DIR: &str = "web/dist";
/// 构造整个 HTTP 服务路由。
pub fn router() -> Router<AppState> {
router_with_trusted_proxies(TrustedProxyConfig::default())
}
/// 使用指定可信代理配置构造整个 HTTP 服务路由。
pub fn router_with_trusted_proxies(trusted_proxies: TrustedProxyConfig) -> Router<AppState> {
router_with_assets_and_trusted_proxies(DEFAULT_ASSET_DIST_DIR, trusted_proxies)
router_with_assets(DEFAULT_ASSET_DIST_DIR)
}
/// 使用指定前端构建产物目录构造整个 HTTP 服务路由。
pub fn router_with_assets(dist_dir: impl AsRef<Path>) -> Router<AppState> {
router_with_assets_and_trusted_proxies(dist_dir, TrustedProxyConfig::default())
}
/// 使用指定前端构建产物目录和可信代理配置构造整个 HTTP 服务路由。
pub fn router_with_assets_and_trusted_proxies(
dist_dir: impl AsRef<Path>,
trusted_proxies: TrustedProxyConfig,
) -> Router<AppState> {
Router::new()
.route("/healthz", get(healthz))
.merge(proxy::router::router())
.merge(admin::router::router())
.fallback_service(web::assets::spa_router(dist_dir))
.layer(http_trace_layer())
.layer(middleware::from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
))
.layer(middleware::from_fn(attach_request_id))
}
async fn healthz(State(state): State<AppState>) -> StatusCode {
+18 -6
View File
@@ -4,12 +4,24 @@ use std::{fs, path::Path, str::FromStr, time::Duration};
use sqlx::{sqlite::SqliteConnectOptions, SqlitePool};
const CURRENT_SCHEMA_VERSION: i64 = 1;
const MIGRATIONS: &[Migration] = &[Migration {
version: 1,
name: "initial",
sql: include_str!("migrations/0001_initial.sql"),
}];
const CURRENT_SCHEMA_VERSION: i64 = 3;
const MIGRATIONS: &[Migration] = &[
Migration {
version: 1,
name: "initial",
sql: include_str!("migrations/0001_initial.sql"),
},
Migration {
version: 2,
name: "rename_rotation_strategies",
sql: include_str!("migrations/0002_rename_rotation_strategies.sql"),
},
Migration {
version: 3,
name: "ops_error_logs",
sql: include_str!("migrations/0003_ops_error_logs.sql"),
},
];
struct Migration {
version: i64,
+11 -10
View File
@@ -63,6 +63,15 @@ pub fn format_percent(value: f64) -> String {
}
}
/// 展示百分比,并去掉无意义的尾随 0。
pub fn format_compact_percent(value: f64) -> String {
if value.is_finite() {
format!("{}%", trim_trailing_zeroes(&format!("{value:.1}")))
} else {
"—".to_string()
}
}
/// 将 0 到 1 的比例展示为百分比。
pub fn format_rate(value: Option<f64>) -> String {
value.map_or_else(|| "—".to_string(), |value| format_percent(value * 100.0))
@@ -86,15 +95,6 @@ pub fn format_cost(value: f64) -> String {
format!("${value:.2}")
}
/// 使用六位小数展示美元成本。
pub fn format_precise_cost(value: f64) -> String {
if value.is_finite() {
format!("${value:.6}")
} else {
"—".to_string()
}
}
/// 展示每百万 token 的美元单价。
pub fn format_token_price(value: f64) -> String {
if value.is_finite() {
@@ -160,7 +160,8 @@ fn format_compact_scaled(value: f64, unit: &str) -> String {
format!("{}{unit}", trim_trailing_zeroes(&rounded))
}
fn trim_trailing_zeroes(value: &str) -> String {
/// 去掉小数字符串中无意义的尾随 0 和小数点。
pub fn trim_trailing_zeroes(value: &str) -> String {
value
.trim_end_matches('0')
.trim_end_matches('.')
+2 -20
View File
@@ -49,7 +49,7 @@ create table runtime_settings (
refresh_concurrency integer not null check (refresh_concurrency > 0),
max_concurrent_per_account integer not null check (max_concurrent_per_account > 0),
request_interval_ms integer not null check (request_interval_ms >= 0),
rotation_strategy text not null check (rotation_strategy in ('least_used', 'round_robin', 'sticky')),
rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')),
admin_api_key text,
updated_at text not null
);
@@ -68,7 +68,7 @@ create table accounts (
refresh_token text,
access_token_expires_at text,
next_refresh_at text,
status text not null check (status in ('active', 'expired', 'quota_exhausted', 'refreshing', 'disabled', 'banned')),
status text not null check (status in ('active', 'expired', 'quota_exhausted', 'disabled', 'banned')),
quota_json text,
quota_fetched_at text,
quota_limit_reached integer not null default 0 check (quota_limit_reached in (0, 1)),
@@ -176,24 +176,6 @@ on usage_time_buckets(bucket_start);
create index idx_usage_time_buckets_model_bucket
on usage_time_buckets(model, bucket_start);
-- ============================================
-- Model Account Routes
-- ============================================
create table model_account_routes (
model text not null,
account_id text not null references accounts(id) on delete cascade,
priority integer not null default 0 check (priority >= 0),
enabled integer not null default 1 check (enabled in (0, 1)),
created_at text not null,
updated_at text not null,
primary key (model, account_id)
);
create index idx_model_account_routes_account
on model_account_routes(account_id);
create index idx_model_account_routes_enabled_model
on model_account_routes(enabled, model, priority asc, account_id);
-- ============================================
-- Account Cookies
-- ============================================
@@ -0,0 +1,40 @@
create table runtime_settings_new (
id integer primary key check (id = 1),
model_aliases_json text not null default '{}',
refresh_margin_seconds integer not null check (refresh_margin_seconds > 0),
refresh_concurrency integer not null check (refresh_concurrency > 0),
max_concurrent_per_account integer not null check (max_concurrent_per_account > 0),
request_interval_ms integer not null check (request_interval_ms >= 0),
rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')),
admin_api_key text,
updated_at text not null
);
insert into runtime_settings_new (
id,
model_aliases_json,
refresh_margin_seconds,
refresh_concurrency,
max_concurrent_per_account,
request_interval_ms,
rotation_strategy,
admin_api_key,
updated_at
)
select
id,
model_aliases_json,
refresh_margin_seconds,
refresh_concurrency,
max_concurrent_per_account,
request_interval_ms,
case rotation_strategy
when 'least_used' then 'quota_reset_priority'
else rotation_strategy
end,
admin_api_key,
updated_at
from runtime_settings;
drop table runtime_settings;
alter table runtime_settings_new rename to runtime_settings;
@@ -0,0 +1,31 @@
create table ops_error_logs (
id text primary key,
request_id text,
kind text not null,
account_id text,
route text,
model text,
status_code integer check (status_code is null or (status_code >= 100 and status_code <= 599)),
client_status_code integer check (client_status_code is null or (client_status_code >= 100 and client_status_code <= 599)),
upstream_status_code integer check (upstream_status_code is null or (upstream_status_code >= 100 and upstream_status_code <= 599)),
transport text,
attempt_index integer check (attempt_index is null or attempt_index >= 0),
failure_class text,
response_id text,
upstream_request_id text,
latency_ms integer check (latency_ms is null or latency_ms >= 0),
message text not null,
metadata_json text not null,
created_at text not null
);
create index idx_ops_error_logs_created_id on ops_error_logs(created_at desc, id desc);
create index idx_ops_error_logs_request_id on ops_error_logs(request_id) where request_id is not null;
create index idx_ops_error_logs_account on ops_error_logs(account_id, created_at desc) where account_id is not null;
create index idx_ops_error_logs_route_created on ops_error_logs(route, created_at desc) where route is not null;
create index idx_ops_error_logs_model_created on ops_error_logs(model, created_at desc) where model is not null;
create index idx_ops_error_logs_status_created on ops_error_logs(status_code, created_at desc) where status_code is not null;
create index idx_ops_error_logs_transport_created on ops_error_logs(transport, created_at desc) where transport is not null;
create index idx_ops_error_logs_failure_class on ops_error_logs(failure_class, created_at desc) where failure_class is not null;
create index idx_ops_error_logs_response_id on ops_error_logs(response_id) where response_id is not null;
create index idx_ops_error_logs_upstream_request_id on ops_error_logs(upstream_request_id) where upstream_request_id is not null;
-57
View File
@@ -90,60 +90,3 @@ fn parse_installation_id(raw: &str) -> Option<String> {
fn generate_installation_id() -> String {
Uuid::new_v4().to_string()
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use super::*;
static ENV_LOCK: Mutex<()> = Mutex::new(());
#[test]
fn data_dir_should_use_xdg_data_home_directly() {
let _guard = ENV_LOCK.lock().unwrap();
let dir = tempfile::tempdir().unwrap();
let previous = std::env::var_os("XDG_DATA_HOME");
std::env::set_var("XDG_DATA_HOME", dir.path());
let resolved = data_dir();
restore_env("XDG_DATA_HOME", previous);
assert_eq!(resolved, dir.path());
}
#[test]
fn load_or_create_installation_id_should_persist_under_data_dir_root() {
let dir = tempfile::tempdir().unwrap();
let data_path = dir.path().join(INSTALLATION_ID_FILE_NAME);
let installation_id = load_or_create_installation_id_from_paths(None, Some(&data_path))
.expect("installation id should be generated");
assert_eq!(fs::read_to_string(data_path).unwrap(), installation_id);
}
#[test]
fn load_or_create_installation_id_should_seed_data_dir_from_codex_desktop_id() {
let dir = tempfile::tempdir().unwrap();
let codex_path = dir.path().join(".codex").join(INSTALLATION_ID_FILE_NAME);
let data_path = dir.path().join("data").join(INSTALLATION_ID_FILE_NAME);
let desktop_id = "018f8f6b-1d7b-7b7c-b9c8-8c9f4c6d0e1a";
persist_installation_id(&codex_path, desktop_id).unwrap();
let installation_id =
load_or_create_installation_id_from_paths(Some(&codex_path), Some(&data_path))
.expect("installation id should be seeded");
assert_eq!(installation_id, desktop_id);
assert_eq!(fs::read_to_string(data_path).unwrap(), desktop_id);
}
fn restore_env(key: &str, value: Option<std::ffi::OsString>) {
if let Some(value) = value {
std::env::set_var(key, value);
} else {
std::env::remove_var(key);
}
}
}
-671
View File
@@ -1,671 +0,0 @@
//! OpenAI Chat Completions 调度服务。
use std::sync::Arc;
use std::time::Instant;
use chrono::Utc;
use serde_json::Value;
use thiserror::Error;
use crate::{
admin::monitoring::{
usage_record_model::{ResponseUsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
},
proxy::dispatch::{
auth_recovery::trigger_refresh_after_auth_failure,
cloudflare::{
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
CloudflareRecovery,
},
errors::{
auth_failure_account_status, backend_transport_name, is_auth_upstream_error,
is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error,
is_rate_limit_upstream_error, rate_limit_cooldown_until, upstream_error_body,
upstream_error_http_status,
},
upstream::{
create_response_with_account, verify_acquired_quota_if_required,
QuotaVerificationContext, QuotaVerificationDecision,
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
},
usage_events::{
reasoning_effort_from_request, record_dispatch_error_event, record_response_event,
DispatchErrorUsageRecord,
},
},
proxy::openai::chat::ChatStreamTranslationError,
upstream::accounts::{
model::AccountStatus,
pool::{AccountAcquireRequest, RuntimeAccountPoolService},
token_refresh::RuntimeTokenRefreshService,
},
upstream::{
models::{catalog::ModelCatalog, service::ModelService},
protocol::responses::{apply_response_model_options, CodexResponsesRequest},
token_client::OpenAiTokenClient,
transport::{
backend_transport_for_response_request, is_banned_upstream_error, CodexBackendClient,
CodexClientError,
},
},
};
/// OpenAI Chat Completions 调度服务。
#[derive(Clone)]
pub struct ChatDispatchService {
account_pool: Arc<RuntimeAccountPoolService>,
models: Arc<ModelService>,
codex: Arc<CodexBackendClient>,
usage_records: Arc<AdminUsageRecordService>,
token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
installation_id: Option<String>,
cloudflare: CloudflareRecovery,
}
impl ChatDispatchService {
pub(crate) fn new(
account_pool: Arc<RuntimeAccountPoolService>,
models: Arc<ModelService>,
codex: Arc<CodexBackendClient>,
usage_records: Arc<AdminUsageRecordService>,
token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
installation_id: Option<String>,
cloudflare: CloudflareRecovery,
) -> Self {
Self {
account_pool,
models,
codex,
usage_records,
token_refresh,
installation_id,
cloudflare,
}
}
/// 调度非流式 Chat Completions 请求到 Codex Responses 上游。
pub async fn complete(
&self,
request_id: &str,
mut request: CodexResponsesRequest,
requested_model: &str,
) -> Result<Value, ChatDispatchError> {
let started_at = Instant::now();
let catalog = self.models.catalog().await;
let parsed_model = catalog.parse_model_name(requested_model);
let display_model = ModelCatalog::build_display_model_name(&parsed_model);
apply_response_model_options(&mut request, &parsed_model);
let include_reasoning = request
.reasoning
.as_ref()
.and_then(|reasoning| reasoning.get("effort"))
.and_then(Value::as_str)
.is_some_and(|effort| !effort.trim().is_empty());
let tuple_schema = request.tuple_schema.clone();
let mut excluded_account_ids = Vec::new();
let mut rate_limited_count = 0usize;
let mut last_rate_limit_error = None;
let mut quota_exhausted_count = 0usize;
let mut last_quota_error = None;
let mut expired_count = 0usize;
let mut last_auth_error = None;
let mut disabled_count = 0usize;
let mut last_disabled_auth_error = None;
let mut banned_count = 0usize;
let mut last_banned_auth_error = None;
let mut last_banned_status_code: Option<u16> = None;
let mut cloudflare_challenge_count = 0usize;
let mut last_cloudflare_challenge_error = None;
let mut cloudflare_path_block_count = 0usize;
let mut last_cloudflare_path_block_error = None;
let mut model_unsupported_count = 0usize;
let mut last_model_unsupported_error = None;
let mut model_unsupported_retry_used = false;
let mut quota_verify_attempts = 0usize;
let mut last_failed_account_id = None;
macro_rules! return_dispatch_error {
($error:expr) => {{
let error = $error;
self.record_chat_dispatch_error(
request_id,
requested_model,
started_at,
last_failed_account_id.as_deref(),
Some(backend_transport_name(
backend_transport_for_response_request(&request),
)),
&error,
)
.await;
return Err(error);
}};
($error:expr, account_id: $account_id:expr, transport: $transport:expr) => {{
let error = $error;
self.record_chat_dispatch_error(
request_id,
requested_model,
started_at,
$account_id,
$transport,
&error,
)
.await;
return Err(error);
}};
}
let (account, 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 {
Some(acquired) => acquired,
None if quota_exhausted_count > 0 => {
return_dispatch_error!(ChatDispatchError::QuotaExhausted {
count: quota_exhausted_count,
upstream_error: last_quota_error.unwrap_or_default(),
});
}
None if rate_limited_count > 0 => {
return_dispatch_error!(ChatDispatchError::RateLimited {
count: rate_limited_count,
upstream_error: last_rate_limit_error.unwrap_or_default(),
});
}
None if expired_count > 0 => {
return_dispatch_error!(ChatDispatchError::Expired {
count: expired_count,
upstream_error: last_auth_error.unwrap_or_default(),
});
}
None if disabled_count > 0 => {
return_dispatch_error!(ChatDispatchError::Disabled {
count: disabled_count,
upstream_error: last_disabled_auth_error.unwrap_or_default(),
});
}
None if banned_count > 0 => {
return_dispatch_error!(ChatDispatchError::Banned {
count: banned_count,
upstream_error: last_banned_auth_error.unwrap_or_default(),
status_code: last_banned_status_code.unwrap_or(403),
});
}
None if cloudflare_challenge_count > 0 => {
return_dispatch_error!(ChatDispatchError::CloudflareChallenge {
count: cloudflare_challenge_count,
upstream_error: last_cloudflare_challenge_error.unwrap_or_default(),
});
}
None if cloudflare_path_block_count > 0 => {
return_dispatch_error!(ChatDispatchError::CloudflarePathBlocked {
count: cloudflare_path_block_count,
upstream_error: last_cloudflare_path_block_error.unwrap_or_default(),
});
}
None if model_unsupported_count > 0 => {
return_dispatch_error!(ChatDispatchError::ModelUnsupported {
count: model_unsupported_count,
upstream_error: last_model_unsupported_error.unwrap_or_default(),
});
}
None => return_dispatch_error!(ChatDispatchError::NoActiveAccount),
};
let acquired = match verify_acquired_quota_if_required(
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,
)
.await
{
QuotaVerificationDecision::Ready(acquired) => *acquired,
QuotaVerificationDecision::RetryWithAnotherAccount => {
rate_limited_count += 1;
last_rate_limit_error = Some(QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string());
continue;
}
QuotaVerificationDecision::MaxAttemptsReached => {
return_dispatch_error!(ChatDispatchError::RateLimited {
count: rate_limited_count + 1,
upstream_error: QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string(),
});
}
};
self.account_pool.wait_for_request_interval(&acquired).await;
let account = acquired.account;
let release_account_id = account.id.clone();
let response_result = create_response_with_account(
&self.codex,
self.installation_id.as_deref(),
&self.cloudflare,
&request,
request_id,
&account,
started_at,
)
.await;
self.account_pool.release(&release_account_id).await;
match response_result {
Ok(response) => break (account, response),
Err(error) if is_rate_limit_upstream_error(&error) => {
rate_limited_count += 1;
last_rate_limit_error = Some(upstream_error_body(&error));
last_failed_account_id = Some(release_account_id.clone());
let cooldown_until = rate_limit_cooldown_until(&error, Utc::now());
self.account_pool
.mark_quota_limited_until(&release_account_id, cooldown_until)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_quota_exhausted_upstream_error(&error) => {
quota_exhausted_count += 1;
last_quota_error = Some(upstream_error_body(&error));
last_failed_account_id = Some(release_account_id.clone());
self.account_pool
.set_status(&release_account_id, AccountStatus::QuotaExhausted)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_auth_upstream_error(&error) => {
let upstream_error = upstream_error_body(&error);
let account_status = auth_failure_account_status(&error);
last_failed_account_id = Some(release_account_id.clone());
match account_status {
AccountStatus::Disabled => {
disabled_count += 1;
last_disabled_auth_error = Some(upstream_error);
}
AccountStatus::Banned => {
banned_count += 1;
last_banned_status_code = Some(upstream_error_http_status(&error));
last_banned_auth_error = Some(upstream_error);
}
_ => {
expired_count += 1;
last_auth_error = Some(upstream_error);
}
}
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
cloudflare_challenge_count += 1;
last_cloudflare_challenge_error =
Some(cloudflare_challenge_error_message().to_string());
last_failed_account_id = Some(release_account_id.clone());
self.cloudflare
.apply_challenge(self.account_pool.as_ref(), &release_account_id)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_path_block_upstream_error(&error) => {
cloudflare_path_block_count += 1;
last_cloudflare_path_block_error =
Some(cloudflare_path_block_error_message().to_string());
last_failed_account_id = Some(release_account_id.clone());
self.cloudflare
.apply_path_block(self.account_pool.as_ref(), &release_account_id)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_model_unsupported_upstream_error(&error) => {
let upstream_error = upstream_error_body(&error);
if model_unsupported_retry_used {
return_dispatch_error!(
ChatDispatchError::ModelUnsupported {
count: model_unsupported_count + 1,
upstream_error,
},
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(backend_transport_for_response_request(
&request
)))
);
}
model_unsupported_count += 1;
last_model_unsupported_error = Some(upstream_error);
last_failed_account_id = Some(release_account_id.clone());
model_unsupported_retry_used = true;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_banned_upstream_error(&error) => {
banned_count += 1;
last_banned_status_code = Some(upstream_error_http_status(&error));
last_banned_auth_error = Some(upstream_error_body(&error));
last_failed_account_id = Some(release_account_id.clone());
self.account_pool
.set_status(&release_account_id, AccountStatus::Banned)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) => {
return_dispatch_error!(
ChatDispatchError::Upstream(error),
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(backend_transport_for_response_request(
&request
)))
);
}
}
};
let account_id = account.id.clone();
let body = match crate::proxy::openai::chat::chat_completion_from_codex_sse(
&response.body,
&display_model,
include_reasoning,
tuple_schema.as_ref(),
) {
Ok(Some(body)) => body,
Ok(None) => {
return_dispatch_error!(
ChatDispatchError::EmptyUpstreamResponse,
account_id: Some(&account_id),
transport: Some(backend_transport_name(response.transport))
);
}
Err(error) => {
return_dispatch_error!(
ChatDispatchError::InvalidSse(error),
account_id: Some(&account_id),
transport: Some(backend_transport_name(response.transport))
);
}
};
let response_id = body.get("id").and_then(Value::as_str);
self.cloudflare.reset_account_recovery(&account_id).await;
self.account_pool
.sync_passive_rate_limit_headers(&account, &response.rate_limit_headers)
.await;
if let Some(ref usage) = response.usage {
self.account_pool
.record_token_usage(&account_id, &request.model, usage)
.await;
}
let mut metadata = serde_json::json!({
"responseId": response_id,
"stream": false,
"transport": backend_transport_name(response.transport),
"firstTokenMs": response.first_token_ms,
"usage": response.usage,
});
if let (Some(object), Some(decision)) =
(metadata.as_object_mut(), response.websocket_pool_decision)
{
object.insert("websocketPool".to_string(), decision.metadata_value());
}
record_response_event(ResponseUsageRecord {
usage_records: &self.usage_records,
request_id,
account_id: &account_id,
route: "/v1/chat/completions",
model: &display_model,
requested_model: Some(requested_model),
client_ip: request.client_ip.as_deref(),
client_user_agent: request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(&request),
service_tier: request.service_tier.as_deref(),
started_at,
status_code: 200,
level: UsageRecordLevel::Info,
message: "v1 chat completions completed",
metadata,
rate_limit_headers: &response.rate_limit_headers,
})
.await;
Ok(body)
}
async fn record_chat_dispatch_error(
&self,
request_id: &str,
requested_model: &str,
started_at: Instant,
account_id: Option<&str>,
transport: Option<&str>,
error: &ChatDispatchError,
) {
record_chat_dispatch_error_event(ChatDispatchErrorEventRecord {
usage_records: &self.usage_records,
request_id,
account_id,
route: "/v1/chat/completions",
model: requested_model,
started_at,
transport,
error,
})
.await;
}
}
/// Chat Completions 调度错误。
#[derive(Debug, Error)]
pub enum ChatDispatchError {
#[error("failed to list runtime accounts")]
AccountStore,
#[error("no active account is available")]
NoActiveAccount,
#[error("all accounts exhausted by quota")]
QuotaExhausted {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by rate limit")]
RateLimited {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by expired auth")]
Expired {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by disabled auth")]
Disabled {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by banned auth")]
Banned {
count: usize,
upstream_error: String,
status_code: u16,
},
#[error("all accounts exhausted by Cloudflare challenge")]
CloudflareChallenge {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by Cloudflare path-block")]
CloudflarePathBlocked {
count: usize,
upstream_error: String,
},
#[error("all accounts do not support the requested model")]
ModelUnsupported {
count: usize,
upstream_error: String,
},
#[error("upstream request failed: {0}")]
Upstream(#[from] CodexClientError),
#[error("invalid upstream SSE response: {0}")]
InvalidSse(ChatStreamTranslationError),
#[error("upstream response did not include response.completed")]
EmptyUpstreamResponse,
}
impl ChatDispatchError {
pub fn http_status_code(&self) -> u16 {
match self {
Self::NoActiveAccount | Self::AccountStore => 503,
Self::QuotaExhausted { .. } => 429,
Self::RateLimited { .. } => 429,
Self::Expired { .. } | Self::Disabled { .. } => 401,
Self::Banned { status_code, .. } => *status_code,
Self::CloudflareChallenge { .. }
| Self::CloudflarePathBlocked { .. }
| Self::InvalidSse(_)
| Self::EmptyUpstreamResponse => 502,
Self::ModelUnsupported { .. } => 400,
Self::Upstream(error) => upstream_error_http_status(error),
}
}
}
// ====================================================================
// Event recording helpers
// ====================================================================
struct ChatDispatchErrorEventRecord<'a> {
usage_records: &'a AdminUsageRecordService,
request_id: &'a str,
account_id: Option<&'a str>,
route: &'a str,
model: &'a str,
started_at: Instant,
transport: Option<&'a str>,
error: &'a ChatDispatchError,
}
async fn record_chat_dispatch_error_event(record: ChatDispatchErrorEventRecord<'_>) {
let mut metadata = serde_json::json!({
"route": record.route,
"apiKind": "chat",
"stream": false,
"failed": true,
"errorKind": "dispatch",
"error": record.error.to_string(),
});
if let Some(object) = metadata.as_object_mut() {
enrich_chat_error_metadata(object, record.error);
if let Some(transport) = record.transport {
object.insert("transport".to_string(), serde_json::json!(transport));
}
}
record_dispatch_error_event(DispatchErrorUsageRecord {
usage_records: record.usage_records,
request_id: record.request_id,
account_id: record.account_id,
route: record.route,
model: record.model,
started_at: record.started_at,
status_code: i64::from(record.error.http_status_code()),
message: "chat dispatch failed",
metadata,
})
.await;
}
fn enrich_chat_error_metadata(
object: &mut serde_json::Map<String, Value>,
error: &ChatDispatchError,
) {
let (failure_class, exhausted_count, upstream_error, upstream_status) = match error {
ChatDispatchError::AccountStore => ("account_store", None, None, None),
ChatDispatchError::NoActiveAccount => ("no_available_accounts", None, None, None),
ChatDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
"quota_exhausted",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::RateLimited {
count,
upstream_error,
} => (
"rate_limited",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::Expired {
count,
upstream_error,
} => ("expired", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::Disabled {
count,
upstream_error,
} => ("disabled", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::Banned {
count,
upstream_error,
..
} => ("banned", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
"cloudflare_challenge",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
"cloudflare_path_blocked",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
"model_unsupported",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::Upstream(error) => {
let upstream_status = match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
};
(
"upstream",
None,
Some(upstream_error_body(error)),
upstream_status,
)
}
ChatDispatchError::InvalidSse(_) => ("invalid_sse", None, None, None),
ChatDispatchError::EmptyUpstreamResponse => ("empty_upstream_response", None, None, None),
};
object.insert(
"failureClass".to_string(),
Value::String(failure_class.to_string()),
);
if let Some(count) = exhausted_count {
object.insert("exhaustedCount".to_string(), serde_json::json!(count));
}
if let Some(error) = upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = upstream_status {
object.insert("upstreamStatus".to_string(), serde_json::json!(status));
}
}
+171 -1
View File
@@ -1,12 +1,182 @@
//! Shared upstream error classification for dispatch routes.
use chrono::{DateTime, Duration, Utc};
use serde_json::{json, Map, Value};
use crate::{
proxy::dispatch::exhaustion::{ExhaustedAccountKind, ExhaustedAccountRef},
upstream::accounts::model::AccountStatus,
upstream::transport::{is_banned_auth_signal, CodexBackendTransport, CodexClientError},
upstream::transport::{
is_banned_auth_signal, CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
},
};
#[derive(Clone, Debug)]
pub(crate) struct DispatchErrorMetadata {
pub failure_class: DispatchFailureClass,
pub exhausted_count: Option<usize>,
pub upstream_error: Option<String>,
pub upstream_status: Option<u16>,
pub diagnostics: Option<CodexUpstreamDiagnostics>,
}
#[derive(Clone, Copy, Debug)]
pub(crate) enum DispatchFailureClass {
NoAvailableAccounts,
QuotaExhausted,
RateLimited,
Expired,
Disabled,
Banned,
CloudflareChallenge,
CloudflarePathBlocked,
ModelUnsupported,
Upstream,
InvalidSse,
MissingCompleted,
EmptyUpstreamResponse,
ResponseFailed,
}
impl DispatchFailureClass {
pub(crate) fn as_str(self) -> &'static str {
match self {
Self::NoAvailableAccounts => "no_available_accounts",
Self::QuotaExhausted => "quota_exhausted",
Self::RateLimited => "rate_limited",
Self::Expired => "expired",
Self::Disabled => "disabled",
Self::Banned => "banned",
Self::CloudflareChallenge => "cloudflare_challenge",
Self::CloudflarePathBlocked => "cloudflare_path_blocked",
Self::ModelUnsupported => "model_unsupported",
Self::Upstream => "upstream",
Self::InvalidSse => "invalid_sse",
Self::MissingCompleted => "missing_completed",
Self::EmptyUpstreamResponse => "empty_upstream_response",
Self::ResponseFailed => "response_failed",
}
}
}
impl DispatchErrorMetadata {
pub(crate) fn no_available_accounts() -> Self {
Self {
failure_class: DispatchFailureClass::NoAvailableAccounts,
exhausted_count: None,
upstream_error: None,
upstream_status: None,
diagnostics: None,
}
}
pub(crate) fn exhausted_ref(exhausted: ExhaustedAccountRef<'_>) -> Self {
Self {
failure_class: DispatchFailureClass::from(exhausted.kind),
exhausted_count: Some(exhausted.count),
upstream_error: Some(exhausted.upstream_error.to_string()),
upstream_status: None,
diagnostics: None,
}
}
pub(crate) fn upstream(error: &CodexClientError) -> Self {
Self {
failure_class: DispatchFailureClass::Upstream,
exhausted_count: None,
upstream_error: Some(upstream_error_body(error)),
upstream_status: match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
},
diagnostics: upstream_error_diagnostics(error).cloned(),
}
}
pub(crate) fn simple(failure_class: DispatchFailureClass) -> Self {
Self {
failure_class,
exhausted_count: None,
upstream_error: None,
upstream_status: None,
diagnostics: None,
}
}
}
impl From<ExhaustedAccountKind> for DispatchFailureClass {
fn from(kind: ExhaustedAccountKind) -> Self {
match kind {
ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted,
ExhaustedAccountKind::RateLimited => Self::RateLimited,
ExhaustedAccountKind::Expired => Self::Expired,
ExhaustedAccountKind::Disabled => Self::Disabled,
ExhaustedAccountKind::Banned => Self::Banned,
ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge,
ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked,
ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported,
}
}
}
pub(crate) fn insert_dispatch_error_metadata(
object: &mut Map<String, Value>,
metadata: DispatchErrorMetadata,
) {
object.insert(
"failureClass".to_string(),
Value::String(metadata.failure_class.as_str().to_string()),
);
if let Some(count) = metadata.exhausted_count {
object.insert("exhaustedCount".to_string(), json!(count));
}
if let Some(error) = metadata.upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = metadata.upstream_status {
object.insert("upstreamStatus".to_string(), json!(status));
}
if let Some(diagnostics) = metadata.diagnostics {
insert_upstream_diagnostics_metadata(object, &diagnostics);
}
}
pub(crate) fn insert_upstream_diagnostics_metadata(
object: &mut Map<String, Value>,
diagnostics: &CodexUpstreamDiagnostics,
) {
if diagnostics.is_empty() {
return;
}
if let Some(status_code) = diagnostics.status_code {
object.insert("upstreamStatus".to_string(), json!(status_code));
}
if let Some(request_id) = &diagnostics.request_id {
object.insert(
"upstreamRequestId".to_string(),
Value::String(request_id.clone()),
);
}
if !diagnostics.trace_headers.is_empty() {
object.insert(
"upstreamTraceHeaders".to_string(),
json!(diagnostics.trace_headers),
);
}
if let Some(cf_ray) = diagnostics.cf_ray() {
object.insert("cfRay".to_string(), Value::String(cf_ray.to_string()));
}
}
pub(crate) fn upstream_error_diagnostics(
error: &CodexClientError,
) -> Option<&CodexUpstreamDiagnostics> {
match error {
CodexClientError::Upstream { diagnostics, .. } => Some(diagnostics),
_ => None,
}
}
pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
+222
View File
@@ -0,0 +1,222 @@
//! 账号耗尽状态聚合。
use crate::upstream::accounts::model::AccountStatus;
#[derive(Clone, Copy, Debug)]
pub(crate) enum ExhaustedAccountKind {
QuotaExhausted,
RateLimited,
Expired,
Disabled,
Banned,
CloudflareChallenge,
CloudflarePathBlocked,
ModelUnsupported,
}
impl ExhaustedAccountKind {
pub(crate) fn message_reason(self) -> &'static str {
match self {
Self::QuotaExhausted => "quota-exhausted",
Self::RateLimited => "rate-limited",
Self::Expired => "expired",
Self::Disabled => "disabled",
Self::Banned => "banned",
Self::CloudflareChallenge => "cloudflare-challenge",
Self::CloudflarePathBlocked => "cloudflare-path-block",
Self::ModelUnsupported => "model-unsupported",
}
}
}
#[derive(Clone, Debug)]
pub(crate) struct ExhaustedAccount {
pub kind: ExhaustedAccountKind,
pub count: usize,
pub upstream_error: String,
pub status_code: Option<u16>,
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct ExhaustedAccountRef<'a> {
pub kind: ExhaustedAccountKind,
pub count: usize,
pub upstream_error: &'a str,
}
#[derive(Default)]
pub(crate) struct AccountExhaustionTracker {
last_kind: Option<ExhaustedAccountKind>,
last_account_id: Option<String>,
quota_exhausted: ExhaustedCounter,
rate_limited: ExhaustedCounter,
expired: ExhaustedCounter,
disabled: ExhaustedCounter,
banned: ExhaustedCounter,
cloudflare_challenge: ExhaustedCounter,
cloudflare_path_blocked: ExhaustedCounter,
model_unsupported: ExhaustedCounter,
model_unsupported_retry_used: bool,
}
#[derive(Default)]
struct ExhaustedCounter {
count: usize,
upstream_error: Option<String>,
status_code: Option<u16>,
}
impl AccountExhaustionTracker {
pub(crate) fn last_account_id(&self) -> Option<&str> {
self.last_account_id.as_deref()
}
pub(crate) fn record_rate_limited(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::RateLimited,
upstream_error,
None,
);
}
pub(crate) fn record_quota_exhausted(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::QuotaExhausted,
upstream_error,
None,
);
}
pub(crate) fn record_auth_failure(
&mut self,
account_id: Option<&str>,
account_status: AccountStatus,
upstream_error: impl Into<String>,
status_code: Option<u16>,
) {
let kind = match account_status {
AccountStatus::Disabled => ExhaustedAccountKind::Disabled,
AccountStatus::Banned => ExhaustedAccountKind::Banned,
_ => ExhaustedAccountKind::Expired,
};
self.record(account_id, kind, upstream_error, status_code);
}
pub(crate) fn record_cloudflare_challenge(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::CloudflareChallenge,
upstream_error,
None,
);
}
pub(crate) fn record_cloudflare_path_blocked(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::CloudflarePathBlocked,
upstream_error,
None,
);
}
pub(crate) fn record_model_unsupported(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.model_unsupported_retry_used = true;
self.record(
account_id,
ExhaustedAccountKind::ModelUnsupported,
upstream_error,
None,
);
}
pub(crate) fn model_unsupported_retry_exhausted(
&self,
upstream_error: impl Into<String>,
) -> Option<ExhaustedAccount> {
self.model_unsupported_retry_used.then(|| ExhaustedAccount {
kind: ExhaustedAccountKind::ModelUnsupported,
count: self.model_unsupported.count + 1,
upstream_error: upstream_error.into(),
status_code: None,
})
}
pub(crate) fn last_exhausted(&self) -> Option<ExhaustedAccount> {
let kind = self.last_kind?;
let counter = self.counter(kind);
Some(ExhaustedAccount {
kind,
count: counter.count,
upstream_error: counter.upstream_error.clone().unwrap_or_default(),
status_code: counter.status_code,
})
}
fn record(
&mut self,
account_id: Option<&str>,
kind: ExhaustedAccountKind,
upstream_error: impl Into<String>,
status_code: Option<u16>,
) {
self.last_kind = Some(kind);
if let Some(account_id) = account_id {
self.last_account_id = Some(account_id.to_string());
}
let counter = self.counter_mut(kind);
counter.count += 1;
counter.upstream_error = Some(upstream_error.into());
if let Some(status_code) = status_code {
counter.status_code = Some(status_code);
}
}
fn counter(&self, kind: ExhaustedAccountKind) -> &ExhaustedCounter {
match kind {
ExhaustedAccountKind::QuotaExhausted => &self.quota_exhausted,
ExhaustedAccountKind::RateLimited => &self.rate_limited,
ExhaustedAccountKind::Expired => &self.expired,
ExhaustedAccountKind::Disabled => &self.disabled,
ExhaustedAccountKind::Banned => &self.banned,
ExhaustedAccountKind::CloudflareChallenge => &self.cloudflare_challenge,
ExhaustedAccountKind::CloudflarePathBlocked => &self.cloudflare_path_blocked,
ExhaustedAccountKind::ModelUnsupported => &self.model_unsupported,
}
}
fn counter_mut(&mut self, kind: ExhaustedAccountKind) -> &mut ExhaustedCounter {
match kind {
ExhaustedAccountKind::QuotaExhausted => &mut self.quota_exhausted,
ExhaustedAccountKind::RateLimited => &mut self.rate_limited,
ExhaustedAccountKind::Expired => &mut self.expired,
ExhaustedAccountKind::Disabled => &mut self.disabled,
ExhaustedAccountKind::Banned => &mut self.banned,
ExhaustedAccountKind::CloudflareChallenge => &mut self.cloudflare_challenge,
ExhaustedAccountKind::CloudflarePathBlocked => &mut self.cloudflare_path_blocked,
ExhaustedAccountKind::ModelUnsupported => &mut self.model_unsupported,
}
}
}
@@ -20,8 +20,8 @@ impl ImplicitResumeSnapshot {
/// 捕获当前请求中隐式续接会改写的字段。
pub fn capture(request: &CodexResponsesRequest) -> Self {
Self {
input: request.input.clone(),
previous_response_id: request.previous_response_id.clone(),
input: request.input().to_vec(),
previous_response_id: request.previous_response_id().map(ToString::to_string),
turn_state: request.turn_state.clone(),
use_websocket: request.use_websocket,
force_http_sse: request.force_http_sse,
@@ -30,8 +30,8 @@ impl ImplicitResumeSnapshot {
/// 恢复隐式续接前的请求字段。
pub fn restore(self, request: &mut CodexResponsesRequest) {
request.input = self.input;
request.previous_response_id = self.previous_response_id;
request.set_input(self.input);
request.set_previous_response_id(self.previous_response_id);
request.turn_state = self.turn_state;
request.use_websocket = self.use_websocket;
request.force_http_sse = self.force_http_sse;
+1 -1
View File
@@ -1,9 +1,9 @@
//! 请求调度模块(账号选择、回退、恢复、会话亲和性)。
mod auth_recovery;
pub mod chat;
pub mod cloudflare;
pub mod errors;
mod exhaustion;
pub mod implicit_resume;
pub mod reasoning_replay;
pub mod responses;
@@ -207,9 +207,6 @@ fn sanitize_reasoning_item(value: &Value) -> Option<Value> {
}
sanitized.insert("summary".to_string(), Value::Array(summary));
sanitized.insert("encrypted_content".to_string(), json!(encrypted_content));
if let Some(content) = object.get("content").and_then(sanitize_content) {
sanitized.insert("content".to_string(), Value::Array(content));
}
Some(Value::Object(sanitized))
}
@@ -250,22 +247,6 @@ fn sanitize_summary(value: &Value) -> Option<Vec<Value>> {
)
}
fn sanitize_content(value: &Value) -> Option<Vec<Value>> {
let content = value
.as_array()?
.iter()
.filter_map(|part| {
let object = part.as_object()?;
if object.get("type").and_then(Value::as_str) != Some("reasoning_text") {
return None;
}
let text = object.get("text")?.as_str()?;
Some(json!({"type": "reasoning_text", "text": text}))
})
.collect::<Vec<_>>();
(!content.is_empty()).then_some(content)
}
fn replay_item_key(item: &Value) -> Option<String> {
match item.get("type").and_then(Value::as_str)? {
"reasoning" => item
@@ -26,8 +26,7 @@ pub(super) async fn record_response_affinity(
usage: Option<TokenUsage>,
) {
let Some(conversation_id) = request
.prompt_cache_key
.as_deref()
.prompt_cache_key()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
@@ -52,7 +51,7 @@ pub(super) async fn record_response_affinity(
turn_state: turn_state
.filter(|value| !value.trim().is_empty())
.or_else(|| request.turn_state.clone()),
instructions_hash: Some(hash_instructions(Some(&request.instructions))),
instructions_hash: Some(hash_instructions(Some(request.instructions()))),
input_tokens: usage.map(|usage| usage.input_tokens),
function_call_ids: metadata.function_call_ids,
variant_hash: Some(variant_hash.clone()),
@@ -86,11 +85,7 @@ pub(super) async fn evict_reasoning_replay(
account_id: &str,
) {
let variant_hash = compute_variant_hash(request);
let conversation_id = request
.prompt_cache_key
.as_deref()
.unwrap_or("")
.to_string();
let conversation_id = request.prompt_cache_key().unwrap_or("").to_string();
reasoning_replay.lock().await.evict_by_identity(
account_id,
&conversation_id,
+203 -104
View File
@@ -1,8 +1,14 @@
use serde_json::{json, Value};
use serde_json::Value;
use thiserror::Error;
use crate::{
proxy::dispatch::errors::{upstream_error_body, upstream_error_http_status},
proxy::dispatch::{
errors::{
insert_dispatch_error_metadata, upstream_error_http_status, DispatchErrorMetadata,
DispatchFailureClass,
},
exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
},
upstream::{
protocol::{responses::ResponsesSseFailure, sse::SseError},
transport::CodexClientError,
@@ -14,8 +20,6 @@ use super::sse_failure::{sse_failure_error_body, stream_failure_http_status};
/// Responses 调度错误。
#[derive(Debug, Error)]
pub enum ResponseDispatchError {
#[error("failed to list runtime accounts")]
AccountStore,
#[error("no active account is available")]
NoActiveAccount,
#[error("all accounts exhausted by quota")]
@@ -72,9 +76,47 @@ pub enum ResponseDispatchError {
}
impl ResponseDispatchError {
pub(crate) fn from_exhausted_account(exhausted: ExhaustedAccount) -> Self {
match exhausted.kind {
ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::RateLimited => Self::RateLimited {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Expired => Self::Expired {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Disabled => Self::Disabled {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Banned => Self::Banned {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
status_code: exhausted.status_code.unwrap_or(403),
},
ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
}
}
pub fn http_status_code(&self) -> u16 {
match self {
Self::NoActiveAccount | Self::AccountStore => 503,
Self::NoActiveAccount => 503,
Self::QuotaExhausted { .. } => 429,
Self::RateLimited { .. } => 429,
Self::Expired { .. } | Self::Disabled { .. } => 401,
@@ -89,6 +131,161 @@ impl ResponseDispatchError {
Self::Upstream(error) => upstream_error_http_status(error),
}
}
pub(crate) fn metadata(&self) -> DispatchErrorMetadata {
match self {
Self::NoActiveAccount => DispatchErrorMetadata::no_available_accounts(),
Self::QuotaExhausted {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::QuotaExhausted,
*count,
upstream_error,
),
Self::RateLimited {
count,
upstream_error,
} => {
Self::exhausted_metadata(ExhaustedAccountKind::RateLimited, *count, upstream_error)
}
Self::Expired {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Expired, *count, upstream_error),
Self::Disabled {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Disabled, *count, upstream_error),
Self::Banned {
count,
upstream_error,
..
} => Self::exhausted_metadata(ExhaustedAccountKind::Banned, *count, upstream_error),
Self::CloudflareChallenge {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflareChallenge,
*count,
upstream_error,
),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflarePathBlocked,
*count,
upstream_error,
),
Self::ModelUnsupported {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::ModelUnsupported,
*count,
upstream_error,
),
Self::Upstream(error) => DispatchErrorMetadata::upstream(error),
Self::InvalidSse(_) => DispatchErrorMetadata::simple(DispatchFailureClass::InvalidSse),
Self::MissingCompleted => {
DispatchErrorMetadata::simple(DispatchFailureClass::MissingCompleted)
}
Self::EmptyUpstreamResponse => {
DispatchErrorMetadata::simple(DispatchFailureClass::EmptyUpstreamResponse)
}
Self::Failed(failure) => DispatchErrorMetadata {
failure_class: DispatchFailureClass::ResponseFailed,
exhausted_count: None,
upstream_error: Some(sse_failure_error_body(failure)),
upstream_status: None,
diagnostics: None,
},
}
}
pub(crate) fn exhausted_account(&self) -> Option<ExhaustedAccountRef<'_>> {
match self {
Self::QuotaExhausted {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::QuotaExhausted,
count: *count,
upstream_error,
}),
Self::RateLimited {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::RateLimited,
count: *count,
upstream_error,
}),
Self::Expired {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Expired,
count: *count,
upstream_error,
}),
Self::Disabled {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Disabled,
count: *count,
upstream_error,
}),
Self::Banned {
count,
upstream_error,
..
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Banned,
count: *count,
upstream_error,
}),
Self::CloudflareChallenge {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflareChallenge,
count: *count,
upstream_error,
}),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflarePathBlocked,
count: *count,
upstream_error,
}),
Self::ModelUnsupported {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::ModelUnsupported,
count: *count,
upstream_error,
}),
_ => None,
}
}
fn exhausted_metadata(
kind: ExhaustedAccountKind,
count: usize,
upstream_error: &str,
) -> DispatchErrorMetadata {
DispatchErrorMetadata::exhausted_ref(ExhaustedAccountRef {
kind,
count,
upstream_error,
})
}
}
/// Responses live SSE body stream error.
@@ -132,103 +329,5 @@ pub(super) fn enrich_response_dispatch_error_metadata(
let Some(object) = metadata.as_object_mut() else {
return;
};
let (failure_class, exhausted_count, upstream_error, upstream_status) = match error {
ResponseDispatchError::AccountStore => ("account_store", None, None, None),
ResponseDispatchError::NoActiveAccount => ("no_available_accounts", None, None, None),
ResponseDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
"quota_exhausted",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::RateLimited {
count,
upstream_error,
} => (
"rate_limited",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::Expired {
count,
upstream_error,
} => ("expired", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::Disabled {
count,
upstream_error,
} => ("disabled", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::Banned {
count,
upstream_error,
..
} => ("banned", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
"cloudflare_challenge",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
"cloudflare_path_blocked",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
"model_unsupported",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::Upstream(error) => {
let upstream_status = match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
};
(
"upstream",
None,
Some(upstream_error_body(error)),
upstream_status,
)
}
ResponseDispatchError::InvalidSse(_) => ("invalid_sse", None, None, None),
ResponseDispatchError::MissingCompleted => ("missing_completed", None, None, None),
ResponseDispatchError::EmptyUpstreamResponse => {
("empty_upstream_response", None, None, None)
}
ResponseDispatchError::Failed(failure) => (
"response_failed",
None,
Some(sse_failure_error_body(failure)),
None,
),
};
object.insert(
"failureClass".to_string(),
Value::String(failure_class.to_string()),
);
if let Some(count) = exhausted_count {
object.insert("exhaustedCount".to_string(), json!(count));
}
if let Some(error) = upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = upstream_status {
object.insert("upstreamStatus".to_string(), json!(status));
}
insert_dispatch_error_metadata(object, error.metadata());
}
@@ -1,24 +1,32 @@
use std::time::Instant;
use serde_json::{json, Value};
use serde_json::{json, Map, Value};
use crate::{
admin::monitoring::{
usage_record_model::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
ops_error_model::OpsErrorLog,
ops_error_service::AdminOpsErrorLogService,
usage_record_model::{UsageRecord, UsageRecordLevel},
},
infra::time::elapsed_millis_i64,
proxy::dispatch::{
errors::{backend_transport_name, upstream_error_http_status},
errors::{
backend_transport_name, insert_dispatch_error_metadata,
insert_upstream_diagnostics_metadata, upstream_error_http_status,
DispatchErrorMetadata,
},
usage_events::{
enrich_event_route_metadata, enrich_usage_record_identity,
event_kind as response_event_kind, reasoning_effort_from_request,
record_dispatch_error_event, record_response_event, DispatchErrorUsageRecord,
record_dispatch_error_event, DispatchErrorLogRecord,
},
},
upstream::{
protocol::responses::{CodexResponsesRequest, ResponsesSseFailure},
transport::{CodexBackendTransport, CodexClientError, WebSocketPoolDecision},
transport::{
CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics,
WebSocketPoolDecision,
},
},
};
@@ -28,10 +36,11 @@ use super::{
},
sse_failure::{status_code_for_stream_failure, stream_failure_metadata},
stream_lifecycle::{latest_response_id, LiveResponseStreamContext},
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
};
pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
pub(super) usage_records: &'a AdminUsageRecordService,
pub(super) ops_errors: &'a AdminOpsErrorLogService,
pub(super) request_id: &'a str,
pub(super) account_id: &'a str,
pub(super) account_email: Option<&'a str>,
@@ -40,11 +49,14 @@ pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
pub(super) started_at: Instant,
pub(super) stream: bool,
pub(super) transport: CodexBackendTransport,
pub(super) request: &'a CodexResponsesRequest,
pub(super) error: &'a CodexClientError,
pub(super) trace: &'a ResponseDispatchTrace,
pub(super) attempt: Option<&'a ResponseDispatchAttempt>,
}
pub(super) struct ResponseStreamFailureEventRecord<'a> {
pub(super) usage_records: &'a AdminUsageRecordService,
pub(super) ops_errors: &'a AdminOpsErrorLogService,
pub(super) request_id: &'a str,
pub(super) account_id: &'a str,
pub(super) route: &'a str,
@@ -54,12 +66,15 @@ pub(super) struct ResponseStreamFailureEventRecord<'a> {
pub(super) transport: CodexBackendTransport,
pub(super) request: &'a CodexResponsesRequest,
pub(super) failure: &'a ResponsesSseFailure,
pub(super) diagnostics: &'a CodexUpstreamDiagnostics,
pub(super) rate_limit_headers: &'a [(String, String)],
pub(super) prefetched: &'a [u8],
pub(super) trace: &'a ResponseDispatchTrace,
pub(super) attempt: &'a ResponseDispatchAttempt,
}
pub(super) struct ResponseDispatchErrorEventRecord<'a> {
pub(super) usage_records: &'a AdminUsageRecordService,
pub(super) ops_errors: &'a AdminOpsErrorLogService,
pub(super) request_id: &'a str,
pub(super) account_id: Option<&'a str>,
pub(super) route: &'a str,
@@ -88,8 +103,16 @@ pub(super) async fn record_response_dispatch_error_event(
record.transport,
);
enrich_response_dispatch_error_metadata(&mut metadata, record.error);
record_dispatch_error_event(DispatchErrorUsageRecord {
usage_records: record.usage_records,
if let Some(object) = metadata.as_object_mut() {
insert_response_status_metadata_object(
object,
i64::from(record.error.http_status_code()),
dispatch_error_client_status_code(record.error, record.stream, record.compact),
dispatch_error_upstream_status_code(record.error),
);
}
record_dispatch_error_event(DispatchErrorLogRecord {
ops_errors: record.ops_errors,
request_id: record.request_id,
account_id: record.account_id,
route: record.route,
@@ -105,6 +128,7 @@ pub(super) async fn record_response_dispatch_error_event(
pub(super) async fn record_response_upstream_error_event(
record: ResponseUpstreamErrorEventRecord<'_>,
) {
let event_status_code = i64::from(upstream_error_http_status(record.error));
let mut metadata = dispatch_error_metadata(
record.error,
record.stream,
@@ -112,18 +136,23 @@ pub(super) async fn record_response_upstream_error_event(
Some(backend_transport_name(record.transport)),
);
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
"v1.response",
UsageRecordLevel::Error,
"v1 responses upstream request failed",
);
let mut event = OpsErrorLog::new("v1.response", "v1 responses upstream request failed");
event.request_id = Some(record.request_id.to_string());
event.account_id = Some(record.account_id.to_string());
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(i64::from(upstream_error_http_status(record.error)));
event.status_code = Some(event_status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
if let Some(object) = metadata.as_object_mut() {
insert_dispatch_error_metadata(object, DispatchErrorMetadata::upstream(record.error));
insert_response_status_metadata_object(
object,
event_status_code,
upstream_failure_client_status_code(record.stream),
Some(event_status_code),
);
insert_response_request_summary_object(object, record.request, record.transport);
insert_response_trace_metadata_object(object, record.trace.attempts(), record.attempt);
if let Some(account_email) = record
.account_email
.map(str::trim)
@@ -135,8 +164,17 @@ pub(super) async fn record_response_upstream_error_event(
);
}
}
enrich_usage_record_identity(
&mut metadata,
Some(record.request.model()),
record.model,
record.request.client_ip.as_deref(),
record.request.client_user_agent.as_deref(),
reasoning_effort_from_request(record.request),
record.request.service_tier(),
);
event.metadata = metadata;
if let Err(error) = record.usage_records.record(event).await {
if let Err(error) = record.ops_errors.record(event).await {
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record upstream error event");
}
}
@@ -144,6 +182,7 @@ pub(super) async fn record_response_upstream_error_event(
pub(super) async fn record_prefetched_response_stream_failure_event(
record: ResponseStreamFailureEventRecord<'_>,
) {
let event_status_code = status_code_for_stream_failure(record.failure);
let mut metadata = stream_failure_metadata(record.failure, None);
if let Some(object) = metadata.as_object_mut() {
if record.transport == CodexBackendTransport::WebSocket {
@@ -152,31 +191,57 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
Value::String("websocket".to_string()),
);
}
insert_response_status_metadata_object(
object,
event_status_code,
200,
diagnostics_status_code(record.diagnostics),
);
insert_response_request_summary_object(object, record.request, record.transport);
object.insert("requestBody".to_string(), json!(record.request));
object.insert(
"responseBody".to_string(),
Value::String(String::from_utf8_lossy(record.prefetched).to_string()),
);
insert_upstream_diagnostics_metadata(object, record.diagnostics);
insert_response_trace_metadata_object(
object,
record.trace.attempts(),
Some(record.attempt),
);
}
enrich_event_route_metadata(&mut metadata, record.route);
enrich_usage_record_identity(
&mut metadata,
Some(record.requested_model),
record.model,
record.request.client_ip.as_deref(),
record.request.client_user_agent.as_deref(),
reasoning_effort_from_request(record.request),
record.request.service_tier(),
);
let mut event = OpsErrorLog::new(
response_event_kind(record.route),
"v1 responses stream failed",
);
event.request_id = Some(record.request_id.to_string());
event.account_id = Some(record.account_id.to_string());
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(event_status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
if !record.rate_limit_headers.is_empty() {
if let Some(object) = metadata.as_object_mut() {
object.insert(
"rateLimitHeaders".to_string(),
serde_json::json!(record.rate_limit_headers),
);
}
}
event.metadata = metadata;
if let Err(error) = record.ops_errors.record(event).await {
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record prefetched stream error event");
}
record_response_event(ResponseUsageRecord {
usage_records: record.usage_records,
request_id: record.request_id,
account_id: record.account_id,
route: record.route,
model: record.model,
requested_model: Some(record.requested_model),
client_ip: record.request.client_ip.as_deref(),
client_user_agent: record.request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(record.request),
service_tier: record.request.service_tier.as_deref(),
started_at: record.started_at,
status_code: status_code_for_stream_failure(record.failure),
level: UsageRecordLevel::Error,
message: "v1 responses stream failed",
metadata,
rate_limit_headers: record.rate_limit_headers,
})
.await;
}
fn ensure_stream_metadata_flag(metadata: &mut Value) {
@@ -193,6 +258,7 @@ fn enrich_live_response_stream_metadata(
context: &LiveResponseStreamContext,
rate_limit_headers: &[(String, String)],
metadata: &mut Value,
status_code: i64,
body: &str,
) {
let Some(object) = metadata.as_object_mut() else {
@@ -211,6 +277,15 @@ fn enrich_live_response_stream_metadata(
.entry("websocketPool".to_string())
.or_insert_with(|| decision.metadata_value());
}
insert_response_status_metadata_object(
object,
status_code,
200,
diagnostics_status_code(&context.diagnostics),
);
insert_response_request_summary_object(object, &context.request, context.transport);
insert_upstream_diagnostics_metadata(object, &context.diagnostics);
insert_response_trace_metadata_object(object, &context.attempts, Some(&context.attempt));
object
.entry("requestBody".to_string())
.or_insert_with(|| serde_json::json!(context.request));
@@ -219,6 +294,181 @@ fn enrich_live_response_stream_metadata(
.or_insert_with(|| Value::String(body.to_string()));
}
pub(super) fn insert_response_upstream_diagnostics(
metadata: &mut Value,
diagnostics: &CodexUpstreamDiagnostics,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
insert_upstream_diagnostics_metadata(object, diagnostics);
}
pub(super) fn insert_response_status_metadata(
metadata: &mut Value,
event_status_code: i64,
client_status_code: i64,
upstream_status_code: Option<i64>,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
insert_response_status_metadata_object(
object,
event_status_code,
client_status_code,
upstream_status_code,
);
}
pub(super) fn insert_response_trace_metadata(
metadata: &mut Value,
trace: &ResponseDispatchTrace,
current_attempt: Option<&ResponseDispatchAttempt>,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
insert_response_trace_metadata_object(object, trace.attempts(), current_attempt);
}
fn insert_response_attempt_metadata_object(
object: &mut serde_json::Map<String, Value>,
attempt: Option<&ResponseDispatchAttempt>,
) {
let Some(attempt) = attempt else {
return;
};
object.insert("attemptIndex".to_string(), json!(attempt.index()));
object.insert(
"attemptAccountId".to_string(),
Value::String(attempt.account_id().to_string()),
);
}
fn insert_response_trace_metadata_object(
object: &mut Map<String, Value>,
attempts: &[ResponseDispatchAttempt],
current_attempt: Option<&ResponseDispatchAttempt>,
) {
insert_response_attempt_metadata_object(object, current_attempt);
object.insert("attemptCount".to_string(), json!(attempts.len()));
object.insert(
"attempts".to_string(),
Value::Array(
attempts
.iter()
.map(|attempt| {
json!({
"index": attempt.index(),
"accountId": attempt.account_id(),
})
})
.collect(),
),
);
}
fn insert_response_status_metadata_object(
object: &mut Map<String, Value>,
event_status_code: i64,
client_status_code: i64,
upstream_status_code: Option<i64>,
) {
object.insert("eventStatusCode".to_string(), json!(event_status_code));
object.insert("clientStatusCode".to_string(), json!(client_status_code));
if let Some(upstream_status_code) = upstream_status_code {
object.insert(
"upstreamStatusCode".to_string(),
json!(upstream_status_code),
);
}
}
fn insert_response_request_summary_object(
object: &mut Map<String, Value>,
request: &CodexResponsesRequest,
transport: CodexBackendTransport,
) {
object.insert(
"requestSummary".to_string(),
response_request_summary(request, transport),
);
}
fn response_request_summary(
request: &CodexResponsesRequest,
transport: CodexBackendTransport,
) -> Value {
let body = request.body();
let input = body.get("input");
let tools = body.get("tools");
json!({
"model": request.model(),
"stream": request.stream(),
"store": request.store(),
"transport": backend_transport_name(transport),
"inputType": json_value_kind(input),
"inputItemsCount": input.and_then(Value::as_array).map(Vec::len),
"toolsType": json_value_kind(tools),
"toolsCount": tools.and_then(Value::as_array).map(Vec::len),
"topLevelFields": body.keys().cloned().collect::<Vec<_>>(),
"previousResponseIdPresent": request.previous_response_id().is_some(),
"serviceTier": request.service_tier(),
"localTransport": {
"useWebsocket": request.use_websocket,
"forceHttpSse": request.force_http_sse,
},
})
}
fn json_value_kind(value: Option<&Value>) -> &'static str {
match value {
None => "missing",
Some(Value::Null) => "null",
Some(Value::Bool(_)) => "boolean",
Some(Value::Number(_)) => "number",
Some(Value::String(_)) => "string",
Some(Value::Array(_)) => "array",
Some(Value::Object(_)) => "object",
}
}
fn diagnostics_status_code(diagnostics: &CodexUpstreamDiagnostics) -> Option<i64> {
diagnostics.status_code.map(i64::from)
}
fn dispatch_error_upstream_status_code(error: &ResponseDispatchError) -> Option<i64> {
match error {
ResponseDispatchError::Upstream(error) => {
Some(i64::from(upstream_error_http_status(error)))
}
_ => None,
}
}
fn dispatch_error_client_status_code(
error: &ResponseDispatchError,
stream: bool,
compact: bool,
) -> i64 {
if stream {
return 200;
}
if matches!(error, ResponseDispatchError::Upstream(_)) && !compact {
return 502;
}
i64::from(error.http_status_code())
}
fn upstream_failure_client_status_code(stream: bool) -> i64 {
if stream {
200
} else {
502
}
}
pub(super) async fn record_live_response_stream_event(
context: &LiveResponseStreamContext,
status_code: i64,
@@ -230,15 +480,14 @@ pub(super) async fn record_live_response_stream_event(
) {
ensure_stream_metadata_flag(&mut metadata);
enrich_event_route_metadata(&mut metadata, &context.route);
enrich_live_response_stream_metadata(context, rate_limit_headers, &mut metadata, body);
enrich_live_response_stream_metadata(
context,
rate_limit_headers,
&mut metadata,
status_code,
body,
);
log_live_response_stream_finalized(context, status_code, level, message, &metadata, body);
let mut event = UsageRecord::new(response_event_kind(&context.route), level, message);
event.request_id = Some(context.request_id.clone());
event.account_id = Some(context.account_id.clone());
event.route = Some(context.route.clone());
event.model = Some(context.display_model.clone());
event.status_code = Some(status_code);
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
enrich_usage_record_identity(
&mut metadata,
Some(&context.requested_model),
@@ -246,14 +495,40 @@ pub(super) async fn record_live_response_stream_event(
context.client_ip.as_deref(),
context.request.client_user_agent.as_deref(),
reasoning_effort_from_request(&context.request),
context.request.service_tier.as_deref(),
context.request.service_tier(),
);
event.metadata = metadata;
if let Err(error) = context.usage_records.record(event).await {
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream event");
if is_success_usage_event(level, status_code) {
let mut event = UsageRecord::new(response_event_kind(&context.route), level, message);
event.request_id = Some(context.request_id.clone());
event.account_id = Some(context.account_id.clone());
event.route = Some(context.route.clone());
event.model = Some(context.display_model.clone());
event.status_code = Some(status_code);
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
event.metadata = metadata;
if let Err(error) = context.usage_records.record(event).await {
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream event");
}
} else {
let mut event = OpsErrorLog::new(response_event_kind(&context.route), message);
event.request_id = Some(context.request_id.clone());
event.account_id = Some(context.account_id.clone());
event.route = Some(context.route.clone());
event.model = Some(context.display_model.clone());
event.status_code = Some(status_code);
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
event.metadata = metadata;
if let Err(error) = context.ops_errors.record(event).await {
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream error event");
}
}
}
fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool {
level != UsageRecordLevel::Error && status_code < 400
}
fn log_live_response_stream_finalized(
context: &LiveResponseStreamContext,
status_code: i64,
@@ -8,3 +8,4 @@ mod prefetch;
pub(crate) mod service;
mod sse_failure;
mod stream_lifecycle;
mod trace;
File diff suppressed because it is too large Load Diff
@@ -5,7 +5,8 @@ use tokio::sync::Mutex;
use crate::{
admin::monitoring::{
usage_record_model::UsageRecordLevel, usage_record_service::AdminUsageRecordService,
ops_error_service::AdminOpsErrorLogService, usage_record_model::UsageRecordLevel,
usage_record_service::AdminUsageRecordService,
},
infra::time::elapsed_millis_i64,
proxy::{
@@ -24,7 +25,7 @@ use crate::{
},
transport::{
CodexBackendTransport, CodexRateLimitHeaderUpdates, CodexTurnStateUpdate,
WebSocketPoolDecision,
CodexUpstreamDiagnostics, WebSocketPoolDecision,
},
},
};
@@ -40,6 +41,7 @@ use super::{
stream_failure_metadata, stream_failure_source, synthetic_stream_disconnected_detail,
STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE,
},
trace::ResponseDispatchAttempt,
};
use crate::proxy::dispatch::cloudflare::CloudflareRecovery;
@@ -48,6 +50,7 @@ pub(super) struct LiveResponseStreamContext {
pub(super) session_affinity: Arc<RuntimeSessionAffinityService>,
pub(super) reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
pub(super) usage_records: Arc<AdminUsageRecordService>,
pub(super) ops_errors: Arc<AdminOpsErrorLogService>,
pub(super) cloudflare: CloudflareRecovery,
pub(super) account_id: String,
pub(super) account_plan_type: Option<String>,
@@ -65,6 +68,9 @@ pub(super) struct LiveResponseStreamContext {
pub(super) turn_state_update: Option<CodexTurnStateUpdate>,
pub(super) websocket_pool_decision: Option<WebSocketPoolDecision>,
pub(super) turn_state: Option<String>,
pub(super) diagnostics: CodexUpstreamDiagnostics,
pub(super) attempt: ResponseDispatchAttempt,
pub(super) attempts: Vec<ResponseDispatchAttempt>,
pub(super) started_at: Instant,
}
@@ -0,0 +1,39 @@
//! Request-level trace state for Responses dispatch.
#[derive(Debug, Default)]
pub(super) struct ResponseDispatchTrace {
next_attempt_index: i64,
attempts: Vec<ResponseDispatchAttempt>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct ResponseDispatchAttempt {
index: i64,
account_id: String,
}
impl ResponseDispatchTrace {
pub(super) fn start_attempt(&mut self, account_id: &str) -> ResponseDispatchAttempt {
let attempt = ResponseDispatchAttempt {
index: self.next_attempt_index,
account_id: account_id.to_string(),
};
self.next_attempt_index += 1;
self.attempts.push(attempt.clone());
attempt
}
pub(super) fn attempts(&self) -> &[ResponseDispatchAttempt] {
&self.attempts
}
}
impl ResponseDispatchAttempt {
pub(super) fn index(&self) -> i64 {
self.index
}
pub(super) fn account_id(&self) -> &str {
&self.account_id
}
}
+62 -32
View File
@@ -106,6 +106,38 @@ impl SessionAffinityMap {
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<String> {
self.latest_entry_by_conversation(conversation_id, max_age, variant_hash, now)
.map(|(response_id, _)| response_id.clone())
}
pub fn lookup_latest_account_by_conversation(
&self,
conversation_id: &str,
max_age: Option<Duration>,
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<String> {
self.latest_entry_by_conversation(conversation_id, max_age, variant_hash, now)
.map(|(_, entry)| entry.account_id.clone())
}
pub fn forget(&mut self, response_id: &str) -> bool {
self.entries.remove(response_id).is_some()
}
fn active_entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<&SessionAffinityEntry> {
self.entries
.get(response_id)
.filter(|entry| self.entry_is_active(entry, now))
}
fn latest_entry_by_conversation(
&self,
conversation_id: &str,
max_age: Option<Duration>,
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<(&String, &SessionAffinityEntry)> {
self.entries
.iter()
.filter(|(_, entry)| entry.conversation_id == conversation_id)
@@ -118,25 +150,6 @@ impl SessionAffinityMap {
max_age.is_none_or(|max_age| now.signed_duration_since(entry.created_at) <= max_age)
})
.max_by_key(|(_, entry)| entry.created_at)
.map(|(response_id, _)| response_id.clone())
}
pub fn forget(&mut self, response_id: &str) -> bool {
self.entries.remove(response_id).is_some()
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
fn active_entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<&SessionAffinityEntry> {
self.entries
.get(response_id)
.filter(|entry| self.entry_is_active(entry, now))
}
fn stored_record_is_active(&self, record: &StoredSessionAffinity, now: DateTime<Utc>) -> bool {
@@ -169,12 +182,9 @@ fn compute_variant_hash_with_identity(
request: &CodexResponsesRequest,
identity: Option<&str>,
) -> String {
let tools_json = request
.tools
.as_ref()
.map_or_else(|| "[]".to_string(), |tools| tools_json(tools));
let tools_json = request.tools().map_or_else(|| "[]".to_string(), tools_json);
let mut hasher = Sha256::new();
hasher.update(request.instructions.as_bytes());
hasher.update(request.instructions().as_bytes());
hasher.update(b"\0");
hasher.update(tools_json.as_bytes());
if let Some(identity) = identity
@@ -263,23 +273,28 @@ pub fn build_conversation_identity(
/// 确保请求拥有上游可复用的 prompt cache key。
pub fn ensure_prompt_cache_key(request: &mut CodexResponsesRequest) {
if let Some(existing) = request
.prompt_cache_key
.as_deref()
.prompt_cache_key()
.map(str::trim)
.filter(|value| !value.is_empty())
{
request.prompt_cache_key = Some(existing.to_string());
let existing = existing.to_string();
request.set_prompt_cache_key(Some(existing));
return;
}
request.prompt_cache_key =
Some(derive_stable_conversation_key(request).unwrap_or_else(|| Uuid::new_v4().to_string()));
request.set_prompt_cache_key(Some(
derive_stable_conversation_key(request).unwrap_or_else(|| Uuid::new_v4().to_string()),
));
}
/// 按原版 `stable-conversation-key.ts` 的规则派生稳定 conversation key。
pub fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option<String> {
let instructions = request.instructions.chars().take(2000).collect::<String>();
let first_user_text = first_user_text(&request.input);
let instructions = request
.instructions()
.chars()
.take(2000)
.collect::<String>();
let first_user_text = first_user_text(request.input());
let normalized_first_user_text = normalize_conversation_anchor_text(&first_user_text);
let first_user_text = if normalized_first_user_text.is_empty() {
first_user_text
@@ -291,7 +306,7 @@ pub fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option
}
let mut hasher = Sha256::new();
hasher.update(request.model.as_bytes());
hasher.update(request.model().as_bytes());
hasher.update(b"\0");
hasher.update(instructions.as_bytes());
hasher.update(b"\0");
@@ -636,6 +651,21 @@ impl RuntimeSessionAffinityService {
.lookup_latest_response_by_conversation(conversation_id, max_age, variant_hash, now)
}
pub async fn lookup_latest_account_by_conversation(
&self,
conversation_id: &str,
max_age: Option<Duration>,
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<String> {
self.map.read().await.lookup_latest_account_by_conversation(
conversation_id,
max_age,
variant_hash,
now,
)
}
pub async fn forget(&self, response_id: &str) -> bool {
let in_memory = self.map.write().await.forget(response_id);
match self.store.forget(response_id).await {
+12 -4
View File
@@ -34,6 +34,7 @@ pub(crate) const QUOTA_VERIFY_LIMIT_REACHED_MESSAGE: &str =
pub(crate) enum QuotaVerificationDecision {
Ready(Box<AcquiredAccount>),
RetryWithAnotherAccount,
RequiredAccountUnavailable,
MaxAttemptsReached,
}
@@ -45,6 +46,7 @@ pub(crate) struct QuotaVerificationContext<'a> {
pub request_id: &'a str,
pub excluded_account_ids: &'a mut Vec<String>,
pub verify_attempts: &'a mut usize,
pub allow_retry_with_another_account: bool,
}
pub(crate) async fn verify_acquired_quota_if_required(
@@ -101,7 +103,10 @@ pub(crate) async fn verify_acquired_quota_if_required(
.apply_quota_snapshot(&account_id, &quota)
.await;
if quota_snapshot_limit_reached(&quota) {
context.account_pool.release(&account_id).await;
context
.account_pool
.release_without_request_usage(&account_id)
.await;
context.excluded_account_ids.push(account_id.clone());
*context.verify_attempts += 1;
let max_attempts_reached = *context.verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS;
@@ -111,9 +116,12 @@ pub(crate) async fn verify_acquired_quota_if_required(
quota_verify_required = true,
quota_verify_result = "limit_reached",
verify_attempts = *context.verify_attempts,
retry_with_another_account = !max_attempts_reached,
retry_with_another_account = context.allow_retry_with_another_account && !max_attempts_reached,
"quota verification reported exhausted account before upstream request"
);
if !context.allow_retry_with_another_account {
return QuotaVerificationDecision::RequiredAccountUnavailable;
}
if max_attempts_reached {
return QuotaVerificationDecision::MaxAttemptsReached;
}
@@ -152,7 +160,7 @@ pub(crate) async fn create_response_with_account(
.cookie_header_for_request(&account.id, "/codex/responses")
.await;
let identity = build_conversation_identity(
request.prompt_cache_key.as_deref(),
request.prompt_cache_key(),
request.codex_window_id.as_deref(),
&account.id,
);
@@ -192,7 +200,7 @@ async fn create_response_stream_with_account(
.cookie_header_for_request(&account.id, "/codex/responses")
.await;
let identity = build_conversation_identity(
request.prompt_cache_key.as_deref(),
request.prompt_cache_key(),
request.codex_window_id.as_deref(),
&account.id,
);
+14 -23
View File
@@ -6,15 +6,16 @@ use serde_json::Value;
use crate::{
admin::monitoring::{
usage_record_model::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
ops_error_model::OpsErrorLog,
ops_error_service::AdminOpsErrorLogService,
usage_record_model::{ResponseUsageRecord, UsageRecord},
},
infra::time::elapsed_millis_i64,
upstream::protocol::responses::{CodexCompactRequest, CodexResponsesRequest},
};
pub(super) struct DispatchErrorUsageRecord<'a> {
pub usage_records: &'a AdminUsageRecordService,
pub(super) struct DispatchErrorLogRecord<'a> {
pub ops_errors: &'a AdminOpsErrorLogService,
pub request_id: &'a str,
pub account_id: Option<&'a str>,
pub route: &'a str,
@@ -25,15 +26,11 @@ pub(super) struct DispatchErrorUsageRecord<'a> {
pub metadata: Value,
}
pub(super) async fn record_dispatch_error_event(record: DispatchErrorUsageRecord<'_>) {
pub(super) async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
event_kind(record.route),
UsageRecordLevel::Error,
record.message,
);
let mut event = OpsErrorLog::new(event_kind(record.route), record.message);
event.request_id = Some(record.request_id.to_string());
event.account_id = record.account_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
@@ -42,7 +39,7 @@ pub(super) async fn record_dispatch_error_event(record: DispatchErrorUsageRecord
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.metadata = metadata;
if let Err(error) = record.usage_records.record(event).await {
if let Err(error) = record.ops_errors.record(event).await {
tracing::warn!(
account_id = record.account_id.unwrap_or(""),
error = %error,
@@ -88,19 +85,13 @@ pub(super) async fn record_response_event(record: ResponseUsageRecord<'_>) {
}
pub(super) fn event_kind(route: &str) -> &'static str {
if route == "/v1/chat/completions" {
"v1.chat"
} else {
"v1.response"
}
let _ = route;
"v1.response"
}
pub(super) fn api_kind(route: &str) -> &'static str {
if route == "/v1/chat/completions" {
"chat"
} else {
"responses"
}
let _ = route;
"responses"
}
pub(super) fn enrich_event_route_metadata(metadata: &mut Value, route: &str) {
@@ -149,11 +140,11 @@ pub(super) fn enrich_usage_record_identity(
}
pub(super) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning.as_ref())
reasoning_effort_from_value(request.reasoning())
}
pub(super) fn reasoning_effort_from_compact_request(request: &CodexCompactRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning.as_ref())
reasoning_effort_from_value(request.reasoning())
}
fn insert_trimmed_string(
+1 -1
View File
@@ -1,7 +1,7 @@
//! OpenAI 兼容代理入口。
//!
//! 本模块负责:
//! - OpenAI API 协议转换(Chat Completions / Responses / Models)
//! - OpenAI API 协议转换(Responses / Models)
//! - 请求调度(账号选择、回退、恢复)
//! - 会话亲和性与 reasoning replay
File diff suppressed because it is too large Load Diff
+74 -180
View File
@@ -8,7 +8,7 @@ use axum::{
use serde_json::{json, Value};
use crate::proxy::{
dispatch::{chat::ChatDispatchError, responses::errors::ResponseDispatchError},
dispatch::{errors::DispatchFailureClass, responses::errors::ResponseDispatchError},
openai::responses::response_failed_sse_event as encode_response_failed_sse_event,
};
@@ -101,19 +101,9 @@ pub fn openai_error_response(
)
}
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()
}
pub fn chat_stream_dispatch_openai_error(error: &ResponseDispatchError) -> OpenAiErrorDetails {
response_dispatch_openai_error(error)
}
pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Response {
match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
ResponseDispatchError::NoActiveAccount => {
responses_no_available_accounts_response().into_response()
}
error => response_dispatch_openai_error_response(&error),
@@ -122,7 +112,7 @@ pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Respon
pub fn responses_compact_dispatch_error_response(error: ResponseDispatchError) -> Response {
match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
ResponseDispatchError::NoActiveAccount => {
responses_no_available_accounts_response().into_response()
}
error => {
@@ -203,11 +193,6 @@ pub fn model_not_found_response() -> (StatusCode, Json<Value>) {
)
}
/// Chat Completions 请求无效。
pub fn invalid_chat_completion_request_response() -> (StatusCode, Json<Value>) {
invalid_openai_request_response("Invalid chat completion request")
}
/// Responses 请求无效。
pub fn invalid_responses_request_response() -> (StatusCode, Json<Value>) {
invalid_openai_request_response("Invalid responses request")
@@ -222,88 +207,6 @@ fn invalid_openai_request_response(message: &str) -> (StatusCode, Json<Value>) {
)
}
fn chat_dispatch_openai_error(error: &ChatDispatchError) -> OpenAiErrorDetails {
let status = chat_dispatch_status(error);
let (message, kind) = match error {
ChatDispatchError::NoActiveAccount | ChatDispatchError::AccountStore => (
NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned(),
UPSTREAM_UNAVAILABLE_ERROR,
),
ChatDispatchError::Upstream(_) => (
UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
UPSTREAM_ERROR,
),
ChatDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "quota-exhausted", upstream_error, false),
INSUFFICIENT_QUOTA_ERROR,
),
ChatDispatchError::RateLimited {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "rate-limited", upstream_error, false),
RATE_LIMIT_ERROR,
),
ChatDispatchError::Expired {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "expired", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::Disabled {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "disabled", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::Banned {
count,
upstream_error,
..
} => (
exhausted_dispatch_message(*count, "banned", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "cloudflare-challenge", upstream_error, false),
UPSTREAM_ERROR,
),
ChatDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "cloudflare-path-block", upstream_error, false),
UPSTREAM_ERROR,
),
ChatDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "model-unsupported", upstream_error, false),
MODEL_NOT_FOUND_ERROR,
),
ChatDispatchError::InvalidSse(_) | ChatDispatchError::EmptyUpstreamResponse => (
INVALID_UPSTREAM_CODEX_RESPONSE_MESSAGE.to_owned(),
INVALID_UPSTREAM_RESPONSE_ERROR,
),
};
OpenAiErrorDetails {
status,
message,
error_type: kind.error_type,
code: kind.code,
}
}
fn response_dispatch_openai_error(error: &ResponseDispatchError) -> OpenAiErrorDetails {
let DispatchHttpError { status, message } = response_dispatch_http_error(
error,
@@ -321,27 +224,10 @@ fn response_dispatch_openai_error(error: &ResponseDispatchError) -> OpenAiErrorD
}
fn response_dispatch_error_kind(error: &ResponseDispatchError) -> OpenAiErrorKind {
match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
UPSTREAM_UNAVAILABLE_ERROR
}
ResponseDispatchError::QuotaExhausted { .. } => INSUFFICIENT_QUOTA_ERROR,
ResponseDispatchError::RateLimited { .. } => RATE_LIMIT_ERROR,
ResponseDispatchError::Expired { .. }
| ResponseDispatchError::Disabled { .. }
| ResponseDispatchError::Banned { .. } => AUTHENTICATION_ERROR,
ResponseDispatchError::ModelUnsupported { .. } => MODEL_NOT_FOUND_ERROR,
ResponseDispatchError::InvalidSse(_)
| ResponseDispatchError::MissingCompleted
| ResponseDispatchError::EmptyUpstreamResponse => INVALID_UPSTREAM_RESPONSE_ERROR,
ResponseDispatchError::Failed(_) => responses_error_kind_for_status(status_from_u16(
error.http_status_code(),
StatusCode::BAD_GATEWAY,
)),
ResponseDispatchError::Upstream(_)
| ResponseDispatchError::CloudflareChallenge { .. }
| ResponseDispatchError::CloudflarePathBlocked { .. } => UPSTREAM_ERROR,
}
dispatch_failure_openai_error_kind(
error.metadata().failure_class,
status_from_u16(error.http_status_code(), StatusCode::BAD_GATEWAY),
)
}
pub fn response_dispatch_http_error(
@@ -372,13 +258,6 @@ fn responses_error_kind_for_status(status: StatusCode) -> OpenAiErrorKind {
}
}
fn chat_dispatch_status(error: &ChatDispatchError) -> StatusCode {
match error {
ChatDispatchError::Upstream(_) => StatusCode::BAD_GATEWAY,
_ => status_from_u16(error.http_status_code(), StatusCode::BAD_GATEWAY),
}
}
fn response_dispatch_status(
error: &ResponseDispatchError,
mode: ResponseDispatchStatusMode,
@@ -395,60 +274,75 @@ fn response_dispatch_message(
error: &ResponseDispatchError,
style: ResponseDispatchMessageStyle,
) -> String {
match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned()
}
ResponseDispatchError::Upstream(_) => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
ResponseDispatchError::QuotaExhausted {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "quota-exhausted", upstream_error, false),
ResponseDispatchError::RateLimited {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "rate-limited", upstream_error, false),
ResponseDispatchError::Expired {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "expired", upstream_error, false),
ResponseDispatchError::Disabled {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "disabled", upstream_error, false),
ResponseDispatchError::Banned {
count,
upstream_error,
..
} => exhausted_dispatch_message(*count, "banned", upstream_error, false),
ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "cloudflare-challenge", upstream_error, false),
ResponseDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => exhausted_dispatch_message(
*count,
"cloudflare-path-block",
upstream_error,
matches!(style, ResponseDispatchMessageStyle::ResponsesStream),
),
ResponseDispatchError::ModelUnsupported {
count,
upstream_error,
} => exhausted_dispatch_message(
*count,
"model-unsupported",
upstream_error,
matches!(style, ResponseDispatchMessageStyle::ResponsesStream),
),
ResponseDispatchError::InvalidSse(_)
| ResponseDispatchError::MissingCompleted
| ResponseDispatchError::EmptyUpstreamResponse => {
let metadata = error.metadata();
if let Some(exhausted) = error.exhausted_account() {
return exhausted_dispatch_message(
exhausted.count,
exhausted.kind.message_reason(),
exhausted.upstream_error,
matches!(
(metadata.failure_class, style),
(
DispatchFailureClass::CloudflarePathBlocked
| DispatchFailureClass::ModelUnsupported,
ResponseDispatchMessageStyle::ResponsesStream,
)
),
);
}
dispatch_failure_message(
metadata.failure_class,
match style {
ResponseDispatchMessageStyle::Standard => DispatchFailureMessageStyle::Standard,
ResponseDispatchMessageStyle::ResponsesStream => {
DispatchFailureMessageStyle::ResponsesStream
}
},
)
}
#[derive(Clone, Copy)]
enum DispatchFailureMessageStyle {
Standard,
ResponsesStream,
}
fn dispatch_failure_message(
failure_class: DispatchFailureClass,
_style: DispatchFailureMessageStyle,
) -> String {
match failure_class {
DispatchFailureClass::NoAvailableAccounts => NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned(),
DispatchFailureClass::Upstream => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
DispatchFailureClass::InvalidSse
| DispatchFailureClass::MissingCompleted
| DispatchFailureClass::EmptyUpstreamResponse => {
INVALID_UPSTREAM_CODEX_RESPONSE_MESSAGE.to_owned()
}
ResponseDispatchError::Failed(_) => UPSTREAM_CODEX_RESPONSE_FAILED_MESSAGE.to_owned(),
DispatchFailureClass::ResponseFailed => UPSTREAM_CODEX_RESPONSE_FAILED_MESSAGE.to_owned(),
_ => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
}
}
fn dispatch_failure_openai_error_kind(
failure_class: DispatchFailureClass,
status: StatusCode,
) -> OpenAiErrorKind {
match failure_class {
DispatchFailureClass::NoAvailableAccounts => UPSTREAM_UNAVAILABLE_ERROR,
DispatchFailureClass::QuotaExhausted => INSUFFICIENT_QUOTA_ERROR,
DispatchFailureClass::RateLimited => RATE_LIMIT_ERROR,
DispatchFailureClass::Expired
| DispatchFailureClass::Disabled
| DispatchFailureClass::Banned => AUTHENTICATION_ERROR,
DispatchFailureClass::ModelUnsupported => MODEL_NOT_FOUND_ERROR,
DispatchFailureClass::InvalidSse
| DispatchFailureClass::MissingCompleted
| DispatchFailureClass::EmptyUpstreamResponse => INVALID_UPSTREAM_RESPONSE_ERROR,
DispatchFailureClass::ResponseFailed => responses_error_kind_for_status(status),
DispatchFailureClass::Upstream
| DispatchFailureClass::CloudflareChallenge
| DispatchFailureClass::CloudflarePathBlocked => UPSTREAM_ERROR,
}
}
-1
View File
@@ -1,6 +1,5 @@
//! OpenAI 兼容 API 模块。
pub mod chat;
pub mod errors;
pub mod models;
pub mod responses;
+194 -519
View File
@@ -1,4 +1,5 @@
//! OpenAI Responses API 类型、协议转换与 HTTP 处理器。
//! OpenAI Responses API 透明代理入口:解析原始 JSON body,提取调度元数据,
//! 做最小 patch 后交给调度层,请求语义原样透传上游。
use axum::{
body::{Body, Bytes},
@@ -7,20 +8,17 @@ use axum::{
response::{IntoResponse, Response},
Extension, Json,
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use std::net::IpAddr;
use crate::{
http::middleware::request_id::RequestId,
http::middleware::request_id::{ClientIp, RequestId},
proxy::{
auth::authorize_client_api_key,
dispatch::responses::{errors::ResponseDispatchError, service::ResponseDispatchStream},
},
runtime::state::AppState,
upstream::protocol::{
responses::{ensure_reasoning_include, CodexCompactRequest, CodexResponsesRequest},
schema::prepare_schema,
responses::{CodexCompactRequest, CodexResponsesRequest},
sse::{encode_sse_event, sse_body_has_done},
},
};
@@ -37,419 +35,162 @@ use super::{
const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent";
// ====================================================================
// 协议类型
// ====================================================================
/// 代理内部 transport 提示字段:用于本地传输决策,不转发上游。
const TRANSPORT_ONLY_KEYS: [&str; 1] = ["use_websocket"];
/// Responses API 请求体。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct OpenAiResponsesRequest {
pub model: String,
#[serde(default = "default_responses_input")]
pub input: Value,
#[serde(default)]
pub instructions: Option<String>,
#[serde(default)]
pub reasoning: Option<Value>,
#[serde(default)]
pub tools: Option<Value>,
#[serde(default)]
pub service_tier: Option<String>,
#[serde(default)]
pub tool_choice: Option<Value>,
#[serde(default)]
pub parallel_tool_calls: Option<bool>,
#[serde(default)]
pub text: Option<Value>,
#[serde(default)]
pub generate: Option<bool>,
#[serde(default)]
pub prompt_cache_key: Option<String>,
#[serde(default)]
pub include: Option<Value>,
#[serde(default)]
pub client_metadata: Option<Value>,
#[serde(default)]
pub previous_response_id: Option<String>,
#[serde(default, rename = "turnState", alias = "turn_state")]
pub turn_state: Option<String>,
#[serde(default, rename = "turnMetadata", alias = "turn_metadata")]
pub turn_metadata: Option<String>,
#[serde(default, rename = "betaFeatures", alias = "beta_features")]
pub beta_features: Option<String>,
#[serde(default)]
pub version: Option<String>,
#[serde(
default,
rename = "includeTimingMetrics",
alias = "include_timing_metrics"
)]
pub include_timing_metrics: Option<String>,
#[serde(default, rename = "codexWindowId", alias = "codex_window_id")]
pub codex_window_id: Option<String>,
#[serde(default, rename = "parentThreadId", alias = "parent_thread_id")]
pub parent_thread_id: Option<String>,
#[serde(default)]
pub use_websocket: Option<bool>,
#[serde(default = "default_responses_stream")]
pub stream: bool,
}
/// Codex 上下文字段:从 body 顶层(camelCase)提取到代理控制状态用于生成请求头,
/// 同时保留在 body 中原样转发上游。
const CONTEXT_HEADER_FIELDS: [(&str, ContextField); 7] = [
("turnState", ContextField::TurnState),
("turnMetadata", ContextField::TurnMetadata),
("betaFeatures", ContextField::BetaFeatures),
("version", ContextField::Version),
("includeTimingMetrics", ContextField::IncludeTimingMetrics),
("codexWindowId", ContextField::CodexWindowId),
("parentThreadId", ContextField::ParentThreadId),
];
fn default_responses_input() -> Value {
Value::Array(Vec::new())
}
fn default_responses_stream() -> bool {
true
#[derive(Clone, Copy)]
enum ContextField {
TurnState,
TurnMetadata,
BetaFeatures,
Version,
IncludeTimingMetrics,
CodexWindowId,
ParentThreadId,
}
// ====================================================================
// 协议转换
// 透明代理请求构造
// ====================================================================
/// 将 OpenAI Responses 请求转换为 Codex Responses 请求。
pub fn translate_response_to_codex(request: OpenAiResponsesRequest) -> CodexResponsesRequest {
let prepared_text = prepare_text_format(request.text, true);
let client_metadata = sanitize_client_metadata(request.client_metadata);
let mut codex_request = CodexResponsesRequest::new_http_sse(
request.model,
request.instructions.unwrap_or_default(),
sanitize_responses_input(request.input),
);
codex_request.previous_response_id = request.previous_response_id;
codex_request.turn_state = request.turn_state;
codex_request.turn_metadata = first_request_string(
request.turn_metadata,
client_metadata.as_ref(),
"x-codex-turn-metadata",
);
codex_request.beta_features = first_request_string(
request.beta_features,
client_metadata.as_ref(),
"x-codex-beta-features",
);
codex_request.include_timing_metrics = first_request_string(
request.include_timing_metrics,
client_metadata.as_ref(),
"x-responsesapi-include-timing-metrics",
);
codex_request.version = non_empty_string(request.version);
codex_request.codex_window_id = first_request_string(
request.codex_window_id,
client_metadata.as_ref(),
"x-codex-window-id",
);
codex_request.parent_thread_id = first_request_string(
request.parent_thread_id,
client_metadata.as_ref(),
"x-codex-parent-thread-id",
);
codex_request.reasoning = responses_reasoning(request.reasoning);
codex_request.tools = non_empty_array(request.tools);
codex_request.tool_choice = request.tool_choice;
codex_request.parallel_tool_calls = request.parallel_tool_calls;
codex_request.text = prepared_text.text;
codex_request.tuple_schema = prepared_text.tuple_schema;
codex_request.generate = request.generate;
codex_request.service_tier = non_empty_string(request.service_tier);
codex_request.prompt_cache_key = non_empty_string(request.prompt_cache_key);
codex_request.explicit_prompt_cache_key = codex_request.prompt_cache_key.is_some();
codex_request.include = string_array(request.include);
codex_request.client_metadata = client_metadata;
ensure_reasoning_include(&mut codex_request);
match request.use_websocket {
Some(true) => codex_request.use_websocket = true,
Some(false) => codex_request.force_http_sse = true,
/// 从客户端原始 Responses JSON body 构造上游请求(透明代理)。
///
/// 只做代理职责范围内的最小处理:
/// - 剥离 transport-only 字段(`use_websocket`),仅用于本地传输决策。
/// - 从 body 与请求头提取 Codex 上下文透传字段到代理控制状态(body 中保留原值)。
/// - review route / 合法请求头时,往 `client_metadata` 注入 forced subagent。
///
/// 其余字段——`input`、`reasoning`、`text`、`tools`、`tool_choice`、`include`、
/// `client_metadata`、`service_tier` 以及一切未知顶层字段——全部原样透传,不重写。
pub fn build_codex_request(
mut body: Map<String, Value>,
headers: &HeaderMap,
forced_subagent: Option<&str>,
) -> CodexResponsesRequest {
let use_websocket = body.get("use_websocket").and_then(Value::as_bool);
for key in TRANSPORT_ONLY_KEYS {
body.remove(key);
}
if let Some(subagent) = forced_subagent
.map(ToString::to_string)
.or_else(|| openai_subagent_from_headers(headers))
{
inject_subagent_metadata(&mut body, &subagent);
}
let explicit_prompt_cache_key = body
.get("prompt_cache_key")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty());
let mut request = CodexResponsesRequest::from_body(body);
request.explicit_prompt_cache_key = explicit_prompt_cache_key;
match use_websocket {
Some(true) => request.use_websocket = true,
Some(false) => request.force_http_sse = true,
None => {}
}
codex_request
apply_context_header_fields(&mut request, headers);
request.client_user_agent = user_agent_from_headers(headers);
request
}
/// 将 OpenAI Responses 请求转换为 Codex compact 请求。
pub fn translate_response_to_compact(request: OpenAiResponsesRequest) -> CodexCompactRequest {
/// 从客户端原始 Responses JSON body 构造 Codex compact 请求。
///
/// compact 端点只返回一次性 JSON(代理按非流式读取),故剥离 `stream` 这一
/// transport 控制字段,避免上游返回代理无法解析的 SSE 形态。业务语义字段
/// (`reasoning`/`text`/`store`/`prompt_cache_key`/`previous_response_id`/
/// `include`/`client_metadata`/未知字段)一律原样透传。
pub fn build_compact_request(
mut body: Map<String, Value>,
headers: &HeaderMap,
) -> CodexCompactRequest {
body.remove(COMPACT_STREAM_KEY);
for key in TRANSPORT_ONLY_KEYS {
body.remove(key);
}
CodexCompactRequest {
model: request.model,
input: sanitize_responses_input(request.input),
instructions: request.instructions.unwrap_or_default(),
tools: non_empty_array(request.tools),
parallel_tool_calls: request.parallel_tool_calls,
reasoning: compact_reasoning(request.reasoning),
text: prepare_text_format(request.text, false).text,
body,
client_ip: None,
client_user_agent: None,
client_user_agent: user_agent_from_headers(headers),
}
}
struct PreparedTextFormat {
text: Option<Value>,
tuple_schema: Option<Value>,
}
/// compact 端点不支持流式,转发前剥离 `stream`(transport 控制字段,非业务语义)。
const COMPACT_STREAM_KEY: &str = "stream";
fn prepare_text_format(text: Option<Value>, prepare_tuple_schema: bool) -> PreparedTextFormat {
let Some(Value::Object(text)) = text else {
return PreparedTextFormat {
text: None,
tuple_schema: None,
};
};
let Some(Value::Object(format)) = text.get("format") else {
return PreparedTextFormat {
text: None,
tuple_schema: None,
};
};
let Some(format_type) = format.get("type").and_then(Value::as_str) else {
return PreparedTextFormat {
text: None,
tuple_schema: None,
};
};
let mut sanitized_format = Map::new();
sanitized_format.insert("type".to_string(), Value::String(format_type.to_string()));
if let Some(name) = format.get("name").and_then(Value::as_str) {
sanitized_format.insert("name".to_string(), Value::String(name.to_string()));
fn inject_subagent_metadata(body: &mut Map<String, Value>, subagent: &str) {
let metadata = body
.entry("client_metadata")
.or_insert_with(|| Value::Object(Map::new()));
if !metadata.is_object() {
*metadata = Value::Object(Map::new());
}
let mut tuple_schema = None;
if let Some(Value::Object(schema)) = format.get("schema") {
let schema = Value::Object(schema.clone());
let schema = if prepare_tuple_schema {
let prepared = prepare_schema(schema);
tuple_schema = prepared.original_schema;
prepared.schema
} else {
schema
};
sanitized_format.insert("schema".to_string(), schema);
}
if let Some(strict) = format.get("strict").and_then(Value::as_bool) {
sanitized_format.insert("strict".to_string(), Value::Bool(strict));
}
PreparedTextFormat {
text: Some(json!({"format": sanitized_format})),
tuple_schema,
}
}
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 => vec![value],
}
}
fn sanitize_codex_input_items(input: Vec<Value>) -> Vec<Value> {
input
.into_iter()
.filter_map(|item| {
if let Value::String(text) = item {
return Some(responses_input_text_message(&text));
}
let Value::Object(object) = item else {
return Some(item);
};
match object.get("type").and_then(Value::as_str) {
Some("reasoning") => sanitize_reasoning_item(&object),
Some("compaction") => sanitize_compaction_item(&object),
_ => Some(Value::Object(object)),
}
})
.collect()
}
fn responses_input_text_message(text: &str) -> Value {
json!({
"type": "message",
"role": "user",
"content": [
{
"type": "input_text",
"text": text
}
]
})
}
fn sanitize_reasoning_item(item: &Map<String, Value>) -> Option<Value> {
let id = non_empty_str(item.get("id"))?;
let summary = sanitize_summary(item.get("summary"))?;
let mut sanitized = Map::new();
sanitized.insert("type".to_string(), Value::String("reasoning".to_string()));
sanitized.insert("id".to_string(), Value::String(id.to_string()));
sanitized.insert("summary".to_string(), Value::Array(summary));
if let Some(status) = item
.get("status")
.and_then(Value::as_str)
.filter(|status| matches!(*status, "in_progress" | "completed" | "incomplete"))
{
sanitized.insert("status".to_string(), Value::String(status.to_string()));
}
if let Some(encrypted_content) = non_empty_str(item.get("encrypted_content")) {
sanitized.insert(
"encrypted_content".to_string(),
Value::String(encrypted_content.to_string()),
if let Some(object) = metadata.as_object_mut() {
object.insert(
OPENAI_SUBAGENT_HEADER.to_string(),
Value::String(subagent.to_string()),
);
}
if let Some(content) = sanitize_reasoning_content(item.get("content")) {
sanitized.insert("content".to_string(), Value::Array(content));
}
/// 从 body 与请求头提取 Codex 上下文透传字段到代理控制状态。
///
/// body 中的原值优先;缺失时回退请求头。提取只填充代理控制状态用于加请求头,
/// **不修改 body**——这些字段在 body 中原样保留转发上游。
fn apply_context_header_fields(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
for (body_key, field) in CONTEXT_HEADER_FIELDS {
let value = body_context_string(request.body(), body_key)
.or_else(|| header_string(headers, field.header_name()));
field.assign(request, value);
}
Some(Value::Object(sanitized))
}
fn sanitize_summary(value: Option<&Value>) -> Option<Vec<Value>> {
let Value::Array(parts) = value? else {
return None;
};
Some(
parts
.iter()
.filter_map(|part| {
let Value::Object(part) = part else {
return None;
};
if part.get("type").and_then(Value::as_str) != Some("summary_text") {
return None;
}
let text = part.get("text").and_then(Value::as_str)?;
Some(json!({"type": "summary_text", "text": text}))
})
.collect(),
)
}
fn sanitize_reasoning_content(value: Option<&Value>) -> Option<Vec<Value>> {
let Value::Array(parts) = value? else {
return None;
};
let content = parts
.iter()
.filter_map(|part| {
let Value::Object(part) = part else {
return None;
};
if part.get("type").and_then(Value::as_str) != Some("reasoning_text") {
return None;
}
let text = part.get("text").and_then(Value::as_str)?;
Some(json!({"type": "reasoning_text", "text": text}))
})
.collect::<Vec<_>>();
(!content.is_empty()).then_some(content)
}
fn sanitize_compaction_item(item: &Map<String, Value>) -> Option<Value> {
let encrypted_content = non_empty_str(item.get("encrypted_content"))?;
let mut sanitized = Map::new();
sanitized.insert("type".to_string(), Value::String("compaction".to_string()));
sanitized.insert(
"encrypted_content".to_string(),
Value::String(encrypted_content.to_string()),
);
if let Some(id) = non_empty_str(item.get("id")) {
sanitized.insert("id".to_string(), Value::String(id.to_string()));
}
Some(Value::Object(sanitized))
}
fn responses_reasoning(reasoning: Option<Value>) -> Option<Value> {
let Value::Object(input) = reasoning? else {
return None;
};
let effort = input.get("effort").and_then(Value::as_str);
let summary = input
.get("summary")
fn body_context_string(body: &Map<String, Value>, key: &str) -> Option<String> {
body.get(key)
.and_then(Value::as_str)
.unwrap_or("auto");
let mut output = Map::new();
output.insert("summary".to_string(), Value::String(summary.to_string()));
if let Some(effort) = effort {
output.insert("effort".to_string(), Value::String(effort.to_string()));
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
impl ContextField {
fn header_name(self) -> &'static str {
match self {
Self::TurnState => "x-codex-turn-state",
Self::TurnMetadata => "x-codex-turn-metadata",
Self::BetaFeatures => "x-codex-beta-features",
Self::Version => "version",
Self::IncludeTimingMetrics => "x-responsesapi-include-timing-metrics",
Self::CodexWindowId => "x-codex-window-id",
Self::ParentThreadId => "x-codex-parent-thread-id",
}
}
Some(Value::Object(output))
}
fn compact_reasoning(reasoning: Option<Value>) -> Option<Value> {
let Value::Object(input) = reasoning? else {
return None;
};
let mut output = Map::new();
if let Some(effort) = input.get("effort").and_then(Value::as_str) {
output.insert("effort".to_string(), Value::String(effort.to_string()));
fn assign(self, request: &mut CodexResponsesRequest, value: Option<String>) {
match self {
Self::TurnState => request.turn_state = value,
Self::TurnMetadata => request.turn_metadata = value,
Self::BetaFeatures => request.beta_features = value,
Self::Version => request.version = value,
Self::IncludeTimingMetrics => request.include_timing_metrics = value,
Self::CodexWindowId => request.codex_window_id = value,
Self::ParentThreadId => request.parent_thread_id = value,
}
}
if let Some(summary) = input.get("summary").and_then(Value::as_str) {
output.insert("summary".to_string(), Value::String(summary.to_string()));
}
(!output.is_empty()).then_some(Value::Object(output))
}
fn sanitize_client_metadata(client_metadata: Option<Value>) -> Option<Value> {
let Value::Object(input) = client_metadata? else {
return None;
};
let metadata = input
.into_iter()
.filter_map(|(key, value)| {
value
.as_str()
.map(|value| (key, Value::String(value.to_string())))
})
.collect::<Map<_, _>>();
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
fn first_request_string(
direct: Option<String>,
client_metadata: Option<&Value>,
metadata_key: &str,
) -> Option<String> {
non_empty_string(direct).or_else(|| metadata_string(client_metadata, metadata_key))
}
fn metadata_string(client_metadata: Option<&Value>, key: &str) -> Option<String> {
client_metadata?
.as_object()?
.get(key)?
.as_str()
.and_then(|value| {
let trimmed = value.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
})
}
fn non_empty_array(value: Option<Value>) -> Option<Vec<Value>> {
let Value::Array(values) = value? else {
return None;
};
(!values.is_empty()).then_some(values)
}
fn string_array(value: Option<Value>) -> Option<Vec<String>> {
let Value::Array(values) = value? else {
return None;
};
values
.into_iter()
.map(|value| match value {
Value::String(value) => Some(value),
_ => None,
})
.collect()
}
fn non_empty_string(value: Option<String>) -> Option<String> {
value.and_then(|value| {
let trimmed = value.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
})
}
fn non_empty_str(value: Option<&Value>) -> Option<&str> {
let value = value?.as_str()?;
(!value.trim().is_empty()).then_some(value)
}
/// 编码 OpenAI Responses `response.failed` SSE 事件。
@@ -495,22 +236,34 @@ pub fn response_failed_sse_event_with_id(
pub async fn responses(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_responses(state, request_id, headers, body, "/v1/responses", None).await
}
/// `POST /v1/responses/review`
pub async fn review_responses(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
client_ip: Option<Extension<ClientIp>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_responses(
state,
request_id,
client_ip_string(client_ip),
headers,
body,
"/v1/responses",
None,
)
.await
}
/// `POST /v1/responses/review`
pub async fn review_responses(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
client_ip: Option<Extension<ClientIp>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_responses(
state,
request_id,
client_ip_string(client_ip),
headers,
body,
"/v1/responses/review",
@@ -523,15 +276,29 @@ pub async fn review_responses(
pub async fn compact_responses(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
client_ip: Option<Extension<ClientIp>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_compact_responses(state, request_id, headers, body).await
handle_compact_responses(
state,
request_id,
client_ip_string(client_ip),
headers,
body,
)
.await
}
/// 从可选的 `ClientIp` extension 提取 IP 字符串。
fn client_ip_string(client_ip: Option<Extension<ClientIp>>) -> Option<String> {
client_ip.map(|Extension(client_ip)| client_ip.as_str().to_string())
}
async fn handle_responses(
state: AppState,
request_id: RequestId,
client_ip: Option<String>,
headers: HeaderMap,
body: Bytes,
route: &str,
@@ -541,19 +308,10 @@ async fn handle_responses(
return missing_client_api_key_response().into_response();
}
let Ok(mut openai_request) = parse_openai_responses_request(&body) else {
let Some(body) = parse_responses_body(&body) else {
return invalid_responses_request_response().into_response();
};
apply_responses_header_context(&mut openai_request, &headers);
if openai_request.use_websocket.is_none() {
openai_request.use_websocket = Some(true);
}
if let Some(subagent) = forced_subagent {
force_openai_subagent(&mut openai_request, subagent);
} else if let Some(subagent) = openai_subagent_from_headers(&headers) {
force_openai_subagent(&mut openai_request, &subagent);
}
let model = match validated_responses_model(&state, &openai_request.model).await {
let model = match validated_responses_model(&state, request_model(&body)).await {
Ok(model) => model,
Err(ResponsesModelValidationError::InvalidRequest) => {
return invalid_responses_request_response().into_response();
@@ -562,9 +320,13 @@ async fn handle_responses(
return model_not_found_response().into_response();
}
};
let stream = openai_request.stream;
let mut codex_request = translate_response_to_codex(openai_request);
attach_client_context(&mut codex_request, &headers);
let mut codex_request = build_codex_request(body, &headers, forced_subagent);
codex_request.client_ip = client_ip;
// 客户端未显式指定 transport 时默认偏好 WebSocket。
if !codex_request.force_http_sse {
codex_request.use_websocket = true;
}
let stream = codex_request.stream();
if stream {
return match state
@@ -592,6 +354,7 @@ async fn handle_responses(
async fn handle_compact_responses(
state: AppState,
request_id: RequestId,
client_ip: Option<String>,
headers: HeaderMap,
body: Bytes,
) -> Response {
@@ -599,10 +362,10 @@ async fn handle_compact_responses(
return missing_client_api_key_response().into_response();
}
let Ok(openai_request) = parse_openai_responses_request(&body) else {
let Some(body) = parse_responses_body(&body) else {
return invalid_responses_request_response().into_response();
};
let model = match validated_responses_model(&state, &openai_request.model).await {
let model = match validated_responses_model(&state, request_model(&body)).await {
Ok(model) => model,
Err(ResponsesModelValidationError::InvalidRequest) => {
return invalid_responses_request_response().into_response();
@@ -611,9 +374,8 @@ async fn handle_compact_responses(
return model_not_found_response().into_response();
}
};
let mut compact_request = translate_response_to_compact(openai_request);
compact_request.client_ip = client_ip_from_headers(&headers);
compact_request.client_user_agent = user_agent_from_headers(&headers);
let mut compact_request = build_compact_request(body, &headers);
compact_request.client_ip = client_ip;
match state
.services
@@ -626,22 +388,18 @@ async fn handle_compact_responses(
}
}
fn force_openai_subagent(request: &mut OpenAiResponsesRequest, subagent: &str) {
let mut metadata = match request.client_metadata.take() {
Some(Value::Object(metadata)) => metadata,
_ => serde_json::Map::new(),
};
metadata.insert(
OPENAI_SUBAGENT_HEADER.to_string(),
Value::String(subagent.to_string()),
);
request.client_metadata = Some(Value::Object(metadata));
/// 解析请求体:必须是 JSON object,否则视为 invalid request。
fn parse_responses_body(body: &Bytes) -> Option<Map<String, Value>> {
match serde_json::from_slice::<Value>(body) {
Ok(Value::Object(map)) => Some(map),
_ => None,
}
}
fn parse_openai_responses_request(
body: &Bytes,
) -> Result<OpenAiResponsesRequest, serde_json::Error> {
serde_json::from_slice(body)
fn request_model(body: &Map<String, Value>) -> &str {
body.get("model")
.and_then(Value::as_str)
.unwrap_or_default()
}
async fn validated_responses_model(
@@ -665,93 +423,10 @@ enum ResponsesModelValidationError {
ModelNotFound,
}
fn apply_responses_header_context(request: &mut OpenAiResponsesRequest, headers: &HeaderMap) {
fill_string_from_header(&mut request.turn_state, headers, "x-codex-turn-state");
fill_string_from_header(&mut request.turn_metadata, headers, "x-codex-turn-metadata");
fill_string_from_header(&mut request.beta_features, headers, "x-codex-beta-features");
fill_string_from_header(
&mut request.include_timing_metrics,
headers,
"x-responsesapi-include-timing-metrics",
);
fill_string_from_header(&mut request.version, headers, "version");
fill_string_from_header(&mut request.codex_window_id, headers, "x-codex-window-id");
fill_string_from_header(
&mut request.parent_thread_id,
headers,
"x-codex-parent-thread-id",
);
}
fn fill_string_from_header(field: &mut Option<String>, headers: &HeaderMap, name: &str) {
if field
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
{
return;
}
*field = header_string(headers, name);
}
pub(crate) fn attach_client_context(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
request.client_ip = client_ip_from_headers(headers);
request.client_user_agent = user_agent_from_headers(headers);
}
pub(crate) fn user_agent_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, "user-agent")
}
pub(crate) fn client_ip_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, "cf-connecting-ip")
.or_else(|| header_string(headers, "x-real-ip"))
.or_else(|| forwarded_client_ip(headers))
.and_then(|value| normalize_client_ip(&value))
}
fn forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
let value = header_string(headers, "x-forwarded-for")?;
let parts = value
.split(',')
.map(str::trim)
.filter(|part| !part.is_empty())
.collect::<Vec<_>>();
parts
.iter()
.find(|part| !is_private_ip(part))
.or_else(|| parts.first())
.map(|part| (*part).to_string())
}
fn normalize_client_ip(value: &str) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
return None;
}
trimmed
.parse::<std::net::SocketAddr>()
.map(|addr| addr.ip().to_string())
.ok()
.or_else(|| Some(trimmed.to_string()))
}
fn is_private_ip(value: &str) -> bool {
let Some(ip) = normalize_client_ip(value).and_then(|value| value.parse::<IpAddr>().ok()) else {
return false;
};
match ip {
IpAddr::V4(ip) => {
ip.is_private()
|| ip.is_loopback()
|| ip.is_link_local()
|| ip.octets()[0] == 100 && (64..=127).contains(&ip.octets()[1])
}
IpAddr::V6(ip) => ip.is_loopback() || ip.is_unique_local() || ip.is_unicast_link_local(),
}
}
fn openai_subagent_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, OPENAI_SUBAGENT_HEADER).and_then(|value| {
if matches!(
+2 -2
View File
@@ -1,6 +1,7 @@
//! OpenAI 兼容 API 路由。
use axum::{
extract::DefaultBodyLimit,
routing::{get, post},
Router,
};
@@ -8,7 +9,6 @@ use axum::{
use crate::runtime::state::AppState;
use super::{
chat::chat_completions,
models::{model_catalog, model_detail, model_info, models},
responses::{compact_responses, responses, review_responses},
};
@@ -19,9 +19,9 @@ pub fn router() -> Router<AppState> {
.route("/v1/responses", post(responses))
.route("/v1/responses/review", post(review_responses))
.route("/v1/responses/compact", post(compact_responses))
.route("/v1/chat/completions", post(chat_completions))
.route("/v1/models", get(models))
.route("/v1/models/catalog", get(model_catalog))
.route("/v1/models/{model_id}", get(model_detail))
.route("/v1/models/{model_id}/info", get(model_info))
.layer(DefaultBodyLimit::disable())
}
+4 -5
View File
@@ -5,7 +5,6 @@ use chrono::Utc;
use crate::config::schema::AppConfig;
use crate::config::settings::RuntimeSettingsService;
use crate::http::middleware::request_id::TrustedProxyConfig;
use crate::infra::database::connect_sqlite;
use crate::infra::logging::{init_tracing, LogGuard, RotationConfig};
use crate::infra::paths::{ensure_data_dir, load_or_create_installation_id};
@@ -13,7 +12,7 @@ use crate::runtime::services::{BackgroundTaskStores, Services};
use crate::runtime::shutdown::{restart_executable_path, shutdown_signal};
use crate::runtime::state::{AppState, RuntimeConfig};
use crate::runtime::tasks::coordinator::TaskCoordinator;
use crate::upstream::fingerprint::{Fingerprint, FingerprintRepository};
use crate::upstream::fingerprint::{Fingerprint, FingerprintRepository, RuntimeFingerprint};
const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS";
@@ -81,6 +80,7 @@ async fn build_application(
let runtime_fingerprint = fingerprint_repository
.ensure_current_seed(&default_fingerprint)
.await?;
let runtime_fingerprint = RuntimeFingerprint::new(runtime_fingerprint);
let runtime_config = RuntimeConfig::from(config.clone());
let services = Services::try_with_installation_id(
&config,
@@ -112,11 +112,10 @@ async fn build_application(
);
let task_coordinator = TaskCoordinator::start(&runtime_config, &services);
let trusted_proxies = TrustedProxyConfig::from_entries(&config.server.trusted_proxies)?;
let app_state = AppState { config, services };
let app_state = AppState { services };
Ok((
crate::http::router::router_with_trusted_proxies(trusted_proxies).with_state(app_state),
crate::http::router::router().with_state(app_state),
task_coordinator,
))
}
+33 -31
View File
@@ -6,11 +6,10 @@ use crate::{
admin::monitoring::{
account_usage_service::AdminUsageService,
account_usage_store::SqliteUsageStore,
ops_error_service::AdminOpsErrorLogService,
ops_error_store::SqliteOpsErrorLogStore,
usage_record_service::AdminUsageRecordService,
usage_record_store::{
SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPACITY,
DEFAULT_USAGE_RECORD_CAPTURE_BODY,
},
usage_record_store::{SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPTURE_BODY},
},
admin::{
accounts::service::{AdminAccountService, AdminAccountServiceParts},
@@ -25,7 +24,6 @@ use crate::{
settings::{account_pool_options_from_config, RuntimeSettingsService},
},
proxy::dispatch::{
chat::ChatDispatchService,
responses::service::{ResponseDispatchService, ResponseDispatchServiceParts},
session_affinity::{RuntimeSessionAffinityService, SqliteSessionAffinityStore},
},
@@ -39,7 +37,7 @@ use crate::{
},
},
upstream::{
fingerprint::{Fingerprint, FingerprintRepository},
fingerprint::{FingerprintRepository, RuntimeFingerprint},
models::{
config::ModelConfig,
service::{ModelService, ModelServiceError},
@@ -95,15 +93,15 @@ pub struct Services {
pub refresh_policy: RuntimeRefreshPolicy,
pub admin_accounts: StdArc<AdminAccountService>,
pub usage_records: StdArc<AdminUsageRecordService>,
pub ops_errors: StdArc<AdminOpsErrorLogService>,
pub usage: StdArc<AdminUsageService>,
pub account_pool: StdArc<RuntimeAccountPoolService>,
pub token_refresh: StdArc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
pub chat: StdArc<ChatDispatchService>,
pub responses: StdArc<ResponseDispatchService>,
pub session_affinity: StdArc<RuntimeSessionAffinityService>,
pub codex: StdArc<CodexBackendClient>,
pub websocket_pool: Option<StdArc<CodexWebSocketPool>>,
pub fingerprint: Fingerprint,
pub fingerprint: RuntimeFingerprint,
pub installation_id: Option<String>,
pub background_tasks: BackgroundTaskStores,
}
@@ -112,7 +110,6 @@ pub struct Services {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UsageRecordOptions {
pub enabled: bool,
pub capacity: u32,
pub capture_body: bool,
}
@@ -120,14 +117,17 @@ impl UsageRecordOptions {
pub fn from_config(config: &AppConfig) -> Self {
Self {
enabled: config.logging.enabled,
capacity: DEFAULT_USAGE_RECORD_CAPACITY,
capture_body: DEFAULT_USAGE_RECORD_CAPTURE_BODY,
}
}
}
impl Services {
pub fn new(config: &AppConfig, stores: BackgroundTaskStores, fingerprint: Fingerprint) -> Self {
pub fn new(
config: &AppConfig,
stores: BackgroundTaskStores,
fingerprint: RuntimeFingerprint,
) -> Self {
Self::try_new(config, stores, fingerprint)
.expect("failed to build runtime services with configured TLS transport")
}
@@ -135,7 +135,7 @@ impl Services {
pub fn try_new(
config: &AppConfig,
stores: BackgroundTaskStores,
fingerprint: Fingerprint,
fingerprint: RuntimeFingerprint,
) -> Result<Self, CustomCaError> {
Self::try_with_installation_id(config, stores, fingerprint, None)
}
@@ -143,7 +143,7 @@ impl Services {
pub(crate) fn try_with_installation_id(
config: &AppConfig,
stores: BackgroundTaskStores,
fingerprint: Fingerprint,
fingerprint: RuntimeFingerprint,
installation_id: Option<String>,
) -> Result<Self, CustomCaError> {
Self::try_with_installation_id_and_usage_record_options(
@@ -158,7 +158,7 @@ impl Services {
pub fn try_with_installation_id_and_usage_record_options(
config: &AppConfig,
stores: BackgroundTaskStores,
fingerprint: Fingerprint,
fingerprint: RuntimeFingerprint,
installation_id: Option<String>,
usage_record_options: UsageRecordOptions,
) -> Result<Self, CustomCaError> {
@@ -170,21 +170,36 @@ impl Services {
enabled: true,
max_age: std::time::Duration::from_millis(config.ws_pool.max_age_ms),
max_per_account: config.ws_pool.max_per_account,
first_token_timeout: (config.ws_pool.first_token_timeout_ms > 0).then(|| {
std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms)
}),
..CodexWebSocketPoolConfig::default()
}))
});
let codex = {
let websocket_first_token_timeout = (config.ws_pool.first_token_timeout_ms > 0)
.then(|| std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms));
let client = CodexBackendClient::new(
build_reqwest_client(config.tls.force_http11)?,
config.api.base_url.clone(),
fingerprint.clone(),
);
)
.with_websocket_first_token_timeout(websocket_first_token_timeout);
if let Some(pool) = &websocket_pool {
StdArc::new(client.with_websocket_pool(pool.clone()))
} else {
StdArc::new(client)
}
};
let usage_records = StdArc::new(AdminUsageRecordService::new(
stores.usage_records.clone(),
usage_record_options.enabled,
usage_record_options.capture_body,
));
let ops_errors = StdArc::new(AdminOpsErrorLogService::new(
SqliteOpsErrorLogStore::new(stores.usage_records.pool().clone()),
usage_record_options.capture_body,
));
let account_pool = StdArc::new(RuntimeAccountPoolService::new(
account_store_trait.clone(),
account_pool_options_from_config(config),
@@ -239,6 +254,7 @@ impl Services {
models: models.clone(),
account_pool: account_pool.clone(),
token_refresher: StdArc::new(token_client),
refresh_leases: stores.refresh_leases.clone(),
oauth: crate::admin::accounts::service::oauth::AccountOAuthService::new(
reqwest::Client::new(),
config.auth.oauth_client_id.clone(),
@@ -247,12 +263,6 @@ impl Services {
refresh_policy: refresh_policy.clone(),
installation_id: installation_id.clone(),
}));
let usage_records = StdArc::new(AdminUsageRecordService::new(
stores.usage_records.clone(),
usage_record_options.enabled,
usage_record_options.capacity,
usage_record_options.capture_body,
));
let usage_store = SqliteUsageStore::new(stores.accounts.pool().clone());
let usage = StdArc::new(AdminUsageService::new(usage_store));
let session_affinity = StdArc::new(RuntimeSessionAffinityService::new(
@@ -261,21 +271,13 @@ impl Services {
let cloudflare_recovery =
crate::proxy::dispatch::cloudflare::CloudflareRecovery::new(stores.cookies.clone());
let chat = StdArc::new(ChatDispatchService::new(
account_pool.clone(),
models.clone(),
codex.clone(),
usage_records.clone(),
token_refresh.clone(),
installation_id.clone(),
cloudflare_recovery.clone(),
));
let responses = StdArc::new(ResponseDispatchService::new(ResponseDispatchServiceParts {
account_pool: account_pool.clone(),
models: models.clone(),
codex: codex.clone(),
session_affinity: session_affinity.clone(),
usage_records: usage_records.clone(),
ops_errors: ops_errors.clone(),
token_refresh: token_refresh.clone(),
installation_id: installation_id.clone(),
cloudflare: cloudflare_recovery,
@@ -291,10 +293,10 @@ impl Services {
refresh_policy,
admin_accounts,
usage_records,
ops_errors,
usage,
account_pool,
token_refresh,
chat,
responses,
session_affinity,
codex,
-1
View File
@@ -6,7 +6,6 @@ use crate::runtime::services::Services;
/// HTTP handler 共享的运行时状态。
#[derive(Clone)]
pub struct AppState {
pub config: AppConfig,
pub services: Services,
}
+32 -36
View File
@@ -1,12 +1,14 @@
//! 通用过期数据清理任务骨架。
use std::{fmt::Display, future::Future, time::Duration};
use std::{fmt::Display, future::Future};
use chrono::{DateTime, Utc};
use tokio::time::interval;
use tracing::{debug, info, warn};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// 支持按时间清理过期记录的存储。
pub(crate) trait ExpiredCleanupStore: Send + Sync + 'static {
@@ -47,39 +49,12 @@ where
}
pub(crate) fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(
interval_secs = self.interval_secs,
"{}", self.messages.started
);
loop {
tokio::select! {
_ = ticker.tick() => {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "{}", self.messages.deleted);
}
Ok(_) => {
debug!("{}", self.messages.empty);
}
Err(error) => {
warn!(error = %error, "{}", self.messages.failed);
}
}
}
_ = shutdown_rx.recv() => {
info!("{}", self.messages.stopped);
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
let config = PeriodicTaskConfig::new(
self.interval_secs,
self.messages.started,
self.messages.stopped,
);
spawn_periodic_task(self, config)
}
pub(crate) async fn cleanup_once(&self) -> Result<u64, S::Error> {
@@ -90,3 +65,24 @@ where
self.store.delete_expired_at(now).await
}
}
impl<S> PeriodicTaskRunner for CleanupTask<S>
where
S: ExpiredCleanupStore,
{
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "{}", self.messages.deleted);
}
Ok(_) => {
debug!("{}", self.messages.empty);
}
Err(error) => {
warn!(error = %error, "{}", self.messages.failed);
}
}
})
}
}
+52 -22
View File
@@ -1,6 +1,6 @@
//! 后台任务协调器。
use std::{path::PathBuf, sync::Arc};
use std::{path::PathBuf, sync::Arc, time::Duration};
use tokio::task::JoinHandle;
@@ -86,16 +86,17 @@ impl TaskCoordinator {
.with_account_pool(services.account_pool.clone())
.start(),
);
coordinator.push(
"fingerprint_update",
coordinator.push("fingerprint_update", {
let fingerprint = services.fingerprint.snapshot();
start_fingerprint_update_task(
Some(stores.fingerprints.clone()),
stores.fingerprints.clone(),
services.fingerprint.clone(),
CODEX_DESKTOP_APPCAST_URL.to_string(),
PathBuf::from(DEFAULT_EXTRACTED_FINGERPRINT_PATH),
services.fingerprint.app_version.clone(),
services.fingerprint.build_number.clone(),
),
);
fingerprint.app_version,
fingerprint.build_number,
)
});
if let Some(pool) = &services.websocket_pool {
coordinator.push(
"websocket_pool",
@@ -120,22 +121,21 @@ impl TaskCoordinator {
/// 后台任务关闭句柄。
pub enum SchedulerHandle {
/// 通过 channel 发送关闭信号。
Channel(tokio::sync::mpsc::Sender<()>),
/// 直接持有 tokio 任务句柄。
JoinHandle(JoinHandle<()>),
Channel {
shutdown_tx: tokio::sync::mpsc::Sender<()>,
handle: JoinHandle<()>,
},
/// 关闭上游 WebSocket 连接池。
WebSocketPool(Arc<crate::upstream::transport::CodexWebSocketPool>),
}
impl SchedulerHandle {
/// 使用关闭 channel 构造句柄。
pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self {
Self::Channel(shutdown_tx)
}
/// 使用 `JoinHandle` 构造句柄。
pub(crate) fn from_join_handle(handle: JoinHandle<()>) -> Self {
Self::JoinHandle(handle)
pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>, handle: JoinHandle<()>) -> Self {
Self::Channel {
shutdown_tx,
handle,
}
}
/// 使用 WebSocket 连接池构造句柄。
@@ -146,11 +146,12 @@ impl SchedulerHandle {
/// 关闭任务。
pub async fn shutdown(self) {
match self {
Self::Channel(shutdown_tx) => {
Self::Channel {
shutdown_tx,
handle,
} => {
let _ = shutdown_tx.send(()).await;
}
Self::JoinHandle(handle) => {
handle.abort();
wait_for_task_shutdown(handle).await;
}
Self::WebSocketPool(pool) => {
pool.shutdown().await;
@@ -158,3 +159,32 @@ impl SchedulerHandle {
}
}
}
const TASK_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
async fn wait_for_task_shutdown(handle: JoinHandle<()>) {
let mut handle = handle;
let timeout = tokio::time::sleep(TASK_SHUTDOWN_TIMEOUT);
tokio::pin!(timeout);
tokio::select! {
result = &mut handle => {
log_task_shutdown_result(result);
}
() = &mut timeout => {
tracing::warn!(timeout_secs = TASK_SHUTDOWN_TIMEOUT.as_secs(), "等待后台任务关闭超时");
handle.abort();
log_task_shutdown_result(handle.await);
}
}
}
fn log_task_shutdown_result(result: Result<(), tokio::task::JoinError>) {
match result {
Ok(()) => {}
Err(error) if error.is_cancelled() => {}
Err(error) => {
tracing::warn!(error = %error, "后台任务关闭时异常退出");
}
}
}
@@ -2,9 +2,16 @@
use std::path::PathBuf;
use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker};
use tracing::{info, warn};
use super::coordinator::SchedulerHandle;
use crate::upstream::fingerprint::{
FingerprintRepository, RuntimeFingerprint, UpdateChecker, APPCAST_POLL_INTERVAL,
};
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// Codex Desktop 官方 appcast 地址。
pub const CODEX_DESKTOP_APPCAST_URL: &str =
@@ -14,7 +21,8 @@ pub const DEFAULT_EXTRACTED_FINGERPRINT_PATH: &str = "data/extracted-fingerprint
/// 启动指纹自动更新后台任务。
pub fn start_fingerprint_update_task(
repository: Option<FingerprintRepository>,
repository: FingerprintRepository,
runtime_fingerprint: RuntimeFingerprint,
appcast_url: String,
extracted_fingerprint_path: PathBuf,
current_version: String,
@@ -28,5 +36,50 @@ pub fn start_fingerprint_update_task(
current_version,
current_build,
);
SchedulerHandle::from_join_handle(update_checker.start_background_checker())
spawn_periodic_task(
FingerprintUpdateTask {
checker: update_checker,
runtime_fingerprint,
first_tick: true,
},
PeriodicTaskConfig::new(
APPCAST_POLL_INTERVAL.as_secs(),
"fingerprint 后台版本检查器已启动",
"fingerprint 后台版本检查器已关闭",
),
)
}
struct FingerprintUpdateTask {
checker: UpdateChecker,
runtime_fingerprint: RuntimeFingerprint,
first_tick: bool,
}
impl PeriodicTaskRunner for FingerprintUpdateTask {
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
let first_tick = std::mem::take(&mut self.first_tick);
match self.checker.check_and_apply_update().await {
Ok(Some(fingerprint)) => {
let app_version = fingerprint.app_version.clone();
let build_number = fingerprint.build_number.clone();
self.runtime_fingerprint.replace(fingerprint);
info!(
app_version = %app_version,
build_number = %build_number,
"fingerprint 运行时版本已更新"
);
}
Ok(None) => {}
Err(error) => {
if first_tick {
warn!(error = %error, "fingerprint 首次版本检查失败");
} else {
warn!(error = %error, "fingerprint 定期版本检查失败");
}
}
}
})
}
}
+1
View File
@@ -5,6 +5,7 @@ pub mod cookie_cleanup;
pub mod coordinator;
pub mod fingerprint_update;
pub mod model_refresh;
mod periodic;
pub mod quota_refresh;
pub mod session_affinity_cleanup;
pub mod session_cleanup;
+70 -64
View File
@@ -3,7 +3,8 @@
use std::{sync::Arc, time::Duration};
use chrono::Utc;
use tokio::time::{interval, sleep};
use tokio::sync::mpsc;
use tokio::time::sleep;
use tracing::{info, warn};
use crate::upstream::accounts::pool::RuntimeAccountPoolService;
@@ -11,7 +12,10 @@ use crate::upstream::models::service::{
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup},
};
/// 模型刷新任务接线器。
pub struct ModelRefreshTask {
@@ -54,25 +58,68 @@ impl ModelRefreshTask {
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
let config = PeriodicTaskConfig::new(
self.refresh_interval_secs,
"模型刷新任务已启动",
"模型刷新任务正在关闭",
)
.wait_first_interval();
spawn_periodic_task(self, config)
}
tokio::spawn(async move {
info!("模型刷新任务已启动");
/// 执行一次模型刷新并同步账号池模型计划约束。
pub async fn refresh_once(&self) -> Result<ModelRefreshResult, ModelServiceError> {
let request_id = uuid::Uuid::new_v4().to_string();
let selected_accounts = self.account_pool.distinct_plan_accounts(Utc::now()).await;
let plan_accounts = selected_accounts
.iter()
.map(|selected| ModelRefreshPlanAccount {
plan_type: selected.plan_type.clone(),
account: selected.account.clone(),
})
.collect::<Vec<_>>();
sleep(Duration::from_millis(self.initial_delay_ms)).await;
let result = self
.model_service
.refresh_backend_models_with_installation_id(
&plan_accounts,
&request_id,
self.installation_id.as_deref(),
)
.await;
for selected in &selected_accounts {
self.account_pool
.release_without_request_usage(&selected.account.id)
.await;
}
let result = result?;
let routing = self.model_service.model_plan_routing().await?;
self.account_pool
.apply_model_plan_routing(routing.allowlist, routing.fetched_plan_types)
.await;
Ok(result)
}
}
impl PeriodicTaskRunner for ModelRefreshTask {
fn before_loop<'a>(
&'a mut self,
shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(async move {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
}
let mut attempt = 0;
let mut has_fetched_once = false;
while attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return;
}
() = async {} => {}
}
match self.refresh_once().await {
Ok(result) => {
info!(
@@ -101,7 +148,7 @@ impl ModelRefreshTask {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return;
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
@@ -114,56 +161,15 @@ impl ModelRefreshTask {
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
}
let mut ticker = interval(Duration::from_secs(self.refresh_interval_secs));
ticker.tick().await;
loop {
tokio::select! {
_ = ticker.tick() => {
if let Err(error) = self.refresh_once().await {
warn!(error = ?error, "刷新模型列表失败");
}
}
_ = shutdown_rx.recv() => {
info!("模型刷新任务正在关闭");
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
PeriodicTaskStartup::Continue
})
}
/// 执行一次模型刷新并同步账号池模型计划约束。
pub async fn refresh_once(&self) -> Result<ModelRefreshResult, ModelServiceError> {
let request_id = uuid::Uuid::new_v4().to_string();
let selected_accounts = self.account_pool.distinct_plan_accounts(Utc::now()).await;
let plan_accounts = selected_accounts
.iter()
.map(|selected| ModelRefreshPlanAccount {
plan_type: selected.plan_type.clone(),
account: selected.account.clone(),
})
.collect::<Vec<_>>();
let result = self
.model_service
.refresh_backend_models_with_installation_id(
&plan_accounts,
&request_id,
self.installation_id.as_deref(),
)
.await;
for selected in &selected_accounts {
self.account_pool.release(&selected.account.id).await;
}
let result = result?;
let routing = self.model_service.model_plan_routing().await?;
self.account_pool
.apply_model_plan_routing(routing.allowlist, routing.fetched_plan_types)
.await;
Ok(result)
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
if let Err(error) = self.refresh_once().await {
warn!(error = ?error, "刷新模型列表失败");
}
})
}
}
+105
View File
@@ -0,0 +1,105 @@
//! 周期性后台任务运行器。
use std::{future::Future, pin::Pin, time::Duration};
use tokio::sync::mpsc;
use tokio::time::interval;
use tracing::info;
use super::coordinator::SchedulerHandle;
pub(crate) type TaskFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
/// 周期任务进入循环前的决策。
pub(crate) enum PeriodicTaskStartup {
/// 继续进入周期循环。
Continue,
/// 任务已自行结束,不再进入周期循环。
Stop,
}
/// 周期任务调度配置。
pub(crate) struct PeriodicTaskConfig {
interval_secs: u64,
started_message: &'static str,
stopped_message: &'static str,
run_immediately: bool,
}
impl PeriodicTaskConfig {
pub(crate) fn new(
interval_secs: u64,
started_message: &'static str,
stopped_message: &'static str,
) -> Self {
Self {
interval_secs,
started_message,
stopped_message,
run_immediately: true,
}
}
pub(crate) fn wait_first_interval(mut self) -> Self {
self.run_immediately = false;
self
}
}
/// 周期任务行为。
pub(crate) trait PeriodicTaskRunner: Send + 'static {
fn before_loop<'a>(
&'a mut self,
_shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(std::future::ready(PeriodicTaskStartup::Continue))
}
fn tick(&mut self) -> TaskFuture<'_, ()>;
fn shutdown(&mut self) -> TaskFuture<'_, ()> {
Box::pin(std::future::ready(()))
}
}
/// 启动周期任务,并统一处理关闭信号与 JoinHandle。
pub(crate) fn spawn_periodic_task<T>(mut task: T, config: PeriodicTaskConfig) -> SchedulerHandle
where
T: PeriodicTaskRunner,
{
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
let handle = tokio::spawn(async move {
info!(
interval_secs = config.interval_secs,
"{}", config.started_message
);
if matches!(
task.before_loop(&mut shutdown_rx).await,
PeriodicTaskStartup::Stop
) {
return;
}
let mut ticker = interval(Duration::from_secs(config.interval_secs));
if !config.run_immediately {
ticker.tick().await;
}
loop {
tokio::select! {
_ = ticker.tick() => {
task.tick().await;
}
_ = shutdown_rx.recv() => {
task.shutdown().await;
info!("{}", config.stopped_message);
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx, handle)
}
+39 -36
View File
@@ -1,8 +1,8 @@
//! 配额刷新后台任务接线器。
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{collections::HashMap, sync::Arc};
use tokio::time::interval;
use tokio::time::Instant;
use tracing::{debug, info, warn};
use crate::{
@@ -15,12 +15,16 @@ use crate::{
upstream::transport::CodexBackendClient,
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// 主动配额刷新后台任务。
pub struct QuotaRefreshTask {
service: RuntimeQuotaRefreshService,
interval_secs: u64,
last_refreshed: HashMap<String, Instant>,
}
impl QuotaRefreshTask {
@@ -38,6 +42,7 @@ impl QuotaRefreshTask {
min_refresh_interval_secs,
),
interval_secs,
last_refreshed: HashMap::new(),
}
}
@@ -61,42 +66,40 @@ impl QuotaRefreshTask {
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
let config = PeriodicTaskConfig::new(
self.interval_secs,
"quota 刷新任务已启动",
"quota 刷新任务已关闭",
);
spawn_periodic_task(self, config)
}
}
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
let mut last_refreshed = HashMap::new();
info!(interval_secs = self.interval_secs, "quota 刷新任务已启动");
loop {
tokio::select! {
_ = ticker.tick() => {
match self.service.refresh_locked_accounts(&mut last_refreshed).await {
Ok(summary) if summary.refreshed > 0 => {
log_refreshed_summary(summary);
}
Ok(summary) => {
debug!(
candidates = summary.candidates,
skipped_recent = summary.skipped_recent,
failed = summary.failed,
"没有需要刷新的 quota 锁定或待验证账号"
);
}
Err(error) => {
warn!(error = %error, "quota 刷新任务失败");
}
}
}
_ = shutdown_rx.recv() => {
info!("quota 刷新任务已关闭");
break;
}
impl PeriodicTaskRunner for QuotaRefreshTask {
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self
.service
.refresh_locked_accounts(&mut self.last_refreshed)
.await
{
Ok(summary) if summary.refreshed > 0 => {
log_refreshed_summary(summary);
}
Ok(summary) => {
debug!(
candidates = summary.candidates,
skipped_cooldown = summary.skipped_cooldown,
skipped_recent = summary.skipped_recent,
failed = summary.failed,
"没有需要刷新的 quota 锁定或待验证账号"
);
}
Err(error) => {
warn!(error = %error, "quota 刷新任务失败");
}
}
});
SchedulerHandle::new(shutdown_tx)
})
}
}
+43 -41
View File
@@ -1,8 +1,5 @@
//! 令牌刷新后台任务接线器。
use std::time::Duration;
use tokio::time::interval;
use tracing::{debug, info, warn};
use crate::upstream::accounts::{
@@ -12,7 +9,10 @@ use crate::upstream::accounts::{
},
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
const DEFAULT_INTERVAL_SECS: u64 = 60;
@@ -53,46 +53,48 @@ where
/// 启动后台任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
let config = PeriodicTaskConfig::new(
self.interval_secs,
"token 刷新任务已启动",
"token 刷新任务已关闭",
);
spawn_periodic_task(self, config)
}
}
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(interval_secs = self.interval_secs, "token 刷新任务已启动");
loop {
tokio::select! {
_ = ticker.tick() => {
match self.service.schedule_account_timers_once().await {
Ok(summary) if summary.changed() > 0 => {
info!(
scheduled = summary.scheduled,
immediate = summary.immediate,
recovery_scheduled = summary.recovery_scheduled,
replaced = summary.replaced,
"token 刷新定时器已调度"
);
}
Ok(summary) => {
debug!(
scanned = summary.scanned,
skipped = summary.skipped,
"没有需要调度刷新的 token"
);
}
Err(error) => {
warn!(error = %error, "token 刷新定时器调度失败");
}
}
}
_ = shutdown_rx.recv() => {
self.service.clear_scheduled_timers().await;
info!("token 刷新任务已关闭");
break;
}
impl<C> PeriodicTaskRunner for TokenRefreshTask<C>
where
C: TokenRefresher,
{
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self.service.schedule_account_timers_once().await {
Ok(summary) if summary.changed() > 0 => {
info!(
scheduled = summary.scheduled,
immediate = summary.immediate,
recovery_scheduled = summary.recovery_scheduled,
replaced = summary.replaced,
"token 刷新定时器已调度"
);
}
Ok(summary) => {
debug!(
scanned = summary.scanned,
skipped = summary.skipped,
"没有需要调度刷新的 token"
);
}
Err(error) => {
warn!(error = %error, "token 刷新定时器调度失败");
}
}
});
})
}
SchedulerHandle::new(shutdown_tx)
fn shutdown(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
self.service.clear_scheduled_timers().await;
})
}
}
-20
View File
@@ -8,26 +8,6 @@ use thiserror::Error;
use tokio::sync::RwLock;
use uuid::Uuid;
/// 单个 Cookie 条目。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CookieEntry {
/// Cookie 名称。
pub name: String,
/// Cookie 值。
pub value: String,
}
/// 将 Cookie 条目组装为 HTTP Cookie 头。
pub fn cookie_header(cookies: &[CookieEntry]) -> Option<String> {
let header = cookies
.iter()
.filter(|cookie| !cookie.name.trim().is_empty())
.map(|cookie| format!("{}={}", cookie.name.trim(), cookie.value))
.collect::<Vec<_>>()
.join("; ");
(!header.is_empty()).then_some(header)
}
const DEFAULT_COOKIE_DOMAIN: &str = "chatgpt.com";
const CAPTURABLE_COOKIES: &[&str] = &["cf_clearance"];
+25 -32
View File
@@ -141,61 +141,55 @@ fn explicit_account_import_source(
fn parse_cpr_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::Cpr,
entries: parse_cpr_entries(payload)?,
entries: parse_account_entries(
payload,
cpr_account_import_entry_from_value,
Some(ACCOUNT_IMPORT_CONTAINER_KEYS),
)?,
})
}
fn parse_cpr_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
ensure_account_import_keys(payload, ACCOUNT_IMPORT_CONTAINER_KEYS)?;
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, cpr_account_import_entry_from_value);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, cpr_account_import_entry_from_value);
}
cpr_account_import_entry_from_value(payload).map(option_to_vec)
}
fn parse_sub2api_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::Sub2Api,
entries: parse_sub2api_entries(payload)?,
entries: parse_account_entries(payload, sub2api_account_import_entry_from_value, None)?,
})
}
fn parse_sub2api_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, sub2api_account_import_entry_from_value);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, sub2api_account_import_entry_from_value);
}
sub2api_account_import_entry_from_value(payload).map(option_to_vec)
}
fn parse_cli_proxy_api_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::CliProxyApi,
entries: parse_cli_proxy_api_entries(payload)?,
entries: parse_account_entries(
payload,
cli_proxy_api_account_import_entry_from_value,
None,
)?,
})
}
fn parse_cli_proxy_api_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
type AccountImportEntryParser = fn(&Value) -> Result<Option<AccountImportEntry>, &'static str>;
fn parse_account_entries(
payload: &Value,
parser: AccountImportEntryParser,
container_keys: Option<&[&str]>,
) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
if let Some(container_keys) = container_keys {
ensure_account_import_keys(payload, container_keys)?;
}
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, cli_proxy_api_account_import_entry_from_value);
return parse_entries(accounts, parser);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, cli_proxy_api_account_import_entry_from_value);
return parse_entries(accounts, parser);
}
cli_proxy_api_account_import_entry_from_value(payload).map(option_to_vec)
parser(payload).map(option_to_vec)
}
fn parse_entries(
accounts: &[Value],
parser: fn(&Value) -> Result<Option<AccountImportEntry>, &'static str>,
parser: AccountImportEntryParser,
) -> Result<Vec<AccountImportEntry>, &'static str> {
let mut entries = Vec::new();
for account in accounts {
@@ -619,7 +613,6 @@ pub fn parse_account_status(status: &str) -> Result<AccountStatus, &'static str>
"disabled" => Ok(AccountStatus::Disabled),
"expired" => Ok(AccountStatus::Expired),
"quota_exhausted" => Ok(AccountStatus::QuotaExhausted),
"refreshing" => Ok(AccountStatus::Refreshing),
"banned" => Ok(AccountStatus::Banned),
_ => Err("invalid account status"),
}
+14 -10
View File
@@ -12,24 +12,28 @@ pub enum AccountStatus {
Expired,
/// 当前配额已经耗尽。
QuotaExhausted,
/// 账号正处于刷新流程中。
Refreshing,
/// 账号被显式禁用。
Disabled,
/// 账号被上游封禁。
Banned,
}
impl AccountStatus {
/// 返回账号状态的持久化/API 字符串。
pub fn as_str(self) -> &'static str {
match self {
Self::Active => "active",
Self::Expired => "expired",
Self::QuotaExhausted => "quota_exhausted",
Self::Disabled => "disabled",
Self::Banned => "banned",
}
}
}
impl std::fmt::Display for AccountStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Active => write!(f, "active"),
Self::Expired => write!(f, "expired"),
Self::QuotaExhausted => write!(f, "quota_exhausted"),
Self::Refreshing => write!(f, "refreshing"),
Self::Disabled => write!(f, "disabled"),
Self::Banned => write!(f, "banned"),
}
f.write_str(self.as_str())
}
}
+287 -251
View File
@@ -1,7 +1,6 @@
//! 账号池调度策略。
use std::{
cmp::Ordering,
collections::{BTreeMap, BTreeSet, VecDeque},
sync::Arc as StdArc,
time::Duration as StdDuration,
@@ -11,7 +10,6 @@ use chrono::{DateTime, Duration, Utc};
use indexmap::IndexMap;
use thiserror::Error;
use crate::infra::time::parse_rfc3339_utc;
use crate::upstream::accounts::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at};
use crate::upstream::accounts::store::AccountStore;
use crate::upstream::accounts::token_refresh::{jwt_expiry, JwtExpiry};
@@ -23,39 +21,12 @@ use crate::upstream::accounts::{
use crate::upstream::protocol::events::{
parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage,
};
use crate::upstream::scheduler::{
candidates::{self, CandidateFilter, CandidateRequest},
AccountScheduler, ScoreWeights,
};
/// Token usage statistics.
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct TokenUsage {
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,
}
/// 账号轮转策略。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RotationStrategy {
/// 优先选择当前窗口用量最低的账号。
LeastUsed,
/// 按候选账号顺序循环选择。
RoundRobin,
/// 优先选择最近使用过的账号。
Sticky,
}
impl RotationStrategy {
fn as_str(self) -> &'static str {
match self {
Self::LeastUsed => "least_used",
Self::RoundRobin => "round_robin",
Self::Sticky => "sticky",
}
}
}
pub use crate::upstream::scheduler::RotationStrategy;
/// 账号池配置。
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -74,8 +45,6 @@ pub struct AccountPoolOptions {
pub model_plan_allowlist: BTreeMap<String, Vec<String>>,
/// 已成功拉取过模型列表的订阅计划。
pub fetched_model_plan_types: BTreeSet<String>,
/// 模型到指定账号 ID 的显式路由。
pub model_account_routes: BTreeMap<String, Vec<String>>,
}
impl Default for AccountPoolOptions {
@@ -83,12 +52,11 @@ impl Default for AccountPoolOptions {
Self {
max_concurrent_per_account: 3,
stale_slot_ttl: Duration::minutes(5),
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::Smart,
skip_quota_limited: true,
tier_priority: Vec::new(),
model_plan_allowlist: BTreeMap::new(),
fetched_model_plan_types: BTreeSet::new(),
model_account_routes: BTreeMap::new(),
}
}
}
@@ -100,6 +68,8 @@ pub struct AccountAcquireRequest {
pub model: String,
/// 本次请求需要排除的账号 ID。
pub exclude_account_ids: Vec<String>,
/// 本次请求必须使用的账号 ID,用于同请求内的同账号重试。
pub required_account_id: Option<String>,
/// 会话亲和性建议的优先账号 ID。
pub preferred_account_id: Option<String>,
/// 本次调度使用的当前时间。
@@ -112,6 +82,7 @@ impl AccountAcquireRequest {
Self {
model: model.into(),
exclude_account_ids: Vec::new(),
required_account_id: None,
preferred_account_id: None,
now,
}
@@ -131,6 +102,12 @@ impl AccountAcquireRequest {
self.preferred_account_id = Some(account_id.into());
self
}
/// 设置本次调度必须使用的账号 ID。
pub fn with_required_account_id(mut self, account_id: impl Into<String>) -> Self {
self.required_account_id = Some(account_id.into());
self
}
}
/// 成功获取到的账号及调度元数据。
@@ -154,13 +131,19 @@ pub struct DistinctPlanAccount {
#[derive(Debug)]
struct AccountAcquireWithStatusRefresh {
acquired: Option<AcquiredAccount>,
expired_account_ids: Vec<String>,
refreshed_accounts: Vec<RefreshedAccountState>,
}
#[derive(Debug)]
struct DistinctPlanAccountsWithStatusRefresh {
accounts: Vec<DistinctPlanAccount>,
expired_account_ids: Vec<String>,
refreshed_accounts: Vec<RefreshedAccountState>,
}
#[derive(Debug)]
struct AccountCapacitySummaryWithStatusRefresh {
summary: AccountCapacitySummary,
refreshed_accounts: Vec<RefreshedAccountState>,
}
/// 账号池容量摘要。
@@ -176,6 +159,71 @@ pub struct AccountCapacitySummary {
pub available_slots: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct RuntimeAccountStateSnapshot {
status: AccountStatus,
quota_limit_reached: bool,
quota_verify_required: bool,
quota_cooldown_until: Option<DateTime<Utc>>,
cloudflare_cooldown_until: Option<DateTime<Utc>>,
window_request_count: u64,
window_input_tokens: u64,
window_output_tokens: u64,
window_cached_tokens: u64,
window_image_input_tokens: u64,
window_image_output_tokens: u64,
window_image_request_count: u64,
window_image_request_failed_count: u64,
window_started_at: Option<DateTime<Utc>>,
window_reset_at: Option<DateTime<Utc>>,
limit_window_seconds: Option<u64>,
}
impl From<&Account> for RuntimeAccountStateSnapshot {
fn from(account: &Account) -> Self {
Self {
status: account.status,
quota_limit_reached: account.quota_limit_reached,
quota_verify_required: account.quota_verify_required,
quota_cooldown_until: account.quota_cooldown_until,
cloudflare_cooldown_until: account.cloudflare_cooldown_until,
window_request_count: account.window_request_count,
window_input_tokens: account.window_input_tokens,
window_output_tokens: account.window_output_tokens,
window_cached_tokens: account.window_cached_tokens,
window_image_input_tokens: account.window_image_input_tokens,
window_image_output_tokens: account.window_image_output_tokens,
window_image_request_count: account.window_image_request_count,
window_image_request_failed_count: account.window_image_request_failed_count,
window_started_at: account.window_started_at,
window_reset_at: account.window_reset_at,
limit_window_seconds: account.limit_window_seconds,
}
}
}
impl RuntimeAccountStateSnapshot {
fn usage_window_changed(&self, other: &Self) -> bool {
self.window_request_count != other.window_request_count
|| self.window_input_tokens != other.window_input_tokens
|| self.window_output_tokens != other.window_output_tokens
|| self.window_cached_tokens != other.window_cached_tokens
|| self.window_image_input_tokens != other.window_image_input_tokens
|| self.window_image_output_tokens != other.window_image_output_tokens
|| self.window_image_request_count != other.window_image_request_count
|| self.window_image_request_failed_count != other.window_image_request_failed_count
|| self.window_started_at != other.window_started_at
|| self.window_reset_at != other.window_reset_at
|| self.limit_window_seconds != other.limit_window_seconds
}
}
#[derive(Debug, Clone)]
struct RefreshedAccountState {
account: Account,
sync_usage_window: bool,
}
/// 运行时窗口用量增量。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountWindowUsageDelta {
@@ -195,13 +243,17 @@ pub struct AccountWindowUsageDelta {
pub image_request_failed: bool,
}
/// 纯内存账号池,负责账号调度和运行时状态维护。
/// 纯内存账号池:负责账号存储、在途槽位与运行时状态维护。
///
/// 账号**选择**已剥离到 [`crate::upstream::scheduler`]:候选过滤走
/// [`candidates::filter`],策略选择走 [`AccountScheduler::select`]。账号池只提供选择
/// 所需的只读视图(候选切片 + `slot_count`),并持有调度器实例。
#[derive(Debug)]
pub struct AccountPool {
accounts: IndexMap<String, Account>,
slots: BTreeMap<String, VecDeque<AccountSlot>>,
options: AccountPoolOptions,
rotation_cursor: usize,
scheduler: AccountScheduler,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -230,7 +282,7 @@ impl AccountPool {
accounts: IndexMap::new(),
slots: BTreeMap::new(),
options,
rotation_cursor: 0,
scheduler: AccountScheduler::new(ScoreWeights::default()),
}
}
@@ -244,10 +296,11 @@ impl AccountPool {
self.accounts.get(account_id).cloned()
}
/// 移除账号及其在途槽位。
/// 移除账号及其在途槽位与运行时反馈。
pub fn remove(&mut self, account_id: &str) -> bool {
let removed = self.accounts.shift_remove(account_id).is_some();
self.slots.remove(account_id);
self.scheduler.forget_feedback(account_id);
removed
}
@@ -255,22 +308,30 @@ impl AccountPool {
pub fn clear(&mut self) {
self.accounts.clear();
self.slots.clear();
self.rotation_cursor = 0;
self.scheduler.reset_cursor();
self.scheduler.clear_feedback();
}
/// 替换账号池运行参数。
pub fn set_options(&mut self, options: AccountPoolOptions) {
if self.options.rotation_strategy != options.rotation_strategy {
self.rotation_cursor = 0;
self.scheduler.reset_cursor();
}
self.options = options;
}
/// 回灌一次请求结果到调度器的运行时反馈(供 Smart 打分)。
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
self.scheduler
.report_feedback(account_id, success, first_token_ms);
}
/// 更新账号状态。
pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool {
let Some(account) = self.accounts.get_mut(account_id) else {
return false;
};
let status = status_after_quota_limit(status, account.quota_limit_reached);
account.status = status;
if status != AccountStatus::Active {
self.slots.remove(account_id);
@@ -293,6 +354,7 @@ impl AccountPool {
.unwrap_or(cooldown_until);
account.quota_limit_reached = true;
account.quota_cooldown_until = Some(final_cooldown_until);
account.status = status_after_quota_limit(account.status, true);
account.window_reset_at = account
.window_reset_at
.filter(|existing| *existing > final_cooldown_until)
@@ -419,46 +481,42 @@ impl AccountPool {
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 selected = if let Some(preferred_account_id) = &request.preferred_account_id {
let preferred = candidates
.iter()
.find(|account| account.id == *preferred_account_id)
.cloned();
preferred.map(|account| {
let refreshed_accounts = self.refresh_account_statuses(request.now);
let candidates = self.filter_candidates(request);
// 亲和优先:required 硬锁定该账号;preferred 若在候选内则直接选中;
// 二者都不命中时,回退到调度器按当前策略选择。
let selected = if let Some(required_account_id) = &request.required_account_id {
candidates.first().cloned().map(|account| {
(
account,
self.previous_slot_at(preferred_account_id),
"preferred",
self.previous_slot_at(required_account_id),
"required",
)
})
} else if let Some(preferred_account_id) = &request.preferred_account_id {
candidates
.iter()
.find(|account| account.id == *preferred_account_id)
.cloned()
.map(|account| {
(
account,
self.previous_slot_at(preferred_account_id),
"preferred",
)
})
} else {
None
}
.or_else(|| {
Some(match self.options.rotation_strategy {
RotationStrategy::LeastUsed => {
let account = self.select_least_used(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous, "least_used")
}
RotationStrategy::RoundRobin => {
let account = self.select_round_robin(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous, "round_robin")
}
RotationStrategy::Sticky => {
let account = Self::select_sticky(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous, "sticky")
}
})
let account = self.select_candidate(&candidates, request.now)?;
let previous = self.previous_slot_at(&account.id);
Some((account, previous, self.options.rotation_strategy.as_str()))
});
let Some((selected, previous_slot_at, selection_source)) = selected else {
return AccountAcquireWithStatusRefresh {
acquired: None,
expired_account_ids,
refreshed_accounts,
};
};
tracing::info!(
@@ -467,6 +525,7 @@ impl AccountPool {
selection_source,
account_id = %selected.id,
candidate_count = candidates.len(),
slot_count = self.slot_count(&selected.id),
request_count = selected.request_count,
window_request_count = selected.window_request_count,
window_reset_at = ?selected.window_reset_at,
@@ -474,7 +533,6 @@ impl AccountPool {
quota_limit_reached = selected.quota_limit_reached,
quota_cooldown_until = ?selected.quota_cooldown_until,
previous_slot_at = ?previous_slot_at,
rotation_cursor = self.rotation_cursor,
"account selected for upstream request"
);
@@ -485,13 +543,15 @@ impl AccountPool {
account: selected,
previous_slot_at,
}),
expired_account_ids,
refreshed_accounts,
}
}
fn refresh_account_statuses(&mut self, now: DateTime<Utc>) -> Vec<String> {
fn refresh_account_statuses(&mut self, now: DateTime<Utc>) -> Vec<RefreshedAccountState> {
let mut refreshed_accounts = Vec::new();
let mut expired_account_ids = Vec::new();
for (account_id, account) in &mut self.accounts {
let before = RuntimeAccountStateSnapshot::from(&*account);
if account.status == AccountStatus::Active && access_token_expired(account, now) {
account.status = AccountStatus::Expired;
expired_account_ids.push(account_id.clone());
@@ -499,15 +559,39 @@ impl AccountPool {
refresh_quota_window(account, now);
refresh_cloudflare_cooldown(account, now);
}
let after = RuntimeAccountStateSnapshot::from(&*account);
if before != after {
let sync_usage_window = before.usage_window_changed(&after);
refreshed_accounts.push(RefreshedAccountState {
account: account.clone(),
sync_usage_window,
});
}
}
for account_id in &expired_account_ids {
self.slots.remove(account_id);
}
expired_account_ids
refreshed_accounts
}
/// 释放指定账号的一个在途槽位。
pub fn release(&mut self, account_id: &str) -> Option<ReleasedAccountUsage> {
self.release_slot(account_id, true)
}
/// 释放指定账号的一个在途槽位,不累计请求用量。
pub fn release_without_request_usage(
&mut self,
account_id: &str,
) -> Option<ReleasedAccountUsage> {
self.release_slot(account_id, false)
}
fn release_slot(
&mut self,
account_id: &str,
record_request_usage: bool,
) -> Option<ReleasedAccountUsage> {
self.accounts.get(account_id)?;
let mut remove_slots = false;
@@ -519,7 +603,9 @@ impl AccountPool {
if remove_slots {
self.slots.remove(account_id);
}
self.mark_request_usage(account_id, Utc::now());
if record_request_usage {
self.mark_request_usage(account_id, Utc::now());
}
Some(ReleasedAccountUsage {
model: slot.and_then(|slot| slot.model),
})
@@ -527,8 +613,15 @@ impl AccountPool {
/// 计算账号池容量摘要。
pub fn capacity_summary(&mut self, now: DateTime<Utc>) -> AccountCapacitySummary {
self.capacity_summary_with_status_refresh(now).summary
}
fn capacity_summary_with_status_refresh(
&mut self,
now: DateTime<Utc>,
) -> AccountCapacitySummaryWithStatusRefresh {
self.cleanup_stale_slots(now);
self.refresh_account_statuses(now);
let refreshed_accounts = self.refresh_account_statuses(now);
let is_capacity_account = |account: &Account| {
account.status == AccountStatus::Active
&& AccountService::quota_available_at(account, now, self.options.skip_quota_limited)
@@ -550,11 +643,14 @@ impl AccountPool {
.map(|(_, slots)| slots.len().min(self.options.max_concurrent_per_account))
.sum();
AccountCapacitySummary {
max_concurrent_per_account: self.options.max_concurrent_per_account,
total_slots,
used_slots,
available_slots: total_slots.saturating_sub(used_slots),
AccountCapacitySummaryWithStatusRefresh {
summary: AccountCapacitySummary {
max_concurrent_per_account: self.options.max_concurrent_per_account,
total_slots,
used_slots,
available_slots: total_slots.saturating_sub(used_slots),
},
refreshed_accounts,
}
}
@@ -569,44 +665,51 @@ impl AccountPool {
now: DateTime<Utc>,
) -> DistinctPlanAccountsWithStatusRefresh {
self.cleanup_stale_slots(now);
let expired_account_ids = self.refresh_account_statuses(now);
let refreshed_accounts = self.refresh_account_statuses(now);
let mut by_plan = IndexMap::<String, Vec<Account>>::new();
for account in self.accounts.values() {
if !self.is_model_refresh_available(account, now) {
continue;
let max_concurrent = self.options.max_concurrent_per_account;
let skip_quota_limited = self.options.skip_quota_limited;
{
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
for account in self.accounts.values() {
if !candidates::is_model_refresh_available(
account,
max_concurrent,
skip_quota_limited,
&slot_count,
now,
) {
continue;
}
let Some(plan_type) = account.plan_type.as_ref() else {
continue;
};
by_plan
.entry(plan_type.clone())
.or_default()
.push(account.clone());
}
let Some(plan_type) = account.plan_type.as_ref() else {
continue;
};
by_plan
.entry(plan_type.clone())
.or_default()
.push(account.clone());
}
let strategy = self.options.rotation_strategy;
let mut accounts = Vec::new();
for (plan_type, group) in by_plan {
let selected = match self.options.rotation_strategy {
RotationStrategy::LeastUsed => self.select_least_used(&group),
RotationStrategy::RoundRobin => self.select_round_robin(&group),
RotationStrategy::Sticky => Self::select_sticky(&group),
};
let Some(account) = selected else {
let Some(account) = self.select_candidate(&group, now) else {
continue;
};
tracing::info!(
plan_type,
rotation_strategy = self.options.rotation_strategy.as_str(),
rotation_strategy = strategy.as_str(),
account_id = %account.id,
candidate_count = group.len(),
slot_count = self.slot_count(&account.id),
request_count = account.request_count,
window_request_count = account.window_request_count,
window_reset_at = ?account.window_reset_at,
last_used_at = account.last_used_at.as_deref().unwrap_or_default(),
quota_limit_reached = account.quota_limit_reached,
quota_cooldown_until = ?account.quota_cooldown_until,
rotation_cursor = self.rotation_cursor,
"account selected for model refresh"
);
self.push_slot(&account.id, now, None);
@@ -615,77 +718,10 @@ impl AccountPool {
DistinctPlanAccountsWithStatusRefresh {
accounts,
expired_account_ids,
refreshed_accounts,
}
}
fn select_least_used(&mut self, candidates: &[Account]) -> Option<Account> {
let mut sorted = candidates.iter().collect::<Vec<_>>();
sorted.sort_by(|a, b| compare_least_used(a, b));
let best = *sorted.first()?;
let tied_count = sorted
.iter()
.take_while(|account| compare_least_used(best, account) == Ordering::Equal)
.count();
let index = self.rotation_cursor % tied_count;
self.rotation_cursor = self.rotation_cursor.wrapping_add(1);
Some((*sorted[index]).clone())
}
fn select_sticky(candidates: &[Account]) -> Option<Account> {
let mut selected = candidates.first()?;
for candidate in &candidates[1..] {
if compare_last_used(
candidate.last_used_at.as_deref(),
selected.last_used_at.as_deref(),
) == Ordering::Greater
{
selected = candidate;
}
}
Some((*selected).clone())
}
fn select_round_robin(&mut self, candidates: &[Account]) -> Option<Account> {
if candidates.is_empty() {
return None;
}
self.rotation_cursor %= candidates.len();
let index = self.rotation_cursor;
self.rotation_cursor = self.rotation_cursor.wrapping_add(1);
Some(candidates[index].clone())
}
fn candidates(&self, request: &AccountAcquireRequest) -> Vec<Account> {
let mut candidates = self
.accounts
.values()
.filter(|account| self.is_base_available(account, request))
.cloned()
.collect::<Vec<_>>();
if let Some(best_tier) = self.best_available_tier(&candidates) {
candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str()));
}
candidates
}
fn is_base_available(&self, account: &Account, request: &AccountAcquireRequest) -> bool {
account.status == AccountStatus::Active
&& self.slot_count(&account.id) < self.options.max_concurrent_per_account
&& AccountService::quota_available_at(
account,
request.now,
self.options.skip_quota_limited,
)
&& self.is_model_account_allowed(account, &request.model)
&& self.is_model_allowed(account, &request.model)
&& !request
.exclude_account_ids
.iter()
.any(|account_id| account_id == &account.id)
&& AccountService::cloudflare_available_at(account, request.now)
}
fn slot_count(&self, account_id: &str) -> usize {
self.slots.get(account_id).map_or(0, VecDeque::len)
}
@@ -707,13 +743,6 @@ impl AccountPool {
});
}
fn is_model_refresh_available(&self, account: &Account, now: DateTime<Utc>) -> bool {
account.status == AccountStatus::Active
&& self.slot_count(&account.id) < self.options.max_concurrent_per_account
&& AccountService::quota_available_at(account, now, self.options.skip_quota_limited)
&& AccountService::cloudflare_available_at(account, now)
}
fn mark_request_usage(&mut self, account_id: &str, now: DateTime<Utc>) -> Option<Account> {
let account = self.accounts.get_mut(account_id)?;
account.last_used_at = Some(now.to_rfc3339());
@@ -763,33 +792,33 @@ impl AccountPool {
});
}
fn is_model_allowed(&self, account: &Account, model: &str) -> bool {
let Some(allowed_plans) = self.options.model_plan_allowlist.get(model) else {
return true;
};
let Some(plan_type) = account.plan_type.as_deref() else {
return false;
};
allowed_plans.iter().any(|plan| plan == plan_type)
|| !self.options.fetched_model_plan_types.contains(plan_type)
/// 过滤出本次请求的可用候选集合(委托 [`candidates::filter`])。
fn filter_candidates(&self, request: &AccountAcquireRequest) -> Vec<Account> {
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
candidates::filter(
self.accounts.values(),
&CandidateFilter {
max_concurrent_per_account: self.options.max_concurrent_per_account,
skip_quota_limited: self.options.skip_quota_limited,
tier_priority: &self.options.tier_priority,
model_plan_allowlist: &self.options.model_plan_allowlist,
fetched_model_plan_types: &self.options.fetched_model_plan_types,
slot_count: &slot_count,
},
&CandidateRequest {
model: &request.model,
exclude_account_ids: &request.exclude_account_ids,
required_account_id: request.required_account_id.as_deref(),
now: request.now,
},
)
}
fn is_model_account_allowed(&self, account: &Account, model: &str) -> bool {
let Some(account_ids) = self.options.model_account_routes.get(model) else {
return true;
};
account_ids
.iter()
.any(|account_id| account_id == &account.id)
}
fn best_available_tier(&self, candidates: &[Account]) -> Option<String> {
self.options.tier_priority.iter().find_map(|tier| {
candidates
.iter()
.find(|account| account.plan_type.as_deref() == Some(tier.as_str()))
.and_then(|account| account.plan_type.clone())
})
/// 按当前策略从候选集合中选择一个账号(委托 [`AccountScheduler::select`])。
fn select_candidate(&self, candidates: &[Account], now: DateTime<Utc>) -> Option<Account> {
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
self.scheduler
.select(self.options.rotation_strategy, candidates, &slot_count, now)
}
}
@@ -874,7 +903,14 @@ impl RuntimeAccountPoolService {
/// 读取账号池容量摘要。
pub async fn capacity_summary(&self, now: DateTime<Utc>) -> AccountCapacitySummary {
self.pool.lock().await.capacity_summary(now)
let refresh = self
.pool
.lock()
.await
.capacity_summary_with_status_refresh(now);
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.summary
}
/// 使用当前时间读取账号池容量摘要。
@@ -884,10 +920,29 @@ impl RuntimeAccountPoolService {
/// 释放账号在途槽位。
pub async fn release(&self, account_id: &str) {
let released = self.pool.lock().await.release(account_id);
self.release_slot(account_id, true).await;
}
/// 释放账号在途槽位,不累计请求用量。
pub async fn release_without_request_usage(&self, account_id: &str) {
self.release_slot(account_id, false).await;
}
async fn release_slot(&self, account_id: &str, record_request_usage: bool) {
let released = if record_request_usage {
self.pool.lock().await.release(account_id)
} else {
self.pool
.lock()
.await
.release_without_request_usage(account_id)
};
let Some(released) = released else {
return;
};
if !record_request_usage {
return;
}
if let Err(error) = self.store.record_request(account_id).await {
tracing::warn!(
account_id,
@@ -964,7 +1019,7 @@ impl RuntimeAccountPoolService {
/// 按请求上下文获取可用账号。
pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option<AcquiredAccount> {
let refresh = self.pool.lock().await.acquire_with_status_refresh(request);
self.persist_expired_statuses(refresh.expired_account_ids)
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.acquired
}
@@ -976,22 +1031,22 @@ impl RuntimeAccountPoolService {
.lock()
.await
.distinct_plan_accounts_with_status_refresh(now);
self.persist_expired_statuses(refresh.expired_account_ids)
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.accounts
}
async fn persist_expired_statuses(&self, account_ids: Vec<String>) {
for account_id in account_ids {
async fn persist_runtime_account_states(&self, refreshed_states: Vec<RefreshedAccountState>) {
for refreshed in refreshed_states {
if let Err(error) = self
.store
.set_status(&account_id, AccountStatus::Expired)
.sync_runtime_account_state(&refreshed.account, refreshed.sync_usage_window)
.await
{
tracing::warn!(
account_id,
account_id = %refreshed.account.id,
error = %error,
"failed to persist runtime-expired account status"
"failed to persist refreshed runtime account state"
);
}
}
@@ -1191,6 +1246,9 @@ impl RuntimeAccountPoolService {
let limit_reached = quota_snapshot_limit_reached(quota);
let reset_at = quota_snapshot_reset_at(quota);
let cooldown_until = limit_reached.then_some(reset_at).flatten();
let limit_window_seconds = reset_at.and_then(|_| {
crate::upstream::accounts::quota::quota_snapshot_limit_window_seconds(quota)
});
let quota_json = quota.to_string();
let persisted = match self
.store
@@ -1207,15 +1265,8 @@ impl RuntimeAccountPoolService {
false
}
};
let in_memory =
self.pool
.lock()
.await
.apply_quota_state(account_id, limit_reached, cooldown_until);
if let Some(reset_at) = reset_at {
let limit_window_seconds =
crate::upstream::accounts::quota::quota_snapshot_limit_window_seconds(quota);
if let Err(error) = self
.store
.sync_rate_limit_window(account_id, reset_at, limit_window_seconds)
@@ -1227,13 +1278,17 @@ impl RuntimeAccountPoolService {
"failed to persist verified quota window"
);
}
self.pool.lock().await.sync_rate_limit_window(
account_id,
reset_at,
limit_window_seconds,
);
}
let in_memory = {
let mut pool = self.pool.lock().await;
let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until);
if let Some(reset_at) = reset_at {
pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds);
}
in_memory
};
persisted || in_memory
}
@@ -1323,6 +1378,13 @@ fn refresh_quota_window(account: &mut Account, now: DateTime<Utc>) {
}
}
fn status_after_quota_limit(status: AccountStatus, quota_limit_reached: bool) -> AccountStatus {
match (status, quota_limit_reached) {
(AccountStatus::Active, true) => AccountStatus::QuotaExhausted,
(status, _) => status,
}
}
fn reset_window_counters(account: &mut Account) {
account.window_request_count = 0;
account.window_input_tokens = 0;
@@ -1365,29 +1427,3 @@ fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime<Utc>) {
account.cloudflare_cooldown_until = None;
}
}
fn compare_least_used(a: &Account, b: &Account) -> Ordering {
a.quota_limit_reached
.cmp(&b.quota_limit_reached)
.then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at))
.then_with(|| a.request_count.cmp(&b.request_count))
.then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref()))
}
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) | (None, Some(_)) => Ordering::Equal,
}
}
fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering {
last_used_millis(a).cmp(&last_used_millis(b))
}
fn last_used_millis(value: Option<&str>) -> i64 {
value
.and_then(|value| parse_rfc3339_utc(value).ok())
.map(|datetime| datetime.timestamp_millis())
.unwrap_or(0)
}
+73 -36
View File
@@ -10,7 +10,7 @@ const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
/// 判断配额快照是否已经触顶。
pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
core_snapshot_blocked(quota)
snapshot_blocked(quota)
|| spend_control_limit_reached(quota)
|| monthly_limit_reached(quota)
|| credits_overage_limit_reached(quota)
@@ -20,8 +20,10 @@ pub fn quota_snapshot_limit_reached(quota: &Value) -> bool {
pub fn quota_snapshot_reset_at(quota: &Value) -> Option<DateTime<Utc>> {
let reset_at = if monthly_limit_reached(quota) || spend_control_limit_reached(quota) {
monthly_limit_reset_at(quota)
.or_else(|| blocking_core_bucket(quota).and_then(bucket_reset_at))
.or_else(|| blocking_bucket(quota).and_then(bucket_reset_at))
.or_else(|| core_bucket_with_reset_at(quota).and_then(bucket_reset_at))
} else if let Some(reset_at) = blocking_bucket(quota).and_then(bucket_reset_at) {
Some(reset_at)
} else {
core_bucket_with_reset_at(quota)
.and_then(bucket_reset_at)
@@ -35,9 +37,14 @@ pub fn quota_snapshot_limit_window_seconds(quota: &Value) -> Option<u64> {
if monthly_limit_reached(quota) || spend_control_limit_reached(quota) {
return monthly_limit_window_seconds(quota).or(Some(MONTH_WINDOW_SECONDS));
}
core_bucket_with_window_minutes(quota)
blocking_bucket(quota)
.and_then(bucket_window_minutes)
.and_then(|minutes| minutes.checked_mul(60))
.or_else(|| {
core_bucket_with_window_minutes(quota)
.and_then(bucket_window_minutes)
.and_then(|minutes| minutes.checked_mul(60))
})
.or_else(|| monthly_limit_window_seconds(quota))
}
@@ -132,6 +139,8 @@ fn window_from_rate_limit(window: Option<&Value>) -> Value {
"reset_at": window.get("reset_at").and_then(positive_i64),
"window_minutes": window_minutes,
"limit_reached": used_percent >= 100.0,
"used": first_present(window, &["used_tokens", "used", "used_credits", "usage", "consumed_tokens", "consumed"]),
"limit": first_present(window, &["limit_tokens", "limit", "limit_credits", "quota", "total_tokens", "total"]),
})
}
@@ -201,8 +210,10 @@ fn monthly_limit_from_snapshots(snapshots: &[Value]) -> Option<Value> {
"reset_at": bucket.get("reset_at").cloned().unwrap_or(Value::Null),
"window_minutes": bucket.get("window_minutes").cloned().unwrap_or(Value::Null),
"limit_reached": bucket.get("limit_reached").cloned().unwrap_or(Value::Bool(false)),
"used_credits": Value::Null,
"limit_credits": Value::Null,
"used": bucket.get("used").cloned().unwrap_or(Value::Null),
"limit": bucket.get("limit").cloned().unwrap_or(Value::Null),
"used_credits": bucket.get("used_credits").cloned().unwrap_or(Value::Null),
"limit_credits": bucket.get("limit_credits").cloned().unwrap_or(Value::Null),
}));
}
}
@@ -214,6 +225,12 @@ fn remaining_percent(used_percent: f64) -> i64 {
(100.0 - used_percent.clamp(0.0, 100.0)).round() as i64
}
fn first_present(value: &Value, keys: &[&str]) -> Value {
keys.iter()
.find_map(|key| value.get(*key).filter(|value| !value.is_null()).cloned())
.unwrap_or(Value::Null)
}
fn rate_limit_blocked(rate_limit: &Value) -> bool {
rate_limit
.get("limit_reached")
@@ -256,41 +273,66 @@ fn monthly_limit_reached(quota: &Value) -> bool {
.pointer("/monthly_limit/limit_reached")
.and_then(Value::as_bool)
.unwrap_or(false)
|| quota
.pointer("/monthly_limit/used_percent")
.and_then(number_value)
.is_some_and(|used_percent| used_percent >= 100.0)
}
fn core_snapshot_blocked(quota: &Value) -> bool {
fn quota_bucket_limit_reached(bucket: &Value) -> bool {
!bucket.is_null()
&& (bucket
.get("limit_reached")
.and_then(Value::as_bool)
.unwrap_or(false)
|| bucket
.get("used_percent")
.and_then(number_value)
.is_some_and(|used_percent| used_percent >= 100.0))
}
fn quota_snapshot_blocked(snapshot: &Value) -> bool {
snapshot
.get("blocked")
.and_then(Value::as_bool)
.unwrap_or(false)
|| snapshot
.get("allowed")
.and_then(Value::as_bool)
.is_some_and(|allowed| !allowed)
|| snapshot_buckets(snapshot).any(quota_bucket_limit_reached)
}
fn snapshot_blocked(quota: &Value) -> bool {
quota
.get("snapshots")
.and_then(Value::as_array)
.into_iter()
.flatten()
.any(|snapshot| {
snapshot.get("source").and_then(Value::as_str) == Some("core")
&& snapshot
.get("blocked")
.and_then(Value::as_bool)
.unwrap_or(false)
.any(quota_snapshot_blocked)
}
fn blocking_bucket(quota: &Value) -> Option<&Value> {
let snapshots = quota.get("snapshots").and_then(Value::as_array)?;
snapshots
.iter()
.filter(|snapshot| quota_snapshot_blocked(snapshot))
.flat_map(snapshot_buckets)
.find(|bucket| quota_bucket_limit_reached(bucket))
.or_else(|| {
snapshots
.iter()
.filter(|snapshot| quota_snapshot_blocked(snapshot))
.flat_map(snapshot_buckets)
.find(|bucket| bucket_reset_at(bucket).is_some())
})
}
fn blocking_core_bucket(quota: &Value) -> Option<&Value> {
quota
.get("snapshots")
.and_then(Value::as_array)?
.iter()
.filter(|snapshot| snapshot.get("source").and_then(Value::as_str) == Some("core"))
.flat_map(|snapshot| {
["primary", "secondary"]
.into_iter()
.filter_map(|key| snapshot.get(key))
})
.find(|bucket| {
!bucket.is_null()
&& bucket
.get("limit_reached")
.and_then(Value::as_bool)
.unwrap_or(false)
})
fn snapshot_buckets(snapshot: &Value) -> impl Iterator<Item = &Value> {
["primary", "secondary"]
.into_iter()
.filter_map(|key| snapshot.get(key))
.filter(|bucket| !bucket.is_null())
}
fn core_bucket_with_reset_at(quota: &Value) -> Option<&Value> {
@@ -308,12 +350,7 @@ fn core_buckets<'a>(quota: &'a Value) -> impl Iterator<Item = &'a Value> + 'a {
.into_iter()
.flatten()
.filter(|snapshot| snapshot.get("source").and_then(Value::as_str) == Some("core"))
.flat_map(|snapshot| {
["primary", "secondary"]
.into_iter()
.filter_map(|key| snapshot.get(key))
})
.filter(|bucket| !bucket.is_null())
.flat_map(snapshot_buckets)
}
fn bucket_reset_at(bucket: &Value) -> Option<i64> {
+23 -3
View File
@@ -2,6 +2,7 @@
use std::{collections::HashMap, sync::Arc, time::Duration};
use chrono::Utc;
use thiserror::Error;
use tokio::time::{sleep, Instant};
use tracing::warn;
@@ -22,6 +23,7 @@ use crate::upstream::{
const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60;
const DEFAULT_REQUEST_SPACING_SECS: u64 = 3;
const COOLDOWN_REFRESH_GRACE_SECS: i64 = 30;
/// 运行时 quota refresh 服务。
pub struct RuntimeQuotaRefreshService {
@@ -109,6 +111,7 @@ impl RuntimeQuotaRefreshService {
};
let min_interval = Duration::from_secs(self.min_refresh_interval_secs);
let now = Instant::now();
let now_utc = Utc::now();
let mut candidates = accounts
.into_iter()
@@ -121,9 +124,18 @@ impl RuntimeQuotaRefreshService {
while let Some(account) = candidates.next() {
summary.candidates += 1;
if last_refreshed
.get(&account.id)
.is_some_and(|last| now.duration_since(*last) < min_interval)
let cooldown_ready = match account.quota_cooldown_until {
Some(cooldown_until) if !quota_cooldown_ready(cooldown_until, now_utc) => {
summary.skipped_cooldown += 1;
continue;
}
Some(_) => true,
None => false,
};
if !cooldown_ready
&& last_refreshed
.get(&account.id)
.is_some_and(|last| now.duration_since(*last) < min_interval)
{
summary.skipped_recent += 1;
continue;
@@ -221,6 +233,12 @@ impl RuntimeQuotaRefreshService {
}
}
fn quota_cooldown_ready(cooldown_until: chrono::DateTime<Utc>, now: chrono::DateTime<Utc>) -> bool {
cooldown_until
.checked_add_signed(chrono::Duration::seconds(COOLDOWN_REFRESH_GRACE_SECS))
.is_some_and(|ready_at| now >= ready_at)
}
/// 单次配额刷新摘要。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct QuotaRefreshSummary {
@@ -228,6 +246,8 @@ pub struct QuotaRefreshSummary {
pub scanned: usize,
/// 配额锁定或待验证候选账号数。
pub candidates: usize,
/// 配额冷却宽限期内跳过的账号数。
pub skipped_cooldown: usize,
/// 成功刷新账号数。
pub refreshed: usize,
/// 最近刷新过而跳过的账号数。
+100 -53
View File
@@ -21,7 +21,7 @@ mod rows;
use queries::*;
use rows::{
count_account_metadata, get_pool_account, list_pool_accounts, map_account_store_error,
metadata_from_row, model_usage_from_row, optional_positive_i64_to_u64, optional_update_value,
metadata_from_row, optional_positive_i64_to_u64, optional_update_value,
push_account_metadata_search, quota_plan_type, quota_snapshot_from_row, sqlite_usage_delta,
status_to_db, stored_account_from_row, to_page, u64_to_i64_saturating,
};
@@ -196,27 +196,6 @@ impl AccountMetadataUpdate {
}
}
/// 账号模型维度用量记录。
#[derive(Debug, Clone)]
pub struct AccountModelUsageRecord {
/// 账号 ID。
pub account_id: String,
/// 模型 ID。
pub model: String,
/// 历史请求总数。
pub request_count: i64,
/// 历史错误数。
pub error_count: i64,
/// 累计输入 token。
pub input_tokens: i64,
/// 累计输出 token。
pub output_tokens: i64,
/// 累计缓存 token。
pub cached_tokens: i64,
/// 最近使用时间。
pub last_used_at: Option<DateTime<Utc>>,
}
/// 账号配额快照。
#[derive(Debug, Clone)]
pub struct AccountQuotaSnapshot {
@@ -322,13 +301,6 @@ pub trait AccountStore: Send + Sync + 'static {
/// 读取账号当前配额 JSON。
async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult<Option<String>>;
/// 更新账号当前配额 JSON。
async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> AccountStoreResult<bool>;
/// 应用已经验证过的账号配额快照。
async fn apply_quota_snapshot(
&self,
@@ -338,6 +310,13 @@ pub trait AccountStore: Send + Sync + 'static {
cooldown_until: Option<DateTime<Utc>>,
) -> AccountStoreResult<bool>;
/// 同步运行时自然刷新出来的账号状态。
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool>;
/// 同步账号当前 rate-limit 统计窗口。
async fn sync_rate_limit_window(
&self,
@@ -561,6 +540,7 @@ impl SqliteAccountStore {
.bind(optional_update_value(&update.plan_type))
.bind(update.status.is_some())
.bind(status)
.bind(status)
.bind(&now)
.bind(account_id)
.execute(&self.pool)
@@ -711,6 +691,69 @@ impl SqliteAccountStore {
Ok(result.rows_affected() > 0)
}
/// 同步运行时自然刷新出来的账号状态。
pub async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> SqliteAccountStoreResult<bool> {
let now = Utc::now().to_rfc3339();
let quota_limit_reached = if account.quota_limit_reached { 1 } else { 0 };
let quota_verify_required = if account.quota_verify_required { 1 } else { 0 };
let quota_cooldown_until = account.quota_cooldown_until.map(|dt| dt.to_rfc3339());
let cloudflare_cooldown_until = account.cloudflare_cooldown_until.map(|dt| dt.to_rfc3339());
let mut tx = self.pool.begin().await?;
let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL)
.bind(quota_limit_reached)
.bind(&now)
.bind(quota_limit_reached)
.bind(status_to_db(account.status))
.bind(quota_limit_reached)
.bind(&now)
.bind(quota_limit_reached)
.bind(quota_verify_required)
.bind(&now)
.bind(quota_verify_required)
.bind(&quota_cooldown_until)
.bind(&now)
.bind(&quota_cooldown_until)
.bind(&cloudflare_cooldown_until)
.bind(&now)
.bind(&cloudflare_cooldown_until)
.bind(&now)
.bind(&account.id)
.execute(&mut *tx)
.await?;
if result.rows_affected() == 0 {
tx.commit().await?;
return Ok(false);
}
if sync_usage_window {
sqlx::query(SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL)
.bind(&account.id)
.bind(u64_to_i64_saturating(account.window_request_count))
.bind(u64_to_i64_saturating(account.window_input_tokens))
.bind(u64_to_i64_saturating(account.window_output_tokens))
.bind(u64_to_i64_saturating(account.window_cached_tokens))
.bind(u64_to_i64_saturating(account.window_image_input_tokens))
.bind(u64_to_i64_saturating(account.window_image_output_tokens))
.bind(u64_to_i64_saturating(account.window_image_request_count))
.bind(u64_to_i64_saturating(
account.window_image_request_failed_count,
))
.bind(account.window_started_at.map(|dt| dt.to_rfc3339()))
.bind(account.window_reset_at.map(|dt| dt.to_rfc3339()))
.bind(account.limit_window_seconds.map(u64_to_i64_saturating))
.bind(&now)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
Ok(true)
}
/// 标记账号 Cloudflare 冷却期。
pub async fn set_cloudflare_cooldown_until(
&self,
@@ -734,12 +777,24 @@ impl SqliteAccountStore {
status: AccountStatus,
) -> SqliteAccountStoreResult<bool> {
let now = Utc::now().to_rfc3339();
let result = sqlx::query("update accounts set status = ?, updated_at = ? where id = ?")
.bind(status_to_db(status))
.bind(&now)
.bind(account_id)
.execute(&self.pool)
.await?;
let status = status_to_db(status);
let result = sqlx::query(
r"
update accounts
set
status = case
when ? = 'active' and quota_limit_reached = 1 then 'quota_exhausted'
else ?
end,
updated_at = ?
where id = ?",
)
.bind(status)
.bind(status)
.bind(&now)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
@@ -804,14 +859,6 @@ impl SqliteAccountStore {
Ok(())
}
/// 列出模型维度用量。
pub async fn list_model_usage(&self) -> SqliteAccountStoreResult<Vec<AccountModelUsageRecord>> {
let rows = sqlx::query(LIST_MODEL_USAGE_SQL)
.fetch_all(&self.pool)
.await?;
rows.iter().map(model_usage_from_row).collect()
}
/// 读取配额快照列表。
pub async fn list_quota_snapshots(
&self,
@@ -1054,16 +1101,6 @@ impl AccountStore for SqliteAccountStore {
.map_err(|error| map_account_store_error(&error))
}
async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> AccountStoreResult<bool> {
SqliteAccountStore::update_quota_json(self, account_id, quota_json)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn apply_quota_snapshot(
&self,
account_id: &str,
@@ -1082,6 +1119,16 @@ impl AccountStore for SqliteAccountStore {
.map_err(|error| map_account_store_error(&error))
}
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool> {
SqliteAccountStore::sync_runtime_account_state(self, account, sync_usage_window)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn sync_rate_limit_window(
&self,
account_id: &str,
+77 -15
View File
@@ -165,7 +165,13 @@ set
chatgpt_user_id = case when ? then ? else chatgpt_user_id end,
label = case when ? then ? else label end,
plan_type = case when ? then ? else plan_type end,
status = case when ? then ? else status end,
status = case
when ? then case
when ? = 'active' and quota_limit_reached = 1 then 'quota_exhausted'
else ?
end
else status
end,
updated_at = ?
where id = ?";
@@ -327,19 +333,6 @@ on conflict(account_id, model) do update set
cached_tokens = cached_tokens + excluded.cached_tokens,
last_used_at = excluded.last_used_at";
pub(super) const LIST_MODEL_USAGE_SQL: &str = r"
select
account_id,
model,
request_count,
error_count,
input_tokens,
output_tokens,
cached_tokens,
last_used_at
from account_model_usage
order by account_id asc, request_count desc, last_used_at desc, model asc";
pub(super) const LIST_QUOTA_SNAPSHOTS_SQL: &str = r"
select
id,
@@ -367,7 +360,7 @@ set
quota_fetched_at = ?,
plan_type = coalesce(?, plan_type),
status = case
when status in ('disabled', 'banned', 'expired', 'refreshing') then status
when status in ('disabled', 'banned', 'expired') then status
when ? = 1 then 'quota_exhausted'
when status = 'quota_exhausted' then 'active'
else status
@@ -426,12 +419,79 @@ on conflict(account_id) do update set
pub(super) const MARK_QUOTA_LIMITED_UNTIL_SQL: &str = r"
update accounts
set
status = case
when status in ('disabled', 'banned', 'expired') then status
else 'quota_exhausted'
end,
quota_limit_reached = 1,
quota_verify_required = 0,
quota_cooldown_until = ?,
updated_at = ?
where id = ?";
pub(super) const SYNC_RUNTIME_ACCOUNT_STATE_SQL: &str = r"
update accounts
set
status = case
when status in ('disabled', 'banned') then status
when (
case
when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached
else ?
end
) = 1 then 'quota_exhausted'
else ?
end,
quota_limit_reached = case
when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached
else ?
end,
quota_verify_required = case
when ? = 1 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_verify_required
else ?
end,
quota_cooldown_until = case
when ? is null and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_cooldown_until
else ?
end,
cloudflare_cooldown_until = case
when ? is null and cloudflare_cooldown_until is not null and cloudflare_cooldown_until > ? then cloudflare_cooldown_until
else ?
end,
updated_at = ?
where id = ?";
pub(super) const SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL: &str = r"
insert into account_usage (
account_id,
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
) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
on conflict(account_id) do update set
window_request_count = excluded.window_request_count,
window_input_tokens = excluded.window_input_tokens,
window_output_tokens = excluded.window_output_tokens,
window_cached_tokens = excluded.window_cached_tokens,
window_image_input_tokens = excluded.window_image_input_tokens,
window_image_output_tokens = excluded.window_image_output_tokens,
window_image_request_count = excluded.window_image_request_count,
window_image_request_failed_count = excluded.window_image_request_failed_count,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = excluded.limit_window_seconds
where account_usage.window_reset_at is null
or account_usage.window_reset_at <= excluded.window_reset_at
or account_usage.window_reset_at <= ?";
pub(super) const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r"
update accounts
set
@@ -452,6 +512,7 @@ set
next_refresh_at = ?,
status = case
when status in ('disabled', 'banned') then status
when quota_limit_reached = 1 then 'quota_exhausted'
else ?
end,
updated_at = ?
@@ -469,6 +530,7 @@ set
next_refresh_at = ?,
status = case
when status in ('disabled', 'banned') then status
when quota_limit_reached = 1 then 'quota_exhausted'
else ?
end,
updated_at = ?
+3 -29
View File
@@ -17,9 +17,8 @@ use crate::{
use super::{
queries::{GET_POOL_ACCOUNT_SQL, LIST_POOL_ACCOUNTS_SQL},
AccountModelUsageRecord, AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore,
SqliteAccountStoreError, SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata,
UsageDelta,
AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore, SqliteAccountStoreError,
SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata, UsageDelta,
};
pub(super) async fn list_pool_accounts(
@@ -230,23 +229,6 @@ pub(super) fn quota_plan_type(quota_json: &str) -> Option<String> {
.map(ToString::to_string)
}
pub(super) fn model_usage_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountModelUsageRecord> {
Ok(AccountModelUsageRecord {
account_id: row.get("account_id"),
model: row.get("model"),
request_count: row.get("request_count"),
error_count: row.get("error_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
last_used_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("last_used_at").as_deref(),
)?,
})
}
pub(super) fn quota_snapshot_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountQuotaSnapshot> {
@@ -261,14 +243,7 @@ pub(super) fn quota_snapshot_from_row(
}
pub(super) fn status_to_db(status: AccountStatus) -> &'static str {
match status {
AccountStatus::Active => "active",
AccountStatus::Expired => "expired",
AccountStatus::QuotaExhausted => "quota_exhausted",
AccountStatus::Refreshing => "refreshing",
AccountStatus::Disabled => "disabled",
AccountStatus::Banned => "banned",
}
status.as_str()
}
pub(super) fn optional_update_value(value: &Option<Option<String>>) -> Option<&str> {
@@ -280,7 +255,6 @@ pub(super) fn status_from_db(value: &str) -> SqliteAccountStoreResult<AccountSta
"active" => Ok(AccountStatus::Active),
"expired" => Ok(AccountStatus::Expired),
"quota_exhausted" => Ok(AccountStatus::QuotaExhausted),
"refreshing" => Ok(AccountStatus::Refreshing),
"disabled" => Ok(AccountStatus::Disabled),
"banned" => Ok(AccountStatus::Banned),
other => Err(SqliteAccountStoreError::InvalidStatus(other.to_string())),
@@ -1,7 +1,7 @@
//! 刷新租约存储、JWT 解码与账号 claims 验证。
use std::{
collections::hash_map::DefaultHasher,
collections::{hash_map::DefaultHasher, HashSet},
hash::{Hash, Hasher},
sync::{Arc, RwLock},
};
@@ -259,9 +259,6 @@ pub enum RefreshFailure {
/// 刷新令牌无效或已过期。
#[error("refresh token is invalid or expired")]
InvalidGrant,
/// 账号配额耗尽。
#[error("account quota is exhausted")]
QuotaExhausted,
/// 账号被上游封禁。
#[error("account is banned")]
Banned,
@@ -350,7 +347,6 @@ where
match self.client.refresh(refresh_token).await {
Ok(token_pair) => Ok(apply_token_pair(account, token_pair)),
Err(RefreshFailure::QuotaExhausted) => Err(RefreshError::Transport),
Err(RefreshFailure::RetryableTransport) => Err(RefreshError::RetryableTransport),
Err(RefreshFailure::Transport) => Err(RefreshError::Transport),
Err(error) => Ok(apply_refresh_failure(account, error)),
@@ -373,22 +369,53 @@ where
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> bool {
if account.status != AccountStatus::Active {
return false;
}
match trigger {
RefreshTrigger::Unauthorized => true,
RefreshTrigger::BeforeExpiry => account
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now + self.refresh_margin()),
RefreshTrigger::Unauthorized => can_refresh_after_auth_failure(account.status),
RefreshTrigger::BeforeExpiry => {
can_refresh_before_expiry(account.status)
&& self.should_refresh_before_expiry(account, now)
}
}
}
fn refresh_margin(&self) -> Duration {
fn should_refresh_before_expiry(&self, account: &Account, now: DateTime<Utc>) -> bool {
if let Some(next_refresh_at) = account.next_refresh_at {
return next_refresh_at <= now;
}
account
.access_token_expires_at
.is_some_and(|expires_at| self.refresh_at(account, expires_at) <= now)
}
fn refresh_at(&self, account: &Account, expires_at: DateTime<Utc>) -> DateTime<Utc> {
let policy = self.policy.snapshot();
let seconds = policy.refresh_margin_seconds.min(86_400 * 7) as i64;
Duration::seconds(seconds)
jittered_refresh_at(
account.id.as_str(),
expires_at,
policy.refresh_margin_seconds,
)
}
}
fn can_refresh_after_auth_failure(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Active | AccountStatus::Expired | AccountStatus::QuotaExhausted
)
}
fn can_refresh_before_expiry(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Active | AccountStatus::QuotaExhausted
)
}
fn status_after_successful_token_refresh(status: AccountStatus) -> AccountStatus {
match status {
AccountStatus::Expired => AccountStatus::Active,
status => status,
}
}
@@ -402,7 +429,7 @@ pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
refreshed.refresh_token = Some(refresh_token);
}
refreshed.status = AccountStatus::Active;
refreshed.status = status_after_successful_token_refresh(account.status);
refreshed
}
@@ -411,7 +438,6 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco
let mut updated = account.clone();
updated.status = match failure {
RefreshFailure::InvalidGrant => AccountStatus::Expired,
RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted,
RefreshFailure::Banned => AccountStatus::Banned,
RefreshFailure::RetryableTransport => AccountStatus::Active,
RefreshFailure::Transport => AccountStatus::Active,
@@ -423,7 +449,7 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco
// 刷新租约存储(SQLite)
// ---------------------------------------------------------------------------
use sqlx::SqlitePool;
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
use thiserror::Error;
/// SQLite 刷新租约存储结果。
@@ -490,6 +516,35 @@ where account_refresh_leases.expires_at <= ?
.await?;
Ok(result.rows_affected() > 0)
}
/// 返回给定账号集合中仍有效的刷新租约。
pub async fn active_account_ids(
&self,
account_ids: &[String],
now: DateTime<Utc>,
) -> RefreshLeaseStoreResult<HashSet<String>> {
if account_ids.is_empty() {
return Ok(HashSet::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(
"select account_id from account_refresh_leases where datetime(expires_at) > datetime(",
);
builder.push_bind(now.to_rfc3339());
builder.push(") and account_id in (");
let mut separated = builder.separated(", ");
for account_id in account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder
.build_query_as::<(String,)>()
.fetch_all(&self.pool)
.await?;
Ok(rows.into_iter().map(|(account_id,)| account_id).collect())
}
}
mod runtime;
@@ -127,6 +127,41 @@ where
self
}
/// 返回给定账号集合中当前正在执行 token 刷新的账号 ID。
pub async fn refreshing_account_ids(
&self,
account_ids: &[String],
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<HashSet<String>> {
if account_ids.is_empty() {
return Ok(HashSet::new());
}
let requested_ids = account_ids
.iter()
.map(String::as_str)
.collect::<HashSet<_>>();
let mut refreshing = self
.in_flight
.lock()
.await
.iter()
.filter(|account_id| requested_ids.contains(account_id.as_str()))
.cloned()
.collect::<HashSet<_>>();
if let Some(refresh_leases) = self.refresh_leases.as_ref() {
refreshing.extend(
refresh_leases
.active_account_ids(account_ids, now)
.await
.map_err(TokenRefreshServiceError::Lease)?,
);
}
Ok(refreshing)
}
/// 执行一次账号级刷新定时器调度。
pub async fn schedule_account_timers_once(
&self,
@@ -165,7 +200,6 @@ where
}
let timer = match account.status {
AccountStatus::Refreshing => Some((Duration::ZERO, RefreshTrigger::Unauthorized)),
AccountStatus::Expired => {
let delay = match account.next_refresh_at {
Some(next_refresh_at) if next_refresh_at > now => {
@@ -543,7 +577,7 @@ where
account: &Account,
now: DateTime<Utc>,
) -> Option<RefreshTrigger> {
if account.status == AccountStatus::Refreshing {
if account.status == AccountStatus::Expired {
return Some(RefreshTrigger::Unauthorized);
}
@@ -562,7 +596,7 @@ where
let mut permanent_failures = 0;
for attempt_index in 0..max_attempts {
let Some(attempt) = self.prepare_refresh_attempt(account, now).await? else {
let Some(attempt) = self.prepare_refresh_attempt(account, trigger, now).await? else {
return Ok(TokenRefreshOutcome::Skipped);
};
@@ -577,32 +611,34 @@ where
.refresh_account_at(&attempt, trigger, now)
.await
{
Ok(mut updated) if account_requires_token_write(&attempt, &updated) => {
updated.access_token_expires_at =
jwt_expiration(&updated.access_token).or(updated.access_token_expires_at);
updated.next_refresh_at = self.computed_refresh_at(&updated);
persist_token_update(&self.store, &updated).await?;
return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated)));
}
Ok(updated) if attempt.status != updated.status => {
if is_permanent_refresh_status(updated.status) {
permanent_failures += 1;
if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD
&& self.sleep_before_retry(&attempt.id, attempt_index).await
{
continue;
}
Ok(updated) if is_permanent_refresh_failure_status(updated.status) => {
permanent_failures += 1;
if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD
&& self.sleep_before_retry(&attempt.id, attempt_index).await
{
continue;
}
persist_status(&self.store, &attempt.id, updated.status).await?;
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
Ok(_) => return Ok(TokenRefreshOutcome::Skipped),
Ok(mut updated) => {
updated.access_token_expires_at =
jwt_expiration(&updated.access_token).or(updated.access_token_expires_at);
updated.next_refresh_at = self.computed_refresh_at(&updated);
persist_token_update(&self.store, &updated).await?;
return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated)));
}
Err(RefreshError::RetryableTransport | RefreshError::Transport)
if self.sleep_before_retry(&attempt.id, attempt_index).await => {}
Err(RefreshError::RetryableTransport | RefreshError::Transport) => {
persist_status(&self.store, &attempt.id, AccountStatus::Active).await?;
persist_status(
&self.store,
&attempt.id,
status_after_temporary_refresh_failure(attempt.status),
)
.await?;
let recovery_delay = stable_jittered_duration(
&attempt.id,
Duration::from_secs(RECOVERY_DELAY_SECONDS as u64),
@@ -624,6 +660,7 @@ where
async fn prepare_refresh_attempt(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<Option<Account>> {
let Some(stored) = self.store.get(&account.id).await? else {
@@ -631,7 +668,7 @@ where
account.id.clone(),
));
};
let mut attempt = stored_account_to_refresh_account(stored);
let attempt = stored_account_to_refresh_account(stored);
if matches!(
attempt.status,
AccountStatus::Disabled | AccountStatus::Banned
@@ -641,11 +678,13 @@ where
return Ok(None);
}
if attempt.status != AccountStatus::Refreshing {
persist_status(&self.store, &attempt.id, AccountStatus::Refreshing).await?;
if !self
.scheduler
.should_refresh_account_at(&attempt, trigger, now)
{
return Ok(None);
}
attempt.status = AccountStatus::Active;
Ok(Some(attempt))
}
@@ -839,20 +878,16 @@ async fn persist_next_refresh_at(
}
}
fn account_requires_token_write(previous: &Account, updated: &Account) -> bool {
previous.access_token != updated.access_token
|| previous.refresh_token != updated.refresh_token
|| (updated.status == AccountStatus::Active && previous.status != updated.status)
fn is_permanent_refresh_failure_status(status: AccountStatus) -> bool {
matches!(status, AccountStatus::Expired | AccountStatus::Banned)
}
fn is_permanent_refresh_status(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Expired
| AccountStatus::QuotaExhausted
| AccountStatus::Disabled
| AccountStatus::Banned
)
fn status_after_temporary_refresh_failure(status: AccountStatus) -> AccountStatus {
match status {
AccountStatus::Expired | AccountStatus::QuotaExhausted => status,
AccountStatus::Active => AccountStatus::Active,
AccountStatus::Disabled | AccountStatus::Banned => status,
}
}
fn refresh_lease_owner() -> String {
+63 -44
View File
@@ -1,6 +1,10 @@
//! 指纹类型、更新与持久化。
use std::{path::PathBuf, sync::Arc, time::Duration};
use std::{
path::PathBuf,
sync::{Arc, RwLock},
time::Duration,
};
use chrono::{DateTime, Utc};
use indexmap::IndexMap;
@@ -8,7 +12,7 @@ use serde::{Deserialize, Serialize};
use sqlx::{Row, SqlitePool};
use thiserror::Error;
use tokio::sync::Mutex;
use tracing::{info, warn};
use tracing::info;
use uuid::Uuid;
/// 运行时当前指纹槽位。
@@ -17,7 +21,7 @@ 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_hours(72);
pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_hours(72);
type AppcastFields = (Option<String>, Option<String>, Option<String>);
@@ -89,6 +93,38 @@ impl Fingerprint {
}
}
/// 运行时共享指纹快照。
#[derive(Debug, Clone)]
pub struct RuntimeFingerprint {
current: Arc<RwLock<Arc<Fingerprint>>>,
}
impl RuntimeFingerprint {
pub fn new(fingerprint: Fingerprint) -> Self {
Self {
current: Arc::new(RwLock::new(Arc::new(fingerprint))),
}
}
pub fn current(&self) -> Arc<Fingerprint> {
self.current
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
pub fn snapshot(&self) -> Fingerprint {
self.current().as_ref().clone()
}
pub fn replace(&self, fingerprint: Fingerprint) {
*self
.current
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::new(fingerprint);
}
}
// ---------------------------------------------------------------------------
// 指纹历史记录与仓储
// ---------------------------------------------------------------------------
@@ -345,7 +381,7 @@ fn json_error(error: serde_json::Error) -> sqlx::Error {
/// Appcast 指纹更新检查器。
#[derive(Clone)]
pub struct UpdateChecker {
repository: Option<FingerprintRepository>,
repository: FingerprintRepository,
client: reqwest::Client,
appcast_url: String,
extracted_fingerprint_path: PathBuf,
@@ -360,7 +396,7 @@ struct InternalState {
impl UpdateChecker {
/// 使用自定义 client 和路径构造检查器。
pub fn with_client(
repository: Option<FingerprintRepository>,
repository: FingerprintRepository,
client: reqwest::Client,
appcast_url: impl Into<String>,
extracted_fingerprint_path: PathBuf,
@@ -406,18 +442,18 @@ impl UpdateChecker {
}
/// 检查并在需要时应用更新。
pub async fn check_and_apply_update(&self) -> Result<bool, UpdateError> {
pub async fn check_and_apply_update(&self) -> Result<Option<Fingerprint>, UpdateError> {
let update_state = self.check_for_update().await?;
if !update_state.update_available {
return Ok(false);
return Ok(None);
}
let (Some(version), Some(build)) = (
update_state.latest_version.as_deref(),
update_state.latest_build.as_deref(),
) else {
return Ok(false);
return Ok(None);
};
let state = self.state.lock().await;
@@ -433,36 +469,13 @@ impl UpdateChecker {
"发现新的 fingerprint 版本"
);
self.apply_version_update(version, build).await?;
let updated_fingerprint = self.apply_version_update(version, build).await?;
let mut state = self.state.lock().await;
state.current_version = version.to_string();
state.current_build = build.to_string();
Ok(true)
}
/// 启动后台轮询任务。
pub fn start_background_checker(self) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
let mut ticker = tokio::time::interval(POLL_INTERVAL);
info!(
interval_secs = POLL_INTERVAL.as_secs(),
"fingerprint 后台版本检查器已启动"
);
if let Err(error) = self.check_and_apply_update().await {
warn!(error = %error, "fingerprint 首次版本检查失败");
}
loop {
ticker.tick().await;
if let Err(error) = self.check_and_apply_update().await {
warn!(error = %error, "fingerprint 定期版本检查失败");
}
}
})
Ok(Some(updated_fingerprint))
}
async fn fetch_appcast(&self) -> Result<String, UpdateError> {
@@ -480,20 +493,26 @@ impl UpdateChecker {
response.text().await.map_err(UpdateError::from)
}
async fn apply_version_update(&self, version: &str, build: &str) -> Result<(), UpdateError> {
async fn apply_version_update(
&self,
version: &str,
build: &str,
) -> Result<Fingerprint, UpdateError> {
let chromium_version =
load_matching_chromium_version(&self.extracted_fingerprint_path, version, build);
if let Some(repository) = &self.repository {
repository
.update_current_version(version, build, chromium_version.as_deref())
.await?;
repository
.insert_update_history(version, build, chromium_version.as_deref(), None)
.await?;
}
Ok(())
self.repository
.update_current_version(version, build, chromium_version.as_deref())
.await?;
self.repository
.insert_update_history(version, build, chromium_version.as_deref(), None)
.await?;
let updated = self
.repository
.load_current()
.await?
.ok_or(sqlx::Error::RowNotFound)?;
Ok(updated)
}
}
+1
View File
@@ -4,5 +4,6 @@ pub mod accounts;
pub mod fingerprint;
pub mod models;
pub mod protocol;
pub mod scheduler;
pub mod token_client;
pub mod transport;
+8 -89
View File
@@ -9,19 +9,6 @@ use super::{
snapshot::ModelPlanSnapshot,
};
/// 解析模型名后得到的标准化结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedModelName {
/// 最终模型 ID。
pub model_id: String,
/// 推理强度后缀。
pub reasoning_effort: Option<String>,
/// 服务层级后缀。
pub service_tier: Option<String>,
}
const SERVICE_TIER_SUFFIXES: [&str; 2] = ["fast", "flex"];
const REASONING_EFFORT_SUFFIXES: [&str; 6] = ["none", "minimal", "low", "medium", "high", "xhigh"];
const BUILTIN_CODEX_MODELS: [(&str, &str, bool); 8] = [
("gpt-5.5", "GPT-5.5", true),
("gpt-5.4", "GPT-5.4", false),
@@ -170,10 +157,9 @@ impl ModelCatalog {
.cloned()
}
/// 按对外模型名查询模型信息,支持别名和已知后缀。
/// 按对外模型名查询模型信息,支持别名映射。
pub fn model_info_for_name(&self, input: &str) -> Option<CodexModelInfo> {
let parsed = self.parse_model_name(input);
self.model_info(&parsed.model_id)
self.model_info(&self.resolve_model_id(input))
}
/// 返回 model -> plans allowlist。
@@ -192,54 +178,14 @@ impl ModelCatalog {
if trimmed.is_empty() {
return false;
}
if self.model_aliases.contains_key(trimmed) || self.model_info(trimmed).is_some() {
return true;
}
let stripped = strip_known_model_suffixes(trimmed);
if stripped.model_name == trimmed
|| (stripped.reasoning_effort.is_none() && stripped.service_tier.is_none())
{
return false;
}
self.model_aliases.contains_key(&stripped.model_name)
|| self.model_info(&stripped.model_name).is_some()
self.model_aliases.contains_key(trimmed) || self.model_info(trimmed).is_some()
}
/// 解析外部传入的模型名,提取别名、推理强度和服务层级后缀。
pub fn parse_model_name(&self, input: &str) -> ParsedModelName {
let trimmed = input.trim();
if self.model_aliases.contains_key(trimmed) || self.model_info(trimmed).is_some() {
return ParsedModelName {
model_id: self.resolve_model_id(trimmed),
reasoning_effort: None,
service_tier: None,
};
}
let stripped = strip_known_model_suffixes(trimmed);
ParsedModelName {
model_id: self.resolve_model_id(&stripped.model_name),
reasoning_effort: stripped.reasoning_effort,
service_tier: stripped.service_tier,
}
}
/// 将标准化模型名重新拼成展示名。
pub fn build_display_model_name(parsed: &ParsedModelName) -> String {
let mut name = parsed.model_id.clone();
if let Some(reasoning_effort) = &parsed.reasoning_effort {
name.push('-');
name.push_str(reasoning_effort);
}
if let Some(service_tier) = &parsed.service_tier {
name.push('-');
name.push_str(service_tier);
}
name
}
fn resolve_model_id(&self, input: &str) -> String {
/// 沿 alias 链将外部模型名解析为真实 model ID。
///
/// 模型名只承载 model 身份;reasoning effort、service tier 等运行参数由请求
/// body 携带并原样透传上游。
pub fn resolve_model_id(&self, input: &str) -> String {
self.resolve_alias_chain(input)
}
@@ -260,33 +206,6 @@ impl ModelCatalog {
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct StrippedModelName {
model_name: String,
reasoning_effort: Option<String>,
service_tier: Option<String>,
}
fn strip_known_model_suffixes(input: &str) -> StrippedModelName {
let mut remaining = input.trim().to_string();
let service_tier = take_suffix(&mut remaining, &SERVICE_TIER_SUFFIXES);
let reasoning_effort = take_suffix(&mut remaining, &REASONING_EFFORT_SUFFIXES);
StrippedModelName {
model_name: remaining,
reasoning_effort,
service_tier,
}
}
fn take_suffix(remaining: &mut String, suffixes: &[&str]) -> Option<String> {
let suffix = suffixes
.iter()
.find(|suffix| remaining.ends_with(&format!("-{suffix}")))?;
let truncate_to = remaining.len() - suffix.len() - 1;
remaining.truncate(truncate_to);
Some((*suffix).to_string())
}
fn normalize_aliases(input: &BTreeMap<String, String>) -> BTreeMap<String, String> {
input
.iter()
+14 -8
View File
@@ -43,11 +43,17 @@ pub enum ModelServiceError {
#[error("no active accounts available for model refresh")]
NoAccounts,
/// 快照写入失败。
#[error("failed to store model snapshot")]
StoreSnapshot,
#[error("failed to store model snapshot: {source}")]
StoreSnapshot {
#[source]
source: ModelSnapshotStoreError,
},
/// 刷新后重新读取快照失败。
#[error("failed to load model snapshots")]
LoadSnapshots,
#[error("failed to load model snapshots: {source}")]
LoadSnapshots {
#[source]
source: ModelSnapshotStoreError,
},
/// 所有计划都刷新失败。
#[error("all model refresh plans failed")]
AllPlansFailed(ModelRefreshResult),
@@ -251,10 +257,10 @@ impl ModelService {
}
}
fn map_store_snapshot_error(_: ModelSnapshotStoreError) -> ModelServiceError {
ModelServiceError::StoreSnapshot
fn map_store_snapshot_error(source: ModelSnapshotStoreError) -> ModelServiceError {
ModelServiceError::StoreSnapshot { source }
}
fn map_load_snapshots_error(_: ModelSnapshotStoreError) -> ModelServiceError {
ModelServiceError::LoadSnapshots
fn map_load_snapshots_error(source: ModelSnapshotStoreError) -> ModelServiceError {
ModelServiceError::LoadSnapshots { source }
}
@@ -7,7 +7,7 @@ use thiserror::Error;
use super::{info::CodexModelInfo, snapshot::ModelPlanSnapshot};
/// 模型快照存储错误。
#[derive(Debug, Error)]
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ModelSnapshotStoreError {
/// 底层存储失败。
#[error("model snapshot store operation failed: {message}")]
+4 -5
View File
@@ -540,11 +540,10 @@ fn quota_snapshot_from_windows(
if primary.is_none() && secondary.is_none() {
return None;
}
let blocked = limit_reached.unwrap_or_else(|| {
allowed.is_some_and(|allowed| !allowed)
|| primary.is_some_and(|window| window.used_percent >= 100.0)
|| secondary.is_some_and(|window| window.used_percent >= 100.0)
});
let blocked = limit_reached.unwrap_or(false)
|| allowed.is_some_and(|allowed| !allowed)
|| primary.is_some_and(|window| window.used_percent >= 100.0)
|| secondary.is_some_and(|window| window.used_percent >= 100.0);
Some(json!({
"source": source,
"limit_name": limit_name,
+245 -189
View File
@@ -1,112 +1,71 @@
use std::{collections::BTreeSet, time::Instant};
use serde::{Deserialize, Serialize};
use serde::Serialize;
use serde_json::{json, Map, Value};
use crate::{
infra::time::elapsed_millis_i64,
upstream::{
models::catalog::ParsedModelName,
protocol::{
schema::reconvert_tuple_values,
sse::{parse_sse_events, SseError, SseEvent},
},
upstream::protocol::{
schema::reconvert_tuple_values,
sse::{parse_sse_events, SseError, SseEvent},
},
};
/// Codex Responses 上游请求体。
#[derive(Debug, Clone, Serialize, Deserialize)]
///
/// 发往上游的 Responses 请求。`body` 持有客户端原始 JSON object,逐字段(含顺序、
/// 含未知字段)透传上游,是上游请求体的唯一来源;`use_websocket`/`force_http_sse`
/// 仅用于本地传输选择,不写入 body。常用字段通过访问器方法读写。
/// 代理只做最小 patch(`stream`/`store` 由传输层控制,`input` 字符串做官方兼容转换,
/// implicit resume 时替换 `input`,模型后缀路由时写 `model`/`reasoning`/`service_tier`)。
///
/// 其余字段是代理控制状态,不进上游 body(原 `#[serde(skip)]` 字段)。
#[derive(Debug, Clone)]
pub struct CodexResponsesRequest {
/// 模型名。
pub model: String,
/// 指令文本。
pub instructions: String,
/// 输入消息与结构化条目。
pub input: Vec<Value>,
/// 是否流式返回。
pub stream: bool,
/// 是否要求上游存储响应。
pub store: bool,
#[serde(skip_serializing_if = "Option::is_none")]
/// reasoning 配置。
pub reasoning: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 工具定义。
pub tools: Option<Vec<Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 工具选择策略。
pub tool_choice: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 是否允许并行工具调用。
pub parallel_tool_calls: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 输出文本格式配置。
pub text: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 图片生成开关。
pub generate: Option<bool>,
#[serde(skip)]
/// 上游请求体(唯一真相源)。
body: Map<String, Value>,
/// tuple schema 原始定义,仅供响应重构时使用。
pub tuple_schema: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// service tier。
pub service_tier: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 用于显式续链的前一个 response ID。
pub previous_response_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 提示缓存键。
pub prompt_cache_key: Option<String>,
#[serde(skip)]
/// 是否由客户端显式提供了 prompt cache key。
pub explicit_prompt_cache_key: bool,
#[serde(skip)]
/// 客户端会话 ID。
pub client_conversation_id: Option<String>,
#[serde(skip)]
/// 变体身份键。
pub variant_identity: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
/// include 列表。
pub include: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 传给上游的 client metadata。
pub client_metadata: Option<Value>,
#[serde(skip)]
/// 代理侧识别的客户端 IP,仅用于管理端使用记录展示。
pub client_ip: Option<String>,
#[serde(skip)]
/// 客户端 User-Agent,仅用于管理端使用记录展示。
pub client_user_agent: Option<String>,
#[serde(skip)]
/// 是否偏好 WebSocket 传输。
pub use_websocket: bool,
#[serde(skip)]
/// 是否强制 HTTP SSE。
pub force_http_sse: bool,
#[serde(skip)]
/// turn state 透传头。
pub turn_state: Option<String>,
#[serde(skip)]
/// turn metadata 透传头。
pub turn_metadata: Option<String>,
#[serde(skip)]
/// beta features 透传头。
pub beta_features: Option<String>,
#[serde(skip)]
/// 客户端版本头。
pub version: Option<String>,
#[serde(skip)]
/// timing metrics 透传头。
pub include_timing_metrics: Option<String>,
#[serde(skip)]
/// codex window id。
pub codex_window_id: Option<String>,
#[serde(skip)]
/// 父线程 ID。
pub parent_thread_id: Option<String>,
}
impl Serialize for CodexResponsesRequest {
/// 上游 body 序列化即原始 `body` map(HTTP SSE 直发;WebSocket 在外层前置 `type`)。
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.body.serialize(serializer)
}
}
/// Codex Responses 请求的上游传输决策。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodexTransport {
@@ -124,7 +83,7 @@ pub fn transport_for_request(request: &CodexResponsesRequest) -> CodexTransport
return CodexTransport::HttpSse;
}
if request.previous_response_id.is_some() {
if request.previous_response_id().is_some() {
return CodexTransport::WebSocketRequired;
}
@@ -567,112 +526,22 @@ fn failure_code(value: &Value) -> Option<String> {
.map(ToString::to_string)
}
/// 将模型名后缀和模型配置应用到 Codex Responses 上游请求。
pub fn apply_response_model_options(
request: &mut CodexResponsesRequest,
parsed_model: &ParsedModelName,
) {
request.model = parsed_model.model_id.clone();
apply_reasoning_options(request, parsed_model);
apply_service_tier_options(request, parsed_model);
}
fn apply_reasoning_options(request: &mut CodexResponsesRequest, parsed_model: &ParsedModelName) {
let existing_reasoning = request.reasoning.take();
let existing_object = match existing_reasoning {
Some(Value::Object(object)) => object,
Some(_) | None => Map::new(),
};
let effort = existing_object
.get("effort")
.and_then(Value::as_str)
.map(ToString::to_string)
.or_else(|| non_empty_string(parsed_model.reasoning_effort.as_deref()));
if effort.is_none() && existing_object.is_empty() {
request.reasoning = None;
return;
}
let summary = existing_object
.get("summary")
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.unwrap_or("auto");
let mut reasoning = Map::new();
reasoning.insert("summary".to_string(), Value::String(summary.to_string()));
if let Some(effort) = effort {
reasoning.insert("effort".to_string(), Value::String(effort));
}
request.reasoning = Some(Value::Object(reasoning));
ensure_reasoning_include(request);
}
pub(crate) fn ensure_reasoning_include(request: &mut CodexResponsesRequest) {
if request.reasoning.is_none() {
return;
}
if request
.include
.as_ref()
.is_some_and(|include| !include.is_empty())
{
return;
}
request.include = Some(vec!["reasoning.encrypted_content".to_string()]);
}
fn apply_service_tier_options(request: &mut CodexResponsesRequest, parsed_model: &ParsedModelName) {
request.service_tier = request
.service_tier
.take()
.and_then(|value| non_empty_string(Some(&value)))
.or_else(|| non_empty_string(parsed_model.service_tier.as_deref()))
.map(normalize_service_tier_for_upstream);
}
fn normalize_service_tier_for_upstream(service_tier: String) -> String {
if service_tier == "fast" {
"priority".to_string()
} else {
service_tier
}
}
fn non_empty_string(value: Option<&str>) -> Option<String> {
value.and_then(|value| {
let trimmed = value.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
})
}
impl CodexResponsesRequest {
/// 构造默认的 HTTP SSE 请求。
pub fn new_http_sse(
model: impl Into<String>,
instructions: impl Into<String>,
input: Vec<Value>,
) -> Self {
/// 从客户端原始 Responses JSON object 构造上游请求。
///
/// 只做最小规范化:`input` 缺省为 `[]`、`stream` 缺省 `true`、`store` 缺省 `false`。
/// 其余字段(含未知字段)原样保留在 `body` 中透传上游。
pub fn from_body(mut body: Map<String, Value>) -> Self {
body.entry("input")
.or_insert_with(|| Value::Array(Vec::new()));
body.entry("stream").or_insert(Value::Bool(true));
body.entry("store").or_insert(Value::Bool(false));
Self {
model: model.into(),
instructions: instructions.into(),
input,
stream: true,
store: false,
reasoning: None,
tools: None,
tool_choice: None,
parallel_tool_calls: None,
text: None,
generate: None,
body,
tuple_schema: None,
service_tier: None,
previous_response_id: None,
prompt_cache_key: None,
explicit_prompt_cache_key: false,
client_conversation_id: None,
variant_identity: None,
include: None,
client_metadata: None,
client_ip: None,
client_user_agent: None,
use_websocket: false,
@@ -687,9 +556,180 @@ impl CodexResponsesRequest {
}
}
/// 构造默认的 HTTP SSE 请求(测试与内部构造用)。
pub fn new_http_sse(
model: impl Into<String>,
instructions: impl Into<String>,
input: Vec<Value>,
) -> Self {
let mut body = Map::new();
body.insert("model".to_string(), Value::String(model.into()));
body.insert(
"instructions".to_string(),
Value::String(instructions.into()),
);
body.insert("input".to_string(), Value::Array(input));
Self::from_body(body)
}
/// 上游 body 的只读视图。
pub fn body(&self) -> &Map<String, Value> {
&self.body
}
// --- body 字段类型化访问器(上游语义字段,透传不重写)---
/// 模型名。
pub fn model(&self) -> &str {
self.body
.get("model")
.and_then(Value::as_str)
.unwrap_or_default()
}
/// 设置模型名(模型后缀路由归一)。
pub fn set_model(&mut self, model: impl Into<String>) {
self.body
.insert("model".to_string(), Value::String(model.into()));
}
/// 指令文本(缺省空串)。
pub fn instructions(&self) -> &str {
self.body
.get("instructions")
.and_then(Value::as_str)
.unwrap_or_default()
}
/// 输入条目切片(非数组时为空)。
pub fn input(&self) -> &[Value] {
self.body
.get("input")
.and_then(Value::as_array)
.map_or(&[], Vec::as_slice)
}
/// 替换输入条目(implicit resume / 字符串输入兼容)。
pub fn set_input(&mut self, input: Vec<Value>) {
self.body.insert("input".to_string(), Value::Array(input));
}
/// 是否流式返回。
pub fn stream(&self) -> bool {
self.body
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(true)
}
/// 设置流式标志。
pub fn set_stream(&mut self, stream: bool) {
self.body.insert("stream".to_string(), Value::Bool(stream));
}
/// 是否要求上游存储响应。
pub fn store(&self) -> bool {
self.body
.get("store")
.and_then(Value::as_bool)
.unwrap_or(false)
}
/// reasoning 配置(透传,不规整)。
pub fn reasoning(&self) -> Option<&Value> {
self.body.get("reasoning")
}
/// 工具定义数组(非数组或空时 None)。
pub fn tools(&self) -> Option<&[Value]> {
self.body
.get("tools")
.and_then(Value::as_array)
.filter(|tools| !tools.is_empty())
.map(Vec::as_slice)
}
/// include 列表(透传原值)。
pub fn include(&self) -> Option<&Value> {
self.body.get("include")
}
/// service tier(透传原值)。
pub fn service_tier(&self) -> Option<&str> {
self.body.get("service_tier").and_then(Value::as_str)
}
/// 设置 service tier(模型后缀路由 / 归一)。
pub fn set_service_tier(&mut self, service_tier: Option<String>) {
match service_tier {
Some(value) => {
self.body
.insert("service_tier".to_string(), Value::String(value));
}
None => {
self.body.remove("service_tier");
}
}
}
/// 前一个 response ID。
pub fn previous_response_id(&self) -> Option<&str> {
self.body
.get("previous_response_id")
.and_then(Value::as_str)
}
/// 设置 / 清除前一个 response ID。
pub fn set_previous_response_id(&mut self, previous_response_id: Option<String>) {
match previous_response_id {
Some(value) => {
self.body
.insert("previous_response_id".to_string(), Value::String(value));
}
None => {
self.body.remove("previous_response_id");
}
}
}
/// 提示缓存键。
pub fn prompt_cache_key(&self) -> Option<&str> {
self.body.get("prompt_cache_key").and_then(Value::as_str)
}
/// 设置提示缓存键。
pub fn set_prompt_cache_key(&mut self, prompt_cache_key: Option<String>) {
match prompt_cache_key {
Some(value) => {
self.body
.insert("prompt_cache_key".to_string(), Value::String(value));
}
None => {
self.body.remove("prompt_cache_key");
}
}
}
/// client metadata(透传原值)。
pub fn client_metadata(&self) -> Option<&Value> {
self.body.get("client_metadata")
}
/// 设置 / 合并 client metadata。
pub fn set_client_metadata(&mut self, client_metadata: Option<Value>) {
match client_metadata {
Some(value) => {
self.body.insert("client_metadata".to_string(), value);
}
None => {
self.body.remove("client_metadata");
}
}
}
/// 判断请求是否声明了图片生成工具。
pub fn expects_image_generation(&self) -> bool {
self.tools.as_deref().is_some_and(|tools| {
self.tools().is_some_and(|tools| {
tools
.iter()
.any(|tool| tool.get("type").and_then(Value::as_str) == Some("image_generation"))
@@ -698,30 +738,46 @@ impl CodexResponsesRequest {
}
/// Codex compact 端点请求体。
#[derive(Debug, Clone, Serialize, Deserialize)]
///
/// 发往上游的 Responses compact 请求。`body` 持有客户端原始 JSON object
/// (已剥离 compact 上游不接受的字段),逐字段透传上游;
/// `client_ip`/`client_user_agent` 仅供管理端使用记录展示,不进上游 body。
#[derive(Debug, Clone)]
pub struct CodexCompactRequest {
/// 模型名。
pub model: String,
/// 输入消息与结构化条目。
pub input: Vec<Value>,
/// 指令文本。
pub instructions: String,
#[serde(skip_serializing_if = "Option::is_none")]
/// 工具定义。
pub tools: Option<Vec<Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 是否允许并行工具调用。
pub parallel_tool_calls: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
/// reasoning 配置。
pub reasoning: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 输出文本格式配置。
pub text: Option<Value>,
#[serde(skip)]
/// 上游请求体(唯一真相源)。
pub body: Map<String, Value>,
/// 代理侧识别的客户端 IP,仅用于管理端使用记录展示。
pub client_ip: Option<String>,
#[serde(skip)]
/// 客户端 User-Agent,仅用于管理端使用记录展示。
pub client_user_agent: Option<String>,
}
impl Serialize for CodexCompactRequest {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.body.serialize(serializer)
}
}
impl CodexCompactRequest {
/// 模型名。
pub fn model(&self) -> &str {
self.body
.get("model")
.and_then(Value::as_str)
.unwrap_or_default()
}
/// 设置模型名(模型后缀路由归一)。
pub fn set_model(&mut self, model: impl Into<String>) {
self.body
.insert("model".to_string(), Value::String(model.into()));
}
/// reasoning 配置(透传原值)。
pub fn reasoning(&self) -> Option<&Value> {
self.body.get("reasoning")
}
}
-124
View File
@@ -2,24 +2,6 @@
use serde_json::{Map, Value};
pub(crate) struct PreparedSchema {
pub(crate) schema: Value,
pub(crate) original_schema: Option<Value>,
}
pub(crate) fn prepare_schema(schema: Value) -> PreparedSchema {
let original_schema = has_tuple_schemas(&schema).then(|| schema.clone());
let mut schema = schema;
if original_schema.is_some() {
convert_tuple_schemas(&mut schema);
}
inject_additional_properties(&mut schema);
PreparedSchema {
schema,
original_schema,
}
}
pub(crate) fn reconvert_tuple_values(data: Value, schema: &Value) -> Value {
reconvert_tuple_values_with_root(data, schema, schema)
}
@@ -145,109 +127,3 @@ fn schema_array_has_tuple_schema(object: &Map<String, Value>, key: &str) -> bool
};
values.iter().any(has_tuple_schemas)
}
fn convert_tuple_schemas(value: &mut Value) {
let Value::Object(object) = value else {
return;
};
if let Some(Value::Array(prefix_items)) = object.remove("prefixItems") {
let mut properties = Map::new();
let mut required = Vec::new();
for (index, mut item) in prefix_items.into_iter().enumerate() {
convert_tuple_schemas(&mut item);
let key = index.to_string();
properties.insert(key.clone(), item);
required.push(Value::String(key));
}
object.insert("type".to_string(), Value::String("object".to_string()));
object.insert("properties".to_string(), Value::Object(properties));
object.insert("required".to_string(), Value::Array(required));
object.insert("additionalProperties".to_string(), Value::Bool(false));
object.remove("items");
return;
}
convert_object_values(object, "properties");
convert_object_value(object, "items");
convert_schema_array(object, "oneOf");
convert_schema_array(object, "anyOf");
convert_schema_array(object, "allOf");
convert_object_values(object, "$defs");
convert_object_values(object, "definitions");
convert_object_value(object, "if");
convert_object_value(object, "then");
convert_object_value(object, "else");
convert_object_value(object, "not");
}
fn convert_object_value(object: &mut Map<String, Value>, key: &str) {
if let Some(value) = object.get_mut(key) {
convert_tuple_schemas(value);
}
}
fn convert_object_values(object: &mut Map<String, Value>, key: &str) {
let Some(Value::Object(values)) = object.get_mut(key) else {
return;
};
for value in values.values_mut() {
convert_tuple_schemas(value);
}
}
fn convert_schema_array(object: &mut Map<String, Value>, key: &str) {
let Some(Value::Array(values)) = object.get_mut(key) else {
return;
};
for value in values {
convert_tuple_schemas(value);
}
}
fn inject_additional_properties(value: &mut Value) {
let Value::Object(object) = value else {
return;
};
if object.get("type").and_then(Value::as_str) == Some("object")
&& !object.contains_key("additionalProperties")
{
object.insert("additionalProperties".to_string(), Value::Bool(false));
}
inject_object_values(object, "properties");
inject_object_value(object, "items");
inject_schema_array(object, "oneOf");
inject_schema_array(object, "anyOf");
inject_schema_array(object, "allOf");
inject_object_values(object, "$defs");
inject_object_values(object, "definitions");
inject_object_value(object, "if");
inject_object_value(object, "then");
inject_object_value(object, "else");
inject_object_value(object, "not");
}
fn inject_object_value(object: &mut Map<String, Value>, key: &str) {
if let Some(value) = object.get_mut(key) {
inject_additional_properties(value);
}
}
fn inject_object_values(object: &mut Map<String, Value>, key: &str) {
let Some(Value::Object(values)) = object.get_mut(key) else {
return;
};
for value in values.values_mut() {
inject_additional_properties(value);
}
}
fn inject_schema_array(object: &mut Map<String, Value>, key: &str) {
let Some(Value::Array(values)) = object.get_mut(key) else {
return;
};
for value in values {
inject_additional_properties(value);
}
}
+30 -209
View File
@@ -1,4 +1,4 @@
use serde::{ser::SerializeMap, Deserialize, Serialize};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::upstream::protocol::responses::{
@@ -48,17 +48,6 @@ pub struct WebSocketAuditArtifact {
pub opening: Option<OpeningAuditSnapshot>,
/// 首个 `response.create` payload 快照。
pub payload: Option<PayloadAuditSnapshot>,
/// 失败快照。
pub error: Option<WebSocketAuditErrorSnapshot>,
}
/// WebSocket audit error snapshot.
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct WebSocketAuditErrorSnapshot {
/// 错误分类。
pub classification: String,
/// 错误消息。
pub message: String,
}
/// 从 WebSocket 错误帧推导出的上游错误分类。
@@ -153,31 +142,13 @@ fn is_internal_websocket_event(event: &str) -> bool {
event == "codex.rate_limits"
}
/// 判断 WebSocket 事件是否应在转发前剥离。
///
/// 仅剥离传输层内部帧:`response.metadata` 承载 turn_state,由上层提取转存到
/// 会话状态,不下发客户端。业务事件一律原样转发,不做 schema 校验。
/// (`codex.rate_limits` 内部事件由 `is_internal_websocket_event` 单独剥离。)
fn websocket_event_should_skip(raw: &str) -> bool {
websocket_metadata_turn_state(raw).is_some()
|| websocket_event_shape_parse_error(raw)
|| websocket_response_completed_missing_response(raw)
|| websocket_response_created_missing_response(raw)
|| websocket_response_output_text_delta_missing_delta(raw)
|| websocket_delta_event_missing_official_required_fields(raw)
|| websocket_output_item_event_missing_item(raw)
|| websocket_output_item_event_non_object_item(raw)
|| websocket_output_item_event_invalid_item_type_tag(raw)
|| websocket_output_item_event_invalid_metadata(raw)
|| websocket_message_output_item_event_invalid_required_fields(raw)
|| websocket_agent_message_output_item_event_invalid_required_fields(raw)
|| websocket_reasoning_output_item_event_invalid_required_fields(raw)
|| websocket_function_call_output_item_event_invalid_required_fields(raw)
|| websocket_function_call_output_payload_item_event_invalid_required_fields(raw)
|| websocket_custom_tool_call_output_item_event_invalid_required_fields(raw)
|| websocket_custom_tool_call_output_payload_item_event_invalid_required_fields(raw)
|| websocket_tool_search_call_output_item_event_invalid_required_fields(raw)
|| websocket_tool_search_output_item_event_invalid_required_fields(raw)
|| websocket_local_shell_call_output_item_event_invalid_required_fields(raw)
|| websocket_web_search_call_output_item_event_invalid_required_fields(raw)
|| websocket_image_generation_call_output_item_event_invalid_required_fields(raw)
|| websocket_compaction_output_item_event_invalid_required_fields(raw)
|| websocket_reasoning_summary_part_added_missing_summary_index(raw)
}
/// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。
@@ -863,6 +834,17 @@ pub fn is_terminal_websocket_event(event: &str) -> bool {
matches!(event, "response.completed" | "response.failed" | "error")
}
/// 判断是否为「首个内容帧之前的纯生命周期帧」。
///
/// `response.created` / `response.in_progress` 仅表示上游已受理请求,尚未开始产出
/// 真实内容(含 reasoning)。首 token 熔断计时不应被它们解除,否则连接一收到
/// `response.created`(通常 1s 内)就被误判为「已开始产出」,使真正的首内容帧
/// 可以无限拖延而不触发保护。其余任何帧(reasoning/delta/output_item 等)到达即视为
/// 上游已开始真实产出,解除熔断——对慢思考模型友好。
pub fn is_pre_content_lifecycle_event(event: &str) -> bool {
matches!(event, "response.created" | "response.in_progress")
}
/// 将 WebSocket 错误帧映射为上游 HTTP 状态分类。
pub fn classify_websocket_error_frame(raw: &str) -> Option<ClassifiedWebSocketError> {
let value = serde_json::from_str::<Value>(raw).ok()?;
@@ -978,7 +960,6 @@ pub fn websocket_audit_artifact_from_attempt(
fallback_allowed: http_sse_fallback_allowed(request),
opening: Some(opening),
payload: Some(payload),
error: None,
}
}
@@ -991,195 +972,35 @@ fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static st
}
/// 生成 Responses WebSocket `response.create` payload。
///
/// payload = `{"type": "response.create"}` 加上原始上游 body 的全部字段
/// (保持插入顺序,含未知字段),逐字段原样透传。
pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Value {
let mut body = Map::new();
insert_value(
&mut body,
"type",
let mut payload = Map::new();
payload.insert(
"type".to_string(),
Value::String("response.create".to_string()),
);
insert_value(&mut body, "model", Value::String(request.model.clone()));
insert_value(
&mut body,
"instructions",
Value::String(request.instructions.clone()),
);
insert_value(&mut body, "input", Value::Array(request.input.clone()));
insert_value(&mut body, "store", Value::Bool(request.store));
insert_value(&mut body, "stream", Value::Bool(request.stream));
if let Some(previous_response_id) = &request.previous_response_id {
insert_value(
&mut body,
"previous_response_id",
Value::String(previous_response_id.clone()),
);
for (key, value) in request.body() {
payload.insert(key.clone(), value.clone());
}
if let Some(reasoning) = &request.reasoning {
insert_value(&mut body, "reasoning", reasoning.clone());
}
if let Some(tools) = non_empty_tools(request) {
insert_value(&mut body, "tools", Value::Array(tools.to_vec()));
}
insert_value(
&mut body,
"tool_choice",
request
.tool_choice
.clone()
.unwrap_or(Value::String("auto".to_string())),
);
insert_value(
&mut body,
"parallel_tool_calls",
Value::Bool(request.parallel_tool_calls.unwrap_or(true)),
);
if let Some(text) = &request.text {
insert_value(&mut body, "text", text.clone());
}
if let Some(service_tier) = &request.service_tier {
insert_value(
&mut body,
"service_tier",
Value::String(service_tier.clone()),
);
}
if let Some(prompt_cache_key) = &request.prompt_cache_key {
insert_value(
&mut body,
"prompt_cache_key",
Value::String(prompt_cache_key.clone()),
);
}
if let Some(include) = non_empty_include(request) {
insert_value(
&mut body,
"include",
include
.iter()
.cloned()
.map(Value::String)
.collect::<Vec<_>>()
.into(),
);
}
if let Some(client_metadata) = &request.client_metadata {
insert_value(&mut body, "client_metadata", client_metadata.clone());
}
Value::Object(body)
Value::Object(payload)
}
/// 生成 Responses WebSocket `response.create` 文本帧内容。
pub fn websocket_response_create_payload_text(
request: &CodexResponsesRequest,
) -> Result<String, serde_json::Error> {
serde_json::to_string(&OrderedResponseCreatePayload(request))
}
struct OrderedResponseCreatePayload<'a>(&'a CodexResponsesRequest);
impl Serialize for OrderedResponseCreatePayload<'_> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let request = self.0;
let mut map = serializer.serialize_map(Some(websocket_payload_keys(request).len()))?;
map.serialize_entry("type", "response.create")?;
map.serialize_entry("model", &request.model)?;
map.serialize_entry("instructions", &request.instructions)?;
map.serialize_entry("input", &request.input)?;
map.serialize_entry("store", &request.store)?;
map.serialize_entry("stream", &request.stream)?;
if let Some(previous_response_id) = &request.previous_response_id {
map.serialize_entry("previous_response_id", previous_response_id)?;
}
if let Some(reasoning) = &request.reasoning {
map.serialize_entry("reasoning", reasoning)?;
}
if let Some(tools) = non_empty_tools(request) {
map.serialize_entry("tools", tools)?;
}
if let Some(tool_choice) = &request.tool_choice {
map.serialize_entry("tool_choice", tool_choice)?;
} else {
map.serialize_entry("tool_choice", "auto")?;
}
map.serialize_entry(
"parallel_tool_calls",
&request.parallel_tool_calls.unwrap_or(true),
)?;
if let Some(text) = &request.text {
map.serialize_entry("text", text)?;
}
if let Some(service_tier) = &request.service_tier {
map.serialize_entry("service_tier", service_tier)?;
}
if let Some(prompt_cache_key) = &request.prompt_cache_key {
map.serialize_entry("prompt_cache_key", prompt_cache_key)?;
}
if let Some(include) = non_empty_include(request) {
map.serialize_entry("include", include)?;
}
if let Some(client_metadata) = &request.client_metadata {
map.serialize_entry("client_metadata", client_metadata)?;
}
map.end()
}
}
fn insert_value(body: &mut Map<String, Value>, key: &str, value: Value) {
body.insert(key.to_string(), value);
serde_json::to_string(&websocket_response_create_payload(request))
}
fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec<String> {
let mut keys = vec![
"type".to_string(),
"model".to_string(),
"instructions".to_string(),
"input".to_string(),
"store".to_string(),
"stream".to_string(),
];
if request.previous_response_id.is_some() {
keys.push("previous_response_id".to_string());
}
if request.reasoning.is_some() {
keys.push("reasoning".to_string());
}
if non_empty_tools(request).is_some() {
keys.push("tools".to_string());
}
keys.push("tool_choice".to_string());
keys.push("parallel_tool_calls".to_string());
if request.text.is_some() {
keys.push("text".to_string());
}
if request.service_tier.is_some() {
keys.push("service_tier".to_string());
}
if request.prompt_cache_key.is_some() {
keys.push("prompt_cache_key".to_string());
}
if non_empty_include(request).is_some() {
keys.push("include".to_string());
}
if request.client_metadata.is_some() {
keys.push("client_metadata".to_string());
}
let mut keys = Vec::with_capacity(request.body().len() + 1);
keys.push("type".to_string());
keys.extend(request.body().keys().cloned());
keys
}
fn non_empty_tools(request: &CodexResponsesRequest) -> Option<&[Value]> {
request.tools.as_deref().filter(|tools| !tools.is_empty())
}
fn non_empty_include(request: &CodexResponsesRequest) -> Option<&[String]> {
request
.include
.as_deref()
.filter(|include| !include.is_empty())
}
fn redact_payload_body(body: Value) -> Value {
let Value::Object(body) = body else {
return body;
@@ -0,0 +1,137 @@
//! 候选账号过滤。
//!
//! 所有调度策略共享的前置:从账号池中筛出「本次请求可用」的候选集合。这是选择流程
//! 的第一步,与具体策略无关——策略只在过滤后的候选之间做选择。
//!
//! 过滤规则(全部满足才入选):
//!
//! 1. 账号状态为 [`AccountStatus::Active`];
//! 2. 在途槽位数未达单账号并发上限;
//! 3. 配额可用([`AccountService::quota_available_at`],受 `skip_quota_limited` 控制);
//! 4. 模型在账号订阅计划的允许列表内;
//! 5. 不在本次请求的排除列表内;
//! 6. 不处于 Cloudflare 冷却期。
//!
//! 之后再按订阅层级优先级(`tier_priority`)收窄到最高可用层级;若请求指定了
//! `required_account_id`,则只保留该账号(跳过层级收窄)。
use chrono::{DateTime, Utc};
use crate::upstream::accounts::model::{Account, AccountStatus};
use crate::upstream::accounts::service::AccountService;
/// 候选过滤所需的账号池策略参数(只读视图)。
pub struct CandidateFilter<'a> {
/// 单账号最大并发请求数。
pub max_concurrent_per_account: usize,
/// 是否跳过已标记配额耗尽的账号。
pub skip_quota_limited: bool,
/// 订阅层级优先级,越靠前优先级越高。
pub tier_priority: &'a [String],
/// 模型到允许订阅计划的映射。
pub model_plan_allowlist: &'a std::collections::BTreeMap<String, Vec<String>>,
/// 已成功拉取过模型列表的订阅计划。
pub fetched_model_plan_types: &'a std::collections::BTreeSet<String>,
/// 按账号 ID 读取当前在途槽位数。
pub slot_count: &'a dyn Fn(&str) -> usize,
}
/// 本次账号获取请求的过滤条件。
pub struct CandidateRequest<'a> {
/// 本次请求使用的模型。
pub model: &'a str,
/// 本次请求需要排除的账号 ID。
pub exclude_account_ids: &'a [String],
/// 本次请求必须使用的账号 ID(存在时跳过层级收窄,只保留该账号)。
pub required_account_id: Option<&'a str>,
/// 本次调度使用的当前时间。
pub now: DateTime<Utc>,
}
/// 从账号池筛出本次请求的可用候选集合。
pub fn filter<'a>(
accounts: impl Iterator<Item = &'a Account>,
filter: &CandidateFilter<'_>,
request: &CandidateRequest<'_>,
) -> Vec<Account> {
let mut candidates = accounts
.filter(|account| is_base_available(account, filter, request))
.cloned()
.collect::<Vec<_>>();
if let Some(required_account_id) = request.required_account_id {
candidates.retain(|account| account.id == required_account_id);
return candidates;
}
if let Some(best_tier) = best_available_tier(&candidates, filter.tier_priority) {
candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str()));
}
candidates
}
/// 判断账号是否满足全部基础可用条件。
fn is_base_available(
account: &Account,
filter: &CandidateFilter<'_>,
request: &CandidateRequest<'_>,
) -> bool {
account.status == AccountStatus::Active
&& (filter.slot_count)(&account.id) < filter.max_concurrent_per_account
&& AccountService::quota_available_at(account, request.now, filter.skip_quota_limited)
&& is_model_allowed(
account,
request.model,
filter.model_plan_allowlist,
filter.fetched_model_plan_types,
)
&& !request
.exclude_account_ids
.iter()
.any(|account_id| account_id == &account.id)
&& AccountService::cloudflare_available_at(account, request.now)
}
/// 判断模型是否允许该账号的订阅计划使用。
fn is_model_allowed(
account: &Account,
model: &str,
model_plan_allowlist: &std::collections::BTreeMap<String, Vec<String>>,
fetched_model_plan_types: &std::collections::BTreeSet<String>,
) -> bool {
let Some(allowed_plans) = model_plan_allowlist.get(model) else {
return true;
};
let Some(plan_type) = account.plan_type.as_deref() else {
return false;
};
allowed_plans.iter().any(|plan| plan == plan_type)
|| !fetched_model_plan_types.contains(plan_type)
}
/// 在候选集合中找出订阅层级优先级最高的可用层级。
fn best_available_tier(candidates: &[Account], tier_priority: &[String]) -> Option<String> {
for tier in tier_priority {
let matched = candidates
.iter()
.any(|account| account.plan_type.as_deref() == Some(tier.as_str()));
if matched {
return Some(tier.clone());
}
}
None
}
/// 判断账号是否可用于模型列表刷新(不涉及排除列表 / 模型允许列表 / 层级收窄)。
pub fn is_model_refresh_available(
account: &Account,
max_concurrent_per_account: usize,
skip_quota_limited: bool,
slot_count: &dyn Fn(&str) -> usize,
now: DateTime<Utc>,
) -> bool {
account.status == AccountStatus::Active
&& slot_count(&account.id) < max_concurrent_per_account
&& AccountService::quota_available_at(account, now, skip_quota_limited)
&& AccountService::cloudflare_available_at(account, now)
}
+171
View File
@@ -0,0 +1,171 @@
//! 账号运行时反馈:错误率与 TTFT 的 EWMA 统计。
//!
//! 每个账号维护两个指数加权移动平均(EWMA):错误率与首 token 延迟(TTFT,
//! 毫秒)。请求结束后经 [`FeedbackStats::report`] 回灌,直接进 [`super::score`]
//! 打分并可驱动 sticky escape。统计以 `AtomicU64`(存 `f64` 位模式)+ CAS 无锁更新,
//! 对齐 sub2api 的设计,避免选择热路径上的锁与 DB 往返。
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::RwLock;
/// EWMA 平滑系数(新样本权重)。与 sub2api 一致取 0.2。
const EWMA_ALPHA: f64 = 0.2;
/// 单账号的运行时 EWMA 统计。
///
/// `*_bits` 存 `f64` 的位模式;`NaN`(全 1 的哨兵)表示"尚无样本"。
#[derive(Debug)]
struct AccountFeedback {
/// 错误率 EWMA,范围 `[0, 1]`。
error_rate_bits: AtomicU64,
/// TTFT EWMA,单位毫秒。
ttft_ms_bits: AtomicU64,
}
/// 表示"尚无样本"的哨兵位模式。
const EMPTY_BITS: u64 = f64::NAN.to_bits();
impl Default for AccountFeedback {
fn default() -> Self {
Self {
error_rate_bits: AtomicU64::new(EMPTY_BITS),
ttft_ms_bits: AtomicU64::new(EMPTY_BITS),
}
}
}
impl AccountFeedback {
/// 读取当前错误率 EWMA;无样本时返回 `None`。
fn error_rate(&self) -> Option<f64> {
load_sample(&self.error_rate_bits)
}
/// 读取当前 TTFT EWMA(毫秒);无样本时返回 `None`。
fn ttft_ms(&self) -> Option<f64> {
load_sample(&self.ttft_ms_bits)
}
/// 回灌一次请求结果。
fn report(&self, success: bool, first_token_ms: Option<u64>) {
let error_sample = if success { 0.0 } else { 1.0 };
update_ewma(&self.error_rate_bits, error_sample);
if let Some(ms) = first_token_ms {
update_ewma(&self.ttft_ms_bits, ms as f64);
}
}
}
/// 读取 EWMA 样本;哨兵(NaN)返回 `None`。
fn load_sample(bits: &AtomicU64) -> Option<f64> {
let value = f64::from_bits(bits.load(Ordering::Relaxed));
if value.is_nan() {
None
} else {
Some(value)
}
}
/// 无锁 CAS 更新 EWMA:首个样本直接落值,其后按 alpha 混合。
fn update_ewma(bits: &AtomicU64, sample: f64) {
let mut current = bits.load(Ordering::Relaxed);
loop {
let previous = f64::from_bits(current);
let next = if previous.is_nan() {
sample
} else {
EWMA_ALPHA * sample + (1.0 - EWMA_ALPHA) * previous
};
match bits.compare_exchange_weak(
current,
next.to_bits(),
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => return,
Err(observed) => current = observed,
}
}
}
/// 账号 EWMA 反馈的读数快照,供打分使用。
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct FeedbackSample {
/// 错误率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
pub error_rate: Option<f64>,
/// TTFT EWMA(毫秒);无样本时为 `None`。
pub ttft_ms: Option<f64>,
}
/// 账号运行时反馈存储:按账号 ID 聚合 EWMA 统计。
///
/// 选择热路径只读(`sample`),请求结束后写(`report`)。内部按账号分片为独立的
/// 原子对象,`RwLock` 仅保护 map 结构本身(插入新账号时短暂写锁)。
#[derive(Debug, Default)]
pub struct FeedbackStats {
accounts: RwLock<HashMap<String, AccountFeedback>>,
}
impl FeedbackStats {
/// 创建空的反馈存储。
pub fn new() -> Self {
Self::default()
}
/// 读取指定账号的 EWMA 快照;账号无记录时返回全 `None`。
pub fn sample(&self, account_id: &str) -> FeedbackSample {
let accounts = self
.accounts
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match accounts.get(account_id) {
Some(feedback) => FeedbackSample {
error_rate: feedback.error_rate(),
ttft_ms: feedback.ttft_ms(),
},
None => FeedbackSample::default(),
}
}
/// 回灌一次请求结果:成功/失败 + 首 token 延迟(毫秒)。
pub fn report(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
// 快路径:账号已存在时只需读锁 + 原子更新。
{
let accounts = self
.accounts
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(feedback) = accounts.get(account_id) {
feedback.report(success, first_token_ms);
return;
}
}
// 慢路径:首次见到该账号,取写锁插入后回灌。
let mut accounts = self
.accounts
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
accounts
.entry(account_id.to_string())
.or_default()
.report(success, first_token_ms);
}
/// 移除指定账号的反馈记录(账号被删除时调用)。
pub fn remove(&self, account_id: &str) {
let mut accounts = self
.accounts
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
accounts.remove(account_id);
}
/// 清空全部反馈记录(账号池清空时调用)。
pub fn clear(&self) {
let mut accounts = self
.accounts
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
accounts.clear();
}
}
+112
View File
@@ -0,0 +1,112 @@
//! 智能账号调度模块。
//!
//! 账号选择的唯一归属地。把「选择」从账号池存储中剥离,集中于此:
//!
//! - [`candidates`] —— 候选过滤(可用性 / 层级 / 模型允许列表),所有策略共享的前置。
//! - [`strategy`] —— 全部选择策略,每个策略一个文件([`strategy::smart`] /
//! [`strategy::quota_reset`] / [`strategy::round_robin`] / [`strategy::sticky`])。
//! - [`feedback`] —— 跨策略的运行时 EWMA 反馈(错误率 / TTFT),喂给 Smart 打分。
//!
//! [`AccountScheduler`] 是对外门面:持有轮转游标、打分权重与 EWMA 反馈存储,接收
//! **已过滤的候选切片**并返回选中账号。它不 lock 账号池、不刷状态、不写库——存储与
//! 状态维护仍归 `AccountPool`,二者经候选切片与 `slot_count` 闭包解耦。
pub mod candidates;
pub mod feedback;
pub mod strategy;
use std::sync::Mutex;
use chrono::{DateTime, Utc};
use crate::upstream::accounts::model::Account;
pub use candidates::{CandidateFilter, CandidateRequest};
pub use feedback::{FeedbackSample, FeedbackStats};
pub use strategy::smart::rank_candidates;
pub use strategy::{RotationStrategy, ScoreBreakdown, ScoreWeights};
/// 账号调度器:账号选择的对外门面。
///
/// 持有轮转游标(并列破并列用)、打分权重(Smart 用)与 EWMA 反馈存储。选择热路径
/// 只读反馈与权重,请求结束后经 [`AccountScheduler::report_feedback`] 回灌。游标用
/// [`Mutex`] 保护(选择需可变,但 `AccountScheduler` 以共享引用参与并发请求)。
#[derive(Debug)]
pub struct AccountScheduler {
/// 轮转游标:在并列最优账号之间做公平轮转。
cursor: Mutex<usize>,
/// Smart 策略的打分权重。
weights: ScoreWeights,
/// 账号运行时 EWMA 反馈(错误率 / TTFT)。
feedback: FeedbackStats,
}
impl Default for AccountScheduler {
fn default() -> Self {
Self::new(ScoreWeights::default())
}
}
impl AccountScheduler {
/// 使用指定打分权重创建调度器。
pub fn new(weights: ScoreWeights) -> Self {
Self {
cursor: Mutex::new(0),
weights,
feedback: FeedbackStats::new(),
}
}
/// 替换打分权重(策略切换或热更时调用)。
pub fn set_weights(&mut self, weights: ScoreWeights) {
self.weights = weights;
}
/// 重置轮转游标(策略切换时调用,保证轮转从头开始)。
pub fn reset_cursor(&self) {
*self.lock_cursor() = 0;
}
/// 从已过滤的候选集合中按指定策略选择一个账号。
///
/// `candidates` 必须是 [`candidates::filter`] 的输出;`slot_count` 按账号 ID 读取
/// 当前在途槽位数(Smart 打分与候选无关的运行时状态由调用方提供)。
pub fn select(
&self,
strategy: RotationStrategy,
candidates: &[Account],
slot_count: &dyn Fn(&str) -> usize,
now: DateTime<Utc>,
) -> Option<Account> {
let input = strategy::SelectionInput {
candidates,
slot_count,
feedback: &self.feedback,
weights: &self.weights,
now,
};
let mut cursor = self.lock_cursor();
strategy::select(strategy, &input, &mut cursor)
}
/// 回灌一次请求结果到运行时 EWMA 反馈(错误率 / TTFT),供 Smart 打分。
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
self.feedback.report(account_id, success, first_token_ms);
}
/// 移除指定账号的运行时反馈(账号被删除时调用)。
pub fn forget_feedback(&self, account_id: &str) {
self.feedback.remove(account_id);
}
/// 清空全部运行时反馈(账号池清空时调用)。
pub fn clear_feedback(&self) {
self.feedback.clear();
}
fn lock_cursor(&self) -> std::sync::MutexGuard<'_, usize> {
self.cursor
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
@@ -0,0 +1,142 @@
//! 账号调度策略层。
//!
//! 所有账号选择策略集中于此,每个策略一个文件,互不依赖:
//!
//! - [`smart`] —— 加权线性打分 + EWMA 反馈(“智能调度”)。
//! - [`quota_reset`] —— 优先消耗即将重置的额度窗口。
//! - [`round_robin`] —— 按候选顺序循环。
//! - [`sticky`] —— 优先最近使用过的账号。
//!
//! 本模块只承担三件事:定义 [`RotationStrategy`] 枚举、把选择请求分派到对应策略、
//! 提供各策略共享的 tie-break 比较原语([`select_by`] / [`compare_window_reset`] /
//! [`compare_last_used`])。策略本身不触碰账号池存储:在途槽位数经 `slot_count`
//! 闭包按需读取,EWMA 反馈与打分权重经 [`SelectionInput`] 传入。
use std::cmp::Ordering;
use chrono::{DateTime, Utc};
use crate::infra::time::parse_rfc3339_utc;
use crate::upstream::accounts::model::Account;
use crate::upstream::scheduler::feedback::FeedbackStats;
pub mod quota_reset;
pub mod round_robin;
pub mod smart;
pub mod sticky;
pub use smart::{ScoreBreakdown, ScoreWeights};
/// 账号轮转策略。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RotationStrategy {
/// 加权线性打分调度:多因子(负载/窗口/错误率/TTFT/重置)归一化后加权求和,
/// 配合 EWMA 反馈。
Smart,
/// 优先消耗即将重置的额度窗口。
QuotaResetPriority,
/// 按候选账号顺序循环选择。
RoundRobin,
/// 优先选择最近使用过的账号。
Sticky,
}
impl RotationStrategy {
/// 返回策略的持久化/日志字符串。
pub fn as_str(self) -> &'static str {
match self {
Self::Smart => "smart",
Self::QuotaResetPriority => "quota_reset_priority",
Self::RoundRobin => "round_robin",
Self::Sticky => "sticky",
}
}
}
/// 策略选择的运行时输入。
///
/// `slot_count` 按账号 ID 读取当前在途槽位数;`feedback` / `weights` 仅 [`smart`]
/// 使用,其余策略忽略;`now` 用于窗口重置相关的打分。
pub struct SelectionInput<'a> {
/// 已过滤的候选账号集合。
pub candidates: &'a [Account],
/// 按账号 ID 读取当前在途槽位数。
pub slot_count: &'a dyn Fn(&str) -> usize,
/// 账号 EWMA 反馈存储(Smart 打分用)。
pub feedback: &'a FeedbackStats,
/// 打分权重(Smart 打分用)。
pub weights: &'a ScoreWeights,
/// 当前时刻(窗口重置因子用)。
pub now: DateTime<Utc>,
}
/// 按指定策略从候选集合中选择一个账号。
///
/// `cursor` 是账号池持有的轮转游标,用于在并列最优账号之间做公平轮转;不同策略对它
/// 的读写方式不同(详见各策略模块),本函数只负责透传可变引用。
pub fn select(
strategy: RotationStrategy,
input: &SelectionInput<'_>,
cursor: &mut usize,
) -> Option<Account> {
match strategy {
RotationStrategy::Smart => smart::select(input, cursor),
RotationStrategy::QuotaResetPriority => quota_reset::select(input, cursor),
RotationStrategy::RoundRobin => round_robin::select(input, cursor),
RotationStrategy::Sticky => sticky::select(input, cursor),
}
}
// ====================================================================
// 共享 tie-break 原语
// ====================================================================
/// 按比较函数升序排序后,在并列最优(`Ordering::Equal`)的账号之间用轮转游标破并列。
///
/// 这是 [`smart`] 与 [`quota_reset`] 共享的选择骨架:先按各自的比较规则排序,再在得分
/// 相同的账号之间轮转,保证并列账号被公平地依次选中。
pub(crate) fn select_by(
candidates: &[Account],
cursor: &mut usize,
compare: impl Fn(&Account, &Account) -> Ordering,
) -> Option<Account> {
let mut sorted = candidates.iter().collect::<Vec<_>>();
sorted.sort_by(|a, b| compare(a, b));
let best = *sorted.first()?;
let tied_count = sorted
.iter()
.take_while(|account| compare(best, account) == Ordering::Equal)
.count();
let index = *cursor % tied_count;
*cursor = cursor.wrapping_add(1);
Some((*sorted[index]).clone())
}
/// 计算账号当前窗口的总 token 数(输入 + 输出 + 图片输入 + 图片输出)。
pub(crate) fn account_window_token_count(account: &Account) -> u64 {
account
.window_input_tokens
.saturating_add(account.window_output_tokens)
.saturating_add(account.window_image_input_tokens)
.saturating_add(account.window_image_output_tokens)
}
/// 比较两个窗口重置时间:越早重置越优先;缺失重置时间视为并列(不惩罚)。
pub(crate) 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) | (None, Some(_)) => Ordering::Equal,
}
}
/// 比较两个账号的最近使用时间:越久未用越优先(升序,最小者排前)。
pub(crate) fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering {
last_used_millis(a).cmp(&last_used_millis(b))
}
fn last_used_millis(value: Option<&str>) -> i64 {
value
.and_then(|value| parse_rfc3339_utc(value).ok())
.map(|datetime| datetime.timestamp_millis())
.unwrap_or(0)
}
@@ -0,0 +1,29 @@
//! QuotaResetPriority 策略:优先消耗即将重置的额度窗口。
//!
//! 排序规则(升序,越靠前越优先):
//! 1. 未封顶配额优先于已封顶(`quota_limit_reached` false < true);
//! 2. 窗口重置时间越早越优先;
//! 3. 历史请求数越少越优先;
//! 4. 最近使用时间越久越优先(LRU)。
//!
//! 并列最优的账号之间由 [`super::select_by`] 用轮转游标破并列。
use std::cmp::Ordering;
use crate::upstream::accounts::model::Account;
use super::{compare_last_used, compare_window_reset, select_by, SelectionInput};
/// QuotaResetPriority 策略选择。
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
select_by(input.candidates, cursor, compare_quota_reset_priority)
}
/// 配额重置优先级比较:升序,最优(最应先用)排在最前。
fn compare_quota_reset_priority(a: &Account, b: &Account) -> Ordering {
a.quota_limit_reached
.cmp(&b.quota_limit_reached)
.then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at))
.then_with(|| a.request_count.cmp(&b.request_count))
.then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref()))
}
@@ -0,0 +1,20 @@
//! RoundRobin 策略:按候选顺序循环选择。
//!
//! 最简单的策略:忽略账号负载与用量,纯粹按候选集合的顺序用轮转游标依次选中,
//! 保证请求在候选账号间均匀轮转。
use crate::upstream::accounts::model::Account;
use super::SelectionInput;
/// 按候选顺序做纯轮转选择。
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
let candidates = input.candidates;
if candidates.is_empty() {
return None;
}
*cursor %= candidates.len();
let index = *cursor;
*cursor = cursor.wrapping_add(1);
Some(candidates[index].clone())
}
@@ -0,0 +1,291 @@
//! Smart 策略:加权线性打分 + EWMA 反馈。
//!
//! 参考 sub2api 的调度设计:把多个归一化因子(负载、窗口 token 用量、窗口缓存、窗口
//! 请求数、错误率、TTFT、窗口重置)按可配置权重线性求和,得分越高越优先。相较字典序
//! 元组比较,加权打分是更平滑的负载均衡,不会让单一维度无条件压制其余维度,同时通过
//! EWMA 反馈感知账号健康度与延迟。
//!
//! 各归一化因子都落在 `[0, 1]` 且方向统一(越大越好),因此权重之间可直接比较。配额
//! 已封顶(`quota_limit_reached`)的账号会被施加一个强惩罚,确保只要还有未封顶的候选就
//! 不会选中它。并列最高分的账号之间由 [`super::select_by`] 用轮转游标破并列。
use crate::upstream::accounts::model::Account;
use crate::upstream::scheduler::feedback::{FeedbackSample, FeedbackStats};
use super::{account_window_token_count, SelectionInput};
/// 配额封顶账号的惩罚分。远大于各因子加权和的理论上限(各因子 ∈ [0,1],权重和通常为
/// 个位数),保证未封顶账号总是优先于封顶账号。
const QUOTA_LIMITED_PENALTY: f64 = 1_000.0;
/// 打分权重。每个字段对应一个归一化因子的权重,权重为 0 即关闭该维度。
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ScoreWeights {
/// 在途槽位压力(越空闲越优先)。
pub load: f64,
/// 当前窗口实际 token 用量(用得越少越优先)。
pub window: f64,
/// 当前窗口缓存 token 用量(缓存越少越优先,权重小,仅在实际用量相当时区分)。
pub cached: f64,
/// 当前窗口请求数(请求越少越优先)。
pub window_requests: f64,
/// 错误率 EWMA(越健康越优先)。
pub error_rate: f64,
/// 首 token 延迟 EWMA(越快越优先)。
pub ttft: f64,
/// 窗口重置临近度(越快重置越优先,use-it-or-lose-it,默认关闭)。
pub reset: f64,
}
impl Default for ScoreWeights {
fn default() -> Self {
// 默认权重对齐 sub2api 的取向:负载与健康度为主,窗口用量次之,TTFT 再次,
// 窗口重置默认关闭(避免过早榨干即将重置的账号)。cached 权重远小于 window,
// 确保实际 token 用量主导、缓存仅作细粒度区分。
Self {
load: 1.0,
window: 0.8,
cached: 0.05,
window_requests: 0.6,
error_rate: 1.0,
ttft: 0.5,
reset: 0.0,
}
}
}
/// 单个候选账号的打分明细,供日志与可观测使用。
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ScoreBreakdown {
/// 各因子加权后的总分(已含配额封顶惩罚)。
pub total: f64,
/// 负载因子(归一化后,未乘权重)。
pub load: f64,
/// 窗口实际 token 用量因子。
pub window: f64,
/// 窗口缓存 token 用量因子。
pub cached: f64,
/// 窗口请求数因子。
pub window_requests: f64,
/// 错误率因子。
pub error_rate: f64,
/// TTFT 因子。
pub ttft: f64,
/// 窗口重置因子。
pub reset: f64,
/// 是否被施加配额封顶惩罚。
pub quota_limited: bool,
}
/// Smart 策略选择:按加权得分降序排序,在并列最高分之间用轮转游标破并列。
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
let ranked = rank_candidates(
input.candidates,
input.weights,
input.now.timestamp(),
input.slot_count,
input.feedback,
);
let best_total = ranked.first()?.1.total;
// 并列最高分的账号数量(浮点相等用位相等判定,得分由相同输入算出,稳定)。
let tied_count = ranked
.iter()
.take_while(|(_, breakdown)| breakdown.total.to_bits() == best_total.to_bits())
.count();
let index = *cursor % tied_count;
*cursor = cursor.wrapping_add(1);
let account_index = ranked[index].0;
Some(input.candidates[account_index].clone())
}
/// 对候选集合按加权得分排序(降序),返回 `(账号索引, 打分明细)` 列表。
///
/// 只做打分与排序,最终的并列破并列交给 [`select`]。公开供测试直接断言排序结果。
pub fn rank_candidates(
candidates: &[Account],
weights: &ScoreWeights,
now_secs: i64,
slot_count: &dyn Fn(&str) -> usize,
feedback: &FeedbackStats,
) -> Vec<(usize, ScoreBreakdown)> {
let samples = candidates
.iter()
.map(|account| feedback.sample(&account.id))
.collect::<Vec<_>>();
let ctx = NormalizationContext::build(candidates, &samples, now_secs, slot_count);
let mut ranked = candidates
.iter()
.zip(&samples)
.enumerate()
.map(|(index, (account, sample))| {
(
index,
score_account(account, sample, weights, &ctx, now_secs, slot_count),
)
})
.collect::<Vec<_>>();
// 降序:得分高在前。NaN 不应出现(因子均有界),保守用 total_cmp。
ranked.sort_by(|a, b| b.1.total.total_cmp(&a.1.total));
ranked
}
/// 候选集合的归一化上下文:预先算好各维度的 min/max,供逐账号归一化。
struct NormalizationContext {
max_slot_count: f64,
max_window_tokens: f64,
max_window_cached: f64,
max_window_requests: f64,
min_ttft_ms: f64,
max_ttft_ms: f64,
min_reset_secs: f64,
max_reset_secs: f64,
}
impl NormalizationContext {
fn build(
candidates: &[Account],
samples: &[FeedbackSample],
now_secs: i64,
slot_count: &dyn Fn(&str) -> usize,
) -> Self {
let mut max_slot_count = 0.0_f64;
let mut max_window_tokens = 0.0_f64;
let mut max_window_cached = 0.0_f64;
let mut max_window_requests = 0.0_f64;
let mut min_ttft_ms = f64::INFINITY;
let mut max_ttft_ms = 0.0_f64;
let mut min_reset_secs = f64::INFINITY;
let mut max_reset_secs = 0.0_f64;
for (account, sample) in candidates.iter().zip(samples) {
max_slot_count = max_slot_count.max(slot_count(&account.id) as f64);
max_window_tokens = max_window_tokens.max(account_window_token_count(account) as f64);
max_window_cached = max_window_cached.max(account.window_cached_tokens as f64);
max_window_requests = max_window_requests.max(account.window_request_count as f64);
if let Some(ttft) = sample.ttft_ms {
if ttft > 0.0 {
min_ttft_ms = min_ttft_ms.min(ttft);
max_ttft_ms = max_ttft_ms.max(ttft);
}
}
if let Some(reset_at) = account.window_reset_at {
let secs = (reset_at.timestamp() - now_secs).max(0) as f64;
min_reset_secs = min_reset_secs.min(secs);
max_reset_secs = max_reset_secs.max(secs);
}
}
Self {
max_slot_count,
max_window_tokens,
max_window_cached,
max_window_requests,
min_ttft_ms: if min_ttft_ms.is_finite() {
min_ttft_ms
} else {
0.0
},
max_ttft_ms,
min_reset_secs: if min_reset_secs.is_finite() {
min_reset_secs
} else {
0.0
},
max_reset_secs,
}
}
}
/// 把 `value` 从 `[min, max]` 线性映射到 `[0, 1]`;区间退化时返回中性值 `1.0`
/// (表示该维度对所有候选无区分度,不应影响排序)。
fn normalize_ascending_good(value: f64, min: f64, max: f64) -> f64 {
if max <= min {
return 1.0;
}
((value - min) / (max - min)).clamp(0.0, 1.0)
}
/// 计算指定账号在候选集合中的加权得分明细。
fn score_account(
account: &Account,
sample: &FeedbackSample,
weights: &ScoreWeights,
ctx: &NormalizationContext,
now_secs: i64,
slot_count: &dyn Fn(&str) -> usize,
) -> ScoreBreakdown {
// 负载:槽位越少越好。归一化到 [0,1] 后取反(1 - x),使“越空闲得分越高”。
let load =
1.0 - normalize_ascending_good(slot_count(&account.id) as f64, 0.0, ctx.max_slot_count);
// 窗口实际 token 用量:token 越少越好。
let window = 1.0
- normalize_ascending_good(
account_window_token_count(account) as f64,
0.0,
ctx.max_window_tokens,
);
// 窗口缓存 token 用量:缓存越少越好(权重小,仅在实际用量相当时细分)。
let cached = 1.0
- normalize_ascending_good(
account.window_cached_tokens as f64,
0.0,
ctx.max_window_cached,
);
// 窗口请求数:请求越少越好。
let window_requests = 1.0
- normalize_ascending_good(
account.window_request_count as f64,
0.0,
ctx.max_window_requests,
);
// 错误率 EWMA:本身即 [0,1],越低越健康。无样本时视为 0(满分健康)。
let error_rate = 1.0 - sample.error_rate.unwrap_or(0.0).clamp(0.0, 1.0);
// TTFT:延迟越低越好。无样本的账号视为中性满分,避免冷启动被压。
let ttft = match sample.ttft_ms {
Some(ms) if ms > 0.0 => {
1.0 - normalize_ascending_good(ms, ctx.min_ttft_ms, ctx.max_ttft_ms)
}
_ => 1.0,
};
// 窗口重置:越快重置越优先(use-it-or-lose-it)。无 reset 时间视为中性满分。
let reset = match account.window_reset_at {
Some(reset_at) => {
let secs = (reset_at.timestamp() - now_secs).max(0) as f64;
1.0 - normalize_ascending_good(secs, ctx.min_reset_secs, ctx.max_reset_secs)
}
None => 1.0,
};
let mut total = weights.load * load
+ weights.window * window
+ weights.cached * cached
+ weights.window_requests * window_requests
+ weights.error_rate * error_rate
+ weights.ttft * ttft
+ weights.reset * reset;
// 配额封顶硬惩罚:确保只要还有未封顶候选就不会选中封顶账号。
if account.quota_limit_reached {
total -= QUOTA_LIMITED_PENALTY;
}
ScoreBreakdown {
total,
load,
window,
cached,
window_requests,
error_rate,
ttft,
reset,
quota_limited: account.quota_limit_reached,
}
}
@@ -0,0 +1,28 @@
//! Sticky 策略:优先选择最近使用过的账号。
//!
//! 选中 `last_used_at` 最大(最近使用)的账号,倾向于持续复用同一账号(利于上游的
//! prompt 缓存命中)。该策略不使用轮转游标——最近使用时间本身即确定性顺序。
use std::cmp::Ordering;
use crate::upstream::accounts::model::Account;
use super::{compare_last_used, SelectionInput};
/// 选中最近使用过的账号(`last_used_at` 最大者)。
///
/// `_cursor` 未使用:Sticky 由 `last_used_at` 决定顺序,不参与轮转破并列。
pub fn select(input: &SelectionInput<'_>, _cursor: &mut usize) -> Option<Account> {
let candidates = input.candidates;
let mut selected = candidates.first()?;
for candidate in &candidates[1..] {
if compare_last_used(
candidate.last_used_at.as_deref(),
selected.last_used_at.as_deref(),
) == Ordering::Greater
{
selected = candidate;
}
}
Some(selected.clone())
}
+110 -61
View File
@@ -22,9 +22,9 @@ use thiserror::Error;
use tokio_tungstenite::tungstenite::handshake::client::generate_key;
use uuid::Uuid;
use crate::upstream::fingerprint::Fingerprint;
use crate::upstream::fingerprint::{Fingerprint, RuntimeFingerprint};
use crate::upstream::models::backend_entry::BackendModelEntry;
use crate::upstream::protocol::events::extract_sse_usage;
use crate::upstream::protocol::events::{extract_sse_usage, retry_after_seconds_from_body};
use crate::upstream::protocol::responses::{
http_sse_fallback_allowed, transport_for_request, CodexCompactRequest, CodexResponsesRequest,
CodexTransport,
@@ -34,6 +34,7 @@ use crate::upstream::protocol::websocket::{
websocket_audit_artifact_from_attempt, websocket_payload_audit_snapshot,
};
use super::diagnostics::CodexUpstreamDiagnostics;
use super::endpoints::{endpoint_url, CODEX_RESPONSES_COMPACT_PATH, CODEX_RESPONSES_PATH};
use super::headers::{
build_ordered_codex_base_headers, insert_optional_header, insert_ordered_headers,
@@ -46,7 +47,9 @@ use super::websocket::{
write_websocket_audit_artifact_from_env, CodexWebSocketConnection, CodexWebSocketExchangeError,
CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketTurnStateUpdate,
};
use super::websocket_pool::{CodexWebSocketPool, CodexWebSocketPoolKey, WebSocketPoolDecision};
use super::websocket_pool::{
CodexWebSocketPool, CodexWebSocketPoolKey, WebSocketPoolDecision, DEFAULT_FIRST_TOKEN_TIMEOUT,
};
// ---------------------------------------------------------------------------
// Constants
@@ -127,6 +130,8 @@ pub enum CodexClientError {
body: String,
/// 推导出的重试秒数。
retry_after_seconds: Option<u64>,
/// 上游诊断元数据。
diagnostics: CodexUpstreamDiagnostics,
/// 上游透传的 `set-cookie` 列表。
set_cookie_headers: Vec<String>,
},
@@ -253,6 +258,8 @@ pub struct CodexBackendResponse {
pub first_token_ms: Option<i64>,
/// WebSocket 连接池决策。
pub websocket_pool_decision: Option<WebSocketPoolDecision>,
/// 上游诊断元数据。
pub diagnostics: CodexUpstreamDiagnostics,
}
/// Codex Responses 实际使用的上游传输。
@@ -300,6 +307,8 @@ pub struct CodexBackendStreamingResponse {
pub turn_state_update: Option<CodexTurnStateUpdate>,
/// WebSocket 连接池决策。
pub websocket_pool_decision: Option<WebSocketPoolDecision>,
/// 上游诊断元数据。
pub diagnostics: CodexUpstreamDiagnostics,
}
/// Codex compact 端点响应。
@@ -311,6 +320,8 @@ pub struct CodexCompactResponse {
pub set_cookie_headers: Vec<String>,
/// 上游透传的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// 上游诊断元数据。
pub diagnostics: CodexUpstreamDiagnostics,
}
// ---------------------------------------------------------------------------
@@ -322,21 +333,32 @@ pub struct CodexCompactResponse {
pub struct CodexBackendClient {
pub(super) client: Client,
pub(super) base_url: String,
fingerprint: Fingerprint,
fingerprint: RuntimeFingerprint,
websocket_pool: Option<Arc<CodexWebSocketPool>>,
websocket_first_token_timeout: Option<Duration>,
}
impl CodexBackendClient {
/// 构造客户端。
pub fn new(client: Client, base_url: impl Into<String>, fingerprint: Fingerprint) -> Self {
pub fn new(
client: Client,
base_url: impl Into<String>,
fingerprint: RuntimeFingerprint,
) -> Self {
Self {
client,
base_url: base_url.into().trim_end_matches('/').to_string(),
fingerprint,
websocket_pool: None,
websocket_first_token_timeout: Some(DEFAULT_FIRST_TOKEN_TIMEOUT),
}
}
pub fn with_websocket_first_token_timeout(mut self, timeout: Option<Duration>) -> Self {
self.websocket_first_token_timeout = timeout.filter(|timeout| !timeout.is_zero());
self
}
/// 为 Responses WebSocket 请求启用连接池。
pub fn with_websocket_pool(mut self, pool: Arc<CodexWebSocketPool>) -> Self {
self.websocket_pool = Some(pool);
@@ -367,6 +389,7 @@ impl CodexBackendClient {
.send()
.await?;
let status = response.status();
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
let turn_state = response_meta::turn_state(response.headers());
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
@@ -379,6 +402,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
diagnostics,
set_cookie_headers,
});
}
@@ -401,6 +425,7 @@ impl CodexBackendClient {
rate_limit_headers,
first_token_ms,
websocket_pool_decision: None,
diagnostics,
})
}
@@ -419,6 +444,7 @@ impl CodexBackendClient {
.send()
.await?;
let status = response.status();
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
let turn_state = response_meta::turn_state(response.headers());
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
@@ -431,6 +457,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
diagnostics,
set_cookie_headers,
});
}
@@ -444,6 +471,7 @@ impl CodexBackendClient {
rate_limit_header_updates: None,
turn_state_update: None,
websocket_pool_decision: None,
diagnostics,
})
}
@@ -590,10 +618,23 @@ impl CodexBackendClient {
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
let exchange = match (self.websocket_pool.as_deref(), pool_key) {
(Some(pool), Some(key)) => {
execute_response_create_request_with_pool(&prepared, Some((pool, key)), started_at)
.await
execute_response_create_request_with_pool(
&prepared,
Some((pool, key)),
started_at,
self.websocket_first_token_timeout,
)
.await
}
_ => {
execute_response_create_request_with_pool(
&prepared,
None,
started_at,
self.websocket_first_token_timeout,
)
.await
}
_ => execute_response_create_request_with_pool(&prepared, None, started_at).await,
}
.map_err(websocket_exchange_error_to_client_error)?;
log_websocket_pool_decision(
@@ -612,6 +653,7 @@ impl CodexBackendClient {
rate_limit_headers: exchange.rate_limit_headers,
first_token_ms: exchange.first_token_ms,
websocket_pool_decision: exchange.pool_decision,
diagnostics: exchange.diagnostics,
})
}
@@ -642,9 +684,21 @@ impl CodexBackendClient {
let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key);
let exchange = match (self.websocket_pool.as_deref(), pool_key) {
(Some(pool), Some(key)) => {
execute_response_create_request_stream_with_pool(&prepared, Some((pool, key))).await
execute_response_create_request_stream_with_pool(
&prepared,
Some((pool, key)),
self.websocket_first_token_timeout,
)
.await
}
_ => {
execute_response_create_request_stream_with_pool(
&prepared,
None,
self.websocket_first_token_timeout,
)
.await
}
_ => execute_response_create_request_stream_with_pool(&prepared, None).await,
}
.map_err(websocket_exchange_error_to_client_error)?;
log_websocket_pool_decision(
@@ -667,6 +721,7 @@ impl CodexBackendClient {
rate_limit_header_updates: Some(exchange.rate_limit_header_updates),
turn_state_update: Some(exchange.turn_state_update),
websocket_pool_decision: exchange.pool_decision,
diagnostics: exchange.diagnostics,
})
}
@@ -678,10 +733,9 @@ impl CodexBackendClient {
) -> Option<CodexWebSocketPoolKey> {
let account_id = pool_account_id.or(context.account_id)?;
let conversation_id = request
.prompt_cache_key
.as_deref()
.prompt_cache_key()
.or(request.client_conversation_id.as_deref())
.or(request.previous_response_id.as_deref())?;
.or(request.previous_response_id())?;
Some(CodexWebSocketPoolKey::new(
&self.base_url,
account_id,
@@ -705,6 +759,7 @@ impl CodexBackendClient {
.await?;
let status = response.status();
let diagnostics = response_meta::diagnostics(Some(status.as_u16()), response.headers());
let set_cookie_headers = response_meta::set_cookie_headers(response.headers());
let rate_limit_headers = response_meta::rate_limit_headers(response.headers());
let retry_after_seconds = retry_after_seconds(response.headers(), None);
@@ -716,6 +771,7 @@ impl CodexBackendClient {
retry_after_seconds: retry_after_seconds
.or_else(|| retry_after_seconds_from_body(&body)),
body,
diagnostics,
set_cookie_headers,
});
}
@@ -728,12 +784,16 @@ impl CodexBackendClient {
"Compact response is not valid JSON: {}",
truncate_for_error(&body)
),
diagnostics: CodexUpstreamDiagnostics::with_status(
StatusCode::BAD_GATEWAY.as_u16(),
),
set_cookie_headers: set_cookie_headers.clone(),
})?;
Ok(CodexCompactResponse {
body: parsed,
set_cookie_headers,
rate_limit_headers,
diagnostics,
})
}
@@ -742,17 +802,18 @@ impl CodexBackendClient {
&self,
context: CodexRequestContext<'_>,
) -> CodexClientResult<Vec<BackendModelEntry>> {
let fingerprint = self.fingerprint.current();
let endpoints = [
format!(
"{}/codex/models?client_version={}",
self.base_url, self.fingerprint.app_version
self.base_url, fingerprint.app_version
),
format!("{}/models", self.base_url),
format!("{}/sentinel/chat-requirements", self.base_url),
];
for endpoint in endpoints {
let headers = self.auxiliary_request_headers(context)?;
let headers = self.auxiliary_request_headers(&fingerprint, context)?;
let response = self.client.get(endpoint).headers(headers).send().await?;
if !response.status().is_success() {
continue;
@@ -768,6 +829,7 @@ impl CodexBackendClient {
status: StatusCode::BAD_GATEWAY,
retry_after_seconds: None,
body: "backend model catalog is unavailable".to_string(),
diagnostics: CodexUpstreamDiagnostics::with_status(StatusCode::BAD_GATEWAY.as_u16()),
set_cookie_headers: Vec::new(),
})
}
@@ -778,7 +840,7 @@ impl CodexBackendClient {
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let mut headers = self.request_headers(context)?;
if let Some(subagent) = openai_subagent_from_metadata(request.client_metadata.as_ref()) {
if let Some(subagent) = openai_subagent_from_metadata(request.client_metadata()) {
headers.insert(
HeaderName::from_static("x-openai-subagent"),
HeaderValue::from_str(&subagent)?,
@@ -788,9 +850,10 @@ impl CodexBackendClient {
}
fn request_headers(&self, context: CodexRequestContext<'_>) -> CodexClientResult<HeaderMap> {
let fingerprint = self.fingerprint.current();
let request_id = context.session_id.unwrap_or(context.request_id);
let ordered_headers = build_ordered_codex_base_headers(
&self.fingerprint,
&fingerprint,
context.access_token,
context.account_id,
);
@@ -840,13 +903,11 @@ impl CodexBackendClient {
fn auxiliary_request_headers(
&self,
fingerprint: &Fingerprint,
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let ordered_headers = build_ordered_codex_base_headers(
&self.fingerprint,
context.access_token,
context.account_id,
);
let ordered_headers =
build_ordered_codex_base_headers(fingerprint, context.access_token, context.account_id);
let mut headers = HeaderMap::new();
insert_ordered_headers(&mut headers, &ordered_headers)?;
@@ -867,10 +928,11 @@ impl CodexBackendClient {
&self,
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let fingerprint = self.fingerprint.current();
let mut headers = HeaderMap::new();
headers.insert(
USER_AGENT,
HeaderValue::from_str(&self.fingerprint.user_agent())?,
HeaderValue::from_str(&fingerprint.user_agent())?,
);
headers.insert(
AUTHORIZATION,
@@ -878,7 +940,7 @@ impl CodexBackendClient {
);
headers.insert(
HeaderName::from_static("originator"),
HeaderValue::from_str(&self.fingerprint.originator)?,
HeaderValue::from_str(&fingerprint.originator)?,
);
insert_optional_header(&mut headers, "chatgpt-account-id", context.account_id)?;
insert_optional_header(&mut headers, "cookie", context.cookie_header)?;
@@ -890,8 +952,9 @@ impl CodexBackendClient {
&self,
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let fingerprint = self.fingerprint.current();
let ordered_headers = build_ordered_codex_base_headers(
&self.fingerprint,
&fingerprint,
context.access_token,
context.account_id,
);
@@ -1012,18 +1075,6 @@ pub(super) fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Op
.or_else(|| body.and_then(retry_after_seconds_from_body))
}
pub(super) fn retry_after_seconds_from_body(body: &str) -> Option<u64> {
let value: Value = serde_json::from_str(body).ok()?;
value
.get("retry_after_seconds")
.or_else(|| value.get("retry_after"))
.and_then(|v| {
v.as_u64()
.or_else(|| v.as_str().and_then(|s| s.parse::<u64>().ok()))
})
.filter(|seconds| *seconds > 0)
}
pub(super) fn truncate_for_error(body: &str) -> String {
body.chars().take(200).collect()
}
@@ -1082,9 +1133,9 @@ fn response_upstream_request(
) -> CodexResponsesRequest {
let mut upstream = request.clone();
if let Some(session_id) = context.session_id {
upstream.prompt_cache_key = Some(session_id.to_string());
upstream.set_prompt_cache_key(Some(session_id.to_string()));
}
upstream.client_metadata = response_client_metadata(request.client_metadata.as_ref(), context);
upstream.set_client_metadata(response_client_metadata(request.client_metadata(), context));
upstream
}
@@ -1095,15 +1146,15 @@ fn websocket_upstream_request(request: &CodexResponsesRequest) -> CodexResponses
}
fn stamp_ws_stream_request_start_ms(request: &mut CodexResponsesRequest) {
let mut metadata = match request.client_metadata.take() {
Some(Value::Object(metadata)) => metadata,
let mut metadata = match request.client_metadata() {
Some(Value::Object(metadata)) => metadata.clone(),
_ => Map::new(),
};
metadata.insert(
X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY.to_string(),
Value::String(now_unix_timestamp_millis().to_string()),
);
request.client_metadata = Some(Value::Object(metadata));
request.set_client_metadata(Some(Value::Object(metadata)));
}
fn now_unix_timestamp_millis() -> u128 {
@@ -1117,14 +1168,12 @@ fn response_client_metadata(
client_metadata: Option<&Value>,
context: CodexRequestContext<'_>,
) -> Option<Value> {
let mut metadata = Map::new();
if let Some(Value::Object(input)) = client_metadata {
for (key, value) in input {
if let Some(value) = value.as_str() {
metadata.insert(key.clone(), Value::String(value.to_string()));
}
}
}
// 以客户端原始 client_metadata 为基础(保留 number/bool/object 等所有值类型),
// 在其上追加代理自身的上下文字段。
let mut metadata = match client_metadata {
Some(Value::Object(input)) => input.clone(),
_ => Map::new(),
};
insert_metadata_string(
&mut metadata,
@@ -1182,17 +1231,17 @@ fn websocket_exchange_error_to_client_error(
error: CodexWebSocketExchangeError,
) -> CodexClientError {
match error {
CodexWebSocketExchangeError::Upstream {
status_code,
retry_after_seconds,
body,
set_cookie_headers,
} => CodexClientError::Upstream {
status: StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_GATEWAY),
body,
retry_after_seconds,
set_cookie_headers,
},
CodexWebSocketExchangeError::Upstream(upstream) => {
let upstream = *upstream;
CodexClientError::Upstream {
status: StatusCode::from_u16(upstream.status_code)
.unwrap_or(StatusCode::BAD_GATEWAY),
body: upstream.body,
retry_after_seconds: upstream.retry_after_seconds,
diagnostics: upstream.diagnostics,
set_cookie_headers: upstream.set_cookie_headers,
}
}
error => CodexClientError::WebSocket(error),
}
}
@@ -0,0 +1,73 @@
//! Upstream response diagnostics captured at the transport boundary.
use reqwest::header::HeaderMap;
const UPSTREAM_REQUEST_ID_HEADERS: &[&str] = &[
"x-request-id",
"x-openai-request-id",
"openai-request-id",
"request-id",
];
const UPSTREAM_TRACE_HEADERS: &[&str] = &[
"x-request-id",
"x-openai-request-id",
"openai-request-id",
"request-id",
"cf-ray",
];
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CodexUpstreamDiagnostics {
pub status_code: Option<u16>,
pub request_id: Option<String>,
pub trace_headers: Vec<(String, String)>,
}
impl CodexUpstreamDiagnostics {
pub fn from_headers(status_code: Option<u16>, headers: &HeaderMap) -> Self {
Self {
status_code,
request_id: first_header(headers, UPSTREAM_REQUEST_ID_HEADERS),
trace_headers: trace_headers(headers),
}
}
pub fn with_status(status_code: u16) -> Self {
Self {
status_code: Some(status_code),
..Self::default()
}
}
pub fn is_empty(&self) -> bool {
self.status_code.is_none() && self.request_id.is_none() && self.trace_headers.is_empty()
}
pub fn cf_ray(&self) -> Option<&str> {
self.trace_headers
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("cf-ray"))
.map(|(_, value)| value.as_str())
}
}
fn first_header(headers: &HeaderMap, names: &[&str]) -> Option<String> {
names.iter().find_map(|name| header_value(headers, name))
}
fn trace_headers(headers: &HeaderMap) -> Vec<(String, String)> {
UPSTREAM_TRACE_HEADERS
.iter()
.filter_map(|name| header_value(headers, name).map(|value| ((*name).to_string(), value)))
.collect()
}
fn header_value(headers: &HeaderMap, name: &str) -> Option<String> {
headers
.get(name)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
+3
View File
@@ -1,6 +1,7 @@
//! Codex HTTP/SSE/WebSocket 上游 transport。
pub mod client;
pub mod diagnostics;
pub mod endpoints;
pub mod headers;
mod response_meta;
@@ -8,6 +9,7 @@ pub mod tls;
pub mod usage;
pub mod websocket;
pub mod websocket_pool;
pub(crate) mod websocket_pump;
pub use self::{
client::{
@@ -18,6 +20,7 @@ pub use self::{
CodexModelCatalogRequest, CodexRateLimitHeaderUpdates, CodexRequestContext,
CodexTurnStateUpdate,
},
diagnostics::CodexUpstreamDiagnostics,
endpoints::{
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,
@@ -6,6 +6,15 @@ use reqwest::header::{HeaderMap, SET_COOKIE};
use crate::upstream::protocol::responses::update_first_response_event_ms;
use super::diagnostics::CodexUpstreamDiagnostics;
pub(super) fn diagnostics(
status_code: Option<u16>,
headers: &HeaderMap,
) -> CodexUpstreamDiagnostics {
CodexUpstreamDiagnostics::from_headers(status_code, headers)
}
pub(super) fn turn_state(headers: &HeaderMap) -> Option<String> {
headers
.get("x-codex-turn-state")

Some files were not shown because too many files have changed in this diff Show More