Compare commits

...
9 Commits
Author SHA1 Message Date
zyycnandCodex 4e527feb38 chore: bump version to 2.1.0
Co-authored-by: Codex <noreply@openai.com>
2026-07-13 01:10:40 +08:00
zyycnandCodex 537c749b73 docs: remove Codex provider setup example
Co-authored-by: Codex <noreply@openai.com>
2026-07-13 01:10:30 +08:00
zyycnandCodex 28b89200f7 feat: refine admin workflows and proxy telemetry
Improve account and API key filtering, sorting, responsive layouts, table loading, usage reporting, and billing semantics across the admin interface.

Remove obsolete compact-route handling, normalize model metadata, and add process-wide draining for HTTP streams and WebSockets during restart.

Co-authored-by: Codex <noreply@openai.com>
2026-07-13 00:57:26 +08:00
zyycnandCodex a8f5c5a54e fix: standardize Docker timezone
Set the runtime image and PostgreSQL service to Asia/Shanghai while preserving UTC timestamptz storage semantics.

Verify the application and PostgreSQL timezone contracts in the container workflow and document the display behavior.

Co-authored-by: Codex <noreply@openai.com>
2026-07-12 21:38:14 +08:00
zyycnandCodex d94bca33dd refactor: remove SQLite migration support
Remove the import-sqlite command, bootstrap importer, migration fixtures and tests, and the SQLx SQLite feature.

Keep the runtime and project documentation focused exclusively on PostgreSQL and Redis storage.

Co-authored-by: Codex <noreply@openai.com>
2026-07-12 21:29:36 +08:00
zyycnandCodex 9b33e60df2 fix: add resilient Responses WebSocket transport
Expose the official Responses WebSocket endpoint so Codex can retry interrupted turns without proxy-side response splicing.

Preserve upstream frames under backpressure, report pump exit reasons, organize Responses handlers by transport, and poll fingerprint updates daily. Keep the warning policy in the backend manifest instead of a repository-wide Cargo config.

Co-authored-by: Codex <noreply@openai.com>
2026-07-12 21:16:04 +08:00
zyycnandCodex 1e99c5037d refactor: migrate backend to Rust 1.97 and edition 2024
Unify the local, CI, release, and container toolchains on Rust 1.97.

Move update configuration, release caching, and event delivery into service-owned state, apply Edition 2024 idioms, and refresh repository ignore rules and documentation.

Co-authored-by: Codex <noreply@openai.com>
2026-07-12 19:10:56 +08:00
zyycnandCodex 26cf170f33 chore(deps): upgrade RustCrypto digest stack
Upgrade hmac and sha2 together so both use the digest 0.11 API.

Co-authored-by: Codex <noreply@openai.com>
2026-07-12 15:47:13 +08:00
dependabot[bot]andCodex 8d54171771 chore(deps): bump Rust builder to 1.97
Co-authored-by: Codex <noreply@openai.com>
2026-07-12 15:38:59 +08:00
243 changed files with 5326 additions and 5684 deletions
+4
View File
@@ -122,6 +122,8 @@ jobs:
compose=(docker compose --project-name "$COMPOSE_PROJECT_NAME" \
-f deploy/docker-compose.yml -f deploy/docker-compose.smoke.yml)
"${compose[@]}" up -d --no-build --wait postgres redis
test "$("${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
psql -At -U codex_proxy -d codex_proxy -c "SHOW TimeZone;" | tr -d '\r')" = Asia/Shanghai
"${compose[@]}" exec -T -e PGPASSWORD="$CPR_POSTGRES_PASSWORD" postgres \
psql -v ON_ERROR_STOP=1 -U codex_proxy -d codex_proxy -c \
"CREATE TABLE cpr_ci_persistence_probe (value text PRIMARY KEY); INSERT INTO cpr_ci_persistence_probe VALUES ('persisted');"
@@ -151,6 +153,8 @@ jobs:
test "$status" = 204
test "$("${compose[@]}" exec -T codex-proxy-rs id -u | tr -d '\r')" = 10001
test "$("${compose[@]}" exec -T codex-proxy-rs id -g | tr -d '\r')" = 10001
test "$("${compose[@]}" exec -T codex-proxy-rs date +%z | tr -d '\r')" = +0800
test "$("${compose[@]}" exec -T codex-proxy-rs cat /etc/timezone | tr -d '\r')" = Asia/Shanghai
"${compose[@]}" exec -T codex-proxy-rs sh -c \
'touch /app/data/.ci-write /app/logs/.ci-write && rm /app/data/.ci-write /app/logs/.ci-write'
+2 -3
View File
@@ -24,7 +24,6 @@ jobs:
CARGO_TERM_COLOR: always
CPR_TEST_DATABASE_URL: postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy
CPR_TEST_REDIS_URL: redis://127.0.0.1:6379
RUSTFLAGS: -D warnings
services:
postgres:
image: postgres:18-bookworm@sha256:d9c83446333daec3f0588cc709adb80c26090b7f9f0f7ec8d43c243385d79818
@@ -56,7 +55,7 @@ jobs:
ref: ${{ inputs.ref }}
- name: Setup Rust
uses: dtolnay/rust-toolchain@eac0f66a48bc4b70a10b9acd4c4e930f835d95ff # 1.95
uses: dtolnay/rust-toolchain@1a3a6d54512beeaffd394f8d516ca16f2c506a20 # 1.97
with:
components: rustfmt,clippy
@@ -71,7 +70,7 @@ jobs:
- name: Rust clippy
working-directory: backend
run: cargo clippy --all-targets --all-features --locked -- -D warnings
run: cargo clippy --all-targets --all-features --locked
- name: Rust tests
working-directory: backend
+2 -2
View File
@@ -205,7 +205,7 @@ jobs:
-v "$HOME/.cargo:/cargo" \
-v "$PWD/..:/workspace" \
-w /workspace/backend \
rust:1.95-bookworm@sha256:6258907abe69656e41cd992e0b705cdcfabcbbe3db374f92ed2d47121282d4a1 \
rust:1.97-bookworm@sha256:7d0723df719e7f213b69dc7c8c595985c3f4b060cfbee4f7bc0e347a86fe3b6a \
cargo build --release --locked --target "${{ matrix.target }}" --bin codex-proxy-rs
cp "target/${{ matrix.target }}/release/${{ matrix.binary }}" "../dist/binaries/${{ matrix.os }}-${{ matrix.arch }}/${{ matrix.binary }}"
working-directory: backend
@@ -239,7 +239,7 @@ jobs:
ref: ${{ needs.release-info.outputs.git_sha }}
- name: Setup Rust
uses: dtolnay/rust-toolchain@eac0f66a48bc4b70a10b9acd4c4e930f835d95ff # 1.95
uses: dtolnay/rust-toolchain@1a3a6d54512beeaffd394f8d516ca16f2c506a20 # 1.97
with:
targets: ${{ matrix.target }}
+1 -1
View File
@@ -26,7 +26,7 @@ jobs:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
- name: Setup Rust
uses: dtolnay/rust-toolchain@eac0f66a48bc4b70a10b9acd4c4e930f835d95ff # 1.95
uses: dtolnay/rust-toolchain@1a3a6d54512beeaffd394f8d516ca16f2c506a20 # 1.97
- name: Cache Rust
uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1
+34 -30
View File
@@ -1,43 +1,47 @@
# Rust build output.
# Rust
/backend/target/
/target/
/dist/
/package/
/checksums.txt
/codex-proxy-rs_*.tar.gz
/codex-proxy-rs_*.zip
*.profraw
*.profdata
# Frontend dependencies and build output.
# Frontend
/frontend/node_modules/
/frontend/dist/
.pnpm-store
/frontend/.vite/
/.pnpm-store/
/frontend/.pnpm-store/
*.tsbuildinfo
# Application runtime state.
*.runtime
*.sqlite
*.sqlite3
*.db
*.db-shm
*.db-wal
# Release artifacts
/dist/
# Runtime data
/.runtime/
*.log
# Local configuration.
/local.yaml
/local.yml
/.env
# Local configuration and secrets
.env
.env.*
!.env.example
!**/.env.example
/.envrc
/config.yaml
/backend/config.yaml
/deploy/.env
# Local workspaces and tooling state.
# Editors and operating systems
/.idea/
/.vscode/
.DS_Store
Thumbs.db
*.swp
*.swo
*~
# Local tools and worktrees
/.codegraph/
/.codex-ws-audit/
/.agents/
/.reasonix/
/.worktrees/
/worktrees/
/.codex-ws-audit/
/.codegraph/
*.tmp
*.swp
*~
*.codegraph
/.agents/
/.cargo/
/.reasonix/
*.tmp
+7 -8
View File
@@ -2,7 +2,7 @@
<h1>Codex Proxy RS</h1>
<p>OpenAI Responses 兼容的 Codex 多账号网关</p>
<p>
<a href="https://www.rust-lang.org/"><img alt="Rust 1.95" src="https://img.shields.io/badge/Rust-1.95-000?style=flat-square&amp;logo=rust" /></a>
<a href="https://www.rust-lang.org/"><img alt="Rust 1.97" src="https://img.shields.io/badge/Rust-1.97-000?style=flat-square&amp;logo=rust" /></a>
<a href="https://github.com/tokio-rs/axum"><img alt="Axum 0.8" src="https://img.shields.io/badge/Axum-0.8-2f7d95?style=flat-square" /></a>
<a href="https://vuejs.org/"><img alt="Vue 3.5" src="https://img.shields.io/badge/Vue-3.5-42b883?style=flat-square&amp;logo=vuedotjs&amp;logoColor=white" /></a>
<a href="https://vite.dev/"><img alt="Vite 8" src="https://img.shields.io/badge/Vite-8-646cff?style=flat-square&amp;logo=vite&amp;logoColor=white" /></a>
@@ -23,17 +23,16 @@ Codex Proxy RS 为多个 Codex 账号提供统一的 Responses API。客户端
## 功能
- 兼容 `/v1/responses`、SSE 流式响应、Review、Compact 和模型查询。
- 支持 HTTP/SSE 与 WebSocket 上游连接。
- 兼容 `/v1/responses`、SSE、官方 Responses WebSocket 和模型查询。
- 客户端与上游均可使用 HTTP/SSE 或 WebSocket。
- 提供 `smart`、额度重置优先、轮询和粘性四种账号调度策略。
- 单个账号发生额度、认证、模型或传输故障时,继续使用账号池中的其他可用账号。
- 使用会话亲和和受限历史重放处理 `previous_response_not_found` 与跨账号续接。
- 为每个账号派生稳定且隔离的 installation、session、thread、turn 和 cache 标识。
- 支持跨账号会话续接,并隔离不同账号的会话状态。
- 提供账号、API Key、模型映射、运行参数、用量、错误记录和系统更新管理界面。
- 使用 PostgreSQL 保存权威业务数据,使用 Redis 保存会话、租约、模型快照和会话亲和数据。
- 内置日志轮转、遥测保留期清理和 PostgreSQL/Redis 健康检查。
代理只补充账号凭据、上游请求头和账号作用域身份,并在必要时执行会话恢复;其余请求与响应尽量保持上游语义。
除账号调度和会话恢复外,请求与响应尽量保持原有语义。
## 快速部署
@@ -135,8 +134,8 @@ curl http://127.0.0.1:8080/v1/responses \
| 路由 | 说明 |
| --- | --- |
| `POST /v1/responses` | Responses 非流式或 SSE 流式请求 |
| `GET /v1/responses` | 官方 Responses WebSocket |
| `POST /v1/responses/review` | Review 请求 |
| `POST /v1/responses/compact` | Compact 请求 |
| `GET /v1/models` | 模型列表 |
| `GET /v1/models/catalog` | 模型目录 |
| `GET /v1/models/{model_id}` | 模型详情 |
@@ -220,7 +219,7 @@ docker compose --env-file deploy/.env -f deploy/docker-compose.yml config
## 架构
模块边界、请求链路、换号条件、previous response 恢复、存储归属和后台任务见 [架构说明](docs/architecture.md)。
模块边界、请求链路、存储归属和后台任务见 [架构说明](docs/architecture.md)。
## License
+13 -14
View File
@@ -363,7 +363,9 @@ dependencies = [
"flate2",
"futures",
"hex",
"hmac 0.12.1",
"hmac",
"hyper",
"hyper-util",
"indexmap",
"ipnet",
"percent-encoding",
@@ -378,7 +380,7 @@ dependencies = [
"semver",
"serde",
"serde_json",
"sha2 0.10.9",
"sha2 0.11.0",
"sqlx",
"subtle",
"tar",
@@ -453,6 +455,12 @@ dependencies = [
"yaml-rust2",
]
[[package]]
name = "const-oid"
version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c"
[[package]]
name = "cookie"
version = "0.18.1"
@@ -640,6 +648,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2"
dependencies = [
"block-buffer 0.12.1",
"const-oid",
"crypto-common 0.2.2",
"ctutils",
]
@@ -1062,16 +1071,7 @@ version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018"
dependencies = [
"hmac 0.13.0",
]
[[package]]
name = "hmac"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
dependencies = [
"digest 0.10.7",
"hmac",
]
[[package]]
@@ -1403,7 +1403,6 @@ version = "0.37.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b1f111c8c41e7c61a49cd34e44c7619462967221a6443b0ec299e0ac30cfb9b1"
dependencies = [
"cc",
"pkg-config",
"vcpkg",
]
@@ -2402,7 +2401,7 @@ dependencies = [
"futures-util",
"hex",
"hkdf",
"hmac 0.13.0",
"hmac",
"itoa",
"log",
"md-5",
+8 -7
View File
@@ -1,9 +1,9 @@
[package]
name = "codex-proxy-rs"
edition = "2021"
edition = "2024"
license = "MIT"
publish = false
rust-version = "1.95"
rust-version = "1.97"
[dependencies]
argon2 = "0.5.3"
@@ -17,7 +17,9 @@ dirs = "6.0.0"
flate2 = "1.1.9"
futures = "0.3.32"
hex = "0.4.3"
hmac = "0.12.1"
hmac = "0.13.0"
hyper = { version = "1.10.1", features = ["http1", "server"] }
hyper-util = { version = "0.1.20", features = ["tokio"] }
indexmap = { version = "2.0", features = ["serde"] }
ipnet = "2.12.0"
percent-encoding = "2.3.2"
@@ -34,8 +36,8 @@ secrecy = "0.10.3"
semver = "1.0.27"
serde = { version = "1.0.228", features = ["derive"] }
serde_json = { version = "1.0.150", features = ["preserve_order"] }
sha2 = "0.10.9"
sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres", "sqlite-bundled", "chrono", "json"] }
sha2 = "0.11.0"
sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres", "chrono", "json"] }
subtle = "2.6.1"
tar = "0.4.44"
thiserror = "2.0.18"
@@ -62,8 +64,7 @@ tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungsteni
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "4fffad30fe373adbdcffab9545e9e9bf4f2fc19f" }
[lints.rust]
future_incompatible = { level = "warn", priority = -1 }
nonstandard_style = { level = "deny", priority = -1 }
warnings = "deny"
unsafe_code = "forbid"
[lints.clippy]
+2 -2
View File
@@ -1,4 +1,4 @@
[toolchain]
channel = "1.95"
channel = "1.97"
components = ["clippy", "rustfmt"]
profile = "minimal"
components = ["rustfmt", "clippy"]
@@ -33,12 +33,12 @@ pub async fn health_check_accounts(
"stagger_ms must be between 500 and 30000",
));
}
if let Some(concurrency) = payload.concurrency {
if !(1..=10).contains(&concurrency) {
return Err(AdminError::bad_request(
"concurrency must be between 1 and 10",
));
}
if let Some(concurrency) = payload.concurrency
&& !(1..=10).contains(&concurrency)
{
return Err(AdminError::bad_request(
"concurrency must be between 1 and 10",
));
}
match state
.services
+23 -17
View File
@@ -1,14 +1,14 @@
//! 账号管理 HTTP 处理器。
use axum::{
Json,
body::Body,
extract::{Query, State},
http::{
header::{CACHE_CONTROL, CONNECTION, CONTENT_TYPE},
StatusCode,
header::{CACHE_CONTROL, CONNECTION, CONTENT_TYPE},
},
response::{IntoResponse, Response},
Json,
};
use std::collections::{HashMap, HashSet};
@@ -17,30 +17,31 @@ use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{
api::AppState,
api::admin::{
response::{
parse_editable_update, AdminEnvelope, AdminError, AdminResponse, BatchDeleteData,
EditableUpdateMessages, PageMeta, ADMIN_OK_CODE, ADMIN_OK_MESSAGE,
ADMIN_OK_CODE, ADMIN_OK_MESSAGE, AdminEnvelope, AdminError, AdminResponse,
BatchDeleteData, EditableUpdateMessages, PageMeta, parse_editable_update,
},
session::AdminAuth,
},
api::AppState,
fleet::{
account::AccountStatus,
account::{AccountListSort, AccountSortField, AccountStatus},
manage::{
quota_view::{
quota_data, AccountQuotaData, AccountQuotaUsageWindow, AccountQuotaWindowLocalUsage,
},
AccountHealthCheck, AccountManageError, AccountRefreshOutcome, AccountRefreshResult,
AccountUpdate, ManagedAccount, OAuthExchangeInput,
quota_view::{
AccountQuotaData, AccountQuotaUsageWindow, AccountQuotaWindowLocalUsage, quota_data,
},
},
refresh::token_refresh_status_eligible,
},
infra::{
format::{
format_cost, format_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
format_billing_amount, format_percent, format_plain_number, format_tokens,
nonnegative_i64_to_u64,
},
json::{clamp_limit, clamp_page, total_pages},
json::{SortDirection, clamp_limit, clamp_page, total_pages},
time::{china_datetime, china_relative_time, china_rfc3339},
},
telemetry::{
@@ -59,6 +60,9 @@ pub(crate) struct AccountsQuery {
page: Option<u32>,
page_size: Option<u32>,
search: Option<String>,
status: Option<String>,
sort_by: Option<String>,
sort_direction: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -282,8 +286,8 @@ struct AdminAccountModelUsageData {
cached_tokens_display: String,
total_tokens: u64,
total_tokens_display: String,
total_cost_usd: f64,
total_cost_usd_display: String,
billing_amount_usd: f64,
billing_amount_usd_display: String,
last_used_at: Option<String>,
last_used_at_display: String,
}
@@ -304,7 +308,7 @@ struct AccountModelUsageRecord {
input_tokens: i64,
output_tokens: i64,
cached_tokens: i64,
total_cost_usd: f64,
billing_amount_usd: f64,
last_used_at: Option<DateTime<Utc>>,
}
@@ -337,7 +341,7 @@ impl AccountListStats {
struct AdminAccountSummaryData {
total: u64,
active: u64,
high_usage: u64,
quota_exhausted: u64,
attention: u64,
}
@@ -505,13 +509,15 @@ pub(crate) async fn accounts(
) -> Result<impl IntoResponse, AdminError> {
let page = clamp_page(params.page.unwrap_or(1));
let page_size = clamp_limit(params.page_size.unwrap_or(50));
let status = account_status_filter(params.status)?;
let sort = account_list_sort(params.sort_by, params.sort_direction)?;
let quota_by_account = quota_snapshots_by_account(&state).await;
let summary = account_summary_data(&state, &quota_by_account).await;
let summary = account_summary_data(&state).await;
match state
.services
.admin_accounts
.list_page(page, page_size, params.search)
.list_page(page, page_size, params.search, status, sort)
.await
{
Ok(page) => {
+52 -22
View File
@@ -202,23 +202,16 @@ pub(super) fn apply_selected_quota_window_usage(
usage
}
pub(super) async fn account_summary_data(
state: &AppState,
quota_by_account: &HashMap<String, AccountQuotaData>,
) -> AdminAccountSummaryData {
pub(super) async fn account_summary_data(state: &AppState) -> AdminAccountSummaryData {
let accounts = list_all_account_metadata(state).await;
let total = accounts.len() as u64;
let active = accounts
.iter()
.filter(|account| account.status == AccountStatus::Active)
.count() as u64;
let high_usage = accounts
let quota_exhausted = accounts
.iter()
.filter(|account| {
quota_by_account
.get(&account.id)
.is_some_and(account_quota_has_high_usage)
})
.filter(|account| account.status == AccountStatus::QuotaExhausted)
.count() as u64;
let attention = accounts
.iter()
@@ -228,7 +221,7 @@ pub(super) async fn account_summary_data(
AdminAccountSummaryData {
total,
active,
high_usage,
quota_exhausted,
attention,
}
}
@@ -259,7 +252,7 @@ pub(super) async fn list_all_account_metadata(state: &AppState) -> Vec<ManagedAc
let Ok(result) = state
.services
.admin_accounts
.list_page(page, ACCOUNT_STATS_PAGE_LIMIT, None)
.list_page(page, ACCOUNT_STATS_PAGE_LIMIT, None, None, None)
.await
else {
return Vec::new();
@@ -273,10 +266,6 @@ pub(super) async fn list_all_account_metadata(state: &AppState) -> Vec<ManagedAc
}
}
pub(super) fn account_quota_has_high_usage(quota: &AccountQuotaData) -> bool {
quota.has_high_usage()
}
pub(super) fn account_summary_needs_attention(status: AccountStatus) -> bool {
matches!(
status,
@@ -284,6 +273,47 @@ pub(super) fn account_summary_needs_attention(status: AccountStatus) -> bool {
)
}
pub(super) fn account_status_filter(
status: Option<String>,
) -> Result<Option<AccountStatus>, AdminError> {
let Some(status) = status.map(|value| value.trim().to_string()) else {
return Ok(None);
};
if status.is_empty() {
return Ok(None);
}
AccountStatus::parse(&status)
.map(Some)
.ok_or_else(|| AdminError::bad_request("Invalid account status"))
}
pub(super) fn account_list_sort(
sort_by: Option<String>,
sort_direction: Option<String>,
) -> Result<Option<AccountListSort>, AdminError> {
let (sort_by, sort_direction) = match (sort_by, sort_direction) {
(None, None) => return Ok(None),
(Some(sort_by), Some(sort_direction)) => (sort_by, sort_direction),
_ => {
return Err(AdminError::bad_request(
"Account sort field and direction must be provided together",
));
}
};
let field = match sort_by.trim() {
"email" => AccountSortField::Email,
"status" => AccountSortField::Status,
"planType" => AccountSortField::PlanType,
"usage" => AccountSortField::Usage,
"lastUsedAt" => AccountSortField::LastUsedAt,
"expiresAt" => AccountSortField::ExpiresAt,
_ => return Err(AdminError::bad_request("Invalid account sort field")),
};
let direction = SortDirection::parse(&sort_direction)
.ok_or_else(|| AdminError::bad_request("Invalid account sort direction"))?;
Ok(Some(AccountListSort { field, direction }))
}
pub(super) async fn list_current_window_model_usage(
state: &AppState,
usage_records: &[AccountUsageRecord],
@@ -323,7 +353,7 @@ pub(super) async fn list_current_window_model_usage(
let input_tokens = row.input_tokens;
let output_tokens = row.output_tokens;
let cached_tokens = row.cached_tokens;
let total_cost_usd = billing::calculate_cost(
let billing_amount_usd = billing::calculate_billing_amount(
nonnegative_i64_to_u64(input_tokens),
nonnegative_i64_to_u64(output_tokens),
nonnegative_i64_to_u64(cached_tokens),
@@ -342,7 +372,7 @@ pub(super) async fn list_current_window_model_usage(
input_tokens: 0,
output_tokens: 0,
cached_tokens: 0,
total_cost_usd: 0.0,
billing_amount_usd: 0.0,
last_used_at: None,
});
record.request_count += request_count;
@@ -350,7 +380,7 @@ pub(super) async fn list_current_window_model_usage(
record.input_tokens += input_tokens;
record.output_tokens += output_tokens;
record.cached_tokens += cached_tokens;
record.total_cost_usd += total_cost_usd;
record.billing_amount_usd += billing_amount_usd;
record.last_used_at = record.last_used_at.max(last_used_at);
}
let records = records_by_model.into_values().collect::<Vec<_>>();
@@ -408,7 +438,7 @@ pub(super) fn model_usage_data(usage: AccountModelUsageRecord) -> AdminAccountMo
} else {
0.0
};
let total_cost_usd = usage.total_cost_usd;
let billing_amount_usd = usage.billing_amount_usd;
AdminAccountModelUsageData {
model: usage.model,
@@ -424,8 +454,8 @@ pub(super) fn model_usage_data(usage: AccountModelUsageRecord) -> AdminAccountMo
cached_tokens_display: format_tokens(cached_tokens),
total_tokens,
total_tokens_display: format_tokens(total_tokens),
total_cost_usd,
total_cost_usd_display: format_cost(total_cost_usd),
billing_amount_usd,
billing_amount_usd_display: format_billing_amount(billing_amount_usd),
last_used_at: usage.last_used_at.map(|value| china_rfc3339(&value)),
last_used_at_display: china_relative_time(usage.last_used_at, Utc::now()),
}
+9 -10
View File
@@ -1,21 +1,21 @@
//! 管理端认证路由。
use axum::{
extract::State,
http::{header::SET_COOKIE, HeaderMap, HeaderValue, StatusCode},
response::{IntoResponse, Response},
Extension, Json,
extract::State,
http::{HeaderMap, HeaderValue, StatusCode, header::SET_COOKIE},
response::{IntoResponse, Response},
};
use serde::{Deserialize, Serialize};
use crate::{
api::{
AppState,
admin::{
response::{AdminEnvelope, AdminError, AdminResponse},
session::{admin_session_cookie, ADMIN_SESSION_COOKIE},
session::{ADMIN_SESSION_COOKIE, admin_session_cookie},
},
middleware::request_id::ClientIp,
AppState,
},
auth::types::SessionError,
infra::time::china_rfc3339,
@@ -107,15 +107,14 @@ pub(crate) async fn logout(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Response, AdminError> {
if let Some(session_id) = admin_session_cookie(&headers) {
if let Err(error) = state
if let Some(session_id) = admin_session_cookie(&headers)
&& let Err(error) = state
.services
.admin_sessions
.delete_session(&session_id)
.await
{
tracing::warn!(error = %error, "failed to revoke admin session during logout");
}
{
tracing::warn!(error = %error, "failed to revoke admin session during logout");
}
let mut response = AdminResponse::new(
+22 -9
View File
@@ -12,15 +12,15 @@ use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{
api::AppState,
api::admin::{
response::{AdminEnvelope, AdminError, AdminResponse},
session::AdminAuth,
usage_routes::{
usage_cost_details, usage_record_models, usage_token_details,
UsageRecordCostDetailsData, UsageRecordTokenDetailsData,
UsageRecordBillingData, UsageRecordTokenDetailsData, usage_billing,
usage_record_models, usage_token_details,
},
},
api::AppState,
fleet::{
account::{Account, AccountStatus},
pool::AccountCapacitySummary,
@@ -36,13 +36,13 @@ use crate::{
telemetry::{
account_usage::query::AccountUsageTimeBucket,
dashboard::{
dashboard_cards, dashboard_health_timeline_data, dashboard_trend_data,
DashboardAccountCounts, DashboardCardsData, DashboardHealthTimelineData,
DashboardTrendData, DashboardTrendKind,
DashboardTrendData, DashboardTrendKind, dashboard_cards,
dashboard_health_timeline_data, dashboard_trend_data,
},
usage::{
query::{UsageQueryFilter, UsageRecordAccountUsage},
types::{metadata_string, UsageRecord},
types::{UsageRecord, metadata_string},
},
},
upstream::openai::fingerprint::Fingerprint,
@@ -102,8 +102,11 @@ struct DashboardUsageRecordData {
client_ip: Option<String>,
user_agent: Option<String>,
reasoning_effort: Option<String>,
compact: bool,
request_kind: Option<String>,
subagent_kind: Option<String>,
token_details: UsageRecordTokenDetailsData,
cost_details: Option<UsageRecordCostDetailsData>,
billing: Option<UsageRecordBillingData>,
first_token_latency_ms_display: String,
latency_ms_display: String,
}
@@ -514,8 +517,15 @@ fn dashboard_usage_record_data(
let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]);
let reasoning_effort =
metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]);
let compact = record
.metadata
.get("compact")
.and_then(Value::as_bool)
.unwrap_or(false);
let request_kind = metadata_string(&record.metadata, &["requestKind", "request_kind"]);
let subagent_kind = metadata_string(&record.metadata, &["subagentKind", "subagent_kind"]);
let token_details = usage_token_details(&record);
let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details);
let billing = usage_billing(&record, upstream_model.as_deref(), &token_details);
let first_token_latency_ms = record.first_token_ms;
DashboardUsageRecordData {
@@ -531,8 +541,11 @@ fn dashboard_usage_record_data(
client_ip,
user_agent,
reasoning_effort,
compact,
request_kind,
subagent_kind,
token_details,
cost_details,
billing,
first_token_latency_ms_display: format_duration_ms(first_token_latency_ms),
latency_ms_display: format_duration_ms(record.latency_ms),
}
+37 -7
View File
@@ -1,26 +1,28 @@
//! 管理端 v1 接口访问 Key 处理器(列表、创建、生命周期)。
use axum::{
Json,
extract::{Query, State},
http::StatusCode,
response::IntoResponse,
Json,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{
api::AppState,
api::admin::response::{
parse_editable_update, AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse,
BatchDeleteData, EditableUpdateMessages,
AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, BatchDeleteData,
EditableUpdateMessages, parse_editable_update,
},
api::admin::session::AdminAuth,
api::AppState,
infra::{
json::{clamp_limit, clamp_page, NumberedPage},
json::{NumberedPage, SortDirection, clamp_limit, clamp_page},
time::{china_relative_time_str, china_rfc3339_str},
},
keys::types::{KeyManageError, ManagedClientApiKey},
keys::types::{
ClientApiKeyListSort, ClientApiKeySortField, KeyManageError, ManagedClientApiKey,
},
};
#[derive(Debug, Deserialize)]
@@ -29,6 +31,8 @@ pub(crate) struct ApiKeysQuery {
page: Option<u32>,
page_size: Option<u32>,
search: Option<String>,
sort_by: Option<String>,
sort_direction: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -120,10 +124,11 @@ pub(crate) async fn api_keys(
) -> Result<impl IntoResponse, AdminError> {
let page = clamp_page(query.page.unwrap_or(1));
let page_size = clamp_limit(query.page_size.unwrap_or(50));
let sort = client_api_key_list_sort(query.sort_by, query.sort_direction)?;
match state
.services
.admin_client_keys
.list_page(page, page_size, query.search)
.list_page(page, page_size, query.search, sort)
.await
{
Ok(page) => {
@@ -142,6 +147,31 @@ pub(crate) async fn api_keys(
}
}
fn client_api_key_list_sort(
sort_by: Option<String>,
sort_direction: Option<String>,
) -> Result<Option<ClientApiKeyListSort>, AdminError> {
let (sort_by, sort_direction) = match (sort_by, sort_direction) {
(None, None) => return Ok(None),
(Some(sort_by), Some(sort_direction)) => (sort_by, sort_direction),
_ => {
return Err(AdminError::bad_request(
"Client API key sort field and direction must be provided together",
));
}
};
let field = match sort_by.trim() {
"name" => ClientApiKeySortField::Name,
"enabled" => ClientApiKeySortField::Enabled,
"createdAt" => ClientApiKeySortField::CreatedAt,
"lastUsedAt" => ClientApiKeySortField::LastUsedAt,
_ => return Err(AdminError::bad_request("Invalid client API key sort field")),
};
let direction = SortDirection::parse(&sort_direction)
.ok_or_else(|| AdminError::bad_request("Invalid client API key sort direction"))?;
Ok(Some(ClientApiKeyListSort { field, direction }))
}
/// `POST /api/admin/keys`
pub(crate) async fn create_api_key(
State(state): State<AppState>,
+2 -2
View File
@@ -8,13 +8,13 @@ use axum::{
use serde::{Deserialize, Serialize};
use crate::{
api::AppState,
api::admin::{
response::{AdminError, AdminPageEnvelope, AdminResponse},
session::AdminAuth,
},
api::AppState,
infra::{
json::{clamp_limit, clamp_page, NumberedPage},
json::{NumberedPage, clamp_limit, clamp_page},
time::china_datetime,
},
telemetry::{ops::query::OpsQueryError, ops::store::OpsErrorFilter, ops::types::OpsErrorLog},
+2 -11
View File
@@ -1,10 +1,10 @@
//! 管理端响应封装与 From 转换。
use crate::infra::json::{total_pages, NumberedPage};
use crate::infra::json::{NumberedPage, total_pages};
use axum::{
Json,
http::StatusCode,
response::{IntoResponse, Response},
Json,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
@@ -16,7 +16,6 @@ const ADMIN_MALFORMED_JSON_CODE: u32 = 40000;
const ADMIN_BAD_REQUEST_CODE: u32 = 40001;
const ADMIN_INVALID_TIME_RANGE_CODE: u32 = 40002;
const ADMIN_INVALID_MODEL_SOURCE_CODE: u32 = 40003;
const ADMIN_INVALID_ENDPOINT_SOURCE_CODE: u32 = 40004;
const ADMIN_SESSION_REQUIRED_CODE: u32 = 40101;
const ADMIN_INVALID_CREDENTIALS_CODE: u32 = 40102;
const ADMIN_INVALID_API_KEY_CODE: u32 = 40103;
@@ -68,14 +67,6 @@ impl AdminError {
)
}
pub fn invalid_endpoint_source(message: impl Into<String>) -> Self {
Self::new(
StatusCode::BAD_REQUEST,
ADMIN_INVALID_ENDPOINT_SOURCE_CODE,
message,
)
}
pub fn admin_session_required() -> Self {
Self::new(
StatusCode::UNAUTHORIZED,
+2 -2
View File
@@ -1,11 +1,11 @@
//! 管理端路由。
use axum::{
http::{header, HeaderValue},
Router,
http::{HeaderValue, header},
middleware,
response::Response,
routing::{get, post},
Router,
};
use crate::api::AppState;
+2 -2
View File
@@ -2,10 +2,10 @@
use axum::{
extract::FromRequestParts,
http::{request::Parts, HeaderMap},
http::{HeaderMap, request::Parts},
};
use crate::api::{admin::response::AdminError, AppState};
use crate::api::{AppState, admin::response::AdminError};
pub(super) const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session";
+1 -1
View File
@@ -6,11 +6,11 @@ use serde_json::Value;
use std::collections::BTreeMap;
use crate::{
api::AppState,
api::admin::{
response::{AdminEnvelope, AdminError, AdminResponse},
session::AdminAuth,
},
api::AppState,
settings::{SettingsError, SettingsPatch, SettingsSnapshot, SettingsValidationError},
};
+58 -33
View File
@@ -6,24 +6,20 @@ use axum::{
extract::{Json, Query, State},
http::StatusCode,
response::{
sse::{Event, KeepAlive, Sse},
IntoResponse,
sse::{Event, KeepAlive, Sse},
},
};
use futures::Stream;
use serde::Deserialize;
use crate::{
api::AppState,
api::admin::{
response::{AdminEnvelope, AdminError, AdminResponse},
session::AdminAuth,
},
api::AppState,
update::{
service::{self, RestartAction},
state::operation_id,
types::UpdateError,
},
update::{service::RestartAction, state::operation_id, types::UpdateError},
};
#[derive(Debug, Deserialize)]
@@ -38,20 +34,30 @@ pub(crate) struct UpdateRequest {
target_version: String,
}
pub(crate) async fn version(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
pub(crate) async fn version(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(service::version_data().await),
AdminEnvelope::ok(state.services.system_update.version_data().await),
))
}
pub(crate) async fn update_detail(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<UpdateDetailQuery>,
) -> Result<impl IntoResponse, AdminError> {
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(service::update_detail(query.refresh.unwrap_or(false)).await),
AdminEnvelope::ok(
state
.services
.system_update
.update_detail(query.refresh.unwrap_or(false))
.await,
),
))
}
@@ -59,29 +65,25 @@ pub(crate) async fn update_event_stream(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, AdminError> {
let receiver = service::subscribe_update_events();
let shutdown = state.services.process_control.subscribe_shutdown();
let receiver = state.services.system_update.subscribe_update_events();
let stream = futures::stream::unfold(
(receiver, shutdown, false),
|(mut receiver, mut shutdown, close_after_send)| async move {
(receiver, false),
|(mut receiver, 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.is_terminal();
let id = message.id().to_string();
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, close_after_send)));
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
match receiver.recv().await {
Ok(message) => {
let close_after_send = message.is_terminal();
let id = message.id().to_string();
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, close_after_send)));
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
Err(tokio::sync::broadcast::error::RecvError::Closed) => return None,
}
}
},
@@ -90,11 +92,15 @@ pub(crate) async fn update_event_stream(
}
pub(crate) async fn perform_update(
State(state): State<AppState>,
_auth: AdminAuth,
payload: Option<Json<UpdateRequest>>,
) -> Result<impl IntoResponse, AdminError> {
let target = payload.map(|Json(payload)| payload.target_version);
let result = service::perform_update(target)
let result = state
.services
.system_update
.perform_update(target)
.await
.map_err(update_error)?;
Ok(AdminResponse::new(
@@ -103,16 +109,31 @@ pub(crate) async fn perform_update(
))
}
pub(crate) async fn update_status(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
let status = service::update_status().map_err(update_error)?;
pub(crate) async fn update_status(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
let status = state
.services
.system_update
.update_status()
.map_err(update_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(status),
))
}
pub(crate) async fn rollback(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
let operation_id = service::rollback().await.map_err(update_error)?;
pub(crate) async fn rollback(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
let operation_id = state
.services
.system_update
.rollback()
.await
.map_err(update_error)?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(serde_json::json!({
@@ -127,7 +148,11 @@ pub(crate) async fn restart(
State(state): State<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
let plan = service::restart_plan().map_err(update_error)?;
let plan = state
.services
.system_update
.restart_plan()
.map_err(update_error)?;
let message = plan.message;
let process_control = state.services.process_control.clone();
tokio::spawn(async move {
+73 -73
View File
@@ -8,28 +8,29 @@ use axum::{
response::IntoResponse,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{
api::AppState,
api::admin::{
response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse},
session::AdminAuth,
},
api::AppState,
infra::{
format::{
format_compact_number, format_cost, format_duration_ms, format_duration_ms_f64,
format_multiplier, format_number, format_percent, format_token_price, format_tokens,
format_billing_amount, format_compact_number, format_duration_ms,
format_duration_ms_f64, format_multiplier, format_number, format_percent,
format_token_price, format_tokens,
},
json::{clamp_limit, clamp_page, NumberedPage},
json::{NumberedPage, clamp_limit, clamp_page},
time::china_datetime,
},
telemetry::{
usage::query::{
usage_record_cost_details, UsageQueryError, UsageQueryFilter, UsageRecordBreakdown,
UsageRecordEndpointSource, UsageRecordModelSource, UsageRecordSummary,
UsageRecordTrendPoint,
UsageQueryError, UsageQueryFilter, UsageRecordBreakdown, UsageRecordModelSource,
UsageRecordSummary, UsageRecordTrendPoint, usage_record_billing,
},
usage::types::{metadata_string, UsageRecord},
usage::types::{UsageRecord, metadata_string},
},
};
@@ -87,8 +88,11 @@ pub(crate) struct UsageRecordData {
client_ip: Option<String>,
user_agent: Option<String>,
reasoning_effort: Option<String>,
compact: bool,
request_kind: Option<String>,
subagent_kind: Option<String>,
token_details: UsageRecordTokenDetailsData,
cost_details: Option<UsageRecordCostDetailsData>,
billing: Option<UsageRecordBillingData>,
first_token_latency_ms: Option<i64>,
first_token_latency_ms_display: String,
latency_ms_display: String,
@@ -111,23 +115,23 @@ pub(crate) struct UsageRecordTokenDetailsData {
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct UsageRecordCostDetailsData {
input_cost: f64,
output_cost: f64,
cache_read_cost: f64,
total_cost: f64,
original_cost: f64,
pub(crate) struct UsageRecordBillingData {
input_amount: f64,
output_amount: f64,
cache_read_amount: f64,
standard_amount: f64,
total_amount: f64,
input_price_per_mtoken: f64,
output_price_per_mtoken: f64,
cache_read_price_per_mtoken: f64,
service_tier: Option<String>,
service_tier_display: String,
multiplier: f64,
input_cost_display: String,
output_cost_display: String,
cache_read_cost_display: String,
total_cost_display: String,
original_cost_display: String,
input_amount_display: String,
output_amount_display: String,
cache_read_amount_display: String,
standard_amount_display: String,
total_amount_display: String,
input_price_display: String,
output_price_display: String,
cache_read_price_display: String,
@@ -158,9 +162,9 @@ struct UsageRecordBreakdownData {
total_tokens_value: u64,
total_tokens_total: String,
total_tokens_total_value: u64,
cost: String,
actual_cost: String,
account_cost: String,
standard_billing_amount: String,
actual_billing_amount: String,
account_billing_amount: String,
average_latency_ms: String,
}
@@ -178,8 +182,8 @@ struct UsageRecordTrendPointData {
total_tokens_value: u64,
cache_hit_rate: String,
cache_hit_rate_value: f64,
cost: String,
actual_cost: String,
standard_billing_amount: String,
actual_billing_amount: String,
average_latency_ms: String,
average_latency_ms_value: Option<f64>,
}
@@ -264,14 +268,13 @@ pub(crate) async fn usage_records_model_distribution(
pub(crate) async fn usage_records_endpoint_distribution(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<UsageRecordDistributionQuery>,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
let source = endpoint_source_from_query(query.source)?;
let filter = filter_from_query(query.records)?;
let filter = filter_from_query(query)?;
match state
.services
.usage_records
.endpoint_distribution(filter, source)
.endpoint_distribution(filter)
.await
{
Ok(items) => Ok(AdminResponse::new(
@@ -390,19 +393,6 @@ fn model_source_from_query(value: Option<String>) -> Result<UsageRecordModelSour
}
}
fn endpoint_source_from_query(
value: Option<String>,
) -> Result<UsageRecordEndpointSource, AdminError> {
match non_empty(value).as_deref().unwrap_or("inbound") {
"inbound" => Ok(UsageRecordEndpointSource::Inbound),
"upstream" => Ok(UsageRecordEndpointSource::Upstream),
"path" => Ok(UsageRecordEndpointSource::Path),
_ => Err(AdminError::invalid_endpoint_source(
"Invalid endpoint distribution source",
)),
}
}
fn optional_datetime(
value: Option<String>,
) -> Result<Option<chrono::DateTime<chrono::Utc>>, AdminError> {
@@ -452,8 +442,15 @@ fn usage_record_data(
let user_agent = metadata_string(&record.metadata, &["userAgent", "user_agent"]);
let reasoning_effort =
metadata_string(&record.metadata, &["reasoningEffort", "reasoning_effort"]);
let compact = record
.metadata
.get("compact")
.and_then(Value::as_bool)
.unwrap_or(false);
let request_kind = metadata_string(&record.metadata, &["requestKind", "request_kind"]);
let subagent_kind = metadata_string(&record.metadata, &["subagentKind", "subagent_kind"]);
let token_details = usage_token_details(&record);
let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details);
let billing = usage_billing(&record, upstream_model.as_deref(), &token_details);
let first_token_latency_ms = record.first_token_ms;
let first_token_latency_ms_display = format_duration_ms(first_token_latency_ms);
let latency_ms_display = format_duration_ms(record.latency_ms);
@@ -466,8 +463,11 @@ fn usage_record_data(
client_ip,
user_agent,
reasoning_effort,
compact,
request_kind,
subagent_kind,
token_details,
cost_details,
billing,
first_token_latency_ms,
first_token_latency_ms_display,
latency_ms_display,
@@ -523,12 +523,12 @@ pub(crate) fn usage_token_details(record: &UsageRecord) -> UsageRecordTokenDetai
}
}
pub(crate) fn usage_cost_details(
pub(crate) fn usage_billing(
record: &UsageRecord,
upstream_model: Option<&str>,
tokens: &UsageRecordTokenDetailsData,
) -> Option<UsageRecordCostDetailsData> {
let cost = usage_record_cost_details(
) -> Option<UsageRecordBillingData> {
let billing = usage_record_billing(
record,
upstream_model,
tokens.input_tokens,
@@ -536,27 +536,27 @@ pub(crate) fn usage_cost_details(
tokens.cached_tokens,
)?;
Some(UsageRecordCostDetailsData {
input_cost: cost.input_cost,
output_cost: cost.output_cost,
cache_read_cost: cost.cache_read_cost,
total_cost: cost.total_cost,
original_cost: cost.original_cost,
input_price_per_mtoken: cost.input_price_per_mtoken,
output_price_per_mtoken: cost.output_price_per_mtoken,
cache_read_price_per_mtoken: cost.cache_read_price_per_mtoken,
service_tier: cost.service_tier,
service_tier_display: cost.service_tier_display,
multiplier: cost.multiplier,
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),
multiplier_display: format_multiplier(cost.multiplier),
Some(UsageRecordBillingData {
input_amount: billing.input_amount,
output_amount: billing.output_amount,
cache_read_amount: billing.cache_read_amount,
standard_amount: billing.standard_amount,
total_amount: billing.total_amount,
input_price_per_mtoken: billing.input_price_per_mtoken,
output_price_per_mtoken: billing.output_price_per_mtoken,
cache_read_price_per_mtoken: billing.cache_read_price_per_mtoken,
service_tier: billing.service_tier,
service_tier_display: billing.service_tier_display,
multiplier: billing.multiplier,
input_amount_display: format_billing_amount(billing.input_amount),
output_amount_display: format_billing_amount(billing.output_amount),
cache_read_amount_display: format_billing_amount(billing.cache_read_amount),
standard_amount_display: format_billing_amount(billing.standard_amount),
total_amount_display: format_billing_amount(billing.total_amount),
input_price_display: format_token_price(billing.input_price_per_mtoken),
output_price_display: format_token_price(billing.output_price_per_mtoken),
cache_read_price_display: format_token_price(billing.cache_read_price_per_mtoken),
multiplier_display: format_multiplier(billing.multiplier),
})
}
@@ -592,9 +592,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_cost(item.cost),
actual_cost: format_cost(item.actual_cost),
account_cost: format_cost(item.account_cost),
standard_billing_amount: format_billing_amount(item.standard_billing_amount),
actual_billing_amount: format_billing_amount(item.actual_billing_amount),
account_billing_amount: format_billing_amount(item.account_billing_amount),
average_latency_ms: format_duration_ms_f64(item.average_latency_ms),
}
}
@@ -631,8 +631,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_cost(point.cost),
actual_cost: format_cost(point.actual_cost),
standard_billing_amount: format_billing_amount(point.standard_billing_amount),
actual_billing_amount: format_billing_amount(point.actual_billing_amount),
average_latency_ms: format_duration_ms_f64(point.average_latency_ms),
average_latency_ms_value: point.average_latency_ms,
}
+1 -1
View File
@@ -4,10 +4,10 @@ use std::{
};
use axum::{
Json, Router,
body::Body,
http::{Request, StatusCode},
response::{IntoResponse, Response},
Json, Router,
};
use serde_json::json;
use tower::service_fn;
+1 -1
View File
@@ -2,7 +2,7 @@
//!
//! 提供 Bearer token 提取与客户端 API key 校验。
use axum::http::{header::AUTHORIZATION, HeaderMap};
use axum::http::{HeaderMap, header::AUTHORIZATION};
use crate::api::AppState;
+78 -7
View File
@@ -1,11 +1,11 @@
//! OpenAI 路由共享错误响应。
use axum::{
Json,
http::StatusCode,
response::{IntoResponse, Response},
Json,
};
use serde_json::{json, Value};
use serde_json::{Value, json};
use crate::{
dispatch::errors::{DispatchFailureClass, ResponseDispatchError},
@@ -117,12 +117,10 @@ pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Respon
pub fn responses_dispatch_error_response_ref(error: &ResponseDispatchError) -> Response {
if let ResponseDispatchError::Upstream(CodexClientError::Upstream { status, body, .. }) = error
&& status.is_client_error()
&& let Ok(body) = serde_json::from_str::<Value>(body)
{
if status.is_client_error() {
if let Ok(body) = serde_json::from_str::<Value>(body) {
return (*status, Json(body)).into_response();
}
}
return (*status, Json(body)).into_response();
}
match error {
@@ -140,6 +138,79 @@ pub fn responses_stream_dispatch_failed_sse_event(error: &ResponseDispatchError)
responses_failed_sse_event(kind, &http_error.message)
}
pub(super) fn responses_websocket_dispatch_error_event(
error: &ResponseDispatchError,
request_id: &str,
) -> String {
if let Some(event) = upstream_client_websocket_error_event(error, request_id) {
return event;
}
let http_error =
response_dispatch_http_error(error, ResponseDispatchMessageStyle::ResponsesStream);
let kind = response_dispatch_error_kind(error);
responses_websocket_error_event(
http_error.status,
kind.error_type,
kind.code,
&http_error.message,
request_id,
)
}
pub(super) fn responses_websocket_error_event(
status: StatusCode,
error_type: &str,
code: &str,
message: &str,
request_id: &str,
) -> String {
json!({
"type": "error",
"status": status.as_u16(),
"error": {
"type": error_type,
"code": code,
"message": message,
},
"headers": {
"x-request-id": request_id,
},
})
.to_string()
}
fn upstream_client_websocket_error_event(
error: &ResponseDispatchError,
request_id: &str,
) -> Option<String> {
let ResponseDispatchError::Upstream(CodexClientError::Upstream { status, body, .. }) = error
else {
return None;
};
if !status.is_client_error() {
return None;
}
let Value::Object(mut payload) = serde_json::from_str::<Value>(body).ok()? else {
return None;
};
if !payload.get("error").is_some_and(Value::is_object) {
return None;
}
payload.insert("type".to_string(), Value::String("error".to_string()));
payload.insert("status".to_string(), Value::Number(status.as_u16().into()));
let headers = payload
.entry("headers".to_string())
.or_insert_with(|| Value::Object(serde_json::Map::new()));
if !headers.is_object() {
*headers = Value::Object(serde_json::Map::new());
}
headers.as_object_mut()?.insert(
"x-request-id".to_string(),
Value::String(request_id.to_string()),
);
Some(Value::Object(payload).to_string())
}
fn responses_no_available_accounts_response() -> (StatusCode, Json<Value>) {
let kind = NO_AVAILABLE_ACCOUNTS_ERROR;
(
-1
View File
@@ -5,4 +5,3 @@ pub mod errors;
pub mod models;
pub mod responses;
pub mod router;
pub mod sse;
+3 -3
View File
@@ -1,15 +1,15 @@
//! OpenAI 模型列表以及模型处理器。
use axum::{
Json,
extract::{Path, State},
http::{HeaderMap, StatusCode},
response::IntoResponse,
Json,
};
use serde_json::{json, Value};
use serde_json::{Value, json};
use crate::{
api::client::auth::authorize_client_api_key, api::AppState, models::catalog::ModelCatalog,
api::AppState, api::client::auth::authorize_client_api_key, models::catalog::ModelCatalog,
};
use super::errors::{missing_client_api_key_response, model_not_found_response};
@@ -1,25 +1,28 @@
//! OpenAI Responses API 透明代理入口:解析原始 JSON body,提取调度元数据,
//! 做最小 patch 后交给调度层,请求语义原样透传上游。
mod sse;
mod websocket;
pub use websocket::responses_websocket;
use axum::{
Extension, Json,
body::{Body, Bytes},
extract::State,
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
Extension, Json,
};
use serde_json::{Map, Value};
use crate::{
api::middleware::request_id::{ClientIp, RequestId},
api::AppState,
api::middleware::request_id::{ClientIp, RequestId},
dispatch::{errors::ResponseDispatchError, service::ResponseDispatchStream},
upstream::openai::protocol::{
responses::{CodexCompactRequest, CodexResponsesRequest},
sse::sse_body_has_done,
},
upstream::openai::protocol::{responses::CodexResponsesRequest, sse::sse_body_has_done},
};
use self::sse::{done_sse_frame, event_stream_response as sse_event_stream_response};
use super::{
errors::{
invalid_responses_request_response, missing_client_api_key_response,
@@ -27,7 +30,6 @@ use super::{
responses_dispatch_error_response_ref, responses_stream_dispatch_failed_sse_event,
},
models::model_catalog_for_state,
sse::{done_sse_frame, event_stream_response as sse_event_stream_response},
};
const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent";
@@ -138,49 +140,6 @@ fn identity_context_string(body: &Map<String, Value>, key: &str) -> Option<Strin
})
}
/// 从客户端原始 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 {
let client_session_id = identity_context_string(&body, "session_id")
.or_else(|| header_string(headers, "session-id"));
let client_thread_id =
identity_context_string(&body, "thread_id").or_else(|| header_string(headers, "thread-id"));
let client_request_id = identity_context_string(&body, "x-client-request-id")
.or_else(|| header_string(headers, "x-client-request-id"));
let client_turn_id = identity_context_string(&body, "turn_id")
.or_else(|| header_string(headers, "x-codex-turn-id"));
let client_window_id = identity_context_string(&body, "x-codex-window-id")
.or_else(|| header_string(headers, "x-codex-window-id"));
let client_parent_thread_id = identity_context_string(&body, "x-codex-parent-thread-id")
.or_else(|| header_string(headers, "x-codex-parent-thread-id"));
body.remove(COMPACT_STREAM_KEY);
for key in TRANSPORT_ONLY_KEYS {
body.remove(key);
}
CodexCompactRequest {
body,
client_ip: None,
client_user_agent: user_agent_from_headers(headers),
client_api_key_id: None,
client_session_id,
client_thread_id,
client_request_id,
client_turn_id,
client_window_id,
client_parent_thread_id,
}
}
/// compact 端点不支持流式,转发前剥离 `stream`(transport 控制字段,非业务语义)。
const COMPACT_STREAM_KEY: &str = "stream";
fn inject_subagent_metadata(body: &mut Map<String, Value>, subagent: &str) {
let metadata = body
.entry("client_metadata")
@@ -286,24 +245,6 @@ pub async fn review_responses(
.await
}
/// `POST /v1/responses/compact`
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,
client_ip.map(|Extension(client_ip)| client_ip.as_str().to_string()),
headers,
body,
)
.await
}
async fn handle_responses(
state: AppState,
request_id: RequestId,
@@ -322,29 +263,39 @@ async fn handle_responses(
let Some(body) = parse_responses_body(&body) else {
return invalid_responses_request_response().into_response();
};
let model = match validated_responses_model(&state, request_model(&body)).await {
Ok(model) => model,
Err(ResponsesModelValidationError::InvalidRequest) => {
let prepared = match prepare_responses_request(
&state,
body,
&headers,
forced_subagent,
client_ip,
client_api_key_id,
)
.await
{
Ok(prepared) => prepared,
Err(ResponsesRequestValidationError::InvalidRequest) => {
return invalid_responses_request_response().into_response();
}
Err(ResponsesModelValidationError::ModelNotFound) => {
Err(ResponsesRequestValidationError::ModelNotFound) => {
return model_not_found_response().into_response();
}
};
let mut codex_request = build_codex_request(body, &headers, forced_subagent);
codex_request.client_ip = client_ip;
codex_request.client_api_key_id = Some(client_api_key_id);
let PreparedResponsesRequest {
mut request,
requested_model,
} = prepared;
// 客户端未显式指定 transport 时默认偏好 WebSocket。
if !codex_request.force_http_sse {
codex_request.use_websocket = true;
if !request.force_http_sse {
request.use_websocket = true;
}
let stream = codex_request.stream();
let stream = request.stream();
if stream {
return match state
.services
.responses
.stream(request_id.as_str(), route, codex_request, &model)
.stream(request_id.as_str(), route, request, &requested_model)
.await
{
Ok(stream) => live_event_stream_response(stream),
@@ -355,50 +306,7 @@ async fn handle_responses(
match state
.services
.responses
.complete(request_id.as_str(), route, codex_request, &model)
.await
{
Ok(result) => apply_safe_response_headers(
(StatusCode::OK, Json(result.body)).into_response(),
result.response_headers,
),
Err(error) => responses_dispatch_error_response(error),
}
}
async fn handle_compact_responses(
state: AppState,
request_id: RequestId,
client_ip: Option<String>,
headers: HeaderMap,
body: Bytes,
) -> Response {
let Some(client_api_key_id) =
crate::api::client::auth::authorized_client_api_key_id(&state, &headers).await
else {
return missing_client_api_key_response().into_response();
};
let Some(body) = parse_responses_body(&body) else {
return invalid_responses_request_response().into_response();
};
let model = match validated_responses_model(&state, request_model(&body)).await {
Ok(model) => model,
Err(ResponsesModelValidationError::InvalidRequest) => {
return invalid_responses_request_response().into_response();
}
Err(ResponsesModelValidationError::ModelNotFound) => {
return model_not_found_response().into_response();
}
};
let mut compact_request = build_compact_request(body, &headers);
compact_request.client_ip = client_ip;
compact_request.client_api_key_id = Some(client_api_key_id);
match state
.services
.responses
.compact(request_id.as_str(), compact_request, &model)
.complete(request_id.as_str(), route, request, &requested_model)
.await
{
Ok(result) => apply_safe_response_headers(
@@ -423,23 +331,47 @@ fn request_model(body: &Map<String, Value>) -> &str {
.unwrap_or_default()
}
pub(super) struct PreparedResponsesRequest {
pub(super) request: CodexResponsesRequest,
pub(super) requested_model: String,
}
pub(super) async fn prepare_responses_request(
state: &AppState,
body: Map<String, Value>,
headers: &HeaderMap,
forced_subagent: Option<&str>,
client_ip: Option<String>,
client_api_key_id: String,
) -> Result<PreparedResponsesRequest, ResponsesRequestValidationError> {
let requested_model = validated_responses_model(state, request_model(&body)).await?;
let mut request = build_codex_request(body, headers, forced_subagent);
request.client_ip = client_ip;
request.client_api_key_id = Some(client_api_key_id);
Ok(PreparedResponsesRequest {
request,
requested_model,
})
}
async fn validated_responses_model(
state: &AppState,
raw_model: &str,
) -> Result<String, ResponsesModelValidationError> {
) -> Result<String, ResponsesRequestValidationError> {
let model = raw_model.trim();
if model.is_empty() {
return Err(ResponsesModelValidationError::InvalidRequest);
return Err(ResponsesRequestValidationError::InvalidRequest);
}
let catalog = model_catalog_for_state(state).await;
if catalog.is_recognized_model_name(model) {
if catalog.is_empty() || catalog.is_recognized_model_name(model) {
Ok(model.to_string())
} else {
Err(ResponsesModelValidationError::ModelNotFound)
Err(ResponsesRequestValidationError::ModelNotFound)
}
}
enum ResponsesModelValidationError {
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum ResponsesRequestValidationError {
InvalidRequest,
ModelNotFound,
}
@@ -449,15 +381,11 @@ pub(crate) fn user_agent_from_headers(headers: &HeaderMap) -> Option<String> {
}
fn openai_subagent_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, OPENAI_SUBAGENT_HEADER).and_then(|value| {
if matches!(
header_string(headers, OPENAI_SUBAGENT_HEADER).filter(|value| {
matches!(
value.as_str(),
"review" | "compact" | "memory_consolidation" | "collab_spawn"
) {
Some(value)
} else {
None
}
)
})
}
@@ -5,15 +5,15 @@
use axum::{
body::Body,
http::{
header::{CACHE_CONTROL, CONTENT_TYPE},
StatusCode,
header::{CACHE_CONTROL, CONTENT_TYPE},
},
response::{IntoResponse, Response},
};
use crate::upstream::openai::protocol::sse::DONE_SSE_FRAME;
use super::errors::openai_error_response;
use crate::api::client::errors::openai_error_response;
/// 编码 OpenAI 流结束标记。
pub fn done_sse_frame() -> &'static str {
@@ -0,0 +1,507 @@
//! 官方 Responses WebSocket 入站协议。
use axum::{
body::Body,
extract::{Request, State},
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
};
use futures::{SinkExt, StreamExt};
use hyper::upgrade::Upgraded;
use hyper_util::rt::TokioIo;
use serde_json::{Map, Value, json};
use thiserror::Error;
use tokio_tungstenite::{
WebSocketStream,
tungstenite::{
Error as WebSocketError, Message,
handshake::server::create_response_with_body,
protocol::{CloseFrame, Role, WebSocketConfig, frame::coding::CloseCode},
},
};
use crate::{
api::{AppState, middleware::request_id::RequestId},
dispatch::service::ResponseDispatchStream,
upstream::openai::protocol::{
sse::{SseEvent, SseEventDecoder},
websocket::{is_terminal_websocket_event, websocket_event_type},
},
};
use super::{PreparedResponsesRequest, ResponsesRequestValidationError, prepare_responses_request};
use crate::api::client::{
auth::authorized_client_api_key_id,
errors::{
missing_client_api_key_response, openai_error_response,
responses_websocket_dispatch_error_event, responses_websocket_error_event,
},
router::MAX_CLIENT_REQUEST_BODY_BYTES,
};
const RESPONSES_ROUTE: &str = "/v1/responses";
const INVALID_REQUEST_MESSAGE: &str = "Invalid responses request";
const MODEL_NOT_FOUND_MESSAGE: &str = "Model not found";
const INVALID_STREAM_MESSAGE: &str = "Responses WebSocket requests require stream=true";
const INVALID_UPSTREAM_EVENT_MESSAGE: &str = "Invalid upstream Responses event";
const STREAM_DISCONNECTED_MESSAGE: &str = "Upstream stream closed before response.completed";
type ResponsesWebSocket = WebSocketStream<TokioIo<Upgraded>>;
/// 将已认证的 `GET /v1/responses` 请求升级为官方 Responses WebSocket 协议。
pub async fn responses_websocket(State(state): State<AppState>, mut request: Request) -> Response {
let headers = request.headers().clone();
let Some(client_api_key_id) = authorized_client_api_key_id(&state, &headers).await else {
return missing_client_api_key_response().into_response();
};
let connection_request_id = request
.extensions()
.get::<RequestId>()
.cloned()
.unwrap_or_else(RequestId::generate);
let client_ip = request
.extensions()
.get::<crate::api::middleware::request_id::ClientIp>()
.map(|client_ip| client_ip.as_str().to_string());
let connection_id = connection_request_id.as_str().to_string();
let response = match create_response_with_body(&request, Body::empty) {
Ok(response) => response,
Err(error) => {
tracing::info!(
websocket_connection_id = %connection_id,
error = %error,
"invalid Responses WebSocket upgrade request"
);
return openai_error_response(
StatusCode::BAD_REQUEST,
"Invalid Responses WebSocket upgrade request",
"invalid_request_error",
"invalid_websocket_upgrade",
)
.into_response();
}
};
let upgrade = hyper::upgrade::on(&mut request);
let connection_drain = state.services.connection_drain.clone();
connection_drain.spawn(async move {
match upgrade.await {
Ok(upgraded) => {
let socket = WebSocketStream::from_raw_socket(
TokioIo::new(upgraded),
Role::Server,
Some(downstream_websocket_config()),
)
.await;
serve_responses_websocket(
socket,
ResponsesWebSocketSession {
state,
headers,
client_ip,
client_api_key_id,
connection_id,
},
)
.await;
}
Err(error) => {
tracing::warn!(
websocket_connection_id = %connection_id,
error = %error,
"Responses WebSocket upgrade failed"
);
}
}
});
response
}
fn downstream_websocket_config() -> WebSocketConfig {
let mut config = WebSocketConfig::default();
config.max_message_size = Some(MAX_CLIENT_REQUEST_BODY_BYTES);
config.max_frame_size = Some(MAX_CLIENT_REQUEST_BODY_BYTES);
config
}
struct ResponsesWebSocketSession {
state: AppState,
headers: HeaderMap,
client_ip: Option<String>,
client_api_key_id: String,
connection_id: String,
}
async fn serve_responses_websocket(
mut socket: ResponsesWebSocket,
session: ResponsesWebSocketSession,
) {
tracing::info!(
websocket_connection_id = %session.connection_id,
"Responses WebSocket connected"
);
let mut request_count = 0u64;
while let Some(message) = socket.next().await {
let payload = match message {
Ok(Message::Text(payload)) => payload,
Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => continue,
Ok(Message::Close(_)) => break,
Ok(Message::Binary(_)) => {
close_for_protocol_violation(
&mut socket,
"Responses WebSocket accepts text frames only",
)
.await;
break;
}
Err(error) => {
tracing::info!(
websocket_connection_id = %session.connection_id,
error = %error,
"Responses WebSocket receive failed"
);
break;
}
};
request_count = request_count.saturating_add(1);
let request_id = RequestId::generate();
let body = match parse_response_create(payload.as_str()) {
Ok(body) => body,
Err(error) => {
tracing::info!(
websocket_connection_id = %session.connection_id,
request_id = %request_id.as_str(),
error = %error,
"Responses WebSocket request rejected"
);
if !send_error(
&mut socket,
StatusCode::BAD_REQUEST,
"invalid_request_error",
"invalid_request",
&error.to_string(),
request_id.as_str(),
)
.await
{
break;
}
continue;
}
};
let prepared = match prepare_responses_request(
&session.state,
body,
&session.headers,
None,
session.client_ip.clone(),
session.client_api_key_id.clone(),
)
.await
{
Ok(prepared) => prepared,
Err(error) => {
let (status, code, message) = match error {
ResponsesRequestValidationError::InvalidRequest => (
StatusCode::BAD_REQUEST,
"invalid_request",
INVALID_REQUEST_MESSAGE,
),
ResponsesRequestValidationError::ModelNotFound => (
StatusCode::NOT_FOUND,
"model_not_found",
MODEL_NOT_FOUND_MESSAGE,
),
};
if !send_error(
&mut socket,
status,
"invalid_request_error",
code,
message,
request_id.as_str(),
)
.await
{
break;
}
continue;
}
};
if !dispatch_response_create(&mut socket, &session, request_id, prepared).await {
break;
}
}
tracing::info!(
websocket_connection_id = %session.connection_id,
request_count,
"Responses WebSocket disconnected"
);
}
async fn dispatch_response_create(
socket: &mut ResponsesWebSocket,
session: &ResponsesWebSocketSession,
request_id: RequestId,
prepared: PreparedResponsesRequest,
) -> bool {
let PreparedResponsesRequest {
mut request,
requested_model,
} = prepared;
if !request.force_http_sse {
request.use_websocket = true;
}
let stream = session
.state
.services
.responses
.stream(
request_id.as_str(),
RESPONSES_ROUTE,
request,
&requested_model,
)
.await;
let stream = match stream {
Ok(stream) => stream,
Err(error) => {
let event = responses_websocket_dispatch_error_event(&error, request_id.as_str());
return send_text(socket, event).await;
}
};
forward_dispatch_stream(socket, stream, request_id.as_str()).await
}
async fn forward_dispatch_stream(
socket: &mut ResponsesWebSocket,
stream: ResponseDispatchStream,
request_id: &str,
) -> bool {
let ResponseDispatchStream {
body,
response_headers,
} = stream;
if !send_text(
socket,
response_metadata_event(request_id, &response_headers),
)
.await
{
return false;
}
let mut body = body;
let mut decoder = SseEventDecoder::default();
let mut terminal_sent = false;
loop {
let next = if terminal_sent {
ActiveResponseInput::Upstream(body.next().await)
} else {
tokio::select! {
next = body.next() => ActiveResponseInput::Upstream(next),
message = socket.next() => ActiveResponseInput::Downstream(message),
}
};
match next {
ActiveResponseInput::Upstream(Some(Ok(chunk))) => {
let events = match decoder.push(&chunk) {
Ok(events) => events,
Err(error) => {
tracing::warn!(request_id, error = %error, "failed to decode Responses SSE for WebSocket forwarding");
let sent = send_error(
socket,
StatusCode::BAD_GATEWAY,
"server_error",
"invalid_upstream_response",
INVALID_UPSTREAM_EVENT_MESSAGE,
request_id,
)
.await;
drop(body);
return sent;
}
};
if !forward_sse_events(socket, events, request_id, &mut terminal_sent).await {
return false;
}
}
ActiveResponseInput::Upstream(Some(Err(error))) => {
tracing::warn!(request_id, error = %error, "Responses dispatch stream failed");
if terminal_sent {
return true;
}
return send_error(
socket,
StatusCode::BAD_GATEWAY,
"server_error",
"upstream_error",
STREAM_DISCONNECTED_MESSAGE,
request_id,
)
.await;
}
ActiveResponseInput::Upstream(None) => {
let events = match decoder.finish() {
Ok(events) => events,
Err(error) => {
tracing::warn!(request_id, error = %error, "failed to finish Responses SSE decoding");
return send_error(
socket,
StatusCode::BAD_GATEWAY,
"server_error",
"invalid_upstream_response",
INVALID_UPSTREAM_EVENT_MESSAGE,
request_id,
)
.await;
}
};
if !forward_sse_events(socket, events, request_id, &mut terminal_sent).await {
return false;
}
if terminal_sent {
return true;
}
return send_error(
socket,
StatusCode::BAD_GATEWAY,
"server_error",
"stream_disconnected",
STREAM_DISCONNECTED_MESSAGE,
request_id,
)
.await;
}
ActiveResponseInput::Downstream(Some(Ok(
Message::Ping(_) | Message::Pong(_) | Message::Frame(_),
))) => {}
ActiveResponseInput::Downstream(Some(Ok(Message::Close(_))) | None) => return false,
ActiveResponseInput::Downstream(Some(Ok(Message::Text(_) | Message::Binary(_)))) => {
close_for_protocol_violation(
socket,
"Only one response.create may be active per connection",
)
.await;
return false;
}
ActiveResponseInput::Downstream(Some(Err(error))) => {
tracing::info!(request_id, error = %error, "Responses WebSocket client disconnected during response");
return false;
}
}
}
}
enum ActiveResponseInput {
Upstream(Option<Result<bytes::Bytes, crate::dispatch::errors::ResponseDispatchStreamError>>),
Downstream(Option<Result<Message, WebSocketError>>),
}
async fn forward_sse_events(
socket: &mut ResponsesWebSocket,
events: Vec<SseEvent>,
request_id: &str,
terminal_sent: &mut bool,
) -> bool {
for event in events {
let Some(event_type) = websocket_event_type(&event.data) else {
tracing::warn!(
request_id,
"Responses SSE event does not contain a valid JSON type"
);
let _ = send_error(
socket,
StatusCode::BAD_GATEWAY,
"server_error",
"invalid_upstream_response",
INVALID_UPSTREAM_EVENT_MESSAGE,
request_id,
)
.await;
return false;
};
if !send_text(socket, event.data).await {
return false;
}
*terminal_sent |= is_terminal_websocket_event(&event_type);
}
true
}
fn response_metadata_event(request_id: &str, response_headers: &[(String, String)]) -> String {
let mut headers = response_headers
.iter()
.map(|(name, value)| (name.clone(), Value::String(value.clone())))
.collect::<Map<String, Value>>();
headers.insert(
"x-request-id".to_string(),
Value::String(request_id.to_string()),
);
json!({
"type": "response.metadata",
"headers": headers,
})
.to_string()
}
fn parse_response_create(payload: &str) -> Result<Map<String, Value>, ResponseCreateFrameError> {
let Value::Object(mut body) = serde_json::from_str::<Value>(payload)? else {
return Err(ResponseCreateFrameError::ExpectedObject);
};
match body.remove("type") {
Some(Value::String(message_type)) if message_type == "response.create" => {}
_ => return Err(ResponseCreateFrameError::UnsupportedType),
}
if body.get("stream").and_then(Value::as_bool) != Some(true) {
return Err(ResponseCreateFrameError::StreamingRequired);
}
Ok(body)
}
#[derive(Debug, Error)]
enum ResponseCreateFrameError {
#[error("Invalid response.create JSON: {0}")]
InvalidJson(#[from] serde_json::Error),
#[error("A response.create payload must be a JSON object")]
ExpectedObject,
#[error("Unsupported Responses WebSocket message type")]
UnsupportedType,
#[error("{INVALID_STREAM_MESSAGE}")]
StreamingRequired,
}
async fn send_error(
socket: &mut ResponsesWebSocket,
status: StatusCode,
error_type: &str,
code: &str,
message: &str,
request_id: &str,
) -> bool {
send_text(
socket,
responses_websocket_error_event(status, error_type, code, message, request_id),
)
.await
}
async fn send_text(socket: &mut ResponsesWebSocket, payload: String) -> bool {
if let Err(error) = socket.send(Message::Text(payload.into())).await {
tracing::info!(error = %error, "Responses WebSocket send failed");
return false;
}
true
}
async fn close_for_protocol_violation(socket: &mut ResponsesWebSocket, reason: &'static str) {
let _ = socket
.send(Message::Close(Some(CloseFrame {
code: CloseCode::Policy,
reason: reason.into(),
})))
.await;
}
+3 -4
View File
@@ -1,16 +1,16 @@
//! OpenAI 兼容 API 路由。
use axum::{
Router,
extract::DefaultBodyLimit,
routing::{get, post},
Router,
};
use crate::api::AppState;
use super::{
models::{model_catalog, model_detail, model_info, models},
responses::{compact_responses, responses, review_responses},
responses::{responses, responses_websocket, review_responses},
};
pub const MAX_CLIENT_REQUEST_BODY_BYTES: usize = 16 * 1024 * 1024;
@@ -18,9 +18,8 @@ pub const MAX_CLIENT_REQUEST_BODY_BYTES: usize = 16 * 1024 * 1024;
/// 构造 OpenAI 兼容 API 路由。
pub fn router() -> Router<AppState> {
Router::new()
.route("/v1/responses", post(responses))
.route("/v1/responses", get(responses_websocket).post(responses))
.route("/v1/responses/review", post(review_responses))
.route("/v1/responses/compact", post(compact_responses))
.route("/v1/models", get(models))
.route("/v1/models/catalog", get(model_catalog))
.route("/v1/models/{model_id}", get(model_detail))
@@ -0,0 +1,69 @@
//! 进程关闭时统一排空入站 HTTP 流与 WebSocket 任务。
use std::future::Future;
use axum::{
body::Body,
extract::{Request, State},
middleware::Next,
response::Response,
};
use futures::StreamExt;
use tokio_util::{sync::CancellationToken, task::TaskTracker};
/// 入站连接排空控制器。
#[derive(Clone)]
pub struct ConnectionDrain {
cancellation: CancellationToken,
connection_tasks: TaskTracker,
}
impl Default for ConnectionDrain {
fn default() -> Self {
Self {
cancellation: CancellationToken::new(),
connection_tasks: TaskTracker::new(),
}
}
}
impl ConnectionDrain {
/// 注册需要随进程关闭而终止的长连接任务。
pub fn spawn(&self, future: impl Future<Output = ()> + Send + 'static) {
let cancellation = self.cancellation.clone();
drop(self.connection_tasks.spawn(async move {
cancellation.run_until_cancelled(future).await;
}));
}
/// 停止所有受管连接,并返回关闭前仍在运行的长连接任务数。
pub fn begin_shutdown(&self) -> usize {
let active_websocket_connections = self.connection_tasks.len();
self.connection_tasks.close();
self.cancellation.cancel();
active_websocket_connections
}
/// 等待所有受管长连接任务退出。
pub async fn wait(&self) {
self.connection_tasks.wait().await;
}
/// 是否已进入连接排空阶段。
pub fn is_shutting_down(&self) -> bool {
self.cancellation.is_cancelled()
}
}
/// 将所有 HTTP 响应体接入统一关闭信号。
pub async fn drain_response_body(
State(connection_drain): State<ConnectionDrain>,
request: Request,
next: Next,
) -> Response {
let response = next.run(request).await;
let (parts, body) = response.into_parts();
let shutdown = connection_drain.cancellation.clone().cancelled_owned();
let body = Body::from_stream(body.into_data_stream().take_until(shutdown));
Response::from_parts(parts, body)
}
+2 -1
View File
@@ -1,4 +1,5 @@
//! HTTP 中间件 —— 请求 ID 与 Tracing。
//! HTTP 中间件 —— 连接排空、请求 ID 与 Tracing。
pub mod connection_drain;
pub mod request_id;
pub mod trace;
+1 -1
View File
@@ -4,7 +4,7 @@ use std::net::{IpAddr, SocketAddr};
use std::sync::LazyLock;
use axum::{
extract::{connect_info::ConnectInfo, Request},
extract::{Request, connect_info::ConnectInfo},
http::{HeaderMap, HeaderValue},
middleware::Next,
response::Response,
+1 -1
View File
@@ -10,7 +10,7 @@ use tower_http::{
classify::{ServerErrorsAsFailures, ServerErrorsFailureClass, SharedClassifier},
trace::TraceLayer,
};
use tracing::{debug, info, warn, Span};
use tracing::{Span, debug, info, warn};
use crate::api::middleware::request_id::{ClientIp, RequestId};
+4 -2
View File
@@ -6,7 +6,7 @@ use std::{
};
use async_trait::async_trait;
use axum::{extract::State, http::StatusCode, middleware, routing::get, Router};
use axum::{Router, extract::State, http::StatusCode, middleware, routing::get};
use crate::{
api::{
@@ -26,6 +26,7 @@ use crate::{
account_usage::query::AccountUsageQueryService, ops::query::OpsQueryService,
usage::query::UsageQueryService,
},
update::service::SystemUpdateService,
upstream::openai::{fingerprint::RuntimeFingerprint, token_client::OpenAiTokenClient},
};
@@ -49,13 +50,14 @@ pub struct ApiServices {
pub session_affinity: Arc<SessionAffinityService>,
pub fingerprint: RuntimeFingerprint,
pub process_control: Arc<dyn ProcessControl>,
pub system_update: Arc<SystemUpdateService>,
pub connection_drain: crate::api::middleware::connection_drain::ConnectionDrain,
}
/// 由 bootstrap 实现的进程生命周期控制端口。
pub trait ProcessControl: Send + Sync + 'static {
fn request_shutdown(&self);
fn request_restart(&self, executable_path: PathBuf);
fn subscribe_shutdown(&self) -> tokio::sync::broadcast::Receiver<()>;
}
/// 由 bootstrap 实现的基础设施健康检查端口。
+1 -1
View File
@@ -1,6 +1,6 @@
use std::{collections::BTreeMap, env, path::Path};
use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
use percent_encoding::{NON_ALPHANUMERIC, utf8_percent_encode};
use serde::{Deserialize, Serialize};
use url::Url;
-460
View File
@@ -1,460 +0,0 @@
use futures::TryStreamExt;
use sqlx::{types::Json, Postgres, Row, SqlitePool, Transaction};
use crate::infra::identity::hash_credential;
use super::{
parse_cookie_expiration, parse_json, parse_optional_json, parse_optional_timestamp,
parse_timestamp, ImportSqliteError, ImportSqliteReport,
};
pub(super) async fn import_core_tables(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
import_admin_users(source, target, report).await?;
import_client_api_keys(source, target, report).await?;
import_runtime_settings(source, target, report).await?;
import_accounts(source, target, report).await?;
import_account_usage(source, target, report).await?;
import_account_cookies(source, target, report).await?;
import_fingerprints(source, target, report).await?;
import_fingerprint_history(source, target, report).await
}
async fn import_admin_users(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, password_hash, created_at, updated_at from admin_users order by id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
sqlx::query(
"insert into admin_users (id, password_hash, created_at, updated_at)
values ($1, $2, $3, $4)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("password_hash")?)
.bind(parse_timestamp(
"admin_users",
"created_at",
row.try_get("created_at")?,
)?)
.bind(parse_timestamp(
"admin_users",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("admin_users");
}
Ok(())
}
async fn import_client_api_keys(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, name, prefix, key, label, enabled, created_at, last_used_at
from client_api_keys order by id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
sqlx::query(
"insert into client_api_keys
(id, name, prefix, key, label, enabled, created_at, last_used_at)
values ($1, $2, $3, $4, $5, $6, $7, $8)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("name")?)
.bind(row.try_get::<String, _>("prefix")?)
.bind(row.try_get::<String, _>("key")?)
.bind(row.try_get::<Option<String>, _>("label")?)
.bind(row.try_get::<i64, _>("enabled")? == 1)
.bind(parse_timestamp(
"client_api_keys",
"created_at",
row.try_get("created_at")?,
)?)
.bind(parse_optional_timestamp(
"client_api_keys",
"last_used_at",
row.try_get("last_used_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("client_api_keys");
}
Ok(())
}
async fn import_runtime_settings(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, model_aliases_json, refresh_margin_seconds, refresh_concurrency,
max_concurrent_per_account, request_interval_ms, rotation_strategy,
admin_api_key, updated_at
from runtime_settings order by id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
let aliases: String = row.try_get("model_aliases_json")?;
let admin_key = row
.try_get::<Option<String>, _>("admin_api_key")?
.and_then(|key| (!key.trim().is_empty()).then(|| hash_credential(&key)));
sqlx::query(
"insert into runtime_settings
(id, model_aliases_json, refresh_margin_seconds, refresh_concurrency,
max_concurrent_per_account, request_interval_ms, rotation_strategy,
admin_api_key_hash, usage_retention_days, ops_error_retention_days,
bucket_retention_days, updated_at)
values ($1, $2, $3, $4, $5, $6, $7, $8, 30, 30, 90, $9)",
)
.bind(row.try_get::<i64, _>("id")?)
.bind(Json(parse_json(
"runtime_settings",
"model_aliases_json",
&aliases,
)?))
.bind(row.try_get::<i64, _>("refresh_margin_seconds")?)
.bind(row.try_get::<i64, _>("refresh_concurrency")?)
.bind(row.try_get::<i64, _>("max_concurrent_per_account")?)
.bind(row.try_get::<i64, _>("request_interval_ms")?)
.bind(row.try_get::<String, _>("rotation_strategy")?)
.bind(admin_key)
.bind(parse_timestamp(
"runtime_settings",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("runtime_settings");
}
Ok(())
}
async fn import_accounts(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, email, chatgpt_account_id, chatgpt_user_id, label, plan_type,
access_token, refresh_token, access_token_expires_at, next_refresh_at,
status, quota_json, quota_fetched_at, quota_limit_reached,
quota_verify_required, quota_cooldown_until, cloudflare_cooldown_until,
added_at, updated_at
from accounts order by id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
let source_status = row.try_get::<String, _>("status")?;
let (status, next_refresh_at) = if source_status == "refreshing" {
report.add_normalized("accounts.refreshing_to_expired");
("expired", None)
} else {
(
source_status.as_str(),
parse_optional_timestamp(
"accounts",
"next_refresh_at",
row.try_get("next_refresh_at")?,
)?,
)
};
let quota = parse_optional_json(
"accounts",
"quota_json",
row.try_get::<Option<String>, _>("quota_json")?,
)?
.map(Json);
sqlx::query(
"insert into accounts
(id, email, chatgpt_account_id, chatgpt_user_id, label, plan_type,
access_token, refresh_token, access_token_expires_at, next_refresh_at,
status, quota_json, quota_fetched_at, quota_limit_reached,
quota_verify_required, quota_cooldown_until, cloudflare_cooldown_until,
added_at, updated_at)
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
$14, $15, $16, $17, $18, $19)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<Option<String>, _>("email")?)
.bind(row.try_get::<Option<String>, _>("chatgpt_account_id")?)
.bind(row.try_get::<Option<String>, _>("chatgpt_user_id")?)
.bind(row.try_get::<Option<String>, _>("label")?)
.bind(row.try_get::<Option<String>, _>("plan_type")?)
.bind(row.try_get::<String, _>("access_token")?)
.bind(row.try_get::<Option<String>, _>("refresh_token")?)
.bind(parse_optional_timestamp(
"accounts",
"access_token_expires_at",
row.try_get("access_token_expires_at")?,
)?)
.bind(next_refresh_at)
.bind(status)
.bind(quota)
.bind(parse_optional_timestamp(
"accounts",
"quota_fetched_at",
row.try_get("quota_fetched_at")?,
)?)
.bind(row.try_get::<i64, _>("quota_limit_reached")? == 1)
.bind(row.try_get::<i64, _>("quota_verify_required")? == 1)
.bind(parse_optional_timestamp(
"accounts",
"quota_cooldown_until",
row.try_get("quota_cooldown_until")?,
)?)
.bind(parse_optional_timestamp(
"accounts",
"cloudflare_cooldown_until",
row.try_get("cloudflare_cooldown_until")?,
)?)
.bind(parse_timestamp(
"accounts",
"added_at",
row.try_get("added_at")?,
)?)
.bind(parse_timestamp(
"accounts",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("accounts");
}
Ok(())
}
async fn import_account_usage(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query("select * from account_usage order by account_id").fetch(source);
while let Some(row) = rows.try_next().await? {
sqlx::query(
"insert into account_usage
(account_id, request_count, empty_response_count, input_tokens, output_tokens,
cached_tokens, reasoning_tokens, total_tokens, image_input_tokens,
image_output_tokens, image_request_count, image_request_failed_count,
window_request_count, window_input_tokens, window_output_tokens,
window_cached_tokens, window_image_input_tokens, window_image_output_tokens,
window_image_request_count, window_image_request_failed_count,
window_started_at, window_reset_at, limit_window_seconds, last_used_at)
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12,
$13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24)",
)
.bind(row.try_get::<String, _>("account_id")?)
.bind(row.try_get::<i64, _>("request_count")?)
.bind(row.try_get::<i64, _>("empty_response_count")?)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<i64, _>("cached_tokens")?)
.bind(row.try_get::<i64, _>("reasoning_tokens")?)
.bind(row.try_get::<i64, _>("total_tokens")?)
.bind(row.try_get::<i64, _>("image_input_tokens")?)
.bind(row.try_get::<i64, _>("image_output_tokens")?)
.bind(row.try_get::<i64, _>("image_request_count")?)
.bind(row.try_get::<i64, _>("image_request_failed_count")?)
.bind(row.try_get::<i64, _>("window_request_count")?)
.bind(row.try_get::<i64, _>("window_input_tokens")?)
.bind(row.try_get::<i64, _>("window_output_tokens")?)
.bind(row.try_get::<i64, _>("window_cached_tokens")?)
.bind(row.try_get::<i64, _>("window_image_input_tokens")?)
.bind(row.try_get::<i64, _>("window_image_output_tokens")?)
.bind(row.try_get::<i64, _>("window_image_request_count")?)
.bind(row.try_get::<i64, _>("window_image_request_failed_count")?)
.bind(parse_optional_timestamp(
"account_usage",
"window_started_at",
row.try_get("window_started_at")?,
)?)
.bind(parse_optional_timestamp(
"account_usage",
"window_reset_at",
row.try_get("window_reset_at")?,
)?)
.bind(row.try_get::<Option<i64>, _>("limit_window_seconds")?)
.bind(parse_optional_timestamp(
"account_usage",
"last_used_at",
row.try_get("last_used_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("account_usage");
}
Ok(())
}
async fn import_account_cookies(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, account_id, domain, name, value, path, expires_at, updated_at
from account_cookies order by id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
let expires_at = row
.try_get::<Option<String>, _>("expires_at")?
.and_then(|value| match parse_cookie_expiration(&value) {
Some(timestamp) => Some(timestamp),
None => {
report.add_cookie_expiration_parse_failure();
None
}
});
sqlx::query(
"insert into account_cookies
(id, account_id, domain, name, value, path, expires_at, updated_at)
values ($1, $2, $3, $4, $5, $6, $7, $8)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("account_id")?)
.bind(row.try_get::<String, _>("domain")?)
.bind(row.try_get::<String, _>("name")?)
.bind(row.try_get::<String, _>("value")?)
.bind(row.try_get::<String, _>("path")?)
.bind(expires_at)
.bind(parse_timestamp(
"account_cookies",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("account_cookies");
}
Ok(())
}
async fn import_fingerprints(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query("select * from fingerprints order by id").fetch(source);
while let Some(row) = rows.try_next().await? {
let default_headers: String = row.try_get("default_headers_json")?;
let header_order: String = row.try_get("header_order_json")?;
sqlx::query(
"insert into fingerprints
(id, originator, app_version, build_number, platform, arch, chromium_version,
user_agent_template, default_headers_json, header_order_json, source,
created_at, updated_at)
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("originator")?)
.bind(row.try_get::<String, _>("app_version")?)
.bind(row.try_get::<String, _>("build_number")?)
.bind(row.try_get::<String, _>("platform")?)
.bind(row.try_get::<String, _>("arch")?)
.bind(row.try_get::<String, _>("chromium_version")?)
.bind(row.try_get::<String, _>("user_agent_template")?)
.bind(Json(parse_json(
"fingerprints",
"default_headers_json",
&default_headers,
)?))
.bind(Json(parse_json(
"fingerprints",
"header_order_json",
&header_order,
)?))
.bind(row.try_get::<String, _>("source")?)
.bind(parse_timestamp(
"fingerprints",
"created_at",
row.try_get("created_at")?,
)?)
.bind(parse_timestamp(
"fingerprints",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("fingerprints");
}
Ok(())
}
async fn import_fingerprint_history(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows =
sqlx::query("select * from fingerprint_update_history order by id").fetch(source);
while let Some(row) = rows.try_next().await? {
let manifest = parse_optional_json(
"fingerprint_update_history",
"manifest_json",
row.try_get::<Option<String>, _>("manifest_json")?,
)?
.map(Json);
sqlx::query(
"insert into fingerprint_update_history
(id, current_fingerprint_id, app_version, build_number, chromium_version,
source, manifest_json, created_at)
values ($1, $2, $3, $4, $5, $6, $7, $8)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("current_fingerprint_id")?)
.bind(row.try_get::<String, _>("app_version")?)
.bind(row.try_get::<String, _>("build_number")?)
.bind(row.try_get::<Option<String>, _>("chromium_version")?)
.bind(row.try_get::<String, _>("source")?)
.bind(manifest)
.bind(parse_timestamp(
"fingerprint_update_history",
"created_at",
row.try_get("created_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("fingerprint_update_history");
}
Ok(())
}
pub(super) async fn count_discarded_runtime_tables(
source: &SqlitePool,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let (sessions, affinities, leases, snapshots, model_usage): (i64, i64, i64, i64, i64) =
sqlx::query_as(
"select
(select count(*) from admin_sessions),
(select count(*) from session_affinities),
(select count(*) from account_refresh_leases),
(select count(*) from model_plan_snapshots),
(select count(*) from account_model_usage)",
)
.fetch_one(source)
.await?;
report.add_discarded("admin_sessions", sessions.max(0) as u64);
report.add_discarded("session_affinities", affinities.max(0) as u64);
report.add_discarded("account_refresh_leases", leases.max(0) as u64);
report.add_discarded("model_plan_snapshots", snapshots.max(0) as u64);
report.add_discarded("account_model_usage", model_usage.max(0) as u64);
Ok(())
}
-232
View File
@@ -1,232 +0,0 @@
//! 一次性 SQLite v3 到 PostgreSQL 终态库导入。
mod core;
mod telemetry;
use std::{collections::BTreeMap, fmt, path::Path};
use chrono::{DateTime, Utc};
use serde_json::Value;
use sqlx::{
sqlite::{SqliteConnectOptions, SqlitePoolOptions},
PgPool, Postgres, SqlitePool, Transaction,
};
use thiserror::Error;
const SQLITE_SCHEMA_VERSION: i64 = 3;
/// SQLite 历史库导入结果。
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct ImportSqliteReport {
imported: BTreeMap<&'static str, u64>,
normalized: BTreeMap<&'static str, u64>,
discarded: BTreeMap<&'static str, u64>,
cookie_expiration_parse_failures: u64,
}
impl ImportSqliteReport {
fn add_imported(&mut self, table: &'static str) {
*self.imported.entry(table).or_default() += 1;
}
fn add_discarded(&mut self, reason: &'static str, count: u64) {
if count > 0 {
*self.discarded.entry(reason).or_default() += count;
}
}
fn add_normalized(&mut self, reason: &'static str) {
*self.normalized.entry(reason).or_default() += 1;
}
fn add_cookie_expiration_parse_failure(&mut self) {
self.cookie_expiration_parse_failures += 1;
}
/// 返回指定目标表导入的行数。
pub fn imported_rows(&self, table: &str) -> u64 {
self.imported.get(table).copied().unwrap_or_default()
}
/// 返回指定原因丢弃的行数。
pub fn discarded_rows(&self, reason: &str) -> u64 {
self.discarded.get(reason).copied().unwrap_or_default()
}
/// 返回指定规则规范化的行数。
pub fn normalized_rows(&self, reason: &str) -> u64 {
self.normalized.get(reason).copied().unwrap_or_default()
}
/// 返回无法解析并降级为会话 Cookie 的过期时间数量。
pub fn cookie_expiration_parse_failures(&self) -> u64 {
self.cookie_expiration_parse_failures
}
}
impl fmt::Display for ImportSqliteReport {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(formatter, "SQLite v3 import completed")?;
for (table, count) in &self.imported {
writeln!(formatter, " imported {table}: {count}")?;
}
for (reason, count) in &self.normalized {
writeln!(formatter, " normalized {reason}: {count}")?;
}
for (reason, count) in &self.discarded {
writeln!(formatter, " discarded {reason}: {count}")?;
}
write!(
formatter,
" cookie expiration parse failures: {}",
self.cookie_expiration_parse_failures
)
}
}
/// SQLite 历史库导入错误。
#[derive(Debug, Error)]
pub enum ImportSqliteError {
#[error("database operation failed: {0}")]
Database(#[from] sqlx::Error),
#[error("unsupported SQLite schema version {actual:?}; expected 3")]
UnsupportedSourceVersion { actual: Option<i64> },
#[error("target PostgreSQL database is not empty")]
TargetNotEmpty,
#[error("invalid timestamp in {table}.{column}: {value}")]
InvalidTimestamp {
table: &'static str,
column: &'static str,
value: String,
},
#[error("invalid JSON in {table}.{column}: {source}")]
InvalidJson {
table: &'static str,
column: &'static str,
#[source]
source: serde_json::Error,
},
}
/// 将只读 SQLite v3 历史库一次性导入空的 PostgreSQL 终态库。
pub async fn import_sqlite(
target: &PgPool,
source_path: impl AsRef<Path>,
) -> Result<ImportSqliteReport, ImportSqliteError> {
let options = SqliteConnectOptions::new()
.filename(source_path.as_ref())
.read_only(true)
.create_if_missing(false);
let source = SqlitePoolOptions::new()
.max_connections(1)
.connect_with(options)
.await?;
validate_source_version(&source).await?;
let mut transaction = target.begin().await?;
reject_nonempty_target(&mut transaction).await?;
let mut report = ImportSqliteReport::default();
core::import_core_tables(&source, &mut transaction, &mut report).await?;
telemetry::import_telemetry_tables(&source, &mut transaction, &mut report).await?;
core::count_discarded_runtime_tables(&source, &mut report).await?;
transaction.commit().await?;
source.close().await;
Ok(report)
}
async fn validate_source_version(source: &SqlitePool) -> Result<(), ImportSqliteError> {
let (actual,): (Option<i64>,) = sqlx::query_as("select max(version) from schema_migrations")
.fetch_one(source)
.await?;
if actual != Some(SQLITE_SCHEMA_VERSION) {
return Err(ImportSqliteError::UnsupportedSourceVersion { actual });
}
Ok(())
}
async fn reject_nonempty_target(
transaction: &mut Transaction<'_, Postgres>,
) -> Result<(), ImportSqliteError> {
let (rows,): (i64,) = sqlx::query_as(
"select
(select count(*) from admin_users)
+ (select count(*) from client_api_keys)
+ (select count(*) from runtime_settings)
+ (select count(*) from accounts)
+ (select count(*) from account_usage)
+ (select count(*) from usage_records)
+ (select count(*) from ops_error_logs)
+ (select count(*) from request_time_buckets)
+ (select count(*) from account_cookies)
+ (select count(*) from fingerprints)
+ (select count(*) from fingerprint_update_history)",
)
.fetch_one(&mut **transaction)
.await?;
if rows != 0 {
return Err(ImportSqliteError::TargetNotEmpty);
}
Ok(())
}
fn parse_timestamp(
table: &'static str,
column: &'static str,
value: String,
) -> Result<DateTime<Utc>, ImportSqliteError> {
DateTime::parse_from_rfc3339(&value)
.map(|timestamp| timestamp.with_timezone(&Utc))
.map_err(|_| ImportSqliteError::InvalidTimestamp {
table,
column,
value,
})
}
fn parse_optional_timestamp(
table: &'static str,
column: &'static str,
value: Option<String>,
) -> Result<Option<DateTime<Utc>>, ImportSqliteError> {
value
.map(|value| parse_timestamp(table, column, value))
.transpose()
}
fn parse_cookie_expiration(value: &str) -> Option<DateTime<Utc>> {
DateTime::parse_from_rfc3339(value)
.or_else(|_| DateTime::parse_from_rfc2822(value))
.ok()
.map(|timestamp| timestamp.with_timezone(&Utc))
}
fn parse_json(
table: &'static str,
column: &'static str,
value: &str,
) -> Result<Value, ImportSqliteError> {
serde_json::from_str(value).map_err(|source| ImportSqliteError::InvalidJson {
table,
column,
source,
})
}
fn parse_optional_json(
table: &'static str,
column: &'static str,
value: Option<String>,
) -> Result<Option<Value>, ImportSqliteError> {
value
.map(|value| parse_json(table, column, &value))
.transpose()
}
fn normalized_dimension(value: Option<String>) -> String {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| "__unknown__".to_string())
}
@@ -1,330 +0,0 @@
use futures::TryStreamExt;
use serde_json::Value;
use sqlx::{types::Json, Postgres, Row, SqlitePool, Transaction};
use super::{
normalized_dimension, parse_json, parse_timestamp, ImportSqliteError, ImportSqliteReport,
};
pub(super) async fn import_telemetry_tables(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
import_mixed_usage_records(source, target, report).await?;
import_ops_error_logs(source, target, report).await?;
import_request_time_buckets(source, target, report).await
}
async fn import_mixed_usage_records(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, request_id, kind, level, account_id, route, model, status_code,
transport, attempt_index, upstream_status_code, failure_class,
response_id, upstream_request_id, latency_ms, message, metadata_json,
created_at
from usage_records order by created_at, id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
let level: String = row.try_get("level")?;
let account_id: Option<String> = row.try_get("account_id")?;
let model: Option<String> = row.try_get("model")?;
let status_code: Option<i64> = row.try_get("status_code")?;
let metadata_text: String = row.try_get("metadata_json")?;
let metadata = parse_json("usage_records", "metadata_json", &metadata_text)?;
if level == "error" {
insert_usage_error(target, &row, metadata).await?;
report.add_imported("ops_error_logs");
continue;
}
let valid_account = account_id
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty());
let valid_model = model
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty());
let valid_status = status_code.is_some_and(|status| (200..=399).contains(&status));
if !valid_account || !valid_model || !valid_status {
report.add_discarded("usage_record_noise", 1);
continue;
}
insert_success_record(
target,
&row,
account_id.expect("validated account id"),
model.expect("validated model"),
status_code.expect("validated status code"),
metadata,
)
.await?;
report.add_imported("usage_records");
}
Ok(())
}
async fn insert_success_record(
target: &mut Transaction<'_, Postgres>,
row: &sqlx::sqlite::SqliteRow,
account_id: String,
model: String,
status_code: i64,
mut metadata: Value,
) -> Result<(), ImportSqliteError> {
let requested_model = metadata_string(&metadata, &["requestedModel", "requested_model"]);
let upstream_model = metadata_string(&metadata, &["upstreamModel", "upstream_model"]);
let service_tier = metadata_string(&metadata, &["serviceTier", "service_tier"]);
let first_token_ms = metadata_nonnegative(&metadata, &["firstTokenMs", "first_token_ms"]);
let input_tokens = usage_number(&metadata, &["inputTokens", "input_tokens"]);
let output_tokens = usage_number(&metadata, &["outputTokens", "output_tokens"]);
let cached_tokens = usage_number(&metadata, &["cachedTokens", "cached_tokens"]);
let reasoning_tokens = usage_number(&metadata, &["reasoningTokens", "reasoning_tokens"]);
strip_promoted_success_fields(&mut metadata);
sqlx::query(
"insert into usage_records
(id, request_id, client_api_key_id, kind, route, provider, account_id, model,
requested_model, upstream_model, service_tier, status_code, transport,
attempt_index, response_id, upstream_request_id, latency_ms, first_token_ms,
input_tokens, output_tokens, cached_tokens, reasoning_tokens, message,
metadata_json, created_at)
values ($1, $2, null, $3, $4, 'openai', $5, $6, $7, $8, $9, $10,
$11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<Option<String>, _>("request_id")?)
.bind(row.try_get::<String, _>("kind")?)
.bind(row.try_get::<Option<String>, _>("route")?)
.bind(account_id)
.bind(model)
.bind(requested_model)
.bind(upstream_model)
.bind(service_tier)
.bind(status_code as i32)
.bind(row.try_get::<Option<String>, _>("transport")?)
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
.bind(row.try_get::<Option<String>, _>("response_id")?)
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
.bind(first_token_ms)
.bind(input_tokens)
.bind(output_tokens)
.bind(cached_tokens)
.bind(reasoning_tokens)
.bind(row.try_get::<String, _>("message")?)
.bind(Json(metadata))
.bind(parse_timestamp(
"usage_records",
"created_at",
row.try_get("created_at")?,
)?)
.execute(&mut **target)
.await?;
Ok(())
}
async fn insert_usage_error(
target: &mut Transaction<'_, Postgres>,
row: &sqlx::sqlite::SqliteRow,
metadata: Value,
) -> Result<(), ImportSqliteError> {
sqlx::query(
"insert into ops_error_logs
(id, request_id, client_api_key_id, kind, provider, 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 ($1, $2, null, $3, 'openai', $4, $5, $6, $7, null, $8, $9,
$10, $11, $12, $13, $14, $15, $16, $17)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<Option<String>, _>("request_id")?)
.bind(row.try_get::<String, _>("kind")?)
.bind(row.try_get::<Option<String>, _>("account_id")?)
.bind(row.try_get::<Option<String>, _>("route")?)
.bind(row.try_get::<Option<String>, _>("model")?)
.bind(optional_status(row.try_get("status_code")?))
.bind(optional_status(row.try_get("upstream_status_code")?))
.bind(row.try_get::<Option<String>, _>("transport")?)
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
.bind(row.try_get::<Option<String>, _>("failure_class")?)
.bind(row.try_get::<Option<String>, _>("response_id")?)
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
.bind(row.try_get::<String, _>("message")?)
.bind(Json(metadata))
.bind(parse_timestamp(
"usage_records",
"created_at",
row.try_get("created_at")?,
)?)
.execute(&mut **target)
.await?;
Ok(())
}
async fn import_ops_error_logs(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows =
sqlx::query("select * from ops_error_logs order by created_at, id").fetch(source);
while let Some(row) = rows.try_next().await? {
let metadata_text: String = row.try_get("metadata_json")?;
sqlx::query(
"insert into ops_error_logs
(id, request_id, client_api_key_id, kind, provider, 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 ($1, $2, null, $3, 'openai', $4, $5, $6, $7, $8, $9, $10,
$11, $12, $13, $14, $15, $16, $17, $18)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<Option<String>, _>("request_id")?)
.bind(row.try_get::<String, _>("kind")?)
.bind(row.try_get::<Option<String>, _>("account_id")?)
.bind(row.try_get::<Option<String>, _>("route")?)
.bind(row.try_get::<Option<String>, _>("model")?)
.bind(optional_status(row.try_get("status_code")?))
.bind(optional_status(row.try_get("client_status_code")?))
.bind(optional_status(row.try_get("upstream_status_code")?))
.bind(row.try_get::<Option<String>, _>("transport")?)
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
.bind(row.try_get::<Option<String>, _>("failure_class")?)
.bind(row.try_get::<Option<String>, _>("response_id")?)
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
.bind(row.try_get::<String, _>("message")?)
.bind(Json(parse_json(
"ops_error_logs",
"metadata_json",
&metadata_text,
)?))
.bind(parse_timestamp(
"ops_error_logs",
"created_at",
row.try_get("created_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("ops_error_logs");
}
Ok(())
}
async fn import_request_time_buckets(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows =
sqlx::query("select * from usage_time_buckets order by bucket_start").fetch(source);
while let Some(row) = rows.try_next().await? {
let request_count: i64 = row.try_get("request_count")?;
let error_count: i64 = row.try_get("error_count")?;
let min_latency_ms: i64 = row.try_get("min_latency_ms")?;
sqlx::query(
"insert into request_time_buckets
(bucket_start, provider, account_id, model, service_tier, success_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, 'openai', $2, $3, $4, $5, $6, $7, $8, $9, $10, $11,
$12, $13, $14, $15, $16)",
)
.bind(parse_timestamp(
"usage_time_buckets",
"bucket_start",
row.try_get("bucket_start")?,
)?)
.bind(normalized_dimension(row.try_get("account_id")?))
.bind(normalized_dimension(row.try_get("model")?))
.bind(normalized_dimension(row.try_get("service_tier")?))
.bind(request_count.saturating_sub(error_count).max(0))
.bind(error_count)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<i64, _>("cached_tokens")?)
.bind(row.try_get::<i64, _>("first_token_latency_sum")?)
.bind(row.try_get::<i64, _>("first_token_latency_count")?)
.bind(row.try_get::<i64, _>("latency_sum")?)
.bind(row.try_get::<i64, _>("latency_count")?)
.bind(row.try_get::<i64, _>("max_latency_ms")?)
.bind((min_latency_ms != 0).then_some(min_latency_ms))
.bind(parse_timestamp(
"usage_time_buckets",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("request_time_buckets");
}
Ok(())
}
fn optional_status(status: Option<i64>) -> Option<i32> {
status.and_then(|status| i32::try_from(status).ok())
}
fn metadata_string(metadata: &Value, keys: &[&str]) -> Option<String> {
keys.iter().find_map(|key| {
metadata
.get(*key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
})
}
fn metadata_nonnegative(metadata: &Value, keys: &[&str]) -> Option<i64> {
keys.iter()
.find_map(|key| metadata.get(*key).and_then(Value::as_i64))
.filter(|value| *value >= 0)
}
fn usage_number(metadata: &Value, keys: &[&str]) -> Option<i64> {
metadata
.get("usage")
.and_then(|usage| metadata_nonnegative(usage, keys))
.or_else(|| metadata_nonnegative(metadata, keys))
}
fn strip_promoted_success_fields(metadata: &mut Value) {
let Some(object) = metadata.as_object_mut() else {
return;
};
for key in [
"usage",
"route",
"apiKind",
"requestedModel",
"requested_model",
"upstreamModel",
"upstream_model",
"serviceTier",
"service_tier",
"statusCode",
"transport",
"attemptIndex",
"responseId",
"upstreamRequestId",
"openaiRequestId",
"firstTokenMs",
"first_token_ms",
] {
object.remove(key);
}
}
-1
View File
@@ -1,7 +1,6 @@
//! 进程装配、共享状态与后台任务。
pub mod config;
pub mod import_sqlite;
pub mod services;
pub mod shutdown;
pub mod state;
+55 -13
View File
@@ -26,7 +26,7 @@ use crate::{
infra::{
database::connect,
identity::AccountPseudonymizer,
logging::{init_tracing, LogError, LogGuard, RotationConfig, TracingConfig},
logging::{LogError, LogGuard, RotationConfig, TracingConfig, init_tracing},
paths::{ensure_data_dir, load_or_create_identity_secret},
redis::RedisConnection,
},
@@ -36,7 +36,7 @@ use crate::{
store::{ModelSnapshotStore, RedisModelSnapshotStore},
types::ModelConfig,
},
settings::{service::SettingsService, SettingsSnapshot},
settings::{SettingsSnapshot, service::SettingsService},
telemetry::{
account_usage::query::AccountUsageQueryService,
account_usage::store::{AccountUsageStore, PgAccountUsageStore},
@@ -45,14 +45,15 @@ use crate::{
ops::store::PgOpsErrorLogStore,
recorder::Recorder,
usage::query::UsageQueryService,
usage::store::{PgUsageRecordStore, DEFAULT_USAGE_RECORD_CAPTURE_BODY},
usage::store::{DEFAULT_USAGE_RECORD_CAPTURE_BODY, PgUsageRecordStore},
},
update::service::SystemUpdateService,
upstream::openai::{
fingerprint::{Fingerprint, PgFingerprintStore, RuntimeFingerprint},
token_client::{default_openai_token_client, OpenAiTokenClient, TokenClientConfig},
token_client::{OpenAiTokenClient, TokenClientConfig, default_openai_token_client},
transport::{
build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient,
CodexWebSocketPool, CodexWebSocketPoolConfig,
CodexBackendClient, CodexModelCatalogClient, CodexWebSocketPool,
CodexWebSocketPoolConfig, build_reqwest_client, tls::CustomCaError,
},
},
};
@@ -117,6 +118,8 @@ pub struct Services {
pub websocket_pool: Option<Arc<CodexWebSocketPool>>,
pub fingerprint: RuntimeFingerprint,
pub account_pseudonymizer: Arc<AccountPseudonymizer>,
pub system_update: Arc<SystemUpdateService>,
pub connection_drain: crate::api::middleware::connection_drain::ConnectionDrain,
pub background_tasks: BackgroundTaskStores,
pub(crate) account_pool_static: AccountPoolStaticSettings,
}
@@ -304,6 +307,8 @@ impl Services {
account_identity,
cloudflare: cloudflare_recovery,
}));
let system_update = Arc::new(SystemUpdateService::from_env());
let connection_drain = crate::api::middleware::connection_drain::ConnectionDrain::default();
Ok(Self {
database,
@@ -327,6 +332,8 @@ impl Services {
websocket_pool,
fingerprint,
account_pseudonymizer,
system_update,
connection_drain,
background_tasks: stores,
account_pool_static,
})
@@ -362,6 +369,8 @@ impl From<&Services> for crate::api::router::ApiServices {
session_affinity: services.session_affinity.clone(),
fingerprint: services.fingerprint.clone(),
process_control: Arc::new(crate::bootstrap::shutdown::RuntimeProcessControl),
system_update: services.system_update.clone(),
connection_drain: services.connection_drain.clone(),
}
}
}
@@ -444,7 +453,8 @@ pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
let port = config.server.port;
let log_guard = init_logging(&config)?;
let listener = tokio::net::TcpListener::bind((host.as_str(), port)).await?;
let (router, task_coordinator) = build_router(config).await?;
let (router, task_coordinator, connection_drain) = build_router(config).await?;
let mut task_coordinator = Some(task_coordinator);
let serve_result = {
let (graceful_shutdown_tx, graceful_shutdown_rx) = tokio::sync::oneshot::channel();
@@ -461,8 +471,28 @@ pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
result = &mut serve => result,
() = crate::bootstrap::shutdown::shutdown_signal() => {
let _ = graceful_shutdown_tx.send(());
let active_websocket_connections = connection_drain.begin_shutdown();
tracing::info!(
active_websocket_connections,
"正在排空入站连接并关闭运行时任务"
);
let coordinator = task_coordinator.take();
let connection_drain = connection_drain.clone();
let runtime_shutdown = async move {
let shutdown_tasks = async move {
if let Some(coordinator) = coordinator {
coordinator.shutdown().await;
}
};
tokio::join!(shutdown_tasks, connection_drain.wait());
};
let drain_runtime = async {
let (serve_result, ()) = tokio::join!(&mut serve, runtime_shutdown);
serve_result
};
tokio::pin!(drain_runtime);
tokio::select! {
result = &mut serve => result,
result = &mut drain_runtime => result,
() = tokio::time::sleep(HTTP_DRAIN_TIMEOUT) => {
tracing::warn!(
timeout_secs = HTTP_DRAIN_TIMEOUT.as_secs(),
@@ -479,7 +509,14 @@ pub async fn run() -> Result<(), Box<dyn Error + Send + Sync>> {
}
};
task_coordinator.shutdown().await;
if let Some(coordinator) = task_coordinator.take() {
let active_websocket_connections = connection_drain.begin_shutdown();
tracing::info!(
active_websocket_connections,
"HTTP 服务已退出,正在关闭其余运行时任务"
);
tokio::join!(coordinator.shutdown(), connection_drain.wait());
}
serve_result?;
if let Some(executable_path) = crate::bootstrap::shutdown::restart_executable_path() {
@@ -502,6 +539,7 @@ async fn build_router(
(
Router,
crate::bootstrap::tasks::coordinator::TaskCoordinator,
crate::api::middleware::connection_drain::ConnectionDrain,
),
Box<dyn Error + Send + Sync>,
> {
@@ -566,11 +604,15 @@ async fn build_router(
let task_coordinator =
crate::bootstrap::tasks::coordinator::TaskCoordinator::start(&runtime_config, &services);
let app_state = crate::api::AppState::from(&services);
let connection_drain = services.connection_drain.clone();
let router = crate::api::router::router().with_state(app_state).layer(
axum::middleware::from_fn_with_state(
connection_drain.clone(),
crate::api::middleware::connection_drain::drain_response_body,
),
);
Ok((
crate::api::router::router().with_state(app_state),
task_coordinator,
))
Ok((router, task_coordinator, connection_drain))
}
fn init_logging(config: &AppConfig) -> Result<LogGuard, LogError> {
+2 -11
View File
@@ -3,13 +3,13 @@
use std::{
path::PathBuf,
sync::{
atomic::{AtomicBool, Ordering},
OnceLock,
atomic::{AtomicBool, Ordering},
},
};
use tokio::signal;
use tokio::sync::broadcast::{self, Receiver};
use tokio::sync::broadcast;
static SHUTDOWN_REQUESTS: OnceLock<broadcast::Sender<()>> = OnceLock::new();
static RESTART_REQUESTED: AtomicBool = AtomicBool::new(false);
@@ -26,10 +26,6 @@ impl crate::api::router::ProcessControl for RuntimeProcessControl {
fn request_restart(&self, executable_path: PathBuf) {
request_process_restart(executable_path);
}
fn subscribe_shutdown(&self) -> Receiver<()> {
shutdown_subscription()
}
}
/// 等待进程关闭信号(Ctrl+C 或 SIGTERM)。
@@ -64,11 +60,6 @@ pub fn request_shutdown() {
let _ = shutdown_sender().send(());
}
/// 订阅进程关闭请求。
pub fn shutdown_subscription() -> Receiver<()> {
shutdown_sender().subscribe()
}
/// 请求进程完成优雅关闭后以新二进制替换当前进程。
pub fn request_process_restart(executable_path: PathBuf) {
let _ = RESTART_EXECUTABLE_PATH.set(executable_path);
+1 -1
View File
@@ -7,7 +7,7 @@ use tracing::{debug, info, warn};
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
};
/// 支持按时间清理过期记录的存储。
+2 -2
View File
@@ -7,8 +7,8 @@ use tokio::task::JoinHandle;
use super::{
cookie_cleanup::CookieCleanupTask,
fingerprint_update::{
start_fingerprint_update_task, CODEX_DESKTOP_APPCAST_URL,
DEFAULT_EXTRACTED_FINGERPRINT_PATH,
CODEX_DESKTOP_APPCAST_URL, DEFAULT_EXTRACTED_FINGERPRINT_PATH,
start_fingerprint_update_task,
},
model_refresh::ModelRefreshTask,
quota_refresh::QuotaRefreshTask,
@@ -5,12 +5,12 @@ use std::path::PathBuf;
use tracing::{info, warn};
use crate::upstream::openai::fingerprint::{
PgFingerprintStore, RuntimeFingerprint, UpdateChecker, APPCAST_POLL_INTERVAL,
APPCAST_POLL_INTERVAL, PgFingerprintStore, RuntimeFingerprint, UpdateChecker,
};
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
};
/// Codex Desktop 官方 appcast 地址。
+1 -1
View File
@@ -15,7 +15,7 @@ use crate::models::service::{
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup},
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup, spawn_periodic_task},
};
/// 模型刷新任务接线器。
+1 -1
View File
@@ -19,7 +19,7 @@ use crate::{
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
};
/// 主动配额刷新后台任务。
@@ -13,7 +13,7 @@ use crate::{
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, TaskFuture},
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, TaskFuture, spawn_periodic_task},
};
const RETENTION_TRIM_INTERVAL_SECS: u64 = 60 * 60;
+1 -1
View File
@@ -10,7 +10,7 @@ use crate::upstream::openai::token_client::TokenRefresher;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
periodic::{PeriodicTaskConfig, PeriodicTaskRunner, spawn_periodic_task},
};
const DEFAULT_INTERVAL_SECS: u64 = 60;
+1 -22
View File
@@ -2,7 +2,7 @@
use crate::{
infra::identity::AccountPseudonymizer,
upstream::openai::protocol::responses::{CodexCompactRequest, CodexResponsesRequest},
upstream::openai::protocol::responses::CodexResponsesRequest,
};
/// 持久化密钥驱动的账号身份隔离服务。
@@ -84,27 +84,6 @@ impl AccountIdentityService {
)
}
pub fn scope_compact(
&self,
request: &CodexCompactRequest,
account_id: &str,
proxy_request_id: &str,
) -> AccountScopedIdentity {
self.scope_values(
account_id,
proxy_request_id,
IdentityValues {
prompt_cache_key: request.prompt_cache_key(),
session_id: request.client_session_id.as_deref(),
thread_id: request.client_thread_id.as_deref(),
client_request_id: request.client_request_id.as_deref(),
turn_id: request.client_turn_id.as_deref(),
window_id: request.client_window_id.as_deref(),
parent_thread_id: request.client_parent_thread_id.as_deref(),
},
)
}
pub fn scope_auxiliary(
&self,
account_id: &str,
+5 -6
View File
@@ -10,7 +10,7 @@ use crate::{
},
upstream::openai::protocol::{
events::TokenUsage,
responses::{completed_response_metadata, CodexResponsesRequest},
responses::{CodexResponsesRequest, completed_response_metadata},
},
};
@@ -146,12 +146,11 @@ fn build_variant_identity(request: &CodexResponsesRequest) -> Option<String> {
if let Some(window_id) = non_empty_str(request.codex_window_id.as_deref()) {
parts.push(format!("window:{window_id}"));
}
if request.explicit_prompt_cache_key
|| non_empty_str(request.client_conversation_id.as_deref()).is_some()
if (request.explicit_prompt_cache_key
|| non_empty_str(request.client_conversation_id.as_deref()).is_some())
&& let Some(anchor) = derive_stable_conversation_key(request)
{
if let Some(anchor) = derive_stable_conversation_key(request) {
parts.push(format!("anchor:{anchor}"));
}
parts.push(format!("anchor:{anchor}"));
}
(!parts.is_empty()).then(|| parts.join("\0"))
+1 -1
View File
@@ -1,7 +1,7 @@
//! Shared upstream error classification for dispatch routes.
use chrono::{DateTime, Duration, Utc};
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use thiserror::Error;
use crate::{
+22 -14
View File
@@ -1,19 +1,20 @@
use std::time::Instant;
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use crate::{
dispatch::errors::{
backend_transport_name, insert_dispatch_error_metadata,
insert_upstream_diagnostics_metadata, upstream_error_http_status, DispatchErrorMetadata,
DispatchErrorMetadata, backend_transport_name, insert_dispatch_error_metadata,
insert_upstream_diagnostics_metadata, upstream_error_http_status,
},
infra::time::elapsed_millis_i64,
telemetry::{
ops::types::OpsErrorLog,
recorder::{
enrich_event_route_metadata, enrich_usage_record_identity,
DispatchErrorLogRecord, Recorder, enrich_event_route_metadata,
enrich_response_request_semantics, enrich_usage_record_identity,
event_kind as response_event_kind, reasoning_effort_from_request,
record_dispatch_error_event, DispatchErrorLogRecord, Recorder,
record_dispatch_error_event,
},
usage::types::UsageRecordLevel,
},
@@ -28,10 +29,10 @@ use crate::{
use super::{
errors::{
dispatch_error_metadata, enrich_response_dispatch_error_metadata, ResponseDispatchError,
ResponseDispatchError, dispatch_error_metadata, enrich_response_dispatch_error_metadata,
},
stream::{
live::{latest_response_id, LiveResponseStreamContext},
live::{LiveResponseStreamContext, latest_response_id},
sse_failure::{status_code_for_stream_failure, stream_failure_metadata},
trace::{ResponseDispatchAttempt, ResponseDispatchTrace},
},
@@ -177,6 +178,7 @@ pub(super) async fn record_response_upstream_error_event(
reasoning_effort_from_request(record.request),
record.request.service_tier(),
);
enrich_response_request_semantics(&mut metadata, record.request);
event.metadata = metadata;
if let Err(error) = record.recorder.record_error(event).await {
tracing::error!(account_id = %record.account_id, error = %error, "failed to record upstream error event");
@@ -224,6 +226,7 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
reasoning_effort_from_request(record.request),
record.request.service_tier(),
);
enrich_response_request_semantics(&mut metadata, record.request);
let mut event = OpsErrorLog::new(
response_event_kind(record.route),
"v1 responses stream failed",
@@ -236,13 +239,13 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
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),
);
}
if !record.rate_limit_headers.is_empty()
&& 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.recorder.record_error(event).await {
@@ -409,10 +412,14 @@ fn response_request_summary(
let body = request.body();
let input = body.get("input");
let tools = body.get("tools");
let semantics = request.semantics();
json!({
"model": request.model(),
"stream": request.stream(),
"store": request.store(),
"compact": semantics.compact,
"requestKind": semantics.request_kind,
"subagentKind": semantics.subagent_kind,
"transport": backend_transport_name(transport),
"inputType": json_value_kind(input),
"inputItemsCount": input.and_then(Value::as_array).map(Vec::len),
@@ -514,6 +521,7 @@ pub(super) async fn record_live_response_stream_event(
reasoning_effort_from_request(&context.request),
context.request.service_tier(),
);
enrich_response_request_semantics(&mut metadata, &context.request);
if is_success_usage_event(level, status_code) {
crate::telemetry::recorder::record_response_event(
@@ -12,9 +12,9 @@ use crate::{
recovery::{
auth::{auth_failure_account_status, is_auth_upstream_error},
cloudflare::{
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
CloudflareRecovery,
CloudflareRecovery, cloudflare_challenge_error_message,
cloudflare_path_block_error_message, is_cloudflare_challenge_upstream_error,
is_cloudflare_path_block_upstream_error,
},
exhaustion::AccountExhaustionTracker,
},
@@ -26,7 +26,7 @@ use crate::{
fleet::{account::AccountStatus, pool::AccountPoolService},
upstream::openai::{
protocol::responses::ResponsesSseFailure,
transport::{is_banned_upstream_error, CodexClientError},
transport::{CodexClientError, is_banned_upstream_error},
},
};
+1 -1
View File
@@ -2,7 +2,7 @@
use crate::{
fleet::account::AccountStatus,
upstream::openai::transport::{is_banned_auth_signal, CodexClientError},
upstream::openai::transport::{CodexClientError, is_banned_auth_signal},
};
pub(crate) fn is_auth_upstream_error(error: &CodexClientError) -> bool {
+1 -1
View File
@@ -5,9 +5,9 @@ use serde_json::Value;
use crate::{
dispatch::affinity::{
SessionAffinityEntry, SessionAffinityService,
service::{MAX_REPLAY_DEPTH, MAX_REPLAY_SESSION_BYTES, MAX_REPLAY_SNAPSHOT_BYTES},
types::ResponseReplaySnapshot,
SessionAffinityEntry, SessionAffinityService,
},
upstream::openai::protocol::responses::{CodexResponsesRequest, StreamCommitPolicy},
};
+22 -17
View File
@@ -1,7 +1,7 @@
//! Responses 创建编排与调度服务。
//!
//! 包含了将 OpenAI 请求调度到 Codex 上游账号的完整逻辑,包括:
//! - 响应创建(非流式 / 流式 / compact)
//! - 响应创建(非流式 / 流式)
//! - 会话亲和性与隐式续接
//! - reasoning replay
//! - 账号回退与错误恢复
@@ -12,7 +12,7 @@ use std::{pin::Pin, sync::Arc, time::Instant};
use bytes::Bytes;
use chrono::Utc;
use futures::stream::Stream;
use serde_json::{json, Value};
use serde_json::{Value, json};
use crate::{
dispatch::{
@@ -31,9 +31,9 @@ use crate::{
history::HistoryRecoveryPlan,
},
upstream_call::{
create_response_with_account_retrying_5xx, verify_acquired_quota_if_required,
AccountUpstreamContext, QuotaVerificationContext, QuotaVerificationDecision,
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
AccountUpstreamContext, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, QuotaVerificationContext,
QuotaVerificationDecision, create_response_with_account_retrying_5xx,
verify_acquired_quota_if_required,
},
},
fleet::{
@@ -42,16 +42,19 @@ use crate::{
},
models::service::ModelService,
telemetry::{
recorder::{reasoning_effort_from_request, record_response_event, Recorder},
recorder::{
Recorder, enrich_response_request_semantics, reasoning_effort_from_request,
record_response_event,
},
usage::types::ResponseUsageRecord,
},
upstream::openai::{
protocol::{
events::TokenUsage,
responses::{response_from_codex_sse, CodexResponsesRequest, CollectedResponse},
responses::{CodexResponsesRequest, CollectedResponse, response_from_codex_sse},
},
transport::{
backend_transport_for_response_request, CodexBackendClient, CodexBackendResponse,
CodexBackendClient, CodexBackendResponse, backend_transport_for_response_request,
},
},
};
@@ -59,11 +62,11 @@ use crate::{
use super::{
errors::{ResponseDispatchError, ResponseDispatchStreamError},
recording::{
insert_response_status_metadata, insert_response_trace_metadata,
insert_response_upstream_diagnostics, insert_websocket_pool_decision,
record_response_dispatch_error_event, record_response_upstream_error_event,
ResponseDispatchErrorDetails, ResponseDispatchErrorEventRecord,
ResponseUpstreamErrorEventRecord,
ResponseUpstreamErrorEventRecord, insert_response_status_metadata,
insert_response_trace_metadata, insert_response_upstream_diagnostics,
insert_websocket_pool_decision, record_response_dispatch_error_event,
record_response_upstream_error_event,
},
stream::{
sse_failure::is_history_recovery_sse_failure,
@@ -157,6 +160,7 @@ impl ResponseDispatchService {
let catalog = self.models.catalog().await;
let display_model = catalog.resolve_model_id(requested_model);
request.set_model(display_model.clone());
let compact = request.semantics().compact;
let tuple_schema = request.tuple_schema.clone();
let image_generation_requested = request.expects_image_generation();
let now = Utc::now();
@@ -216,7 +220,7 @@ impl ResponseDispatchService {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: exhausted_accounts.last_account_id(),
stream: false,
compact: false,
compact,
transport: Some(backend_transport_name(
backend_transport_for_response_request(&request),
)),
@@ -309,7 +313,7 @@ impl ResponseDispatchService {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: Some(&release_account_id),
stream: false,
compact: false,
compact,
transport: Some(backend_transport_name(response.transport)),
},
&error,
@@ -470,6 +474,7 @@ impl ResponseDispatchService {
insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics);
insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt));
insert_websocket_pool_decision(&mut metadata, response.websocket_pool_decision);
enrich_response_request_semantics(&mut metadata, &request);
record_response_event(ResponseUsageRecord {
recorder: &self.recorder,
request_id,
@@ -511,7 +516,7 @@ impl ResponseDispatchService {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: Some(&account.id),
stream: false,
compact: false,
compact,
transport: Some(backend_transport_name(response.transport)),
},
&error,
@@ -530,7 +535,7 @@ impl ResponseDispatchService {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: Some(&account.id),
stream: false,
compact: false,
compact,
transport: Some(backend_transport_name(response.transport)),
},
&error,
@@ -549,7 +554,7 @@ impl ResponseDispatchService {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: Some(&account.id),
stream: false,
compact: false,
compact,
transport: Some(backend_transport_name(response.transport)),
},
&error,
+12 -11
View File
@@ -7,14 +7,14 @@ use crate::{
affinity::prepare_variant_identity,
attempts::AccountAttemptLedger,
errors::{
backend_transport_name, is_continuation_busy_error, is_history_recovery_upstream_error,
is_retryable_account_transport_error, upstream_error_body,
upstream_error_set_cookie_headers, ResponseDispatchError,
ResponseDispatchError, backend_transport_name, is_continuation_busy_error,
is_history_recovery_upstream_error, is_retryable_account_transport_error,
upstream_error_body, upstream_error_set_cookie_headers,
},
recording::{
record_prefetched_response_stream_failure_event, record_response_upstream_error_event,
ResponseDispatchErrorDetails, ResponseStreamFailureEventRecord,
ResponseUpstreamErrorEventRecord,
ResponseUpstreamErrorEventRecord, record_prefetched_response_stream_failure_event,
record_response_upstream_error_event,
},
recovery::{
account_failure::{isolate_rotatable_account_failure, isolate_sse_account_failure},
@@ -23,7 +23,7 @@ use crate::{
},
service::{ResponseDispatchService, ResponseDispatchStream},
stream::{
live::{spawn_live_response_stream, LiveResponseStreamContext},
live::{LiveResponseStreamContext, spawn_live_response_stream},
prefetch::prefetch_first_sse_chunk,
sse_failure::{
first_sse_failure, is_history_recovery_sse_failure, sse_failure_error_body,
@@ -31,9 +31,9 @@ use crate::{
trace::ResponseDispatchTrace,
},
upstream_call::{
create_response_stream_with_account_retrying_5xx, verify_acquired_quota_if_required,
AccountUpstreamContext, QuotaVerificationContext, QuotaVerificationDecision,
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
AccountUpstreamContext, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, QuotaVerificationContext,
QuotaVerificationDecision, create_response_stream_with_account_retrying_5xx,
verify_acquired_quota_if_required,
},
},
fleet::pool::AccountAcquireRequest,
@@ -57,6 +57,7 @@ impl ResponseDispatchService {
let display_model = catalog.resolve_model_id(requested_model);
request.set_model(display_model.clone());
request.set_stream(true);
let compact = request.semantics().compact;
let tuple_schema = request.tuple_schema.clone();
let now = Utc::now();
prepare_variant_identity(&mut request);
@@ -86,7 +87,7 @@ impl ResponseDispatchService {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: exhausted_accounts.last_account_id(),
stream: true,
compact: false,
compact,
transport: Some(backend_transport_name(
backend_transport_for_response_request(&request),
)),
@@ -107,7 +108,7 @@ impl ResponseDispatchService {
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: $account_id,
stream: true,
compact: false,
compact,
transport: $transport,
},
&error,
+8 -9
View File
@@ -6,14 +6,14 @@ use std::{
};
use bytes::Bytes;
use futures::{stream::Stream, StreamExt};
use futures::{StreamExt, stream::Stream};
use serde_json::Value;
use tokio::sync::{mpsc, oneshot};
use crate::{
dispatch::{
affinity::{resolve::record_response_affinity, SessionAffinityService},
errors::{backend_transport_name, ResponseDispatchStreamError},
affinity::{SessionAffinityService, resolve::record_response_affinity},
errors::{ResponseDispatchStreamError, backend_transport_name},
recording::{
insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state,
record_live_response_stream_event,
@@ -28,13 +28,13 @@ use crate::{
protocol::{
events::extract_sse_usage,
responses::{
CodexResponsesRequest, CollectedResponse,
reconvert_responses_sse_event_tuple_values, response_from_codex_sse,
response_sse_event_is_terminal, update_first_response_output_ms,
CodexResponsesRequest, CollectedResponse,
},
sse::{
encode_sse_event, parse_sse_events, response_failed_sse_event_with_id,
sse_body_has_done, sse_frame_end, DONE_SSE_FRAME,
DONE_SSE_FRAME, encode_sse_event, parse_sse_events,
response_failed_sse_event_with_id, sse_body_has_done, sse_frame_end,
},
},
transport::{
@@ -47,9 +47,8 @@ use crate::{
use super::{
sse_failure::{
status_code_for_stream_failure, stream_failure_metadata, stream_failure_source,
synthetic_stream_disconnected_detail, STREAM_DISCONNECTED_CODE,
STREAM_DISCONNECTED_MESSAGE,
STREAM_DISCONNECTED_CODE, STREAM_DISCONNECTED_MESSAGE, status_code_for_stream_failure,
stream_failure_metadata, stream_failure_source, synthetic_stream_disconnected_detail,
},
trace::ResponseDispatchAttempt,
};
+1 -1
View File
@@ -2,7 +2,7 @@ use bytes::Bytes;
use futures::StreamExt;
use crate::upstream::openai::{
protocol::sse::{sse_frame_end, SseError},
protocol::sse::{SseError, sse_frame_end},
transport::CodexBackendSseStream,
};
+1 -1
View File
@@ -6,7 +6,7 @@ use crate::{
upstream::openai::{
protocol::{
events::TokenUsage,
responses::{response_from_codex_sse, CollectedResponse, ResponsesSseFailure},
responses::{CollectedResponse, ResponsesSseFailure, response_from_codex_sse},
sse::SseError,
},
transport::is_banned_auth_signal,
+8 -451
View File
@@ -2,57 +2,26 @@
use std::{future::Future, time::Instant};
use chrono::Utc;
use serde_json::{json, Value};
use serde_json::Value;
use crate::{
dispatch::{
affinity::{AccountIdentityService, AccountScopedIdentity},
attempts::AccountAttemptLedger,
errors::{
is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error,
is_rate_limit_upstream_error, is_retryable_account_transport_error,
is_retryable_upstream_5xx_error, rate_limit_cooldown_until, upstream_error_body,
upstream_error_http_status, upstream_error_set_cookie_headers, ResponseDispatchError,
},
recording::{
insert_response_status_metadata, insert_response_trace_metadata,
insert_response_upstream_diagnostics, record_response_dispatch_error_event,
ResponseDispatchErrorEventRecord,
},
recovery::{
auth::{auth_failure_account_status, is_auth_upstream_error},
cloudflare::{
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
CloudflareRecovery,
},
exhaustion::AccountExhaustionTracker,
},
service::{ResponseDispatchResponse, ResponseDispatchService},
stream::trace::ResponseDispatchTrace,
affinity::AccountIdentityService, errors::is_retryable_upstream_5xx_error,
recovery::cloudflare::CloudflareRecovery,
},
fleet::{
account::{Account, AccountStatus},
pool::{AccountAcquireRequest, AccountLease, AccountPoolService},
account::Account,
pool::{AccountLease, AccountPoolService},
quota::{
quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds,
quota_snapshot_reset_at,
},
},
telemetry::{
recorder::{reasoning_effort_from_compact_request, record_response_event},
usage::types::ResponseUsageRecord,
},
upstream::openai::{
protocol::{
events::extract_usage,
responses::{CodexCompactRequest, CodexResponsesRequest},
},
protocol::responses::CodexResponsesRequest,
transport::{
is_banned_upstream_error, CodexBackendClient, CodexBackendResponse,
CodexBackendStreamingResponse, CodexClientError, CodexCompactResponse,
CodexRequestContext,
CodexBackendClient, CodexBackendResponse, CodexBackendStreamingResponse,
CodexClientError, CodexRequestContext,
},
},
};
@@ -254,133 +223,6 @@ async fn create_response_stream_with_account(
.await
}
async fn create_compact_response_with_account(
context: AccountUpstreamContext<'_>,
request: &CodexCompactRequest,
) -> Result<CodexCompactResponse, CodexClientError> {
let cookie_header = context
.cloudflare
.cookie_header_for_request(&context.account.id, "/codex/responses/compact")
.await;
let identity =
context
.account_identity
.scope_compact(request, &context.account.id, context.request_id);
let upstream_request = compact_upstream_request(request, &identity);
context
.codex
.create_compact_response(
&upstream_request,
CodexRequestContext {
access_token: &context.account.access_token,
account_id: context.account.account_id.as_deref(),
request_id: context.request_id,
turn_state: None,
turn_metadata: None,
beta_features: None,
include_timing_metrics: None,
version: None,
codex_window_id: identity.window_id.as_deref(),
parent_thread_id: identity.parent_thread_id.as_deref(),
cookie_header: cookie_header.as_deref(),
installation_id: Some(&identity.installation_id),
session_id: identity.session_id.as_deref(),
thread_id: identity.thread_id.as_deref(),
prompt_cache_key: identity.prompt_cache_key.as_deref(),
client_request_id: Some(&identity.client_request_id),
turn_id: identity.turn_id.as_deref(),
},
)
.await
}
fn compact_upstream_request(
request: &CodexCompactRequest,
identity: &AccountScopedIdentity,
) -> CodexCompactRequest {
let mut upstream = request.clone();
for (key, value) in [
("prompt_cache_key", identity.prompt_cache_key.as_deref()),
("session_id", identity.session_id.as_deref()),
("thread_id", identity.thread_id.as_deref()),
("turn_id", identity.turn_id.as_deref()),
(
"x-client-request-id",
Some(identity.client_request_id.as_str()),
),
("x-codex-window-id", identity.window_id.as_deref()),
(
"x-codex-parent-thread-id",
identity.parent_thread_id.as_deref(),
),
(
"x-codex-installation-id",
Some(identity.installation_id.as_str()),
),
] {
replace_existing_identity_field(&mut upstream.body, key, value);
}
let metadata = match upstream.body.get("client_metadata") {
Some(Value::Object(metadata)) => Some(metadata.clone()),
None => Some(serde_json::Map::new()),
Some(_) => None,
};
if let Some(mut metadata) = metadata {
for (key, value) in [
(
"x-codex-installation-id",
Some(identity.installation_id.as_str()),
),
("session_id", identity.session_id.as_deref()),
("thread_id", identity.thread_id.as_deref()),
(
"x-client-request-id",
Some(identity.client_request_id.as_str()),
),
("turn_id", identity.turn_id.as_deref()),
("x-codex-window-id", identity.window_id.as_deref()),
(
"x-codex-parent-thread-id",
identity.parent_thread_id.as_deref(),
),
] {
match value.filter(|value| !value.trim().is_empty()) {
Some(value) => {
metadata.insert(key.to_string(), Value::String(value.to_string()));
}
None => {
metadata.remove(key);
}
}
}
if !metadata.is_empty() {
upstream
.body
.insert("client_metadata".to_string(), Value::Object(metadata));
}
}
upstream
}
fn replace_existing_identity_field(
body: &mut serde_json::Map<String, Value>,
key: &str,
value: Option<&str>,
) {
if !body.get(key).is_some_and(Value::is_string) {
return;
}
match value.filter(|value| !value.trim().is_empty()) {
Some(value) => {
body.insert(key.to_string(), Value::String(value.to_string()));
}
None => {
body.remove(key);
}
}
}
const MAX_UPSTREAM_5XX_RETRIES_PER_ACCOUNT: usize = 2;
async fn retry_upstream_5xx<T, F, Fut>(
@@ -474,288 +316,3 @@ pub(crate) async fn create_response_stream_with_account_retrying_5xx(
)
.await
}
pub(crate) async fn create_compact_response_with_account_retrying_5xx(
context: AccountUpstreamContext<'_>,
request: &CodexCompactRequest,
) -> Result<CodexCompactResponse, CodexClientError> {
retry_upstream_5xx(
context.request_id,
&context.account.id,
Some("compact"),
"upstream compact request failed with retryable 5xx",
|| create_compact_response_with_account(context, request),
)
.await
}
impl ResponseDispatchService {
/// 调度 Responses compact 请求到 Codex compact 上游。
pub async fn compact(
&self,
request_id: &str,
mut request: CodexCompactRequest,
requested_model: &str,
) -> Result<ResponseDispatchResponse, ResponseDispatchError> {
let started_at = Instant::now();
let client_api_key_id = request.client_api_key_id.clone();
let catalog = self.models.catalog().await;
let display_model = catalog.resolve_model_id(requested_model);
request.set_model(display_model.clone());
let mut exhausted_accounts = AccountExhaustionTracker::default();
let mut trace = ResponseDispatchTrace::default();
let acquire_request = AccountAcquireRequest::new(request.model(), Utc::now());
let mut candidates =
AccountAttemptLedger::freeze(&self.account_pool, &acquire_request).await;
loop {
let acquired = match candidates.acquire_next(&self.account_pool).await {
Some(acquired) => acquired,
None => {
tracing::info!(
candidate_count = candidates.candidate_count(),
attempted = candidates.attempted_count(),
state_excluded = candidates.state_excluded_count(),
"Responses compact account candidate ledger exhausted"
);
let error = exhausted_accounts
.last_exhausted()
.map(ResponseDispatchError::from_exhausted_account)
.unwrap_or(ResponseDispatchError::NoActiveAccount);
self.record_compact_dispatch_error(
request_id,
client_api_key_id.as_deref(),
requested_model,
started_at,
exhausted_accounts.last_account_id(),
&error,
)
.await;
return Err(error);
}
};
let acquired_account_id = acquired.account.id.clone();
let acquired = match verify_acquired_quota_if_required(
QuotaVerificationContext {
account_pool: self.account_pool.as_ref(),
codex: self.codex.as_ref(),
cloudflare: &self.cloudflare,
account_identity: &self.account_identity,
request_id,
},
acquired,
)
.await
{
QuotaVerificationDecision::Ready(acquired) => *acquired,
QuotaVerificationDecision::RetryWithAnotherAccount => {
exhausted_accounts.record_rate_limited(
Some(&acquired_account_id),
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
);
continue;
}
};
let account = acquired.account.clone();
let release_account_id = account.id.clone();
let attempt = trace.start_attempt(&release_account_id);
let response_result = create_compact_response_with_account_retrying_5xx(
AccountUpstreamContext {
codex: &self.codex,
account_identity: &self.account_identity,
cloudflare: &self.cloudflare,
request_id,
account: &account,
},
&request,
)
.await;
acquired.complete().await;
if let Err(error) = &response_result {
self.cloudflare
.capture_set_cookie_headers(
&release_account_id,
upstream_error_set_cookie_headers(error),
)
.await;
}
match response_result {
Ok(response) => {
self.cloudflare
.capture_set_cookie_headers(
&release_account_id,
&response.set_cookie_headers,
)
.await;
self.cloudflare.reset_account_recovery(&account.id).await;
self.account_pool
.sync_passive_rate_limit_headers(&account, &response.rate_limit_headers)
.await;
let usage = extract_usage(&response.body);
if let Some(usage) = usage {
self.account_pool
.record_token_usage(&account.id, &usage)
.await;
}
let mut metadata = json!({
"stream": false,
"compact": true,
"usage": usage,
"effectiveModel": response.response_metadata.effective_model.as_deref(),
"modelsEtag": response.response_metadata.models_etag.as_deref(),
"reasoningIncluded": response.response_metadata.reasoning_included,
});
insert_response_status_metadata(
&mut metadata,
200,
200,
response.diagnostics.status_code.map(i64::from),
);
insert_response_upstream_diagnostics(&mut metadata, &response.diagnostics);
insert_response_trace_metadata(&mut metadata, &trace, Some(&attempt));
record_response_event(ResponseUsageRecord {
recorder: &self.recorder,
request_id,
client_api_key_id: request.client_api_key_id.as_deref(),
account_id: &account.id,
route: "/v1/responses/compact",
model: response
.response_metadata
.effective_model
.as_deref()
.unwrap_or(&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_compact_request(&request),
service_tier: None,
started_at,
status_code: 200,
message: "v1 responses compact completed",
metadata,
rate_limit_headers: &response.rate_limit_headers,
})
.await;
self.models
.observe_models_etag(response.response_metadata.models_etag.as_deref());
return Ok(ResponseDispatchResponse {
body: response.body,
response_headers: response.response_metadata.client_headers,
});
}
Err(error) if is_rate_limit_upstream_error(&error) => {
exhausted_accounts.record_rate_limited(
Some(&release_account_id),
upstream_error_body(&error),
);
let cooldown_until = rate_limit_cooldown_until(&error, Utc::now());
self.account_pool
.mark_quota_limited_until(&release_account_id, cooldown_until)
.await;
}
Err(error) if is_quota_exhausted_upstream_error(&error) => {
exhausted_accounts.record_quota_exhausted(
Some(&release_account_id),
upstream_error_body(&error),
);
self.account_pool
.set_status(&release_account_id, AccountStatus::QuotaExhausted)
.await;
}
Err(error) if is_auth_upstream_error(&error) => {
let upstream_error = upstream_error_body(&error);
let account_status = auth_failure_account_status(&error);
exhausted_accounts.record_auth_failure(
Some(&release_account_id),
account_status,
upstream_error,
Some(upstream_error_http_status(&error)),
);
self.account_pool
.set_status(&release_account_id, account_status)
.await;
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
exhausted_accounts.record_cloudflare_challenge(
Some(&release_account_id),
cloudflare_challenge_error_message(),
);
self.cloudflare
.apply_challenge(self.account_pool.as_ref(), &release_account_id)
.await;
}
Err(error) if is_cloudflare_path_block_upstream_error(&error) => {
exhausted_accounts.record_cloudflare_path_blocked(
Some(&release_account_id),
cloudflare_path_block_error_message(),
);
self.cloudflare
.apply_path_block(self.account_pool.as_ref(), &release_account_id)
.await;
}
Err(error) if is_model_unsupported_upstream_error(&error) => {
let upstream_error = upstream_error_body(&error);
exhausted_accounts
.record_model_unsupported(Some(&release_account_id), upstream_error);
}
Err(error) if is_banned_upstream_error(&error) => {
exhausted_accounts.record_auth_failure(
Some(&release_account_id),
AccountStatus::Banned,
upstream_error_body(&error),
Some(upstream_error_http_status(&error)),
);
self.account_pool
.set_status(&release_account_id, AccountStatus::Banned)
.await;
}
Err(error) => {
if is_retryable_account_transport_error(&error) {
exhausted_accounts.record_upstream_unavailable(
Some(&release_account_id),
upstream_error_body(&error),
);
continue;
}
let error = ResponseDispatchError::Upstream(error);
self.record_compact_dispatch_error(
request_id,
client_api_key_id.as_deref(),
requested_model,
started_at,
Some(&release_account_id),
&error,
)
.await;
return Err(error);
}
}
}
}
async fn record_compact_dispatch_error(
&self,
request_id: &str,
client_api_key_id: Option<&str>,
requested_model: &str,
started_at: Instant,
account_id: Option<&str>,
error: &ResponseDispatchError,
) {
record_response_dispatch_error_event(ResponseDispatchErrorEventRecord {
recorder: &self.recorder,
request_id,
client_api_key_id,
account_id,
route: "/v1/responses/compact",
model: requested_model,
started_at,
stream: false,
compact: true,
transport: Some("http"),
error,
})
.await;
}
}
+30
View File
@@ -3,6 +3,24 @@
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::infra::json::SortDirection;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AccountSortField {
Email,
Status,
PlanType,
Usage,
LastUsedAt,
ExpiresAt,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountListSort {
pub field: AccountSortField,
pub direction: SortDirection,
}
/// 账号当前状态。
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum AccountStatus {
@@ -29,6 +47,18 @@ impl AccountStatus {
Self::Banned => "banned",
}
}
/// 解析持久化/API 使用的账号状态字符串。
pub fn parse(value: &str) -> Option<Self> {
match value.trim().to_ascii_lowercase().as_str() {
"active" => Some(Self::Active),
"expired" => Some(Self::Expired),
"quota_exhausted" => Some(Self::QuotaExhausted),
"disabled" => Some(Self::Disabled),
"banned" => Some(Self::Banned),
_ => None,
}
}
}
impl std::fmt::Display for AccountStatus {
+1 -1
View File
@@ -4,7 +4,7 @@ use chrono::{DateTime, Utc};
use serde_json::{Map, Value};
use crate::fleet::account::AccountStatus;
use crate::fleet::refresh::{jwt_expiry, JwtExpiry};
use crate::fleet::refresh::{JwtExpiry, jwt_expiry};
/// 账号导入条目。
#[derive(Debug, Clone)]
+1 -1
View File
@@ -1,8 +1,8 @@
use secrecy::ExposeSecret;
use super::{
types::{ExportedAccount, ExportedAccounts},
AccountManageError, AccountManageService,
types::{ExportedAccount, ExportedAccounts},
};
use crate::{fleet::store::StoredAccount, infra::time::china_rfc3339_str};
+9 -10
View File
@@ -2,12 +2,12 @@ use chrono::Utc;
use secrecy::SecretString;
use super::{
types::{
import_quota_plan_type, import_status_from_usage_error, import_usage_plan_type,
import_usage_string, AccountManageError, ImportSupplementalAccountInfo,
ImportSupplementalNeeds, ImportedAccountState, ImportedAccounts, ResolvedImportTokens,
},
AccountManageService,
types::{
AccountManageError, ImportSupplementalAccountInfo, ImportSupplementalNeeds,
ImportedAccountState, ImportedAccounts, ResolvedImportTokens, import_quota_plan_type,
import_status_from_usage_error, import_usage_plan_type, import_usage_string,
},
};
impl AccountManageService {
@@ -304,15 +304,14 @@ impl AccountManageService {
let chatgpt_account_id = crate::fleet::import::normalize_nonempty_str(account_id);
let chatgpt_user_id = crate::fleet::import::normalize_nonempty_str(user_id);
if let Some(chatgpt_account_id) = chatgpt_account_id {
if let Some(existing) = self
if let Some(chatgpt_account_id) = chatgpt_account_id
&& let Some(existing) = self
.store
.find_by_chatgpt_identity(chatgpt_account_id, chatgpt_user_id)
.await
.map_err(|_| AccountManageError::Inspect)?
{
return Ok(existing.id);
}
{
return Ok(existing.id);
}
Ok(provided_id.unwrap_or_else(|| crate::fleet::import::normalized_account_id(None)))
+14 -9
View File
@@ -1,27 +1,30 @@
use std::{
collections::{hash_map::DefaultHasher, HashSet},
collections::{HashSet, hash_map::DefaultHasher},
hash::{Hash, Hasher},
};
use chrono::{DateTime, Duration, Utc};
use futures::{stream, StreamExt};
use futures::{StreamExt, stream};
use secrecy::{ExposeSecret, SecretString};
use tokio::time::sleep;
use uuid::Uuid;
use crate::{
fleet::{account::AccountStatus, store::AccountStore},
fleet::{
account::{AccountListSort, AccountStatus},
store::AccountStore,
},
infra::{json::NumberedPage, time::elapsed_millis_i64},
upstream::openai::token_client::RefreshFailure,
};
use super::{
types::{
parse_account_status, stored_to_admin_metadata, AccountHealthCheck, AccountManageError,
AccountRefreshOutcome, AccountRefreshResult, AccountUpdate, BatchDeleteAccounts,
ManagedAccount, ManualCreateTokens,
},
AccountManageService,
types::{
AccountHealthCheck, AccountManageError, AccountRefreshOutcome, AccountRefreshResult,
AccountUpdate, BatchDeleteAccounts, ManagedAccount, ManualCreateTokens,
parse_account_status, stored_to_admin_metadata,
},
};
const ADMIN_REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60;
@@ -32,10 +35,12 @@ impl AccountManageService {
page: u32,
page_size: u32,
search: Option<String>,
status: Option<AccountStatus>,
sort: Option<AccountListSort>,
) -> Result<NumberedPage<ManagedAccount>, AccountManageError> {
let page = self
.store
.list_metadata_page(page, page_size, search.as_deref())
.list_metadata_page(page, page_size, search.as_deref(), status, sort)
.await
.map_err(|_| AccountManageError::List)?;
Ok(NumberedPage {
+2 -2
View File
@@ -1,6 +1,6 @@
use std::{collections::HashMap, fmt::Write as _, sync::Arc};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, Duration, Utc};
use rand::Rng;
use reqwest::{Client, StatusCode, Url};
@@ -9,8 +9,8 @@ use sha2::{Digest, Sha256};
use tokio::sync::Mutex;
use super::{
types::{AccountManageError, OAuthAuthorizeResult, OAuthExchangeInput},
AccountManageService,
types::{AccountManageError, OAuthAuthorizeResult, OAuthExchangeInput},
};
const OPENAI_OAUTH_AUTHORIZE_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize";
+47 -30
View File
@@ -3,7 +3,7 @@ use std::{convert::Infallible, pin::Pin};
use bytes::Bytes;
use futures::{Stream, StreamExt};
use secrecy::ExposeSecret;
use serde_json::{json, Value};
use serde_json::{Value, json};
use tokio::sync::mpsc;
use crate::{
@@ -12,15 +12,15 @@ use crate::{
upstream::openai::{
protocol::{
responses::{CodexResponsesRequest, ResponsesSseFailure},
sse::{encode_sse_event, parse_sse_events, SseEvent},
sse::{SseEvent, encode_sse_event, parse_sse_events},
},
transport::{
is_banned_auth_signal, is_banned_upstream_error, CodexClientError, CodexRequestContext,
CodexClientError, CodexRequestContext, is_banned_auth_signal, is_banned_upstream_error,
},
},
};
use super::{types::AccountManageError, AccountManageService};
use super::{AccountManageService, types::AccountManageError};
pub(super) type AccountTestStream = Pin<Box<dyn Stream<Item = Result<Bytes, Infallible>> + Send>>;
@@ -50,6 +50,16 @@ impl AccountTestOutcome {
status,
}
}
fn attach_account_status(&mut self, status: AccountStatus) {
let Some(event) = self.event.as_object_mut() else {
return;
};
event.insert(
"accountStatus".to_string(),
Value::String(status.as_str().to_string()),
);
}
}
impl AccountManageService {
@@ -166,7 +176,7 @@ impl AccountManageService {
turn_id: None,
};
let outcome = match codex.create_response_stream(&request, context).await {
let mut outcome = match codex.create_response_stream(&request, context).await {
Ok(response) => process_upstream_test_stream(response.body, &tx).await,
Err(error) => AccountTestOutcome::error(
error.to_string(),
@@ -174,10 +184,12 @@ impl AccountManageService {
),
};
if let Some(status) = outcome.status {
service
if let Some(status) = outcome.status
&& let Some(status) = service
.apply_connection_test_status(&stored_account_id, status)
.await;
.await
{
outcome.attach_account_status(status);
}
send_test_event(&tx, outcome.event).await;
});
@@ -190,26 +202,30 @@ impl AccountManageService {
Ok(Box::pin(stream))
}
async fn apply_connection_test_status(&self, account_id: &str, status: AccountStatus) {
async fn apply_connection_test_status(
&self,
account_id: &str,
status: AccountStatus,
) -> Option<AccountStatus> {
let current = match self.store.get(account_id).await {
Ok(Some(account)) => account,
Ok(None) => return,
Ok(None) => return None,
Err(error) => {
tracing::warn!(
account_id,
error = %error,
"failed to inspect account after connection test"
);
return;
return None;
}
};
if current.status == AccountStatus::Disabled {
return;
return Some(AccountStatus::Disabled);
}
match self.store.set_status(account_id, status).await {
Ok(true) => {}
Ok(false) => return,
Ok(false) => return None,
Err(error) => {
tracing::warn!(
account_id,
@@ -217,21 +233,22 @@ impl AccountManageService {
error = %error,
"failed to persist account status after connection test"
);
return;
return None;
}
}
if matches!(status, AccountStatus::Expired | AccountStatus::Banned) {
if let Err(error) = self.store.set_next_refresh_at(account_id, None).await {
tracing::warn!(
account_id,
error = %error,
"failed to clear token refresh schedule after connection test"
);
}
if matches!(status, AccountStatus::Expired | AccountStatus::Banned)
&& let Err(error) = self.store.set_next_refresh_at(account_id, None).await
{
tracing::warn!(
account_id,
error = %error,
"failed to clear token refresh schedule after connection test"
);
}
self.sync_account_pool_best_effort(account_id, "connection test")
.await;
Some(status)
}
}
@@ -297,10 +314,10 @@ async fn process_upstream_test_stream(
}
}
if !buffer.trim().is_empty() {
if let Some(outcome) = process_sse_frame(tx, &buffer).await {
return outcome;
}
if !buffer.trim().is_empty()
&& let Some(outcome) = process_sse_frame(tx, &buffer).await
{
return outcome;
}
AccountTestOutcome::error("Stream ended before response.completed", None)
@@ -332,10 +349,10 @@ async fn process_sse_event(
};
match value.get("type").and_then(Value::as_str) {
Some("response.output_text.delta") => {
if let Some(delta) = value.get("delta").and_then(Value::as_str) {
if !delta.is_empty() {
send_test_event(tx, json!({ "type": "content", "text": delta })).await;
}
if let Some(delta) = value.get("delta").and_then(Value::as_str)
&& !delta.is_empty()
{
send_test_event(tx, json!({ "type": "content", "text": delta })).await;
}
None
}
+1 -1
View File
@@ -1,6 +1,6 @@
use secrecy::ExposeSecret;
use super::{types::AccountManageError, AccountManageService};
use super::{AccountManageService, types::AccountManageError};
impl AccountManageService {
async fn usage_cookie_header(&self, account_id: &str) -> Option<String> {
-6
View File
@@ -107,12 +107,6 @@ impl Default for AccountQuotaData {
}
impl AccountQuotaData {
pub(crate) fn has_high_usage(&self) -> bool {
self.windows
.iter()
.any(|window| window.used_percent.is_some_and(|percent| percent >= 80.0))
}
pub(crate) fn usage_windows(&self) -> Vec<AccountQuotaUsageWindow> {
self.windows
.iter()
+19 -21
View File
@@ -13,10 +13,10 @@ use indexmap::IndexMap;
use thiserror::Error;
use crate::fleet::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at};
use crate::fleet::refresh::{jwt_expiry, JwtExpiry};
use crate::fleet::refresh::{JwtExpiry, jwt_expiry};
use crate::fleet::scheduler::{
candidates::{self, CandidateFilter, CandidateRequest},
AccountScheduler, ScoreWeights,
candidates::{self, CandidateFilter, CandidateRequest},
};
use crate::fleet::store::AccountStore;
use crate::fleet::{
@@ -27,7 +27,7 @@ use crate::telemetry::account_usage::store::{
AccountUsageDelta, AccountUsageSnapshot, AccountUsageStore, AccountUsageWindow,
};
use crate::upstream::openai::protocol::events::{
parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage,
TokenUsage as CodexTokenUsage, parse_rate_limit_headers, rate_limit_quota,
};
pub use crate::fleet::scheduler::RotationStrategy;
@@ -457,21 +457,20 @@ impl AccountPoolService {
"failed to persist refreshed runtime account state"
);
}
if refreshed.sync_usage_window {
if let Err(error) = self
if refreshed.sync_usage_window
&& let Err(error) = self
.usage_store
.sync_runtime_window(
&refreshed.account.id,
account_usage_window(&refreshed.account),
)
.await
{
tracing::error!(
account_id = %refreshed.account.id,
error = %error,
"failed to persist refreshed runtime account usage window"
);
}
{
tracing::error!(
account_id = %refreshed.account.id,
error = %error,
"failed to persist refreshed runtime account usage window"
);
}
}
}
@@ -641,18 +640,17 @@ impl AccountPoolService {
}
};
if let Some(reset_at) = reset_at {
if let Err(error) = self
if let Some(reset_at) = reset_at
&& let Err(error) = self
.usage_store
.sync_rate_limit_window(account_id, reset_at, limit_window_seconds)
.await
{
tracing::error!(
account_id,
error = %error,
"failed to persist verified quota window"
);
}
{
tracing::error!(
account_id,
error = %error,
"failed to persist verified quota window"
);
}
let in_memory = {
+4 -5
View File
@@ -451,13 +451,12 @@ impl AccountPool {
});
let mut account_ids = Vec::with_capacity(ordered.len());
if let Some(preferred_account_id) = request.preferred_account_id.as_deref() {
if let Some(index) = ordered
if let Some(preferred_account_id) = request.preferred_account_id.as_deref()
&& let Some(index) = ordered
.iter()
.position(|account| account.id == preferred_account_id)
{
account_ids.push(ordered.remove(index).id);
}
{
account_ids.push(ordered.remove(index).id);
}
account_ids.extend(ordered.into_iter().map(|account| account.id));
AccountCandidateSnapshotWithStatusRefresh {
+1 -1
View File
@@ -4,7 +4,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
use chrono::Utc;
use thiserror::Error;
use tokio::time::{sleep, Instant};
use tokio::time::{Instant, sleep};
use tracing::warn;
use crate::fleet::{
+1 -1
View File
@@ -3,7 +3,7 @@
use std::time::Duration;
use chrono::{DateTime, Utc};
use serde_json::{json, Value};
use serde_json::{Value, json};
const MONTH_WINDOW_MINUTES: u64 = 43_200;
const MONTH_WINDOW_SECONDS: u64 = 2_592_000;
+2 -2
View File
@@ -7,10 +7,10 @@ use std::{
time::Duration as StdDuration,
};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use chrono::{DateTime, Duration, TimeZone, Utc};
use serde_json::{Map, Value};
use tokio::sync::{watch, Semaphore};
use tokio::sync::{Semaphore, watch};
use crate::fleet::account::{Account, AccountStatus};
use crate::upstream::openai::token_client::{RefreshFailure, TokenPair, TokenRefresher};
+6 -6
View File
@@ -5,8 +5,8 @@ use std::{
future::Future,
pin::Pin,
sync::{
atomic::{AtomicBool, Ordering},
Arc, Mutex as StdMutex,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
@@ -30,12 +30,12 @@ use crate::upstream::openai::token_client::TokenRefresher;
use super::{
lease::{RedisRefreshLeaseStore, RedisRefreshLeaseStoreError},
policy::{
default_refresh_retry_delays, duration_until, is_permanent_refresh_failure_status,
jittered_refresh_at, jwt_expiration, same_refresh_time, scheduled_at_from_delay,
stable_jittered_duration, token_refresh_status_eligible, RefreshError, RefreshPolicy,
RefreshScheduler, RuntimeRefreshPolicy, TokenTimerSummary,
PERMANENT_FAILURE_CONFIRMATION_THRESHOLD, RECOVERY_DELAY_JITTER, RECOVERY_DELAY_SECONDS,
RETRY_DELAY_JITTER,
RETRY_DELAY_JITTER, RefreshError, RefreshPolicy, RefreshScheduler, RuntimeRefreshPolicy,
TokenTimerSummary, default_refresh_retry_delays, duration_until,
is_permanent_refresh_failure_status, jittered_refresh_at, jwt_expiration,
same_refresh_time, scheduled_at_from_delay, stable_jittered_duration,
token_refresh_status_eligible,
},
};
+2 -6
View File
@@ -6,8 +6,8 @@
//! 对齐 sub2api 的设计,避免选择热路径上的锁与 DB 往返。
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::RwLock;
use std::sync::atomic::{AtomicU64, Ordering};
/// EWMA 平滑系数(新样本权重)。与 sub2api 一致取 0.2。
const EWMA_ALPHA: f64 = 0.2;
@@ -59,11 +59,7 @@ impl AccountFeedback {
/// 读取 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)
}
if value.is_nan() { None } else { Some(value) }
}
/// 无锁 CAS 更新 EWMA:首个样本直接落值,其后按 alpha 混合。
@@ -12,7 +12,7 @@ use std::cmp::Ordering;
use crate::fleet::account::Account;
use crate::fleet::scheduler::{
compare_last_used, compare_window_reset, order_by, select_by, SelectionInput,
SelectionInput, compare_last_used, compare_window_reset, order_by, select_by,
};
/// QuotaResetPriority 策略选择。
@@ -11,7 +11,7 @@
use crate::fleet::account::Account;
use crate::fleet::scheduler::feedback::{FeedbackSample, FeedbackStats};
use crate::fleet::scheduler::{account_window_token_count, SelectionInput};
use crate::fleet::scheduler::{SelectionInput, account_window_token_count};
/// 配额封顶账号的惩罚分。远大于各因子加权和的理论上限(各因子 ∈ [0,1],权重和通常为
/// 个位数),保证未封顶账号总是优先于封顶账号。
@@ -202,11 +202,11 @@ impl NormalizationContext {
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(ttft) = sample.ttft_ms
&& 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 {
@@ -6,7 +6,7 @@
use std::cmp::Ordering;
use crate::fleet::account::Account;
use crate::fleet::scheduler::{compare_last_used, SelectionInput};
use crate::fleet::scheduler::{SelectionInput, compare_last_used};
/// 选中最近使用过的账号(`last_used_at` 最大者)。
///
+4 -4
View File
@@ -7,8 +7,8 @@ use serde_json::Value;
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
use thiserror::Error;
use crate::fleet::account::{Account, AccountStatus};
use crate::infra::json::{page_offset, NumberedPage};
use crate::fleet::account::{Account, AccountListSort, AccountSortField, AccountStatus};
use crate::infra::json::{NumberedPage, SortDirection, page_offset};
mod queries;
mod rows;
@@ -16,7 +16,7 @@ mod rows;
use queries::*;
use rows::{
count_account_metadata, get_pool_account, list_pool_accounts, map_account_store_error,
metadata_from_row, optional_update_value, push_account_metadata_search, quota_plan_type,
metadata_from_row, optional_update_value, push_account_metadata_filter, quota_plan_type,
quota_snapshot_from_row, status_to_db, stored_account_from_row,
};
@@ -249,7 +249,7 @@ pub trait AccountStore: Send + Sync + 'static {
/// 更新账号状态。
async fn set_status(&self, account_id: &str, status: AccountStatus)
-> AccountStoreResult<bool>;
-> AccountStoreResult<bool>;
/// 读取账号当前配额 JSON。
async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult<Option<String>>;
+44 -3
View File
@@ -315,15 +315,18 @@ impl PgAccountStore {
page: u32,
page_size: u32,
search: Option<&str>,
status: Option<AccountStatus>,
sort: Option<AccountListSort>,
) -> PgAccountStoreResult<NumberedPage<StoredAccountMetadata>> {
let page_size = page_size.clamp(1, 200);
let search = search.map(str::trim).filter(|value| !value.is_empty());
let total = count_account_metadata(&self.pool, search).await?;
let total = count_account_metadata(&self.pool, search, status).await?;
let offset = page_offset(page, page_size);
let mut builder = QueryBuilder::<Postgres>::new(LIST_ACCOUNT_METADATA_SELECT_SQL);
push_account_metadata_search(&mut builder, search);
builder.push(" order by added_at desc, id desc limit ");
push_account_metadata_filter(&mut builder, search, status);
push_account_metadata_order(&mut builder, sort);
builder.push(" limit ");
builder.push_bind(i64::from(page_size));
builder.push(" offset ");
builder.push_bind(offset.min(i64::MAX as u64) as i64);
@@ -373,3 +376,41 @@ impl PgAccountStore {
.map(|value| value.0.to_string()))
}
}
fn push_account_metadata_order(
builder: &mut QueryBuilder<Postgres>,
sort: Option<AccountListSort>,
) {
let Some(sort) = sort else {
builder.push(" order by added_at desc, id desc");
return;
};
builder.push(" order by ");
match sort.field {
AccountSortField::Email => builder.push("lower(coalesce(email, id))"),
AccountSortField::Status => builder.push(
"case status when 'active' then 0 when 'quota_exhausted' then 1 \
when 'expired' then 2 when 'disabled' then 3 when 'banned' then 4 else 5 end",
),
AccountSortField::PlanType => builder.push("lower(coalesce(plan_type, ''))"),
AccountSortField::Usage => builder.push(
"(select max((quota_values.value #>> '{}')::double precision) \
from jsonb_path_query(coalesce(accounts.quota_json, '{}'::jsonb), \
'$.**.used_percent') as quota_values(value) \
where jsonb_typeof(quota_values.value) = 'number')",
),
AccountSortField::LastUsedAt => builder.push(
"(select account_usage.last_used_at from account_usage \
where account_usage.account_id = accounts.id)",
),
AccountSortField::ExpiresAt => builder.push("access_token_expires_at"),
};
let direction = match sort.direction {
SortDirection::Asc => " asc",
SortDirection::Desc => " desc",
};
builder.push(direction);
builder.push(" nulls last, id");
builder.push(direction);
}
+26 -27
View File
@@ -2,7 +2,7 @@
use secrecy::SecretString;
use serde_json::Value;
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use sqlx::{PgPool, Postgres, QueryBuilder, Row, postgres::PgRow};
use crate::{
fleet::{
@@ -13,9 +13,9 @@ use crate::{
};
use super::{
queries::{GET_POOL_ACCOUNT_SQL, LIST_POOL_ACCOUNTS_SQL},
AccountQuotaSnapshot, AccountStoreError, PgAccountStore, PgAccountStoreError,
PgAccountStoreResult, StoredAccount, StoredAccountMetadata,
queries::{GET_POOL_ACCOUNT_SQL, LIST_POOL_ACCOUNTS_SQL},
};
pub(super) async fn list_pool_accounts(
@@ -176,43 +176,42 @@ pub(super) fn optional_update_value(value: &Option<Option<String>>) -> Option<&s
}
pub(super) fn status_from_db(value: &str) -> PgAccountStoreResult<AccountStatus> {
match value {
"active" => Ok(AccountStatus::Active),
"expired" => Ok(AccountStatus::Expired),
"quota_exhausted" => Ok(AccountStatus::QuotaExhausted),
"disabled" => Ok(AccountStatus::Disabled),
"banned" => Ok(AccountStatus::Banned),
other => Err(PgAccountStoreError::InvalidStatus(other.to_string())),
}
AccountStatus::parse(value).ok_or_else(|| PgAccountStoreError::InvalidStatus(value.to_string()))
}
pub(super) async fn count_account_metadata(
pool: &PgPool,
search: Option<&str>,
status: Option<AccountStatus>,
) -> PgAccountStoreResult<u64> {
let mut builder = QueryBuilder::<Postgres>::new("select count(*) from accounts");
push_account_metadata_search(&mut builder, search);
push_account_metadata_filter(&mut builder, search, status);
let (total,): (i64,) = builder.build_query_as().fetch_one(pool).await?;
Ok(nonnegative_i64_to_u64(total))
}
pub(super) fn push_account_metadata_search(
pub(super) fn push_account_metadata_filter(
builder: &mut QueryBuilder<Postgres>,
search: Option<&str>,
status: Option<AccountStatus>,
) {
let Some(search) = search else {
return;
};
let pattern = format!("%{search}%");
builder.push(" where id ilike ");
builder.push_bind(pattern.clone());
builder.push(" or email ilike ");
builder.push_bind(pattern.clone());
builder.push(" or label ilike ");
builder.push_bind(pattern.clone());
builder.push(" or chatgpt_account_id ilike ");
builder.push_bind(pattern.clone());
builder.push(" or chatgpt_user_id ilike ");
builder.push_bind(pattern);
builder.push(" where true");
if let Some(search) = search {
let pattern = format!("%{search}%");
builder.push(" and (id ilike ");
builder.push_bind(pattern.clone());
builder.push(" or email ilike ");
builder.push_bind(pattern.clone());
builder.push(" or label ilike ");
builder.push_bind(pattern.clone());
builder.push(" or chatgpt_account_id ilike ");
builder.push_bind(pattern.clone());
builder.push(" or chatgpt_user_id ilike ");
builder.push_bind(pattern);
builder.push(")");
}
if let Some(status) = status {
builder.push(" and status = ");
builder.push_bind(status.as_str());
}
}
+1 -1
View File
@@ -3,7 +3,7 @@
use std::time::Duration;
use sha2::{Digest, Sha256};
use sqlx::{postgres::PgPoolOptions, PgPool, Postgres, Row, Transaction};
use sqlx::{PgPool, Postgres, Row, Transaction, postgres::PgPoolOptions};
const MIGRATIONS: &[Migration] = &[Migration {
version: 1,
+2 -2
View File
@@ -77,8 +77,8 @@ pub fn format_rate(value: Option<f64>) -> String {
value.map_or_else(|| "—".to_string(), |value| format_percent(value * 100.0))
}
/// 使用紧凑货币格式展示美元成本。
pub fn format_cost(value: f64) -> String {
/// 使用紧凑货币格式展示美元计费金额。
pub fn format_billing_amount(value: f64) -> String {
if !value.is_finite() {
return "—".to_string();
}
+3 -3
View File
@@ -1,11 +1,11 @@
//! 身份认证相关原语:管理员密码哈希与 API Key 生成。
use argon2::{
password_hash::{rand_core::OsRng, PasswordHash, PasswordHasher, PasswordVerifier, SaltString},
Argon2,
password_hash::{PasswordHash, PasswordHasher, PasswordVerifier, SaltString, rand_core::OsRng},
};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use hmac::{Hmac, Mac};
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
use hmac::{Hmac, KeyInit, Mac};
use rand::Rng;
use sha2::{Digest, Sha256};
use thiserror::Error;
+20
View File
@@ -2,6 +2,26 @@
use serde::Serialize;
/// 列表排序方向。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SortDirection {
/// 升序。
Asc,
/// 降序。
Desc,
}
impl SortDirection {
/// 解析管理端查询参数中的排序方向。
pub fn parse(value: &str) -> Option<Self> {
match value.trim().to_ascii_lowercase().as_str() {
"asc" => Some(Self::Asc),
"desc" => Some(Self::Desc),
_ => None,
}
}
}
// ---------------------------------------------------------------------------
// 分页
// ---------------------------------------------------------------------------
+11 -9
View File
@@ -16,10 +16,10 @@ use chrono::{DateTime, Duration, NaiveDate, Utc};
use thiserror::Error;
use tracing_appender::non_blocking::{ErrorCounter, NonBlockingBuilder, WorkerGuard};
use tracing_subscriber::{
EnvFilter,
fmt::{self as tracing_fmt, format::Writer, time::FormatTime},
layer::SubscriberExt,
util::SubscriberInitExt,
EnvFilter,
};
use crate::infra::time::{china_date, china_rfc3339_millis};
@@ -459,14 +459,16 @@ impl DroppedLogMonitor {
let (shutdown_tx, shutdown_rx) = mpsc::channel();
let handle = thread::Builder::new()
.name("tracing-drop-monitor".to_string())
.spawn(move || loop {
match shutdown_rx.recv_timeout(DROP_MONITOR_INTERVAL) {
Ok(()) | Err(mpsc::RecvTimeoutError::Disconnected) => {
report_dropped_logs(&mut counters, &mut io::stderr().lock());
break;
}
Err(mpsc::RecvTimeoutError::Timeout) => {
report_dropped_logs(&mut counters, &mut io::stderr().lock());
.spawn(move || {
loop {
match shutdown_rx.recv_timeout(DROP_MONITOR_INTERVAL) {
Ok(()) | Err(mpsc::RecvTimeoutError::Disconnected) => {
report_dropped_logs(&mut counters, &mut io::stderr().lock());
break;
}
Err(mpsc::RecvTimeoutError::Timeout) => {
report_dropped_logs(&mut counters, &mut io::stderr().lock());
}
}
}
})?;
+1 -1
View File
@@ -1,6 +1,6 @@
//! Redis 运行态连接与键前缀。
use redis::{aio::ConnectionManager, AsyncCommands};
use redis::{AsyncCommands, aio::ConnectionManager};
/// Redis 运行态连接。
#[derive(Clone)]
+4 -2
View File
@@ -5,7 +5,8 @@ use crate::infra::json::NumberedPage;
use super::{
store::PgClientKeyStore,
types::{
parse_client_key_status, BatchDeleteClientApiKeys, KeyManageError, ManagedClientApiKey,
BatchDeleteClientApiKeys, ClientApiKeyListSort, KeyManageError, ManagedClientApiKey,
parse_client_key_status,
},
};
@@ -36,10 +37,11 @@ impl KeyManageService {
page: u32,
page_size: u32,
search: Option<String>,
sort: Option<ClientApiKeyListSort>,
) -> Result<NumberedPage<ManagedClientApiKey>, KeyManageError> {
let page = self
.store
.list_page(page, page_size, search.as_deref())
.list_page(page, page_size, search.as_deref(), sort)
.await
.map_err(|_| KeyManageError::List)?;
Ok(NumberedPage {
+1 -1
View File
@@ -3,8 +3,8 @@
use std::{
collections::BTreeMap,
sync::{
atomic::{AtomicBool, Ordering},
Arc, Mutex,
atomic::{AtomicBool, Ordering},
},
time::Duration,
};
+27 -2
View File
@@ -9,8 +9,9 @@ use uuid::Uuid;
use crate::infra::{
identity::generate_client_api_key,
json::{page_offset, NumberedPage},
json::{NumberedPage, page_offset},
};
use crate::keys::types::{ClientApiKeyListSort, ClientApiKeySortField};
const CLIENT_KEY_SELECT_SQL: &str = r"select
id, name, label, prefix, key, enabled, created_at, last_used_at
@@ -96,6 +97,7 @@ values ($1, $2, null, $3, $4, true, $5, null)",
page: u32,
page_size: u32,
search: Option<&str>,
sort: Option<ClientApiKeyListSort>,
) -> Result<NumberedPage<StoredClientApiKey>, PgClientKeyStoreError> {
let page = page.max(1);
let page_size = page_size.clamp(1, 200);
@@ -112,7 +114,8 @@ values ($1, $2, null, $3, $4, true, $5, null)",
let mut builder = QueryBuilder::<Postgres>::new(CLIENT_KEY_SELECT_SQL);
push_search(&mut builder, search);
builder.push(" order by created_at desc, id desc limit ");
push_order(&mut builder, sort);
builder.push(" limit ");
builder.push_bind(i64::from(page_size));
builder.push(" offset ");
builder.push_bind(page_offset(page, page_size).min(i64::MAX as u64) as i64);
@@ -213,6 +216,28 @@ where keys.id = touched.id",
}
}
fn push_order(builder: &mut QueryBuilder<Postgres>, sort: Option<ClientApiKeyListSort>) {
let Some(sort) = sort else {
builder.push(" order by created_at desc, id desc");
return;
};
builder.push(" order by ");
match sort.field {
ClientApiKeySortField::Name => builder.push("lower(name)"),
ClientApiKeySortField::Enabled => builder.push("enabled"),
ClientApiKeySortField::CreatedAt => builder.push("created_at"),
ClientApiKeySortField::LastUsedAt => builder.push("last_used_at"),
};
let direction = match sort.direction {
crate::infra::json::SortDirection::Asc => " asc",
crate::infra::json::SortDirection::Desc => " desc",
};
builder.push(direction);
builder.push(" nulls last, id");
builder.push(direction);
}
fn push_search(builder: &mut QueryBuilder<Postgres>, search: Option<&str>) {
let Some(search) = search else {
return;
+16
View File
@@ -2,8 +2,24 @@
use thiserror::Error;
use crate::infra::json::SortDirection;
use super::store::StoredClientApiKey;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ClientApiKeySortField {
Name,
Enabled,
CreatedAt,
LastUsedAt,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ClientApiKeyListSort {
pub field: ClientApiKeySortField,
pub direction: SortDirection,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ManagedClientApiKey {
pub id: String,
+2 -15
View File
@@ -1,4 +1,4 @@
use std::{env, error::Error, io, path::PathBuf};
use std::{env, error::Error, io};
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
@@ -8,19 +8,6 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
reject_extra_arguments(arguments)?;
codex_proxy_rs::bootstrap::services::run().await
}
Some("import-sqlite") => {
let source = arguments
.next()
.map(PathBuf::from)
.ok_or_else(|| invalid_cli("usage: codex-proxy-rs import-sqlite <path.sqlite>"))?;
reject_extra_arguments(arguments)?;
let config = codex_proxy_rs::bootstrap::config::AppConfig::load()?;
let pool = codex_proxy_rs::infra::database::connect(&config.database.url).await?;
let report =
codex_proxy_rs::bootstrap::import_sqlite::import_sqlite(&pool, source).await?;
println!("{report}");
Ok(())
}
Some("rebuild-buckets") => {
reject_extra_arguments(arguments)?;
let config = codex_proxy_rs::bootstrap::config::AppConfig::load()?;
@@ -35,7 +22,7 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync>> {
Ok(())
}
Some(command) => Err(invalid_cli(&format!(
"unknown command {command:?}; expected serve, import-sqlite, or rebuild-buckets"
"unknown command {command:?}; expected serve or rebuild-buckets"
))
.into()),
}
+11 -56
View File
@@ -7,17 +7,6 @@ use super::types::{
ReasoningEffortInfo,
};
const BUILTIN_CODEX_MODELS: [(&str, &str, bool); 8] = [
("gpt-5.5", "GPT-5.5", true),
("gpt-5.4", "GPT-5.4", false),
("gpt-5.4-mini", "GPT-5.4 Mini", false),
("gpt-5.3-codex", "GPT-5.3 Codex", false),
("gpt-5.3-codex-spark", "GPT-5.3 Codex Spark", false),
("codex-auto-review", "Codex Auto Review", false),
("gpt-5.2", "GPT-5.2", false),
("gpt-image-1", "GPT Image 1", false),
];
/// 模型目录聚合。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelCatalog {
@@ -55,11 +44,11 @@ impl ModelPlanSnapshot {
}
impl ModelCatalog {
/// 从静态配置构造模型目录。
/// 从运行时别名构造尚未刷新上游快照的空模型目录。
pub fn from_config(config: &ModelConfig) -> Self {
let model_aliases = normalize_aliases(&config.model_aliases);
Self {
models: builtin_codex_models(),
models: Vec::new(),
model_aliases,
model_plan_index: BTreeMap::new(),
fetched_plan_types: BTreeSet::new(),
@@ -101,7 +90,7 @@ impl ModelCatalog {
.collect::<Vec<_>>();
if models.is_empty() {
return Self {
models: builtin_codex_models(),
models,
model_aliases,
model_plan_index,
fetched_plan_types,
@@ -230,39 +219,6 @@ fn normalize_aliases(input: &BTreeMap<String, String>) -> BTreeMap<String, Strin
.collect()
}
fn builtin_codex_models() -> Vec<CodexModelInfo> {
BUILTIN_CODEX_MODELS
.into_iter()
.map(|(id, display_name, is_default)| CodexModelInfo {
id: id.to_string(),
display_name: display_name.to_string(),
description: String::new(),
is_default,
supported_reasoning_efforts: default_reasoning_efforts(),
default_reasoning_effort: "medium".to_string(),
input_modalities: vec!["text".to_string()],
output_modalities: vec!["text".to_string()],
supports_personality: false,
upgrade: None,
source: "builtin".to_string(),
context_window: None,
max_context_window: None,
max_output_tokens: None,
truncation_policy_limit: None,
})
.collect()
}
fn default_reasoning_efforts() -> Vec<ReasoningEffortInfo> {
["minimal", "low", "medium", "high"]
.into_iter()
.map(|reasoning_effort| ReasoningEffortInfo {
reasoning_effort: reasoning_effort.to_string(),
description: String::new(),
})
.collect()
}
fn normalize_backend_model(raw: BackendModelEntry) -> CodexModelInfo {
let id = first_non_empty([raw.slug.as_deref(), raw.id.as_deref(), raw.name.as_deref()])
.unwrap_or("unknown")
@@ -279,22 +235,21 @@ fn normalize_backend_model(raw: BackendModelEntry) -> CodexModelInfo {
} else {
raw.supported_reasoning_efforts
};
let mut supported_reasoning_efforts = reasoning_entries
let supported_reasoning_efforts = reasoning_entries
.into_iter()
.filter_map(normalize_backend_reasoning_effort)
.collect::<Vec<_>>();
if supported_reasoning_efforts.is_empty() {
supported_reasoning_efforts.push(ReasoningEffortInfo {
reasoning_effort: "medium".to_string(),
description: "Default".to_string(),
});
}
let default_reasoning_effort = first_non_empty([
raw.default_reasoning_effort.as_deref(),
raw.default_reasoning_level.as_deref(),
])
.unwrap_or("medium")
.to_string();
.map(ToString::to_string)
.or_else(|| {
supported_reasoning_efforts
.first()
.map(|effort| effort.reasoning_effort.clone())
})
.unwrap_or_default();
CodexModelInfo {
id,
+1 -1
View File
@@ -3,7 +3,7 @@
use std::{collections::BTreeMap, sync::Arc};
use subtle::ConstantTimeEq;
use tokio::sync::{watch, Mutex};
use tokio::sync::{Mutex, watch};
use crate::infra::identity::{generate_admin_api_key, hash_credential};
+1 -1
View File
@@ -3,7 +3,7 @@
use std::collections::BTreeMap;
use chrono::Utc;
use sqlx::{postgres::PgRow, PgPool, Row};
use sqlx::{PgPool, Row, postgres::PgRow};
use super::types::{ManagementApiKeyStatus, SettingsError, SettingsSnapshot};
+5 -5
View File
@@ -163,7 +163,7 @@ pub struct AccountUsageTimeBucket {
pub latency_count: i64,
pub max_latency_ms: i64,
pub min_latency_ms: i64,
pub cost_usd: Option<f64>,
pub billing_amount_usd: Option<f64>,
}
impl From<UsageSummary> for AccountUsageSummary {
@@ -187,7 +187,7 @@ impl From<UsageSummary> for AccountUsageSummary {
impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
fn from(record: UsageTimeBucketRecord) -> Self {
let cost_usd = time_bucket_cost_usd(&record);
let billing_amount_usd = time_bucket_billing_amount_usd(&record);
Self {
bucket_start: record.bucket_start,
model: record.model,
@@ -203,12 +203,12 @@ impl From<UsageTimeBucketRecord> for AccountUsageTimeBucket {
latency_count: record.latency_count,
max_latency_ms: record.max_latency_ms,
min_latency_ms: record.min_latency_ms,
cost_usd,
billing_amount_usd,
}
}
}
fn time_bucket_cost_usd(record: &UsageTimeBucketRecord) -> Option<f64> {
fn time_bucket_billing_amount_usd(record: &UsageTimeBucketRecord) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(record.input_tokens);
let output_tokens = nonnegative_i64_to_u64(record.output_tokens);
let cached_tokens = nonnegative_i64_to_u64(record.cached_tokens);
@@ -221,7 +221,7 @@ fn time_bucket_cost_usd(record: &UsageTimeBucketRecord) -> Option<f64> {
return None;
}
Some(billing::calculate_cost(
Some(billing::calculate_billing_amount(
input_tokens,
output_tokens,
cached_tokens,

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