diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c642f6a82..1dea14e58 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,6 +16,29 @@ concurrency: jobs: backend: runs-on: ubuntu-latest + services: + postgres: + image: postgres:18-alpine + env: + POSTGRES_DB: codex_proxy + POSTGRES_USER: codex_proxy + POSTGRES_PASSWORD: codex_proxy + ports: + - 5432:5432 + options: >- + --health-cmd "pg_isready -U codex_proxy -d codex_proxy" + --health-interval 5s + --health-timeout 5s + --health-retries 10 + redis: + image: redis:8-alpine + ports: + - 6379:6379 + options: >- + --health-cmd "redis-cli ping" + --health-interval 5s + --health-timeout 5s + --health-retries 10 steps: - name: Checkout uses: actions/checkout@v7 @@ -38,6 +61,9 @@ jobs: working-directory: backend run: cargo clippy --all-targets --all-features --locked -- -D warnings + - name: Architecture boundaries + run: backend/scripts/check_architecture.sh + - name: Rust tests working-directory: backend run: cargo test --test main --locked diff --git a/README.md b/README.md index 63ca35664..3b25776b7 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,8 @@ - + +

@@ -17,15 +18,16 @@ Codex Proxy RS 是一个单进程网关： -- Rust/Axum 后端提供 Responses 兼容代理、管理端 API、SQLite 持久化和静态前端托管。 +- Rust/Axum 后端提供 Responses 兼容代理、管理端 API、PostgreSQL 持久化、Redis 运行态存储和静态前端托管。 - Vue 管理端负责账号导入、API Key、用量统计、请求记录、运行参数、模型映射和系统更新。 -- 运行数据默认写入仓库 `.runtime/`；Docker 部署也使用这套宿主机目录挂载。 +- 应用本地文件默认写入仓库 `.runtime/`；PostgreSQL 与 Redis 数据由各自服务管理。 ## 环境 - Rust 1.95 - Node 24 或兼容版本，pnpm 11（只在开发前端或构建镜像时需要） -- Docker / Docker Compose（部署或验证镜像时需要） +- PostgreSQL 18、Redis 8 +- Docker / Docker Compose（推荐用于启动依赖、部署或验证镜像） ## 本地运行 @@ -35,11 +37,14 @@ Codex Proxy RS 是一个单进程网关： cp deploy/config.example.yaml .runtime/config.yaml ``` -如果直接用本地二进制而不是 Docker，确保 `.runtime/config.yaml` 里的路径是宿主机路径： +如果直接用本地二进制而不是 Docker，确保 `.runtime/config.yaml` 使用宿主机可访问的连接地址与日志路径： ```yaml database: - url: 'sqlite://.runtime/data/codex-proxy-rs.sqlite' + url: 'postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy' + +redis: + url: 'redis://127.0.0.1:6379/' logging: directory: .runtime/logs @@ -59,12 +64,14 @@ CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.tom ## Docker 部署 -Compose 默认把宿主机 `.runtime` 映射进容器： +Compose 同时启动 PostgreSQL、Redis 和主服务。主服务默认把宿主机 `.runtime` 映射进容器： - `.runtime/config.yaml` -> `/app/config.yaml` -- `.runtime/data` -> `/app/data` +- `.runtime/data` -> `/app/data`（installation id 与在线更新状态，不含数据库） - `.runtime/logs` -> `/app/logs` +PostgreSQL 与 Redis 分别使用 `postgres-data`、`redis-data` 命名卷。 + 初始化配置： ```bash @@ -77,7 +84,10 @@ Docker 配置文件里路径保持容器路径： ```yaml database: - url: 'sqlite:///app/data/codex-proxy-rs.sqlite' + url: 'postgres://codex_proxy:codex_proxy@postgres:5432/codex_proxy' + +redis: + url: 'redis://redis:6379/' logging: directory: /app/logs @@ -142,20 +152,19 @@ CPR_CONFIG_FILE=.runtime/config.yaml cargo run --manifest-path backend/Cargo.tom 管理端功能： - 账号导入、OAuth 授权、连接测试、quota 刷新、token 刷新。 -- 客户端 API Key 管理。 +- 客户端 API Key 管理，完整 Key 可长期复制并导入 CCSwitch。 - Dashboard、请求记录、Token 用量、模型用量和账号额度。 - 运行参数、模型别名、账号选择策略。 - 版本检查、在线更新、更新日志、重启和回滚。 ## 配置与运行数据 -配置文件只承载启动必需项。账号、API Key、模型映射、账号选择策略和多数运行参数由管理端写入 SQLite。 +配置文件只承载启动必需项。账号、API Key、模型映射、账号选择策略和多数运行参数由管理端写入 PostgreSQL；管理会话、刷新租约、会话亲和与模型缓存位于 Redis。 `.runtime` 约定： ```text .runtime/config.yaml # 本地或 Docker 启动配置，路径按运行环境写 -.runtime/data/codex-proxy-rs.sqlite # SQLite 数据库 .runtime/data/installation_id # 上游 installation id .runtime/data/update-state.json # 在线更新状态 .runtime/data/update.lock # 在线更新锁 @@ -184,7 +193,7 @@ release/publish 1.0.4 ## 项目结构 ```text -backend/ Rust/Axum 后端、SQLite migration、集成测试 +backend/ Rust/Axum 后端、PostgreSQL migration、SQLite 导入工具、集成测试 frontend/ Vue 3 管理端 deploy/ Dockerfile、Compose、部署配置模板 docs/ 架构和维护文档 diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 909757c40..79fe04aef 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -47,6 +47,21 @@ dependencies = [ "libc", ] +[[package]] +name = "arc-swap" +version = "1.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" +dependencies = [ + "rustversion", +] + +[[package]] +name = "arcstr" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" + [[package]] name = "argon2" version = "0.5.3" @@ -87,6 +102,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener", + "event-listener-strategy", + "pin-project-lite", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -171,6 +197,15 @@ dependencies = [ "tracing", ] +[[package]] +name = "backon" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" +dependencies = [ + "fastrand", +] + [[package]] name = "base64" version = "0.22.1" @@ -332,6 +367,7 @@ dependencies = [ "ipnet", "rand 0.10.1", "rand_core 0.6.4", + "redis", "reqwest", "rustls", "rustls-native-certs", @@ -358,6 +394,20 @@ dependencies = [ "wiremock", ] +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "futures-core", + "memchr", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "compression-codecs" version = "0.4.38" @@ -755,6 +805,16 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener", + "pin-project-lite", +] + [[package]] name = "fastrand" version = "2.4.1" @@ -1882,6 +1942,33 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "redis" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fa6f8e4b491d7a8ef3a9550a4d71969bd0064f46e32b8dbbcc7fc60dad94fed" +dependencies = [ + "arc-swap", + "arcstr", + "async-lock", + "backon", + "bytes", + "cfg-if", + "combine", + "futures-channel", + "futures-util", + "itoa", + "percent-encoding", + "pin-project-lite", + "ryu", + "sha1_smol", + "socket2", + "tokio", + "tokio-util", + "url", + "xxhash-rust", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2254,6 +2341,12 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sha1_smol" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d" + [[package]] name = "sha2" version = "0.10.9" @@ -3577,6 +3670,12 @@ dependencies = [ "rustix", ] +[[package]] +name = "xxhash-rust" +version = "0.8.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d93c89cdc2d3a63c3ec48ffe926931bdc069eafa8e4402fe6d8f790c9d1e576" + [[package]] name = "yaml-rust2" version = "0.11.0" diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 4aca96cbd..98405de14 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -4,7 +4,6 @@ edition = "2021" license = "MIT" publish = false rust-version = "1.95" -build = "build/build.rs" [dependencies] argon2 = "0.5.3" @@ -22,6 +21,7 @@ indexmap = { version = "2.0", features = ["serde"] } ipnet = "2.12.0" rand = "0.10.1" rand_core = { version = "0.6.4", features = ["getrandom"] } +redis = { version = "1.3.0", default-features = false, features = ["tokio-comp", "connection-manager", "script"] } reqwest = { version = "=0.12.28", default-features = false, features = ["json", "stream", "cookies", "rustls-tls-native-roots", "gzip", "brotli", "zstd", "deflate", "http2"] } rustls = { version = "=0.23.36", default-features = false, features = ["ring", "std", "tls12"] } rustls-native-certs = "0.8.4" @@ -31,7 +31,7 @@ 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", features = ["runtime-tokio", "tls-rustls", "sqlite", "chrono", "uuid"] } +sqlx = { version = "0.9.0", features = ["runtime-tokio", "tls-rustls", "postgres", "sqlite", "chrono", "uuid", "json"] } subtle = "2.6.1" tar = "0.4.44" thiserror = "2.0.18" diff --git a/backend/build/build.rs b/backend/build.rs similarity index 100% rename from backend/build/build.rs rename to backend/build.rs diff --git a/backend/scripts/check_architecture.sh b/backend/scripts/check_architecture.sh new file mode 100755 index 000000000..5fdf124c5 --- /dev/null +++ b/backend/scripts/check_architecture.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +set -euo pipefail + +cd "$(dirname "$0")/.." + +fail=0 + +deny() { + local description=$1 + local pattern=$2 + shift 2 + if grep -rn --include='*.rs' -E "$pattern" "$@"; then + echo "ARCH VIOLATION: $description" >&2 + fail=1 + fi +} + +domain_paths=( + src/infra src/upstream src/telemetry src/keys src/auth src/settings src/update + src/accounts src/models src/dispatch +) + +deny "domain -> api/bootstrap" 'crate::(api|bootstrap)\b' "${domain_paths[@]}" +deny "upstream -> domain" 'crate::(accounts|dispatch|telemetry|keys|auth|settings|update|models)\b' src/upstream +deny "axum outside api/bootstrap" '\baxum(::|\b)' "${domain_paths[@]}" +deny "sqlx in api" '\bsqlx(::|\b)' src/api +deny "redis in api" '\bredis::' src/api +deny "legacy module path" 'crate::(admin|proxy|runtime|web|http|config)::' src +deny "Admin-prefixed domain type" '\bAdmin[A-Z][A-Za-z0-9]*\b' "${domain_paths[@]}" +deny "Repository naming" '\b[A-Za-z0-9]*Repository\b' src +deny "persistence repository naming" '\b[A-Za-z0-9_]*repository[A-Za-z0-9_]*\b' \ + src/accounts src/auth src/dispatch src/keys src/models src/settings src/telemetry \ + src/upstream/openai/fingerprint +deny "Arc alias" '\bArc[[:space:]]+as[[:space:]]+' src + +banned=$(find src \( -name 'importing.rs' -o -name 'exporting.rs' -o -name 'testing.rs' \ + -o -name 'util.rs' -o -name 'helper.rs' -o -name 'common.rs' -o -name 'misc.rs' \)) +if [ -n "$banned" ]; then + echo "ARCH VIOLATION: banned file names: $banned" >&2 + fail=1 +fi + +while read -r file; do + lines=$(wc -l < "$file") + if [ "$lines" -gt 50 ]; then + echo "ARCH VIOLATION: $file has $lines lines (mod.rs <= 50)" >&2 + fail=1 + fi +done < <(find src -name mod.rs) + +while read -r file; do + lines=$(wc -l < "$file") + if [ "$lines" -gt 800 ]; then + echo "ARCH VIOLATION: $file has $lines lines (> 800)" >&2 + fail=1 + fi +done < <(find src -name '*.rs' ! -name mod.rs) + +exit "$fail" diff --git a/backend/src/upstream/accounts/model.rs b/backend/src/accounts/account.rs similarity index 100% rename from backend/src/upstream/accounts/model.rs rename to backend/src/accounts/account.rs diff --git a/backend/src/accounts/cookies/mod.rs b/backend/src/accounts/cookies/mod.rs new file mode 100644 index 000000000..f46bf4eca --- /dev/null +++ b/backend/src/accounts/cookies/mod.rs @@ -0,0 +1,5 @@ +//! 账号 Cookie 存储。 + +mod store; + +pub use store::*; diff --git a/backend/src/upstream/accounts/cookies.rs b/backend/src/accounts/cookies/store.rs similarity index 85% rename from backend/src/upstream/accounts/cookies.rs rename to backend/src/accounts/cookies/store.rs index 9752ee167..02f569fdc 100644 --- a/backend/src/upstream/accounts/cookies.rs +++ b/backend/src/accounts/cookies/store.rs @@ -3,7 +3,7 @@ use std::{collections::HashMap, sync::Arc}; use chrono::{DateTime, Duration, Utc}; -use sqlx::{Row, SqlitePool}; +use sqlx::{PgPool, Row}; use thiserror::Error; use tokio::sync::RwLock; use uuid::Uuid; @@ -11,35 +11,31 @@ use uuid::Uuid; const DEFAULT_COOKIE_DOMAIN: &str = "chatgpt.com"; const CAPTURABLE_COOKIES: &[&str] = &["cf_clearance"]; -/// SQLite Cookie 存储错误。 +/// PostgreSQL Cookie 存储错误。 #[derive(Debug, Error)] -pub enum SqliteCookieStoreError { +pub enum PgCookieStoreError { /// 数据库错误。 - #[error("sqlite cookie store database error: {0}")] + #[error("PostgreSQL cookie store database error: {0}")] Database(#[from] sqlx::Error), } -/// SQLite Cookie 存储结果。 -pub type SqliteCookieStoreResult = Result; +/// PostgreSQL Cookie 存储结果。 +pub type PgCookieStoreResult = Result; -/// SQLite Cookie 存储。 +/// PostgreSQL Cookie 存储。 #[derive(Clone)] -pub struct SqliteCookieStore { - pool: SqlitePool, +pub struct PgCookieStore { + pool: PgPool, } -impl SqliteCookieStore { +impl PgCookieStore { /// 构造 Cookie 存储。 - pub fn new(pool: SqlitePool) -> Self { + pub fn new(pool: PgPool) -> Self { Self { pool } } /// 捕获上游 `Set-Cookie` 响应头中允许持久化的 Cookie。 - pub async fn capture_set_cookie( - &self, - account_id: &str, - raw: &str, - ) -> SqliteCookieStoreResult<()> { + pub async fn capture_set_cookie(&self, account_id: &str, raw: &str) -> PgCookieStoreResult<()> { let Some(parsed) = parse_set_cookie(raw) else { return Ok(()); }; @@ -54,7 +50,7 @@ impl SqliteCookieStore { &self, account_id: &str, request_domain: &str, - ) -> SqliteCookieStoreResult> { + ) -> PgCookieStoreResult> { self.cookie_header_for_request(account_id, request_domain, "/") .await } @@ -65,9 +61,9 @@ impl SqliteCookieStore { account_id: &str, request_domain: &str, request_path: &str, - ) -> SqliteCookieStoreResult> { + ) -> PgCookieStoreResult> { let rows = sqlx::query( - "select domain, name, value, path, expires_at from account_cookies where account_id = ?", + "select domain, name, value, path, expires_at from account_cookies where account_id = $1", ) .bind(account_id) .fetch_all(&self.pool) @@ -83,8 +79,8 @@ impl SqliteCookieStore { if !path_matches(request_path, &path) { continue; } - let expires_at = row.get::, _>("expires_at"); - if cookie_is_expired(expires_at.as_deref(), now) { + let expires_at = row.get::>, _>("expires_at"); + if expires_at.is_some_and(|expires_at| expires_at <= now) { continue; } let name = row.get::("name"); @@ -114,8 +110,8 @@ impl SqliteCookieStore { } /// 删除账号全部 Cookie。 - pub async fn delete_account_cookies(&self, account_id: &str) -> SqliteCookieStoreResult { - let result = sqlx::query("delete from account_cookies where account_id = ?") + pub async fn delete_account_cookies(&self, account_id: &str) -> PgCookieStoreResult { + let result = sqlx::query("delete from account_cookies where account_id = $1") .bind(account_id) .execute(&self.pool) .await?; @@ -123,11 +119,11 @@ impl SqliteCookieStore { } /// 删除指定时间之前过期的 Cookie。 - pub async fn cleanup_expired(&self, now: DateTime) -> SqliteCookieStoreResult { + pub async fn cleanup_expired(&self, now: DateTime) -> PgCookieStoreResult { let result = sqlx::query( - "delete from account_cookies where expires_at is not null and expires_at <= ?", + "delete from account_cookies where expires_at is not null and expires_at <= $1", ) - .bind(now.to_rfc3339()) + .bind(now) .execute(&self.pool) .await?; Ok(result.rows_affected()) @@ -137,10 +133,10 @@ impl SqliteCookieStore { &self, account_id: &str, parsed: ParsedCookie, - ) -> SqliteCookieStoreResult<()> { - let now = Utc::now().to_rfc3339(); + ) -> PgCookieStoreResult<()> { + let now = Utc::now(); sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?) on conflict(account_id, domain, name, path) do update set value = excluded.value, expires_at = excluded.expires_at, updated_at = excluded.updated_at", + "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7, $8) on conflict(account_id, domain, name, path) do update set value = excluded.value, expires_at = excluded.expires_at, updated_at = excluded.updated_at", ) .bind(Uuid::new_v4().to_string()) .bind(account_id) @@ -162,7 +158,7 @@ struct ParsedCookie { name: String, value: String, path: String, - expires_at: Option, + expires_at: Option>, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -197,7 +193,7 @@ fn parse_set_cookie(raw: &str) -> Option { } break; } - "expires" => expires_at = Some(value.trim().to_string()), + "expires" => expires_at = parse_cookie_expires_at(value.trim()), _ => {} } } @@ -211,12 +207,12 @@ fn parse_set_cookie(raw: &str) -> Option { }) } -fn max_age_expires_at(seconds: i64) -> String { +fn max_age_expires_at(seconds: i64) -> DateTime { let now = Utc::now(); if seconds <= 0 { - return (now - Duration::seconds(1)).to_rfc3339(); + return now - Duration::seconds(1); } - (now + Duration::seconds(seconds.min(i32::MAX as i64))).to_rfc3339() + now + Duration::seconds(seconds.min(i32::MAX as i64)) } fn domain_matches(request_domain: &str, cookie_domain: &str) -> bool { @@ -256,12 +252,6 @@ fn normalize_cookie_path(path: &str) -> String { } } -fn cookie_is_expired(expires_at: Option<&str>, now: DateTime) -> bool { - expires_at - .and_then(parse_cookie_expires_at) - .is_some_and(|expires_at| expires_at <= now) -} - fn parse_cookie_expires_at(value: &str) -> Option> { DateTime::parse_from_rfc2822(value) .or_else(|_| DateTime::parse_from_rfc3339(value)) diff --git a/backend/src/upstream/accounts/importing.rs b/backend/src/accounts/import.rs similarity index 99% rename from backend/src/upstream/accounts/importing.rs rename to backend/src/accounts/import.rs index bb2d3f893..fc750b327 100644 --- a/backend/src/upstream/accounts/importing.rs +++ b/backend/src/accounts/import.rs @@ -3,8 +3,8 @@ use chrono::{DateTime, Utc}; use serde_json::{Map, Value}; -use crate::upstream::accounts::model::AccountStatus; -use crate::upstream::accounts::token_refresh::{jwt_expiry, JwtExpiry}; +use crate::accounts::account::AccountStatus; +use crate::accounts::refresh::{jwt_expiry, JwtExpiry}; /// 账号导入条目。 #[derive(Debug, Clone)] diff --git a/backend/src/admin/accounts/service/cookies.rs b/backend/src/accounts/manage/cookies.rs similarity index 81% rename from backend/src/admin/accounts/service/cookies.rs rename to backend/src/accounts/manage/cookies.rs index 29b0df90a..f301f37a1 100644 --- a/backend/src/admin/accounts/service/cookies.rs +++ b/backend/src/accounts/manage/cookies.rs @@ -1,6 +1,6 @@ -use super::AdminAccountService; +use super::AccountManageService; -impl AdminAccountService { +impl AccountManageService { pub(crate) async fn usage_cookie_header(&self, account_id: &str) -> Option { self.cookies .cookie_header_for_request(account_id, "chatgpt.com", "/codex/usage") diff --git a/backend/src/admin/accounts/service/exporting.rs b/backend/src/accounts/manage/export.rs similarity index 77% rename from backend/src/admin/accounts/service/exporting.rs rename to backend/src/accounts/manage/export.rs index 298d43591..ce6a25b6b 100644 --- a/backend/src/admin/accounts/service/exporting.rs +++ b/backend/src/accounts/manage/export.rs @@ -1,15 +1,15 @@ use secrecy::ExposeSecret; use super::{ - contracts::{ExportedAccount, ExportedAccounts}, - AdminAccountError, AdminAccountService, + types::{ExportedAccount, ExportedAccounts}, + AccountManageError, AccountManageService, }; -use crate::{infra::time::china_rfc3339_str, upstream::accounts::store::StoredAccount}; +use crate::{accounts::store::StoredAccount, infra::time::china_rfc3339_str}; -impl AdminAccountService { - pub async fn export(&self, ids: Vec) -> Result { +impl AccountManageService { + pub async fn export(&self, ids: Vec) -> Result { if ids.is_empty() { - return Err(AdminAccountError::EmptyIds); + return Err(AccountManageError::EmptyIds); } let accounts = self.export_accounts_by_id(ids).await?; @@ -22,16 +22,16 @@ impl AdminAccountService { async fn export_accounts_by_id( &self, ids: Vec, - ) -> Result, AdminAccountError> { + ) -> Result, AccountManageError> { let mut accounts = Vec::with_capacity(ids.len()); for id in ids { let Some(account) = self .store .get(&id) .await - .map_err(|_| AdminAccountError::Export)? + .map_err(|_| AccountManageError::Export)? else { - return Err(AdminAccountError::NotFound); + return Err(AccountManageError::NotFound); }; accounts.push(account); } diff --git a/backend/src/admin/accounts/service/importing.rs b/backend/src/accounts/manage/import.rs similarity index 71% rename from backend/src/admin/accounts/service/importing.rs rename to backend/src/accounts/manage/import.rs index d36e0615d..dbb0a4acb 100644 --- a/backend/src/admin/accounts/service/importing.rs +++ b/backend/src/accounts/manage/import.rs @@ -2,25 +2,25 @@ use chrono::Utc; use secrecy::SecretString; use super::{ - contracts::{ + types::{ import_quota_plan_type, import_status_from_usage_error, import_usage_plan_type, - import_usage_string, AdminAccountError, ImportSupplementalAccountInfo, + import_usage_string, AccountManageError, ImportSupplementalAccountInfo, ImportSupplementalNeeds, ImportedAccountState, ImportedAccounts, ResolvedImportTokens, }, - AdminAccountService, + AccountManageService, }; -impl AdminAccountService { +impl AccountManageService { pub async fn import( &self, data: serde_json::Value, - ) -> Result { - let parsed = crate::upstream::accounts::importing::parse_account_import_payload(&data) - .map_err(|_| AdminAccountError::NoImportableAccounts)?; + ) -> Result { + let parsed = crate::accounts::import::parse_account_import_payload(&data) + .map_err(|_| AccountManageError::NoImportableAccounts)?; let source_format = parsed.source.as_str(); let entries = parsed.entries; if entries.is_empty() { - return Err(AdminAccountError::NoImportableAccounts); + return Err(AccountManageError::NoImportableAccounts); } let mut imported = 0u32; @@ -43,46 +43,45 @@ impl AdminAccountService { } async fn import_entry( &self, - entry: crate::upstream::accounts::importing::AccountImportEntry, - source: crate::upstream::accounts::importing::AccountImportSource, - ) -> Result { + entry: crate::accounts::import::AccountImportEntry, + source: crate::accounts::import::AccountImportSource, + ) -> Result { let Some(resolved_tokens) = self .resolve_import_tokens(entry.token, entry.refresh_token) .await? else { return Ok(ImportedAccountState::Skipped); }; - let label = crate::upstream::accounts::importing::normalize_label(entry.label); + let label = crate::accounts::import::normalize_label(entry.label); if label .as_ref() .is_some_and(|label| label.chars().count() > 64) { - return Err(AdminAccountError::LabelTooLong); + return Err(AccountManageError::LabelTooLong); } let access_token_expires_at = entry .access_token_expires_at .as_deref() - .map(crate::upstream::accounts::importing::parse_account_import_datetime) + .map(crate::accounts::import::parse_account_import_datetime) .transpose() - .map_err(|_| AdminAccountError::InvalidAccessTokenExpiresAt)?; + .map_err(|_| AccountManageError::InvalidAccessTokenExpiresAt)?; let quota_fetched_at = entry .quota_fetched_at .as_deref() - .map(crate::upstream::accounts::importing::parse_account_import_datetime) + .map(crate::accounts::import::parse_account_import_datetime) .transpose() - .map_err(|_| AdminAccountError::InvalidAccessTokenExpiresAt)?; + .map_err(|_| AccountManageError::InvalidAccessTokenExpiresAt)?; let mut quota_json = entry .cached_quota .as_ref() .map(serde_json::Value::to_string); let mut quota_fetched_at = quota_fetched_at; let quota_verify_required = entry.quota_verify_required.unwrap_or(false); - let parsed_status = crate::upstream::accounts::importing::parse_account_import_status( - entry.status.as_deref(), - ) - .map_err(|error| AdminAccountError::InvalidStatus(error.to_string()))?; - let mut status = crate::upstream::accounts::importing::normalized_imported_account_status( + let parsed_status = + crate::accounts::import::parse_account_import_status(entry.status.as_deref()) + .map_err(|error| AccountManageError::InvalidStatus(error.to_string()))?; + let mut status = crate::accounts::import::normalized_imported_account_status( parsed_status, source, &resolved_tokens.access_token, @@ -95,27 +94,21 @@ impl AdminAccountService { let claims = resolved_tokens.claims.as_ref(); let mut plan_type = claims .and_then(|claims| claims.plan_type.clone()) - .or_else(|| crate::upstream::accounts::importing::normalize_nonempty(entry.plan_type)); + .or_else(|| crate::accounts::import::normalize_nonempty(entry.plan_type)); if plan_type.is_none() { plan_type = entry.cached_quota.as_ref().and_then(import_quota_plan_type); } - let email = claims.and_then(|claims| claims.email.clone()).or_else(|| { - crate::upstream::accounts::importing::normalize_nonempty(entry.email.clone()) - }); + let email = claims + .and_then(|claims| claims.email.clone()) + .or_else(|| crate::accounts::import::normalize_nonempty(entry.email.clone())); let chatgpt_account_id = claims .and_then(|claims| claims.account_id.as_deref()) .or_else(|| { - crate::upstream::accounts::importing::normalize_nonempty_str( - entry.account_id.as_deref(), - ) + crate::accounts::import::normalize_nonempty_str(entry.account_id.as_deref()) }); let chatgpt_user_id = claims .and_then(|claims| claims.user_id.as_deref()) - .or_else(|| { - crate::upstream::accounts::importing::normalize_nonempty_str( - entry.user_id.as_deref(), - ) - }); + .or_else(|| crate::accounts::import::normalize_nonempty_str(entry.user_id.as_deref())); let supplemental = self .import_supplemental_account_info( &resolved_tokens.access_token, @@ -162,7 +155,7 @@ impl AdminAccountService { if let Some(supplemental_status) = supplemental.status { status = supplemental_status; } - let account = crate::upstream::accounts::store::NewAccount { + let account = crate::accounts::store::NewAccount { id: account_id.clone(), email, account_id: chatgpt_account_id, @@ -180,32 +173,32 @@ impl AdminAccountService { Ok(Some(_)) => { let updated = self .store - .update_from_import(crate::upstream::accounts::store::ImportedAccountUpdate { + .update_from_import(crate::accounts::store::ImportedAccountUpdate { account, quota_json, quota_fetched_at, quota_verify_required, }) .await - .map_err(|_| AdminAccountError::Import)?; + .map_err(|_| AccountManageError::Import)?; if !updated { - return Err(AdminAccountError::NotFound); + return Err(AccountManageError::NotFound); } self.store .set_next_refresh_at(&account_id, next_refresh_at) .await - .map_err(|_| AdminAccountError::Import)?; + .map_err(|_| AccountManageError::Import)?; } Ok(None) => { self.store .insert(account) .await - .map_err(|_| AdminAccountError::Import)?; + .map_err(|_| AccountManageError::Import)?; self.store .set_next_refresh_at(&account_id, next_refresh_at) .await - .map_err(|_| AdminAccountError::Import)?; + .map_err(|_| AccountManageError::Import)?; if quota_json.is_some() || quota_fetched_at.is_some() || quota_verify_required { self.store @@ -216,10 +209,10 @@ impl AdminAccountService { quota_verify_required, ) .await - .map_err(|_| AdminAccountError::Import)?; + .map_err(|_| AccountManageError::Import)?; } } - Err(_) => return Err(AdminAccountError::Inspect), + Err(_) => return Err(AccountManageError::Inspect), } Ok(ImportedAccountState::Imported(account_id)) @@ -235,7 +228,7 @@ impl AdminAccountService { } let request_id = uuid::Uuid::new_v4().to_string(); - let context = crate::upstream::transport::CodexRequestContext { + let context = crate::upstream::openai::transport::CodexRequestContext { access_token, account_id, request_id: &request_id, @@ -253,7 +246,7 @@ impl AdminAccountService { match self.codex.fetch_usage(context).await { Ok(raw) => { - let normalized = crate::upstream::accounts::quota::quota_from_usage(&raw); + let normalized = crate::accounts::quota::quota_from_usage(&raw); ImportSupplementalAccountInfo { account_id: import_usage_string(&raw, "account_id"), user_id: import_usage_string(&raw, "user_id"), @@ -280,12 +273,12 @@ impl AdminAccountService { &self, token: Option, refresh_token: Option, - ) -> Result, AdminAccountError> { - let refresh_token = crate::upstream::accounts::importing::normalize_nonempty(refresh_token); - let Some(access_token) = crate::upstream::accounts::importing::normalize_nonempty( + ) -> Result, AccountManageError> { + let refresh_token = crate::accounts::import::normalize_nonempty(refresh_token); + let Some(access_token) = crate::accounts::import::normalize_nonempty( token .as_deref() - .map(crate::upstream::accounts::importing::normalize_bearer_token), + .map(crate::accounts::import::normalize_bearer_token), ) else { let Some(existing_refresh_token) = refresh_token else { return Ok(None); @@ -300,10 +293,9 @@ impl AdminAccountService { })); }; - if let Ok(claims) = crate::upstream::accounts::token_refresh::manual_account_claims( - &access_token, - Utc::now(), - ) { + if let Ok(claims) = + crate::accounts::refresh::manual_account_claims(&access_token, Utc::now()) + { return Ok(Some(ResolvedImportTokens { access_token, refresh_token, @@ -332,32 +324,30 @@ impl AdminAccountService { id: Option<&str>, account_id: Option<&str>, user_id: Option<&str>, - ) -> Result { - let provided_id = crate::upstream::accounts::importing::normalize_nonempty_str(id) - .map(ToString::to_string); + ) -> Result { + let provided_id = + crate::accounts::import::normalize_nonempty_str(id).map(ToString::to_string); if let Some(id) = provided_id.as_deref() { match self.store.get(id).await { Ok(Some(_)) => return Ok(id.to_string()), Ok(None) => {} - Err(_) => return Err(AdminAccountError::Inspect), + Err(_) => return Err(AccountManageError::Inspect), } } - let chatgpt_account_id = - crate::upstream::accounts::importing::normalize_nonempty_str(account_id); - let chatgpt_user_id = crate::upstream::accounts::importing::normalize_nonempty_str(user_id); + let chatgpt_account_id = crate::accounts::import::normalize_nonempty_str(account_id); + let chatgpt_user_id = crate::accounts::import::normalize_nonempty_str(user_id); if let Some(chatgpt_account_id) = chatgpt_account_id { if let Some(existing) = self .store .find_by_chatgpt_identity(chatgpt_account_id, chatgpt_user_id) .await - .map_err(|_| AdminAccountError::Inspect)? + .map_err(|_| AccountManageError::Inspect)? { return Ok(existing.id); } } - Ok(provided_id - .unwrap_or_else(|| crate::upstream::accounts::importing::normalized_account_id(None))) + Ok(provided_id.unwrap_or_else(|| crate::accounts::import::normalized_account_id(None))) } } diff --git a/backend/src/admin/accounts/service/lifecycle.rs b/backend/src/accounts/manage/lifecycle.rs similarity index 70% rename from backend/src/admin/accounts/service/lifecycle.rs rename to backend/src/accounts/manage/lifecycle.rs index 15d37998d..d14d916fa 100644 --- a/backend/src/admin/accounts/service/lifecycle.rs +++ b/backend/src/accounts/manage/lifecycle.rs @@ -10,41 +10,38 @@ use tokio::time::sleep; use uuid::Uuid; use crate::{ + accounts::{account::AccountStatus, store::AccountStore}, infra::{ json::{NumberedPage, Page}, time::elapsed_millis_i64, }, - upstream::accounts::{model::AccountStatus, store::AccountStore}, + upstream::openai::token_client::RefreshFailure, }; use super::{ - contracts::{ - parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountHealthCheck, - AdminAccountMetadata, AdminAccountRefreshOutcome, AdminAccountRefreshResult, - AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens, + types::{ + parse_account_status, stored_to_admin_metadata, AccountHealthCheck, AccountManageError, + AccountRefreshOutcome, AccountRefreshResult, AccountUpdate, BatchDeleteAccounts, + ManagedAccount, ManualCreateTokens, }, - AdminAccountService, + AccountManageService, }; const ADMIN_REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60; -impl AdminAccountService { +impl AccountManageService { pub async fn list( &self, cursor: Option, limit: u32, - ) -> Result, AdminAccountError> { + ) -> Result, AccountManageError> { let page = self .store .list_metadata(cursor, limit) .await - .map_err(|_| AdminAccountError::List)?; + .map_err(|_| AccountManageError::List)?; Ok(Page { - items: page - .items - .into_iter() - .map(AdminAccountMetadata::from) - .collect(), + items: page.items.into_iter().map(ManagedAccount::from).collect(), next_cursor: page.next_cursor, }) } @@ -54,18 +51,14 @@ impl AdminAccountService { page: u32, page_size: u32, search: Option, - ) -> Result, AdminAccountError> { + ) -> Result, AccountManageError> { let page = self .store .list_metadata_page(page, page_size, search.as_deref()) .await - .map_err(|_| AdminAccountError::List)?; + .map_err(|_| AccountManageError::List)?; Ok(NumberedPage { - items: page - .items - .into_iter() - .map(AdminAccountMetadata::from) - .collect(), + items: page.items.into_iter().map(ManagedAccount::from).collect(), total: page.total, page: page.page, page_size: page.page_size, @@ -75,31 +68,26 @@ impl AdminAccountService { pub async fn get( &self, account_id: &str, - ) -> Result, AdminAccountError> { + ) -> Result, AccountManageError> { self.store .get_metadata(account_id) .await - .map(|account| account.map(AdminAccountMetadata::from)) - .map_err(|_| AdminAccountError::Inspect) + .map(|account| account.map(ManagedAccount::from)) + .map_err(|_| AccountManageError::Inspect) } pub async fn create( &self, token: Option, refresh_token: Option, - ) -> Result { - let provided_refresh_token = - crate::upstream::accounts::importing::normalize_nonempty(refresh_token); - let tokens = if let Some(access_token) = - crate::upstream::accounts::importing::normalize_nonempty( - token - .as_deref() - .map(crate::upstream::accounts::importing::normalize_bearer_token), - ) { - let claims = crate::upstream::accounts::token_refresh::manual_account_claims( - &access_token, - Utc::now(), - ) - .map_err(AdminAccountError::InvalidToken)?; + ) -> Result { + let provided_refresh_token = crate::accounts::import::normalize_nonempty(refresh_token); + let tokens = if let Some(access_token) = crate::accounts::import::normalize_nonempty( + token + .as_deref() + .map(crate::accounts::import::normalize_bearer_token), + ) { + let claims = crate::accounts::refresh::manual_account_claims(&access_token, Utc::now()) + .map_err(AccountManageError::InvalidToken)?; ManualCreateTokens { access_token, refresh_token_for_new: provided_refresh_token.clone(), @@ -117,7 +105,7 @@ impl AdminAccountService { claims: refreshed.claims, } } else { - return Err(AdminAccountError::TokenRequired); + return Err(AccountManageError::TokenRequired); }; let claims = tokens.claims; @@ -125,7 +113,7 @@ impl AdminAccountService { self.store .find_by_chatgpt_identity(account_id, claims.user_id.as_deref()) .await - .map_err(|_| AdminAccountError::Inspect)? + .map_err(|_| AccountManageError::Inspect)? } else { None }; @@ -135,7 +123,7 @@ impl AdminAccountService { .store .update_from_claims( &existing.id, - crate::upstream::accounts::store::AccountClaimsUpdate { + crate::accounts::store::AccountClaimsUpdate { email: claims.email.clone(), account_id: claims.account_id.clone(), user_id: claims.user_id.clone(), @@ -152,19 +140,19 @@ impl AdminAccountService { tokens.refresh_token_for_new.is_some() || existing.refresh_token.is_some(), ), - status: crate::upstream::accounts::model::AccountStatus::Active, + status: crate::accounts::account::AccountStatus::Active, }, ) .await - .map_err(|_| AdminAccountError::UpdateClaims)?; + .map_err(|_| AccountManageError::UpdateClaims)?; if !updated { - return Err(AdminAccountError::NotFound); + return Err(AccountManageError::NotFound); } existing.id } else { - let id = crate::upstream::accounts::importing::normalized_account_id(None); + let id = crate::accounts::import::normalized_account_id(None); self.store - .insert(crate::upstream::accounts::store::NewAccount { + .insert(crate::accounts::store::NewAccount { id: id.clone(), email: claims.email.clone(), account_id: claims.account_id.clone(), @@ -176,11 +164,11 @@ impl AdminAccountService { .refresh_token_for_new .map(|token| SecretString::new(token.into())), access_token_expires_at: Some(claims.expires_at), - status: crate::upstream::accounts::model::AccountStatus::Active, + status: crate::accounts::account::AccountStatus::Active, added_at: None, }) .await - .map_err(|_| AdminAccountError::Import)?; + .map_err(|_| AccountManageError::Import)?; id }; @@ -189,15 +177,15 @@ impl AdminAccountService { self.store .get(&account_id) .await - .map_err(|_| AdminAccountError::Inspect)? + .map_err(|_| AccountManageError::Inspect)? .map(stored_to_admin_metadata) - .ok_or(AdminAccountError::NotFound) + .ok_or(AccountManageError::NotFound) } pub async fn update_account( &self, account_id: &str, - update: AdminAccountUpdate, - ) -> Result, AdminAccountError> { + update: AccountUpdate, + ) -> Result, AccountManageError> { if !update.any() { return self.get(account_id).await; } @@ -207,7 +195,7 @@ impl AdminAccountService { .and_then(|label| label.as_ref()) .is_some_and(|label| label.chars().count() > 64) { - return Err(AdminAccountError::LabelTooLong); + return Err(AccountManageError::LabelTooLong); } let status = update @@ -220,7 +208,7 @@ impl AdminAccountService { .store .update_metadata( account_id, - crate::upstream::accounts::store::AccountMetadataUpdate { + crate::accounts::store::AccountMetadataUpdate { email: None, account_id: None, user_id: None, @@ -230,7 +218,7 @@ impl AdminAccountService { }, ) .await - .map_err(|_| AdminAccountError::UpdateMetadata)?; + .map_err(|_| AccountManageError::UpdateMetadata)?; if !updated { return Ok(None); } @@ -245,9 +233,9 @@ impl AdminAccountService { pub async fn batch_delete( &self, ids: Vec, - ) -> Result { + ) -> Result { if ids.is_empty() { - return Err(AdminAccountError::EmptyIds); + return Err(AccountManageError::EmptyIds); } let mut deleted = 0u32; let mut not_found = Vec::new(); @@ -256,9 +244,16 @@ impl AdminAccountService { Ok(true) => { deleted += 1; self.account_pool.remove_account(&id).await; + if let Err(error) = self.session_affinity.forget_account(&id).await { + tracing::warn!( + account_id = id, + error = %error, + "failed to remove Redis affinities for deleted account" + ); + } } Ok(false) => not_found.push(id), - Err(_) => return Err(AdminAccountError::Delete), + Err(_) => return Err(AccountManageError::Delete), } } Ok(BatchDeleteAccounts { deleted, not_found }) @@ -266,20 +261,20 @@ impl AdminAccountService { pub async fn refresh_account( &self, account_id: &str, - ) -> Result { + ) -> Result { let account = match self.store.get(account_id).await { Ok(Some(account)) => account, - Ok(None) => return Err(AdminAccountError::NotFound), - Err(_) => return Err(AdminAccountError::Inspect), + Ok(None) => return Err(AccountManageError::NotFound), + Err(_) => return Err(AccountManageError::Inspect), }; let previous_status = account.status; let started_at = std::time::Instant::now(); - let skipped = |error: &'static str| AdminAccountRefreshResult { + let skipped = |error: &'static str| AccountRefreshResult { id: account.id.clone(), email: account.email.clone(), previous_status, - outcome: AdminAccountRefreshOutcome::Skipped, + outcome: AccountRefreshOutcome::Skipped, error: Some(error.to_string()), duration_ms: elapsed_millis_i64(started_at), }; @@ -314,7 +309,7 @@ impl AdminAccountService { lease_now, ) .await - .map_err(|_| AdminAccountError::RefreshLease)?; + .map_err(|_| AccountManageError::RefreshLease)?; if !lease_acquired { return Ok(skipped("refresh already in progress")); } @@ -322,7 +317,7 @@ impl AdminAccountService { let refresh_result = self.refresh_tokens_from_refresh_token(&refresh_token).await; let release_result = self.refresh_leases.release(account_id, &lease_owner).await; if release_result.is_err() { - return Err(AdminAccountError::RefreshLease); + return Err(AccountManageError::RefreshLease); } match refresh_result { @@ -331,7 +326,7 @@ impl AdminAccountService { .store .update_from_claims( account_id, - crate::upstream::accounts::store::AccountClaimsUpdate { + crate::accounts::store::AccountClaimsUpdate { email: refreshed.claims.email, account_id: refreshed.claims.account_id.or(account.account_id), user_id: refreshed.claims.user_id, @@ -345,55 +340,55 @@ impl AdminAccountService { account_id, refreshed.claims.expires_at, )), - status: crate::upstream::accounts::model::AccountStatus::Active, + status: account.status, }, ) .await - .map_err(|_| AdminAccountError::UpdateClaims)?; + .map_err(|_| AccountManageError::UpdateClaims)?; if !updated { - return Err(AdminAccountError::NotFound); + return Err(AccountManageError::NotFound); } self.sync_account_pool(account_id).await?; - Ok(AdminAccountRefreshResult { + Ok(AccountRefreshResult { id: account.id, email: account.email, previous_status, - outcome: AdminAccountRefreshOutcome::Alive, + outcome: AccountRefreshOutcome::Alive, error: None, duration_ms: elapsed_millis_i64(started_at), }) } - Err(AdminAccountError::RefreshTokenExchange(failure)) => { + Err(AccountManageError::RefreshTokenExchange(failure)) => { if manual_refresh_failure_is_permanent(&failure) { - let status = crate::upstream::accounts::model::AccountStatus::Expired; + let status = crate::accounts::account::AccountStatus::Expired; let updated = self .store .set_status(account_id, status) .await - .map_err(|_| AdminAccountError::UpdateStatus)?; + .map_err(|_| AccountManageError::UpdateStatus)?; if !updated { - return Err(AdminAccountError::NotFound); + return Err(AccountManageError::NotFound); } - if crate::upstream::accounts::importing::refresh_failure_status_clears_next_refresh_at( + if crate::accounts::import::refresh_failure_status_clears_next_refresh_at( status, ) { let cleared = self .store .set_next_refresh_at(account_id, None) .await - .map_err(|_| AdminAccountError::UpdateStatus)?; + .map_err(|_| AccountManageError::UpdateStatus)?; if !cleared { - return Err(AdminAccountError::NotFound); + return Err(AccountManageError::NotFound); } } self.sync_account_pool_best_effort(account_id, "account refresh failure") .await; } - Ok(AdminAccountRefreshResult { + Ok(AccountRefreshResult { id: account.id, email: account.email, previous_status, - outcome: AdminAccountRefreshOutcome::Dead, + outcome: AccountRefreshOutcome::Dead, error: Some(format!("token refresh exchange failed: {failure}")), duration_ms: elapsed_millis_i64(started_at), }) @@ -407,12 +402,12 @@ impl AdminAccountService { ids: Option>, stagger_ms: u64, concurrency: usize, - ) -> Result { + ) -> Result { let accounts = self .store .list_pool_accounts() .await - .map_err(|_| AdminAccountError::List)?; + .map_err(|_| AccountManageError::List)?; let ids = ids.map(|ids| ids.into_iter().collect::>()); let candidate_ids = accounts .into_iter() @@ -437,12 +432,12 @@ impl AdminAccountService { .into_iter() .collect::, _>>()?; - Ok(AdminAccountHealthCheck { results }) + Ok(AccountHealthCheck { results }) } } fn create_next_refresh_at( - service: &AdminAccountService, + service: &AccountManageService, account_id: &str, expires_at: DateTime, has_refresh_token: bool, @@ -465,13 +460,9 @@ fn stable_jittered_millis(account_id: &str, base_millis: u64, variance: f64) -> .clamp(0.0, u64::MAX as f64) as u64 } -fn manual_refresh_failure_is_permanent( - failure: &crate::upstream::accounts::token_refresh::RefreshFailure, -) -> bool { +fn manual_refresh_failure_is_permanent(failure: &RefreshFailure) -> bool { match failure { - crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant - | crate::upstream::accounts::token_refresh::RefreshFailure::Banned => true, - crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport - | crate::upstream::accounts::token_refresh::RefreshFailure::Transport => false, + RefreshFailure::InvalidGrant | RefreshFailure::Banned => true, + RefreshFailure::RetryableTransport | RefreshFailure::Transport => false, } } diff --git a/backend/src/accounts/manage/mod.rs b/backend/src/accounts/manage/mod.rs new file mode 100644 index 000000000..d41046aae --- /dev/null +++ b/backend/src/accounts/manage/mod.rs @@ -0,0 +1,15 @@ +//! 账号管理用例。 + +mod cookies; +mod export; +mod import; +mod lifecycle; +pub(crate) mod oauth; +mod probe; +mod quota; +pub mod quota_view; +mod service; +mod types; + +pub use service::*; +pub use types::*; diff --git a/backend/src/admin/accounts/service/oauth.rs b/backend/src/accounts/manage/oauth.rs similarity index 81% rename from backend/src/admin/accounts/service/oauth.rs rename to backend/src/accounts/manage/oauth.rs index b176ee3c5..545f9e1e1 100644 --- a/backend/src/admin/accounts/service/oauth.rs +++ b/backend/src/accounts/manage/oauth.rs @@ -9,8 +9,8 @@ use sha2::{Digest, Sha256}; use tokio::sync::Mutex; use super::{ - contracts::{AdminAccountError, OAuthAuthorizeResult, OAuthExchangeInput}, - AdminAccountService, + types::{AccountManageError, OAuthAuthorizeResult, OAuthExchangeInput}, + AccountManageService, }; const OPENAI_OAUTH_AUTHORIZE_ENDPOINT: &str = "https://auth.openai.com/oauth/authorize"; @@ -36,7 +36,7 @@ impl AccountOAuthService { } } - async fn authorize(&self) -> Result { + async fn authorize(&self) -> Result { let now = Utc::now(); let expires_at = now + Duration::minutes(OAUTH_SESSION_TTL_MINUTES); let session_id = random_url_token(32); @@ -55,7 +55,7 @@ impl AccountOAuthService { drop(sessions); let mut auth_url = - Url::parse(OPENAI_OAUTH_AUTHORIZE_ENDPOINT).map_err(|_| AdminAccountError::Import)?; + Url::parse(OPENAI_OAUTH_AUTHORIZE_ENDPOINT).map_err(|_| AccountManageError::Import)?; auth_url .query_pairs_mut() .append_pair("response_type", "code") @@ -78,11 +78,11 @@ impl AccountOAuthService { async fn exchange( &self, input: OAuthExchangeInput, - ) -> Result { + ) -> Result { let (code, state) = oauth_callback_parts(&input)?; let session = self.session(&input.session_id).await?; if session.state != state { - return Err(AdminAccountError::OAuthStateMismatch); + return Err(AccountManageError::OAuthStateMismatch); } let response = self @@ -97,22 +97,22 @@ impl AccountOAuthService { ]) .send() .await - .map_err(|error| AdminAccountError::OAuthCodeExchange(error.to_string()))?; + .map_err(|error| AccountManageError::OAuthCodeExchange(error.to_string()))?; let status = response.status(); let body = response .text() .await - .map_err(|error| AdminAccountError::OAuthCodeExchange(error.to_string()))?; + .map_err(|error| AccountManageError::OAuthCodeExchange(error.to_string()))?; if !status.is_success() { - return Err(AdminAccountError::OAuthCodeExchange(exchange_error_text( + return Err(AccountManageError::OAuthCodeExchange(exchange_error_text( status, &body, ))); } let tokens = serde_json::from_str::(&body) - .map_err(|_| AdminAccountError::OAuthCodeExchange("invalid token response".into()))?; + .map_err(|_| AccountManageError::OAuthCodeExchange("invalid token response".into()))?; if tokens.access_token.trim().is_empty() { - return Err(AdminAccountError::OAuthCodeExchange( + return Err(AccountManageError::OAuthCodeExchange( "empty access token".into(), )); } @@ -121,14 +121,14 @@ impl AccountOAuthService { Ok(tokens) } - async fn session(&self, session_id: &str) -> Result { + async fn session(&self, session_id: &str) -> Result { let now = Utc::now(); let mut sessions = self.sessions.lock().await; sessions.retain(|_, session| session.expires_at > now); sessions .get(session_id) .cloned() - .ok_or(AdminAccountError::OAuthSessionInvalid) + .ok_or(AccountManageError::OAuthSessionInvalid) } async fn remove_session(&self, session_id: &str) { @@ -136,15 +136,15 @@ impl AccountOAuthService { } } -impl AdminAccountService { - pub async fn oauth_authorize(&self) -> Result { +impl AccountManageService { + pub async fn oauth_authorize(&self) -> Result { self.oauth.authorize().await } pub async fn oauth_exchange( &self, input: OAuthExchangeInput, - ) -> Result { + ) -> Result { let tokens = self.oauth.exchange(input).await?; self.import(serde_json::json!({ "sourceFormat": "cpr", @@ -170,21 +170,23 @@ struct AccountOAuthTokens { refresh_token: Option, } -fn oauth_callback_parts(input: &OAuthExchangeInput) -> Result<(String, String), AdminAccountError> { +fn oauth_callback_parts( + input: &OAuthExchangeInput, +) -> Result<(String, String), AccountManageError> { if let Some(callback_url) = input.callback_url.as_deref() { let (code, state) = callback_url_parts(callback_url)?; return Ok((code, state)); } - let code = normalized(input.code.as_deref()).ok_or(AdminAccountError::OAuthCallbackInvalid)?; + let code = normalized(input.code.as_deref()).ok_or(AccountManageError::OAuthCallbackInvalid)?; let state = - normalized(input.state.as_deref()).ok_or(AdminAccountError::OAuthCallbackInvalid)?; + normalized(input.state.as_deref()).ok_or(AccountManageError::OAuthCallbackInvalid)?; Ok((code, state)) } -fn callback_url_parts(callback_url: &str) -> Result<(String, String), AdminAccountError> { +fn callback_url_parts(callback_url: &str) -> Result<(String, String), AccountManageError> { let url = - Url::parse(callback_url.trim()).map_err(|_| AdminAccountError::OAuthCallbackInvalid)?; + Url::parse(callback_url.trim()).map_err(|_| AccountManageError::OAuthCallbackInvalid)?; let mut code = None; let mut state = None; for (key, value) in url.query_pairs() { @@ -195,7 +197,7 @@ fn callback_url_parts(callback_url: &str) -> Result<(String, String), AdminAccou } } code.zip(state) - .ok_or(AdminAccountError::OAuthCallbackInvalid) + .ok_or(AccountManageError::OAuthCallbackInvalid) } fn normalized(value: Option<&str>) -> Option { diff --git a/backend/src/admin/accounts/service/testing.rs b/backend/src/accounts/manage/probe.rs similarity index 94% rename from backend/src/admin/accounts/service/testing.rs rename to backend/src/accounts/manage/probe.rs index 23e88caaa..6a31123cf 100644 --- a/backend/src/admin/accounts/service/testing.rs +++ b/backend/src/accounts/manage/probe.rs @@ -7,18 +7,18 @@ use serde_json::{json, Value}; use tokio::sync::mpsc; use crate::{ - proxy::dispatch::{ + accounts::{ + account::{Account, AccountStatus}, + store::StoredAccount, + }, + dispatch::{ errors::{ auth_failure_account_status, is_auth_upstream_error, is_quota_exhausted_upstream_error, }, responses::sse_failure::sse_failure_account_status, }, - upstream::{ - accounts::{ - model::{Account, AccountStatus}, - store::StoredAccount, - }, - models::{info::CodexModelInfo, service::ModelRefreshPlanAccount}, + models::{info::CodexModelInfo, service::ModelRefreshPlanAccount}, + upstream::openai::{ protocol::{ responses::{CodexResponsesRequest, ResponsesSseFailure}, sse::{encode_sse_event, parse_sse_events, SseEvent}, @@ -27,7 +27,7 @@ use crate::{ }, }; -use super::{contracts::AdminAccountError, AdminAccountService}; +use super::{types::AccountManageError, AccountManageService}; pub(super) type AccountTestStream = Pin> + Send>>; @@ -59,17 +59,17 @@ impl AccountTestOutcome { } } -impl AdminAccountService { +impl AccountManageService { pub async fn account_models( &self, account_id: &str, - ) -> Result, AdminAccountError> { + ) -> Result, AccountManageError> { let account = self .store .get(account_id) .await - .map_err(|_| AdminAccountError::Inspect)? - .ok_or(AdminAccountError::NotFound)?; + .map_err(|_| AccountManageError::Inspect)? + .ok_or(AccountManageError::NotFound)?; let plan_type = account_plan_type(&account); let mut models = self.models.catalog().await.models_for_plan(&plan_type); if models.is_empty() { @@ -78,7 +78,7 @@ impl AdminAccountService { } let models = models.iter().map(account_model_option).collect::>(); if models.is_empty() { - return Err(AdminAccountError::NoModels); + return Err(AccountManageError::NoModels); } Ok(models) } @@ -111,13 +111,13 @@ impl AdminAccountService { &self, account_id: &str, model: String, - ) -> Result { + ) -> Result { let account = self .store .get(account_id) .await - .map_err(|_| AdminAccountError::Inspect)? - .ok_or(AdminAccountError::NotFound)?; + .map_err(|_| AccountManageError::Inspect)? + .ok_or(AccountManageError::NotFound)?; let token = account.access_token.expose_secret().to_string(); let upstream_account_id = account.account_id.clone(); @@ -317,7 +317,7 @@ fn test_responses_request(model: String) -> CodexResponsesRequest { } async fn process_upstream_test_stream( - mut body: crate::upstream::transport::CodexBackendSseStream, + mut body: crate::upstream::openai::transport::CodexBackendSseStream, tx: &mpsc::Sender, ) -> AccountTestOutcome { let mut buffer = String::new(); diff --git a/backend/src/admin/accounts/service/quota.rs b/backend/src/accounts/manage/quota.rs similarity index 72% rename from backend/src/admin/accounts/service/quota.rs rename to backend/src/accounts/manage/quota.rs index e2a823b79..6dfeab9cc 100644 --- a/backend/src/admin/accounts/service/quota.rs +++ b/backend/src/accounts/manage/quota.rs @@ -1,23 +1,23 @@ use secrecy::ExposeSecret; -use super::{contracts::AdminAccountError, AdminAccountService}; +use super::{types::AccountManageError, AccountManageService}; -impl AdminAccountService { +impl AccountManageService { pub async fn account_quota( &self, account_id: &str, - ) -> Result { + ) -> Result { let stored = self .store .get(account_id) .await - .map_err(|_| AdminAccountError::NotFound)? - .ok_or(AdminAccountError::NotFound)?; + .map_err(|_| AccountManageError::NotFound)? + .ok_or(AccountManageError::NotFound)?; let request_id = uuid::Uuid::new_v4().to_string(); let token = stored.access_token.expose_secret().to_string(); let cookie_header = self.usage_cookie_header(account_id).await; - let context = crate::upstream::transport::CodexRequestContext { + let context = crate::upstream::openai::transport::CodexRequestContext { access_token: &token, account_id: stored.account_id.as_deref(), request_id: &request_id, @@ -36,8 +36,8 @@ impl AdminAccountService { .codex .fetch_usage(context) .await - .map_err(|e| AdminAccountError::FetchQuota(e.to_string()))?; - let normalized = crate::upstream::accounts::quota::quota_from_usage(&raw); + .map_err(|e| AccountManageError::FetchQuota(e.to_string()))?; + let normalized = crate::accounts::quota::quota_from_usage(&raw); self.account_pool .apply_quota_snapshot(account_id, &normalized) .await; diff --git a/backend/src/admin/accounts/quota_view.rs b/backend/src/accounts/manage/quota_view.rs similarity index 87% rename from backend/src/admin/accounts/quota_view.rs rename to backend/src/accounts/manage/quota_view.rs index 8a008d855..33ef5759f 100644 --- a/backend/src/admin/accounts/quota_view.rs +++ b/backend/src/accounts/manage/quota_view.rs @@ -17,14 +17,14 @@ const MONTH_WINDOW_SECONDS: u64 = 2_592_000; #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -pub(crate) struct AdminAccountQuotaData { +pub(crate) struct AccountQuotaData { refreshed_at_display: String, - windows: Vec, + windows: Vec, } #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -struct AdminAccountQuotaWindowData { +struct AccountQuotaWindowData { key: String, group: String, window_seconds: Option, @@ -32,7 +32,7 @@ struct AdminAccountQuotaWindowData { used_percent: Option, used_percent_display: String, #[serde(skip_serializing_if = "Option::is_none")] - local_usage: Option, + local_usage: Option, reset_at_display: String, window_used_display: String, #[serde(skip)] @@ -41,7 +41,7 @@ struct AdminAccountQuotaWindowData { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -struct AdminAccountQuotaWindowLocalUsageData { +struct AccountQuotaWindowLocalUsageData { request_count: i64, request_count_display: String, input_tokens: i64, @@ -55,14 +55,14 @@ struct AdminAccountQuotaWindowLocalUsageData { } #[derive(Debug, Clone)] -pub(crate) struct AdminAccountQuotaUsageWindow { +pub(crate) struct AccountQuotaUsageWindow { pub(crate) key: String, pub(crate) start: DateTime, pub(crate) end: DateTime, pub(crate) window_seconds: u64, } -impl AdminAccountQuotaUsageWindow { +impl AccountQuotaUsageWindow { pub(crate) fn duration_seconds(&self) -> u64 { self.end .signed_duration_since(self.start) @@ -72,24 +72,15 @@ impl AdminAccountQuotaUsageWindow { } #[derive(Debug, Clone, Copy, Default)] -pub(crate) struct AdminAccountQuotaWindowLocalUsage { +pub(crate) struct AccountQuotaWindowLocalUsage { pub(crate) request_count: i64, pub(crate) input_tokens: i64, pub(crate) output_tokens: i64, pub(crate) cached_tokens: i64, } -impl AdminAccountQuotaWindowLocalUsage { - pub(crate) fn add(&mut self, other: Self) { - self.request_count = self.request_count.saturating_add(other.request_count); - self.input_tokens = self.input_tokens.saturating_add(other.input_tokens); - self.output_tokens = self.output_tokens.saturating_add(other.output_tokens); - self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens); - } -} - -impl From for AdminAccountQuotaWindowLocalUsageData { - fn from(usage: AdminAccountQuotaWindowLocalUsage) -> Self { +impl From for AccountQuotaWindowLocalUsageData { + fn from(usage: AccountQuotaWindowLocalUsage) -> Self { let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens); Self { request_count: usage.request_count, @@ -106,7 +97,7 @@ impl From for AdminAccountQuotaWindowLocalUsa } } -impl Default for AdminAccountQuotaData { +impl Default for AccountQuotaData { fn default() -> Self { Self { refreshed_at_display: "-".to_string(), @@ -115,23 +106,23 @@ impl Default for AdminAccountQuotaData { } } -impl AdminAccountQuotaData { +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 { + pub(crate) fn usage_windows(&self) -> Vec { self.windows .iter() - .filter_map(AdminAccountQuotaWindowData::usage_window) + .filter_map(AccountQuotaWindowData::usage_window) .collect() } pub(crate) fn apply_local_usage( &mut self, - usage_by_window: &HashMap, + usage_by_window: &HashMap, ) { for window in &mut self.windows { window.local_usage = window.usage_window().map(|_| { @@ -145,13 +136,13 @@ impl AdminAccountQuotaData { } } -impl AdminAccountQuotaWindowData { - fn usage_window(&self) -> Option { +impl AccountQuotaWindowData { + fn usage_window(&self) -> Option { let end = self.reset_at?; let window_seconds = self.window_seconds?; let seconds = i64::try_from(window_seconds).ok()?; let start = end.checked_sub_signed(Duration::seconds(seconds))?; - (start <= end).then(|| AdminAccountQuotaUsageWindow { + (start <= end).then(|| AccountQuotaUsageWindow { key: self.key.clone(), start, end, @@ -160,20 +151,17 @@ impl AdminAccountQuotaWindowData { } } -pub(crate) fn quota_data( - quota_json: &str, - fetched_at: Option>, -) -> AdminAccountQuotaData { +pub(crate) fn quota_data(quota_json: &str, fetched_at: Option>) -> AccountQuotaData { let quota = serde_json::from_str::(quota_json).unwrap_or(Value::Null); let windows = quota_windows("a); - AdminAccountQuotaData { + AccountQuotaData { refreshed_at_display: china_relative_time(fetched_at, Utc::now()), windows, } } -fn quota_windows(quota: &Value) -> Vec { +fn quota_windows(quota: &Value) -> Vec { let mut windows = Vec::new(); let has_monthly_limit = push_monthly_quota_window(&mut windows, quota.get("monthly_limit")); @@ -188,7 +176,7 @@ fn quota_windows(quota: &Value) -> Vec { } fn push_monthly_quota_window( - windows: &mut Vec, + windows: &mut Vec, monthly_limit: Option<&Value>, ) -> bool { let Some(monthly_limit) = monthly_limit.filter(|value| !value.is_null()) else { @@ -215,7 +203,7 @@ fn push_monthly_quota_window( .get("key") .and_then(Value::as_str) .unwrap_or("monthly"); - windows.push(AdminAccountQuotaWindowData { + windows.push(AccountQuotaWindowData { key: quota_key_segment(key), group: "monthly".to_string(), window_seconds, @@ -233,7 +221,7 @@ fn push_monthly_quota_window( } fn push_snapshot_quota_windows( - windows: &mut Vec, + windows: &mut Vec, snapshot: &Value, skip_core_monthly: bool, ) { @@ -269,7 +257,7 @@ fn push_snapshot_quota_windows( } fn push_quota_window( - windows: &mut Vec, + windows: &mut Vec, source_key: &str, role: &str, label_prefix: Option<&str>, @@ -304,7 +292,7 @@ fn push_quota_window( .map(|value| format!("{value} · {base_label}")) .unwrap_or(base_label); - windows.push(AdminAccountQuotaWindowData { + windows.push(AccountQuotaWindowData { key: unique_quota_window_key(windows, source_key, role, window_seconds), group: quota_window_group(window_seconds).to_string(), window_seconds, @@ -350,7 +338,7 @@ fn snapshot_label(snapshot: &Value) -> Option { } fn unique_quota_window_key( - windows: &[AdminAccountQuotaWindowData], + windows: &[AccountQuotaWindowData], source_key: &str, role: &str, window_seconds: Option, @@ -392,7 +380,7 @@ fn quota_key_segment(value: &str) -> String { } } -fn quota_window_sort_key(window: &AdminAccountQuotaWindowData) -> (u8, u64, String) { +fn quota_window_sort_key(window: &AccountQuotaWindowData) -> (u8, u64, String) { let group_order = match window.group.as_str() { "monthly" => 0, "shortTerm" => 1, diff --git a/backend/src/admin/accounts/service/mod.rs b/backend/src/accounts/manage/service.rs similarity index 54% rename from backend/src/admin/accounts/service/mod.rs rename to backend/src/accounts/manage/service.rs index 1130c9688..28c2879e3 100644 --- a/backend/src/admin/accounts/service/mod.rs +++ b/backend/src/accounts/manage/service.rs @@ -1,66 +1,58 @@ //! 管理端账号业务服务。 -mod contracts; -mod cookies; -mod exporting; -mod importing; -mod lifecycle; -pub(crate) mod oauth; -mod quota; -mod testing; - -use std::sync::Arc as StdArc; +use std::sync::Arc; use chrono::{DateTime, Utc}; use uuid::Uuid; +use super::oauth; + use crate::{ - upstream::accounts::{ - cookies::SqliteCookieStore, + accounts::{ + cookies::PgCookieStore, pool::RuntimeAccountPoolService, - store::SqliteAccountStore, - token_refresh::{RefreshLeaseStore, RuntimeRefreshPolicy, TokenRefresher}, + refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy}, + store::PgAccountStore, }, - upstream::models::service::ModelService, - upstream::transport::CodexBackendClient, + dispatch::affinity::RedisSessionAffinityStore, + models::service::ModelService, + upstream::openai::{token_client::TokenRefresher, transport::CodexBackendClient}, }; -pub use contracts::{ - AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, AdminAccountRefreshOutcome, - AdminAccountRefreshResult, AdminAccountUpdate, BatchDeleteAccounts, ExportedAccounts, - ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, -}; +use super::types::{AccountManageError, RefreshedAccountTokens}; #[derive(Clone)] -pub struct AdminAccountService { - pub store: SqliteAccountStore, - pub(crate) cookies: SqliteCookieStore, - pub(crate) codex: StdArc, - pub(crate) models: StdArc, - pub(crate) account_pool: StdArc, - pub(crate) token_refresher: StdArc, - pub(crate) refresh_leases: RefreshLeaseStore, +pub struct AccountManageService { + pub store: PgAccountStore, + pub(crate) cookies: PgCookieStore, + pub(crate) codex: Arc, + pub(crate) models: Arc, + pub(crate) account_pool: Arc, + pub(crate) token_refresher: Arc, + pub(crate) refresh_leases: RedisRefreshLeaseStore, + pub(crate) session_affinity: RedisSessionAffinityStore, pub(crate) refresh_lease_owner_prefix: String, pub(crate) oauth: oauth::AccountOAuthService, pub(crate) refresh_policy: RuntimeRefreshPolicy, pub(crate) installation_id: Option, } -pub struct AdminAccountServiceParts { - pub store: SqliteAccountStore, - pub cookies: SqliteCookieStore, - pub codex: StdArc, - pub models: StdArc, - pub account_pool: StdArc, - pub token_refresher: StdArc, - pub refresh_leases: RefreshLeaseStore, +pub struct AccountManageServiceParts { + pub store: PgAccountStore, + pub cookies: PgCookieStore, + pub codex: Arc, + pub models: Arc, + pub account_pool: Arc, + pub token_refresher: Arc, + pub refresh_leases: RedisRefreshLeaseStore, + pub session_affinity: RedisSessionAffinityStore, pub oauth: oauth::AccountOAuthService, pub refresh_policy: RuntimeRefreshPolicy, pub installation_id: Option, } -impl AdminAccountService { - pub fn new(parts: AdminAccountServiceParts) -> Self { +impl AccountManageService { + pub fn new(parts: AccountManageServiceParts) -> Self { Self { store: parts.store, cookies: parts.cookies, @@ -69,6 +61,7 @@ impl AdminAccountService { account_pool: parts.account_pool, token_refresher: parts.token_refresher, refresh_leases: parts.refresh_leases, + session_affinity: parts.session_affinity, refresh_lease_owner_prefix: format!( "admin-account-refresh:{}", Uuid::new_v4().simple() @@ -84,33 +77,30 @@ impl AdminAccountService { account_id: &str, expires_at: DateTime, ) -> DateTime { - crate::upstream::accounts::token_refresh::jittered_refresh_at( + crate::accounts::refresh::jittered_refresh_at( account_id, expires_at, self.refresh_policy.refresh_margin_seconds(), ) } - async fn refresh_tokens_from_refresh_token( + pub(super) async fn refresh_tokens_from_refresh_token( &self, refresh_token: &str, - ) -> Result { + ) -> Result { let token_pair = self .token_refresher .refresh(refresh_token) .await - .map_err(AdminAccountError::RefreshTokenExchange)?; - let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( - crate::upstream::accounts::importing::normalize_bearer_token(&token_pair.access_token), + .map_err(AccountManageError::RefreshTokenExchange)?; + let access_token = crate::accounts::import::normalize_nonempty(Some( + crate::accounts::import::normalize_bearer_token(&token_pair.access_token), )) - .ok_or(AdminAccountError::TokenRequired)?; - let claims = crate::upstream::accounts::token_refresh::manual_account_claims( - &access_token, - Utc::now(), - ) - .map_err(AdminAccountError::InvalidToken)?; + .ok_or(AccountManageError::TokenRequired)?; + let claims = crate::accounts::refresh::manual_account_claims(&access_token, Utc::now()) + .map_err(AccountManageError::InvalidToken)?; - Ok(contracts::RefreshedAccountTokens { + Ok(RefreshedAccountTokens { access_token, refresh_token: token_pair.refresh_token, claims, @@ -120,20 +110,16 @@ impl AdminAccountService { pub(crate) async fn sync_account_pool( &self, account_id: &str, - ) -> Result<(), AdminAccountError> { + ) -> Result<(), AccountManageError> { self.account_pool - .sync_account_from_repository(account_id) + .sync_account_from_store(account_id) .await .map(|_| ()) - .map_err(|_| AdminAccountError::SyncAccountPool) + .map_err(|_| AccountManageError::SyncAccountPool) } pub(crate) async fn sync_account_pool_best_effort(&self, account_id: &str, operation: &str) { - if let Err(error) = self - .account_pool - .sync_account_from_repository(account_id) - .await - { + if let Err(error) = self.account_pool.sync_account_from_store(account_id).await { tracing::warn!( account_id, operation, diff --git a/backend/src/admin/accounts/service/contracts.rs b/backend/src/accounts/manage/types.rs similarity index 83% rename from backend/src/admin/accounts/service/contracts.rs rename to backend/src/accounts/manage/types.rs index 316b46be7..c8f230f97 100644 --- a/backend/src/admin/accounts/service/contracts.rs +++ b/backend/src/accounts/manage/types.rs @@ -5,15 +5,15 @@ use serde::Serialize; use thiserror::Error; use crate::{ - upstream::accounts::{ - model::{AccountStatus, AccountStatus as AcctStatus}, + accounts::{ + account::{AccountStatus, AccountStatus as AcctStatus}, store::StoredAccountMetadata, }, - upstream::transport::CodexClientError, + upstream::openai::{token_client::RefreshFailure, transport::CodexClientError}, }; #[derive(Debug, Clone, serde::Serialize)] -pub struct AdminAccountMetadata { +pub struct ManagedAccount { pub id: String, pub email: Option, pub account_id: Option, @@ -27,7 +27,7 @@ pub struct AdminAccountMetadata { pub updated_at: DateTime, } -impl From for AdminAccountMetadata { +impl From for ManagedAccount { fn from(m: StoredAccountMetadata) -> Self { let added_at = m .added_at @@ -54,7 +54,7 @@ impl From for AdminAccountMetadata { } #[derive(Debug, Error)] -pub enum AdminAccountError { +pub enum AccountManageError { #[error("failed to list accounts")] List, #[error("failed to import accounts")] @@ -100,7 +100,7 @@ pub enum AdminAccountError { #[error("invalid token: {0}")] InvalidToken(&'static str), #[error("token refresh exchange failed: {0}")] - RefreshTokenExchange(crate::upstream::accounts::token_refresh::RefreshFailure), + RefreshTokenExchange(RefreshFailure), #[error("oauth session is invalid or expired")] OAuthSessionInvalid, #[error("oauth callback is missing code or state")] @@ -116,57 +116,57 @@ pub enum AdminAccountError { } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum AdminAccountRefreshOutcome { +pub enum AccountRefreshOutcome { Alive, Dead, Skipped, } #[derive(Debug, Clone)] -pub struct AdminAccountRefreshResult { +pub struct AccountRefreshResult { pub id: String, pub email: Option, pub previous_status: AcctStatus, - pub outcome: AdminAccountRefreshOutcome, + pub outcome: AccountRefreshOutcome, pub error: Option, pub duration_ms: i64, } #[derive(Debug, Clone, Default)] -pub struct AdminAccountHealthCheck { - pub results: Vec, +pub struct AccountHealthCheck { + pub results: Vec, } -impl AdminAccountHealthCheck { +impl AccountHealthCheck { pub fn alive(&self) -> usize { self.results .iter() - .filter(|result| result.outcome == AdminAccountRefreshOutcome::Alive) + .filter(|result| result.outcome == AccountRefreshOutcome::Alive) .count() } pub fn dead(&self) -> usize { self.results .iter() - .filter(|result| result.outcome == AdminAccountRefreshOutcome::Dead) + .filter(|result| result.outcome == AccountRefreshOutcome::Dead) .count() } pub fn skipped(&self) -> usize { self.results .iter() - .filter(|result| result.outcome == AdminAccountRefreshOutcome::Skipped) + .filter(|result| result.outcome == AccountRefreshOutcome::Skipped) .count() } } #[derive(Debug, Clone, Default)] -pub struct AdminAccountUpdate { +pub struct AccountUpdate { pub label: Option>, pub status: Option, } -impl AdminAccountUpdate { +impl AccountUpdate { pub fn any(&self) -> bool { self.label.is_some() || self.status.is_some() } @@ -229,21 +229,21 @@ pub(super) struct ManualCreateTokens { pub(super) access_token: String, pub(super) refresh_token_for_new: Option, pub(super) refresh_token_for_existing: Option, - pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims, + pub(super) claims: crate::accounts::refresh::ManualAccountClaims, } #[derive(Debug, Clone)] pub(super) struct RefreshedAccountTokens { pub(super) access_token: String, pub(super) refresh_token: Option, - pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims, + pub(super) claims: crate::accounts::refresh::ManualAccountClaims, } #[derive(Debug, Clone)] pub(super) struct ResolvedImportTokens { pub(super) access_token: String, pub(super) refresh_token: Option, - pub(super) claims: Option, + pub(super) claims: Option, } #[derive(Debug, Clone, Default)] @@ -278,10 +278,8 @@ pub(super) enum ImportedAccountState { Skipped, } -pub(super) fn stored_to_admin_metadata( - s: crate::upstream::accounts::store::StoredAccount, -) -> AdminAccountMetadata { - AdminAccountMetadata::from(crate::upstream::accounts::store::StoredAccountMetadata { +pub(super) fn stored_to_admin_metadata(s: crate::accounts::store::StoredAccount) -> ManagedAccount { + ManagedAccount::from(crate::accounts::store::StoredAccountMetadata { id: s.id, email: s.email, account_id: s.account_id, @@ -307,7 +305,7 @@ pub(super) fn import_usage_string(usage: &serde_json::Value, key: &str) -> Optio usage .get(key) .and_then(serde_json::Value::as_str) - .and_then(|value| crate::upstream::accounts::importing::normalize_nonempty_str(Some(value))) + .and_then(|value| crate::accounts::import::normalize_nonempty_str(Some(value))) .map(ToString::to_string) } @@ -324,18 +322,18 @@ fn normalized_plan_type(value: &str) -> Option { } pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option { - if crate::upstream::transport::is_banned_upstream_error(error) { + if crate::upstream::openai::transport::is_banned_upstream_error(error) { Some(AccountStatus::Banned) } else { None } } -pub(super) fn parse_account_status(status: &str) -> Result { +pub(super) fn parse_account_status(status: &str) -> Result { match status.trim().to_ascii_lowercase().as_str() { "active" => Ok(AcctStatus::Active), "disabled" => Ok(AcctStatus::Disabled), - _ => Err(AdminAccountError::InvalidStatus( + _ => Err(AccountManageError::InvalidStatus( status.trim().to_ascii_lowercase(), )), } diff --git a/backend/src/accounts/mod.rs b/backend/src/accounts/mod.rs new file mode 100644 index 000000000..e72bb1666 --- /dev/null +++ b/backend/src/accounts/mod.rs @@ -0,0 +1,13 @@ +//! 账号、调度、配额与刷新领域。 + +pub mod account; +pub mod cookies; +pub mod import; +pub mod manage; +pub mod pool; +pub mod quota; +pub mod refresh; +pub mod scheduler; +pub mod service; +pub mod store; +mod window; diff --git a/backend/src/upstream/accounts/pool.rs b/backend/src/accounts/pool.rs similarity index 53% rename from backend/src/upstream/accounts/pool.rs rename to backend/src/accounts/pool.rs index 4b410d63f..6703b7d64 100644 --- a/backend/src/upstream/accounts/pool.rs +++ b/backend/src/accounts/pool.rs @@ -2,7 +2,7 @@ use std::{ collections::{BTreeMap, BTreeSet, VecDeque}, - sync::Arc as StdArc, + sync::Arc, time::Duration as StdDuration, }; @@ -10,123 +10,27 @@ use chrono::{DateTime, Duration, Utc}; use indexmap::IndexMap; use thiserror::Error; -use crate::upstream::accounts::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at}; -use crate::upstream::accounts::store::AccountStore; -use crate::upstream::accounts::token_refresh::{jwt_expiry, JwtExpiry}; -use crate::upstream::accounts::window::should_reset_usage_window; -use crate::upstream::accounts::{ - model::{Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta}, - service::AccountService, -}; -use crate::upstream::protocol::events::{ - parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage, -}; -use crate::upstream::scheduler::{ +use crate::accounts::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at}; +use crate::accounts::refresh::{jwt_expiry, JwtExpiry}; +use crate::accounts::scheduler::{ candidates::{self, CandidateFilter, CandidateRequest}, AccountScheduler, ScoreWeights, }; +use crate::accounts::store::AccountStore; +use crate::accounts::window::should_reset_usage_window; +use crate::accounts::{ + account::{Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta}, + service::AccountService, +}; +use crate::upstream::openai::protocol::events::{ + parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage, +}; -pub use crate::upstream::scheduler::RotationStrategy; +pub use crate::accounts::scheduler::RotationStrategy; /// 账号池配置。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AccountPoolOptions { - /// 单账号允许的最大并发请求数。 - pub max_concurrent_per_account: usize, - /// 在途槽位未释放时的过期时间。 - pub stale_slot_ttl: Duration, - /// 账号选择轮转策略。 - pub rotation_strategy: RotationStrategy, - /// 是否跳过已标记配额耗尽的账号。 - pub skip_quota_limited: bool, - /// 订阅层级优先级，越靠前优先级越高。 - pub tier_priority: Vec, - /// 模型到允许订阅计划的映射。 - pub model_plan_allowlist: BTreeMap>, - /// 已成功拉取过模型列表的订阅计划。 - pub fetched_model_plan_types: BTreeSet, -} - -impl Default for AccountPoolOptions { - fn default() -> Self { - Self { - max_concurrent_per_account: 3, - stale_slot_ttl: Duration::minutes(5), - rotation_strategy: RotationStrategy::Smart, - skip_quota_limited: true, - tier_priority: Vec::new(), - model_plan_allowlist: BTreeMap::new(), - fetched_model_plan_types: BTreeSet::new(), - } - } -} - -/// 账号获取请求。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AccountAcquireRequest { - /// 本次请求使用的模型。 - pub model: String, - /// 本次请求需要排除的账号 ID。 - pub exclude_account_ids: Vec, - /// 本次请求必须使用的账号 ID，用于同请求内的同账号重试。 - pub required_account_id: Option, - /// 会话亲和性建议的优先账号 ID。 - pub preferred_account_id: Option, - /// 本次调度使用的当前时间。 - pub now: DateTime, -} - -impl AccountAcquireRequest { - /// 构造账号获取请求。 - pub fn new(model: impl Into, now: DateTime) -> Self { - Self { - model: model.into(), - exclude_account_ids: Vec::new(), - required_account_id: None, - preferred_account_id: None, - now, - } - } - - /// 设置需要排除的账号 ID。 - pub fn with_exclude_account_ids( - mut self, - account_ids: impl IntoIterator>, - ) -> Self { - self.exclude_account_ids = account_ids.into_iter().map(Into::into).collect(); - self - } - - /// 设置会话亲和性建议的优先账号 ID。 - pub fn with_preferred_account_id(mut self, account_id: impl Into) -> Self { - self.preferred_account_id = Some(account_id.into()); - self - } - - /// 设置本次调度必须使用的账号 ID。 - pub fn with_required_account_id(mut self, account_id: impl Into) -> Self { - self.required_account_id = Some(account_id.into()); - self - } -} - -/// 成功获取到的账号及调度元数据。 -#[derive(Debug, Clone)] -pub struct AcquiredAccount { - /// 被选中的账号快照。 - pub account: Account, - /// 同一账号上一个在途槽位的创建时间。 - pub previous_slot_at: Option>, -} - -/// 模型刷新时每个订阅计划选中的账号。 -#[derive(Debug, Clone)] -pub struct DistinctPlanAccount { - /// 订阅计划类型。 - pub plan_type: String, - /// 被选中的账号快照。 - pub account: Account, -} +mod types; +pub use types::*; #[derive(Debug)] struct AccountAcquireWithStatusRefresh { @@ -245,7 +149,7 @@ pub struct AccountWindowUsageDelta { /// 纯内存账号池：负责账号存储、在途槽位与运行时状态维护。 /// -/// 账号**选择**已剥离到 [`crate::upstream::scheduler`]：候选过滤走 +/// 账号**选择**已剥离到 [`crate::accounts::scheduler`]：候选过滤走 /// [`candidates::filter`]，策略选择走 [`AccountScheduler::select`]。账号池只提供选择 /// 所需的只读视图（候选切片 + `slot_count`），并持有调度器实例。 #[derive(Debug)] @@ -838,592 +742,9 @@ fn access_token_expired(account: &Account, now: DateTime) -> bool { } /// 运行时账号池服务。 -#[derive(Clone)] -pub struct RuntimeAccountPoolService { - pool: StdArc>, - store: StdArc, - request_interval: StdArc>, -} - -impl RuntimeAccountPoolService { - /// 构造运行时账号池服务。 - pub fn new( - store: StdArc, - options: AccountPoolOptions, - request_interval_ms: u64, - ) -> Self { - Self { - pool: StdArc::new(tokio::sync::Mutex::new(AccountPool::with_options(options))), - store, - request_interval: StdArc::new(std::sync::RwLock::new(StdDuration::from_millis( - request_interval_ms, - ))), - } - } - - /// 更新账号池运行参数。 - pub async fn apply_options(&self, mut options: AccountPoolOptions, request_interval_ms: u64) { - let mut pool = self.pool.lock().await; - options.model_plan_allowlist = pool.options.model_plan_allowlist.clone(); - options.fetched_model_plan_types = pool.options.fetched_model_plan_types.clone(); - pool.set_options(options); - *self - .request_interval - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) = - StdDuration::from_millis(request_interval_ms); - } - - /// 更新模型到可用账号计划的调度约束。 - pub(crate) async fn apply_model_plan_routing( - &self, - model_plan_allowlist: BTreeMap>, - fetched_model_plan_types: BTreeSet, - ) { - let mut pool = self.pool.lock().await; - pool.options.model_plan_allowlist = model_plan_allowlist; - pool.options.fetched_model_plan_types = fetched_model_plan_types; - } - - /// 从持久化账号恢复运行时账号池。 - pub async fn restore_from_repository(&self) -> Result { - let accounts = self - .store - .list_pool_accounts() - .await - .map_err(|_| RuntimeAccountPoolError::Generic)?; - let count = accounts.len(); - let mut pool = self.pool.lock().await; - pool.clear(); - for account in accounts { - pool.insert(account); - } - Ok(count) - } - - /// 读取账号池容量摘要。 - pub async fn capacity_summary(&self, now: DateTime) -> AccountCapacitySummary { - let refresh = self - .pool - .lock() - .await - .capacity_summary_with_status_refresh(now); - self.persist_runtime_account_states(refresh.refreshed_accounts) - .await; - refresh.summary - } - - /// 使用当前时间读取账号池容量摘要。 - pub async fn capacity_summary_now(&self) -> AccountCapacitySummary { - self.capacity_summary(Utc::now()).await - } - - /// 释放账号在途槽位。 - pub async fn release(&self, account_id: &str) { - self.release_slot(account_id, true).await; - } - - /// 释放账号在途槽位，不累计请求用量。 - pub async fn release_without_request_usage(&self, account_id: &str) { - self.release_slot(account_id, false).await; - } - - async fn release_slot(&self, account_id: &str, record_request_usage: bool) { - let released = if record_request_usage { - self.pool.lock().await.release(account_id) - } else { - self.pool - .lock() - .await - .release_without_request_usage(account_id) - }; - let Some(released) = released else { - return; - }; - if !record_request_usage { - return; - } - if let Err(error) = self.store.record_request(account_id).await { - tracing::warn!( - account_id, - error = %error, - "failed to persist account request usage" - ); - } - let Some(model) = released.model else { - tracing::debug!( - account_id, - "released account without active slot; skipped model request usage persistence" - ); - return; - }; - if let Err(error) = self - .store - .record_model_usage_delta( - account_id, - &model, - AccountModelUsageDelta { - requests: 1, - ..AccountModelUsageDelta::default() - }, - ) - .await - { - tracing::warn!( - account_id, - model, - error = %error, - "failed to persist account model request usage" - ); - } - } - - /// 更新账号状态并同步内存池。 - pub async fn set_status(&self, account_id: &str, status: AccountStatus) -> bool { - let persisted = match self.store.set_status(account_id, status).await { - Ok(persisted) => persisted, - Err(error) => { - tracing::warn!( - account_id, - error = %error, - "failed to persist account status" - ); - false - } - }; - let in_memory = self.pool.lock().await.set_status(account_id, status); - persisted || in_memory - } - - /// 标记账号 quota 冷却状态。 - pub async fn mark_quota_limited_until(&self, account_id: &str, until: DateTime) -> bool { - let persisted = match self.store.mark_quota_limited_until(account_id, until).await { - Ok(persisted) => persisted, - Err(error) => { - tracing::warn!( - account_id, - error = %error, - "failed to persist quota cooldown state" - ); - false - } - }; - let in_memory = self - .pool - .lock() - .await - .mark_quota_limited_until(account_id, until); - persisted || in_memory - } - - /// 按请求上下文获取可用账号。 - pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option { - let refresh = self.pool.lock().await.acquire_with_status_refresh(request); - self.persist_runtime_account_states(refresh.refreshed_accounts) - .await; - refresh.acquired - } - - /// 按订阅计划各选一个可用账号，用于刷新模型列表。 - pub async fn distinct_plan_accounts(&self, now: DateTime) -> Vec { - let refresh = self - .pool - .lock() - .await - .distinct_plan_accounts_with_status_refresh(now); - self.persist_runtime_account_states(refresh.refreshed_accounts) - .await; - refresh.accounts - } - - async fn persist_runtime_account_states(&self, refreshed_states: Vec) { - for refreshed in refreshed_states { - if let Err(error) = self - .store - .sync_runtime_account_state(&refreshed.account, refreshed.sync_usage_window) - .await - { - tracing::warn!( - account_id = %refreshed.account.id, - error = %error, - "failed to persist refreshed runtime account state" - ); - } - } - } - - /// 按账号上一个在途槽位控制请求间隔。 - pub async fn wait_for_request_interval(&self, acquired: &AcquiredAccount) { - let request_interval = *self - .request_interval - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if request_interval.is_zero() { - return; - } - let Some(previous_slot_at) = acquired.previous_slot_at else { - return; - }; - let elapsed = Utc::now() - .signed_duration_since(previous_slot_at) - .to_std() - .unwrap_or_default(); - if elapsed < request_interval { - tokio::time::sleep(request_interval - elapsed).await; - } - } - - /// 记录上游 token 用量。 - pub async fn record_token_usage(&self, account_id: &str, model: &str, usage: &CodexTokenUsage) { - self.record_response_usage(account_id, model, *usage, false) - .await; - } - - /// 记录 Responses 请求用量。 - pub async fn record_response_usage( - &self, - account_id: &str, - model: &str, - usage: CodexTokenUsage, - image_generation_requested: bool, - ) { - let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0; - let image_request_failed = image_generation_requested && !image_request_succeeded; - let persisted_usage = AccountUsageDelta { - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - reasoning_tokens: usage.reasoning_tokens, - total_tokens: usage.total_tokens, - image_input_tokens: usage.image_input_tokens, - image_output_tokens: usage.image_output_tokens, - image_requests: u64::from(image_request_succeeded), - image_request_failures: u64::from(image_request_failed), - ..AccountUsageDelta::default() - }; - if let Err(error) = self - .store - .record_usage_delta(account_id, persisted_usage) - .await - { - tracing::warn!( - account_id, - error = %error, - "failed to persist account token usage" - ); - } - if let Err(error) = self - .store - .record_model_usage_delta( - account_id, - model, - AccountModelUsageDelta { - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - ..AccountModelUsageDelta::default() - }, - ) - .await - { - tracing::warn!( - account_id, - model, - error = %error, - "failed to persist account model token usage" - ); - } - self.pool.lock().await.record_window_token_usage( - account_id, - AccountWindowUsageDelta { - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cached_tokens: usage.cached_tokens, - image_input_tokens: usage.image_input_tokens, - image_output_tokens: usage.image_output_tokens, - image_request_succeeded, - image_request_failed, - }, - ); - } - - /// 记录空响应尝试。 - pub async fn record_empty_response_attempt( - &self, - account_id: &str, - model: &str, - image_generation_requested: bool, - ) { - let usage = AccountUsageDelta { - empty_responses: 1, - image_request_failures: u64::from(image_generation_requested), - ..AccountUsageDelta::default() - }; - if let Err(error) = self.store.record_usage_delta(account_id, usage).await { - tracing::warn!( - account_id, - error = %error, - "failed to persist empty response usage" - ); - } - if let Err(error) = self - .store - .record_model_usage_delta( - account_id, - model, - AccountModelUsageDelta { - errors: 1, - ..AccountModelUsageDelta::default() - }, - ) - .await - { - tracing::warn!( - account_id, - model, - error = %error, - "failed to persist account model empty response usage" - ); - } - if image_generation_requested { - self.pool.lock().await.record_window_token_usage( - account_id, - AccountWindowUsageDelta { - image_request_failed: true, - ..AccountWindowUsageDelta::default() - }, - ); - } - } - - /// 根据上游返回的 rate-limit header 被动同步 quota 状态。 - pub async fn sync_passive_rate_limit_headers( - &self, - account: &Account, - headers: &[(String, String)], - ) { - self.sync_passive_rate_limit_headers_for_account( - &account.id, - account.plan_type.as_deref(), - headers, - ) - .await; - } - - /// 根据上游返回的 rate-limit header 被动同步指定账号 quota 状态。 - pub async fn sync_passive_rate_limit_headers_for_account( - &self, - account_id: &str, - plan_type: Option<&str>, - headers: &[(String, String)], - ) { - let Some(rate_limits) = parse_rate_limit_headers(headers) else { - return; - }; - let existing_quota = match self.store.get_quota_json(account_id).await { - Ok(Some(quota_json)) => serde_json::from_str::("a_json).ok(), - Ok(None) => None, - Err(error) => { - tracing::warn!( - account_id = %account_id, - error = %error, - "failed to read existing quota json before passive rate-limit sync" - ); - None - } - }; - let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref()); - self.apply_quota_snapshot(account_id, "a).await; - } - - /// 读取运行时账号快照。 - pub async fn account_snapshot(&self, account_id: &str) -> Option { - self.pool.lock().await.get(account_id) - } - - /// 应用 quota 快照到持久化存储和运行时账号池。 - pub async fn apply_quota_snapshot(&self, account_id: &str, quota: &serde_json::Value) -> bool { - let limit_reached = quota_snapshot_limit_reached(quota); - let reset_at = quota_snapshot_reset_at(quota); - let cooldown_until = limit_reached.then_some(reset_at).flatten(); - let limit_window_seconds = reset_at.and_then(|_| { - crate::upstream::accounts::quota::quota_snapshot_limit_window_seconds(quota) - }); - let quota_json = quota.to_string(); - let persisted = match self - .store - .apply_quota_snapshot(account_id, "a_json, limit_reached, cooldown_until) - .await - { - Ok(persisted) => persisted, - Err(error) => { - tracing::warn!( - account_id, - error = %error, - "failed to persist verified quota snapshot" - ); - false - } - }; - - if let Some(reset_at) = reset_at { - if let Err(error) = self - .store - .sync_rate_limit_window(account_id, reset_at, limit_window_seconds) - .await - { - tracing::warn!( - account_id, - error = %error, - "failed to persist verified quota window" - ); - } - } - - let in_memory = { - let mut pool = self.pool.lock().await; - let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until); - if let Some(reset_at) = reset_at { - pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds); - } - in_memory - }; - - persisted || in_memory - } - - /// 从运行时账号池移除账号。 - pub async fn remove_account(&self, account_id: &str) -> bool { - self.pool.lock().await.remove(account_id) - } - - /// 从仓储同步单个账号到运行时账号池。 - pub async fn sync_account_from_repository( - &self, - account_id: &str, - ) -> Result { - let account = self - .store - .get_pool_account(account_id) - .await - .map_err(|_| RuntimeAccountPoolError::Generic)?; - let mut pool = self.pool.lock().await; - if let Some(account) = account { - pool.insert(account); - return Ok(true); - } - Ok(pool.remove(account_id)) - } - - /// 设置 Cloudflare 冷却状态。 - pub async fn set_cloudflare_cooldown_until( - &self, - account_id: &str, - cooldown_until: DateTime, - ) -> bool { - let persisted = match self - .store - .set_cloudflare_cooldown_until(account_id, cooldown_until) - .await - { - Ok(persisted) => persisted, - Err(error) => { - tracing::warn!( - account_id, - error = %error, - "failed to persist Cloudflare cooldown state" - ); - false - } - }; - let in_memory = self - .pool - .lock() - .await - .set_cloudflare_cooldown_until(account_id, cooldown_until); - persisted || in_memory - } -} - -/// 运行时账号池错误。 -#[derive(Debug, Error)] -pub enum RuntimeAccountPoolError { - /// 通用账号池错误。 - #[error("account pool error")] - Generic, -} - -fn refresh_quota_window(account: &mut Account, now: DateTime) { - let window_expired = account - .window_reset_at - .is_some_and(|reset_at| now >= reset_at); - let cooldown_expired = account - .quota_cooldown_until - .is_some_and(|cooldown_until| now >= cooldown_until); - - if window_expired { - reset_window_counters(account); - account.window_started_at = Some(now); - account.window_reset_at = - next_window_reset_at(account.window_reset_at, account.limit_window_seconds, now); - } - - if account.quota_limit_reached && (cooldown_expired || window_expired) { - account.quota_verify_required = true; - account.quota_limit_reached = false; - account.quota_cooldown_until = None; - if account.status == AccountStatus::QuotaExhausted { - account.status = AccountStatus::Active; - } - } -} - -fn status_after_quota_limit(status: AccountStatus, quota_limit_reached: bool) -> AccountStatus { - match (status, quota_limit_reached) { - (AccountStatus::Active, true) => AccountStatus::QuotaExhausted, - (status, _) => status, - } -} - -fn reset_window_counters(account: &mut Account) { - account.window_request_count = 0; - account.window_input_tokens = 0; - account.window_output_tokens = 0; - account.window_cached_tokens = 0; - account.window_image_input_tokens = 0; - account.window_image_output_tokens = 0; - account.window_image_request_count = 0; - account.window_image_request_failed_count = 0; -} - -fn next_window_reset_at( - reset_at: Option>, - limit_window_seconds: Option, - now: DateTime, -) -> Option> { - let reset_at = reset_at?; - let window_seconds = limit_window_seconds?; - if window_seconds == 0 { - return None; - } - - let elapsed_seconds = now - .signed_duration_since(reset_at) - .num_seconds() - .max(0) - .cast_unsigned(); - let windows_to_advance = elapsed_seconds / window_seconds + 1; - let advance_seconds = window_seconds - .saturating_mul(windows_to_advance) - .min(i64::MAX as u64); - Some(reset_at + Duration::seconds(advance_seconds as i64)) -} - -fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime) { - if account - .cloudflare_cooldown_until - .is_some_and(|cooldown_until| now >= cooldown_until) - { - account.cloudflare_cooldown_until = None; - } -} +mod runtime; +pub use runtime::*; +use runtime::{ + refresh_cloudflare_cooldown, refresh_quota_window, reset_window_counters, + status_after_quota_limit, +}; diff --git a/backend/src/accounts/pool/runtime.rs b/backend/src/accounts/pool/runtime.rs new file mode 100644 index 000000000..f50532501 --- /dev/null +++ b/backend/src/accounts/pool/runtime.rs @@ -0,0 +1,593 @@ +use super::*; + +#[derive(Clone)] +pub struct RuntimeAccountPoolService { + pool: Arc>, + store: Arc, + request_interval: Arc>, +} + +impl RuntimeAccountPoolService { + /// 构造运行时账号池服务。 + pub fn new( + store: Arc, + options: AccountPoolOptions, + request_interval_ms: u64, + ) -> Self { + Self { + pool: Arc::new(tokio::sync::Mutex::new(AccountPool::with_options(options))), + store, + request_interval: Arc::new(std::sync::RwLock::new(StdDuration::from_millis( + request_interval_ms, + ))), + } + } + + /// 更新账号池运行参数。 + pub async fn apply_options(&self, mut options: AccountPoolOptions, request_interval_ms: u64) { + let mut pool = self.pool.lock().await; + options.model_plan_allowlist = pool.options.model_plan_allowlist.clone(); + options.fetched_model_plan_types = pool.options.fetched_model_plan_types.clone(); + pool.set_options(options); + *self + .request_interval + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + StdDuration::from_millis(request_interval_ms); + } + + /// 更新模型到可用账号计划的调度约束。 + pub(crate) async fn apply_model_plan_routing( + &self, + model_plan_allowlist: BTreeMap>, + fetched_model_plan_types: BTreeSet, + ) { + let mut pool = self.pool.lock().await; + pool.options.model_plan_allowlist = model_plan_allowlist; + pool.options.fetched_model_plan_types = fetched_model_plan_types; + } + + /// 从持久化账号恢复运行时账号池。 + pub async fn restore_from_store(&self) -> Result { + let accounts = self + .store + .list_pool_accounts() + .await + .map_err(|_| RuntimeAccountPoolError::Generic)?; + let count = accounts.len(); + let mut pool = self.pool.lock().await; + pool.clear(); + for account in accounts { + pool.insert(account); + } + Ok(count) + } + + /// 读取账号池容量摘要。 + pub async fn capacity_summary(&self, now: DateTime) -> AccountCapacitySummary { + let refresh = self + .pool + .lock() + .await + .capacity_summary_with_status_refresh(now); + self.persist_runtime_account_states(refresh.refreshed_accounts) + .await; + refresh.summary + } + + /// 使用当前时间读取账号池容量摘要。 + pub async fn capacity_summary_now(&self) -> AccountCapacitySummary { + self.capacity_summary(Utc::now()).await + } + + /// 释放账号在途槽位。 + pub async fn release(&self, account_id: &str) { + self.release_slot(account_id, true).await; + } + + /// 释放账号在途槽位，不累计请求用量。 + pub async fn release_without_request_usage(&self, account_id: &str) { + self.release_slot(account_id, false).await; + } + + async fn release_slot(&self, account_id: &str, record_request_usage: bool) { + let released = if record_request_usage { + self.pool.lock().await.release(account_id) + } else { + self.pool + .lock() + .await + .release_without_request_usage(account_id) + }; + let Some(released) = released else { + return; + }; + if !record_request_usage { + return; + } + if let Err(error) = self.store.record_request(account_id).await { + tracing::warn!( + account_id, + error = %error, + "failed to persist account request usage" + ); + } + let Some(model) = released.model else { + tracing::debug!( + account_id, + "released account without active slot; skipped model request usage persistence" + ); + return; + }; + if let Err(error) = self + .store + .record_model_usage_delta( + account_id, + &model, + AccountModelUsageDelta { + requests: 1, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id, + model, + error = %error, + "failed to persist account model request usage" + ); + } + } + + /// 更新账号状态并同步内存池。 + pub async fn set_status(&self, account_id: &str, status: AccountStatus) -> bool { + let persisted = match self.store.set_status(account_id, status).await { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist account status" + ); + false + } + }; + let in_memory = self.pool.lock().await.set_status(account_id, status); + persisted || in_memory + } + + /// 标记账号 quota 冷却状态。 + pub async fn mark_quota_limited_until(&self, account_id: &str, until: DateTime) -> bool { + let persisted = match self.store.mark_quota_limited_until(account_id, until).await { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist quota cooldown state" + ); + false + } + }; + let in_memory = self + .pool + .lock() + .await + .mark_quota_limited_until(account_id, until); + persisted || in_memory + } + + /// 按请求上下文获取可用账号。 + pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option { + let refresh = self.pool.lock().await.acquire_with_status_refresh(request); + self.persist_runtime_account_states(refresh.refreshed_accounts) + .await; + refresh.acquired + } + + /// 按订阅计划各选一个可用账号，用于刷新模型列表。 + pub async fn distinct_plan_accounts(&self, now: DateTime) -> Vec { + let refresh = self + .pool + .lock() + .await + .distinct_plan_accounts_with_status_refresh(now); + self.persist_runtime_account_states(refresh.refreshed_accounts) + .await; + refresh.accounts + } + + async fn persist_runtime_account_states(&self, refreshed_states: Vec) { + for refreshed in refreshed_states { + if let Err(error) = self + .store + .sync_runtime_account_state(&refreshed.account, refreshed.sync_usage_window) + .await + { + tracing::warn!( + account_id = %refreshed.account.id, + error = %error, + "failed to persist refreshed runtime account state" + ); + } + } + } + + /// 按账号上一个在途槽位控制请求间隔。 + pub async fn wait_for_request_interval(&self, acquired: &AcquiredAccount) { + let request_interval = *self + .request_interval + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if request_interval.is_zero() { + return; + } + let Some(previous_slot_at) = acquired.previous_slot_at else { + return; + }; + let elapsed = Utc::now() + .signed_duration_since(previous_slot_at) + .to_std() + .unwrap_or_default(); + if elapsed < request_interval { + tokio::time::sleep(request_interval - elapsed).await; + } + } + + /// 记录上游 token 用量。 + pub async fn record_token_usage(&self, account_id: &str, model: &str, usage: &CodexTokenUsage) { + self.record_response_usage(account_id, model, *usage, false) + .await; + } + + /// 记录 Responses 请求用量。 + pub async fn record_response_usage( + &self, + account_id: &str, + model: &str, + usage: CodexTokenUsage, + image_generation_requested: bool, + ) { + let image_request_succeeded = image_generation_requested && usage.image_output_tokens > 0; + let image_request_failed = image_generation_requested && !image_request_succeeded; + let persisted_usage = AccountUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + reasoning_tokens: usage.reasoning_tokens, + total_tokens: usage.total_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_requests: u64::from(image_request_succeeded), + image_request_failures: u64::from(image_request_failed), + ..AccountUsageDelta::default() + }; + if let Err(error) = self + .store + .record_usage_delta(account_id, persisted_usage) + .await + { + tracing::warn!( + account_id, + error = %error, + "failed to persist account token usage" + ); + } + if let Err(error) = self + .store + .record_model_usage_delta( + account_id, + model, + AccountModelUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id, + model, + error = %error, + "failed to persist account model token usage" + ); + } + self.pool.lock().await.record_window_token_usage( + account_id, + AccountWindowUsageDelta { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + image_input_tokens: usage.image_input_tokens, + image_output_tokens: usage.image_output_tokens, + image_request_succeeded, + image_request_failed, + }, + ); + } + + /// 记录空响应尝试。 + pub async fn record_empty_response_attempt( + &self, + account_id: &str, + model: &str, + image_generation_requested: bool, + ) { + let usage = AccountUsageDelta { + empty_responses: 1, + image_request_failures: u64::from(image_generation_requested), + ..AccountUsageDelta::default() + }; + if let Err(error) = self.store.record_usage_delta(account_id, usage).await { + tracing::warn!( + account_id, + error = %error, + "failed to persist empty response usage" + ); + } + if let Err(error) = self + .store + .record_model_usage_delta( + account_id, + model, + AccountModelUsageDelta { + errors: 1, + ..AccountModelUsageDelta::default() + }, + ) + .await + { + tracing::warn!( + account_id, + model, + error = %error, + "failed to persist account model empty response usage" + ); + } + if image_generation_requested { + self.pool.lock().await.record_window_token_usage( + account_id, + AccountWindowUsageDelta { + image_request_failed: true, + ..AccountWindowUsageDelta::default() + }, + ); + } + } + + /// 根据上游返回的 rate-limit header 被动同步 quota 状态。 + pub async fn sync_passive_rate_limit_headers( + &self, + account: &Account, + headers: &[(String, String)], + ) { + self.sync_passive_rate_limit_headers_for_account( + &account.id, + account.plan_type.as_deref(), + headers, + ) + .await; + } + + /// 根据上游返回的 rate-limit header 被动同步指定账号 quota 状态。 + pub async fn sync_passive_rate_limit_headers_for_account( + &self, + account_id: &str, + plan_type: Option<&str>, + headers: &[(String, String)], + ) { + let Some(rate_limits) = parse_rate_limit_headers(headers) else { + return; + }; + let existing_quota = match self.store.get_quota_json(account_id).await { + Ok(Some(quota_json)) => serde_json::from_str::("a_json).ok(), + Ok(None) => None, + Err(error) => { + tracing::warn!( + account_id = %account_id, + error = %error, + "failed to read existing quota json before passive rate-limit sync" + ); + None + } + }; + let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref()); + self.apply_quota_snapshot(account_id, "a).await; + } + + /// 读取运行时账号快照。 + pub async fn account_snapshot(&self, account_id: &str) -> Option { + self.pool.lock().await.get(account_id) + } + + /// 应用 quota 快照到持久化存储和运行时账号池。 + pub async fn apply_quota_snapshot(&self, account_id: &str, quota: &serde_json::Value) -> bool { + let limit_reached = quota_snapshot_limit_reached(quota); + let reset_at = quota_snapshot_reset_at(quota); + let cooldown_until = limit_reached.then_some(reset_at).flatten(); + let limit_window_seconds = reset_at + .and_then(|_| crate::accounts::quota::quota_snapshot_limit_window_seconds(quota)); + let quota_json = quota.to_string(); + let persisted = match self + .store + .apply_quota_snapshot(account_id, "a_json, limit_reached, cooldown_until) + .await + { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist verified quota snapshot" + ); + false + } + }; + + if let Some(reset_at) = reset_at { + if let Err(error) = self + .store + .sync_rate_limit_window(account_id, reset_at, limit_window_seconds) + .await + { + tracing::warn!( + account_id, + error = %error, + "failed to persist verified quota window" + ); + } + } + + let in_memory = { + let mut pool = self.pool.lock().await; + let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until); + if let Some(reset_at) = reset_at { + pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds); + } + in_memory + }; + + persisted || in_memory + } + + /// 从运行时账号池移除账号。 + pub async fn remove_account(&self, account_id: &str) -> bool { + self.pool.lock().await.remove(account_id) + } + + /// 从仓储同步单个账号到运行时账号池。 + pub async fn sync_account_from_store( + &self, + account_id: &str, + ) -> Result { + let account = self + .store + .get_pool_account(account_id) + .await + .map_err(|_| RuntimeAccountPoolError::Generic)?; + let mut pool = self.pool.lock().await; + if let Some(account) = account { + pool.insert(account); + return Ok(true); + } + Ok(pool.remove(account_id)) + } + + /// 设置 Cloudflare 冷却状态。 + pub async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> bool { + let persisted = match self + .store + .set_cloudflare_cooldown_until(account_id, cooldown_until) + .await + { + Ok(persisted) => persisted, + Err(error) => { + tracing::warn!( + account_id, + error = %error, + "failed to persist Cloudflare cooldown state" + ); + false + } + }; + let in_memory = self + .pool + .lock() + .await + .set_cloudflare_cooldown_until(account_id, cooldown_until); + persisted || in_memory + } +} + +/// 运行时账号池错误。 +#[derive(Debug, Error)] +pub enum RuntimeAccountPoolError { + /// 通用账号池错误。 + #[error("account pool error")] + Generic, +} + +pub(super) fn refresh_quota_window(account: &mut Account, now: DateTime) { + let window_expired = account + .window_reset_at + .is_some_and(|reset_at| now >= reset_at); + let cooldown_expired = account + .quota_cooldown_until + .is_some_and(|cooldown_until| now >= cooldown_until); + + if window_expired { + reset_window_counters(account); + account.window_started_at = Some(now); + account.window_reset_at = + next_window_reset_at(account.window_reset_at, account.limit_window_seconds, now); + } + + if account.quota_limit_reached && (cooldown_expired || window_expired) { + account.quota_verify_required = true; + account.quota_limit_reached = false; + account.quota_cooldown_until = None; + if account.status == AccountStatus::QuotaExhausted { + account.status = AccountStatus::Active; + } + } +} + +pub(super) fn status_after_quota_limit( + status: AccountStatus, + quota_limit_reached: bool, +) -> AccountStatus { + match (status, quota_limit_reached) { + (AccountStatus::Active, true) => AccountStatus::QuotaExhausted, + (status, _) => status, + } +} + +pub(super) fn reset_window_counters(account: &mut Account) { + account.window_request_count = 0; + account.window_input_tokens = 0; + account.window_output_tokens = 0; + account.window_cached_tokens = 0; + account.window_image_input_tokens = 0; + account.window_image_output_tokens = 0; + account.window_image_request_count = 0; + account.window_image_request_failed_count = 0; +} + +fn next_window_reset_at( + reset_at: Option>, + limit_window_seconds: Option, + now: DateTime, +) -> Option> { + let reset_at = reset_at?; + let window_seconds = limit_window_seconds?; + if window_seconds == 0 { + return None; + } + + let elapsed_seconds = now + .signed_duration_since(reset_at) + .num_seconds() + .max(0) + .cast_unsigned(); + let windows_to_advance = elapsed_seconds / window_seconds + 1; + let advance_seconds = window_seconds + .saturating_mul(windows_to_advance) + .min(i64::MAX as u64); + Some(reset_at + Duration::seconds(advance_seconds as i64)) +} + +pub(super) fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime) { + if account + .cloudflare_cooldown_until + .is_some_and(|cooldown_until| now >= cooldown_until) + { + account.cloudflare_cooldown_until = None; + } +} diff --git a/backend/src/accounts/pool/types.rs b/backend/src/accounts/pool/types.rs new file mode 100644 index 000000000..01f4951d3 --- /dev/null +++ b/backend/src/accounts/pool/types.rs @@ -0,0 +1,100 @@ +use super::*; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountPoolOptions { + /// 单账号允许的最大并发请求数。 + pub max_concurrent_per_account: usize, + /// 在途槽位未释放时的过期时间。 + pub stale_slot_ttl: Duration, + /// 账号选择轮转策略。 + pub rotation_strategy: RotationStrategy, + /// 是否跳过已标记配额耗尽的账号。 + pub skip_quota_limited: bool, + /// 订阅层级优先级，越靠前优先级越高。 + pub tier_priority: Vec, + /// 模型到允许订阅计划的映射。 + pub model_plan_allowlist: BTreeMap>, + /// 已成功拉取过模型列表的订阅计划。 + pub fetched_model_plan_types: BTreeSet, +} + +impl Default for AccountPoolOptions { + fn default() -> Self { + Self { + max_concurrent_per_account: 3, + stale_slot_ttl: Duration::minutes(5), + rotation_strategy: RotationStrategy::Smart, + skip_quota_limited: true, + tier_priority: Vec::new(), + model_plan_allowlist: BTreeMap::new(), + fetched_model_plan_types: BTreeSet::new(), + } + } +} + +/// 账号获取请求。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AccountAcquireRequest { + /// 本次请求使用的模型。 + pub model: String, + /// 本次请求需要排除的账号 ID。 + pub exclude_account_ids: Vec, + /// 本次请求必须使用的账号 ID，用于同请求内的同账号重试。 + pub required_account_id: Option, + /// 会话亲和性建议的优先账号 ID。 + pub preferred_account_id: Option, + /// 本次调度使用的当前时间。 + pub now: DateTime, +} + +impl AccountAcquireRequest { + /// 构造账号获取请求。 + pub fn new(model: impl Into, now: DateTime) -> Self { + Self { + model: model.into(), + exclude_account_ids: Vec::new(), + required_account_id: None, + preferred_account_id: None, + now, + } + } + + /// 设置需要排除的账号 ID。 + pub fn with_exclude_account_ids( + mut self, + account_ids: impl IntoIterator>, + ) -> Self { + self.exclude_account_ids = account_ids.into_iter().map(Into::into).collect(); + self + } + + /// 设置会话亲和性建议的优先账号 ID。 + pub fn with_preferred_account_id(mut self, account_id: impl Into) -> Self { + self.preferred_account_id = Some(account_id.into()); + self + } + + /// 设置本次调度必须使用的账号 ID。 + pub fn with_required_account_id(mut self, account_id: impl Into) -> Self { + self.required_account_id = Some(account_id.into()); + self + } +} + +/// 成功获取到的账号及调度元数据。 +#[derive(Debug, Clone)] +pub struct AcquiredAccount { + /// 被选中的账号快照。 + pub account: Account, + /// 同一账号上一个在途槽位的创建时间。 + pub previous_slot_at: Option>, +} + +/// 模型刷新时每个订阅计划选中的账号。 +#[derive(Debug, Clone)] +pub struct DistinctPlanAccount { + /// 订阅计划类型。 + pub plan_type: String, + /// 被选中的账号快照。 + pub account: Account, +} diff --git a/backend/src/accounts/quota/mod.rs b/backend/src/accounts/quota/mod.rs new file mode 100644 index 000000000..0421c320d --- /dev/null +++ b/backend/src/accounts/quota/mod.rs @@ -0,0 +1,7 @@ +//! 账号配额解析与刷新。 + +mod runtime; +mod service; + +pub use runtime::*; +pub use service::*; diff --git a/backend/src/upstream/accounts/quota/runtime.rs b/backend/src/accounts/quota/runtime.rs similarity index 92% rename from backend/src/upstream/accounts/quota/runtime.rs rename to backend/src/accounts/quota/runtime.rs index 267856906..5843150e4 100644 --- a/backend/src/upstream/accounts/quota/runtime.rs +++ b/backend/src/accounts/quota/runtime.rs @@ -7,18 +7,18 @@ use thiserror::Error; use tokio::time::{sleep, Instant}; use tracing::warn; -use crate::upstream::accounts::{ - model::AccountStatus, +use crate::accounts::{ + account::AccountStatus, pool::RuntimeAccountPoolService, quota::{ quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at, }, - store::{AccountStore, AccountStoreError, SqliteAccountStore, SqliteAccountStoreError}, + store::{AccountStore, AccountStoreError, PgAccountStore, PgAccountStoreError}, }; -use crate::upstream::{ - accounts::cookies::SqliteCookieStore, - transport::{CodexBackendClient, CodexClientError, CodexRequestContext}, +use crate::{ + accounts::cookies::PgCookieStore, + upstream::openai::transport::{CodexBackendClient, CodexClientError, CodexRequestContext}, }; const MIN_REFRESH_INTERVAL_SECS: u64 = 30 * 60; @@ -27,18 +27,18 @@ const COOLDOWN_REFRESH_GRACE_SECS: i64 = 30; /// 运行时 quota refresh 服务。 pub struct RuntimeQuotaRefreshService { - store: SqliteAccountStore, + store: PgAccountStore, codex: Arc, min_refresh_interval_secs: u64, request_spacing: Duration, installation_id: Option, - cookie_store: Option, + cookie_store: Option, account_pool: Option>, } impl RuntimeQuotaRefreshService { /// 构造默认 quota refresh 服务。 - pub fn new(store: SqliteAccountStore, codex: Arc) -> Self { + pub fn new(store: PgAccountStore, codex: Arc) -> Self { Self { store, codex, @@ -56,7 +56,7 @@ impl RuntimeQuotaRefreshService { /// 使用自定义最小刷新间隔构造 quota refresh 服务。 pub fn with_min_refresh_interval_secs( - store: SqliteAccountStore, + store: PgAccountStore, codex: Arc, min_refresh_interval_secs: u64, ) -> Self { @@ -84,7 +84,7 @@ impl RuntimeQuotaRefreshService { } /// 设置 usage 请求可复用的账号 Cookie 存储。 - pub fn with_cookie_store(mut self, cookie_store: SqliteCookieStore) -> Self { + pub fn with_cookie_store(mut self, cookie_store: PgCookieStore) -> Self { self.cookie_store = Some(cookie_store); self } @@ -223,7 +223,7 @@ impl RuntimeQuotaRefreshService { let Some(account_pool) = &self.account_pool else { return; }; - if let Err(error) = account_pool.sync_account_from_repository(account_id).await { + if let Err(error) = account_pool.sync_account_from_store(account_id).await { warn!( account_id = %account_id, error = %error, @@ -266,7 +266,7 @@ pub enum QuotaRefreshServiceError { AccountStore(#[from] AccountStoreError), /// 账号写入失败。 #[error("failed to persist quota refresh result: {0}")] - Store(#[from] SqliteAccountStoreError), + Store(#[from] PgAccountStoreError), /// 上游 usage 请求失败。 #[error("failed to fetch quota usage: {0}")] Codex(#[from] CodexClientError), diff --git a/backend/src/upstream/accounts/quota/mod.rs b/backend/src/accounts/quota/service.rs similarity index 98% rename from backend/src/upstream/accounts/quota/mod.rs rename to backend/src/accounts/quota/service.rs index c81096d83..0888ed2d0 100644 --- a/backend/src/upstream/accounts/quota/mod.rs +++ b/backend/src/accounts/quota/service.rs @@ -410,10 +410,3 @@ fn trimmed_str(value: &Value) -> Option<&str> { fn window_minutes_match(actual: u64, expected: u64) -> bool { actual > 0 && actual.abs_diff(expected) <= expected / 20 } - -mod runtime; - -pub use runtime::{ - QuotaRefreshServiceError, QuotaRefreshServiceResult, QuotaRefreshSummary, - RuntimeQuotaRefreshService, -}; diff --git a/backend/src/accounts/refresh/lease.rs b/backend/src/accounts/refresh/lease.rs new file mode 100644 index 000000000..11a4f71e7 --- /dev/null +++ b/backend/src/accounts/refresh/lease.rs @@ -0,0 +1,118 @@ +//! Redis token 刷新租约存储。 + +use std::collections::HashSet; + +use chrono::{DateTime, Utc}; +use thiserror::Error; + +use crate::infra::redis::RedisConnection; + +/// Redis 刷新租约存储结果。 +pub type RedisRefreshLeaseStoreResult = Result; + +/// Redis 刷新租约存储错误。 +#[derive(Debug, Error)] +pub enum RedisRefreshLeaseStoreError { + #[error("Redis refresh lease operation failed: {0}")] + Redis(#[from] redis::RedisError), + #[error("failed to encode refresh lease owner: {0}")] + Json(#[from] serde_json::Error), +} + +/// Redis 刷新租约存储。 +#[derive(Clone)] +pub struct RedisRefreshLeaseStore { + redis: RedisConnection, +} + +impl RedisRefreshLeaseStore { + /// 构造存储。 + pub fn new(redis: RedisConnection) -> Self { + Self { redis } + } + + /// 尝试获取账号刷新租约。 + pub async fn try_acquire( + &self, + account_id: &str, + owner: &str, + expires_at: DateTime, + now: DateTime, + ) -> RedisRefreshLeaseStoreResult { + let ttl_ms = expires_at.signed_duration_since(now).num_milliseconds(); + if ttl_ms <= 0 { + return Ok(false); + } + let key = self.lease_key(account_id); + let owner = serde_json::to_string(owner)?; + let mut connection = self.redis.manager(); + let acquired: i64 = redis::Script::new( + r" +if redis.call('SET', KEYS[1], ARGV[1], 'NX', 'PX', ARGV[2]) then + return 1 +end +if redis.call('GET', KEYS[1]) == ARGV[1] then + redis.call('PEXPIRE', KEYS[1], ARGV[2]) + return 1 +end +return 0", + ) + .key(key) + .arg(owner) + .arg(ttl_ms) + .invoke_async(&mut connection) + .await?; + Ok(acquired == 1) + } + + /// 释放账号刷新租约。 + pub async fn release( + &self, + account_id: &str, + owner: &str, + ) -> RedisRefreshLeaseStoreResult { + let mut connection = self.redis.manager(); + let released: i64 = redis::Script::new( + r" +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0", + ) + .key(self.lease_key(account_id)) + .arg(serde_json::to_string(owner)?) + .invoke_async(&mut connection) + .await?; + Ok(released == 1) + } + + /// 返回给定账号集合中仍有效的刷新租约。 + pub async fn active_account_ids( + &self, + account_ids: &[String], + _now: DateTime, + ) -> RedisRefreshLeaseStoreResult> { + if account_ids.is_empty() { + return Ok(HashSet::new()); + } + + let keys = account_ids + .iter() + .map(|account_id| self.lease_key(account_id)) + .collect::>(); + let mut connection = self.redis.manager(); + let values: Vec> = redis::cmd("MGET") + .arg(&keys) + .query_async(&mut connection) + .await?; + Ok(account_ids + .iter() + .zip(values) + .filter_map(|(account_id, value)| value.map(|_| account_id.clone())) + .collect()) + } + + fn lease_key(&self, account_id: &str) -> String { + self.redis.key(&format!("lease:refresh:{account_id}")) + } +} diff --git a/backend/src/accounts/refresh/mod.rs b/backend/src/accounts/refresh/mod.rs new file mode 100644 index 000000000..2cdc80f73 --- /dev/null +++ b/backend/src/accounts/refresh/mod.rs @@ -0,0 +1,9 @@ +//! Token 刷新策略、租约与运行时服务。 + +mod lease; +mod runtime; +mod service; + +pub use lease::*; +pub use runtime::*; +pub use service::*; diff --git a/backend/src/upstream/accounts/token_refresh/runtime.rs b/backend/src/accounts/refresh/runtime.rs similarity index 85% rename from backend/src/upstream/accounts/token_refresh/runtime.rs rename to backend/src/accounts/refresh/runtime.rs index c75aa7146..5508d4bba 100644 --- a/backend/src/upstream/accounts/token_refresh/runtime.rs +++ b/backend/src/accounts/refresh/runtime.rs @@ -16,18 +16,18 @@ use tokio::{sync::Mutex, task::JoinHandle, time::sleep}; use tracing::{debug, info, warn}; use uuid::Uuid; -use crate::upstream::accounts::{ - model::{Account, AccountStatus}, - store::{ - AccountClaimsUpdate, AccountStore, SqliteAccountStore, SqliteAccountStoreError, - StoredAccount, +use crate::accounts::{ + account::{Account, AccountStatus}, + refresh::{ + jittered_refresh_at, jwt_expiration, token_refresh_status_eligible, RedisRefreshLeaseStore, + RedisRefreshLeaseStoreError, RefreshError, RefreshPolicy, RefreshScheduler, + RuntimeRefreshPolicy, }, - token_refresh::{ - jittered_refresh_at, jwt_expiration, token_refresh_status_eligible, RefreshError, - RefreshLeaseStore, RefreshLeaseStoreError, RefreshPolicy, RefreshScheduler, - RuntimeRefreshPolicy, TokenRefresher, + store::{ + AccountClaimsUpdate, AccountStore, PgAccountStore, PgAccountStoreError, StoredAccount, }, }; +use crate::upstream::openai::token_client::TokenRefresher; const REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60; const MAX_REFRESH_ATTEMPTS: usize = 5; @@ -43,10 +43,10 @@ pub struct RuntimeTokenRefreshService where C: TokenRefresher, { - store: SqliteAccountStore, + store: PgAccountStore, scheduler: Arc>, policy: RuntimeRefreshPolicy, - refresh_leases: Option, + refresh_leases: Option, lease_owner: String, retry_delays: Vec, in_flight: Arc>>, @@ -88,11 +88,7 @@ where C: TokenRefresher, { /// 构造运行时 token refresh 服务。 - pub fn new( - store: SqliteAccountStore, - policy: impl Into, - client: C, - ) -> Self { + pub fn new(store: PgAccountStore, policy: impl Into, client: C) -> Self { let policy = policy.into(); Self { store, @@ -107,7 +103,7 @@ where } /// 使用刷新租约存储保护账号刷新。 - pub fn with_refresh_lease_store(mut self, refresh_leases: RefreshLeaseStore) -> Self { + pub fn with_refresh_lease_store(mut self, refresh_leases: RedisRefreshLeaseStore) -> Self { self.refresh_leases = Some(refresh_leases); self } @@ -433,9 +429,7 @@ where now: DateTime, ) -> TokenRefreshServiceResult { let Some(account) = self.store.get(account_id).await? else { - return Err(TokenRefreshServiceError::AccountNotFound( - account_id.to_string(), - )); + return Ok(TokenRefreshOutcome::Skipped); }; let account = stored_account_to_refresh_account(account); if self.clear_refresh_schedule_if_ineligible(&account).await? { @@ -654,75 +648,11 @@ where } } -enum TokenRefreshOutcome { - Refreshed(Box), - StatusUpdated, - Skipped, - Failed, -} - -/// 单次刷新扫描摘要。 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct TokenRefreshSummary { - /// 扫描账号数。 - pub scanned: usize, - /// 成功刷新 token 的账号数。 - pub refreshed: usize, - /// 仅更新状态的账号数。 - pub status_updated: usize, - /// 无需刷新的账号数。 - pub skipped: usize, - /// 传输失败账号数。 - pub failed: usize, -} - -/// 账号级刷新定时器调度摘要。 -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct TokenTimerSummary { - /// 扫描账号数。 - pub scanned: usize, - /// 已调度的未来刷新定时器数。 - pub scheduled: usize, - /// 已调度的立即刷新定时器数。 - pub immediate: usize, - /// 无需调度的账号数。 - pub skipped: usize, - /// 被替换的既有定时器数。 - pub replaced: usize, -} - -impl TokenTimerSummary { - /// 返回本轮调度发生变化的定时器数量。 - pub fn changed(&self) -> usize { - self.scheduled + self.immediate - } -} - -/// 令牌刷新任务错误。 -#[derive(Debug, Error)] -pub enum TokenRefreshServiceError { - /// 账号存储读取失败。 - #[error("failed to list accounts for token refresh: {0}")] - AccountStore(#[from] crate::upstream::accounts::store::AccountStoreError), - /// 账号存储写入失败。 - #[error("failed to persist token refresh result: {0}")] - Store(#[from] SqliteAccountStoreError), - /// 刷新租约存储失败。 - #[error("failed to coordinate token refresh lease: {0}")] - Lease(#[from] RefreshLeaseStoreError), - /// core 刷新调度失败。 - #[error("token refresh scheduler failed: {0}")] - Refresh(#[from] RefreshError), - /// 刷新期间账号不存在。 - #[error("account disappeared during token refresh: {0}")] - AccountNotFound(String), -} - -/// 令牌刷新任务结果。 -pub type TokenRefreshServiceResult = Result; +mod types; +use types::*; async fn persist_token_update( - store: &SqliteAccountStore, + store: &PgAccountStore, account: &Account, ) -> TokenRefreshServiceResult<()> { let updated = store @@ -755,7 +685,7 @@ async fn persist_token_update( } async fn persist_status( - store: &SqliteAccountStore, + store: &PgAccountStore, account_id: &str, status: AccountStatus, ) -> TokenRefreshServiceResult<()> { @@ -770,7 +700,7 @@ async fn persist_status( } async fn persist_next_refresh_at( - store: &SqliteAccountStore, + store: &PgAccountStore, account_id: &str, next_refresh_at: Option>, ) -> TokenRefreshServiceResult<()> { @@ -845,44 +775,3 @@ fn same_refresh_time(left: Option>, right: Option>) _ => false, } } - -fn stored_account_to_refresh_account(stored: StoredAccount) -> Account { - Account { - id: stored.id, - email: stored.email, - account_id: stored.account_id, - user_id: stored.user_id, - label: stored.label, - plan_type: stored.plan_type, - access_token: stored.access_token.expose_secret().to_string(), - refresh_token: stored - .refresh_token - .map(|token| token.expose_secret().to_string()), - access_token_expires_at: stored.access_token_expires_at, - next_refresh_at: stored.next_refresh_at, - status: stored.status, - quota_limit_reached: false, - quota_verify_required: false, - quota_cooldown_until: None, - cloudflare_cooldown_until: None, - request_count: 0, - empty_response_count: 0, - image_input_tokens: 0, - image_output_tokens: 0, - image_request_count: 0, - image_request_failed_count: 0, - window_request_count: 0, - window_input_tokens: 0, - window_output_tokens: 0, - window_cached_tokens: 0, - window_image_input_tokens: 0, - window_image_output_tokens: 0, - window_image_request_count: 0, - window_image_request_failed_count: 0, - window_started_at: None, - window_reset_at: None, - limit_window_seconds: None, - added_at: stored.added_at, - last_used_at: None, - } -} diff --git a/backend/src/accounts/refresh/runtime/types.rs b/backend/src/accounts/refresh/runtime/types.rs new file mode 100644 index 000000000..8607e0446 --- /dev/null +++ b/backend/src/accounts/refresh/runtime/types.rs @@ -0,0 +1,109 @@ +use super::*; + +pub(super) enum TokenRefreshOutcome { + Refreshed(Box), + StatusUpdated, + Skipped, + Failed, +} + +/// 单次刷新扫描摘要。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct TokenRefreshSummary { + /// 扫描账号数。 + pub scanned: usize, + /// 成功刷新 token 的账号数。 + pub refreshed: usize, + /// 仅更新状态的账号数。 + pub status_updated: usize, + /// 无需刷新的账号数。 + pub skipped: usize, + /// 传输失败账号数。 + pub failed: usize, +} + +/// 账号级刷新定时器调度摘要。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct TokenTimerSummary { + /// 扫描账号数。 + pub scanned: usize, + /// 已调度的未来刷新定时器数。 + pub scheduled: usize, + /// 已调度的立即刷新定时器数。 + pub immediate: usize, + /// 无需调度的账号数。 + pub skipped: usize, + /// 被替换的既有定时器数。 + pub replaced: usize, +} + +impl TokenTimerSummary { + /// 返回本轮调度发生变化的定时器数量。 + pub fn changed(&self) -> usize { + self.scheduled + self.immediate + } +} + +/// 令牌刷新任务错误。 +#[derive(Debug, Error)] +pub enum TokenRefreshServiceError { + /// 账号存储读取失败。 + #[error("failed to list accounts for token refresh: {0}")] + AccountStore(#[from] crate::accounts::store::AccountStoreError), + /// 账号存储写入失败。 + #[error("failed to persist token refresh result: {0}")] + Store(#[from] PgAccountStoreError), + /// 刷新租约存储失败。 + #[error("failed to coordinate token refresh lease: {0}")] + Lease(#[from] RedisRefreshLeaseStoreError), + /// core 刷新调度失败。 + #[error("token refresh scheduler failed: {0}")] + Refresh(#[from] RefreshError), + /// 刷新期间账号不存在。 + #[error("account disappeared during token refresh: {0}")] + AccountNotFound(String), +} + +/// 令牌刷新任务结果。 +pub type TokenRefreshServiceResult = Result; + +pub(super) fn stored_account_to_refresh_account(stored: StoredAccount) -> Account { + Account { + id: stored.id, + email: stored.email, + account_id: stored.account_id, + user_id: stored.user_id, + label: stored.label, + plan_type: stored.plan_type, + access_token: stored.access_token.expose_secret().to_string(), + refresh_token: stored + .refresh_token + .map(|token| token.expose_secret().to_string()), + access_token_expires_at: stored.access_token_expires_at, + next_refresh_at: stored.next_refresh_at, + status: stored.status, + quota_limit_reached: false, + quota_verify_required: false, + quota_cooldown_until: None, + cloudflare_cooldown_until: None, + request_count: 0, + empty_response_count: 0, + image_input_tokens: 0, + image_output_tokens: 0, + image_request_count: 0, + image_request_failed_count: 0, + window_request_count: 0, + window_input_tokens: 0, + window_output_tokens: 0, + window_cached_tokens: 0, + window_image_input_tokens: 0, + window_image_output_tokens: 0, + window_image_request_count: 0, + window_image_request_failed_count: 0, + window_started_at: None, + window_reset_at: None, + limit_window_seconds: None, + added_at: stored.added_at, + last_used_at: None, + } +} diff --git a/backend/src/upstream/accounts/token_refresh/mod.rs b/backend/src/accounts/refresh/service.rs similarity index 71% rename from backend/src/upstream/accounts/token_refresh/mod.rs rename to backend/src/accounts/refresh/service.rs index 4c10676d6..9fea61628 100644 --- a/backend/src/upstream/accounts/token_refresh/mod.rs +++ b/backend/src/accounts/refresh/service.rs @@ -1,27 +1,18 @@ //! 刷新租约存储、JWT 解码与账号 claims 验证。 use std::{ - collections::{hash_map::DefaultHasher, HashSet}, + collections::hash_map::DefaultHasher, hash::{Hash, Hasher}, sync::{Arc, RwLock}, }; -use async_trait::async_trait; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use chrono::{DateTime, Duration, TimeZone, Utc}; use serde_json::{Map, Value}; use tokio::sync::Semaphore; -use crate::upstream::accounts::model::{Account, AccountStatus}; - -/// 一次刷新返回的 token 对。 -#[derive(Debug, Clone)] -pub struct TokenPair { - /// 新的访问令牌。 - pub access_token: String, - /// 可选的新刷新令牌。 - pub refresh_token: Option, -} +use crate::accounts::account::{Account, AccountStatus}; +use crate::upstream::openai::token_client::{RefreshFailure, TokenPair, TokenRefresher}; // --------------------------------------------------------------------------- // JWT 解码 @@ -244,23 +235,6 @@ fn jittered_refresh_margin_seconds(account_id: &str, margin_seconds: u64) -> i64 .clamp(0.0, i64::MAX as f64) as i64 } -/// 上游刷新失败后的领域结果。 -#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] -pub enum RefreshFailure { - /// 刷新令牌无效或已过期。 - #[error("refresh token is invalid or expired")] - InvalidGrant, - /// 账号被上游封禁。 - #[error("account is banned")] - Banned, - /// 刷新请求在到达服务端前失败，可安全复用当前 refresh token 重试。 - #[error("refresh transport failed before server processing")] - RetryableTransport, - /// 刷新请求在传输层失败，refresh token 可能已经被服务端消费。 - #[error("refresh transport failed after possible server processing")] - Transport, -} - /// 调度器自身的执行错误。 #[derive(Debug, thiserror::Error)] pub enum RefreshError { @@ -275,13 +249,6 @@ pub enum RefreshError { Transport, } -/// 刷新令牌的上游端口。 -#[async_trait] -pub trait TokenRefresher: Send + Sync + 'static { - /// 使用给定刷新令牌换取新的 token 对。 - async fn refresh(&self, refresh_token: &str) -> Result; -} - /// 负责执行单账号刷新策略的调度器。 #[derive(Clone)] pub struct RefreshScheduler { @@ -395,117 +362,7 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco updated.status = match failure { RefreshFailure::InvalidGrant => AccountStatus::Expired, RefreshFailure::Banned => AccountStatus::Banned, - RefreshFailure::RetryableTransport => AccountStatus::Active, - RefreshFailure::Transport => AccountStatus::Active, + RefreshFailure::RetryableTransport | RefreshFailure::Transport => account.status, }; updated } - -// --------------------------------------------------------------------------- -// 刷新租约存储（SQLite） -// --------------------------------------------------------------------------- - -use sqlx::{QueryBuilder, Sqlite, SqlitePool}; -use thiserror::Error; - -/// SQLite 刷新租约存储结果。 -pub type RefreshLeaseStoreResult = Result; - -/// SQLite 刷新租约存储错误。 -#[derive(Debug, Error)] -pub enum RefreshLeaseStoreError { - /// SQLite 操作失败。 - #[error("sqlite refresh lease query failed: {0}")] - Sqlx(#[from] sqlx::Error), -} - -/// SQLite 刷新租约存储。 -#[derive(Clone)] -pub struct RefreshLeaseStore { - pool: SqlitePool, -} - -impl RefreshLeaseStore { - /// 构造存储。 - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - /// 尝试获取账号刷新租约。 - pub async fn try_acquire( - &self, - account_id: &str, - owner: &str, - expires_at: DateTime, - now: DateTime, - ) -> RefreshLeaseStoreResult { - let result = sqlx::query( - r" -insert into account_refresh_leases (account_id, owner, expires_at, updated_at) -values (?, ?, ?, ?) -on conflict(account_id) do update set - owner = excluded.owner, - expires_at = excluded.expires_at, - updated_at = excluded.updated_at -where account_refresh_leases.expires_at <= ? - or account_refresh_leases.owner = ? -", - ) - .bind(account_id) - .bind(owner) - .bind(expires_at.to_rfc3339()) - .bind(now.to_rfc3339()) - .bind(now.to_rfc3339()) - .bind(owner) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 释放账号刷新租约。 - pub async fn release(&self, account_id: &str, owner: &str) -> RefreshLeaseStoreResult { - let result = - sqlx::query("delete from account_refresh_leases where account_id = ? and owner = ?") - .bind(account_id) - .bind(owner) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 返回给定账号集合中仍有效的刷新租约。 - pub async fn active_account_ids( - &self, - account_ids: &[String], - now: DateTime, - ) -> RefreshLeaseStoreResult> { - if account_ids.is_empty() { - return Ok(HashSet::new()); - } - - let mut builder = QueryBuilder::::new( - "select account_id from account_refresh_leases where datetime(expires_at) > datetime(", - ); - builder.push_bind(now.to_rfc3339()); - builder.push(") and account_id in ("); - let mut separated = builder.separated(", "); - for account_id in account_ids { - separated.push_bind(account_id); - } - separated.push_unseparated(")"); - - let rows = builder - .build_query_as::<(String,)>() - .fetch_all(&self.pool) - .await?; - - Ok(rows.into_iter().map(|(account_id,)| account_id).collect()) - } -} - -mod runtime; - -pub use runtime::{ - RuntimeTokenRefreshService, TokenRefreshServiceError, TokenRefreshServiceResult, - TokenRefreshSummary, TokenTimerSummary, -}; diff --git a/backend/src/upstream/scheduler/mod.rs b/backend/src/accounts/scheduler.rs similarity index 98% rename from backend/src/upstream/scheduler/mod.rs rename to backend/src/accounts/scheduler.rs index 802cb86ed..5e914f96e 100644 --- a/backend/src/upstream/scheduler/mod.rs +++ b/backend/src/accounts/scheduler.rs @@ -19,7 +19,7 @@ use std::sync::Mutex; use chrono::{DateTime, Utc}; -use crate::upstream::accounts::model::Account; +use crate::accounts::account::Account; pub use candidates::{CandidateFilter, CandidateRequest}; pub use feedback::{FeedbackSample, FeedbackStats}; diff --git a/backend/src/upstream/scheduler/candidates.rs b/backend/src/accounts/scheduler/candidates.rs similarity index 97% rename from backend/src/upstream/scheduler/candidates.rs rename to backend/src/accounts/scheduler/candidates.rs index 240f20f9f..81b3c27b3 100644 --- a/backend/src/upstream/scheduler/candidates.rs +++ b/backend/src/accounts/scheduler/candidates.rs @@ -17,8 +17,8 @@ use chrono::{DateTime, Utc}; -use crate::upstream::accounts::model::{Account, AccountStatus}; -use crate::upstream::accounts::service::AccountService; +use crate::accounts::account::{Account, AccountStatus}; +use crate::accounts::service::AccountService; /// 候选过滤所需的账号池策略参数（只读视图）。 pub struct CandidateFilter<'a> { diff --git a/backend/src/upstream/scheduler/feedback.rs b/backend/src/accounts/scheduler/feedback.rs similarity index 100% rename from backend/src/upstream/scheduler/feedback.rs rename to backend/src/accounts/scheduler/feedback.rs diff --git a/backend/src/accounts/scheduler/strategy/mod.rs b/backend/src/accounts/scheduler/strategy/mod.rs new file mode 100644 index 000000000..ba837de05 --- /dev/null +++ b/backend/src/accounts/scheduler/strategy/mod.rs @@ -0,0 +1,13 @@ +//! 账号调度策略集合。 + +pub mod quota_reset; +pub mod round_robin; +pub mod smart; +pub mod sticky; +mod types; + +pub use smart::{ScoreBreakdown, ScoreWeights}; +pub(crate) use types::{ + account_window_token_count, compare_last_used, compare_window_reset, select_by, +}; +pub use types::{select, RotationStrategy, SelectionInput}; diff --git a/backend/src/upstream/scheduler/strategy/quota_reset.rs b/backend/src/accounts/scheduler/strategy/quota_reset.rs similarity index 96% rename from backend/src/upstream/scheduler/strategy/quota_reset.rs rename to backend/src/accounts/scheduler/strategy/quota_reset.rs index dbe9d526a..b8d143e63 100644 --- a/backend/src/upstream/scheduler/strategy/quota_reset.rs +++ b/backend/src/accounts/scheduler/strategy/quota_reset.rs @@ -10,7 +10,7 @@ use std::cmp::Ordering; -use crate::upstream::accounts::model::Account; +use crate::accounts::account::Account; use super::{compare_last_used, compare_window_reset, select_by, SelectionInput}; diff --git a/backend/src/upstream/scheduler/strategy/round_robin.rs b/backend/src/accounts/scheduler/strategy/round_robin.rs similarity index 92% rename from backend/src/upstream/scheduler/strategy/round_robin.rs rename to backend/src/accounts/scheduler/strategy/round_robin.rs index 3d3db1cc2..a3f89e41b 100644 --- a/backend/src/upstream/scheduler/strategy/round_robin.rs +++ b/backend/src/accounts/scheduler/strategy/round_robin.rs @@ -3,7 +3,7 @@ //! 最简单的策略：忽略账号负载与用量，纯粹按候选集合的顺序用轮转游标依次选中， //! 保证请求在候选账号间均匀轮转。 -use crate::upstream::accounts::model::Account; +use crate::accounts::account::Account; use super::SelectionInput; diff --git a/backend/src/upstream/scheduler/strategy/smart.rs b/backend/src/accounts/scheduler/strategy/smart.rs similarity index 98% rename from backend/src/upstream/scheduler/strategy/smart.rs rename to backend/src/accounts/scheduler/strategy/smart.rs index 9f4d730b4..483afc6a6 100644 --- a/backend/src/upstream/scheduler/strategy/smart.rs +++ b/backend/src/accounts/scheduler/strategy/smart.rs @@ -9,8 +9,8 @@ //! 已封顶（`quota_limit_reached`）的账号会被施加一个强惩罚，确保只要还有未封顶的候选就 //! 不会选中它。并列最高分的账号之间由 [`super::select_by`] 用轮转游标破并列。 -use crate::upstream::accounts::model::Account; -use crate::upstream::scheduler::feedback::{FeedbackSample, FeedbackStats}; +use crate::accounts::account::Account; +use crate::accounts::scheduler::feedback::{FeedbackSample, FeedbackStats}; use super::{account_window_token_count, SelectionInput}; diff --git a/backend/src/upstream/scheduler/strategy/sticky.rs b/backend/src/accounts/scheduler/strategy/sticky.rs similarity index 95% rename from backend/src/upstream/scheduler/strategy/sticky.rs rename to backend/src/accounts/scheduler/strategy/sticky.rs index c57caa23a..38b118c0d 100644 --- a/backend/src/upstream/scheduler/strategy/sticky.rs +++ b/backend/src/accounts/scheduler/strategy/sticky.rs @@ -5,7 +5,7 @@ use std::cmp::Ordering; -use crate::upstream::accounts::model::Account; +use crate::accounts::account::Account; use super::{compare_last_used, SelectionInput}; diff --git a/backend/src/upstream/scheduler/strategy/mod.rs b/backend/src/accounts/scheduler/strategy/types.rs similarity index 90% rename from backend/src/upstream/scheduler/strategy/mod.rs rename to backend/src/accounts/scheduler/strategy/types.rs index 0e82793ca..3b17eb940 100644 --- a/backend/src/upstream/scheduler/strategy/mod.rs +++ b/backend/src/accounts/scheduler/strategy/types.rs @@ -16,16 +16,11 @@ use std::cmp::Ordering; use chrono::{DateTime, Utc}; +use crate::accounts::account::Account; +use crate::accounts::scheduler::feedback::FeedbackStats; use crate::infra::time::parse_rfc3339_utc; -use crate::upstream::accounts::model::Account; -use crate::upstream::scheduler::feedback::FeedbackStats; -pub mod quota_reset; -pub mod round_robin; -pub mod smart; -pub mod sticky; - -pub use smart::{ScoreBreakdown, ScoreWeights}; +use super::smart::ScoreWeights; /// 账号轮转策略。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -80,10 +75,10 @@ pub fn select( cursor: &mut usize, ) -> Option { match strategy { - RotationStrategy::Smart => smart::select(input, cursor), - RotationStrategy::QuotaResetPriority => quota_reset::select(input, cursor), - RotationStrategy::RoundRobin => round_robin::select(input, cursor), - RotationStrategy::Sticky => sticky::select(input, cursor), + RotationStrategy::Smart => super::smart::select(input, cursor), + RotationStrategy::QuotaResetPriority => super::quota_reset::select(input, cursor), + RotationStrategy::RoundRobin => super::round_robin::select(input, cursor), + RotationStrategy::Sticky => super::sticky::select(input, cursor), } } diff --git a/backend/src/upstream/accounts/service.rs b/backend/src/accounts/service.rs similarity index 94% rename from backend/src/upstream/accounts/service.rs rename to backend/src/accounts/service.rs index 9ae052750..0cc64956d 100644 --- a/backend/src/upstream/accounts/service.rs +++ b/backend/src/accounts/service.rs @@ -2,7 +2,7 @@ use chrono::{DateTime, Utc}; -use crate::upstream::accounts::model::Account; +use crate::accounts::account::Account; /// 账号服务。 #[derive(Debug, Clone, Default)] diff --git a/backend/src/accounts/store.rs b/backend/src/accounts/store.rs new file mode 100644 index 000000000..3e3a60630 --- /dev/null +++ b/backend/src/accounts/store.rs @@ -0,0 +1,477 @@ +//! PostgreSQL 账号仓储适配器。 + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use secrecy::{ExposeSecret, SecretString}; +use serde_json::Value; +use sqlx::{PgPool, Postgres, QueryBuilder, Row}; +use thiserror::Error; + +use crate::accounts::account::{Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta}; +use crate::accounts::window::should_reset_usage_window; +use crate::infra::{ + json::{decode_cursor, page_offset, NumberedPage, Page}, + time::parse_rfc3339_utc, +}; + +mod queries; +mod rows; + +use queries::*; +use rows::{ + count_account_metadata, get_pool_account, list_pool_accounts, map_account_store_error, + metadata_from_row, optional_positive_i64_to_u64, optional_update_value, pg_usage_delta, + push_account_metadata_search, quota_plan_type, quota_snapshot_from_row, status_to_db, + stored_account_from_row, to_page, u64_to_i64_saturating, +}; + +// ============================================================================ +// 错误类型 +// ============================================================================ + +/// PostgreSQL 账号仓储错误。 +#[derive(Debug, Error)] +pub enum PgAccountStoreError { + /// 数据库错误。 + #[error("PostgreSQL account store database error: {0}")] + Database(#[from] sqlx::Error), + /// 时间格式错误。 + #[error("account store timestamp error: {0}")] + Timestamp(#[from] chrono::ParseError), + #[error("account store JSON error: {0}")] + Json(#[from] serde_json::Error), + /// 账号状态非法。 + #[error("PostgreSQL account store status error: {0}")] + InvalidStatus(String), + /// 分页游标非法。 + #[error("invalid account pagination cursor")] + InvalidCursor, +} + +/// PostgreSQL 账号仓储结果。 +pub type PgAccountStoreResult = Result; + +// ============================================================================ +// 数据类型 +// ============================================================================ + +/// 新建账号数据。 +#[derive(Debug)] +pub struct NewAccount { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// access token 明文。 + pub access_token: SecretString, + /// refresh token 明文。 + pub refresh_token: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 初始状态。 + pub status: AccountStatus, + /// 账号添加时间。 + pub added_at: Option>, +} + +/// 通过 JWT claims 更新已有账号。 +#[derive(Debug)] +pub struct AccountClaimsUpdate { + /// 展示邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 计划类型。 + pub plan_type: Option, + /// 新 access token。 + pub access_token: SecretString, + /// 可选的新 refresh token。 + pub refresh_token: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 下一次刷新时间。 + pub next_refresh_at: Option>, + /// 更新后的状态。 + pub status: AccountStatus, +} + +/// 通过导入数据更新已有账号。 +#[derive(Debug)] +pub struct ImportedAccountUpdate { + /// 新建账号数据。 + pub account: NewAccount, + /// 缓存配额 JSON。 + pub quota_json: Option, + /// 配额抓取时间。 + pub quota_fetched_at: Option>, + /// 是否需要执行额外配额校验。 + pub quota_verify_required: bool, +} + +/// 存储中的账号完整记录。 +#[derive(Debug, Clone)] +pub struct StoredAccount { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// access token 明文。 + pub access_token: SecretString, + /// refresh token 明文。 + pub refresh_token: Option, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 下一次刷新时间。 + pub next_refresh_at: Option>, + /// 账号状态。 + pub status: AccountStatus, + /// 添加时间。 + pub added_at: String, + /// 更新时间。 + pub updated_at: String, +} + +/// 存储中的账号元数据（不含 token）。 +#[derive(Debug, Clone)] +pub struct StoredAccountMetadata { + /// 账号 ID。 + pub id: String, + /// 邮箱。 + pub email: Option, + /// ChatGPT 账号 ID。 + pub account_id: Option, + /// ChatGPT 用户 ID。 + pub user_id: Option, + /// 管理标签。 + pub label: Option, + /// 计划类型。 + pub plan_type: Option, + /// 是否保存了 refresh token。 + pub has_refresh_token: bool, + /// access token 过期时间。 + pub access_token_expires_at: Option>, + /// 账号状态。 + pub status: AccountStatus, + /// 添加时间。 + pub added_at: String, + /// 更新时间。 + pub updated_at: String, +} + +#[derive(Debug, Clone, Default)] +pub struct AccountMetadataUpdate { + pub email: Option>, + pub account_id: Option>, + pub user_id: Option>, + pub label: Option>, + pub plan_type: Option>, + pub status: Option, +} + +impl AccountMetadataUpdate { + pub fn any(&self) -> bool { + self.email.is_some() + || self.account_id.is_some() + || self.user_id.is_some() + || self.label.is_some() + || self.plan_type.is_some() + || self.status.is_some() + } +} + +/// 账号配额快照。 +#[derive(Debug, Clone)] +pub struct AccountQuotaSnapshot { + /// 账号 ID。 + pub account_id: String, + /// 邮箱。 + pub email: Option, + /// 配额 JSON。 + pub quota_json: String, + /// 配额抓取时间。 + pub quota_fetched_at: Option>, +} + +/// 用量增量（PostgreSQL 内部表示）。 +#[derive(Debug, Clone, Copy, Default)] +pub struct UsageDelta { + request_count: i64, + empty_response_count: i64, + input_tokens: i64, + output_tokens: i64, + cached_tokens: i64, + reasoning_tokens: i64, + total_tokens: i64, + image_input_tokens: i64, + image_output_tokens: i64, + image_request_count: i64, + image_request_failed_count: i64, + window_request_count: i64, + window_input_tokens: i64, + window_output_tokens: i64, + window_cached_tokens: i64, + window_image_input_tokens: i64, + window_image_output_tokens: i64, + window_image_request_count: i64, + window_image_request_failed_count: i64, +} + +// ============================================================================ +// AccountStore 端口 trait +// ============================================================================ + +/// 账号存储错误。 +#[derive(Debug, Error)] +pub enum AccountStoreError { + /// 底层存储失败。 + #[error("account store operation failed: {message}")] + OperationFailed { + /// 错误说明。 + message: String, + }, +} + +/// 账号存储结果类型。 +pub type AccountStoreResult = Result; + +/// 提供运行时账号列表的端口。 +#[async_trait] +pub trait AccountStore: Send + Sync + 'static { + /// 列出当前账号池可见的账号。 + async fn list_pool_accounts(&self) -> AccountStoreResult>; + + /// 读取单个账号池账号快照。 + async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult> { + let accounts = self.list_pool_accounts().await?; + Ok(accounts + .into_iter() + .find(|account| account.id == account_id)) + } + + /// 标记账号进入配额冷却期。 + async fn mark_quota_limited_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> AccountStoreResult; + + /// 标记账号进入 Cloudflare 冷却期。 + async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> AccountStoreResult; + + /// 更新账号状态。 + async fn set_status(&self, account_id: &str, status: AccountStatus) + -> AccountStoreResult; + + /// 记录账号用量增量。 + async fn record_usage_delta( + &self, + account_id: &str, + usage: AccountUsageDelta, + ) -> AccountStoreResult<()>; + + /// 记录账号模型维度用量增量。 + async fn record_model_usage_delta( + &self, + account_id: &str, + model: &str, + usage: AccountModelUsageDelta, + ) -> AccountStoreResult<()>; + + /// 读取账号当前配额 JSON。 + async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult>; + + /// 应用已经验证过的账号配额快照。 + async fn apply_quota_snapshot( + &self, + account_id: &str, + quota_json: &str, + limit_reached: bool, + cooldown_until: Option>, + ) -> AccountStoreResult; + + /// 同步运行时自然刷新出来的账号状态。 + async fn sync_runtime_account_state( + &self, + account: &Account, + sync_usage_window: bool, + ) -> AccountStoreResult; + + /// 同步账号当前 rate-limit 统计窗口。 + async fn sync_rate_limit_window( + &self, + account_id: &str, + reset_at: DateTime, + limit_window_seconds: Option, + ) -> AccountStoreResult<()>; + + /// 记录账号被用于一次外部请求。 + async fn record_request(&self, account_id: &str) -> AccountStoreResult<()> { + self.record_usage_delta( + account_id, + AccountUsageDelta { + requests: 1, + ..AccountUsageDelta::default() + }, + ) + .await + } +} + +#[derive(Clone)] +pub struct PgAccountStore { + pool: PgPool, +} + +mod write; + +impl PgAccountStore { + /// 设置账号标签。 + pub async fn set_label( + &self, + account_id: &str, + label: Option, + ) -> Result { + let updated_at = Utc::now(); + let result = sqlx::query("update accounts set label = $1, updated_at = $2 where id = $3") + .bind(&label) + .bind(updated_at) + .bind(account_id) + .execute(&self.pool) + .await + .map_err(PgAccountStoreError::Database)?; + Ok(result.rows_affected() > 0) + } +} + +#[async_trait] +impl AccountStore for PgAccountStore { + async fn list_pool_accounts(&self) -> AccountStoreResult> { + list_pool_accounts(self) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult> { + get_pool_account(self, account_id) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn mark_quota_limited_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> AccountStoreResult { + PgAccountStore::mark_quota_limited_until(self, account_id, cooldown_until) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> AccountStoreResult { + PgAccountStore::set_cloudflare_cooldown_until(self, account_id, cooldown_until) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn set_status( + &self, + account_id: &str, + status: AccountStatus, + ) -> AccountStoreResult { + PgAccountStore::set_status(self, account_id, status) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn record_usage_delta( + &self, + account_id: &str, + usage: AccountUsageDelta, + ) -> AccountStoreResult<()> { + self.record_usage(account_id, pg_usage_delta(usage)) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn record_model_usage_delta( + &self, + account_id: &str, + model: &str, + usage: AccountModelUsageDelta, + ) -> AccountStoreResult<()> { + self.record_model_usage(account_id, model, usage) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult> { + PgAccountStore::get_quota_json(self, account_id) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn apply_quota_snapshot( + &self, + account_id: &str, + quota_json: &str, + limit_reached: bool, + cooldown_until: Option>, + ) -> AccountStoreResult { + PgAccountStore::apply_quota_snapshot( + self, + account_id, + quota_json, + limit_reached, + cooldown_until, + ) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn sync_runtime_account_state( + &self, + account: &Account, + sync_usage_window: bool, + ) -> AccountStoreResult { + PgAccountStore::sync_runtime_account_state(self, account, sync_usage_window) + .await + .map_err(|error| map_account_store_error(&error)) + } + + async fn sync_rate_limit_window( + &self, + account_id: &str, + reset_at: DateTime, + limit_window_seconds: Option, + ) -> AccountStoreResult<()> { + PgAccountStore::sync_rate_limit_window(self, account_id, reset_at, limit_window_seconds) + .await + .map_err(|error| map_account_store_error(&error)) + } +} diff --git a/backend/src/upstream/accounts/store/queries.rs b/backend/src/accounts/store/queries.rs similarity index 63% rename from backend/src/upstream/accounts/store/queries.rs rename to backend/src/accounts/store/queries.rs index 1a652896e..a669bb6a2 100644 --- a/backend/src/upstream/accounts/store/queries.rs +++ b/backend/src/accounts/store/queries.rs @@ -1,4 +1,4 @@ -//! SQLite 账号仓储 SQL。 +//! PostgreSQL 账号仓储 SQL。 // ============================================================================ // SQL 常量 @@ -43,7 +43,7 @@ select au.last_used_at as usage_last_used_at from accounts a left join account_usage au on au.account_id = a.id -order by a.rowid asc"; +order by a.added_at asc, a.id asc"; pub(super) const GET_POOL_ACCOUNT_SQL: &str = r" select @@ -84,7 +84,7 @@ select au.last_used_at as usage_last_used_at from accounts a left join account_usage au on au.account_id = a.id -where a.id = ?"; +where a.id = $1"; pub(super) const INSERT_ACCOUNT_SQL: &str = r" insert into accounts ( @@ -100,7 +100,7 @@ insert into accounts ( status, added_at, updated_at -) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; +) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)"; pub(super) const SELECT_STORED_ACCOUNT_BY_ID_SQL: &str = r" select @@ -118,7 +118,7 @@ select added_at, updated_at from accounts -where id = ?"; +where id = $1"; pub(super) const SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL: &str = r" select @@ -136,8 +136,8 @@ select added_at, updated_at from accounts -where chatgpt_account_id = ? - and ((chatgpt_user_id is null and ? is null) or chatgpt_user_id = ?) +where chatgpt_account_id = $1 + and ((chatgpt_user_id is null and $2 is null) or chatgpt_user_id = $3) order by added_at asc limit 1"; @@ -149,31 +149,31 @@ select chatgpt_user_id as user_id, label, plan_type, - case when refresh_token is not null and refresh_token <> '' then 1 else 0 end as has_refresh_token, + (refresh_token is not null and refresh_token <> '') as has_refresh_token, access_token_expires_at, status, added_at, updated_at from accounts -where id = ?"; +where id = $1"; pub(super) const UPDATE_ACCOUNT_METADATA_SQL: &str = r" update accounts set - email = case when ? then ? else email end, - chatgpt_account_id = case when ? then ? else chatgpt_account_id end, - chatgpt_user_id = case when ? then ? else chatgpt_user_id end, - label = case when ? then ? else label end, - plan_type = case when ? then ? else plan_type end, + email = case when $1 then $2 else email end, + chatgpt_account_id = case when $3 then $4 else chatgpt_account_id end, + chatgpt_user_id = case when $5 then $6 else chatgpt_user_id end, + label = case when $7 then $8 else label end, + plan_type = case when $9 then $10 else plan_type end, status = case - when ? then case - when ? = 'active' and quota_limit_reached = 1 then 'quota_exhausted' - else ? + when $11 then case + when $12 = 'active' and quota_limit_reached then 'quota_exhausted' + else $13 end else status end, - updated_at = ? -where id = ?"; + updated_at = $14 +where id = $15"; pub(super) const LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL: &str = r" select @@ -191,10 +191,10 @@ select added_at, updated_at from accounts -where added_at < ? - or (added_at = ? and id < ?) +where added_at < $1 + or (added_at = $2 and id < $3) order by added_at desc, id desc -limit ?"; +limit $4"; pub(super) const LIST_STORED_ACCOUNTS_SQL: &str = r" select @@ -213,7 +213,7 @@ select updated_at from accounts order by added_at desc, id desc -limit ?"; +limit $1"; pub(super) const LIST_ACCOUNT_METADATA_AFTER_CURSOR_SQL: &str = r" select @@ -223,16 +223,16 @@ select chatgpt_user_id as user_id, label, plan_type, - case when refresh_token is not null and refresh_token <> '' then 1 else 0 end as has_refresh_token, + (refresh_token is not null and refresh_token <> '') as has_refresh_token, access_token_expires_at, status, added_at, updated_at from accounts -where added_at < ? - or (added_at = ? and id < ?) +where added_at < $1 + or (added_at = $2 and id < $3) order by added_at desc, id desc -limit ?"; +limit $4"; pub(super) const LIST_ACCOUNT_METADATA_SQL: &str = r" select @@ -242,14 +242,14 @@ select chatgpt_user_id as user_id, label, plan_type, - case when refresh_token is not null and refresh_token <> '' then 1 else 0 end as has_refresh_token, + (refresh_token is not null and refresh_token <> '') as has_refresh_token, access_token_expires_at, status, added_at, updated_at from accounts order by added_at desc, id desc -limit ?"; +limit $1"; pub(super) const LIST_ACCOUNT_METADATA_SELECT_SQL: &str = r" select @@ -259,7 +259,7 @@ select chatgpt_user_id as user_id, label, plan_type, - case when refresh_token is not null and refresh_token <> '' then 1 else 0 end as has_refresh_token, + (refresh_token is not null and refresh_token <> '') as has_refresh_token, access_token_expires_at, status, added_at, @@ -290,27 +290,27 @@ insert into account_usage ( window_image_request_failed_count, window_started_at, last_used_at -) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22) on conflict(account_id) do update set - request_count = request_count + excluded.request_count, - empty_response_count = empty_response_count + excluded.empty_response_count, - input_tokens = input_tokens + excluded.input_tokens, - output_tokens = output_tokens + excluded.output_tokens, - cached_tokens = cached_tokens + excluded.cached_tokens, - reasoning_tokens = reasoning_tokens + excluded.reasoning_tokens, - total_tokens = total_tokens + excluded.total_tokens, - image_input_tokens = image_input_tokens + excluded.image_input_tokens, - image_output_tokens = image_output_tokens + excluded.image_output_tokens, - image_request_count = image_request_count + excluded.image_request_count, - image_request_failed_count = image_request_failed_count + excluded.image_request_failed_count, - window_request_count = window_request_count + excluded.window_request_count, - window_input_tokens = window_input_tokens + excluded.window_input_tokens, - window_output_tokens = window_output_tokens + excluded.window_output_tokens, - window_cached_tokens = window_cached_tokens + excluded.window_cached_tokens, - window_image_input_tokens = window_image_input_tokens + excluded.window_image_input_tokens, - window_image_output_tokens = window_image_output_tokens + excluded.window_image_output_tokens, - window_image_request_count = window_image_request_count + excluded.window_image_request_count, - window_image_request_failed_count = window_image_request_failed_count + excluded.window_image_request_failed_count, + request_count = account_usage.request_count + excluded.request_count, + empty_response_count = account_usage.empty_response_count + excluded.empty_response_count, + input_tokens = account_usage.input_tokens + excluded.input_tokens, + output_tokens = account_usage.output_tokens + excluded.output_tokens, + cached_tokens = account_usage.cached_tokens + excluded.cached_tokens, + reasoning_tokens = account_usage.reasoning_tokens + excluded.reasoning_tokens, + total_tokens = account_usage.total_tokens + excluded.total_tokens, + image_input_tokens = account_usage.image_input_tokens + excluded.image_input_tokens, + image_output_tokens = account_usage.image_output_tokens + excluded.image_output_tokens, + image_request_count = account_usage.image_request_count + excluded.image_request_count, + image_request_failed_count = account_usage.image_request_failed_count + excluded.image_request_failed_count, + window_request_count = account_usage.window_request_count + excluded.window_request_count, + window_input_tokens = account_usage.window_input_tokens + excluded.window_input_tokens, + window_output_tokens = account_usage.window_output_tokens + excluded.window_output_tokens, + window_cached_tokens = account_usage.window_cached_tokens + excluded.window_cached_tokens, + window_image_input_tokens = account_usage.window_image_input_tokens + excluded.window_image_input_tokens, + window_image_output_tokens = account_usage.window_image_output_tokens + excluded.window_image_output_tokens, + window_image_request_count = account_usage.window_image_request_count + excluded.window_image_request_count, + window_image_request_failed_count = account_usage.window_image_request_failed_count + excluded.window_image_request_failed_count, window_started_at = coalesce(account_usage.window_started_at, excluded.window_started_at), last_used_at = excluded.last_used_at"; @@ -324,13 +324,13 @@ insert into account_model_usage ( output_tokens, cached_tokens, last_used_at -) values (?, ?, ?, ?, ?, ?, ?, ?) +) values ($1, $2, $3, $4, $5, $6, $7, $8) on conflict(account_id, model) do update set - request_count = request_count + excluded.request_count, - error_count = error_count + excluded.error_count, - input_tokens = input_tokens + excluded.input_tokens, - output_tokens = output_tokens + excluded.output_tokens, - cached_tokens = cached_tokens + excluded.cached_tokens, + request_count = account_model_usage.request_count + excluded.request_count, + error_count = account_model_usage.error_count + excluded.error_count, + input_tokens = account_model_usage.input_tokens + excluded.input_tokens, + output_tokens = account_model_usage.output_tokens + excluded.output_tokens, + cached_tokens = account_model_usage.cached_tokens + excluded.cached_tokens, last_used_at = excluded.last_used_at"; pub(super) const LIST_QUOTA_SNAPSHOTS_SQL: &str = r" @@ -341,42 +341,41 @@ select quota_fetched_at from accounts where quota_json is not null - and trim(quota_json) <> '' -order by coalesce(quota_fetched_at, '') desc, id desc"; +order by quota_fetched_at desc nulls last, id desc"; pub(super) const UPDATE_QUOTA_JSON_SQL: &str = r" update accounts set - quota_json = ?, - quota_fetched_at = ?, - plan_type = coalesce(?, plan_type), - updated_at = ? -where id = ?"; + quota_json = $1, + quota_fetched_at = $2, + plan_type = coalesce($3, plan_type), + updated_at = $4 +where id = $5"; pub(super) const APPLY_QUOTA_SNAPSHOT_SQL: &str = r" update accounts set - quota_json = ?, - quota_fetched_at = ?, - plan_type = coalesce(?, plan_type), + quota_json = $1, + quota_fetched_at = $2, + plan_type = coalesce($3, plan_type), status = case when status in ('disabled', 'banned', 'expired') then status - when ? = 1 then 'quota_exhausted' + when $4 then 'quota_exhausted' when status = 'quota_exhausted' then 'active' else status end, - quota_limit_reached = ?, - quota_verify_required = 0, - quota_cooldown_until = ?, - updated_at = ? -where id = ?"; + quota_limit_reached = $5, + quota_verify_required = false, + quota_cooldown_until = $6, + updated_at = $7 +where id = $8"; pub(super) const SELECT_RATE_LIMIT_WINDOW_SQL: &str = r" select window_reset_at, limit_window_seconds from account_usage -where account_id = ?"; +where account_id = $1"; pub(super) const SYNC_RATE_LIMIT_WINDOW_RESET_SQL: &str = r" insert into account_usage ( @@ -392,7 +391,7 @@ insert into account_usage ( window_started_at, window_reset_at, limit_window_seconds -) values (?, 0, 0, 0, 0, 0, 0, 0, 0, ?, ?, ?) +) values ($1, 0, 0, 0, 0, 0, 0, 0, 0, $2, $3, $4) on conflict(account_id) do update set window_request_count = 0, window_input_tokens = 0, @@ -411,7 +410,7 @@ insert into account_usage ( account_id, window_reset_at, limit_window_seconds -) values (?, ?, ?) +) values ($1, $2, $3) on conflict(account_id) do update set window_reset_at = excluded.window_reset_at, limit_window_seconds = coalesce(excluded.limit_window_seconds, account_usage.limit_window_seconds)"; @@ -423,11 +422,11 @@ set when status in ('disabled', 'banned', 'expired') then status else 'quota_exhausted' end, - quota_limit_reached = 1, - quota_verify_required = 0, - quota_cooldown_until = ?, - updated_at = ? -where id = ?"; + quota_limit_reached = true, + quota_verify_required = false, + quota_cooldown_until = $1, + updated_at = $2 +where id = $3"; pub(super) const SYNC_RUNTIME_ACCOUNT_STATE_SQL: &str = r" update accounts @@ -436,30 +435,30 @@ set when status in ('disabled', 'banned') then status when ( case - when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached - else ? + when not $1 and quota_cooldown_until is not null and quota_cooldown_until > $2 then quota_limit_reached + else $3 end - ) = 1 then 'quota_exhausted' - else ? + ) then 'quota_exhausted' + else $4 end, quota_limit_reached = case - when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached - else ? + when not $5 and quota_cooldown_until is not null and quota_cooldown_until > $6 then quota_limit_reached + else $7 end, quota_verify_required = case - when ? = 1 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_verify_required - else ? + when $8 and quota_cooldown_until is not null and quota_cooldown_until > $9 then quota_verify_required + else $10 end, quota_cooldown_until = case - when ? is null and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_cooldown_until - else ? + when $11 is null and quota_cooldown_until is not null and quota_cooldown_until > $12 then quota_cooldown_until + else $13 end, cloudflare_cooldown_until = case - when ? is null and cloudflare_cooldown_until is not null and cloudflare_cooldown_until > ? then cloudflare_cooldown_until - else ? + when $14 is null and cloudflare_cooldown_until is not null and cloudflare_cooldown_until > $15 then cloudflare_cooldown_until + else $16 end, - updated_at = ? -where id = ?"; + updated_at = $17 +where id = $18"; pub(super) const SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL: &str = r" insert into account_usage ( @@ -475,7 +474,7 @@ insert into account_usage ( window_started_at, window_reset_at, limit_window_seconds -) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) +) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) on conflict(account_id) do update set window_request_count = excluded.window_request_count, window_input_tokens = excluded.window_input_tokens, @@ -490,108 +489,108 @@ on conflict(account_id) do update set limit_window_seconds = excluded.limit_window_seconds where account_usage.window_reset_at is null or account_usage.window_reset_at <= excluded.window_reset_at - or account_usage.window_reset_at <= ?"; + or account_usage.window_reset_at <= $13"; pub(super) const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r" update accounts set - cloudflare_cooldown_until = ?, - updated_at = ? -where id = ?"; + cloudflare_cooldown_until = $1, + updated_at = $2 +where id = $3"; pub(super) const UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL: &str = r" update accounts set - email = coalesce(?, email), - chatgpt_account_id = coalesce(?, chatgpt_account_id), - chatgpt_user_id = coalesce(?, chatgpt_user_id), - plan_type = coalesce(?, plan_type), - access_token = ?, - refresh_token = ?, - access_token_expires_at = ?, - next_refresh_at = ?, + email = coalesce($1, email), + chatgpt_account_id = coalesce($2, chatgpt_account_id), + chatgpt_user_id = coalesce($3, chatgpt_user_id), + plan_type = coalesce($4, plan_type), + access_token = $5, + refresh_token = $6, + access_token_expires_at = $7, + next_refresh_at = $8, status = case when status in ('disabled', 'banned') then status - when quota_limit_reached = 1 then 'quota_exhausted' - else ? + when quota_limit_reached then 'quota_exhausted' + else $9 end, - updated_at = ? -where id = ?"; + updated_at = $10 +where id = $11"; pub(super) const UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL: &str = r" update accounts set - email = coalesce(?, email), - chatgpt_account_id = coalesce(?, chatgpt_account_id), - chatgpt_user_id = coalesce(?, chatgpt_user_id), - plan_type = coalesce(?, plan_type), - access_token = ?, - access_token_expires_at = ?, - next_refresh_at = ?, + email = coalesce($1, email), + chatgpt_account_id = coalesce($2, chatgpt_account_id), + chatgpt_user_id = coalesce($3, chatgpt_user_id), + plan_type = coalesce($4, plan_type), + access_token = $5, + access_token_expires_at = $6, + next_refresh_at = $7, status = case when status in ('disabled', 'banned') then status - when quota_limit_reached = 1 then 'quota_exhausted' - else ? + when quota_limit_reached then 'quota_exhausted' + else $8 end, - updated_at = ? -where id = ?"; + updated_at = $9 +where id = $10"; pub(super) const SET_NEXT_REFRESH_AT_SQL: &str = r" update accounts set - next_refresh_at = ?, - updated_at = ? -where id = ?"; + next_refresh_at = $1, + updated_at = $2 +where id = $3"; pub(super) const UPDATE_IMPORTED_ACCOUNT_WITH_REFRESH_SQL: &str = r" update accounts set - email = ?, - chatgpt_account_id = ?, - chatgpt_user_id = ?, - label = ?, - plan_type = ?, - access_token = ?, - refresh_token = ?, - access_token_expires_at = ?, - status = ?, - quota_json = coalesce(?, quota_json), - quota_fetched_at = case when ? is null then quota_fetched_at else ? end, - quota_limit_reached = 0, + email = $1, + chatgpt_account_id = $2, + chatgpt_user_id = $3, + label = $4, + plan_type = $5, + access_token = $6, + refresh_token = $7, + access_token_expires_at = $8, + status = $9, + quota_json = coalesce($10, quota_json), + quota_fetched_at = case when $11 is null then quota_fetched_at else $12 end, + quota_limit_reached = false, quota_cooldown_until = null, - quota_verify_required = ?, - updated_at = ? -where id = ?"; + quota_verify_required = $13, + updated_at = $14 +where id = $15"; pub(super) const UPDATE_IMPORTED_ACCOUNT_PRESERVING_REFRESH_SQL: &str = r" update accounts set - email = ?, - chatgpt_account_id = ?, - chatgpt_user_id = ?, - label = ?, - plan_type = ?, - access_token = ?, - access_token_expires_at = ?, - status = ?, - quota_json = coalesce(?, quota_json), - quota_fetched_at = case when ? is null then quota_fetched_at else ? end, - quota_limit_reached = 0, + email = $1, + chatgpt_account_id = $2, + chatgpt_user_id = $3, + label = $4, + plan_type = $5, + access_token = $6, + access_token_expires_at = $7, + status = $8, + quota_json = coalesce($9, quota_json), + quota_fetched_at = case when $10 is null then quota_fetched_at else $11 end, + quota_limit_reached = false, quota_cooldown_until = null, - quota_verify_required = ?, - updated_at = ? -where id = ?"; + quota_verify_required = $12, + updated_at = $13 +where id = $14"; pub(super) const APPLY_IMPORTED_QUOTA_STATE_SQL: &str = r" update accounts set - quota_json = coalesce(?, quota_json), - quota_fetched_at = case when ? is null then quota_fetched_at else ? end, - plan_type = coalesce(?, plan_type), - quota_limit_reached = 0, + quota_json = coalesce($1, quota_json), + quota_fetched_at = case when $2 is null then quota_fetched_at else $3 end, + plan_type = coalesce($4, plan_type), + quota_limit_reached = false, quota_cooldown_until = null, - quota_verify_required = ?, - updated_at = ? -where id = ?"; + quota_verify_required = $5, + updated_at = $6 +where id = $7"; -pub(super) const DELETE_ACCOUNT_SQL: &str = "delete from accounts where id = ?"; +pub(super) const DELETE_ACCOUNT_SQL: &str = "delete from accounts where id = $1"; diff --git a/backend/src/upstream/accounts/store/rows.rs b/backend/src/accounts/store/rows.rs similarity index 67% rename from backend/src/upstream/accounts/store/rows.rs rename to backend/src/accounts/store/rows.rs index f79ba8c37..6e2466324 100644 --- a/backend/src/upstream/accounts/store/rows.rs +++ b/backend/src/accounts/store/rows.rs @@ -1,29 +1,26 @@ -//! SQLite 账号仓储 row 映射与查询辅助。 +//! PostgreSQL 账号仓储 row 映射与查询辅助。 use secrecy::SecretString; use serde_json::Value; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool}; +use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row}; use crate::{ - infra::{ - format::nonnegative_i64_to_u64, json::Page, - time::parse_optional_rfc3339_utc as parse_optional_rfc3339, - }, - upstream::accounts::{ - model::{Account, AccountStatus, AccountUsageDelta}, + accounts::{ + account::{Account, AccountStatus, AccountUsageDelta}, quota::{quota_snapshot_limit_window_seconds, quota_snapshot_reset_at}, }, + infra::{format::nonnegative_i64_to_u64, json::Page}, }; use super::{ queries::{GET_POOL_ACCOUNT_SQL, LIST_POOL_ACCOUNTS_SQL}, - AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore, SqliteAccountStoreError, - SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata, UsageDelta, + AccountQuotaSnapshot, AccountStoreError, PgAccountStore, PgAccountStoreError, + PgAccountStoreResult, StoredAccount, StoredAccountMetadata, UsageDelta, }; pub(super) async fn list_pool_accounts( - store: &SqliteAccountStore, -) -> SqliteAccountStoreResult> { + store: &PgAccountStore, +) -> PgAccountStoreResult> { let rows = sqlx::query(LIST_POOL_ACCOUNTS_SQL) .fetch_all(&store.pool) .await?; @@ -37,9 +34,9 @@ pub(super) async fn list_pool_accounts( } pub(super) async fn get_pool_account( - store: &SqliteAccountStore, + store: &PgAccountStore, account_id: &str, -) -> SqliteAccountStoreResult> { +) -> PgAccountStoreResult> { let row = sqlx::query(GET_POOL_ACCOUNT_SQL) .bind(account_id) .fetch_optional(&store.pool) @@ -47,10 +44,10 @@ pub(super) async fn get_pool_account( row.map(|row| pool_account_from_row(&row)).transpose() } -pub(super) fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult { +pub(super) fn pool_account_from_row(row: &PgRow) -> PgAccountStoreResult { let quota_json = row - .get::, _>("quota_json") - .and_then(|quota_json| serde_json::from_str::("a_json).ok()); + .get::>, _>("quota_json") + .map(|quota_json| quota_json.0); let quota_window_reset_at = quota_json.as_ref().and_then(quota_snapshot_reset_at); let quota_limit_window_seconds = quota_json .as_ref() @@ -65,24 +62,13 @@ pub(super) fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult plan_type: row.get("plan_type"), access_token: row.get("access_token"), refresh_token: row.get("refresh_token"), - access_token_expires_at: parse_optional_rfc3339( - row.get::, _>("access_token_expires_at") - .as_deref(), - )?, - next_refresh_at: parse_optional_rfc3339( - row.get::, _>("next_refresh_at").as_deref(), - )?, + access_token_expires_at: row.get("access_token_expires_at"), + next_refresh_at: row.get("next_refresh_at"), status: status_from_db(&row.get::("status"))?, - quota_limit_reached: row.get::("quota_limit_reached") != 0, - quota_verify_required: row.get::("quota_verify_required") != 0, - quota_cooldown_until: parse_optional_rfc3339( - row.get::, _>("quota_cooldown_until") - .as_deref(), - )?, - cloudflare_cooldown_until: parse_optional_rfc3339( - row.get::, _>("cloudflare_cooldown_until") - .as_deref(), - )?, + quota_limit_reached: row.get("quota_limit_reached"), + quota_verify_required: row.get("quota_verify_required"), + quota_cooldown_until: row.get("quota_cooldown_until"), + cloudflare_cooldown_until: row.get("cloudflare_cooldown_until"), request_count: nonnegative_i64_to_u64(row.get::("usage_request_count")), empty_response_count: nonnegative_i64_to_u64( row.get::("usage_empty_response_count"), @@ -115,25 +101,24 @@ pub(super) fn pool_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult window_image_request_failed_count: nonnegative_i64_to_u64( row.get::("usage_window_image_request_failed_count"), ), - window_started_at: parse_optional_rfc3339( - row.get::, _>("usage_window_started_at") - .as_deref(), - )?, - window_reset_at: parse_optional_rfc3339( - row.get::, _>("usage_window_reset_at") - .as_deref(), - )? - .or(quota_window_reset_at), + window_started_at: row.get("usage_window_started_at"), + window_reset_at: row + .get::>, _>("usage_window_reset_at") + .or(quota_window_reset_at), limit_window_seconds: optional_positive_i64_to_u64( row.get::, _>("usage_limit_window_seconds"), ) .or(quota_limit_window_seconds), - added_at: row.get("added_at"), - last_used_at: row.get("usage_last_used_at"), + added_at: row + .get::, _>("added_at") + .to_rfc3339(), + last_used_at: row + .get::>, _>("usage_last_used_at") + .map(|value| value.to_rfc3339()), }) } -pub(super) fn stored_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult { +pub(super) fn stored_account_from_row(row: &PgRow) -> PgAccountStoreResult { Ok(StoredAccount { id: row.get("id"), email: row.get("email"), @@ -145,22 +130,19 @@ pub(super) fn stored_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResu refresh_token: row .get::, _>("refresh_token") .map(|token| SecretString::new(token.into())), - access_token_expires_at: parse_optional_rfc3339( - row.get::, _>("access_token_expires_at") - .as_deref(), - )?, - next_refresh_at: parse_optional_rfc3339( - row.get::, _>("next_refresh_at").as_deref(), - )?, + access_token_expires_at: row.get("access_token_expires_at"), + next_refresh_at: row.get("next_refresh_at"), status: status_from_db(&row.get::("status"))?, - added_at: row.get("added_at"), - updated_at: row.get("updated_at"), + added_at: row + .get::, _>("added_at") + .to_rfc3339(), + updated_at: row + .get::, _>("updated_at") + .to_rfc3339(), }) } -pub(super) fn metadata_from_row( - row: &SqliteRow, -) -> SqliteAccountStoreResult { +pub(super) fn metadata_from_row(row: &PgRow) -> PgAccountStoreResult { Ok(StoredAccountMetadata { id: row.get("id"), email: row.get("email"), @@ -168,14 +150,15 @@ pub(super) fn metadata_from_row( user_id: row.get("user_id"), label: row.get("label"), plan_type: row.get("plan_type"), - has_refresh_token: row.get::("has_refresh_token") != 0, - access_token_expires_at: parse_optional_rfc3339( - row.get::, _>("access_token_expires_at") - .as_deref(), - )?, + has_refresh_token: row.get("has_refresh_token"), + access_token_expires_at: row.get("access_token_expires_at"), status: status_from_db(&row.get::("status"))?, - added_at: row.get("added_at"), - updated_at: row.get("updated_at"), + added_at: row + .get::, _>("added_at") + .to_rfc3339(), + updated_at: row + .get::, _>("updated_at") + .to_rfc3339(), }) } @@ -185,7 +168,7 @@ pub(super) fn map_account_store_error(error: &impl ToString) -> AccountStoreErro } } -pub(super) fn sqlite_usage_delta(usage: AccountUsageDelta) -> UsageDelta { +pub(super) fn pg_usage_delta(usage: AccountUsageDelta) -> UsageDelta { let request_count = u64_to_i64_saturating(usage.requests); let input_tokens = u64_to_i64_saturating(usage.input_tokens); let output_tokens = u64_to_i64_saturating(usage.output_tokens); @@ -230,15 +213,16 @@ pub(super) fn quota_plan_type(quota_json: &str) -> Option { } pub(super) fn quota_snapshot_from_row( - row: &sqlx::sqlite::SqliteRow, -) -> SqliteAccountStoreResult { + row: &sqlx::postgres::PgRow, +) -> PgAccountStoreResult { Ok(AccountQuotaSnapshot { account_id: row.get("id"), email: row.get("email"), - quota_json: row.get("quota_json"), - quota_fetched_at: parse_optional_rfc3339( - row.get::, _>("quota_fetched_at").as_deref(), - )?, + quota_json: row + .get::, _>("quota_json") + .0 + .to_string(), + quota_fetched_at: row.get("quota_fetched_at"), }) } @@ -250,14 +234,14 @@ pub(super) fn optional_update_value(value: &Option>) -> Option<&s value.as_ref().and_then(|value| value.as_deref()) } -pub(super) fn status_from_db(value: &str) -> SqliteAccountStoreResult { +pub(super) fn status_from_db(value: &str) -> PgAccountStoreResult { 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(SqliteAccountStoreError::InvalidStatus(other.to_string())), + other => Err(PgAccountStoreError::InvalidStatus(other.to_string())), } } @@ -272,17 +256,17 @@ pub(super) fn optional_positive_i64_to_u64(value: Option) -> Option { } pub(super) async fn count_account_metadata( - pool: &SqlitePool, + pool: &PgPool, search: Option<&str>, -) -> SqliteAccountStoreResult { - let mut builder = QueryBuilder::::new("select count(*) from accounts"); +) -> PgAccountStoreResult { + let mut builder = QueryBuilder::::new("select count(*) from accounts"); push_account_metadata_search(&mut builder, search); let (total,): (i64,) = builder.build_query_as().fetch_one(pool).await?; Ok(nonnegative_i64_to_u64(total)) } pub(super) fn push_account_metadata_search( - builder: &mut QueryBuilder, + builder: &mut QueryBuilder, search: Option<&str>, ) { let Some(search) = search else { @@ -290,27 +274,27 @@ pub(super) fn push_account_metadata_search( }; let pattern = format!("%{search}%"); - builder.push(" where id like "); + builder.push(" where id ilike "); builder.push_bind(pattern.clone()); - builder.push(" or email like "); + builder.push(" or email ilike "); builder.push_bind(pattern.clone()); - builder.push(" or label like "); + builder.push(" or label ilike "); builder.push_bind(pattern.clone()); - builder.push(" or chatgpt_account_id like "); + builder.push(" or chatgpt_account_id ilike "); builder.push_bind(pattern.clone()); - builder.push(" or chatgpt_user_id like "); + builder.push(" or chatgpt_user_id ilike "); builder.push_bind(pattern); } pub(super) fn to_page( - rows: &[SqliteRow], + rows: &[PgRow], limit: u32, - mapper: impl Fn(&SqliteRow) -> SqliteAccountStoreResult, + mapper: impl Fn(&PgRow) -> PgAccountStoreResult, cursor_fields: (&str, &str), ) -> Page { let has_more = rows.len() > limit as usize; let mut items: Vec = Vec::with_capacity(limit as usize); - let mut last_row: Option<&SqliteRow> = None; + let mut last_row: Option<&PgRow> = None; for (i, row) in rows.iter().enumerate() { if i >= limit as usize { break; @@ -323,7 +307,9 @@ pub(super) fn to_page( let next_cursor = if has_more { last_row.map(|row| { use sqlx::Row; - let ts: String = row.get(cursor_fields.0); + let ts = row + .get::, _>(cursor_fields.0) + .to_rfc3339(); let id: String = row.get(cursor_fields.1); crate::infra::json::encode_cursor(&ts, &id) }) diff --git a/backend/src/accounts/store/write.rs b/backend/src/accounts/store/write.rs new file mode 100644 index 000000000..a31d6da0c --- /dev/null +++ b/backend/src/accounts/store/write.rs @@ -0,0 +1,660 @@ +use super::*; + +impl PgAccountStore { + /// 构造存储。 + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + /// 返回底层连接池。 + pub fn pool(&self) -> &PgPool { + &self.pool + } + + /// 插入新账号。 + pub async fn insert(&self, account: NewAccount) -> PgAccountStoreResult<()> { + let now = Utc::now(); + let refresh_token = account + .refresh_token + .as_ref() + .map(ExposeSecret::expose_secret); + sqlx::query(INSERT_ACCOUNT_SQL) + .bind(&account.id) + .bind(&account.email) + .bind(&account.account_id) + .bind(&account.user_id) + .bind(&account.label) + .bind(&account.plan_type) + .bind(account.access_token.expose_secret()) + .bind(refresh_token) + .bind(account.access_token_expires_at) + .bind(status_to_db(account.status)) + .bind(account.added_at.unwrap_or(now)) + .bind(now) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 读取单个账号。 + pub async fn get(&self, account_id: &str) -> PgAccountStoreResult> { + let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_ID_SQL) + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(stored_account_from_row).transpose() + } + + /// 通过 ChatGPT 身份查找账号。 + pub async fn find_by_chatgpt_identity( + &self, + chatgpt_account_id: &str, + chatgpt_user_id: Option<&str>, + ) -> PgAccountStoreResult> { + let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL) + .bind(chatgpt_account_id) + .bind(chatgpt_user_id) + .bind(chatgpt_user_id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(stored_account_from_row).transpose() + } + + /// 分页列出所有账号（含 token）。 + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> PgAccountStoreResult> { + let limit = limit.clamp(1, 200); + if let Some(cursor) = cursor { + let (added_at, id) = + decode_cursor(&cursor).ok_or(PgAccountStoreError::InvalidCursor)?; + let added_at = + parse_rfc3339_utc(&added_at).map_err(|_| PgAccountStoreError::InvalidCursor)?; + let rows = sqlx::query(LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL) + .bind(added_at) + .bind(added_at) + .bind(&id) + .bind(i64::from(limit) + 1) + .fetch_all(&self.pool) + .await?; + Ok(to_page( + &rows, + limit, + stored_account_from_row, + ("added_at", "id"), + )) + } else { + let rows = sqlx::query(LIST_STORED_ACCOUNTS_SQL) + .bind(i64::from(limit) + 1) + .fetch_all(&self.pool) + .await?; + Ok(to_page( + &rows, + limit, + stored_account_from_row, + ("added_at", "id"), + )) + } + } + + /// 分页列出账号元数据（不含 token）。 + pub async fn list_metadata( + &self, + cursor: Option, + limit: u32, + ) -> PgAccountStoreResult> { + let limit = limit.clamp(1, 200); + if let Some(cursor) = cursor { + let (added_at, id) = + decode_cursor(&cursor).ok_or(PgAccountStoreError::InvalidCursor)?; + let added_at = + parse_rfc3339_utc(&added_at).map_err(|_| PgAccountStoreError::InvalidCursor)?; + let rows = sqlx::query(LIST_ACCOUNT_METADATA_AFTER_CURSOR_SQL) + .bind(added_at) + .bind(added_at) + .bind(&id) + .bind(i64::from(limit) + 1) + .fetch_all(&self.pool) + .await?; + Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) + } else { + let rows = sqlx::query(LIST_ACCOUNT_METADATA_SQL) + .bind(i64::from(limit) + 1) + .fetch_all(&self.pool) + .await?; + Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) + } + } + + /// 按页码列出账号元数据（不含 token）。 + pub async fn list_metadata_page( + &self, + page: u32, + page_size: u32, + search: Option<&str>, + ) -> PgAccountStoreResult> { + 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 offset = page_offset(page, page_size); + + let mut builder = QueryBuilder::::new(LIST_ACCOUNT_METADATA_SELECT_SQL); + push_account_metadata_search(&mut builder, search); + builder.push(" order by added_at desc, id desc limit "); + builder.push_bind(i64::from(page_size)); + builder.push(" offset "); + builder.push_bind(offset.min(i64::MAX as u64) as i64); + + let rows = builder.build().fetch_all(&self.pool).await?; + let items = rows + .iter() + .map(metadata_from_row) + .collect::>>()?; + + Ok(NumberedPage { + items, + total, + page: page.max(1), + page_size, + }) + } + + /// 读取单个账号元数据（不含 token）。 + pub async fn get_metadata( + &self, + account_id: &str, + ) -> PgAccountStoreResult> { + let row = sqlx::query(SELECT_ACCOUNT_METADATA_BY_ID_SQL) + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + row.as_ref().map(metadata_from_row).transpose() + } + + /// 更新单账号元数据（不含 token）。 + pub async fn update_metadata( + &self, + account_id: &str, + update: AccountMetadataUpdate, + ) -> PgAccountStoreResult { + if !update.any() { + return Ok(false); + } + let now = Utc::now(); + let status = update.status.map(status_to_db); + let result = sqlx::query(UPDATE_ACCOUNT_METADATA_SQL) + .bind(update.email.is_some()) + .bind(optional_update_value(&update.email)) + .bind(update.account_id.is_some()) + .bind(optional_update_value(&update.account_id)) + .bind(update.user_id.is_some()) + .bind(optional_update_value(&update.user_id)) + .bind(update.label.is_some()) + .bind(optional_update_value(&update.label)) + .bind(update.plan_type.is_some()) + .bind(optional_update_value(&update.plan_type)) + .bind(update.status.is_some()) + .bind(status) + .bind(status) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 更新账号 claims（含 refresh token）。 + pub async fn update_from_claims( + &self, + account_id: &str, + update: AccountClaimsUpdate, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let refresh_token = update + .refresh_token + .as_ref() + .map(ExposeSecret::expose_secret); + + let result = if let Some(refresh_token) = refresh_token { + sqlx::query(UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL) + .bind(&update.email) + .bind(&update.account_id) + .bind(&update.user_id) + .bind(&update.plan_type) + .bind(update.access_token.expose_secret()) + .bind(refresh_token) + .bind(update.access_token_expires_at) + .bind(update.next_refresh_at) + .bind(status_to_db(update.status)) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await? + } else { + sqlx::query(UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL) + .bind(&update.email) + .bind(&update.account_id) + .bind(&update.user_id) + .bind(&update.plan_type) + .bind(update.access_token.expose_secret()) + .bind(update.access_token_expires_at) + .bind(update.next_refresh_at) + .bind(status_to_db(update.status)) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await? + }; + Ok(result.rows_affected() > 0) + } + + /// 通过导入数据更新已有账号。 + pub async fn update_from_import( + &self, + update: ImportedAccountUpdate, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let refresh_token = update + .account + .refresh_token + .as_ref() + .map(ExposeSecret::expose_secret); + let quota_json = update + .quota_json + .map(|value| serde_json::from_str::(&value)) + .transpose()? + .map(sqlx::types::Json); + let quota_fetched_at = update.quota_fetched_at; + + let result = if let Some(refresh_token) = refresh_token { + sqlx::query(UPDATE_IMPORTED_ACCOUNT_WITH_REFRESH_SQL) + .bind(&update.account.email) + .bind(&update.account.account_id) + .bind(&update.account.user_id) + .bind(&update.account.label) + .bind(&update.account.plan_type) + .bind(update.account.access_token.expose_secret()) + .bind(refresh_token) + .bind(update.account.access_token_expires_at) + .bind(status_to_db(update.account.status)) + .bind("a_json) + .bind(quota_fetched_at) + .bind(quota_fetched_at) + .bind(update.quota_verify_required) + .bind(now) + .bind(&update.account.id) + .execute(&self.pool) + .await? + } else { + sqlx::query(UPDATE_IMPORTED_ACCOUNT_PRESERVING_REFRESH_SQL) + .bind(&update.account.email) + .bind(&update.account.account_id) + .bind(&update.account.user_id) + .bind(&update.account.label) + .bind(&update.account.plan_type) + .bind(update.account.access_token.expose_secret()) + .bind(update.account.access_token_expires_at) + .bind(status_to_db(update.account.status)) + .bind("a_json) + .bind(quota_fetched_at) + .bind(quota_fetched_at) + .bind(update.quota_verify_required) + .bind(now) + .bind(&update.account.id) + .execute(&self.pool) + .await? + }; + Ok(result.rows_affected() > 0) + } + + /// 设置下一次刷新时间。 + pub async fn set_next_refresh_at( + &self, + account_id: &str, + next_refresh_at: Option>, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let result = sqlx::query(SET_NEXT_REFRESH_AT_SQL) + .bind(next_refresh_at) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 标记账号进入配额冷却期。 + pub async fn mark_quota_limited_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let result = sqlx::query(MARK_QUOTA_LIMITED_UNTIL_SQL) + .bind(cooldown_until) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 同步运行时自然刷新出来的账号状态。 + pub async fn sync_runtime_account_state( + &self, + account: &Account, + sync_usage_window: bool, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let quota_limit_reached = account.quota_limit_reached; + let quota_verify_required = account.quota_verify_required; + let quota_cooldown_until = account.quota_cooldown_until; + let cloudflare_cooldown_until = account.cloudflare_cooldown_until; + let mut tx = self.pool.begin().await?; + let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL) + .bind(quota_limit_reached) + .bind(now) + .bind(quota_limit_reached) + .bind(status_to_db(account.status)) + .bind(quota_limit_reached) + .bind(now) + .bind(quota_limit_reached) + .bind(quota_verify_required) + .bind(now) + .bind(quota_verify_required) + .bind(quota_cooldown_until) + .bind(now) + .bind(quota_cooldown_until) + .bind(cloudflare_cooldown_until) + .bind(now) + .bind(cloudflare_cooldown_until) + .bind(now) + .bind(&account.id) + .execute(&mut *tx) + .await?; + if result.rows_affected() == 0 { + tx.commit().await?; + return Ok(false); + } + + if sync_usage_window { + sqlx::query(SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL) + .bind(&account.id) + .bind(u64_to_i64_saturating(account.window_request_count)) + .bind(u64_to_i64_saturating(account.window_input_tokens)) + .bind(u64_to_i64_saturating(account.window_output_tokens)) + .bind(u64_to_i64_saturating(account.window_cached_tokens)) + .bind(u64_to_i64_saturating(account.window_image_input_tokens)) + .bind(u64_to_i64_saturating(account.window_image_output_tokens)) + .bind(u64_to_i64_saturating(account.window_image_request_count)) + .bind(u64_to_i64_saturating( + account.window_image_request_failed_count, + )) + .bind(account.window_started_at) + .bind(account.window_reset_at) + .bind(account.limit_window_seconds.map(u64_to_i64_saturating)) + .bind(now) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + Ok(true) + } + + /// 标记账号 Cloudflare 冷却期。 + pub async fn set_cloudflare_cooldown_until( + &self, + account_id: &str, + cooldown_until: DateTime, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let result = sqlx::query(SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL) + .bind(cooldown_until) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 更新账号状态。 + pub async fn set_status( + &self, + account_id: &str, + status: AccountStatus, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let status = status_to_db(status); + let result = sqlx::query( + r" +update accounts +set + status = case + when $1 = 'active' and quota_limit_reached then 'quota_exhausted' + else $2 + end, + updated_at = $3 +where id = $4", + ) + .bind(status) + .bind(status) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 记录用量。 + pub async fn record_usage( + &self, + account_id: &str, + delta: UsageDelta, + ) -> PgAccountStoreResult<()> { + let last_used_at = Utc::now(); + sqlx::query(RECORD_USAGE_SQL) + .bind(account_id) + .bind(delta.request_count) + .bind(delta.empty_response_count) + .bind(delta.input_tokens) + .bind(delta.output_tokens) + .bind(delta.cached_tokens) + .bind(delta.reasoning_tokens) + .bind(delta.total_tokens) + .bind(delta.image_input_tokens) + .bind(delta.image_output_tokens) + .bind(delta.image_request_count) + .bind(delta.image_request_failed_count) + .bind(delta.window_request_count) + .bind(delta.window_input_tokens) + .bind(delta.window_output_tokens) + .bind(delta.window_cached_tokens) + .bind(delta.window_image_input_tokens) + .bind(delta.window_image_output_tokens) + .bind(delta.window_image_request_count) + .bind(delta.window_image_request_failed_count) + .bind(last_used_at) + .bind(last_used_at) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 记录模型维度用量。 + pub async fn record_model_usage( + &self, + account_id: &str, + model: &str, + delta: AccountModelUsageDelta, + ) -> PgAccountStoreResult<()> { + let model = model.trim(); + if model.is_empty() { + return Ok(()); + } + let last_used_at = Utc::now(); + sqlx::query(RECORD_MODEL_USAGE_SQL) + .bind(account_id) + .bind(model) + .bind(u64_to_i64_saturating(delta.requests)) + .bind(u64_to_i64_saturating(delta.errors)) + .bind(u64_to_i64_saturating(delta.input_tokens)) + .bind(u64_to_i64_saturating(delta.output_tokens)) + .bind(u64_to_i64_saturating(delta.cached_tokens)) + .bind(last_used_at) + .execute(&self.pool) + .await?; + Ok(()) + } + + /// 读取配额快照列表。 + pub async fn list_quota_snapshots(&self) -> PgAccountStoreResult> { + let rows = sqlx::query(LIST_QUOTA_SNAPSHOTS_SQL) + .fetch_all(&self.pool) + .await?; + rows.iter().map(quota_snapshot_from_row).collect() + } + + /// 读取单账号配额 JSON。 + pub async fn get_quota_json(&self, account_id: &str) -> PgAccountStoreResult> { + let row = sqlx::query("select quota_json from accounts where id = $1") + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + Ok(row + .and_then(|row| row.get::>, _>("quota_json")) + .map(|value| value.0.to_string())) + } + + /// 更新配额 JSON。 + pub async fn update_quota_json( + &self, + account_id: &str, + quota_json: &str, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let plan_type = quota_plan_type(quota_json); + let quota_json = sqlx::types::Json(serde_json::from_str::(quota_json)?); + let result = sqlx::query(UPDATE_QUOTA_JSON_SQL) + .bind(quota_json) + .bind(now) + .bind(&plan_type) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 应用配额快照。 + pub async fn apply_quota_snapshot( + &self, + account_id: &str, + quota_json: &str, + limit_reached: bool, + cooldown_until: Option>, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let cooldown = cooldown_until; + let plan_type = quota_plan_type(quota_json); + let quota_json = sqlx::types::Json(serde_json::from_str::(quota_json)?); + let result = sqlx::query(APPLY_QUOTA_SNAPSHOT_SQL) + .bind(quota_json) + .bind(now) + .bind(&plan_type) + .bind(limit_reached) + .bind(limit_reached) + .bind(cooldown) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 应用导入的配额状态。 + pub async fn apply_imported_quota_state( + &self, + account_id: &str, + quota_json: Option<&str>, + quota_fetched_at: Option>, + quota_verify_required: bool, + ) -> PgAccountStoreResult { + let now = Utc::now(); + let fetched = quota_fetched_at; + let plan_type = quota_json.and_then(quota_plan_type); + let quota_json = quota_json + .map(serde_json::from_str::) + .transpose()? + .map(sqlx::types::Json); + let result = sqlx::query(APPLY_IMPORTED_QUOTA_STATE_SQL) + .bind(quota_json) + .bind(fetched) + .bind(fetched) + .bind(&plan_type) + .bind(quota_verify_required) + .bind(now) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + /// 同步 rate-limit 窗口（含重置）。 + pub async fn sync_rate_limit_window( + &self, + account_id: &str, + reset_at: DateTime, + limit_window_seconds: Option, + ) -> PgAccountStoreResult<()> { + let existing = sqlx::query(SELECT_RATE_LIMIT_WINDOW_SQL) + .bind(account_id) + .fetch_optional(&self.pool) + .await?; + + let should_reset = existing + .as_ref() + .map(|row| { + let existing_reset_at = row.get::>, _>("window_reset_at"); + let existing_limit_window_seconds = + optional_positive_i64_to_u64(row.get::, _>("limit_window_seconds")); + should_reset_usage_window( + existing_reset_at, + existing_limit_window_seconds, + reset_at, + limit_window_seconds, + ) + }) + .unwrap_or_default(); + + if should_reset { + let window_started_at = Utc::now(); + sqlx::query(SYNC_RATE_LIMIT_WINDOW_RESET_SQL) + .bind(account_id) + .bind(window_started_at) + .bind(reset_at) + .bind(limit_window_seconds.map(u64_to_i64_saturating)) + .execute(&self.pool) + .await?; + } else { + sqlx::query(SYNC_RATE_LIMIT_WINDOW_SQL) + .bind(account_id) + .bind(reset_at) + .bind(limit_window_seconds.map(u64_to_i64_saturating)) + .execute(&self.pool) + .await?; + } + Ok(()) + } + + /// 删除单个账号。 + pub async fn delete(&self, account_id: &str) -> PgAccountStoreResult { + let result = sqlx::query(DELETE_ACCOUNT_SQL) + .bind(account_id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } +} diff --git a/backend/src/upstream/accounts/window.rs b/backend/src/accounts/window.rs similarity index 100% rename from backend/src/upstream/accounts/window.rs rename to backend/src/accounts/window.rs diff --git a/backend/src/admin/accounts/mod.rs b/backend/src/admin/accounts/mod.rs deleted file mode 100644 index 34caaa04d..000000000 --- a/backend/src/admin/accounts/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -//! 管理端账号 HTTP 模块。 - -pub(crate) mod quota_view; -pub mod routes; -pub mod service; diff --git a/backend/src/admin/accounts/routes.rs b/backend/src/admin/accounts/routes.rs deleted file mode 100644 index af056b9ab..000000000 --- a/backend/src/admin/accounts/routes.rs +++ /dev/null @@ -1,1481 +0,0 @@ -//! 账号管理 HTTP 处理器。 - -use axum::{ - body::Body, - extract::{Query, State}, - http::{ - header::{CACHE_CONTROL, CONNECTION, CONTENT_TYPE}, - StatusCode, - }, - response::{IntoResponse, Response}, - Json, -}; -use std::collections::{HashMap, HashSet}; - -use chrono::{DateTime, Utc}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use sqlx::{QueryBuilder, Row, Sqlite}; - -use crate::{ - admin::auth::session::AdminAuth, - admin::response::{ - AdminEnvelope, AdminError, AdminResponse, BatchDeleteData, CursorPageMeta, - NumberedPageMeta, PageMeta, ADMIN_OK_CODE, ADMIN_OK_MESSAGE, - }, - admin::update_payload::{parse_editable_update, EditableUpdateMessages}, - admin::{ - accounts::quota_view::{ - quota_data, AdminAccountQuotaData, AdminAccountQuotaUsageWindow, - AdminAccountQuotaWindowLocalUsage, - }, - accounts::service::{ - AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, - AdminAccountRefreshOutcome, AdminAccountRefreshResult, AdminAccountUpdate, - OAuthExchangeInput, - }, - monitoring::{account_usage_service::AdminUsageRecord, billing}, - }, - infra::{ - format::{ - format_cost, format_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64, - }, - json::{clamp_limit, clamp_page, total_pages, Page}, - time::{ - china_datetime, china_quarter_hour_start, china_relative_time, china_rfc3339, - parse_rfc3339_utc, - }, - }, - runtime::state::AppState, - upstream::accounts::{model::AccountStatus, token_refresh::token_refresh_status_eligible}, -}; - -const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200; -const ACCOUNT_EXPORT_CONFIRMATION: &str = "export_sensitive_accounts"; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AccountsQuery { - cursor: Option, - limit: Option, - page: Option, - page_size: Option, - search: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct CreateAccountRequest { - token: Option, - refresh_token: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct BatchDeleteAccountsRequest { - ids: Vec, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AccountActionRequest { - id: String, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AccountIdQuery { - id: String, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AccountExportQuery { - ids: Option, - confirm: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AccountTestQuery { - id: String, - model_id: Option, -} - -#[derive(Debug, Default, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct AccountOAuthExchangeRequest { - session_id: String, - callback_url: Option, - code: Option, - state: Option, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountModelsData { - models: Vec, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountModelData { - id: String, - label: String, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminAccountData { - id: String, - email: Option, - account_id: Option, - user_id: Option, - label: Option, - plan_type: Option, - has_refresh_token: bool, - status: String, - display_status: String, - token_refreshing: bool, - access_token_expires_at: Option, - access_token_expires_at_display: Option, - added_at: String, - added_at_display: String, - updated_at: String, - updated_at_display: String, - quota: AdminAccountQuotaData, - usage: AdminAccountUsageData, -} - -impl From for AdminAccountData { - fn from(a: AdminAccountMetadata) -> Self { - Self::from_parts(a, None, None, Vec::new(), false) - } -} - -impl AdminAccountData { - fn from_parts( - a: AdminAccountMetadata, - usage: Option<&AdminUsageRecord>, - quota: Option, - models: Vec, - token_refreshing: bool, - ) -> Self { - let access_token_expires_at = a.access_token_expires_at.as_ref().map(china_rfc3339); - let access_token_expires_at_display = - a.access_token_expires_at.as_ref().map(china_datetime); - let display_status = account_display_status(a.status, token_refreshing).to_string(); - Self { - id: a.id, - email: a.email, - account_id: a.account_id, - user_id: a.user_id, - label: a.label, - plan_type: a.plan_type, - has_refresh_token: a.has_refresh_token, - status: a.status.as_str().to_string(), - display_status, - token_refreshing, - access_token_expires_at, - access_token_expires_at_display, - added_at: china_rfc3339(&a.added_at), - added_at_display: china_datetime(&a.added_at), - updated_at: china_rfc3339(&a.updated_at), - updated_at_display: china_datetime(&a.updated_at), - quota: quota.unwrap_or_default(), - usage: AdminAccountUsageData::from_usage(usage, models), - } - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminAccountUsageData { - request_count: i64, - request_count_display: String, - empty_response_count: i64, - input_tokens: i64, - input_tokens_display: String, - output_tokens: i64, - output_tokens_display: String, - cached_tokens: i64, - cached_tokens_display: String, - reasoning_tokens: i64, - total_tokens: i64, - total_tokens_display: String, - image_input_tokens: i64, - image_output_tokens: i64, - image_tokens_display: String, - image_request_count: i64, - image_request_failed_count: i64, - created_tokens: i64, - created_tokens_display: String, - read_tokens: i64, - read_tokens_display: String, - last_used_at: Option, - last_used_at_display: String, - models: Vec, -} - -impl AdminAccountUsageData { - fn from_usage( - usage: Option<&AdminUsageRecord>, - models: Vec, - ) -> Self { - let request_count = usage.map_or(0, |usage| usage.window_request_count); - let empty_response_count = 0; - let input_tokens = usage.map_or(0, |usage| usage.window_input_tokens); - let output_tokens = usage.map_or(0, |usage| usage.window_output_tokens); - let cached_tokens = usage.map_or(0, |usage| usage.window_cached_tokens); - let reasoning_tokens = 0; - let total_tokens = input_tokens.saturating_add(output_tokens); - let image_input_tokens = 0; - let image_output_tokens = 0; - let image_request_count = 0; - let image_request_failed_count = 0; - let created_tokens = input_tokens.saturating_sub(cached_tokens); - let read_tokens = cached_tokens; - let last_used_at = usage.and_then(|usage| usage.last_used_at); - - Self { - request_count, - request_count_display: format_plain_number(nonnegative_i64_to_u64(request_count)), - empty_response_count, - input_tokens, - input_tokens_display: format_tokens(nonnegative_i64_to_u64(input_tokens)), - output_tokens, - output_tokens_display: format_tokens(nonnegative_i64_to_u64(output_tokens)), - cached_tokens, - cached_tokens_display: format_tokens(nonnegative_i64_to_u64(cached_tokens)), - reasoning_tokens, - total_tokens, - total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)), - image_input_tokens, - image_output_tokens, - image_tokens_display: format_tokens(nonnegative_i64_to_u64( - image_input_tokens + image_output_tokens, - )), - image_request_count, - image_request_failed_count, - created_tokens, - created_tokens_display: format_tokens(nonnegative_i64_to_u64(created_tokens)), - read_tokens, - read_tokens_display: format_tokens(nonnegative_i64_to_u64(read_tokens)), - last_used_at: last_used_at.map(|value| china_rfc3339(&value)), - last_used_at_display: china_relative_time(last_used_at, Utc::now()), - models, - } - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminAccountModelUsageData { - model: String, - request_count: u64, - request_count_display: String, - success_rate: f64, - success_rate_display: String, - input_tokens: u64, - input_tokens_display: String, - output_tokens: u64, - output_tokens_display: String, - cached_tokens: u64, - cached_tokens_display: String, - total_tokens: u64, - total_tokens_display: String, - total_cost_usd: f64, - total_cost_usd_display: String, - last_used_at: Option, - last_used_at_display: String, -} - -#[derive(Debug, Clone, Default)] -struct AccountListStats { - usage_by_account: HashMap, - quota_by_account: HashMap, - models_by_account: HashMap>, - refreshing_account_ids: HashSet, -} - -struct AccountModelUsageRecord { - account_id: String, - model: String, - request_count: i64, - error_count: i64, - input_tokens: i64, - output_tokens: i64, - cached_tokens: i64, - total_cost_usd: f64, - last_used_at: Option>, -} - -#[derive(Debug, Clone)] -struct AccountQuotaUsageWindow { - account_id: String, - window: AdminAccountQuotaUsageWindow, -} - -impl AccountListStats { - fn data_for(&self, account: AdminAccountMetadata) -> AdminAccountData { - let account_id = account.id.clone(); - let token_refreshing = token_refresh_status_eligible(account.status) - && self.refreshing_account_ids.contains(&account_id); - AdminAccountData::from_parts( - account, - self.usage_by_account.get(&account_id), - self.quota_by_account.get(&account_id).cloned(), - self.models_by_account - .get(&account_id) - .cloned() - .unwrap_or_default(), - token_refreshing, - ) - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminAccountSummaryData { - total: u64, - active: u64, - high_usage: u64, - attention: u64, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminAccountPageData { - items: Vec, - page: PageMeta, - summary: AdminAccountSummaryData, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AdminAccountPageEnvelope { - code: u32, - message: String, - data: AdminAccountPageData, -} - -impl AdminAccountPageEnvelope { - fn cursor(page: Page, limit: u32, summary: AdminAccountSummaryData) -> Self { - Self { - code: ADMIN_OK_CODE, - message: ADMIN_OK_MESSAGE.into(), - data: AdminAccountPageData { - items: page.items, - page: PageMeta::Cursor(CursorPageMeta { - limit, - next_cursor: page.next_cursor, - }), - summary, - }, - } - } - - fn numbered( - page: crate::infra::json::NumberedPage, - summary: AdminAccountSummaryData, - ) -> Self { - Self { - code: ADMIN_OK_CODE, - message: ADMIN_OK_MESSAGE.into(), - data: AdminAccountPageData { - page: PageMeta::Numbered(NumberedPageMeta { - page: page.page, - page_size: page.page_size, - total: page.total, - total_pages: total_pages(page.total, page.page_size), - }), - items: page.items, - summary, - }, - } - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(untagged)] -enum AccountUpdateData { - Account(Box), -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountRefreshData { - id: String, - email: Option, - previous_status: String, - result: &'static str, - error: Option, - duration_ms: i64, -} - -impl From for AccountRefreshData { - fn from(result: AdminAccountRefreshResult) -> Self { - Self { - id: result.id, - email: result.email, - previous_status: result.previous_status.as_str().to_string(), - result: account_refresh_outcome_str(result.outcome), - error: result.error, - duration_ms: result.duration_ms, - } - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountQuotaData { - quota: Value, - raw: Value, - quota_data: AdminAccountQuotaData, - plan_type: Option, - account: AdminAccountData, -} - -impl AccountQuotaData { - fn from_account( - quota: Value, - raw: Value, - quota_data: AdminAccountQuotaData, - account: AdminAccountData, - ) -> Self { - let plan_type = account.plan_type.clone(); - Self { - quota, - raw, - quota_data, - plan_type, - account, - } - } -} - -#[derive(Debug, Clone, Deserialize, Default)] -#[serde(default)] -pub(crate) struct AccountHealthCheckRequest { - ids: Option>, - stagger_ms: Option, - concurrency: Option, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountHealthCheckSummaryData { - total: usize, - alive: usize, - dead: usize, - skipped: usize, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountHealthCheckData { - summary: AccountHealthCheckSummaryData, - results: Vec, -} - -impl From for AccountHealthCheckData { - fn from(result: AdminAccountHealthCheck) -> Self { - let summary = AccountHealthCheckSummaryData { - total: result.results.len(), - alive: result.alive(), - dead: result.dead(), - skipped: result.skipped(), - }; - Self { - summary, - results: result - .results - .into_iter() - .map(AccountRefreshData::from) - .collect(), - } - } -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountImportData { - imported: u32, - skipped: u32, - source_format: String, -} - -#[derive(Debug, Clone, Serialize)] -#[serde(rename_all = "camelCase")] -struct AccountOAuthAuthorizeData { - session_id: String, - auth_url: String, - expires_at: String, - expires_at_display: String, -} - -/// `GET /api/admin/accounts` -pub(crate) async fn accounts( - State(state): State, - _auth: AdminAuth, - Query(params): Query, -) -> Result { - let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50)); - let use_numbered_page = params.page.is_some() || params.page_size.is_some(); - let quota_by_account = quota_snapshots_by_account(&state).await; - let summary = account_summary_data(&state, "a_by_account).await; - - if use_numbered_page { - return match state - .services - .admin_accounts - .list_page(clamp_page(params.page.unwrap_or(1)), limit, params.search) - .await - { - Ok(page) => { - let stats = account_list_stats(&state, &page.items, "a_by_account).await?; - let page = crate::infra::json::NumberedPage { - items: page - .items - .into_iter() - .map(|item| stats.data_for(item)) - .collect(), - total: page.total, - page: page.page, - page_size: page.page_size, - }; - Ok(AdminResponse::new( - StatusCode::OK, - AdminAccountPageEnvelope::numbered(page, summary), - )) - } - Err(error) => Err(account_error(&error)), - }; - } - - match state - .services - .admin_accounts - .list(params.cursor, limit) - .await - { - Ok(page) => { - let stats = account_list_stats(&state, &page.items, "a_by_account).await?; - let page = Page { - items: page - .items - .into_iter() - .map(|item| stats.data_for(item)) - .collect(), - next_cursor: page.next_cursor, - }; - Ok(AdminResponse::new( - StatusCode::OK, - AdminAccountPageEnvelope::cursor(page, limit, summary), - )) - } - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts` -pub(crate) async fn create_account( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - match state - .services - .admin_accounts - .create(payload.token, payload.refresh_token) - .await - { - Ok(account) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AdminAccountData::from(account)), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `GET /api/admin/accounts/export` -pub(crate) async fn export_accounts( - State(state): State, - _auth: AdminAuth, - Query(query): Query, -) -> Result { - let ids = account_export_ids(query.ids.as_deref()); - - if query.confirm.as_deref() != Some(ACCOUNT_EXPORT_CONFIRMATION) { - return Err(AdminError::bad_request( - "account export requires confirm=export_sensitive_accounts", - )); - } - - match state.services.admin_accounts.export(ids).await { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(result), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/refresh` -pub(crate) async fn refresh_account( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - let account_id = payload.id; - match state - .services - .admin_accounts - .refresh_account(&account_id) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountRefreshData::from(result)), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/health-check` -pub(crate) async fn health_check_accounts( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - let stagger_ms = payload.stagger_ms.unwrap_or(3_000); - if !(500..=30_000).contains(&stagger_ms) && stagger_ms != 0 { - return Err(AdminError::bad_request( - "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", - )); - } - } - match state - .services - .admin_accounts - .health_check_accounts(payload.ids, stagger_ms, payload.concurrency.unwrap_or(2)) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountHealthCheckData::from(result)), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `GET /api/admin/accounts/quota` -pub(crate) async fn account_quota( - State(state): State, - _auth: AdminAuth, - Query(query): Query, -) -> Result { - let account_id = query.id; - match state - .services - .admin_accounts - .account_quota(&account_id) - .await - { - Ok(data) => { - let quota = data.get("quota").cloned().unwrap_or(Value::Null); - let raw = data.get("raw").cloned().unwrap_or(Value::Null); - let quota_json = quota.to_string(); - let mut quota_data = quota_data("a_json, Some(Utc::now())); - apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await; - let account = - account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountQuotaData::from_account( - quota, raw, quota_data, account, - )), - )) - } - Err(AdminAccountError::NotFound) => Err(account_not_found()), - Err(AdminAccountError::Inactive(status)) => Err(AdminError::conflict(format!( - "Account is {}, cannot query quota", - status.as_str() - ))), - Err(AdminAccountError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!( - "Failed to fetch quota from Codex API: {msg}" - ))), - Err(e) => Err(account_error(&e)), - } -} - -async fn account_data_for_quota_refresh( - state: &AppState, - account_id: &str, - quota_data: AdminAccountQuotaData, -) -> Result { - let Some(account) = state - .services - .admin_accounts - .get(account_id) - .await - .map_err(|error| account_error(&error))? - else { - return Err(account_not_found()); - }; - let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]); - let stats = - account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await?; - Ok(stats.data_for(account)) -} - -/// `POST /api/admin/accounts/import` -pub(crate) async fn import_accounts( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - match state.services.admin_accounts.import(payload).await { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountImportData { - imported: result.imported, - skipped: result.skipped, - source_format: result.source_format.to_string(), - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/oauth/authorize` -pub(crate) async fn oauth_authorize_account( - State(state): State, - _auth: AdminAuth, -) -> Result { - match state.services.admin_accounts.oauth_authorize().await { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountOAuthAuthorizeData { - session_id: result.session_id, - auth_url: result.auth_url, - expires_at: china_rfc3339(&result.expires_at), - expires_at_display: china_datetime(&result.expires_at), - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/oauth/exchange` -pub(crate) async fn oauth_exchange_account( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - match state - .services - .admin_accounts - .oauth_exchange(OAuthExchangeInput { - session_id: payload.session_id, - callback_url: payload.callback_url, - code: payload.code, - state: payload.state, - }) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountImportData { - imported: result.imported, - skipped: result.skipped, - source_format: result.source_format.to_string(), - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `GET /api/admin/accounts/test?id=...&modelId=...` -pub(crate) async fn test_account_connection( - State(state): State, - Query(query): Query, - _auth: AdminAuth, -) -> Result { - let account_id = query.id; - - let model = query - .model_id - .as_deref() - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToString::to_string) - .ok_or_else(|| AdminError::bad_request("Model is required"))?; - let stream = state - .services - .admin_accounts - .test_connection_stream(&account_id, model) - .await - .map_err(|error| account_error(&error))?; - - Response::builder() - .status(StatusCode::OK) - .header(CONTENT_TYPE, "text/event-stream") - .header(CACHE_CONTROL, "no-cache") - .header(CONNECTION, "keep-alive") - .header("x-accel-buffering", "no") - .body(Body::from_stream(stream)) - .map_err(|_| AdminError::internal("Failed to build account test stream")) -} - -/// `GET /api/admin/accounts/models?id=...` -pub(crate) async fn account_models( - State(state): State, - Query(query): Query, - _auth: AdminAuth, -) -> Result { - let account_id = query.id; - match state - .services - .admin_accounts - .account_models(&account_id) - .await - { - Ok(models) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountModelsData { - models: models - .into_iter() - .map(|model| AccountModelData { - id: model.id, - label: model.label, - }) - .collect(), - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/delete` -pub(crate) async fn batch_delete_accounts( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - match state - .services - .admin_accounts - .batch_delete(payload.ids) - .await - { - Ok(result) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(BatchDeleteData { - deleted: result.deleted, - not_found: result.not_found, - }), - )), - Err(error) => Err(account_error(&error)), - } -} - -/// `POST /api/admin/accounts/update` -pub(crate) async fn update_account( - State(state): State, - _auth: AdminAuth, - Json(payload): Json, -) -> Result { - let ParsedAccountUpdate { id, update } = parse_account_update(&payload)?; - - match state - .services - .admin_accounts - .update_account(&id, update) - .await - { - Ok(Some(account)) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(AccountUpdateData::Account(Box::new(account.into()))), - )), - Ok(None) => Err(account_not_found()), - Err(error) => Err(account_error(&error)), - } -} - -async fn account_list_stats( - state: &AppState, - accounts: &[AdminAccountMetadata], - quota_by_account: &HashMap, -) -> Result { - let account_ids = accounts - .iter() - .map(|account| account.id.clone()) - .collect::>(); - let usage_records = state - .services - .usage - .list_by_account_ids(&account_ids) - .await - .unwrap_or_default(); - let mut quota_by_account = quota_by_account.clone(); - let all_quota_windows = quota_usage_windows_by_account(accounts, "a_by_account); - let quota_window_usage_stats = - quota_window_local_usage_by_account(state, &all_quota_windows).await; - apply_quota_window_local_usage(accounts, &mut quota_by_account, "a_window_usage_stats); - let selected_quota_windows = selected_quota_windows_by_account(accounts, "a_by_account); - let usage_records = usage_records - .into_iter() - .map(|usage| { - apply_selected_quota_window_usage( - usage, - &selected_quota_windows, - "a_window_usage_stats, - ) - }) - .collect::>(); - let models_by_account = list_current_window_model_usage(state, &usage_records).await; - let refreshing_account_ids = state - .services - .token_refresh - .refreshing_account_ids(&account_ids, Utc::now()) - .await - .map_err(|error| AdminError::internal(error.to_string()))?; - - Ok(AccountListStats { - usage_by_account: usage_records - .into_iter() - .map(|usage| (usage.account_id.clone(), usage)) - .collect(), - quota_by_account, - models_by_account, - refreshing_account_ids, - }) -} - -fn account_display_status(status: AccountStatus, token_refreshing: bool) -> &'static str { - if token_refreshing { - "refreshing" - } else { - status.as_str() - } -} - -fn selected_quota_windows_by_account( - accounts: &[AdminAccountMetadata], - quota_by_account: &HashMap, -) -> HashMap { - accounts - .iter() - .filter_map(|account| { - quota_by_account - .get(&account.id) - .and_then(selected_quota_window) - .map(|window| (account.id.clone(), window)) - }) - .collect() -} - -fn selected_quota_window(quota: &AdminAccountQuotaData) -> Option { - quota.usage_windows().into_iter().max_by(|left, right| { - left.duration_seconds() - .cmp(&right.duration_seconds()) - .then_with(|| left.end.cmp(&right.end)) - .then_with(|| left.key.cmp(&right.key)) - }) -} - -fn quota_usage_windows_by_account( - accounts: &[AdminAccountMetadata], - quota_by_account: &HashMap, -) -> Vec { - accounts - .iter() - .filter_map(|account| { - quota_by_account - .get(&account.id) - .map(|quota| (account.id.as_str(), quota)) - }) - .flat_map(|(account_id, quota)| { - quota - .usage_windows() - .into_iter() - .map(move |window| AccountQuotaUsageWindow { - account_id: account_id.to_string(), - window, - }) - }) - .collect() -} - -async fn apply_account_quota_window_local_usage( - state: &AppState, - account_id: &str, - quota: &mut AdminAccountQuotaData, -) { - let windows = quota - .usage_windows() - .into_iter() - .map(|window| AccountQuotaUsageWindow { - account_id: account_id.to_string(), - window, - }) - .collect::>(); - let usage_by_account = quota_window_local_usage_by_account(state, &windows).await; - let empty_usage = HashMap::new(); - quota.apply_local_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage)); -} - -fn apply_quota_window_local_usage( - accounts: &[AdminAccountMetadata], - quota_by_account: &mut HashMap, - usage_by_account: &HashMap>, -) { - for account in accounts { - if let Some(quota) = quota_by_account.get_mut(&account.id) { - let empty_usage = HashMap::new(); - quota.apply_local_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage)); - } - } -} - -async fn quota_window_local_usage_by_account( - state: &AppState, - windows: &[AccountQuotaUsageWindow], -) -> HashMap> { - let Some(min_start) = windows.iter().map(|window| window.window.start).min() else { - return HashMap::new(); - }; - let Some(max_end) = windows.iter().map(|window| window.window.end).max() else { - return HashMap::new(); - }; - - let account_ids = windows - .iter() - .map(|window| window.account_id.as_str()) - .collect::>(); - let mut windows_by_account = HashMap::<&str, Vec<&AdminAccountQuotaUsageWindow>>::new(); - for window in windows { - windows_by_account - .entry(window.account_id.as_str()) - .or_default() - .push(&window.window); - } - let mut builder = QueryBuilder::::new( - "select - account_id, - bucket_start, - coalesce(sum(request_count), 0) as request_count, - coalesce(sum(input_tokens), 0) as input_tokens, - coalesce(sum(output_tokens), 0) as output_tokens, - coalesce(sum(cached_tokens), 0) as cached_tokens - from usage_time_buckets - where account_id in (", - ); - let mut separated = builder.separated(", "); - for account_id in &account_ids { - separated.push_bind(*account_id); - } - separated.push_unseparated(")"); - builder.push(" and bucket_start >= "); - builder.push_bind(china_quarter_hour_start(min_start).to_rfc3339()); - builder.push(" and bucket_start <= "); - builder.push_bind(max_end.to_rfc3339()); - builder.push(" group by account_id, bucket_start"); - - let rows = builder - .build() - .fetch_all(state.services.background_tasks.accounts.pool()) - .await - .unwrap_or_default(); - let mut usage_by_account = - HashMap::>::new(); - - for row in rows { - let account_id: String = row.get("account_id"); - let Ok(bucket_start) = parse_rfc3339_utc(row.get::<&str, _>("bucket_start")) else { - continue; - }; - let bucket_usage = AdminAccountQuotaWindowLocalUsage { - request_count: row.get("request_count"), - input_tokens: row.get("input_tokens"), - output_tokens: row.get("output_tokens"), - cached_tokens: row.get("cached_tokens"), - }; - let Some(account_windows) = windows_by_account.get(account_id.as_str()) else { - continue; - }; - - for window in account_windows.iter().filter(|window| { - bucket_start >= china_quarter_hour_start(window.start) && bucket_start <= window.end - }) { - usage_by_account - .entry(account_id.clone()) - .or_default() - .entry(window.key.clone()) - .or_default() - .add(bucket_usage); - } - } - - usage_by_account -} - -fn apply_selected_quota_window_usage( - mut usage: AdminUsageRecord, - selected_quota_windows: &HashMap, - quota_window_usage_stats: &HashMap>, -) -> AdminUsageRecord { - let Some(window) = selected_quota_windows.get(&usage.account_id) else { - return usage; - }; - let stats = quota_window_usage_stats - .get(&usage.account_id) - .and_then(|usage_by_window| usage_by_window.get(&window.key)) - .copied() - .unwrap_or_default(); - usage.window_request_count = stats.request_count; - usage.window_input_tokens = stats.input_tokens; - usage.window_output_tokens = stats.output_tokens; - usage.window_cached_tokens = stats.cached_tokens; - usage.window_started_at = Some(window.start); - usage.window_reset_at = Some(window.end); - usage.limit_window_seconds = Some(window.window_seconds); - usage -} - -async fn account_summary_data( - state: &AppState, - quota_by_account: &HashMap, -) -> 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 - .iter() - .filter(|account| { - quota_by_account - .get(&account.id) - .is_some_and(account_quota_has_high_usage) - }) - .count() as u64; - let attention = accounts - .iter() - .filter(|account| account_summary_needs_attention(account.status)) - .count() as u64; - - AdminAccountSummaryData { - total, - active, - high_usage, - attention, - } -} - -async fn quota_snapshots_by_account(state: &AppState) -> HashMap { - state - .services - .background_tasks - .accounts - .list_quota_snapshots() - .await - .unwrap_or_default() - .into_iter() - .map(|snapshot| { - ( - snapshot.account_id, - quota_data(&snapshot.quota_json, snapshot.quota_fetched_at), - ) - }) - .collect() -} - -async fn list_all_account_metadata(state: &AppState) -> Vec { - let mut page = 1; - let mut accounts = Vec::new(); - loop { - let Ok(result) = state - .services - .admin_accounts - .list_page(page, ACCOUNT_STATS_PAGE_LIMIT, None) - .await - else { - return Vec::new(); - }; - let total = result.total; - accounts.extend(result.items); - if accounts.len() as u64 >= total || total == 0 { - return accounts; - } - page = page.saturating_add(1); - } -} - -fn account_quota_has_high_usage(quota: &AdminAccountQuotaData) -> bool { - quota.has_high_usage() -} - -fn account_summary_needs_attention(status: AccountStatus) -> bool { - matches!( - status, - AccountStatus::Expired | AccountStatus::Disabled | AccountStatus::Banned - ) -} - -async fn list_current_window_model_usage( - state: &AppState, - usage_records: &[AdminUsageRecord], -) -> HashMap> { - let now = Utc::now(); - let windows = usage_records - .iter() - .filter_map(|usage| { - current_usage_window(usage, now) - .map(|(start, end)| (usage.account_id.clone(), start, end)) - }) - .collect::>(); - if windows.is_empty() { - return HashMap::new(); - } - - let mut builder = QueryBuilder::::new( - "select - account_id, - model, - service_tier, - coalesce(sum(request_count), 0) as request_count, - coalesce(sum(error_count), 0) as error_count, - coalesce(sum(input_tokens), 0) as input_tokens, - coalesce(sum(output_tokens), 0) as output_tokens, - coalesce(sum(cached_tokens), 0) as cached_tokens, - max(bucket_start) as last_used_at - from usage_time_buckets - where model != '' and (", - ); - for (index, (account_id, start, end)) in windows.iter().enumerate() { - if index > 0 { - builder.push(" or "); - } - let bucket_start = china_quarter_hour_start(*start); - builder.push("(account_id = "); - builder.push_bind(account_id); - builder.push(" and bucket_start >= "); - builder.push_bind(bucket_start.to_rfc3339()); - builder.push(" and bucket_start <= "); - builder.push_bind(end.to_rfc3339()); - builder.push(")"); - } - builder.push(") group by account_id, model, service_tier"); - - let rows = builder - .build() - .fetch_all(state.services.background_tasks.accounts.pool()) - .await - .unwrap_or_default(); - let mut records_by_model = HashMap::<(String, String), AccountModelUsageRecord>::new(); - for row in rows { - let account_id: String = row.get("account_id"); - let model: String = row.get("model"); - let request_count: i64 = row.get("request_count"); - let error_count: i64 = row.get("error_count"); - let input_tokens: i64 = row.get("input_tokens"); - let output_tokens: i64 = row.get("output_tokens"); - let cached_tokens: i64 = row.get("cached_tokens"); - let total_cost_usd = billing::calculate_cost( - nonnegative_i64_to_u64(input_tokens), - nonnegative_i64_to_u64(output_tokens), - nonnegative_i64_to_u64(cached_tokens), - &model, - row.get::, _>("service_tier").as_deref(), - ); - let last_used_at = row - .get::, _>("last_used_at") - .and_then(|value| DateTime::parse_from_rfc3339(&value).ok()) - .map(|value| value.with_timezone(&Utc)); - - let record = records_by_model - .entry((account_id.clone(), model.clone())) - .or_insert_with(|| AccountModelUsageRecord { - account_id, - model, - request_count: 0, - error_count: 0, - input_tokens: 0, - output_tokens: 0, - cached_tokens: 0, - total_cost_usd: 0.0, - last_used_at: None, - }); - record.request_count += request_count; - record.error_count += error_count; - record.input_tokens += input_tokens; - record.output_tokens += output_tokens; - record.cached_tokens += cached_tokens; - record.total_cost_usd += total_cost_usd; - record.last_used_at = record.last_used_at.max(last_used_at); - } - let records = records_by_model.into_values().collect::>(); - - models_by_account(records) -} - -fn current_usage_window( - usage: &AdminUsageRecord, - now: DateTime, -) -> Option<(DateTime, DateTime)> { - let start = usage.window_started_at?; - let end = usage.window_reset_at.unwrap_or(now); - (start <= end).then_some((start, end)) -} - -fn models_by_account( - records: Vec, -) -> HashMap> { - let mut by_account = HashMap::>::new(); - for record in records { - by_account - .entry(record.account_id.clone()) - .or_default() - .push(record); - } - - by_account - .into_iter() - .map(|(account_id, mut records)| { - records.sort_by(|a, b| { - b.request_count - .cmp(&a.request_count) - .then_with(|| b.last_used_at.cmp(&a.last_used_at)) - .then_with(|| a.model.cmp(&b.model)) - }); - ( - account_id, - records.into_iter().map(model_usage_data).collect(), - ) - }) - .collect() -} - -fn model_usage_data(usage: AccountModelUsageRecord) -> AdminAccountModelUsageData { - let request_count = nonnegative_i64_to_u64(usage.request_count); - let error_count = nonnegative_i64_to_u64(usage.error_count); - let input_tokens = nonnegative_i64_to_u64(usage.input_tokens); - let output_tokens = nonnegative_i64_to_u64(usage.output_tokens); - let cached_tokens = nonnegative_i64_to_u64(usage.cached_tokens); - let total_tokens = input_tokens + output_tokens; - let success_rate = if request_count > 0 { - ((request_count.saturating_sub(error_count)) as f64 / request_count as f64 * 1000.0).round() - / 10.0 - } else { - 0.0 - }; - let total_cost_usd = usage.total_cost_usd; - - AdminAccountModelUsageData { - model: usage.model, - request_count, - request_count_display: format_plain_number(request_count), - success_rate, - success_rate_display: format_percent(success_rate), - input_tokens, - input_tokens_display: format_tokens(input_tokens), - output_tokens, - output_tokens_display: format_tokens(output_tokens), - cached_tokens, - cached_tokens_display: format_tokens(cached_tokens), - total_tokens, - total_tokens_display: format_tokens(total_tokens), - total_cost_usd, - total_cost_usd_display: format_cost(total_cost_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()), - } -} - -struct ParsedAccountUpdate { - id: String, - update: AdminAccountUpdate, -} - -fn parse_account_update(payload: &Value) -> Result { - let payload = parse_editable_update( - payload, - EditableUpdateMessages { - object_required: "Account update request must be an object", - invalid: "Invalid account update request", - empty_update: "Account update request must include editable fields", - unknown_field_editable: true, - }, - )?; - let update = AdminAccountUpdate { - label: payload.label.map(|label| { - label.and_then(|value| { - let value = value.trim(); - (!value.is_empty()).then(|| value.to_string()) - }) - }), - status: payload.status, - }; - if !update.any() { - return Err(AdminError::bad_request( - "Account update request must include editable fields", - )); - } - Ok(ParsedAccountUpdate { - id: payload.id, - update, - }) -} - -fn account_error(error: &AdminAccountError) -> AdminError { - match error { - AdminAccountError::InvalidStatus(_) - | AdminAccountError::LabelTooLong - | AdminAccountError::EmptyIds - | AdminAccountError::NoImportableAccounts - | AdminAccountError::NoModels - | AdminAccountError::InvalidAccessTokenExpiresAt - | AdminAccountError::TokenRequired - | AdminAccountError::InvalidToken(_) - | AdminAccountError::RefreshTokenExchange(_) - | AdminAccountError::OAuthSessionInvalid - | AdminAccountError::OAuthCallbackInvalid - | AdminAccountError::OAuthStateMismatch - | AdminAccountError::NoValidCookies => AdminError::bad_request(error.to_string()), - AdminAccountError::OAuthCodeExchange(_) => AdminError::bad_gateway(error.to_string()), - AdminAccountError::NotFound => account_not_found(), - AdminAccountError::Inactive(_) => AdminError::conflict(error.to_string()), - _ => AdminError::internal(error.to_string()), - } -} - -fn account_refresh_outcome_str(outcome: AdminAccountRefreshOutcome) -> &'static str { - match outcome { - AdminAccountRefreshOutcome::Alive => "alive", - AdminAccountRefreshOutcome::Dead => "dead", - AdminAccountRefreshOutcome::Skipped => "skipped", - } -} - -fn account_not_found() -> AdminError { - AdminError::not_found("Account not found") -} - -fn account_export_ids(value: Option<&str>) -> Vec { - value - .into_iter() - .flat_map(|ids| ids.split(',')) - .map(str::trim) - .filter(|id| !id.is_empty()) - .map(ToString::to_string) - .collect() -} diff --git a/backend/src/admin/auth/mod.rs b/backend/src/admin/auth/mod.rs deleted file mode 100644 index 87d8bfc98..000000000 --- a/backend/src/admin/auth/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -//! 管理端认证与本地会话模块。 - -pub mod service; -pub mod session; diff --git a/backend/src/admin/auth/service.rs b/backend/src/admin/auth/service.rs deleted file mode 100644 index 92ed4365c..000000000 --- a/backend/src/admin/auth/service.rs +++ /dev/null @@ -1,341 +0,0 @@ -//! 管理员会话存储与认证领域逻辑。 - -use std::{collections::HashMap, sync::Arc}; - -use chrono::{DateTime, Duration, Utc}; -use sqlx::SqlitePool; -use thiserror::Error; -use tokio::sync::Mutex; -use uuid::Uuid; - -use crate::infra::identity::{hash_admin_password, verify_admin_password}; - -const LOGIN_FAILURE_WINDOW: Duration = Duration::minutes(15); -const LOGIN_LOCK_DURATION: Duration = Duration::minutes(15); -const LOGIN_MAX_FAILURES: u32 = 5; - -/// 管理员登录成功后的会话。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminLoginSession { - /// 会话 ID。 - pub session_id: String, - /// 过期时间。 - pub expires_at: DateTime, -} - -/// 管理员会话错误。 -#[derive(Debug, Error)] -pub enum AdminSessionError { - /// 校验管理员会话失败。 - #[error("failed to validate admin session")] - Validate, - /// 计算管理员密码哈希失败。 - #[error("failed to hash admin password")] - HashPassword, - /// 创建默认管理员失败。 - #[error("failed to create default admin user")] - CreateAdmin, - /// 读取管理员用户失败。 - #[error("failed to load admin user")] - LoadAdmin, - /// 校验管理员密码失败。 - #[error("failed to verify admin password")] - VerifyPassword, - /// 创建管理员会话失败。 - #[error("failed to create admin session")] - CreateSession, - /// 删除管理员会话失败。 - #[error("failed to delete admin session")] - DeleteSession, - /// 管理员登录尝试过于频繁。 - #[error("too many admin login attempts")] - LoginThrottled, -} - -/// 管理员会话服务。 -#[derive(Clone)] -pub struct AdminSessionService { - store: SqliteAdminSessionStore, - default_username: String, - session_ttl_minutes: u64, - login_failures: Arc>>, -} - -#[derive(Debug, Clone)] -struct LoginFailureState { - failures: u32, - first_failed_at: DateTime, - locked_until: Option>, -} - -impl AdminSessionService { - /// 构造服务。 - pub fn new( - store: SqliteAdminSessionStore, - default_username: String, - session_ttl_minutes: u64, - ) -> Self { - Self { - store, - default_username, - session_ttl_minutes, - login_failures: Arc::new(Mutex::new(HashMap::new())), - } - } - - /// 校验管理员会话。 - pub async fn validate(&self, session_id: Option<&str>) -> Result { - let Some(session_id) = session_id else { - return Ok(false); - }; - self.store - .validate_session(session_id) - .await - .map_err(|_| AdminSessionError::Validate) - } - - /// 确保默认管理员存在。 - pub async fn ensure_default_admin(&self, password: &str) -> Result { - let password_hash = - hash_admin_password(password).map_err(|_| AdminSessionError::HashPassword)?; - self.store - .ensure_default_admin(&password_hash) - .await - .map_err(|_| AdminSessionError::CreateAdmin) - } - - /// 登录并创建管理员会话。 - pub async fn login( - &self, - source: &str, - username: Option<&str>, - password: &str, - ) -> Result, AdminSessionError> { - let source = login_source_key(source); - self.ensure_login_allowed(&source).await?; - - let username = username.unwrap_or(&self.default_username); - if username != self.default_username.as_str() { - self.record_failed_login(&source).await; - return Ok(None); - } - let Some(admin) = self - .store - .load_first_admin() - .await - .map_err(|_| AdminSessionError::LoadAdmin)? - else { - self.record_failed_login(&source).await; - return Ok(None); - }; - let password_matches = verify_admin_password(password, &admin.password_hash) - .map_err(|_| AdminSessionError::VerifyPassword)?; - if !password_matches { - self.record_failed_login(&source).await; - return Ok(None); - } - self.clear_failed_login(&source).await; - let session_id = format!("sess_{}", Uuid::new_v4().simple()); - let ttl_minutes = self.session_ttl_minutes.min(i64::MAX as u64) as i64; - let expires_at = Utc::now() + Duration::minutes(ttl_minutes); - self.store - .create_session(&session_id, &admin.id, expires_at) - .await - .map_err(|_| AdminSessionError::CreateSession)?; - Ok(Some(AdminLoginSession { - session_id, - expires_at, - })) - } - - /// 删除管理员会话。 - pub async fn delete_session(&self, session_id: &str) -> Result { - self.store - .delete_session(session_id) - .await - .map_err(|_| AdminSessionError::DeleteSession) - } - - async fn ensure_login_allowed(&self, source: &str) -> Result<(), AdminSessionError> { - let now = Utc::now(); - let mut failures = self.login_failures.lock().await; - retain_active_login_failures(&mut failures, now); - if failures - .get(source) - .and_then(|state| state.locked_until) - .is_some_and(|locked_until| locked_until > now) - { - return Err(AdminSessionError::LoginThrottled); - } - Ok(()) - } - - async fn record_failed_login(&self, source: &str) { - let now = Utc::now(); - let mut failures = self.login_failures.lock().await; - retain_active_login_failures(&mut failures, now); - let state = failures - .entry(source.to_string()) - .or_insert_with(|| LoginFailureState { - failures: 0, - first_failed_at: now, - locked_until: None, - }); - if now - state.first_failed_at > LOGIN_FAILURE_WINDOW { - state.failures = 0; - state.first_failed_at = now; - state.locked_until = None; - } - state.failures += 1; - if state.failures >= LOGIN_MAX_FAILURES { - state.locked_until = Some(now + LOGIN_LOCK_DURATION); - } - } - - async fn clear_failed_login(&self, source: &str) { - self.login_failures.lock().await.remove(source); - } -} - -fn login_source_key(source: &str) -> String { - let source = source.trim(); - if source.is_empty() { - "unknown".to_string() - } else { - source.to_string() - } -} - -fn retain_active_login_failures( - failures: &mut HashMap, - now: DateTime, -) { - failures.retain(|_, state| { - state - .locked_until - .is_some_and(|locked_until| locked_until > now) - || now - state.first_failed_at <= LOGIN_FAILURE_WINDOW - }); -} - -// --------------------------------------------------------------------------- -// 已持久化的管理用户 -// --------------------------------------------------------------------------- - -/// 已持久化的管理员用户。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StoredAdminUser { - /// 管理员用户 ID。 - pub id: String, - /// 管理员密码哈希。 - pub password_hash: String, -} - -// --------------------------------------------------------------------------- -// SQLite 管理员会话存储 -// --------------------------------------------------------------------------- - -/// SQLite 管理员会话存储结果。 -pub type SqliteAdminSessionStoreResult = Result; - -/// SQLite 管理员会话存储。 -#[derive(Clone)] -pub struct SqliteAdminSessionStore { - pool: SqlitePool, -} - -impl SqliteAdminSessionStore { - /// 构造存储。 - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - /// 如果还没有管理员用户，则创建默认管理员。 - pub async fn ensure_default_admin( - &self, - password_hash: &str, - ) -> SqliteAdminSessionStoreResult { - let existing: (i64,) = sqlx::query_as("select count(*) from admin_users") - .fetch_one(&self.pool) - .await?; - if existing.0 > 0 { - return Ok(false); - } - - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", - ) - .bind("admin_1") - .bind(password_hash) - .bind(&now) - .bind(&now) - .execute(&self.pool) - .await?; - - Ok(true) - } - - /// 读取第一个管理员用户。 - pub async fn load_first_admin(&self) -> SqliteAdminSessionStoreResult> { - let row = sqlx::query_as::<_, (String, String)>( - "select id, password_hash from admin_users order by created_at asc, id asc limit 1", - ) - .fetch_optional(&self.pool) - .await?; - - Ok(row.map(|(id, password_hash)| StoredAdminUser { id, password_hash })) - } - - /// 创建管理员会话。 - pub async fn create_session( - &self, - session_id: &str, - user_id: &str, - expires_at: DateTime, - ) -> SqliteAdminSessionStoreResult<()> { - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind(user_id) - .bind(expires_at.to_rfc3339()) - .bind(now) - .execute(&self.pool) - .await?; - Ok(()) - } - - /// 校验管理员会话是否存在且未过期。 - pub async fn validate_session(&self, session_id: &str) -> SqliteAdminSessionStoreResult { - let now = Utc::now().to_rfc3339(); - let row = sqlx::query("select 1 from admin_sessions where id = ? and expires_at > ?") - .bind(session_id) - .bind(now) - .fetch_optional(&self.pool) - .await?; - Ok(row.is_some()) - } - - /// 删除已经过期的管理员会话。 - pub async fn cleanup_expired_sessions( - &self, - now: DateTime, - ) -> SqliteAdminSessionStoreResult { - let result = sqlx::query("delete from admin_sessions where expires_at <= ?") - .bind(now.to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(result.rows_affected()) - } - - /// 删除指定的管理员会话。 - pub async fn delete_session(&self, session_id: &str) -> SqliteAdminSessionStoreResult { - let result = sqlx::query("delete from admin_sessions where id = ?") - .bind(session_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } -} diff --git a/backend/src/admin/keys/mod.rs b/backend/src/admin/keys/mod.rs deleted file mode 100644 index 6c986dc17..000000000 --- a/backend/src/admin/keys/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -//! 管理端 v1 接口访问 Key 模块。 - -pub mod routes; -pub mod service; diff --git a/backend/src/admin/keys/service.rs b/backend/src/admin/keys/service.rs deleted file mode 100644 index 954cf00f4..000000000 --- a/backend/src/admin/keys/service.rs +++ /dev/null @@ -1,698 +0,0 @@ -//! 客户端 API Key 端口定义、业务服务与 SQLite 存储适配器。 - -use std::{ - collections::{BTreeMap, HashMap}, - sync::{ - atomic::{AtomicBool, Ordering}, - Arc, Mutex, RwLock, - }, - time::Duration, -}; - -use async_trait::async_trait; -use chrono::Utc; -use sqlx::{Row, SqlitePool}; -use thiserror::Error; -use tokio::time::sleep; -use uuid::Uuid; - -use crate::infra::{ - identity::generate_client_api_key, - json::{decode_cursor, encode_cursor, Page}, -}; - -// --------------------------------------------------------------------------- -// 端口定义 -// --------------------------------------------------------------------------- - -/// 客户端 API Key 存储错误。 -#[derive(Debug, Error)] -pub enum ClientKeyStoreError { - /// 底层存储失败。 - #[error("client key store operation failed: {message}")] - OperationFailed { - /// 错误说明。 - message: String, - }, -} - -/// 客户端 API Key 存储结果类型。 -pub type ClientKeyStoreResult = Result; - -/// 提供客户端 API Key 验证能力的端口。 -#[async_trait] -pub trait ClientKeyStore: Send + Sync + 'static { - /// 验证明文客户端 API Key，并在成功时记录使用时间。 - async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult; -} - -const CLIENT_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1); - -// --------------------------------------------------------------------------- -// 业务服务 -// --------------------------------------------------------------------------- - -/// 客户端 API Key 服务。 -#[derive(Clone)] -pub struct ClientKeyService { - store: Arc, -} - -impl ClientKeyService { - /// 构造服务。 - pub fn new(store: Arc) -> Self { - Self { store } - } - - /// 验证客户端 API Key。 - pub async fn verify(&self, plaintext: &str) -> ClientKeyStoreResult { - if !plaintext.starts_with("sk_") { - return Ok(false); - } - self.store.verify_and_touch(plaintext).await - } - - /// 重新加载运行时内存鉴权表。 - pub async fn reload_from_store(&self) -> ClientKeyStoreResult<()> { - self.store.reload_from_store().await - } - - /// 立即刷写待持久化的最近使用时间。 - pub async fn flush_pending_last_used(&self) { - self.store.flush_pending_last_used().await; - } -} - -// --------------------------------------------------------------------------- -// 管理端业务服务 -// --------------------------------------------------------------------------- - -/// 管理端客户端 API Key 服务。 -#[derive(Clone)] -pub struct AdminClientKeyService { - store: SqliteClientKeyStore, - runtime: Arc, -} - -impl AdminClientKeyService { - /// 构造服务。 - pub fn new(store: SqliteClientKeyStore, runtime: Arc) -> Self { - Self { store, runtime } - } - - /// 创建客户端 API Key。 - pub async fn create(&self, name: &str) -> Result { - let name = name.trim(); - if name.is_empty() { - return Err(AdminClientKeyError::EmptyName); - } - self.store - .create(name) - .await - .inspect(|created| self.runtime.upsert_stored(created)) - .map(AdminClientApiKey::from) - .map_err(|_| AdminClientKeyError::Create) - } - - /// 分页列出客户端 API Key。 - pub async fn list( - &self, - cursor: Option, - limit: u32, - ) -> Result, AdminClientKeyError> { - let page = self - .store - .list(cursor, limit) - .await - .map_err(|_| AdminClientKeyError::List)?; - Ok(Page { - items: page - .items - .into_iter() - .map(AdminClientApiKey::from) - .collect(), - next_cursor: page.next_cursor, - }) - } - - /// 按 ID 读取客户端 API Key 元数据。 - pub async fn get( - &self, - key_id: &str, - ) -> Result, AdminClientKeyError> { - self.store - .get(key_id) - .await - .map(|key| key.map(AdminClientApiKey::from)) - .map_err(|_| AdminClientKeyError::List) - } - - /// 更新客户端 API Key 标签。 - pub async fn update_label( - &self, - key_id: &str, - label: Option, - ) -> Result, AdminClientKeyError> { - if label.as_ref().is_some_and(|l| l.chars().count() > 64) { - return Err(AdminClientKeyError::LabelTooLong); - } - self.store - .set_label(key_id, label) - .await - .map(|key| key.map(AdminClientApiKey::from)) - .map_err(|_| AdminClientKeyError::UpdateLabel) - } - - /// 更新客户端 API Key 启用状态。 - pub async fn update_status( - &self, - key_id: &str, - status: &str, - ) -> Result { - let enabled = parse_client_key_status(status)?; - let updated = self - .store - .set_enabled(key_id, enabled) - .await - .map_err(|_| AdminClientKeyError::UpdateStatus)?; - if updated { - if enabled { - let Some(key) = self - .store - .get(key_id) - .await - .map_err(|_| AdminClientKeyError::UpdateStatus)? - else { - return Ok(false); - }; - self.runtime.upsert_stored(&key); - } else { - self.runtime.remove_by_id(key_id); - } - } - Ok(updated) - } - - /// 批量删除客户端 API Key。 - pub async fn batch_delete( - &self, - ids: Vec, - ) -> Result { - if ids.is_empty() { - return Err(AdminClientKeyError::EmptyIds); - } - let mut deleted = 0u32; - let mut not_found = Vec::new(); - for id in ids { - match self.store.delete(&id).await { - Ok(true) => { - self.runtime.remove_by_id(&id); - deleted += 1; - } - Ok(false) => not_found.push(id), - Err(_) => return Err(AdminClientKeyError::Delete), - } - } - Ok(BatchDeleteClientApiKeys { deleted, not_found }) - } -} - -/// 管理端可见的客户端 API Key。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminClientApiKey { - pub id: String, - pub name: String, - pub label: Option, - pub prefix: String, - pub key: String, - pub enabled: bool, - pub created_at: String, - pub last_used_at: Option, -} - -/// 客户端 API Key 批量删除结果。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct BatchDeleteClientApiKeys { - pub deleted: u32, - pub not_found: Vec, -} - -/// 管理端客户端 API Key 错误。 -#[derive(Debug, Error)] -pub enum AdminClientKeyError { - #[error("failed to list client API keys")] - List, - #[error("failed to create client API key")] - Create, - #[error("failed to delete client API key")] - Delete, - #[error("failed to update client API key label")] - UpdateLabel, - #[error("failed to update client API key status")] - UpdateStatus, - #[error("unsupported client API key status: {0}")] - InvalidStatus(String), - #[error("client API key name is required")] - EmptyName, - #[error("client API key ids are required")] - EmptyIds, - #[error("client API key label must be 64 characters or fewer")] - LabelTooLong, -} - -fn parse_client_key_status(status: &str) -> Result { - match status.trim().to_ascii_lowercase().as_str() { - "active" => Ok(true), - "disabled" => Ok(false), - other => Err(AdminClientKeyError::InvalidStatus(other.to_string())), - } -} - -impl From for AdminClientApiKey { - fn from(key: StoredClientApiKey) -> Self { - Self { - id: key.id, - name: key.name, - label: key.label, - prefix: key.prefix, - key: key.key, - enabled: key.enabled, - created_at: key.created_at, - last_used_at: key.last_used_at, - } - } -} - -// --------------------------------------------------------------------------- -// SQLite 存储适配器 -// --------------------------------------------------------------------------- - -/// SQLite 客户端 API Key 存储错误。 -#[derive(Debug, Error)] -pub enum SqliteClientKeyStoreError { - /// 数据库错误。 - #[error("sqlite client key database error: {0}")] - Database(#[from] sqlx::Error), - /// 分页游标无效。 - #[error("invalid client key pagination cursor")] - InvalidCursor, -} - -/// 已持久化的客户端 API Key 元数据。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StoredClientApiKey { - /// API Key 记录 ID。 - pub id: String, - /// API Key 名称。 - pub name: String, - /// 管理员可见标签。 - pub label: Option, - /// 明文 API Key 的短前缀。 - pub prefix: String, - /// 管理端可复制的完整 API Key。 - pub key: String, - /// 是否允许用于 `/v1` 认证。 - pub enabled: bool, - /// 创建时间。 - pub created_at: String, - /// 最近一次成功使用时间。 - pub last_used_at: Option, -} - -/// SQLite 客户端 API Key 存储。 -#[derive(Clone)] -pub struct SqliteClientKeyStore { - pool: SqlitePool, -} - -impl SqliteClientKeyStore { - /// 构造存储适配器。 - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - /// 暴露底层连接池，供集成测试和运行时组合层复用。 - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - - /// 读取所有已启用客户端 API Key，用于运行时内存鉴权表初始化。 - pub async fn list_enabled(&self) -> Result, SqliteClientKeyStoreError> { - let rows = sqlx::query( - "select id, name, label, prefix, key, enabled, created_at, last_used_at from client_api_keys where enabled = 1 order by created_at asc, id asc", - ) - .fetch_all(&self.pool) - .await?; - - Ok(rows - .into_iter() - .map(|row| Self::key_from_row(&row)) - .collect()) - } - - /// 创建新的本地客户端 API Key。 - pub async fn create( - &self, - name: &str, - ) -> Result { - let generated = generate_client_api_key(); - let id = format!("key_{}", Uuid::new_v4().simple()); - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into client_api_keys (id, name, label, prefix, key, enabled, created_at, last_used_at) values (?, ?, null, ?, ?, 1, ?, null)", - ) - .bind(&id) - .bind(name) - .bind(&generated.prefix) - .bind(&generated.key) - .bind(&now) - .execute(&self.pool) - .await?; - - Ok(StoredClientApiKey { - id, - name: name.to_string(), - label: None, - prefix: generated.prefix, - key: generated.key, - enabled: true, - created_at: now, - last_used_at: None, - }) - } - - /// 按创建时间倒序分页列出客户端 API Key。 - pub async fn list( - &self, - cursor: Option, - limit: u32, - ) -> Result, SqliteClientKeyStoreError> { - let fetch_limit = i64::from(limit) + 1; - let rows = if let Some(cursor) = cursor { - let (created_at, id) = - decode_cursor(&cursor).ok_or(SqliteClientKeyStoreError::InvalidCursor)?; - sqlx::query( - "select id, name, label, prefix, key, enabled, created_at, last_used_at from client_api_keys where created_at < ? or (created_at = ? and id < ?) order by created_at desc, id desc limit ?", - ) - .bind(&created_at) - .bind(created_at) - .bind(id) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - } else { - sqlx::query( - "select id, name, label, prefix, key, enabled, created_at, last_used_at from client_api_keys order by created_at desc, id desc limit ?", - ) - .bind(fetch_limit) - .fetch_all(&self.pool) - .await? - }; - - let has_next = rows.len() > limit as usize; - let take_count = rows.len().min(limit as usize); - let items = rows - .into_iter() - .take(take_count) - .map(|row| Self::key_from_row(&row)) - .collect::>(); - let next_cursor = if has_next { - items - .last() - .map(|key| encode_cursor(&key.created_at, &key.id)) - } else { - None - }; - - Ok(Page { items, next_cursor }) - } - - /// 按 ID 读取客户端 API Key 元数据。 - pub async fn get( - &self, - id: &str, - ) -> Result, SqliteClientKeyStoreError> { - let row = sqlx::query( - "select id, name, label, prefix, key, enabled, created_at, last_used_at from client_api_keys where id = ?", - ) - .bind(id) - .fetch_optional(&self.pool) - .await?; - - Ok(row.map(|row| Self::key_from_row(&row))) - } - - /// 更新客户端 API Key 启用状态。 - pub async fn set_enabled( - &self, - id: &str, - enabled: bool, - ) -> Result { - let result = sqlx::query("update client_api_keys set enabled = ? where id = ?") - .bind(if enabled { 1_i64 } else { 0_i64 }) - .bind(id) - .execute(&self.pool) - .await?; - - Ok(result.rows_affected() > 0) - } - - /// 更新客户端 API Key 标签并返回最新元数据。 - pub async fn set_label( - &self, - id: &str, - label: Option, - ) -> Result, SqliteClientKeyStoreError> { - let result = sqlx::query("update client_api_keys set label = ? where id = ?") - .bind(label) - .bind(id) - .execute(&self.pool) - .await?; - if result.rows_affected() == 0 { - return Ok(None); - } - - self.get(id).await - } - - /// 删除客户端 API Key。 - pub async fn delete(&self, id: &str) -> Result { - let result = sqlx::query("delete from client_api_keys where id = ?") - .bind(id) - .execute(&self.pool) - .await?; - - Ok(result.rows_affected() > 0) - } - - /// 更新最近使用时间。 - pub async fn update_last_used_at( - &self, - id: &str, - last_used_at: &str, - ) -> Result { - let result = sqlx::query("update client_api_keys set last_used_at = ? where id = ?") - .bind(last_used_at) - .bind(id) - .execute(&self.pool) - .await?; - - Ok(result.rows_affected() > 0) - } - - fn key_from_row(row: &sqlx::sqlite::SqliteRow) -> StoredClientApiKey { - StoredClientApiKey { - id: row.get("id"), - name: row.get("name"), - label: row.get("label"), - prefix: row.get("prefix"), - key: row.get("key"), - enabled: row.get::("enabled") != 0, - created_at: row.get("created_at"), - last_used_at: row.get("last_used_at"), - } - } -} - -#[async_trait] -impl ClientKeyStore for RuntimeClientKeyStore { - async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult { - let Some(id) = self.lookup_key_id(plaintext) else { - return Ok(false); - }; - - self.queue_last_used_touch(id); - Ok(true) - } -} - -/// 运行时客户端 API Key 内存鉴权表。 -#[derive(Clone)] -pub struct RuntimeClientKeyStore { - store: SqliteClientKeyStore, - cache: Arc>, - pending_last_used: Arc>>, - flush_scheduled: Arc, - flush_delay: Duration, -} - -#[derive(Default)] -struct RuntimeClientKeyCache { - keys_by_plaintext: HashMap, - plaintext_by_id: HashMap, -} - -impl RuntimeClientKeyStore { - /// 构造运行时内存鉴权表。 - pub fn new(store: SqliteClientKeyStore) -> Self { - Self { - store, - cache: Arc::new(RwLock::new(RuntimeClientKeyCache::default())), - pending_last_used: Arc::new(Mutex::new(BTreeMap::new())), - flush_scheduled: Arc::new(AtomicBool::new(false)), - flush_delay: CLIENT_KEY_LAST_USED_FLUSH_DELAY, - } - } - - /// 从 SQLite 加载当前启用的客户端 API Key。 - pub async fn reload_from_store(&self) -> ClientKeyStoreResult<()> { - let keys = self - .store - .list_enabled() - .await - .map_err(|error| map_client_key_store_error(&error))?; - let mut keys_by_plaintext = HashMap::with_capacity(keys.len()); - let mut plaintext_by_id = HashMap::with_capacity(keys.len()); - for key in keys { - keys_by_plaintext.insert(key.key.clone(), key.id.clone()); - plaintext_by_id.insert(key.id, key.key); - } - - *self - .cache - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) = RuntimeClientKeyCache { - keys_by_plaintext, - plaintext_by_id, - }; - Ok(()) - } - - fn lookup_key_id(&self, plaintext: &str) -> Option { - self.cache - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .keys_by_plaintext - .get(plaintext) - .cloned() - } - - fn upsert_stored(&self, key: &StoredClientApiKey) { - if key.enabled { - self.upsert_enabled_key(&key.id, &key.key); - } else { - self.remove_by_id(&key.id); - } - } - - fn upsert_enabled_key(&self, id: &str, plaintext: &str) { - let mut cache = self - .cache - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(old_plaintext) = cache.plaintext_by_id.remove(id) { - cache.keys_by_plaintext.remove(&old_plaintext); - } - cache - .keys_by_plaintext - .insert(plaintext.to_string(), id.to_string()); - cache - .plaintext_by_id - .insert(id.to_string(), plaintext.to_string()); - } - - fn remove_by_id(&self, id: &str) { - let mut cache = self - .cache - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner); - if let Some(plaintext) = cache.plaintext_by_id.remove(id) { - cache.keys_by_plaintext.remove(&plaintext); - } - drop(cache); - self.pending_last_used - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .remove(id); - } - - fn queue_last_used_touch(&self, id: String) { - self.pending_last_used - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(id, Utc::now().to_rfc3339()); - self.schedule_flush(); - } - - fn schedule_flush(&self) { - if self - .flush_scheduled - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) - .is_err() - { - return; - } - - let runtime = self.clone(); - tokio::spawn(async move { - sleep(runtime.flush_delay).await; - runtime.flush_pending_last_used().await; - }); - } - - /// 立即刷写待持久化的最近使用时间。 - pub async fn flush_pending_last_used(&self) { - let updates = self.take_pending_last_used(); - let mut failed = BTreeMap::new(); - for (id, last_used_at) in updates { - if let Err(error) = self.store.update_last_used_at(&id, &last_used_at).await { - tracing::warn!( - key_id = %id, - error = %error, - "failed to flush client api key last_used_at" - ); - failed.insert(id, last_used_at); - } - } - if !failed.is_empty() { - self.pending_last_used - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .extend(failed); - } - self.flush_scheduled.store(false, Ordering::Release); - if !self - .pending_last_used - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_empty() - { - self.schedule_flush(); - } - } - - fn take_pending_last_used(&self) -> BTreeMap { - std::mem::take( - &mut *self - .pending_last_used - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner), - ) - } -} - -fn map_client_key_store_error(error: &SqliteClientKeyStoreError) -> ClientKeyStoreError { - ClientKeyStoreError::OperationFailed { - message: error.to_string(), - } -} diff --git a/backend/src/admin/mod.rs b/backend/src/admin/mod.rs deleted file mode 100644 index 4708c8030..000000000 --- a/backend/src/admin/mod.rs +++ /dev/null @@ -1,19 +0,0 @@ -//! 管理端 HTTP API。 - -/// 账号 CRUD 处理器。 -pub mod accounts; -/// 管理端账号认证状态与本地管理员会话处理器。 -pub mod auth; -/// v1 接口访问 Key 处理器。 -pub mod keys; -/// 管理面日志、用量与诊断统计。 -pub mod monitoring; -/// 响应封装。 -pub mod response; -/// 管理端路由。 -pub mod router; -/// 设置处理器。 -pub mod settings; -/// 系统版本与在线更新处理器。 -pub mod system; -pub(crate) mod update_payload; diff --git a/backend/src/admin/monitoring/mod.rs b/backend/src/admin/monitoring/mod.rs deleted file mode 100644 index e4d6e0469..000000000 --- a/backend/src/admin/monitoring/mod.rs +++ /dev/null @@ -1,14 +0,0 @@ -//! 管理面监控模块：日志、用量统计和 Dashboard 聚合。 - -pub mod account_usage_service; -pub mod account_usage_store; -pub(crate) mod billing; -pub mod dashboard; -pub mod diagnostics; -pub mod ops_error_model; -pub mod ops_error_service; -pub mod ops_error_store; -pub mod usage_record_model; -pub mod usage_record_routes; -pub mod usage_record_service; -pub mod usage_record_store; diff --git a/backend/src/admin/monitoring/ops_error_service.rs b/backend/src/admin/monitoring/ops_error_service.rs deleted file mode 100644 index c005bc8e5..000000000 --- a/backend/src/admin/monitoring/ops_error_service.rs +++ /dev/null @@ -1,68 +0,0 @@ -//! 管理端运维错误事件服务。 - -use chrono::Utc; -use thiserror::Error; - -use crate::admin::monitoring::{ - ops_error_model::OpsErrorLog, ops_error_store::SqliteOpsErrorLogStore, -}; - -/// 管理端运维错误事件错误。 -#[derive(Debug, Error)] -pub enum AdminOpsErrorLogError { - /// 写入失败。 - #[error("failed to append ops error log")] - Append, - /// 保留期清理失败。 - #[error("failed to trim expired ops error logs")] - Retention, -} - -/// 管理端运维错误事件服务。 -#[derive(Clone)] -pub struct AdminOpsErrorLogService { - store: SqliteOpsErrorLogStore, - capture_body: bool, -} - -impl AdminOpsErrorLogService { - /// 构造管理端运维错误事件服务。 - pub fn new(store: SqliteOpsErrorLogStore, capture_body: bool) -> Self { - Self { - store, - capture_body, - } - } - - /// 记录运维错误事件。 - pub async fn record(&self, mut event: OpsErrorLog) -> Result<(), AdminOpsErrorLogError> { - apply_capture_body_policy(&mut event, self.capture_body); - self.store - .append(&event) - .await - .map_err(|_| AdminOpsErrorLogError::Append)?; - self.store - .trim_to_retention(Utc::now()) - .await - .map_err(|_| AdminOpsErrorLogError::Retention)?; - Ok(()) - } -} - -fn apply_capture_body_policy(event: &mut OpsErrorLog, capture_body: bool) { - if capture_body { - return; - } - let Some(metadata) = event.metadata.as_object_mut() else { - return; - }; - for key in [ - "body", - "rawBody", - "requestBody", - "responseBody", - "upstreamBody", - ] { - metadata.remove(key); - } -} diff --git a/backend/src/admin/monitoring/ops_error_store.rs b/backend/src/admin/monitoring/ops_error_store.rs deleted file mode 100644 index 6392e7594..000000000 --- a/backend/src/admin/monitoring/ops_error_store.rs +++ /dev/null @@ -1,261 +0,0 @@ -//! 运维错误事件存储实现（SQLite）。 - -use chrono::{DateTime, Duration, Utc}; -use serde_json::Value; -use sqlx::SqlitePool; -use thiserror::Error; - -use crate::{ - admin::monitoring::{ - ops_error_model::OpsErrorLog, - usage_record_model::{metadata_i64, metadata_service_tier, metadata_string}, - usage_record_store::USAGE_RECORD_RETENTION_DAYS, - }, - infra::time::china_quarter_hour_start, -}; - -/// SQLite 运维错误事件错误。 -#[derive(Debug, Error)] -pub enum SqliteOpsErrorLogStoreError { - /// 数据库错误。 - #[error("sqlite ops error log database error: {0}")] - Database(#[from] sqlx::Error), -} - -/// SQLite 运维错误事件结果。 -pub type SqliteOpsErrorLogStoreResult = Result; - -/// SQLite 运维错误事件存储。 -#[derive(Clone)] -pub struct SqliteOpsErrorLogStore { - pool: SqlitePool, -} - -impl SqliteOpsErrorLogStore { - /// 构造存储。 - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - /// 写入运维错误事件。 - pub async fn append(&self, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> { - append_event(&self.pool, event).await - } - - /// 按保留期清理运维错误事件。 - pub async fn trim_to_retention(&self, now: DateTime) -> SqliteOpsErrorLogStoreResult { - trim_to_retention(&self.pool, now).await - } - - /// 清空运维错误事件。 - pub async fn clear(&self) -> SqliteOpsErrorLogStoreResult { - let result = sqlx::query("delete from ops_error_logs") - .execute(&self.pool) - .await?; - Ok(result.rows_affected()) - } - - /// 读取错误事件数量。 - pub async fn count(&self) -> SqliteOpsErrorLogStoreResult { - let (total,): (i64,) = sqlx::query_as("select count(*) from ops_error_logs") - .fetch_one(&self.pool) - .await?; - Ok(total.max(0) as u64) - } -} - -#[derive(Debug, Clone)] -struct OpsErrorSummaryFields { - client_status_code: Option, - upstream_status_code: Option, - transport: Option, - attempt_index: Option, - failure_class: Option, - response_id: Option, - upstream_request_id: Option, -} - -fn ops_error_summary_fields(event: &OpsErrorLog) -> OpsErrorSummaryFields { - OpsErrorSummaryFields { - client_status_code: event.client_status_code.or_else(|| { - metadata_i64( - &event.metadata, - &["clientStatusCode", "client_status_code", "clientStatus"], - ) - }), - upstream_status_code: event.upstream_status_code.or_else(|| { - metadata_i64( - &event.metadata, - &[ - "upstreamStatus", - "upstreamStatusCode", - "upstream_status_code", - ], - ) - }), - transport: event - .transport - .clone() - .or_else(|| metadata_string(&event.metadata, &["transport"])), - attempt_index: event - .attempt_index - .or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"])), - failure_class: event - .failure_class - .clone() - .or_else(|| metadata_string(&event.metadata, &["failureClass", "failure_class"])), - response_id: event - .response_id - .clone() - .or_else(|| metadata_string(&event.metadata, &["responseId", "response_id"])), - upstream_request_id: event.upstream_request_id.clone().or_else(|| { - metadata_string( - &event.metadata, - &[ - "upstreamRequestId", - "upstream_request_id", - "openaiRequestId", - ], - ) - }), - } -} - -async fn append_event(pool: &SqlitePool, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> { - let summary = ops_error_summary_fields(event); - sqlx::query( - "insert into ops_error_logs (id, request_id, kind, account_id, route, model, status_code, client_status_code, upstream_status_code, transport, attempt_index, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(&event.id) - .bind(&event.request_id) - .bind(&event.kind) - .bind(&event.account_id) - .bind(&event.route) - .bind(&event.model) - .bind(event.status_code) - .bind(summary.client_status_code) - .bind(summary.upstream_status_code) - .bind(summary.transport) - .bind(summary.attempt_index) - .bind(summary.failure_class) - .bind(summary.response_id) - .bind(summary.upstream_request_id) - .bind(event.latency_ms) - .bind(&event.message) - .bind(event.metadata.to_string()) - .bind(event.created_at.to_rfc3339()) - .execute(pool) - .await?; - append_error_time_bucket(pool, event).await?; - Ok(()) -} - -async fn append_error_time_bucket( - pool: &SqlitePool, - event: &OpsErrorLog, -) -> SqliteOpsErrorLogStoreResult<()> { - let bucket_start = china_quarter_hour_start(event.created_at); - let input_tokens = metadata_usage_i64(&event.metadata, "inputTokens"); - let output_tokens = metadata_usage_i64(&event.metadata, "outputTokens"); - let cached_tokens = metadata_usage_i64(&event.metadata, "cachedTokens"); - let first_token_latency = metadata_first_token_latency(&event.metadata).unwrap_or(0); - let first_token_latency_count = i64::from(first_token_latency > 0); - let latency = event.latency_ms.filter(|value| *value > 0).unwrap_or(0); - let latency_count = i64::from(latency > 0); - let now = Utc::now().to_rfc3339(); - - sqlx::query( - r" -insert into usage_time_buckets ( - bucket_start, - account_id, - model, - service_tier, - request_count, - error_count, - input_tokens, - output_tokens, - cached_tokens, - first_token_latency_sum, - first_token_latency_count, - latency_sum, - latency_count, - max_latency_ms, - min_latency_ms, - updated_at -) values (?, ?, ?, ?, 1, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -on conflict(bucket_start, account_id, model, service_tier) do update set - request_count = request_count + excluded.request_count, - error_count = error_count + excluded.error_count, - input_tokens = input_tokens + excluded.input_tokens, - output_tokens = output_tokens + excluded.output_tokens, - cached_tokens = cached_tokens + excluded.cached_tokens, - first_token_latency_sum = first_token_latency_sum + excluded.first_token_latency_sum, - first_token_latency_count = first_token_latency_count + excluded.first_token_latency_count, - latency_sum = latency_sum + excluded.latency_sum, - latency_count = latency_count + excluded.latency_count, - max_latency_ms = max(max_latency_ms, excluded.max_latency_ms), - min_latency_ms = case - when min_latency_ms = 0 then excluded.min_latency_ms - when excluded.min_latency_ms = 0 then min_latency_ms - else min(min_latency_ms, excluded.min_latency_ms) - end, - updated_at = excluded.updated_at", - ) - .bind(bucket_start.to_rfc3339()) - .bind(event.account_id.as_deref().unwrap_or_default()) - .bind(event.model.as_deref().unwrap_or_default()) - .bind(metadata_service_tier(&event.metadata).unwrap_or_default()) - .bind(input_tokens) - .bind(output_tokens) - .bind(cached_tokens) - .bind(first_token_latency) - .bind(first_token_latency_count) - .bind(latency) - .bind(latency_count) - .bind(latency) - .bind(latency) - .bind(now) - .execute(pool) - .await?; - Ok(()) -} - -fn metadata_usage_i64(metadata: &Value, field: &str) -> i64 { - metadata - .get("usage") - .and_then(|usage| usage.get(field)) - .or_else(|| metadata.get(field)) - .and_then(Value::as_i64) - .filter(|value| *value > 0) - .unwrap_or(0) -} - -fn metadata_first_token_latency(metadata: &Value) -> Option { - [ - "firstTokenMs", - "first_token_ms", - "firstTokenLatencyMs", - "first_token_latency_ms", - ] - .into_iter() - .find_map(|field| { - metadata - .get(field) - .or_else(|| metadata.get("usage").and_then(|usage| usage.get(field))) - .and_then(Value::as_i64) - .filter(|value| *value > 0) - }) -} - -async fn trim_to_retention( - pool: &SqlitePool, - now: DateTime, -) -> SqliteOpsErrorLogStoreResult { - let cutoff = now - Duration::days(USAGE_RECORD_RETENTION_DAYS); - let result = sqlx::query("delete from ops_error_logs where created_at < ?") - .bind(cutoff.to_rfc3339()) - .execute(pool) - .await?; - Ok(result.rows_affected()) -} diff --git a/backend/src/admin/monitoring/usage_record_store.rs b/backend/src/admin/monitoring/usage_record_store.rs deleted file mode 100644 index 3edba85b8..000000000 --- a/backend/src/admin/monitoring/usage_record_store.rs +++ /dev/null @@ -1,1114 +0,0 @@ -//! 使用记录存储实现（SQLite）。 - -use std::collections::{BTreeMap, HashMap}; - -use chrono::{DateTime, Duration, Utc}; -use serde_json::Value; -use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool}; -use thiserror::Error; - -use crate::admin::monitoring::{ - billing, - usage_record_model::{ - metadata_i64, metadata_service_tier, metadata_string, UsageRecord, UsageRecordLevel, - }, -}; -use crate::infra::{ - format::optional_nonnegative_i64_to_u64, - json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page}, - time::{china_quarter_hour_start, parse_rfc3339_utc as parse_rfc3339}, -}; - -/// 使用记录保留天数。 -pub const USAGE_RECORD_RETENTION_DAYS: i64 = 30; -/// 默认不保存请求/响应体，避免把用户内容写入诊断日志。 -pub const DEFAULT_USAGE_RECORD_CAPTURE_BODY: bool = false; - -/// SQLite 使用记录错误。 -#[derive(Debug, Error)] -pub enum SqliteUsageRecordStoreError { - /// 数据库错误。 - #[error("sqlite usage record database error: {0}")] - Database(#[from] sqlx::Error), - /// JSON 错误。 - #[error("sqlite usage record json error: {0}")] - Json(#[from] serde_json::Error), - /// 时间格式错误。 - #[error("sqlite usage record timestamp error: {0}")] - Timestamp(#[from] chrono::ParseError), - /// 事件等级非法。 - #[error("invalid event level: {0}")] - InvalidLevel(String), - /// 分页游标非法。 - #[error("invalid usage record pagination cursor")] - InvalidCursor, -} - -/// SQLite 使用记录结果。 -pub type SqliteUsageRecordStoreResult = Result; - -const USAGE_RECORD_SERVICE_TIER_SQL: &str = - "nullif(trim(json_extract(metadata_json, '$.serviceTier')), '')"; -const USAGE_RECORD_SELECT_SQL: &str = "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"; - -/// 使用记录查询过滤器。 -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct UsageRecordFilter { - /// 事件类别。 - pub kind: Option, - /// 事件等级。 - pub level: Option, - /// 请求 ID。 - pub request_id: Option, - /// 账号 ID。 - pub account_id: Option, - /// 路由。 - pub route: Option, - /// 模型。 - pub model: Option, - /// HTTP 状态码。 - pub status_code: Option, - /// 上游传输方式。 - pub transport: Option, - /// 同一请求内的上游尝试序号。 - pub attempt_index: Option, - /// 上游 HTTP 状态码。 - pub upstream_status_code: Option, - /// 失败分类。 - pub failure_class: Option, - /// 上游响应 ID。 - pub response_id: Option, - /// 上游请求 ID。 - pub upstream_request_id: Option, - /// 搜索关键词。 - pub search: Option, - /// 起始时间。 - pub start_time: Option>, - /// 结束时间。 - pub end_time: Option>, -} - -/// SQLite 使用记录存储。 -#[derive(Clone)] -pub struct SqliteUsageRecordStore { - pool: SqlitePool, -} - -impl SqliteUsageRecordStore { - /// 构造存储。 - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - /// 暴露底层连接池，供运行时组合层复用。 - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - - /// 写入使用记录。 - pub async fn append(&self, event: &UsageRecord) -> SqliteUsageRecordStoreResult<()> { - append_event(&self.pool, event).await - } - - /// 按保留期清理使用记录。 - pub async fn trim_to_retention(&self, now: DateTime) -> SqliteUsageRecordStoreResult { - trim_to_retention(&self.pool, now).await - } - - /// 分页查询使用记录。 - pub async fn list( - &self, - filter: UsageRecordFilter, - cursor: Option, - limit: u32, - ) -> SqliteUsageRecordStoreResult> { - let fetch_limit = i64::from(limit) + 1; - let mut builder = QueryBuilder::::new(USAGE_RECORD_SELECT_SQL); - push_filter(&mut builder, &filter, cursor.as_deref())?; - builder.push(" order by created_at desc, id desc limit "); - builder.push_bind(fetch_limit); - - let rows = builder.build().fetch_all(&self.pool).await?; - let has_next = rows.len() > limit as usize; - let take_count = rows.len().min(limit as usize); - let items = rows - .into_iter() - .take(take_count) - .map(|row| usage_record_from_row(&row)) - .collect::>>()?; - let next_cursor = if has_next { - items - .last() - .map(|event| encode_cursor(&event.created_at.to_rfc3339(), &event.id)) - } else { - None - }; - - Ok(Page { items, next_cursor }) - } - - /// 按页码查询使用记录。 - pub async fn list_page( - &self, - filter: UsageRecordFilter, - page: u32, - page_size: u32, - ) -> SqliteUsageRecordStoreResult> { - let page_size = page_size.clamp(1, 200); - let total = count_usage_records(&self.pool, &filter).await?; - let offset = page_offset(page, page_size); - let mut builder = QueryBuilder::::new(USAGE_RECORD_SELECT_SQL); - push_filter(&mut builder, &filter, None)?; - builder.push(" order by created_at desc, id desc limit "); - builder.push_bind(i64::from(page_size)); - builder.push(" offset "); - builder.push_bind(offset.min(i64::MAX as u64) as i64); - - let rows = builder.build().fetch_all(&self.pool).await?; - let items = rows - .iter() - .map(usage_record_from_row) - .collect::>>()?; - - Ok(NumberedPage { - items, - total, - page: page.max(1), - page_size, - }) - } - - /// 按 ID 读取使用记录。 - pub async fn get(&self, id: &str) -> SqliteUsageRecordStoreResult> { - let mut builder = QueryBuilder::::new(USAGE_RECORD_SELECT_SQL); - builder.push(" where id = "); - builder.push_bind(id); - let row = builder.build().fetch_optional(&self.pool).await?; - row.map(|row| usage_record_from_row(&row)).transpose() - } - - /// 读取使用记录关联账号的邮箱映射。 - pub async fn account_email_map( - &self, - items: &[UsageRecord], - ) -> SqliteUsageRecordStoreResult> { - usage_record_account_email_map(&self.pool, items).await - } - - /// 汇总事件使用记录。 - pub async fn summary( - &self, - filter: UsageRecordFilter, - ) -> SqliteUsageRecordStoreResult { - usage_summary(&self.pool, &filter).await - } - - /// 按账号聚合使用记录。 - pub async fn account_usage( - &self, - filter: UsageRecordFilter, - limit: u32, - ) -> SqliteUsageRecordStoreResult> { - usage_account_usage(&self.pool, &filter, limit).await - } - - /// 按模型来源聚合使用记录分布。 - pub async fn model_distribution( - &self, - filter: UsageRecordFilter, - source: UsageRecordModelSource, - ) -> SqliteUsageRecordStoreResult> { - usage_model_distribution(&self.pool, &filter, source, 8).await - } - - /// 按端点来源聚合使用记录分布。 - pub async fn endpoint_distribution( - &self, - filter: UsageRecordFilter, - source: UsageRecordEndpointSource, - ) -> SqliteUsageRecordStoreResult> { - usage_endpoint_distribution(&self.pool, &filter, source, 8).await - } - - /// 聚合使用记录趋势。 - pub async fn trend( - &self, - filter: UsageRecordFilter, - ) -> SqliteUsageRecordStoreResult> { - usage_trend(&self.pool, &filter).await - } - - /// 清空使用记录。 - pub async fn clear(&self) -> SqliteUsageRecordStoreResult { - let result = sqlx::query("delete from usage_records") - .execute(&self.pool) - .await?; - Ok(result.rows_affected()) - } -} - -#[derive(Debug, Clone)] -struct UsageRecordSummaryFields { - transport: Option, - attempt_index: Option, - upstream_status_code: Option, - failure_class: Option, - response_id: Option, - upstream_request_id: Option, -} - -fn usage_record_summary_fields(event: &UsageRecord) -> UsageRecordSummaryFields { - UsageRecordSummaryFields { - transport: event - .transport - .clone() - .or_else(|| metadata_string(&event.metadata, &["transport"])), - attempt_index: event - .attempt_index - .or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"])), - upstream_status_code: event.upstream_status_code.or_else(|| { - metadata_i64( - &event.metadata, - &[ - "upstreamStatus", - "upstreamStatusCode", - "upstream_status_code", - ], - ) - }), - failure_class: event - .failure_class - .clone() - .or_else(|| metadata_string(&event.metadata, &["failureClass", "failure_class"])), - response_id: event - .response_id - .clone() - .or_else(|| metadata_string(&event.metadata, &["responseId", "response_id"])), - upstream_request_id: event.upstream_request_id.clone().or_else(|| { - metadata_string( - &event.metadata, - &[ - "upstreamRequestId", - "upstream_request_id", - "openaiRequestId", - ], - ) - }), - } -} - -async fn append_event(pool: &SqlitePool, event: &UsageRecord) -> SqliteUsageRecordStoreResult<()> { - let summary = usage_record_summary_fields(event); - sqlx::query( - "insert into usage_records (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) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(&event.id) - .bind(&event.request_id) - .bind(&event.kind) - .bind(level_to_db(event.level)) - .bind(&event.account_id) - .bind(&event.route) - .bind(&event.model) - .bind(event.status_code) - .bind(summary.transport) - .bind(summary.attempt_index) - .bind(summary.upstream_status_code) - .bind(summary.failure_class) - .bind(summary.response_id) - .bind(summary.upstream_request_id) - .bind(event.latency_ms) - .bind(&event.message) - .bind(event.metadata.to_string()) - .bind(event.created_at.to_rfc3339()) - .execute(pool) - .await?; - append_usage_time_bucket(pool, event).await?; - Ok(()) -} - -async fn append_usage_time_bucket( - pool: &SqlitePool, - event: &UsageRecord, -) -> SqliteUsageRecordStoreResult<()> { - let bucket_start = china_quarter_hour_start(event.created_at); - let error_count = i64::from(is_error_usage_record(event)); - let input_tokens = metadata_usage_i64(&event.metadata, "inputTokens"); - let output_tokens = metadata_usage_i64(&event.metadata, "outputTokens"); - let cached_tokens = metadata_usage_i64(&event.metadata, "cachedTokens"); - let first_token_latency = metadata_first_token_latency(&event.metadata).unwrap_or(0); - let first_token_latency_count = i64::from(first_token_latency > 0); - let latency = event.latency_ms.filter(|value| *value > 0).unwrap_or(0); - let latency_count = i64::from(latency > 0); - let now = Utc::now().to_rfc3339(); - - sqlx::query( - r" -insert into usage_time_buckets ( - bucket_start, - account_id, - model, - service_tier, - request_count, - error_count, - input_tokens, - output_tokens, - cached_tokens, - first_token_latency_sum, - first_token_latency_count, - latency_sum, - latency_count, - max_latency_ms, - min_latency_ms, - updated_at -) values (?, ?, ?, ?, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -on conflict(bucket_start, account_id, model, service_tier) do update set - request_count = request_count + excluded.request_count, - error_count = error_count + excluded.error_count, - input_tokens = input_tokens + excluded.input_tokens, - output_tokens = output_tokens + excluded.output_tokens, - cached_tokens = cached_tokens + excluded.cached_tokens, - first_token_latency_sum = first_token_latency_sum + excluded.first_token_latency_sum, - first_token_latency_count = first_token_latency_count + excluded.first_token_latency_count, - latency_sum = latency_sum + excluded.latency_sum, - latency_count = latency_count + excluded.latency_count, - max_latency_ms = max(max_latency_ms, excluded.max_latency_ms), - min_latency_ms = case - when min_latency_ms = 0 then excluded.min_latency_ms - when excluded.min_latency_ms = 0 then min_latency_ms - else min(min_latency_ms, excluded.min_latency_ms) - end, - updated_at = excluded.updated_at", - ) - .bind(bucket_start.to_rfc3339()) - .bind(event.account_id.as_deref().unwrap_or_default()) - .bind(event.model.as_deref().unwrap_or_default()) - .bind(metadata_service_tier(&event.metadata).unwrap_or_default()) - .bind(error_count) - .bind(input_tokens) - .bind(output_tokens) - .bind(cached_tokens) - .bind(first_token_latency) - .bind(first_token_latency_count) - .bind(latency) - .bind(latency_count) - .bind(latency) - .bind(latency) - .bind(now) - .execute(pool) - .await?; - Ok(()) -} - -fn is_error_usage_record(event: &UsageRecord) -> bool { - event.level == UsageRecordLevel::Error || event.status_code.is_some_and(|status| status >= 400) -} - -fn metadata_usage_i64(metadata: &Value, field: &str) -> i64 { - metadata - .get("usage") - .and_then(|usage| usage.get(field)) - .or_else(|| metadata.get(field)) - .and_then(Value::as_i64) - .filter(|value| *value > 0) - .unwrap_or(0) -} - -fn metadata_first_token_latency(metadata: &Value) -> Option { - [ - "firstTokenMs", - "first_token_ms", - "firstTokenLatencyMs", - "first_token_latency_ms", - ] - .into_iter() - .find_map(|field| { - metadata - .get(field) - .or_else(|| metadata.get("usage").and_then(|usage| usage.get(field))) - .and_then(Value::as_i64) - .filter(|value| *value > 0) - }) -} - -async fn trim_to_retention( - pool: &SqlitePool, - now: DateTime, -) -> SqliteUsageRecordStoreResult { - let cutoff = now - Duration::days(USAGE_RECORD_RETENTION_DAYS); - let result = sqlx::query("delete from usage_records where created_at < ?") - .bind(cutoff.to_rfc3339()) - .execute(pool) - .await?; - Ok(result.rows_affected()) -} - -async fn count_usage_records( - pool: &SqlitePool, - filter: &UsageRecordFilter, -) -> SqliteUsageRecordStoreResult { - let mut builder = QueryBuilder::::new("select count(*) from usage_records"); - push_filter(&mut builder, filter, None)?; - let (total,): (i64,) = builder.build_query_as().fetch_one(pool).await?; - Ok(total.max(0).cast_unsigned()) -} - -async fn usage_summary( - pool: &SqlitePool, - filter: &UsageRecordFilter, -) -> SqliteUsageRecordStoreResult { - let mut builder = QueryBuilder::::new( - "select - count(*) as total_requests, - sum(case when level = 'error' or coalesce(status_code, 0) >= 400 then 1 else 0 end) as error_requests, - sum(coalesce(cast(json_extract(metadata_json, '$.usage.inputTokens') as integer), cast(json_extract(metadata_json, '$.inputTokens') as integer), 0)) as input_tokens, - sum(coalesce(cast(json_extract(metadata_json, '$.usage.outputTokens') as integer), cast(json_extract(metadata_json, '$.outputTokens') as integer), 0)) as output_tokens, - sum(coalesce(cast(json_extract(metadata_json, '$.usage.cachedTokens') as integer), cast(json_extract(metadata_json, '$.cachedTokens') as integer), 0)) as cached_tokens, - avg(case when latency_ms > 0 then latency_ms else null end) as average_latency_ms - from usage_records", - ); - push_filter(&mut builder, filter, None)?; - - let row = builder.build().fetch_one(pool).await?; - let input_tokens = optional_nonnegative_i64_to_u64(row.get::, _>("input_tokens")); - let output_tokens = optional_nonnegative_i64_to_u64(row.get::, _>("output_tokens")); - let cached_tokens = optional_nonnegative_i64_to_u64(row.get::, _>("cached_tokens")); - - Ok(UsageRecordSummary { - total_requests: optional_nonnegative_i64_to_u64( - row.get::, _>("total_requests"), - ), - error_requests: optional_nonnegative_i64_to_u64( - row.get::, _>("error_requests"), - ), - input_tokens, - output_tokens, - cached_tokens, - total_tokens: input_tokens + output_tokens, - average_latency_ms: row.get("average_latency_ms"), - }) -} - -async fn usage_account_usage( - pool: &SqlitePool, - filter: &UsageRecordFilter, - limit: u32, -) -> SqliteUsageRecordStoreResult> { - let mut builder = QueryBuilder::::new( - "select - account_id, - sum(coalesce(cast(json_extract(metadata_json, '$.usage.inputTokens') as integer), cast(json_extract(metadata_json, '$.inputTokens') as integer), 0)) as input_tokens, - sum(coalesce(cast(json_extract(metadata_json, '$.usage.outputTokens') as integer), cast(json_extract(metadata_json, '$.outputTokens') as integer), 0)) as output_tokens, - max(created_at) as last_used_at - from usage_records", - ); - push_filter(&mut builder, filter, None)?; - builder.push(" and account_id is not null and trim(account_id) <> ''"); - builder.push(" group by account_id order by last_used_at desc, account_id asc limit "); - builder.push_bind(i64::from(limit.clamp(1, 50))); - - builder - .build() - .fetch_all(pool) - .await? - .into_iter() - .map(|row| { - let input_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("input_tokens")); - let output_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("output_tokens")); - Ok(UsageRecordAccountUsage { - account_id: row.get("account_id"), - input_tokens, - output_tokens, - total_tokens: input_tokens + output_tokens, - last_used_at: parse_rfc3339(&row.get::("last_used_at"))?, - }) - }) - .collect() -} - -async fn usage_record_account_email_map( - pool: &SqlitePool, - items: &[UsageRecord], -) -> SqliteUsageRecordStoreResult> { - let mut account_ids = items - .iter() - .filter_map(|item| item.account_id.as_deref()) - .map(str::to_string) - .collect::>(); - account_ids.sort_unstable(); - account_ids.dedup(); - - if account_ids.is_empty() { - return Ok(HashMap::new()); - } - - let mut builder = QueryBuilder::::new("select id, email from accounts where id in ("); - let mut separated = builder.separated(", "); - for account_id in &account_ids { - separated.push_bind(account_id); - } - separated.push_unseparated(")"); - - let rows = builder.build().fetch_all(pool).await?; - Ok(rows - .into_iter() - .filter_map(|row| { - row.get::, _>("email") - .map(|email| email.trim().to_string()) - .filter(|email| !email.is_empty()) - .map(|email| (row.get::("id"), email)) - }) - .collect()) -} - -async fn usage_model_distribution( - pool: &SqlitePool, - filter: &UsageRecordFilter, - source: UsageRecordModelSource, - limit: u32, -) -> SqliteUsageRecordStoreResult> { - let (group_expr, group_alias) = match source { - UsageRecordModelSource::Requested => ( - "coalesce(nullif(trim(json_extract(metadata_json, '$.requestedModel')), ''), nullif(trim(usage_records.model), ''), '未知模型')", - "requested_model", - ), - UsageRecordModelSource::Upstream => ( - "coalesce( - nullif(trim(json_extract(usage_records.metadata_json, '$.upstreamModel')), ''), - nullif(trim(json_extract(usage_records.metadata_json, '$.requestedModel')), ''), - nullif(trim(usage_records.model), ''), - '未知模型' - )", - "model", - ), - UsageRecordModelSource::Mapping => ( - "case - when coalesce(nullif(trim(json_extract(metadata_json, '$.requestedModel')), ''), nullif(trim(usage_records.model), '')) is not null - then coalesce(nullif(trim(json_extract(metadata_json, '$.requestedModel')), ''), nullif(trim(usage_records.model), ''), '未知模型') - || ' -> ' || - coalesce(nullif(trim(json_extract(metadata_json, '$.upstreamModel')), ''), nullif(trim(json_extract(metadata_json, '$.requestedModel')), ''), nullif(trim(usage_records.model), ''), '未知模型') - else '未知模型' - end", - "model_mapping", - ), - }; - usage_breakdown(pool, filter, group_expr, group_alias, limit).await -} - -async fn usage_endpoint_distribution( - pool: &SqlitePool, - filter: &UsageRecordFilter, - source: UsageRecordEndpointSource, - limit: u32, -) -> SqliteUsageRecordStoreResult> { - let (group_expr, group_alias) = match source { - UsageRecordEndpointSource::Inbound => { - ("coalesce(nullif(route, ''), '未知端点')", "endpoint") - } - UsageRecordEndpointSource::Upstream => ( - "coalesce(nullif(trim(json_extract(metadata_json, '$.upstreamEndpoint')), ''), nullif(trim(json_extract(metadata_json, '$.upstreamRoute')), ''), nullif(trim(json_extract(metadata_json, '$.upstreamPath')), ''), nullif(route, ''), '未知端点')", - "upstream_endpoint", - ), - UsageRecordEndpointSource::Path => ( - "coalesce(nullif(route, ''), '未知端点') || ' -> ' || coalesce(nullif(trim(json_extract(metadata_json, '$.upstreamEndpoint')), ''), nullif(trim(json_extract(metadata_json, '$.upstreamRoute')), ''), nullif(trim(json_extract(metadata_json, '$.upstreamPath')), ''), nullif(route, ''), '未知端点')", - "endpoint_path", - ), - }; - usage_breakdown(pool, filter, group_expr, group_alias, limit).await -} - -async fn usage_breakdown( - pool: &SqlitePool, - filter: &UsageRecordFilter, - group_expr: &str, - group_alias: &str, - limit: u32, -) -> SqliteUsageRecordStoreResult> { - let mut builder = QueryBuilder::::new("select "); - builder.push(group_expr); - builder.push(" as "); - builder.push(group_alias); - builder.push( - ", count(*) as request_count, - sum(coalesce(cast(json_extract(metadata_json, '$.usage.inputTokens') as integer), cast(json_extract(metadata_json, '$.inputTokens') as integer), 0)) as input_tokens, - sum(coalesce(cast(json_extract(metadata_json, '$.usage.outputTokens') as integer), cast(json_extract(metadata_json, '$.outputTokens') as integer), 0)) as output_tokens, - sum(coalesce(cast(json_extract(metadata_json, '$.usage.cachedTokens') as integer), cast(json_extract(metadata_json, '$.cachedTokens') as integer), 0)) as cached_tokens, - coalesce( - nullif(trim(json_extract(metadata_json, '$.upstreamModel')), ''), - nullif(trim(json_extract(metadata_json, '$.requestedModel')), ''), - nullif(trim(usage_records.model), ''), - '未知模型' - ) as billing_model, ", - ); - builder.push(USAGE_RECORD_SERVICE_TIER_SQL); - builder.push( - " as service_tier, - avg(case when latency_ms > 0 then latency_ms else null end) as average_latency_ms - from usage_records", - ); - push_filter(&mut builder, filter, None)?; - builder.push(" group by "); - builder.push(group_alias); - builder.push(", billing_model, service_tier"); - builder.push(" order by request_count desc, "); - builder.push(group_alias); - builder.push(" asc limit "); - builder.push_bind(i64::from(limit.clamp(1, 50) * 8)); - - let rows = builder.build().fetch_all(pool).await?; - let mut items = Vec::::new(); - for row in rows { - let name: String = row.get(group_alias); - let input_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("input_tokens")); - let output_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("output_tokens")); - let cached_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("cached_tokens")); - let cost = usage_breakdown_cost( - input_tokens, - output_tokens, - cached_tokens, - &row.get::("billing_model"), - row.get::, _>("service_tier").as_deref(), - ); - let request_count = - optional_nonnegative_i64_to_u64(row.get::, _>("request_count")); - let average_latency_ms = row.get::, _>("average_latency_ms"); - - if let Some(item) = items.iter_mut().find(|item| item.name == name) { - let previous_latency_count = item.request_count as f64; - let next_latency_count = request_count as f64; - item.request_count += request_count; - item.input_tokens += input_tokens; - item.output_tokens += output_tokens; - item.cached_tokens += cached_tokens; - item.total_tokens += input_tokens + output_tokens; - item.cost += cost; - item.actual_cost += cost; - item.account_cost += cost; - item.average_latency_ms = weighted_average_latency( - item.average_latency_ms, - previous_latency_count, - average_latency_ms, - next_latency_count, - ); - } else { - items.push(UsageRecordBreakdown { - name, - request_count, - input_tokens, - output_tokens, - cached_tokens, - total_tokens: input_tokens + output_tokens, - cost, - actual_cost: cost, - account_cost: cost, - average_latency_ms, - }); - } - } - - items.sort_by(|left, right| { - right - .request_count - .cmp(&left.request_count) - .then_with(|| left.name.cmp(&right.name)) - }); - items.truncate(limit.clamp(1, 50) as usize); - Ok(items) -} - -async fn usage_trend( - pool: &SqlitePool, - filter: &UsageRecordFilter, -) -> SqliteUsageRecordStoreResult> { - let mut builder = QueryBuilder::::new( - "select - substr(datetime(created_at, '+8 hours'), 1, 10) as date, - coalesce(cast(json_extract(metadata_json, '$.usage.inputTokens') as integer), cast(json_extract(metadata_json, '$.inputTokens') as integer), 0) as input_tokens, - coalesce(cast(json_extract(metadata_json, '$.usage.outputTokens') as integer), cast(json_extract(metadata_json, '$.outputTokens') as integer), 0) as output_tokens, - coalesce(cast(json_extract(metadata_json, '$.usage.cacheCreationTokens') as integer), cast(json_extract(metadata_json, '$.cacheCreationTokens') as integer), cast(json_extract(metadata_json, '$.usage.cache_creation_tokens') as integer), cast(json_extract(metadata_json, '$.cache_creation_tokens') as integer), 0) as cache_creation_tokens, - coalesce(cast(json_extract(metadata_json, '$.usage.cachedTokens') as integer), cast(json_extract(metadata_json, '$.cachedTokens') as integer), 0) as cached_tokens, - coalesce( - nullif(trim(json_extract(metadata_json, '$.upstreamModel')), ''), - nullif(trim(json_extract(metadata_json, '$.requestedModel')), ''), - nullif(trim(usage_records.model), ''), - '未知模型' - ) as billing_model, ", - ); - builder.push(USAGE_RECORD_SERVICE_TIER_SQL); - builder.push( - " as service_tier, - latency_ms - from usage_records", - ); - push_filter(&mut builder, filter, None)?; - builder.push(" order by date asc, created_at asc"); - - let rows = builder.build().fetch_all(pool).await?; - let mut days = BTreeMap::::new(); - - for row in rows { - let date: String = row.get("date"); - let input_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("input_tokens")); - let output_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("output_tokens")); - let cache_creation_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("cache_creation_tokens")); - let cached_tokens = - optional_nonnegative_i64_to_u64(row.get::, _>("cached_tokens")); - let cost = usage_breakdown_cost( - input_tokens, - output_tokens, - cached_tokens, - &row.get::("billing_model"), - row.get::, _>("service_tier").as_deref(), - ); - - days.entry(date.clone()) - .or_insert_with(|| UsageTrendAccumulator::new(date)) - .push( - input_tokens, - output_tokens, - cache_creation_tokens, - cached_tokens, - cost, - row.get("latency_ms"), - ); - } - - let mut points = days - .into_values() - .map(UsageTrendAccumulator::into_point) - .collect::>(); - if points.len() > 60 { - points = points.split_off(points.len() - 60); - } - Ok(points) -} - -#[derive(Debug)] -struct UsageTrendAccumulator { - point: UsageRecordTrendPoint, - latency_sum: f64, - latency_count: u64, -} - -impl UsageTrendAccumulator { - fn new(date: String) -> Self { - Self { - point: UsageRecordTrendPoint { - date, - ..UsageRecordTrendPoint::default() - }, - latency_sum: 0.0, - latency_count: 0, - } - } - - fn push( - &mut self, - input_tokens: u64, - output_tokens: u64, - cache_creation_tokens: u64, - cached_tokens: u64, - cost: f64, - latency_ms: Option, - ) { - self.point.input_tokens += input_tokens; - self.point.output_tokens += output_tokens; - self.point.cache_creation_tokens += cache_creation_tokens; - self.point.cached_tokens += cached_tokens; - self.point.total_tokens += input_tokens + output_tokens; - self.point.cost += cost; - self.point.actual_cost += cost; - - if let Some(latency_ms) = latency_ms.filter(|latency_ms| *latency_ms > 0) { - self.latency_sum += latency_ms as f64; - self.latency_count += 1; - } - } - - fn into_point(mut self) -> UsageRecordTrendPoint { - self.point.average_latency_ms = if self.latency_count > 0 { - Some(self.latency_sum / self.latency_count as f64) - } else { - None - }; - self.point - } -} - -fn usage_breakdown_cost( - input_tokens: u64, - output_tokens: u64, - cached_tokens: u64, - model: &str, - service_tier: Option<&str>, -) -> f64 { - billing::calculate_cost( - input_tokens, - output_tokens, - cached_tokens, - model, - service_tier, - ) -} - -fn weighted_average_latency( - current_average: Option, - current_count: f64, - next_average: Option, - next_count: f64, -) -> Option { - match (current_average, next_average) { - (Some(current), Some(next)) if current_count > 0.0 && next_count > 0.0 => { - Some((current * current_count + next * next_count) / (current_count + next_count)) - } - (Some(current), _) => Some(current), - (_, Some(next)) => Some(next), - _ => None, - } -} - -fn push_filter( - builder: &mut QueryBuilder, - filter: &UsageRecordFilter, - cursor: Option<&str>, -) -> SqliteUsageRecordStoreResult<()> { - let mut separated = builder.separated(" and "); - separated.push(" where 1 = 1"); - - if let Some(kind) = filter.kind.as_deref() { - separated.push("kind = "); - separated.push_bind_unseparated(kind); - } - if let Some(level) = filter.level { - separated.push("level = "); - separated.push_bind_unseparated(level_to_db(level)); - } - if let Some(request_id) = filter.request_id.as_deref() { - separated.push("request_id = "); - separated.push_bind_unseparated(request_id); - } - if let Some(account_id) = filter.account_id.as_deref() { - separated.push("account_id = "); - separated.push_bind_unseparated(account_id); - } - if let Some(route) = filter.route.as_deref() { - separated.push("route = "); - separated.push_bind_unseparated(route); - } - if let Some(model) = filter.model.as_deref() { - separated.push("model = "); - separated.push_bind_unseparated(model); - } - if let Some(status_code) = filter.status_code { - separated.push("status_code = "); - separated.push_bind_unseparated(status_code); - } - if let Some(transport) = filter.transport.as_deref() { - separated.push("transport = "); - separated.push_bind_unseparated(transport); - } - if let Some(attempt_index) = filter.attempt_index { - separated.push("attempt_index = "); - separated.push_bind_unseparated(attempt_index); - } - if let Some(upstream_status_code) = filter.upstream_status_code { - separated.push("upstream_status_code = "); - separated.push_bind_unseparated(upstream_status_code); - } - if let Some(failure_class) = filter.failure_class.as_deref() { - separated.push("failure_class = "); - separated.push_bind_unseparated(failure_class); - } - if let Some(response_id) = filter.response_id.as_deref() { - separated.push("response_id = "); - separated.push_bind_unseparated(response_id); - } - if let Some(upstream_request_id) = filter.upstream_request_id.as_deref() { - separated.push("upstream_request_id = "); - separated.push_bind_unseparated(upstream_request_id); - } - if let Some(search) = filter.search.as_deref() { - let pattern = format!("%{search}%"); - separated.push("(message like "); - separated.push_bind_unseparated(pattern.clone()); - separated.push_unseparated(" or metadata_json like "); - separated.push_bind_unseparated(pattern); - separated.push_unseparated(")"); - } - if let Some(start_time) = filter.start_time { - separated.push("created_at >= "); - separated.push_bind_unseparated(start_time.to_rfc3339()); - } - if let Some(end_time) = filter.end_time { - separated.push("created_at <= "); - separated.push_bind_unseparated(end_time.to_rfc3339()); - } - if let Some(cursor) = cursor { - let (created_at, id) = - decode_cursor(cursor).ok_or(SqliteUsageRecordStoreError::InvalidCursor)?; - separated.push("(created_at < "); - separated.push_bind_unseparated(created_at.clone()); - separated.push_unseparated(" or (created_at = "); - separated.push_bind_unseparated(created_at); - separated.push_unseparated(" and id < "); - separated.push_bind_unseparated(id); - separated.push_unseparated("))"); - } - - Ok(()) -} - -fn usage_record_from_row( - row: &sqlx::sqlite::SqliteRow, -) -> SqliteUsageRecordStoreResult { - Ok(UsageRecord { - id: row.get("id"), - request_id: row.get("request_id"), - kind: row.get("kind"), - level: level_from_db(&row.get::("level"))?, - account_id: row.get("account_id"), - route: row.get("route"), - model: row.get("model"), - status_code: row.get("status_code"), - transport: row.get("transport"), - attempt_index: row.get("attempt_index"), - upstream_status_code: row.get("upstream_status_code"), - failure_class: row.get("failure_class"), - response_id: row.get("response_id"), - upstream_request_id: row.get("upstream_request_id"), - latency_ms: row.get("latency_ms"), - message: row.get("message"), - metadata: serde_json::from_str(&row.get::("metadata_json"))?, - created_at: parse_rfc3339(&row.get::("created_at"))?, - }) -} - -fn level_to_db(level: UsageRecordLevel) -> &'static str { - match level { - UsageRecordLevel::Debug => "debug", - UsageRecordLevel::Info => "info", - UsageRecordLevel::Warn => "warn", - UsageRecordLevel::Error => "error", - } -} - -fn level_from_db(value: &str) -> SqliteUsageRecordStoreResult { - match value { - "debug" => Ok(UsageRecordLevel::Debug), - "info" => Ok(UsageRecordLevel::Info), - "warn" => Ok(UsageRecordLevel::Warn), - "error" => Ok(UsageRecordLevel::Error), - other => Err(SqliteUsageRecordStoreError::InvalidLevel(other.to_string())), - } -} - -/// 使用记录汇总。 -#[derive(Debug, Clone, Copy, Default, PartialEq)] -pub struct UsageRecordSummary { - /// 请求总数。 - pub total_requests: u64, - /// 错误请求数。 - pub error_requests: u64, - /// 输入 Token。 - pub input_tokens: u64, - /// 输出 Token。 - pub output_tokens: u64, - /// 缓存命中 Token。 - pub cached_tokens: u64, - /// 总 Token。 - pub total_tokens: u64, - /// 平均耗时。 - pub average_latency_ms: Option, -} - -/// 账号使用记录聚合。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct UsageRecordAccountUsage { - /// 账号 ID。 - pub account_id: String, - /// 输入 Token。 - pub input_tokens: u64, - /// 输出 Token。 - pub output_tokens: u64, - /// 总 Token。 - pub total_tokens: u64, - /// 最近使用时间。 - pub last_used_at: DateTime, -} - -/// 模型分布来源。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum UsageRecordModelSource { - /// 客户端请求模型。 - Requested, - /// 实际上游模型。 - Upstream, - /// 请求模型到上游模型的映射。 - Mapping, -} - -/// 端点分布来源。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum UsageRecordEndpointSource { - /// 入站端点。 - Inbound, - /// 上游端点。 - Upstream, - /// 入站到上游路径映射。 - Path, -} - -/// 使用记录分布项。 -#[derive(Debug, Clone, Default, PartialEq)] -pub struct UsageRecordBreakdown { - /// 名称。 - pub name: String, - /// 请求数。 - pub request_count: u64, - /// 输入 Token。 - pub input_tokens: u64, - /// 输出 Token。 - pub output_tokens: u64, - /// 缓存命中 Token。 - pub cached_tokens: u64, - /// 总 Token。 - pub total_tokens: u64, - /// 标准成本。 - pub cost: f64, - /// 实际成本。 - pub actual_cost: f64, - /// 账号成本。 - pub account_cost: f64, - /// 平均耗时。 - pub average_latency_ms: Option, -} - -/// 使用记录趋势点。 -#[derive(Debug, Clone, Default, PartialEq)] -pub struct UsageRecordTrendPoint { - /// 日期。 - pub date: String, - /// 输入 Token。 - pub input_tokens: u64, - /// 输出 Token。 - pub output_tokens: u64, - /// 缓存创建 Token。 - pub cache_creation_tokens: u64, - /// 缓存命中 Token。 - pub cached_tokens: u64, - /// 总 Token。 - pub total_tokens: u64, - /// 标准成本。 - pub cost: f64, - /// 实际成本。 - pub actual_cost: f64, - /// 平均耗时。 - pub average_latency_ms: Option, -} diff --git a/backend/src/admin/settings/mod.rs b/backend/src/admin/settings/mod.rs deleted file mode 100644 index a42da318f..000000000 --- a/backend/src/admin/settings/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -//! 管理端设置 HTTP 模块。 - -pub mod routes; diff --git a/backend/src/admin/system/mod.rs b/backend/src/admin/system/mod.rs deleted file mode 100644 index 23005fb71..000000000 --- a/backend/src/admin/system/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -//! 系统版本与在线更新处理器。 - -pub mod routes; -mod state; -mod updater; diff --git a/backend/src/admin/system/routes.rs b/backend/src/admin/system/routes.rs deleted file mode 100644 index 0ea259788..000000000 --- a/backend/src/admin/system/routes.rs +++ /dev/null @@ -1,5 +0,0 @@ -//! 管理端系统路由入口。 - -pub(crate) use super::updater::{ - perform_update, restart, rollback, update_detail, update_event_stream, update_status, version, -}; diff --git a/backend/src/api/admin/accounts_routes.rs b/backend/src/api/admin/accounts_routes.rs new file mode 100644 index 000000000..d6ee60b85 --- /dev/null +++ b/backend/src/api/admin/accounts_routes.rs @@ -0,0 +1,606 @@ +//! 账号管理 HTTP 处理器。 + +use axum::{ + body::Body, + extract::{Query, State}, + http::{ + header::{CACHE_CONTROL, CONNECTION, CONTENT_TYPE}, + StatusCode, + }, + response::{IntoResponse, Response}, + Json, +}; +use std::collections::{HashMap, HashSet}; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::{ + accounts::{ + account::AccountStatus, + manage::{ + quota_view::{ + quota_data, AccountQuotaData, AccountQuotaUsageWindow, AccountQuotaWindowLocalUsage, + }, + AccountHealthCheck, AccountManageError, AccountRefreshOutcome, AccountRefreshResult, + AccountUpdate, ManagedAccount, OAuthExchangeInput, + }, + refresh::token_refresh_status_eligible, + }, + api::admin::{ + response::{ + AdminEnvelope, AdminError, AdminResponse, BatchDeleteData, CursorPageMeta, + NumberedPageMeta, PageMeta, ADMIN_OK_CODE, ADMIN_OK_MESSAGE, + }, + session::AdminAuth, + update_payload::{parse_editable_update, EditableUpdateMessages}, + }, + bootstrap::state::AppState, + infra::{ + format::{ + format_cost, format_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64, + }, + json::{clamp_limit, clamp_page, total_pages, Page}, + time::{china_datetime, china_relative_time, china_rfc3339}, + }, + telemetry::{ + account_usage::query::AccountUsageRecord, + billing, + buckets::query::{ModelUsageWindow, UsageBucketWindow}, + }, +}; + +const ACCOUNT_STATS_PAGE_LIMIT: u32 = 200; +const ACCOUNT_EXPORT_CONFIRMATION: &str = "export_sensitive_accounts"; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AccountsQuery { + cursor: Option, + limit: Option, + page: Option, + page_size: Option, + search: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct CreateAccountRequest { + token: Option, + refresh_token: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct BatchDeleteAccountsRequest { + ids: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AccountActionRequest { + id: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AccountIdQuery { + id: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AccountExportQuery { + ids: Option, + confirm: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AccountTestQuery { + id: String, + model_id: Option, +} + +#[derive(Debug, Default, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct AccountOAuthExchangeRequest { + session_id: String, + callback_url: Option, + code: Option, + state: Option, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountModelsData { + models: Vec, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountModelData { + id: String, + label: String, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminAccountData { + id: String, + email: Option, + account_id: Option, + user_id: Option, + label: Option, + plan_type: Option, + has_refresh_token: bool, + status: String, + display_status: String, + token_refreshing: bool, + access_token_expires_at: Option, + access_token_expires_at_display: Option, + added_at: String, + added_at_display: String, + updated_at: String, + updated_at_display: String, + quota: AccountQuotaData, + usage: AdminAccountUsageData, +} + +impl From for AdminAccountData { + fn from(a: ManagedAccount) -> Self { + Self::from_parts(a, None, None, Vec::new(), false) + } +} + +impl AdminAccountData { + fn from_parts( + a: ManagedAccount, + usage: Option<&AccountUsageRecord>, + quota: Option, + models: Vec, + token_refreshing: bool, + ) -> Self { + let access_token_expires_at = a.access_token_expires_at.as_ref().map(china_rfc3339); + let access_token_expires_at_display = + a.access_token_expires_at.as_ref().map(china_datetime); + let display_status = account_display_status(a.status, token_refreshing).to_string(); + Self { + id: a.id, + email: a.email, + account_id: a.account_id, + user_id: a.user_id, + label: a.label, + plan_type: a.plan_type, + has_refresh_token: a.has_refresh_token, + status: a.status.as_str().to_string(), + display_status, + token_refreshing, + access_token_expires_at, + access_token_expires_at_display, + added_at: china_rfc3339(&a.added_at), + added_at_display: china_datetime(&a.added_at), + updated_at: china_rfc3339(&a.updated_at), + updated_at_display: china_datetime(&a.updated_at), + quota: quota.unwrap_or_default(), + usage: AdminAccountUsageData::from_usage(usage, models), + } + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminAccountUsageData { + request_count: i64, + request_count_display: String, + empty_response_count: i64, + input_tokens: i64, + input_tokens_display: String, + output_tokens: i64, + output_tokens_display: String, + cached_tokens: i64, + cached_tokens_display: String, + reasoning_tokens: i64, + total_tokens: i64, + total_tokens_display: String, + image_input_tokens: i64, + image_output_tokens: i64, + image_tokens_display: String, + image_request_count: i64, + image_request_failed_count: i64, + created_tokens: i64, + created_tokens_display: String, + read_tokens: i64, + read_tokens_display: String, + last_used_at: Option, + last_used_at_display: String, + models: Vec, +} + +impl AdminAccountUsageData { + fn from_usage( + usage: Option<&AccountUsageRecord>, + models: Vec, + ) -> Self { + let request_count = usage.map_or(0, |usage| usage.window_request_count); + let empty_response_count = 0; + let input_tokens = usage.map_or(0, |usage| usage.window_input_tokens); + let output_tokens = usage.map_or(0, |usage| usage.window_output_tokens); + let cached_tokens = usage.map_or(0, |usage| usage.window_cached_tokens); + let reasoning_tokens = 0; + let total_tokens = input_tokens.saturating_add(output_tokens); + let image_input_tokens = 0; + let image_output_tokens = 0; + let image_request_count = 0; + let image_request_failed_count = 0; + let created_tokens = input_tokens.saturating_sub(cached_tokens); + let read_tokens = cached_tokens; + let last_used_at = usage.and_then(|usage| usage.last_used_at); + + Self { + request_count, + request_count_display: format_plain_number(nonnegative_i64_to_u64(request_count)), + empty_response_count, + input_tokens, + input_tokens_display: format_tokens(nonnegative_i64_to_u64(input_tokens)), + output_tokens, + output_tokens_display: format_tokens(nonnegative_i64_to_u64(output_tokens)), + cached_tokens, + cached_tokens_display: format_tokens(nonnegative_i64_to_u64(cached_tokens)), + reasoning_tokens, + total_tokens, + total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)), + image_input_tokens, + image_output_tokens, + image_tokens_display: format_tokens(nonnegative_i64_to_u64( + image_input_tokens + image_output_tokens, + )), + image_request_count, + image_request_failed_count, + created_tokens, + created_tokens_display: format_tokens(nonnegative_i64_to_u64(created_tokens)), + read_tokens, + read_tokens_display: format_tokens(nonnegative_i64_to_u64(read_tokens)), + last_used_at: last_used_at.map(|value| china_rfc3339(&value)), + last_used_at_display: china_relative_time(last_used_at, Utc::now()), + models, + } + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminAccountModelUsageData { + model: String, + request_count: u64, + request_count_display: String, + success_rate: f64, + success_rate_display: String, + input_tokens: u64, + input_tokens_display: String, + output_tokens: u64, + output_tokens_display: String, + cached_tokens: u64, + cached_tokens_display: String, + total_tokens: u64, + total_tokens_display: String, + total_cost_usd: f64, + total_cost_usd_display: String, + last_used_at: Option, + last_used_at_display: String, +} + +#[derive(Debug, Clone, Default)] +struct AccountListStats { + usage_by_account: HashMap, + quota_by_account: HashMap, + models_by_account: HashMap>, + refreshing_account_ids: HashSet, +} + +struct AccountModelUsageRecord { + account_id: String, + model: String, + request_count: i64, + error_count: i64, + input_tokens: i64, + output_tokens: i64, + cached_tokens: i64, + total_cost_usd: f64, + last_used_at: Option>, +} + +#[derive(Debug, Clone)] +struct AccountQuotaWindowSelection { + account_id: String, + window: AccountQuotaUsageWindow, +} + +impl AccountListStats { + fn data_for(&self, account: ManagedAccount) -> AdminAccountData { + let account_id = account.id.clone(); + let token_refreshing = token_refresh_status_eligible(account.status) + && self.refreshing_account_ids.contains(&account_id); + AdminAccountData::from_parts( + account, + self.usage_by_account.get(&account_id), + self.quota_by_account.get(&account_id).cloned(), + self.models_by_account + .get(&account_id) + .cloned() + .unwrap_or_default(), + token_refreshing, + ) + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminAccountSummaryData { + total: u64, + active: u64, + high_usage: u64, + attention: u64, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminAccountPageData { + items: Vec, + page: PageMeta, + summary: AdminAccountSummaryData, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminAccountPageEnvelope { + code: u32, + message: String, + data: AdminAccountPageData, +} + +impl AdminAccountPageEnvelope { + fn cursor(page: Page, limit: u32, summary: AdminAccountSummaryData) -> Self { + Self { + code: ADMIN_OK_CODE, + message: ADMIN_OK_MESSAGE.into(), + data: AdminAccountPageData { + items: page.items, + page: PageMeta::Cursor(CursorPageMeta { + limit, + next_cursor: page.next_cursor, + }), + summary, + }, + } + } + + fn numbered( + page: crate::infra::json::NumberedPage, + summary: AdminAccountSummaryData, + ) -> Self { + Self { + code: ADMIN_OK_CODE, + message: ADMIN_OK_MESSAGE.into(), + data: AdminAccountPageData { + page: PageMeta::Numbered(NumberedPageMeta { + page: page.page, + page_size: page.page_size, + total: page.total, + total_pages: total_pages(page.total, page.page_size), + }), + items: page.items, + summary, + }, + } + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(untagged)] +enum AccountUpdateData { + Account(Box), +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountRefreshData { + id: String, + email: Option, + previous_status: String, + result: &'static str, + error: Option, + duration_ms: i64, +} + +impl From for AccountRefreshData { + fn from(result: AccountRefreshResult) -> Self { + Self { + id: result.id, + email: result.email, + previous_status: result.previous_status.as_str().to_string(), + result: account_refresh_outcome_str(result.outcome), + error: result.error, + duration_ms: result.duration_ms, + } + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AdminAccountQuotaResponseData { + quota: Value, + raw: Value, + quota_data: AccountQuotaData, + plan_type: Option, + account: AdminAccountData, +} + +impl AdminAccountQuotaResponseData { + fn from_account( + quota: Value, + raw: Value, + quota_data: AccountQuotaData, + account: AdminAccountData, + ) -> Self { + let plan_type = account.plan_type.clone(); + Self { + quota, + raw, + quota_data, + plan_type, + account, + } + } +} + +#[derive(Debug, Clone, Deserialize, Default)] +#[serde(default)] +pub(crate) struct AccountHealthCheckRequest { + ids: Option>, + stagger_ms: Option, + concurrency: Option, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountHealthCheckSummaryData { + total: usize, + alive: usize, + dead: usize, + skipped: usize, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountHealthCheckData { + summary: AccountHealthCheckSummaryData, + results: Vec, +} + +impl From for AccountHealthCheckData { + fn from(result: AccountHealthCheck) -> Self { + let summary = AccountHealthCheckSummaryData { + total: result.results.len(), + alive: result.alive(), + dead: result.dead(), + skipped: result.skipped(), + }; + Self { + summary, + results: result + .results + .into_iter() + .map(AccountRefreshData::from) + .collect(), + } + } +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountImportData { + imported: u32, + skipped: u32, + source_format: String, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +struct AccountOAuthAuthorizeData { + session_id: String, + auth_url: String, + expires_at: String, + expires_at_display: String, +} + +/// `GET /api/admin/accounts` +pub(crate) async fn accounts( + State(state): State, + _auth: AdminAuth, + Query(params): Query, +) -> Result { + let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50)); + let use_numbered_page = params.page.is_some() || params.page_size.is_some(); + let quota_by_account = quota_snapshots_by_account(&state).await; + let summary = account_summary_data(&state, "a_by_account).await; + + if use_numbered_page { + return match state + .services + .admin_accounts + .list_page(clamp_page(params.page.unwrap_or(1)), limit, params.search) + .await + { + Ok(page) => { + let stats = account_list_stats(&state, &page.items, "a_by_account).await?; + let page = crate::infra::json::NumberedPage { + items: page + .items + .into_iter() + .map(|item| stats.data_for(item)) + .collect(), + total: page.total, + page: page.page, + page_size: page.page_size, + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminAccountPageEnvelope::numbered(page, summary), + )) + } + Err(error) => Err(account_error(&error)), + }; + } + + match state + .services + .admin_accounts + .list(params.cursor, limit) + .await + { + Ok(page) => { + let stats = account_list_stats(&state, &page.items, "a_by_account).await?; + let page = Page { + items: page + .items + .into_iter() + .map(|item| stats.data_for(item)) + .collect(), + next_cursor: page.next_cursor, + }; + Ok(AdminResponse::new( + StatusCode::OK, + AdminAccountPageEnvelope::cursor(page, limit, summary), + )) + } + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts` +pub(crate) async fn create_account( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + match state + .services + .admin_accounts + .create(payload.token, payload.refresh_token) + .await + { + Ok(account) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminAccountData::from(account)), + )), + Err(error) => Err(account_error(&error)), + } +} + +mod actions; +mod stats; + +pub(crate) use actions::*; +use stats::*; diff --git a/backend/src/api/admin/accounts_routes/actions.rs b/backend/src/api/admin/accounts_routes/actions.rs new file mode 100644 index 000000000..4dfd8114e --- /dev/null +++ b/backend/src/api/admin/accounts_routes/actions.rs @@ -0,0 +1,312 @@ +use super::*; + +/// `GET /api/admin/accounts/export` +pub(crate) async fn export_accounts( + State(state): State, + _auth: AdminAuth, + Query(query): Query, +) -> Result { + let ids = account_export_ids(query.ids.as_deref()); + + if query.confirm.as_deref() != Some(ACCOUNT_EXPORT_CONFIRMATION) { + return Err(AdminError::bad_request( + "account export requires confirm=export_sensitive_accounts", + )); + } + + match state.services.admin_accounts.export(ids).await { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(result), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/refresh` +pub(crate) async fn refresh_account( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + let account_id = payload.id; + match state + .services + .admin_accounts + .refresh_account(&account_id) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountRefreshData::from(result)), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/health-check` +pub(crate) async fn health_check_accounts( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + let stagger_ms = payload.stagger_ms.unwrap_or(3_000); + if !(500..=30_000).contains(&stagger_ms) && stagger_ms != 0 { + return Err(AdminError::bad_request( + "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", + )); + } + } + match state + .services + .admin_accounts + .health_check_accounts(payload.ids, stagger_ms, payload.concurrency.unwrap_or(2)) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountHealthCheckData::from(result)), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `GET /api/admin/accounts/quota` +pub(crate) async fn account_quota( + State(state): State, + _auth: AdminAuth, + Query(query): Query, +) -> Result { + let account_id = query.id; + match state + .services + .admin_accounts + .account_quota(&account_id) + .await + { + Ok(data) => { + let quota = data.get("quota").cloned().unwrap_or(Value::Null); + let raw = data.get("raw").cloned().unwrap_or(Value::Null); + let quota_json = quota.to_string(); + let mut quota_data = quota_data("a_json, Some(Utc::now())); + apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await; + let account = + account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account( + quota, raw, quota_data, account, + )), + )) + } + Err(AccountManageError::NotFound) => Err(account_not_found()), + Err(AccountManageError::Inactive(status)) => Err(AdminError::conflict(format!( + "Account is {}, cannot query quota", + status.as_str() + ))), + Err(AccountManageError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!( + "Failed to fetch quota from Codex API: {msg}" + ))), + Err(e) => Err(account_error(&e)), + } +} + +async fn account_data_for_quota_refresh( + state: &AppState, + account_id: &str, + quota_data: AccountQuotaData, +) -> Result { + let Some(account) = state + .services + .admin_accounts + .get(account_id) + .await + .map_err(|error| account_error(&error))? + else { + return Err(account_not_found()); + }; + let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]); + let stats = + account_list_stats(state, std::slice::from_ref(&account), "a_by_account).await?; + Ok(stats.data_for(account)) +} + +/// `POST /api/admin/accounts/import` +pub(crate) async fn import_accounts( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + match state.services.admin_accounts.import(payload).await { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountImportData { + imported: result.imported, + skipped: result.skipped, + source_format: result.source_format.to_string(), + }), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/oauth/authorize` +pub(crate) async fn oauth_authorize_account( + State(state): State, + _auth: AdminAuth, +) -> Result { + match state.services.admin_accounts.oauth_authorize().await { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountOAuthAuthorizeData { + session_id: result.session_id, + auth_url: result.auth_url, + expires_at: china_rfc3339(&result.expires_at), + expires_at_display: china_datetime(&result.expires_at), + }), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/oauth/exchange` +pub(crate) async fn oauth_exchange_account( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + match state + .services + .admin_accounts + .oauth_exchange(OAuthExchangeInput { + session_id: payload.session_id, + callback_url: payload.callback_url, + code: payload.code, + state: payload.state, + }) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountImportData { + imported: result.imported, + skipped: result.skipped, + source_format: result.source_format.to_string(), + }), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `GET /api/admin/accounts/test?id=...&modelId=...` +pub(crate) async fn test_account_connection( + State(state): State, + Query(query): Query, + _auth: AdminAuth, +) -> Result { + let account_id = query.id; + + let model = query + .model_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) + .ok_or_else(|| AdminError::bad_request("Model is required"))?; + let stream = state + .services + .admin_accounts + .test_connection_stream(&account_id, model) + .await + .map_err(|error| account_error(&error))?; + + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(CACHE_CONTROL, "no-cache") + .header(CONNECTION, "keep-alive") + .header("x-accel-buffering", "no") + .body(Body::from_stream(stream)) + .map_err(|_| AdminError::internal("Failed to build account test stream")) +} + +/// `GET /api/admin/accounts/models?id=...` +pub(crate) async fn account_models( + State(state): State, + Query(query): Query, + _auth: AdminAuth, +) -> Result { + let account_id = query.id; + match state + .services + .admin_accounts + .account_models(&account_id) + .await + { + Ok(models) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountModelsData { + models: models + .into_iter() + .map(|model| AccountModelData { + id: model.id, + label: model.label, + }) + .collect(), + }), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/delete` +pub(crate) async fn batch_delete_accounts( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + match state + .services + .admin_accounts + .batch_delete(payload.ids) + .await + { + Ok(result) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(BatchDeleteData { + deleted: result.deleted, + not_found: result.not_found, + }), + )), + Err(error) => Err(account_error(&error)), + } +} + +/// `POST /api/admin/accounts/update` +pub(crate) async fn update_account( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + let ParsedAccountUpdate { id, update } = parse_account_update(&payload)?; + + match state + .services + .admin_accounts + .update_account(&id, update) + .await + { + Ok(Some(account)) => Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountUpdateData::Account(Box::new(account.into()))), + )), + Ok(None) => Err(account_not_found()), + Err(error) => Err(account_error(&error)), + } +} diff --git a/backend/src/api/admin/accounts_routes/stats.rs b/backend/src/api/admin/accounts_routes/stats.rs new file mode 100644 index 000000000..63a04f923 --- /dev/null +++ b/backend/src/api/admin/accounts_routes/stats.rs @@ -0,0 +1,512 @@ +use super::*; + +pub(super) async fn account_list_stats( + state: &AppState, + accounts: &[ManagedAccount], + quota_by_account: &HashMap, +) -> Result { + let account_ids = accounts + .iter() + .map(|account| account.id.clone()) + .collect::>(); + let usage_records = state + .services + .usage + .list_by_account_ids(&account_ids) + .await + .unwrap_or_default(); + let mut quota_by_account = quota_by_account.clone(); + let all_quota_windows = quota_usage_windows_by_account(accounts, "a_by_account); + let quota_window_usage_stats = + quota_window_local_usage_by_account(state, &all_quota_windows).await; + apply_quota_window_local_usage(accounts, &mut quota_by_account, "a_window_usage_stats); + let selected_quota_windows = selected_quota_windows_by_account(accounts, "a_by_account); + let usage_records = usage_records + .into_iter() + .map(|usage| { + apply_selected_quota_window_usage( + usage, + &selected_quota_windows, + "a_window_usage_stats, + ) + }) + .collect::>(); + let models_by_account = list_current_window_model_usage(state, &usage_records).await; + let refreshing_account_ids = state + .services + .token_refresh + .refreshing_account_ids(&account_ids, Utc::now()) + .await + .map_err(|error| AdminError::internal(error.to_string()))?; + + Ok(AccountListStats { + usage_by_account: usage_records + .into_iter() + .map(|usage| (usage.account_id.clone(), usage)) + .collect(), + quota_by_account, + models_by_account, + refreshing_account_ids, + }) +} + +pub(super) fn account_display_status( + status: AccountStatus, + token_refreshing: bool, +) -> &'static str { + if token_refreshing { + "refreshing" + } else { + status.as_str() + } +} + +pub(super) fn selected_quota_windows_by_account( + accounts: &[ManagedAccount], + quota_by_account: &HashMap, +) -> HashMap { + accounts + .iter() + .filter_map(|account| { + quota_by_account + .get(&account.id) + .and_then(selected_quota_window) + .map(|window| (account.id.clone(), window)) + }) + .collect() +} + +pub(super) fn selected_quota_window(quota: &AccountQuotaData) -> Option { + quota.usage_windows().into_iter().max_by(|left, right| { + left.duration_seconds() + .cmp(&right.duration_seconds()) + .then_with(|| left.end.cmp(&right.end)) + .then_with(|| left.key.cmp(&right.key)) + }) +} + +pub(super) fn quota_usage_windows_by_account( + accounts: &[ManagedAccount], + quota_by_account: &HashMap, +) -> Vec { + accounts + .iter() + .filter_map(|account| { + quota_by_account + .get(&account.id) + .map(|quota| (account.id.as_str(), quota)) + }) + .flat_map(|(account_id, quota)| { + quota + .usage_windows() + .into_iter() + .map(move |window| AccountQuotaWindowSelection { + account_id: account_id.to_string(), + window, + }) + }) + .collect() +} + +pub(super) async fn apply_account_quota_window_local_usage( + state: &AppState, + account_id: &str, + quota: &mut AccountQuotaData, +) { + let windows = quota + .usage_windows() + .into_iter() + .map(|window| AccountQuotaWindowSelection { + account_id: account_id.to_string(), + window, + }) + .collect::>(); + let usage_by_account = quota_window_local_usage_by_account(state, &windows).await; + let empty_usage = HashMap::new(); + quota.apply_local_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage)); +} + +pub(super) fn apply_quota_window_local_usage( + accounts: &[ManagedAccount], + quota_by_account: &mut HashMap, + usage_by_account: &HashMap>, +) { + for account in accounts { + if let Some(quota) = quota_by_account.get_mut(&account.id) { + let empty_usage = HashMap::new(); + quota.apply_local_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage)); + } + } +} + +pub(super) async fn quota_window_local_usage_by_account( + state: &AppState, + windows: &[AccountQuotaWindowSelection], +) -> HashMap> { + let queries = windows + .iter() + .map(|window| UsageBucketWindow { + account_id: window.account_id.clone(), + key: window.window.key.clone(), + start: window.window.start, + end: window.window.end, + }) + .collect::>(); + state + .services + .usage + .usage_by_windows(&queries) + .await + .unwrap_or_default() + .into_iter() + .map(|(account_id, usage_by_window)| { + let usage_by_window = usage_by_window + .into_iter() + .map(|(key, usage)| { + ( + key, + AccountQuotaWindowLocalUsage { + request_count: usage.request_count, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cached_tokens: usage.cached_tokens, + }, + ) + }) + .collect(); + (account_id, usage_by_window) + }) + .collect() +} + +pub(super) fn apply_selected_quota_window_usage( + mut usage: AccountUsageRecord, + selected_quota_windows: &HashMap, + quota_window_usage_stats: &HashMap>, +) -> AccountUsageRecord { + let Some(window) = selected_quota_windows.get(&usage.account_id) else { + return usage; + }; + let stats = quota_window_usage_stats + .get(&usage.account_id) + .and_then(|usage_by_window| usage_by_window.get(&window.key)) + .copied() + .unwrap_or_default(); + usage.window_request_count = stats.request_count; + usage.window_input_tokens = stats.input_tokens; + usage.window_output_tokens = stats.output_tokens; + usage.window_cached_tokens = stats.cached_tokens; + usage.window_started_at = Some(window.start); + usage.window_reset_at = Some(window.end); + usage.limit_window_seconds = Some(window.window_seconds); + usage +} + +pub(super) async fn account_summary_data( + state: &AppState, + quota_by_account: &HashMap, +) -> 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 + .iter() + .filter(|account| { + quota_by_account + .get(&account.id) + .is_some_and(account_quota_has_high_usage) + }) + .count() as u64; + let attention = accounts + .iter() + .filter(|account| account_summary_needs_attention(account.status)) + .count() as u64; + + AdminAccountSummaryData { + total, + active, + high_usage, + attention, + } +} + +pub(super) async fn quota_snapshots_by_account( + state: &AppState, +) -> HashMap { + state + .services + .background_tasks + .accounts + .list_quota_snapshots() + .await + .unwrap_or_default() + .into_iter() + .map(|snapshot| { + ( + snapshot.account_id, + quota_data(&snapshot.quota_json, snapshot.quota_fetched_at), + ) + }) + .collect() +} + +pub(super) async fn list_all_account_metadata(state: &AppState) -> Vec { + let mut page = 1; + let mut accounts = Vec::new(); + loop { + let Ok(result) = state + .services + .admin_accounts + .list_page(page, ACCOUNT_STATS_PAGE_LIMIT, None) + .await + else { + return Vec::new(); + }; + let total = result.total; + accounts.extend(result.items); + if accounts.len() as u64 >= total || total == 0 { + return accounts; + } + page = page.saturating_add(1); + } +} + +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, + AccountStatus::Expired | AccountStatus::Disabled | AccountStatus::Banned + ) +} + +pub(super) async fn list_current_window_model_usage( + state: &AppState, + usage_records: &[AccountUsageRecord], +) -> HashMap> { + let now = Utc::now(); + let windows = usage_records + .iter() + .filter_map(|usage| { + current_usage_window(usage, now) + .map(|(start, end)| (usage.account_id.clone(), start, end)) + }) + .collect::>(); + if windows.is_empty() { + return HashMap::new(); + } + + let windows = windows + .into_iter() + .map(|(account_id, start, end)| ModelUsageWindow { + account_id, + start, + end, + }) + .collect::>(); + let rows = state + .services + .usage + .model_usage_by_windows(&windows) + .await + .unwrap_or_default(); + let mut records_by_model = HashMap::<(String, String), AccountModelUsageRecord>::new(); + for row in rows { + let account_id = row.account_id; + let model = row.model; + let request_count = row.request_count; + let error_count = row.error_count; + 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( + nonnegative_i64_to_u64(input_tokens), + nonnegative_i64_to_u64(output_tokens), + nonnegative_i64_to_u64(cached_tokens), + &model, + row.service_tier.as_deref(), + ); + let last_used_at = row.last_used_at; + + let record = records_by_model + .entry((account_id.clone(), model.clone())) + .or_insert_with(|| AccountModelUsageRecord { + account_id, + model, + request_count: 0, + error_count: 0, + input_tokens: 0, + output_tokens: 0, + cached_tokens: 0, + total_cost_usd: 0.0, + last_used_at: None, + }); + record.request_count += request_count; + record.error_count += error_count; + record.input_tokens += input_tokens; + record.output_tokens += output_tokens; + record.cached_tokens += cached_tokens; + record.total_cost_usd += total_cost_usd; + record.last_used_at = record.last_used_at.max(last_used_at); + } + let records = records_by_model.into_values().collect::>(); + + models_by_account(records) +} + +pub(super) fn current_usage_window( + usage: &AccountUsageRecord, + now: DateTime, +) -> Option<(DateTime, DateTime)> { + let start = usage.window_started_at?; + let end = usage.window_reset_at.unwrap_or(now); + (start <= end).then_some((start, end)) +} + +pub(super) fn models_by_account( + records: Vec, +) -> HashMap> { + let mut by_account = HashMap::>::new(); + for record in records { + by_account + .entry(record.account_id.clone()) + .or_default() + .push(record); + } + + by_account + .into_iter() + .map(|(account_id, mut records)| { + records.sort_by(|a, b| { + b.request_count + .cmp(&a.request_count) + .then_with(|| b.last_used_at.cmp(&a.last_used_at)) + .then_with(|| a.model.cmp(&b.model)) + }); + ( + account_id, + records.into_iter().map(model_usage_data).collect(), + ) + }) + .collect() +} + +pub(super) fn model_usage_data(usage: AccountModelUsageRecord) -> AdminAccountModelUsageData { + let request_count = nonnegative_i64_to_u64(usage.request_count); + let error_count = nonnegative_i64_to_u64(usage.error_count); + let input_tokens = nonnegative_i64_to_u64(usage.input_tokens); + let output_tokens = nonnegative_i64_to_u64(usage.output_tokens); + let cached_tokens = nonnegative_i64_to_u64(usage.cached_tokens); + let total_tokens = input_tokens + output_tokens; + let success_rate = if request_count > 0 { + ((request_count.saturating_sub(error_count)) as f64 / request_count as f64 * 1000.0).round() + / 10.0 + } else { + 0.0 + }; + let total_cost_usd = usage.total_cost_usd; + + AdminAccountModelUsageData { + model: usage.model, + request_count, + request_count_display: format_plain_number(request_count), + success_rate, + success_rate_display: format_percent(success_rate), + input_tokens, + input_tokens_display: format_tokens(input_tokens), + output_tokens, + output_tokens_display: format_tokens(output_tokens), + cached_tokens, + cached_tokens_display: format_tokens(cached_tokens), + total_tokens, + total_tokens_display: format_tokens(total_tokens), + total_cost_usd, + total_cost_usd_display: format_cost(total_cost_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()), + } +} + +pub(super) struct ParsedAccountUpdate { + pub(super) id: String, + pub(super) update: AccountUpdate, +} + +pub(super) fn parse_account_update(payload: &Value) -> Result { + let payload = parse_editable_update( + payload, + EditableUpdateMessages { + object_required: "Account update request must be an object", + invalid: "Invalid account update request", + empty_update: "Account update request must include editable fields", + unknown_field_editable: true, + }, + )?; + let update = AccountUpdate { + label: payload.label.map(|label| { + label.and_then(|value| { + let value = value.trim(); + (!value.is_empty()).then(|| value.to_string()) + }) + }), + status: payload.status, + }; + if !update.any() { + return Err(AdminError::bad_request( + "Account update request must include editable fields", + )); + } + Ok(ParsedAccountUpdate { + id: payload.id, + update, + }) +} + +pub(super) fn account_error(error: &AccountManageError) -> AdminError { + match error { + AccountManageError::InvalidStatus(_) + | AccountManageError::LabelTooLong + | AccountManageError::EmptyIds + | AccountManageError::NoImportableAccounts + | AccountManageError::NoModels + | AccountManageError::InvalidAccessTokenExpiresAt + | AccountManageError::TokenRequired + | AccountManageError::InvalidToken(_) + | AccountManageError::RefreshTokenExchange(_) + | AccountManageError::OAuthSessionInvalid + | AccountManageError::OAuthCallbackInvalid + | AccountManageError::OAuthStateMismatch + | AccountManageError::NoValidCookies => AdminError::bad_request(error.to_string()), + AccountManageError::OAuthCodeExchange(_) => AdminError::bad_gateway(error.to_string()), + AccountManageError::NotFound => account_not_found(), + AccountManageError::Inactive(_) => AdminError::conflict(error.to_string()), + _ => AdminError::internal(error.to_string()), + } +} + +pub(super) fn account_refresh_outcome_str(outcome: AccountRefreshOutcome) -> &'static str { + match outcome { + AccountRefreshOutcome::Alive => "alive", + AccountRefreshOutcome::Dead => "dead", + AccountRefreshOutcome::Skipped => "skipped", + } +} + +pub(super) fn account_not_found() -> AdminError { + AdminError::not_found("Account not found") +} + +pub(super) fn account_export_ids(value: Option<&str>) -> Vec { + value + .into_iter() + .flat_map(|ids| ids.split(',')) + .map(str::trim) + .filter(|id| !id.is_empty()) + .map(ToString::to_string) + .collect() +} diff --git a/backend/src/api/admin/dashboard_routes.rs b/backend/src/api/admin/dashboard_routes.rs new file mode 100644 index 000000000..e6807fcf9 --- /dev/null +++ b/backend/src/api/admin/dashboard_routes.rs @@ -0,0 +1,388 @@ +//! 管理端 Dashboard 聚合视图。 + +use std::collections::{HashMap, HashSet}; + +use axum::{ + extract::{Query, State}, + http::StatusCode, + response::IntoResponse, +}; +use chrono::{DateTime, Duration, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::{ + accounts::{ + account::{Account, AccountStatus}, + refresh::token_refresh_status_eligible, + }, + api::admin::{ + response::{AdminEnvelope, AdminError, AdminResponse}, + session::AdminAuth, + usage_routes::{ + usage_cost_details, usage_record_models, usage_token_details, + UsageRecordCostDetailsData, UsageRecordTokenDetailsData, + }, + }, + bootstrap::state::AppState, + infra::{ + format::{ + format_compact_number, format_cost, format_duration_ms, format_percent, format_rate, + format_tokens, nonnegative_i64_to_u64, + }, + time::{ + china_datetime, china_datetime_rfc3339_str, china_day_start, china_hour, + china_hour_start, china_quarter_hour_start, china_relative_time, + }, + }, + telemetry::{ + account_usage::query::AccountUsageTimeBucket, + diagnostics::{ + fingerprint_diagnostics, AccountCapacityDiagnostics, AccountPoolDiagnostics, + }, + usage::{ + query::UsageQueryFilter, + store::{UsageRecordAccountUsage, UsageRecordSummary}, + types::{metadata_string, UsageRecord}, + }, + }, +}; + +const DASHBOARD_ACCOUNT_USAGE_LIMIT: u32 = 4; +const DASHBOARD_USAGE_RECORD_LIMIT: u32 = 10; +const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15; +const DASHBOARD_TIME_BUCKET_SLOTS: i64 = 7 * 24 * 4; +const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4; +const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3; +const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000; +const WEEK_WINDOW_SECONDS: u64 = 604_800; +const MONTH_WINDOW_SECONDS: u64 = 2_592_000; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct DashboardTrendQuery { + kind: Option, +} + +#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +enum DashboardTrendKind { + #[default] + Usage, + Latency, + Errors, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardSummaryData { + cards: DashboardCardsData, + trend: DashboardTrendData, + health_timeline: DashboardHealthTimelineData, + account_usage: Vec, + service_statuses: Vec, + usage_records: Vec, + pool_summary: AccountPoolDiagnostics, + capacity_info: AccountCapacityDiagnostics, + rotation_strategy: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardHealthTimelineData { + title: String, + description: String, + reliability_display: String, + points: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardCardsData { + accounts: DashboardAccountsCardData, + traffic: DashboardTrafficCardData, + tokens: DashboardTokenCardData, + cache: DashboardCacheCardData, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardAccountsCardData { + total: String, + total_value: u64, + enabled: String, + enabled_value: u64, + abnormal: String, + abnormal_value: u64, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardTrafficCardData { + today_requests: String, + today_requests_value: u64, + yesterday_requests_value: u64, + total_requests: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardTokenCardData { + today_tokens: String, + today_tokens_value: u64, + yesterday_tokens_value: u64, + total_tokens: String, + today_cost_usd: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardCacheCardData { + today_hit_rate: String, + today_hit_rate_value: Option, + yesterday_hit_rate_value: Option, + total_hit_rate: String, + total_cached_tokens: String, + average_first_token_latency_ms: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardTrendData { + kind: DashboardTrendKind, + points: Vec, + summary: Vec, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardTrendPointData { + time: String, + requests: String, + requests_value: u64, + input_tokens: String, + input_tokens_value: u64, + output_tokens: String, + output_tokens_value: u64, + cached_tokens: String, + cached_tokens_value: u64, + cache_hit_rate_value: f64, + tokens_value: u64, + errors: String, + errors_value: u64, + latency: String, + latency_value: u64, + max_latency: String, + max_latency_value: u64, + min_latency: String, + min_latency_value: u64, + success_rate: String, + success_rate_value: f64, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardTrendSummaryData { + label: String, + value: String, + ratio: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardAccountUsageData { + id: String, + email: String, + plan_type: Option, + tokens: String, + quota_used_percent: Option, + last_used: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardUsageRecordData { + id: String, + route: Option, + model: String, + status_code: i64, + transport: Option, + created_at_display: String, + account_email: Option, + requested_model: Option, + upstream_model: Option, + client_ip: Option, + user_agent: Option, + reasoning_effort: Option, + token_details: UsageRecordTokenDetailsData, + cost_details: Option, + first_token_latency_ms_display: String, + latency_ms_display: String, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct DashboardServiceStatusData { + label: String, + value: String, + detail: String, + tone: String, +} + +#[derive(Debug, Clone, Copy, Default)] +struct UsageWindow { + requests: u64, + input_tokens: u64, + output_tokens: u64, + cached_tokens: u64, + errors: u64, + first_token_latency_sum: u64, + first_token_latency_count: u64, + max_first_token_bucket_latency: u64, + min_first_token_bucket_latency: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum QuotaWindowPriority { + FiveHour, + Weekly, + Monthly, + Other, +} + +impl UsageWindow { + fn tokens(self) -> u64 { + self.input_tokens + self.output_tokens + } + + fn cache_hit_rate(self) -> Option { + if self.input_tokens > 0 { + Some(self.cached_tokens as f64 / self.input_tokens as f64) + } else { + None + } + } + + fn avg_first_token_latency(self) -> Option { + self.first_token_latency_sum + .checked_div(self.first_token_latency_count) + } + + fn max_first_token_bucket_latency(self) -> Option { + (self.max_first_token_bucket_latency > 0).then_some(self.max_first_token_bucket_latency) + } + + fn min_first_token_bucket_latency(self) -> Option { + (self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency) + } +} + +/// `GET /api/admin/dashboard/summary` +pub(crate) async fn dashboard_summary( + State(state): State, + _auth: AdminAuth, + Query(query): Query, +) -> Result { + let accounts = state + .services + .accounts + .list_pool_accounts() + .await + .unwrap_or_default(); + let capacity = state.services.account_pool.capacity_summary_now().await; + let now = Utc::now(); + let today_filter = today_usage_record_filter(now); + let usage_summary = state + .services + .usage_records + .summary(UsageQueryFilter::default()) + .await + .unwrap_or_else(|_| UsageRecordSummary::default()); + let account_usage_records = state + .services + .usage_records + .account_usage(today_filter.clone(), DASHBOARD_ACCOUNT_USAGE_LIMIT) + .await + .unwrap_or_default(); + let recent_events = state + .services + .usage_records + .list(None, DASHBOARD_USAGE_RECORD_LIMIT, today_filter) + .await + .map(|page| page.items) + .unwrap_or_default(); + + let time_buckets = dashboard_time_buckets(&state, now).await; + let account_ids = accounts + .iter() + .map(|account| account.id.clone()) + .collect::>(); + let refreshing_account_ids = state + .services + .token_refresh + .refreshing_account_ids(&account_ids, now) + .await + .map_err(|error| AdminError::internal(error.to_string()))?; + let pool_summary = account_pool_summary(&accounts, &refreshing_account_ids); + let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default()); + let quota_used_by_account = + account_quota_used_percent_by_id(&state, &account_usage_records).await; + let settings = state.services.settings.current(); + let recent_usage_records = recent_events; + let account_emails = state + .services + .usage_records + .account_email_map(&recent_usage_records) + .await + .map_err(|error| AdminError::usage_record_accounts_failed(error.to_string()))?; + let dashboard_usage_records = + dashboard_usage_record_items(recent_usage_records, &account_emails); + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(DashboardSummaryData { + cards: dashboard_cards(&accounts, &time_buckets, &usage_summary), + trend, + health_timeline: dashboard_health_timeline_data(&time_buckets), + account_usage: account_usage_data( + &accounts, + &account_usage_records, + "a_used_by_account, + now, + ), + service_statuses: service_status_data(&state), + usage_records: dashboard_usage_records, + pool_summary, + capacity_info: AccountCapacityDiagnostics::from(capacity), + rotation_strategy: Some(settings.rotation_strategy.clone()), + }), + )) +} + +fn today_usage_record_filter(now: DateTime) -> UsageQueryFilter { + UsageQueryFilter { + start_time: Some(china_day_start(now)), + end_time: Some(now), + ..UsageQueryFilter::default() + } +} + +/// `GET /api/admin/dashboard/trend?kind=usage|latency|errors` +pub(crate) async fn dashboard_trend( + State(state): State, + _auth: AdminAuth, + Query(query): Query, +) -> Result { + let time_buckets = dashboard_time_buckets(&state, Utc::now()).await; + + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(dashboard_trend_data( + &time_buckets, + query.kind.unwrap_or_default(), + )), + )) +} + +mod data; +use data::*; diff --git a/backend/src/admin/monitoring/dashboard.rs b/backend/src/api/admin/dashboard_routes/data.rs similarity index 62% rename from backend/src/admin/monitoring/dashboard.rs rename to backend/src/api/admin/dashboard_routes/data.rs index eb9466c9b..a3e5d87e3 100644 --- a/backend/src/admin/monitoring/dashboard.rs +++ b/backend/src/api/admin/dashboard_routes/data.rs @@ -1,401 +1,16 @@ -//! 管理端 Dashboard 聚合视图。 +use super::*; -use std::collections::{HashMap, HashSet}; - -use axum::{ - extract::{Query, State}, - http::StatusCode, - response::IntoResponse, -}; -use chrono::{DateTime, Duration, Utc}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; - -use crate::{ - admin::{ - auth::session::AdminAuth, - monitoring::{ - account_usage_service::AdminUsageTimeBucketRecord, - diagnostics::{ - fingerprint_diagnostics, AccountCapacityDiagnostics, AccountPoolDiagnostics, - }, - usage_record_model::{ - metadata_nonnegative_i64, metadata_string, UsageRecord, UsageRecordLevel, - }, - usage_record_routes::{ - usage_cost_details, usage_record_models, usage_token_details, - UsageRecordCostDetailsData, UsageRecordTokenDetailsData, - }, - usage_record_service::AdminUsageRecordFilter, - usage_record_store::{UsageRecordAccountUsage, UsageRecordSummary}, - }, - response::{AdminEnvelope, AdminError, AdminResponse}, - }, - infra::{ - format::{ - format_compact_number, format_cost, format_duration_ms, format_percent, format_rate, - format_tokens, nonnegative_i64_to_u64, - }, - time::{ - china_datetime, china_datetime_rfc3339_str, china_day_start, china_hour, - china_hour_start, china_quarter_hour_start, china_relative_time, - }, - }, - runtime::state::AppState, - upstream::accounts::{ - model::{Account, AccountStatus}, - token_refresh::token_refresh_status_eligible, - }, -}; - -const DASHBOARD_ACCOUNT_USAGE_LIMIT: u32 = 4; -const DASHBOARD_USAGE_RECORD_LIMIT: u32 = 10; -const HEALTH_TIMELINE_SLOT_MINUTES: i64 = 15; -const DASHBOARD_TIME_BUCKET_SLOTS: i64 = 7 * 24 * 4; -const HEALTH_TIMELINE_SLOTS: i64 = 24 * 4; -const HEALTH_TIMELINE_STABLE_SUCCESS_THRESHOLD: u64 = 3; -const FIVE_HOUR_WINDOW_SECONDS: u64 = 18_000; -const WEEK_WINDOW_SECONDS: u64 = 604_800; -const MONTH_WINDOW_SECONDS: u64 = 2_592_000; - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct DashboardTrendQuery { - kind: Option, -} - -#[derive(Debug, Clone, Copy, Default, Deserialize, Serialize, PartialEq, Eq)] -#[serde(rename_all = "lowercase")] -enum DashboardTrendKind { - #[default] - Usage, - Latency, - Errors, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardSummaryData { - cards: DashboardCardsData, - trend: DashboardTrendData, - health_timeline: DashboardHealthTimelineData, - account_usage: Vec, - service_statuses: Vec, - usage_records: Vec, - pool_summary: AccountPoolDiagnostics, - capacity_info: AccountCapacityDiagnostics, - rotation_strategy: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardHealthTimelineData { - title: String, - description: String, - reliability_display: String, - points: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardCardsData { - accounts: DashboardAccountsCardData, - traffic: DashboardTrafficCardData, - tokens: DashboardTokenCardData, - cache: DashboardCacheCardData, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardAccountsCardData { - total: String, - total_value: u64, - enabled: String, - enabled_value: u64, - abnormal: String, - abnormal_value: u64, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardTrafficCardData { - today_requests: String, - today_requests_value: u64, - yesterday_requests_value: u64, - total_requests: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardTokenCardData { - today_tokens: String, - today_tokens_value: u64, - yesterday_tokens_value: u64, - total_tokens: String, - today_cost_usd: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardCacheCardData { - today_hit_rate: String, - today_hit_rate_value: Option, - yesterday_hit_rate_value: Option, - total_hit_rate: String, - total_cached_tokens: String, - average_first_token_latency_ms: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardTrendData { - kind: DashboardTrendKind, - points: Vec, - summary: Vec, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardTrendPointData { - time: String, - requests: String, - requests_value: u64, - input_tokens: String, - input_tokens_value: u64, - output_tokens: String, - output_tokens_value: u64, - cached_tokens: String, - cached_tokens_value: u64, - cache_hit_rate_value: f64, - tokens_value: u64, - errors: String, - errors_value: u64, - latency: String, - latency_value: u64, - max_latency: String, - max_latency_value: u64, - min_latency: String, - min_latency_value: u64, - success_rate: String, - success_rate_value: f64, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardTrendSummaryData { - label: String, - value: String, - ratio: Option, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardAccountUsageData { - id: String, - email: String, - plan_type: Option, - tokens: String, - quota_used_percent: Option, - last_used: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardUsageRecordData { - id: String, - level: UsageRecordLevel, - route: Option, - model: Option, - status_code: Option, - transport: Option, - created_at_display: String, - account_email: Option, - requested_model: Option, - upstream_model: Option, - client_ip: Option, - user_agent: Option, - reasoning_effort: Option, - token_details: UsageRecordTokenDetailsData, - cost_details: Option, - first_token_latency_ms_display: String, - latency_ms_display: String, -} - -#[derive(Debug, Serialize)] -#[serde(rename_all = "camelCase")] -struct DashboardServiceStatusData { - label: String, - value: String, - detail: String, - tone: String, -} - -#[derive(Debug, Clone, Copy, Default)] -struct UsageWindow { - requests: u64, - input_tokens: u64, - output_tokens: u64, - cached_tokens: u64, - errors: u64, - first_token_latency_sum: u64, - first_token_latency_count: u64, - max_first_token_bucket_latency: u64, - min_first_token_bucket_latency: u64, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] -enum QuotaWindowPriority { - FiveHour, - Weekly, - Monthly, - Other, -} - -impl UsageWindow { - fn tokens(self) -> u64 { - self.input_tokens + self.output_tokens - } - - fn cache_hit_rate(self) -> Option { - if self.input_tokens > 0 { - Some(self.cached_tokens as f64 / self.input_tokens as f64) - } else { - None - } - } - - fn avg_first_token_latency(self) -> Option { - self.first_token_latency_sum - .checked_div(self.first_token_latency_count) - } - - fn max_first_token_bucket_latency(self) -> Option { - (self.max_first_token_bucket_latency > 0).then_some(self.max_first_token_bucket_latency) - } - - fn min_first_token_bucket_latency(self) -> Option { - (self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency) - } -} - -/// `GET /api/admin/dashboard/summary` -pub(crate) async fn dashboard_summary( - State(state): State, - _auth: AdminAuth, - Query(query): Query, -) -> Result { - let accounts = state - .services - .accounts - .list_pool_accounts() - .await - .unwrap_or_default(); - let capacity = state.services.account_pool.capacity_summary_now().await; - let now = Utc::now(); - let today_filter = today_usage_record_filter(now); - let usage_summary = state - .services - .usage_records - .summary(AdminUsageRecordFilter::default()) - .await - .unwrap_or_else(|_| UsageRecordSummary::default()); - let account_usage_records = state - .services - .usage_records - .account_usage(today_filter.clone(), DASHBOARD_ACCOUNT_USAGE_LIMIT) - .await - .unwrap_or_default(); - let recent_events = state - .services - .usage_records - .list(None, DASHBOARD_USAGE_RECORD_LIMIT, today_filter) - .await - .map(|page| page.items) - .unwrap_or_default(); - - let time_buckets = dashboard_time_buckets(&state, now).await; - let account_ids = accounts - .iter() - .map(|account| account.id.clone()) - .collect::>(); - let refreshing_account_ids = state - .services - .token_refresh - .refreshing_account_ids(&account_ids, now) - .await - .map_err(|error| AdminError::internal(error.to_string()))?; - let pool_summary = account_pool_summary(&accounts, &refreshing_account_ids); - let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default()); - let quota_used_by_account = - account_quota_used_percent_by_id(&state, &account_usage_records).await; - let settings = state.services.settings.current(); - let recent_usage_records = recent_events; - let account_emails = state - .services - .usage_records - .account_email_map(&recent_usage_records) - .await - .map_err(|error| AdminError::usage_record_accounts_failed(error.to_string()))?; - let dashboard_usage_records = - dashboard_usage_record_items(recent_usage_records, &account_emails); - - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(DashboardSummaryData { - cards: dashboard_cards(&accounts, &time_buckets, &usage_summary), - trend, - health_timeline: dashboard_health_timeline_data(&time_buckets), - account_usage: account_usage_data( - &accounts, - &account_usage_records, - "a_used_by_account, - now, - ), - service_statuses: service_status_data(&state), - usage_records: dashboard_usage_records, - pool_summary, - capacity_info: AccountCapacityDiagnostics::from(capacity), - rotation_strategy: Some(settings.auth.rotation_strategy.clone()), - }), - )) -} - -fn today_usage_record_filter(now: DateTime) -> AdminUsageRecordFilter { - AdminUsageRecordFilter { - start_time: Some(china_day_start(now)), - end_time: Some(now), - ..AdminUsageRecordFilter::default() - } -} - -/// `GET /api/admin/dashboard/trend?kind=usage|latency|errors` -pub(crate) async fn dashboard_trend( - State(state): State, - _auth: AdminAuth, - Query(query): Query, -) -> Result { - let time_buckets = dashboard_time_buckets(&state, Utc::now()).await; - - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(dashboard_trend_data( - &time_buckets, - query.kind.unwrap_or_default(), - )), - )) -} - -fn dashboard_cards( +pub(super) fn dashboard_cards( accounts: &[Account], - buckets: &[AdminUsageTimeBucketRecord], + buckets: &[AccountUsageTimeBucket], usage_summary: &UsageRecordSummary, ) -> DashboardCardsData { dashboard_cards_at(accounts, buckets, usage_summary, Utc::now()) } -fn dashboard_cards_at( +pub(super) fn dashboard_cards_at( accounts: &[Account], - buckets: &[AdminUsageTimeBucketRecord], + buckets: &[AccountUsageTimeBucket], usage_summary: &UsageRecordSummary, now: DateTime, ) -> DashboardCardsData { @@ -462,7 +77,7 @@ fn dashboard_cards_at( } } -fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option { +pub(super) fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option { let input_tokens = summary.input_tokens; if input_tokens > 0 { Some(summary.cached_tokens as f64 / input_tokens as f64) @@ -471,14 +86,14 @@ fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option { } } -fn format_optional_duration_ms(value: Option) -> String { +pub(super) fn format_optional_duration_ms(value: Option) -> String { format_duration_ms(value.and_then(|value| i64::try_from(value).ok())) } -async fn dashboard_time_buckets( +pub(super) async fn dashboard_time_buckets( state: &AppState, now: DateTime, -) -> Vec { +) -> Vec { let current_slot = china_quarter_hour_start(now); let start = current_slot - Duration::minutes(HEALTH_TIMELINE_SLOT_MINUTES * (DASHBOARD_TIME_BUCKET_SLOTS - 1)); @@ -490,15 +105,15 @@ async fn dashboard_time_buckets( .unwrap_or_default() } -fn dashboard_trend_data( - buckets: &[AdminUsageTimeBucketRecord], +pub(super) fn dashboard_trend_data( + buckets: &[AccountUsageTimeBucket], kind: DashboardTrendKind, ) -> DashboardTrendData { dashboard_trend_data_at(buckets, kind, Utc::now()) } -fn dashboard_trend_data_at( - records: &[AdminUsageTimeBucketRecord], +pub(super) fn dashboard_trend_data_at( + records: &[AccountUsageTimeBucket], kind: DashboardTrendKind, now: DateTime, ) -> DashboardTrendData { @@ -572,14 +187,14 @@ fn dashboard_trend_data_at( } } -fn dashboard_health_timeline_data( - buckets: &[AdminUsageTimeBucketRecord], +pub(super) fn dashboard_health_timeline_data( + buckets: &[AccountUsageTimeBucket], ) -> DashboardHealthTimelineData { dashboard_health_timeline_data_at(buckets, Utc::now()) } -fn dashboard_health_timeline_data_at( - records: &[AdminUsageTimeBucketRecord], +pub(super) fn dashboard_health_timeline_data_at( + records: &[AccountUsageTimeBucket], now: DateTime, ) -> DashboardHealthTimelineData { let current_slot = china_quarter_hour_start(now); @@ -645,7 +260,7 @@ fn dashboard_health_timeline_data_at( } } -fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char { +pub(super) fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char { let successes = bucket.requests.saturating_sub(bucket.errors); if is_future { '0' @@ -662,7 +277,7 @@ fn health_tone(bucket: UsageWindow, success_rate: f64, is_future: bool) -> char } } -fn trend_summary( +pub(super) fn trend_summary( kind: DashboardTrendKind, points: &[DashboardTrendPointData], ) -> Vec { @@ -740,7 +355,7 @@ fn trend_summary( } } -fn trend_summary_item( +pub(super) fn trend_summary_item( kind: DashboardTrendKind, label: &str, value: u64, @@ -753,7 +368,7 @@ fn trend_summary_item( } } -fn trend_summary_value_display(kind: DashboardTrendKind, value: u64) -> String { +pub(super) fn trend_summary_value_display(kind: DashboardTrendKind, value: u64) -> String { match kind { DashboardTrendKind::Usage => format_tokens(value), DashboardTrendKind::Latency => format_optional_duration_ms(Some(value)), @@ -761,8 +376,8 @@ fn trend_summary_value_display(kind: DashboardTrendKind, value: u64) -> String { } } -fn usage_window( - records: &[AdminUsageTimeBucketRecord], +pub(super) fn usage_window( + records: &[AccountUsageTimeBucket], start: DateTime, end: DateTime, ) -> UsageWindow { @@ -775,8 +390,8 @@ fn usage_window( window } -fn cost_window( - records: &[AdminUsageTimeBucketRecord], +pub(super) fn cost_window( + records: &[AccountUsageTimeBucket], start: DateTime, end: DateTime, ) -> Option { @@ -793,7 +408,7 @@ fn cost_window( has_usage.then_some(total) } -fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) { +pub(super) fn apply_bucket(window: &mut UsageWindow, record: &AccountUsageTimeBucket) { window.requests += nonnegative_i64_to_u64(record.request_count); window.input_tokens += nonnegative_i64_to_u64(record.input_tokens); window.output_tokens += nonnegative_i64_to_u64(record.output_tokens); @@ -812,13 +427,13 @@ fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) { } } -fn first_token_bucket_latency(record: &AdminUsageTimeBucketRecord) -> Option { +pub(super) fn first_token_bucket_latency(record: &AccountUsageTimeBucket) -> Option { let sum = nonnegative_i64_to_u64(record.first_token_latency_sum); let count = nonnegative_i64_to_u64(record.first_token_latency_count); sum.checked_div(count).filter(|latency| *latency > 0) } -fn account_pool_summary( +pub(super) fn account_pool_summary( accounts: &[Account], refreshing_account_ids: &HashSet, ) -> AccountPoolDiagnostics { @@ -844,7 +459,7 @@ fn account_pool_summary( summary } -fn account_usage_data( +pub(super) fn account_usage_data( accounts: &[Account], usage_records: &[UsageRecordAccountUsage], quota_used_by_account: &HashMap, @@ -873,7 +488,7 @@ fn account_usage_data( email: account .and_then(|account| account.email.as_ref()) .cloned() - .expect("dashboard account usage requires email"), + .unwrap_or_else(|| usage.account_id.clone()), plan_type: account .and_then(|account| account.plan_type.as_ref()) .cloned(), @@ -885,7 +500,7 @@ fn account_usage_data( .collect() } -async fn account_quota_used_percent_by_id( +pub(super) async fn account_quota_used_percent_by_id( state: &AppState, usage_records: &[UsageRecordAccountUsage], ) -> HashMap { @@ -906,7 +521,7 @@ async fn account_quota_used_percent_by_id( quota_used_by_account } -fn quota_used_percent(quota_json: &str) -> Option { +pub(super) fn quota_used_percent(quota_json: &str) -> Option { let quota = serde_json::from_str::(quota_json).ok()?; let mut selected = None; if let Some(monthly_limit) = quota.get("monthly_limit") { @@ -932,7 +547,7 @@ fn quota_used_percent(quota_json: &str) -> Option { selected.map(|(_, used_percent)| used_percent) } -fn select_quota_used_percent( +pub(super) fn select_quota_used_percent( selected: &mut Option<(QuotaWindowPriority, f64)>, priority: QuotaWindowPriority, used_percent: Option, @@ -948,7 +563,10 @@ fn select_quota_used_percent( } } -fn quota_window_priority(window: &Value, fallback: QuotaWindowPriority) -> QuotaWindowPriority { +pub(super) fn quota_window_priority( + window: &Value, + fallback: QuotaWindowPriority, +) -> QuotaWindowPriority { let window_seconds = window .get("window_minutes") .and_then(Value::as_u64) @@ -967,18 +585,18 @@ fn quota_window_priority(window: &Value, fallback: QuotaWindowPriority) -> Quota } } -fn quota_window_matches(actual: u64, expected: u64) -> bool { +pub(super) fn quota_window_matches(actual: u64, expected: u64) -> bool { actual > 0 && actual.abs_diff(expected) <= expected / 20 } -fn percent_value(value: &Value) -> Option { +pub(super) fn percent_value(value: &Value) -> Option { let percent = value .as_f64() .or_else(|| value.as_str().and_then(|value| value.parse::().ok()))?; percent.is_finite().then_some(percent.clamp(0.0, 100.0)) } -fn dashboard_usage_record_items( +pub(super) fn dashboard_usage_record_items( items: Vec, account_emails: &HashMap, ) -> Vec { @@ -988,14 +606,12 @@ fn dashboard_usage_record_items( .collect() } -fn dashboard_usage_record_data( +pub(super) fn dashboard_usage_record_data( record: UsageRecord, account_emails: &HashMap, ) -> DashboardUsageRecordData { - let account_email = record - .account_id - .as_deref() - .and_then(|account_id| account_emails.get(account_id)) + let account_email = account_emails + .get(&record.account_id) .cloned() .or_else(|| metadata_string(&record.metadata, &["accountEmail", "account_email"])); let (requested_model, upstream_model) = usage_record_models(&record); @@ -1003,21 +619,12 @@ 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 token_details = usage_token_details(&record.metadata); + let token_details = usage_token_details(&record); let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details); - let first_token_latency_ms = metadata_nonnegative_i64( - &record.metadata, - &[ - "firstTokenMs", - "first_token_ms", - "firstTokenLatencyMs", - "first_token_latency_ms", - ], - ); + let first_token_latency_ms = record.first_token_ms; DashboardUsageRecordData { id: record.id, - level: record.level, route: record.route, model: record.model, status_code: record.status_code, @@ -1036,7 +643,7 @@ fn dashboard_usage_record_data( } } -fn service_status_data(state: &AppState) -> Vec { +pub(super) fn service_status_data(state: &AppState) -> Vec { let current_fingerprint = state.services.fingerprint.snapshot(); let fingerprint = fingerprint_diagnostics(¤t_fingerprint); vec![ @@ -1081,7 +688,7 @@ fn service_status_data(state: &AppState) -> Vec { ] } -fn empty_dash(value: String) -> String { +pub(super) fn empty_dash(value: String) -> String { if value.trim().is_empty() { "-".to_string() } else { @@ -1089,7 +696,7 @@ fn empty_dash(value: String) -> String { } } -fn format_fingerprint_updated_at(value: Option) -> String { +pub(super) fn format_fingerprint_updated_at(value: Option) -> String { let Some(value) = value else { return "-".to_string(); }; diff --git a/backend/src/admin/keys/routes.rs b/backend/src/api/admin/keys_routes.rs similarity index 90% rename from backend/src/admin/keys/routes.rs rename to backend/src/api/admin/keys_routes.rs index a9d118b95..7e0728910 100644 --- a/backend/src/admin/keys/routes.rs +++ b/backend/src/api/admin/keys_routes.rs @@ -10,17 +10,17 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::{ - admin::auth::session::AdminAuth, - admin::keys::service::{AdminClientApiKey, AdminClientKeyError}, - admin::response::{ + api::admin::response::{ AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, BatchDeleteData, }, - admin::update_payload::{parse_editable_update, EditableUpdateMessages}, + api::admin::session::AdminAuth, + api::admin::update_payload::{parse_editable_update, EditableUpdateMessages}, + bootstrap::state::AppState, infra::{ json::{clamp_limit, Page}, time::{china_relative_time_str, china_rfc3339_str}, }, - runtime::state::AppState, + keys::types::{KeyManageError, ManagedClientApiKey}, }; #[derive(Debug, Deserialize)] @@ -57,8 +57,8 @@ struct ClientApiKeyData { last_used_at_display: String, } -impl From for ClientApiKeyData { - fn from(k: AdminClientApiKey) -> Self { +impl From for ClientApiKeyData { + fn from(k: ManagedClientApiKey) -> Self { Self { id: k.id, name: k.name, @@ -74,12 +74,12 @@ impl From for ClientApiKeyData { } } -fn client_key_error(error: &AdminClientKeyError) -> AdminError { +fn client_key_error(error: &KeyManageError) -> AdminError { match error { - AdminClientKeyError::InvalidStatus(_) - | AdminClientKeyError::EmptyName - | AdminClientKeyError::EmptyIds - | AdminClientKeyError::LabelTooLong => AdminError::bad_request(error.to_string()), + KeyManageError::InvalidStatus(_) + | KeyManageError::EmptyName + | KeyManageError::EmptyIds + | KeyManageError::LabelTooLong => AdminError::bad_request(error.to_string()), _ => AdminError::internal(error.to_string()), } } diff --git a/backend/src/api/admin/mod.rs b/backend/src/api/admin/mod.rs new file mode 100644 index 000000000..dc908655d --- /dev/null +++ b/backend/src/api/admin/mod.rs @@ -0,0 +1,13 @@ +//! 管理端 HTTP API。 + +pub mod accounts_routes; +pub mod dashboard_routes; +pub mod keys_routes; +pub mod ops_routes; +pub mod response; +pub mod router; +pub mod session; +pub mod settings_routes; +pub mod system_routes; +pub mod update_payload; +pub mod usage_routes; diff --git a/backend/src/api/admin/ops_routes.rs b/backend/src/api/admin/ops_routes.rs new file mode 100644 index 000000000..d6aeb6846 --- /dev/null +++ b/backend/src/api/admin/ops_routes.rs @@ -0,0 +1,155 @@ +//! 运维错误明细 HTTP 处理器。 + +use axum::{ + extract::{Query, State}, + http::StatusCode, + response::IntoResponse, +}; +use serde::{Deserialize, Serialize}; + +use crate::{ + api::admin::{ + response::{AdminError, AdminPageEnvelope, AdminResponse}, + session::AdminAuth, + }, + bootstrap::state::AppState, + infra::{ + json::{clamp_limit, clamp_page, NumberedPage, Page}, + time::china_datetime, + }, + telemetry::{ops::query::OpsQueryError, ops::store::OpsErrorFilter, ops::types::OpsErrorLog}, +}; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct OpsErrorsQuery { + cursor: Option, + limit: Option, + page: Option, + page_size: Option, + kind: Option, + client_api_key_id: Option, + provider: Option, + request_id: Option, + account_id: Option, + route: Option, + model: Option, + status_code: Option, + client_status_code: Option, + upstream_status_code: Option, + transport: Option, + attempt_index: Option, + failure_class: Option, + response_id: Option, + upstream_request_id: Option, + search: Option, + start_time: Option, + end_time: Option, +} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct OpsErrorData { + #[serde(flatten)] + error: OpsErrorLog, + created_at_display: String, +} + +/// `GET /api/admin/ops/errors` +pub(crate) async fn ops_errors( + State(state): State, + _auth: AdminAuth, + Query(query): Query, +) -> Result { + let limit = clamp_limit(query.page_size.or(query.limit).unwrap_or(50)); + let page = query.page; + let numbered = page.is_some() || query.page_size.is_some(); + let cursor = query.cursor.clone(); + let filter = filter_from_query(query)?; + + if numbered { + return state + .services + .ops_errors + .list_page(clamp_page(page.unwrap_or(1)), limit, filter) + .await + .map(|page| { + let page = NumberedPage { + items: ops_error_items(page.items), + total: page.total, + page: page.page, + page_size: page.page_size, + }; + AdminResponse::new(StatusCode::OK, AdminPageEnvelope::numbered(page)) + }) + .map_err(log_error); + } + + state + .services + .ops_errors + .list(cursor, limit, filter) + .await + .map(|page| { + let page = Page { + items: ops_error_items(page.items), + next_cursor: page.next_cursor, + }; + AdminResponse::new(StatusCode::OK, AdminPageEnvelope::ok(page, limit)) + }) + .map_err(log_error) +} + +fn filter_from_query(query: OpsErrorsQuery) -> Result { + Ok(OpsErrorFilter { + kind: non_empty(query.kind), + client_api_key_id: non_empty(query.client_api_key_id), + provider: non_empty(query.provider), + request_id: non_empty(query.request_id), + account_id: non_empty(query.account_id), + route: non_empty(query.route), + model: non_empty(query.model), + status_code: query.status_code, + client_status_code: query.client_status_code, + upstream_status_code: query.upstream_status_code, + transport: non_empty(query.transport), + attempt_index: query.attempt_index, + failure_class: non_empty(query.failure_class), + response_id: non_empty(query.response_id), + upstream_request_id: non_empty(query.upstream_request_id), + search: non_empty(query.search), + start_time: optional_datetime(query.start_time)?, + end_time: optional_datetime(query.end_time)?, + }) +} + +fn ops_error_items(items: Vec) -> Vec { + items + .into_iter() + .map(|error| OpsErrorData { + created_at_display: china_datetime(&error.created_at), + error, + }) + .collect() +} + +fn optional_datetime( + value: Option, +) -> Result>, AdminError> { + let Some(value) = non_empty(value) else { + return Ok(None); + }; + chrono::DateTime::parse_from_rfc3339(&value) + .map(|value| Some(value.with_timezone(&chrono::Utc))) + .map_err(|_| AdminError::invalid_time_range("Invalid time range")) +} + +fn non_empty(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +fn log_error(error: OpsQueryError) -> AdminError { + AdminError::internal(error.to_string()) +} diff --git a/backend/src/admin/response.rs b/backend/src/api/admin/response.rs similarity index 100% rename from backend/src/admin/response.rs rename to backend/src/api/admin/response.rs diff --git a/backend/src/admin/router.rs b/backend/src/api/admin/router.rs similarity index 85% rename from backend/src/admin/router.rs rename to backend/src/api/admin/router.rs index 2a55de747..eb9329e2a 100644 --- a/backend/src/admin/router.rs +++ b/backend/src/api/admin/router.rs @@ -5,32 +5,31 @@ use axum::{ Router, }; -use crate::runtime::state::AppState; +use crate::bootstrap::state::AppState; use super::{ - accounts::routes::{ + accounts_routes::{ account_models, account_quota, accounts, batch_delete_accounts, create_account, export_accounts, health_check_accounts, import_accounts, oauth_authorize_account, oauth_exchange_account, refresh_account, test_account_connection, update_account, }, - auth::session::{login, logout, session_status}, - keys::routes::{api_keys, batch_delete_api_keys, create_api_key, update_api_key}, - monitoring::{ - dashboard::{dashboard_summary, dashboard_trend}, - usage_record_routes::{ - clear_usage_records, usage_record_detail, usage_records, - usage_records_endpoint_distribution, usage_records_latency_trend, - usage_records_model_distribution, usage_records_summary, usage_records_token_trend, - }, - }, - settings::routes::{ + dashboard_routes::{dashboard_summary, dashboard_trend}, + keys_routes::{api_keys, batch_delete_api_keys, create_api_key, update_api_key}, + ops_routes::ops_errors, + session::{login, logout, session_status}, + settings_routes::{ admin_api_key_status, delete_admin_api_key, regenerate_admin_api_key, settings, update_settings, }, - system::routes::{ + system_routes::{ perform_update, restart, rollback, update_detail, update_event_stream, update_status, version, }, + usage_routes::{ + clear_usage_records, usage_record_detail, usage_records, + usage_records_endpoint_distribution, usage_records_latency_trend, + usage_records_model_distribution, usage_records_summary, usage_records_token_trend, + }, }; /// 构造管理端路由。 @@ -101,6 +100,7 @@ pub fn router() -> Router { ) .route("/api/admin/usage/records/delete", post(clear_usage_records)) .route("/api/admin/usage/records/detail", get(usage_record_detail)) + .route("/api/admin/ops/errors", get(ops_errors)) .route("/api/admin/keys", get(api_keys).post(create_api_key)) .route("/api/admin/keys/delete", post(batch_delete_api_keys)) .route("/api/admin/keys/update", post(update_api_key)) diff --git a/backend/src/admin/auth/session.rs b/backend/src/api/admin/session.rs similarity index 95% rename from backend/src/admin/auth/session.rs rename to backend/src/api/admin/session.rs index 75e899f1d..8e9e9f3ca 100644 --- a/backend/src/admin/auth/session.rs +++ b/backend/src/api/admin/session.rs @@ -9,13 +9,13 @@ use axum::{ use serde::{Deserialize, Serialize}; use crate::{ - admin::{ - auth::service::AdminSessionError, - response::{AdminEnvelope, AdminError, AdminResponse}, + api::{ + admin::response::{AdminEnvelope, AdminError, AdminResponse}, + middleware::request_id::ClientIp, }, - http::middleware::request_id::ClientIp, + auth::types::SessionError, + bootstrap::state::AppState, infra::time::china_rfc3339, - runtime::state::AppState, }; const ADMIN_SESSION_COOKIE: &str = "cpr_admin_session"; @@ -74,7 +74,7 @@ pub(crate) async fn login( let session = match login_result { Ok(Some(session)) => session, Ok(None) => return Err(AdminError::invalid_admin_credentials()), - Err(AdminSessionError::LoginThrottled) => return Err(AdminError::too_many_login_attempts()), + Err(SessionError::LoginThrottled) => return Err(AdminError::too_many_login_attempts()), Err(_) => return Err(AdminError::internal("Failed to create admin session")), }; diff --git a/backend/src/admin/settings/routes.rs b/backend/src/api/admin/settings_routes.rs similarity index 70% rename from backend/src/admin/settings/routes.rs rename to backend/src/api/admin/settings_routes.rs index 2197e358e..6b37d233e 100644 --- a/backend/src/admin/settings/routes.rs +++ b/backend/src/api/admin/settings_routes.rs @@ -6,14 +6,12 @@ use serde_json::Value; use std::collections::BTreeMap; use crate::{ - admin::auth::session::AdminAuth, - admin::response::{AdminEnvelope, AdminError, AdminResponse}, - config::{ - schema::AppConfig, - settings::{AdminSettingsPatch, RuntimeSettingsError, SettingsServiceError}, + api::admin::{ + response::{AdminEnvelope, AdminError, AdminResponse}, + session::AdminAuth, }, - runtime::state::AppState, - upstream::models::config::ModelConfig, + bootstrap::state::AppState, + settings::{SettingsError, SettingsPatch, SettingsSnapshot, SettingsValidationError}, }; const ALLOWED_SETTINGS_KEYS: [&str; 6] = [ @@ -37,14 +35,14 @@ struct AdminSettingsData { } impl AdminSettingsData { - fn from_config(config: &AppConfig) -> Self { + fn from_settings(settings: &SettingsSnapshot) -> Self { Self { - model_aliases: config.model_aliases.clone(), - refresh_margin_seconds: config.auth.refresh_margin_seconds, - refresh_concurrency: config.auth.refresh_concurrency, - max_concurrent_per_account: config.auth.max_concurrent_per_account, - request_interval_ms: config.auth.request_interval_ms, - rotation_strategy: config.auth.rotation_strategy.clone(), + model_aliases: settings.model_aliases.clone(), + refresh_margin_seconds: settings.refresh_margin_seconds, + refresh_concurrency: settings.refresh_concurrency, + max_concurrent_per_account: settings.max_concurrent_per_account, + request_interval_ms: settings.request_interval_ms, + rotation_strategy: settings.rotation_strategy.clone(), } } } @@ -57,7 +55,7 @@ pub(crate) async fn settings( let config = state.services.settings.current(); Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminSettingsData::from_config(&config)), + AdminEnvelope::ok(AdminSettingsData::from_settings(&config)), )) } @@ -74,13 +72,11 @@ pub(crate) async fn update_settings( .update(patch) .await .map_err(settings_service_error)?; - state.services.models.update_config(ModelConfig { - model_aliases: config.model_aliases.clone(), - }); + state.services.apply_settings(&config).await; Ok(AdminResponse::new( StatusCode::OK, - AdminEnvelope::ok(AdminSettingsData::from_config(&config)), + AdminEnvelope::ok(AdminSettingsData::from_settings(&config)), )) } @@ -138,7 +134,7 @@ pub(crate) async fn delete_admin_api_key( )) } -fn parse_settings_patch(body: &[u8]) -> Result { +fn parse_settings_patch(body: &[u8]) -> Result { let value = serde_json::from_slice::(body) .map_err(|_| AdminError::malformed_json("Malformed JSON body"))?; let object = value @@ -157,15 +153,12 @@ fn parse_settings_patch(body: &[u8]) -> Result { serde_json::from_value(value).map_err(|_| AdminError::bad_request("Invalid settings payload")) } -fn settings_service_error(error: RuntimeSettingsError) -> AdminError { +fn settings_service_error(error: SettingsError) -> AdminError { match error { - RuntimeSettingsError::InvalidField(SettingsServiceError::InvalidField { - field, - message, - }) => AdminError::bad_request(format!("Invalid setting {field}: {message}")), - RuntimeSettingsError::Database(_) - | RuntimeSettingsError::Json(_) - | RuntimeSettingsError::StoredField { .. } => { + SettingsError::InvalidField(SettingsValidationError::InvalidField { field, message }) => { + AdminError::bad_request(format!("Invalid setting {field}: {message}")) + } + SettingsError::Database(_) | SettingsError::Json(_) | SettingsError::StoredField { .. } => { AdminError::settings_persist_failed("Failed to persist settings") } } diff --git a/backend/src/api/admin/system_routes.rs b/backend/src/api/admin/system_routes.rs new file mode 100644 index 000000000..1d306ba3a --- /dev/null +++ b/backend/src/api/admin/system_routes.rs @@ -0,0 +1,151 @@ +//! 管理端系统版本、自更新与重启 HTTP 处理器。 + +use std::{convert::Infallible, time::Duration}; + +use axum::{ + extract::{Json, Query}, + http::StatusCode, + response::{ + sse::{Event, KeepAlive, Sse}, + IntoResponse, + }, +}; +use futures::Stream; +use serde::Deserialize; + +use crate::{ + api::admin::{ + response::{AdminEnvelope, AdminError, AdminResponse}, + session::AdminAuth, + }, + bootstrap::shutdown::{request_process_restart, request_shutdown, shutdown_subscription}, + update::{ + service::{self, RestartAction}, + state::operation_id, + types::UpdateError, + }, +}; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct UpdateDetailQuery { + refresh: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct UpdateRequest { + target_version: String, +} + +pub(crate) async fn version(_auth: AdminAuth) -> Result { + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(service::version_data().await), + )) +} + +pub(crate) async fn update_detail( + _auth: AdminAuth, + Query(query): Query, +) -> Result { + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(service::update_detail(query.refresh.unwrap_or(false)).await), + )) +} + +pub(crate) async fn update_event_stream( + _auth: AdminAuth, +) -> Result>>, AdminError> { + let receiver = service::subscribe_update_events(); + let shutdown = shutdown_subscription(); + let stream = futures::stream::unfold( + (receiver, shutdown, false), + |(mut receiver, mut shutdown, close_after_send)| async move { + if close_after_send { + return None; + } + loop { + tokio::select! { + _ = shutdown.recv() => return None, + received = receiver.recv() => match received { + Ok(message) => { + let close_after_send = message.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, + } + } + } + }, + ); + Ok(Sse::new(stream).keep_alive(KeepAlive::default())) +} + +pub(crate) async fn perform_update( + _auth: AdminAuth, + payload: Option>, +) -> Result { + let target = payload.map(|Json(payload)| payload.target_version); + let result = service::perform_update(target) + .await + .map_err(update_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(result), + )) +} + +pub(crate) async fn update_status(_auth: AdminAuth) -> Result { + let status = service::update_status().map_err(update_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(status), + )) +} + +pub(crate) async fn rollback(_auth: AdminAuth) -> Result { + let operation_id = service::rollback().await.map_err(update_error)?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(serde_json::json!({ + "message": "回滚完成，请重启服务。", + "needRestart": true, + "operationId": operation_id, + })), + )) +} + +pub(crate) async fn restart(_auth: AdminAuth) -> Result { + let plan = service::restart_plan().map_err(update_error)?; + let message = plan.message; + tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(500)).await; + match plan.action { + RestartAction::Exec(executable_path) => request_process_restart(executable_path), + RestartAction::Shutdown => request_shutdown(), + } + }); + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(serde_json::json!({ + "message": message, + "operationId": operation_id("restart"), + })), + )) +} + +fn update_error(error: UpdateError) -> AdminError { + match error { + UpdateError::BadRequest(message) => AdminError::bad_request(message), + UpdateError::Conflict(message) => AdminError::conflict(message), + UpdateError::BadGateway(message) => AdminError::bad_gateway(message), + UpdateError::Internal(message) => AdminError::internal(message), + } +} diff --git a/backend/src/admin/update_payload.rs b/backend/src/api/admin/update_payload.rs similarity index 97% rename from backend/src/admin/update_payload.rs rename to backend/src/api/admin/update_payload.rs index 11584e417..9a52804a0 100644 --- a/backend/src/admin/update_payload.rs +++ b/backend/src/api/admin/update_payload.rs @@ -3,7 +3,7 @@ use serde::Deserialize; use serde_json::Value; -use crate::admin::response::AdminError; +use crate::api::admin::response::AdminError; #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase", deny_unknown_fields)] diff --git a/backend/src/admin/monitoring/usage_record_routes.rs b/backend/src/api/admin/usage_routes.rs similarity index 82% rename from backend/src/admin/monitoring/usage_record_routes.rs rename to backend/src/api/admin/usage_routes.rs index 379fae96e..19842f0c5 100644 --- a/backend/src/admin/monitoring/usage_record_routes.rs +++ b/backend/src/api/admin/usage_routes.rs @@ -8,35 +8,29 @@ use axum::{ response::IntoResponse, }; use serde::{Deserialize, Serialize}; -use serde_json::Value; use crate::{ - admin::monitoring::{ - usage_record_model::{ - metadata_nonnegative_i64, metadata_string, UsageRecord, UsageRecordLevel, - }, - usage_record_service::{ - usage_record_cost_details, AdminUsageRecordError, AdminUsageRecordFilter, - }, - usage_record_store::{ - UsageRecordBreakdown, UsageRecordEndpointSource, UsageRecordModelSource, - UsageRecordSummary, UsageRecordTrendPoint, - }, - }, - admin::{ - auth::session::AdminAuth, + api::admin::{ response::{AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse}, + session::AdminAuth, }, + bootstrap::state::AppState, infra::{ format::{ format_compact_number, format_cost, format_duration_ms, format_duration_ms_f64, - format_multiplier, format_number, format_percent, format_rate, format_token_price, - format_tokens, + format_multiplier, format_number, format_percent, format_token_price, format_tokens, }, json::{clamp_limit, clamp_page, NumberedPage, Page}, time::china_datetime, }, - runtime::state::AppState, + telemetry::{ + usage::query::{usage_record_cost_details, UsageQueryError, UsageQueryFilter}, + usage::store::{ + UsageRecordBreakdown, UsageRecordEndpointSource, UsageRecordModelSource, + UsageRecordSummary, UsageRecordTrendPoint, + }, + usage::types::{metadata_string, UsageRecord}, + }, }; #[derive(Debug, Deserialize)] @@ -47,7 +41,8 @@ pub(crate) struct UsageRecordsQuery { page: Option, page_size: Option, kind: Option, - level: Option, + client_api_key_id: Option, + provider: Option, request_id: Option, account_id: Option, route: Option, @@ -55,8 +50,6 @@ pub(crate) struct UsageRecordsQuery { status_code: Option, transport: Option, attempt_index: Option, - upstream_status_code: Option, - failure_class: Option, response_id: Option, upstream_request_id: Option, search: Option, @@ -147,8 +140,6 @@ pub(crate) struct UsageRecordCostDetailsData { #[serde(rename_all = "camelCase")] struct UsageRecordSummaryData { total_requests: String, - error_requests: String, - error_rate: String, input_tokens: String, output_tokens: String, cached_tokens: String, @@ -183,8 +174,6 @@ struct UsageRecordTrendPointData { input_tokens_value: u64, output_tokens: String, output_tokens_value: u64, - cache_creation_tokens: String, - cache_creation_tokens_value: u64, cached_tokens: String, cached_tokens_value: u64, total_tokens: String, @@ -401,10 +390,11 @@ pub(crate) async fn clear_usage_records( } } -fn filter_from_query(query: UsageRecordsQuery) -> Result { - Ok(AdminUsageRecordFilter { +fn filter_from_query(query: UsageRecordsQuery) -> Result { + Ok(UsageQueryFilter { kind: non_empty(query.kind), - level: level_from_query(query.level).map_err(AdminError::bad_request)?, + client_api_key_id: non_empty(query.client_api_key_id), + provider: non_empty(query.provider), request_id: non_empty(query.request_id), account_id: non_empty(query.account_id), route: non_empty(query.route), @@ -412,8 +402,6 @@ fn filter_from_query(query: UsageRecordsQuery) -> Result AdminError { +fn log_error(error: &UsageQueryError) -> AdminError { match error { - AdminUsageRecordError::List - | AdminUsageRecordError::Get - | AdminUsageRecordError::Clear - | AdminUsageRecordError::Append - | AdminUsageRecordError::Retention => AdminError::internal(error.to_string()), - AdminUsageRecordError::Accounts => { - AdminError::usage_record_accounts_failed(error.to_string()) - } - } -} - -fn level_from_query(value: Option) -> Result, String> { - let Some(value) = non_empty(value) else { - return Ok(None); - }; - match value.as_str() { - "debug" => Ok(Some(UsageRecordLevel::Debug)), - "info" => Ok(Some(UsageRecordLevel::Info)), - "warn" => Ok(Some(UsageRecordLevel::Warn)), - "error" => Ok(Some(UsageRecordLevel::Error)), - other => Err(format!("Unsupported log level: {other}")), + UsageQueryError::List + | UsageQueryError::Get + | UsageQueryError::Clear + | UsageQueryError::Append + | UsageQueryError::Retention => AdminError::internal(error.to_string()), + UsageQueryError::Accounts => AdminError::usage_record_accounts_failed(error.to_string()), } } @@ -503,10 +476,8 @@ fn usage_record_data( record: UsageRecord, account_emails: &HashMap, ) -> UsageRecordData { - let account_email = record - .account_id - .as_deref() - .and_then(|account_id| account_emails.get(account_id)) + let account_email = account_emails + .get(&record.account_id) .cloned() .or_else(|| metadata_string(&record.metadata, &["accountEmail", "account_email"])); let (requested_model, upstream_model) = usage_record_models(&record); @@ -514,17 +485,9 @@ 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 token_details = usage_token_details(&record.metadata); + let token_details = usage_token_details(&record); let cost_details = usage_cost_details(&record, upstream_model.as_deref(), &token_details); - let first_token_latency_ms = metadata_nonnegative_i64( - &record.metadata, - &[ - "firstTokenMs", - "first_token_ms", - "firstTokenLatencyMs", - "first_token_latency_ms", - ], - ); + 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); @@ -546,27 +509,38 @@ fn usage_record_data( } pub(crate) fn usage_record_models(record: &UsageRecord) -> (Option, Option) { - let stored_model = record - .model - .as_deref() - .map(str::trim) - .filter(|model| !model.is_empty()); - let requested_model = metadata_string(&record.metadata, &["requestedModel"]) - .or_else(|| stored_model.map(ToString::to_string)); - let upstream_model = metadata_string(&record.metadata, &["upstreamModel"]) + let stored_model = record.model.trim(); + let requested_model = record + .requested_model + .clone() + .or_else(|| (!stored_model.is_empty()).then(|| stored_model.to_string())); + let upstream_model = record + .upstream_model + .clone() .or_else(|| requested_model.clone()) - .or_else(|| stored_model.map(ToString::to_string)); + .or_else(|| (!stored_model.is_empty()).then(|| stored_model.to_string())); (requested_model, upstream_model) } -pub(crate) fn usage_token_details(metadata: &Value) -> UsageRecordTokenDetailsData { - let input_tokens = metadata_usage_u64(metadata, "inputTokens"); - let output_tokens = metadata_usage_u64(metadata, "outputTokens"); - let cached_tokens = metadata_usage_u64(metadata, "cachedTokens"); - let reasoning_tokens = metadata_usage_u64(metadata, "reasoningTokens"); - let total_tokens = - metadata_usage_u64(metadata, "totalTokens").max(input_tokens.saturating_add(output_tokens)); +pub(crate) fn usage_token_details(record: &UsageRecord) -> UsageRecordTokenDetailsData { + let input_tokens: u64 = record + .input_tokens + .and_then(|value| value.try_into().ok()) + .unwrap_or(0); + let output_tokens: u64 = record + .output_tokens + .and_then(|value| value.try_into().ok()) + .unwrap_or(0); + let cached_tokens: u64 = record + .cached_tokens + .and_then(|value| value.try_into().ok()) + .unwrap_or(0); + let reasoning_tokens: u64 = record + .reasoning_tokens + .and_then(|value| value.try_into().ok()) + .unwrap_or(0); + let total_tokens = input_tokens.saturating_add(output_tokens); UsageRecordTokenDetailsData { input_tokens, @@ -619,32 +593,10 @@ pub(crate) fn usage_cost_details( }) } -fn metadata_usage_u64(metadata: &Value, key: &str) -> u64 { - metadata - .get("usage") - .and_then(|usage| metadata_u64_at(usage, key)) - .or_else(|| metadata_u64_at(metadata, key)) - .unwrap_or_default() -} - -fn metadata_u64_at(value: &Value, key: &str) -> Option { - value - .get(key) - .and_then(|value| value.as_u64().or_else(|| value.as_i64()?.try_into().ok())) -} - impl From for UsageRecordSummaryData { fn from(summary: UsageRecordSummary) -> Self { Self { total_requests: format_compact_number(summary.total_requests), - error_requests: format_compact_number(summary.error_requests), - error_rate: if summary.total_requests > 0 { - format_rate(Some( - summary.error_requests as f64 / summary.total_requests as f64, - )) - } else { - "—".to_string() - }, input_tokens: format_tokens(summary.input_tokens), output_tokens: format_tokens(summary.output_tokens), cached_tokens: format_tokens(summary.cached_tokens), @@ -694,10 +646,7 @@ fn usage_record_breakdown_items(items: Vec) -> Vec for UsageRecordTrendPointData { fn from(point: UsageRecordTrendPoint) -> Self { - let prompt_tokens = point - .input_tokens - .saturating_add(point.cache_creation_tokens) - .saturating_add(point.cached_tokens); + let prompt_tokens = point.input_tokens.saturating_add(point.cached_tokens); let cache_hit_rate_value = if prompt_tokens > 0 { (point.cached_tokens as f64 / prompt_tokens as f64 * 100.0).round() } else { @@ -709,8 +658,6 @@ impl From for UsageRecordTrendPointData { input_tokens_value: point.input_tokens, output_tokens: format_tokens(point.output_tokens), output_tokens_value: point.output_tokens, - cache_creation_tokens: format_tokens(point.cache_creation_tokens), - cache_creation_tokens_value: point.cache_creation_tokens, cached_tokens: format_tokens(point.cached_tokens), cached_tokens_value: point.cached_tokens, total_tokens: format_tokens(point.total_tokens), diff --git a/backend/src/web/assets.rs b/backend/src/api/assets.rs similarity index 100% rename from backend/src/web/assets.rs rename to backend/src/api/assets.rs diff --git a/backend/src/proxy/auth.rs b/backend/src/api/client/auth.rs similarity index 87% rename from backend/src/proxy/auth.rs rename to backend/src/api/client/auth.rs index 1a5a60d72..1c895b6f0 100644 --- a/backend/src/proxy/auth.rs +++ b/backend/src/api/client/auth.rs @@ -4,7 +4,7 @@ use axum::http::{header::AUTHORIZATION, HeaderMap}; -use crate::runtime::state::AppState; +use crate::bootstrap::state::AppState; /// client API key 鉴权失败原因。 #[derive(Debug, Clone, PartialEq, Eq)] @@ -56,11 +56,16 @@ pub fn bearer_client_api_key(headers: &HeaderMap) -> Result<&str, ClientApiKeyAu /// 校验客户端 API key。 pub async fn authorize_client_api_key(state: &AppState, headers: &HeaderMap) -> bool { + authorized_client_api_key_id(state, headers).await.is_some() +} + +/// 校验客户端 API key，成功时返回用于事实归因的稳定 ID。 +pub async fn authorized_client_api_key_id(state: &AppState, headers: &HeaderMap) -> Option { match authorize_client_api_key_result(state, headers).await { - Ok(()) => true, + Ok(key_id) => Some(key_id), Err(error) => { log_client_api_key_auth_failure(&error); - false + None } } } @@ -69,7 +74,7 @@ pub async fn authorize_client_api_key(state: &AppState, headers: &HeaderMap) -> pub async fn authorize_client_api_key_result( state: &AppState, headers: &HeaderMap, -) -> Result<(), ClientApiKeyAuthError> { +) -> Result { let api_key = bearer_client_api_key(headers)?; let verified = state .services @@ -79,9 +84,7 @@ pub async fn authorize_client_api_key_result( .map_err(|error| ClientApiKeyAuthError::StoreUnavailable { message: error.to_string(), })?; - verified - .then_some(()) - .ok_or(ClientApiKeyAuthError::InvalidKey) + verified.ok_or(ClientApiKeyAuthError::InvalidKey) } fn log_client_api_key_auth_failure(error: &ClientApiKeyAuthError) { diff --git a/backend/src/proxy/openai/errors.rs b/backend/src/api/client/errors.rs similarity index 98% rename from backend/src/proxy/openai/errors.rs rename to backend/src/api/client/errors.rs index 1b572952d..054e08a2f 100644 --- a/backend/src/proxy/openai/errors.rs +++ b/backend/src/api/client/errors.rs @@ -7,10 +7,9 @@ use axum::{ }; use serde_json::{json, Value}; -use crate::proxy::{ - dispatch::{errors::DispatchFailureClass, responses::errors::ResponseDispatchError}, - openai::responses::response_failed_sse_event as encode_response_failed_sse_event, -}; +use crate::dispatch::{errors::DispatchFailureClass, responses::errors::ResponseDispatchError}; + +use super::responses::response_failed_sse_event as encode_response_failed_sse_event; const NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE: &str = "No active upstream account is available"; const NO_AVAILABLE_RESPONSES_ACCOUNTS_MESSAGE: &str = diff --git a/backend/src/proxy/openai/mod.rs b/backend/src/api/client/mod.rs similarity index 76% rename from backend/src/proxy/openai/mod.rs rename to backend/src/api/client/mod.rs index 3a8350971..ea188f667 100644 --- a/backend/src/proxy/openai/mod.rs +++ b/backend/src/api/client/mod.rs @@ -1,7 +1,8 @@ //! OpenAI 兼容 API 模块。 +pub mod auth; pub mod errors; pub mod models; pub mod responses; -pub mod routes; +pub mod router; pub mod sse; diff --git a/backend/src/proxy/openai/models.rs b/backend/src/api/client/models.rs similarity index 96% rename from backend/src/proxy/openai/models.rs rename to backend/src/api/client/models.rs index 99696dbc9..3d33d5adc 100644 --- a/backend/src/proxy/openai/models.rs +++ b/backend/src/api/client/models.rs @@ -9,8 +9,8 @@ use axum::{ use serde_json::{json, Value}; use crate::{ - proxy::auth::authorize_client_api_key, runtime::state::AppState, - upstream::models::catalog::ModelCatalog, + api::client::auth::authorize_client_api_key, bootstrap::state::AppState, + models::catalog::ModelCatalog, }; use super::errors::{missing_client_api_key_response, model_not_found_response}; diff --git a/backend/src/proxy/openai/responses.rs b/backend/src/api/client/responses.rs similarity index 95% rename from backend/src/proxy/openai/responses.rs rename to backend/src/api/client/responses.rs index 3b002348b..e5cd4b369 100644 --- a/backend/src/proxy/openai/responses.rs +++ b/backend/src/api/client/responses.rs @@ -11,13 +11,10 @@ use axum::{ use serde_json::{json, Map, Value}; use crate::{ - http::middleware::request_id::{ClientIp, RequestId}, - proxy::{ - auth::authorize_client_api_key, - dispatch::responses::{errors::ResponseDispatchError, service::ResponseDispatchStream}, - }, - runtime::state::AppState, - upstream::protocol::{ + api::middleware::request_id::{ClientIp, RequestId}, + bootstrap::state::AppState, + dispatch::responses::{errors::ResponseDispatchError, service::ResponseDispatchStream}, + upstream::openai::protocol::{ responses::{CodexCompactRequest, CodexResponsesRequest}, sse::{encode_sse_event, sse_body_has_done}, }, @@ -126,6 +123,7 @@ pub fn build_compact_request( body, client_ip: None, client_user_agent: user_agent_from_headers(headers), + client_api_key_id: None, } } @@ -304,9 +302,11 @@ async fn handle_responses( route: &str, forced_subagent: Option<&str>, ) -> Response { - if !authorize_client_api_key(&state, &headers).await { + 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(); @@ -322,6 +322,7 @@ async fn handle_responses( }; 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); // 客户端未显式指定 transport 时默认偏好 WebSocket。 if !codex_request.force_http_sse { codex_request.use_websocket = true; @@ -358,9 +359,11 @@ async fn handle_compact_responses( headers: HeaderMap, body: Bytes, ) -> Response { - if !authorize_client_api_key(&state, &headers).await { + 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(); @@ -376,6 +379,7 @@ async fn handle_compact_responses( }; 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 diff --git a/backend/src/proxy/openai/routes.rs b/backend/src/api/client/router.rs similarity index 95% rename from backend/src/proxy/openai/routes.rs rename to backend/src/api/client/router.rs index 8c36db1bb..34f08f073 100644 --- a/backend/src/proxy/openai/routes.rs +++ b/backend/src/api/client/router.rs @@ -6,7 +6,7 @@ use axum::{ Router, }; -use crate::runtime::state::AppState; +use crate::bootstrap::state::AppState; use super::{ models::{model_catalog, model_detail, model_info, models}, diff --git a/backend/src/proxy/openai/sse.rs b/backend/src/api/client/sse.rs similarity index 95% rename from backend/src/proxy/openai/sse.rs rename to backend/src/api/client/sse.rs index c43d6b597..caa9aeb46 100644 --- a/backend/src/proxy/openai/sse.rs +++ b/backend/src/api/client/sse.rs @@ -11,7 +11,7 @@ use axum::{ response::{IntoResponse, Response}, }; -use crate::upstream::protocol::sse::{encode_sse_event, DONE_SSE_FRAME}; +use crate::upstream::openai::protocol::sse::{encode_sse_event, DONE_SSE_FRAME}; use super::errors::openai_error_response; diff --git a/backend/src/http/middleware/mod.rs b/backend/src/api/middleware/mod.rs similarity index 100% rename from backend/src/http/middleware/mod.rs rename to backend/src/api/middleware/mod.rs diff --git a/backend/src/http/middleware/request_id.rs b/backend/src/api/middleware/request_id.rs similarity index 100% rename from backend/src/http/middleware/request_id.rs rename to backend/src/api/middleware/request_id.rs diff --git a/backend/src/http/middleware/trace.rs b/backend/src/api/middleware/trace.rs similarity index 98% rename from backend/src/http/middleware/trace.rs rename to backend/src/api/middleware/trace.rs index ca2fde150..d8d8ded6e 100644 --- a/backend/src/http/middleware/trace.rs +++ b/backend/src/api/middleware/trace.rs @@ -12,7 +12,7 @@ use tower_http::{ }; use tracing::{info, warn, Span}; -use crate::http::middleware::request_id::RequestId; +use crate::api::middleware::request_id::RequestId; /// HTTP tracing layer type. pub type HttpTraceLayer = TraceLayer< diff --git a/backend/src/api/mod.rs b/backend/src/api/mod.rs new file mode 100644 index 000000000..1f4d7da64 --- /dev/null +++ b/backend/src/api/mod.rs @@ -0,0 +1,7 @@ +//! 入站 HTTP API 与静态资源。 + +pub mod admin; +pub mod assets; +pub mod client; +pub mod middleware; +pub mod router; diff --git a/backend/src/http/router.rs b/backend/src/api/router.rs similarity index 54% rename from backend/src/http/router.rs rename to backend/src/api/router.rs index 3ca1025f1..9338c986d 100644 --- a/backend/src/http/router.rs +++ b/backend/src/api/router.rs @@ -4,11 +4,13 @@ use std::path::Path; use axum::{extract::State, http::StatusCode, middleware, routing::get, Router}; -use crate::admin; -use crate::http::middleware::{request_id::attach_request_id, trace::http_trace_layer}; -use crate::proxy; -use crate::runtime::state::AppState; -use crate::web; +use crate::{ + api::{ + admin, assets, client, + middleware::{request_id::attach_request_id, trace::http_trace_layer}, + }, + bootstrap::state::AppState, +}; /// 默认前端构建产物目录。 pub const DEFAULT_ASSET_DIST_DIR: &str = "web/dist"; @@ -22,22 +24,21 @@ pub fn router() -> Router { pub fn router_with_assets(dist_dir: impl AsRef) -> Router { Router::new() .route("/healthz", get(healthz)) - .merge(proxy::router::router()) + .merge(client::router::router()) .merge(admin::router::router()) - .fallback_service(web::assets::spa_router(dist_dir)) + .fallback_service(assets::spa_router(dist_dir)) .layer(http_trace_layer()) .layer(middleware::from_fn(attach_request_id)) } async fn healthz(State(state): State) -> StatusCode { - match sqlx::query("select 1") - .execute(state.services.background_tasks.client_keys.pool()) - .await - { - Ok(_) => StatusCode::NO_CONTENT, - Err(error) => { - tracing::warn!(error = %error, "health check failed"); - StatusCode::SERVICE_UNAVAILABLE - } + if let Err(error) = crate::infra::database::ping(&state.services.database).await { + tracing::warn!(error = %error, "PostgreSQL health check failed"); + return StatusCode::SERVICE_UNAVAILABLE; } + if let Err(error) = state.services.redis.ping().await { + tracing::warn!(error = %error, "Redis health check failed"); + return StatusCode::SERVICE_UNAVAILABLE; + } + StatusCode::NO_CONTENT } diff --git a/backend/src/auth/mod.rs b/backend/src/auth/mod.rs new file mode 100644 index 000000000..2c7dde78e --- /dev/null +++ b/backend/src/auth/mod.rs @@ -0,0 +1,5 @@ +//! 管理员认证与 Redis 会话。 + +pub mod service; +pub mod store; +pub mod types; diff --git a/backend/src/auth/service.rs b/backend/src/auth/service.rs new file mode 100644 index 000000000..b3030b497 --- /dev/null +++ b/backend/src/auth/service.rs @@ -0,0 +1,192 @@ +//! 管理员认证领域服务。 + +use std::{collections::HashMap, sync::Arc}; + +use chrono::{DateTime, Duration, Utc}; +use tokio::sync::Mutex; + +use crate::infra::identity::{ + generate_admin_session_token, hash_admin_password, verify_admin_password, +}; + +use super::{ + store::{PgAdminUserStore, RedisAdminSessionStore}, + types::{LoginSession, SessionError}, +}; + +const LOGIN_FAILURE_WINDOW: Duration = Duration::minutes(15); +const LOGIN_LOCK_DURATION: Duration = Duration::minutes(15); +const LOGIN_MAX_FAILURES: u32 = 5; + +/// 管理员会话服务。 +#[derive(Clone)] +pub struct SessionService { + users: PgAdminUserStore, + sessions: RedisAdminSessionStore, + default_username: String, + session_ttl_minutes: u64, + login_failures: Arc>>, +} + +#[derive(Debug, Clone)] +struct LoginFailureState { + failures: u32, + first_failed_at: DateTime, + locked_until: Option>, +} + +impl SessionService { + /// 构造服务。 + pub fn new( + users: PgAdminUserStore, + sessions: RedisAdminSessionStore, + default_username: String, + session_ttl_minutes: u64, + ) -> Self { + Self { + users, + sessions, + default_username, + session_ttl_minutes, + login_failures: Arc::new(Mutex::new(HashMap::new())), + } + } + + /// 校验管理员会话。 + pub async fn validate(&self, session_id: Option<&str>) -> Result { + let Some(session_id) = session_id else { + return Ok(false); + }; + self.sessions + .validate_session(session_id) + .await + .map_err(|_| SessionError::Validate) + } + + /// 确保默认管理员存在。 + pub async fn ensure_default_admin(&self, password: &str) -> Result { + let password_hash = + hash_admin_password(password).map_err(|_| SessionError::HashPassword)?; + self.users + .ensure_default_admin(&password_hash) + .await + .map_err(|_| SessionError::CreateAdmin) + } + + /// 登录并创建管理员会话。 + pub async fn login( + &self, + source: &str, + username: Option<&str>, + password: &str, + ) -> Result, SessionError> { + let source = login_source_key(source); + self.ensure_login_allowed(&source).await?; + + let username = username.unwrap_or(&self.default_username); + if username != self.default_username.as_str() { + self.record_failed_login(&source).await; + return Ok(None); + } + let Some(admin) = self + .users + .load_first_admin() + .await + .map_err(|_| SessionError::LoadAdmin)? + else { + self.record_failed_login(&source).await; + return Ok(None); + }; + let password_matches = verify_admin_password(password, &admin.password_hash) + .map_err(|_| SessionError::VerifyPassword)?; + if !password_matches { + self.record_failed_login(&source).await; + return Ok(None); + } + self.clear_failed_login(&source).await; + let session_id = generate_admin_session_token(); + let ttl_minutes = self.session_ttl_minutes.min(i64::MAX as u64) as i64; + let expires_at = Utc::now() + Duration::minutes(ttl_minutes); + self.sessions + .create_session( + &session_id, + &admin.id, + Duration::minutes(ttl_minutes.max(1)), + ) + .await + .map_err(|_| SessionError::CreateSession)?; + Ok(Some(LoginSession { + session_id, + expires_at, + })) + } + + /// 删除管理员会话。 + pub async fn delete_session(&self, session_id: &str) -> Result { + self.sessions + .delete_session(session_id) + .await + .map_err(|_| SessionError::DeleteSession) + } + + async fn ensure_login_allowed(&self, source: &str) -> Result<(), SessionError> { + let now = Utc::now(); + let mut failures = self.login_failures.lock().await; + retain_active_login_failures(&mut failures, now); + if failures + .get(source) + .and_then(|state| state.locked_until) + .is_some_and(|locked_until| locked_until > now) + { + return Err(SessionError::LoginThrottled); + } + Ok(()) + } + + async fn record_failed_login(&self, source: &str) { + let now = Utc::now(); + let mut failures = self.login_failures.lock().await; + retain_active_login_failures(&mut failures, now); + let state = failures + .entry(source.to_string()) + .or_insert_with(|| LoginFailureState { + failures: 0, + first_failed_at: now, + locked_until: None, + }); + if now - state.first_failed_at > LOGIN_FAILURE_WINDOW { + state.failures = 0; + state.first_failed_at = now; + state.locked_until = None; + } + state.failures += 1; + if state.failures >= LOGIN_MAX_FAILURES { + state.locked_until = Some(now + LOGIN_LOCK_DURATION); + } + } + + async fn clear_failed_login(&self, source: &str) { + self.login_failures.lock().await.remove(source); + } +} + +fn login_source_key(source: &str) -> String { + let source = source.trim(); + if source.is_empty() { + "unknown".to_string() + } else { + source.to_string() + } +} + +fn retain_active_login_failures( + failures: &mut HashMap, + now: DateTime, +) { + failures.retain(|_, state| { + state + .locked_until + .is_some_and(|locked_until| locked_until > now) + || now - state.first_failed_at <= LOGIN_FAILURE_WINDOW + }); +} diff --git a/backend/src/auth/store.rs b/backend/src/auth/store.rs new file mode 100644 index 000000000..5de3ca9dd --- /dev/null +++ b/backend/src/auth/store.rs @@ -0,0 +1,137 @@ +//! 管理员用户与会话存储。 + +use chrono::{DateTime, Duration, Utc}; +use redis::AsyncCommands; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use thiserror::Error; + +use crate::infra::{identity::hash_credential, redis::RedisConnection}; + +/// 已持久化的管理员用户。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StoredAdminUser { + /// 管理员用户 ID。 + pub id: String, + /// 管理员密码哈希。 + pub password_hash: String, +} + +/// 管理认证存储错误。 +#[derive(Debug, Error)] +pub enum AuthStoreError { + #[error("PostgreSQL admin user operation failed: {0}")] + Postgres(#[from] sqlx::Error), + #[error("Redis admin session operation failed: {0}")] + Redis(#[from] redis::RedisError), + #[error("invalid Redis admin session value: {0}")] + Json(#[from] serde_json::Error), +} + +pub type AuthStoreResult = Result; + +/// PostgreSQL 管理员用户存储。 +#[derive(Clone)] +pub struct PgAdminUserStore { + pool: PgPool, +} + +impl PgAdminUserStore { + /// 构造存储。 + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + /// 如果还没有管理员用户，则创建默认管理员。 + pub async fn ensure_default_admin(&self, password_hash: &str) -> AuthStoreResult { + let now = Utc::now(); + let result = sqlx::query( + r" +insert into admin_users (id, password_hash, created_at, updated_at) +select $1, $2, $3, $3 +where not exists (select 1 from admin_users) +on conflict do nothing", + ) + .bind("admin_1") + .bind(password_hash) + .bind(now) + .execute(&self.pool) + .await?; + + Ok(result.rows_affected() > 0) + } + + /// 读取第一个管理员用户。 + pub async fn load_first_admin(&self) -> AuthStoreResult> { + let row = sqlx::query_as::<_, (String, String)>( + "select id, password_hash from admin_users order by created_at asc, id asc limit 1", + ) + .fetch_optional(&self.pool) + .await?; + + Ok(row.map(|(id, password_hash)| StoredAdminUser { id, password_hash })) + } +} + +/// Redis 管理员会话存储。 +#[derive(Clone)] +pub struct RedisAdminSessionStore { + redis: RedisConnection, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct StoredAdminSession { + user_id: String, + created_at: DateTime, +} + +impl RedisAdminSessionStore { + pub fn new(redis: RedisConnection) -> Self { + Self { redis } + } + + pub async fn create_session( + &self, + session_id: &str, + user_id: &str, + ttl: Duration, + ) -> AuthStoreResult<()> { + let key = self.session_key(session_id); + let value = serde_json::to_string(&StoredAdminSession { + user_id: user_id.to_string(), + created_at: Utc::now(), + })?; + let ttl_seconds = ttl.num_seconds().max(1) as u64; + let mut connection = self.redis.manager(); + let _: () = redis::cmd("SET") + .arg(key) + .arg(value) + .arg("EX") + .arg(ttl_seconds) + .query_async(&mut connection) + .await?; + Ok(()) + } + + pub async fn validate_session(&self, session_id: &str) -> AuthStoreResult { + let mut connection = self.redis.manager(); + let value: Option = connection.get(self.session_key(session_id)).await?; + let Some(value) = value else { + return Ok(false); + }; + let session: StoredAdminSession = serde_json::from_str(&value)?; + Ok(!session.user_id.trim().is_empty()) + } + + pub async fn delete_session(&self, session_id: &str) -> AuthStoreResult { + let mut connection = self.redis.manager(); + let deleted: u64 = connection.del(self.session_key(session_id)).await?; + Ok(deleted > 0) + } + + fn session_key(&self, session_id: &str) -> String { + self.redis + .key(&format!("admin:session:{}", hash_credential(session_id))) + } +} diff --git a/backend/src/auth/types.rs b/backend/src/auth/types.rs new file mode 100644 index 000000000..2515392de --- /dev/null +++ b/backend/src/auth/types.rs @@ -0,0 +1,42 @@ +//! 管理员认证领域类型。 + +use chrono::{DateTime, Utc}; +use thiserror::Error; + +/// 管理员登录成功后的会话。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LoginSession { + /// 会话 ID。 + pub session_id: String, + /// 过期时间。 + pub expires_at: DateTime, +} + +/// 管理员会话错误。 +#[derive(Debug, Error)] +pub enum SessionError { + /// 校验管理员会话失败。 + #[error("failed to validate admin session")] + Validate, + /// 计算管理员密码哈希失败。 + #[error("failed to hash admin password")] + HashPassword, + /// 创建默认管理员失败。 + #[error("failed to create default admin user")] + CreateAdmin, + /// 读取管理员用户失败。 + #[error("failed to load admin user")] + LoadAdmin, + /// 校验管理员密码失败。 + #[error("failed to verify admin password")] + VerifyPassword, + /// 创建管理员会话失败。 + #[error("failed to create admin session")] + CreateSession, + /// 删除管理员会话失败。 + #[error("failed to delete admin session")] + DeleteSession, + /// 管理员登录尝试过于频繁。 + #[error("too many admin login attempts")] + LoginThrottled, +} diff --git a/backend/src/runtime/bootstrap.rs b/backend/src/bootstrap/application.rs similarity index 60% rename from backend/src/runtime/bootstrap.rs rename to backend/src/bootstrap/application.rs index 933d3ab6a..db114b972 100644 --- a/backend/src/runtime/bootstrap.rs +++ b/backend/src/bootstrap/application.rs @@ -1,18 +1,21 @@ use std::{env, error::Error, time::Duration}; use axum::Router; -use chrono::Utc; -use crate::config::schema::AppConfig; -use crate::config::settings::RuntimeSettingsService; -use crate::infra::database::connect_sqlite; +use crate::bootstrap::config::AppConfig; +use crate::bootstrap::services::{ + apply_settings_to_config, fingerprint_from_config, settings_snapshot_from_config, + BackgroundTaskStores, Services, +}; +use crate::bootstrap::shutdown::{restart_executable_path, shutdown_signal}; +use crate::bootstrap::state::{AppState, RuntimeConfig}; +use crate::bootstrap::tasks::coordinator::TaskCoordinator; +use crate::infra::database::connect; use crate::infra::logging::{init_tracing, LogGuard, RotationConfig}; use crate::infra::paths::{ensure_data_dir, load_or_create_installation_id}; -use crate::runtime::services::{BackgroundTaskStores, Services}; -use crate::runtime::shutdown::{restart_executable_path, shutdown_signal}; -use crate::runtime::state::{AppState, RuntimeConfig}; -use crate::runtime::tasks::coordinator::TaskCoordinator; -use crate::upstream::fingerprint::{Fingerprint, FingerprintRepository, RuntimeFingerprint}; +use crate::infra::redis::RedisConnection; +use crate::settings::service::RuntimeSettingsService; +use crate::upstream::openai::fingerprint::{PgFingerprintStore, RuntimeFingerprint}; const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS"; @@ -49,35 +52,41 @@ pub async fn run() -> Result<(), Box> { } async fn build_application( - config: AppConfig, + mut config: AppConfig, ) -> Result<(Router, TaskCoordinator), Box> { - let pool = connect_sqlite(&config.database.url).await?; - let config = RuntimeSettingsService::load_or_initialize_config(config, &pool).await?; + let pool = connect(&config.database.url).await?; + let redis = RedisConnection::connect(&config.redis.url, "cpr").await?; + let settings = + RuntimeSettingsService::load_or_initialize(settings_snapshot_from_config(&config), &pool) + .await?; + apply_settings_to_config(&mut config, &settings); - let fingerprint_repository = FingerprintRepository::new(pool.clone()); + let fingerprint_store = PgFingerprintStore::new(pool.clone()); let stores = BackgroundTaskStores { - accounts: crate::upstream::accounts::store::SqliteAccountStore::new(pool.clone()), - admin_sessions: crate::admin::auth::service::SqliteAdminSessionStore::new(pool.clone()), - cookies: crate::upstream::accounts::cookies::SqliteCookieStore::new(pool.clone()), - fingerprints: fingerprint_repository.clone(), - session_affinity: crate::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new( + redis: redis.clone(), + accounts: crate::accounts::store::PgAccountStore::new(pool.clone()), + admin_users: crate::auth::store::PgAdminUserStore::new(pool.clone()), + admin_sessions: crate::auth::store::RedisAdminSessionStore::new(redis.clone()), + cookies: crate::accounts::cookies::PgCookieStore::new(pool.clone()), + fingerprints: fingerprint_store.clone(), + session_affinity: crate::dispatch::affinity::RedisSessionAffinityStore::new(redis.clone()), + refresh_leases: crate::accounts::refresh::RedisRefreshLeaseStore::new(redis.clone()), + model_snapshots: crate::models::store::RedisModelSnapshotStore::new(redis.clone()), + client_keys: crate::keys::store::PgClientKeyStore::new(pool.clone()), + usage_records: crate::telemetry::usage::store::PgUsageRecordStore::new(pool.clone()), + ops_errors: crate::telemetry::ops::store::PgOpsErrorLogStore::new(pool.clone()), + account_usage: crate::telemetry::account_usage::store::PgAccountUsageStore::new( pool.clone(), ), - refresh_leases: crate::upstream::accounts::token_refresh::RefreshLeaseStore::new( - pool.clone(), - ), - client_keys: crate::admin::keys::service::SqliteClientKeyStore::new(pool.clone()), - usage_records: crate::admin::monitoring::usage_record_store::SqliteUsageRecordStore::new( - pool, - ), + request_buckets: crate::telemetry::buckets::store::PgRequestBucketStore::new(pool), }; config.admin.validate_default_password()?; let default_admin_password = config.admin.default_password.clone(); let data_dir = ensure_data_dir()?; let installation_id = load_or_create_installation_id(Some(&data_dir))?; - let default_fingerprint = Fingerprint::from_config(&config.fingerprint); - let runtime_fingerprint = fingerprint_repository + let default_fingerprint = fingerprint_from_config(&config.fingerprint); + let runtime_fingerprint = fingerprint_store .ensure_current_seed(&default_fingerprint) .await?; let runtime_fingerprint = RuntimeFingerprint::new(runtime_fingerprint); @@ -99,23 +108,17 @@ async fn build_application( "默认管理员账号已完成初始化检查" ); - let restored_accounts = services.account_pool.restore_from_repository().await?; - tracing::info!(count = restored_accounts, "运行时账号池已从 SQLite 恢复"); - - let restored_session_affinities = services - .session_affinity - .restore_from_repository(Utc::now()) - .await?; + let restored_accounts = services.account_pool.restore_from_store().await?; tracing::info!( - count = restored_session_affinities, - "会话亲和性映射已从 SQLite 恢复" + count = restored_accounts, + "运行时账号池已从 PostgreSQL 恢复" ); let task_coordinator = TaskCoordinator::start(&runtime_config, &services); let app_state = AppState { services }; Ok(( - crate::http::router::router().with_state(app_state), + crate::api::router::router().with_state(app_state), task_coordinator, )) } diff --git a/backend/src/config/schema.rs b/backend/src/bootstrap/config.rs similarity index 89% rename from backend/src/config/schema.rs rename to backend/src/bootstrap/config.rs index 65365e81a..ce56f9296 100644 --- a/backend/src/config/schema.rs +++ b/backend/src/bootstrap/config.rs @@ -1,4 +1,4 @@ -use std::collections::BTreeMap; +use std::{collections::BTreeMap, env, path::Path}; use serde::{Deserialize, Serialize}; @@ -32,6 +32,8 @@ pub struct AppConfig { pub quota: QuotaConfig, /// 数据库配置。 pub database: DatabaseConfig, + /// Redis 运行态存储配置。 + pub redis: RedisConfig, /// TLS 偏好配置。 #[serde(default)] pub tls: TlsConfig, @@ -137,7 +139,15 @@ fn default_skip_exhausted() -> bool { /// 数据库连接配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] pub struct DatabaseConfig { - /// SQLite 连接 URL。 + /// PostgreSQL 连接 URL。 + pub url: String, +} + +/// Redis 连接配置。 +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct RedisConfig { + /// Redis 连接 URL。 pub url: String, } @@ -272,12 +282,10 @@ impl FingerprintHeaderConfig { /// 管理员初始化与会话配置。 #[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] pub struct AdminConfig { /// 会话有效期（分钟）。 pub session_ttl_minutes: u64, - /// 会话清理周期（秒）。 - #[serde(default = "default_session_cleanup_interval")] - pub session_cleanup_interval_secs: u64, /// 默认管理员用户名（首次启动时创建） #[serde(default = "default_admin_username")] pub default_username: String, @@ -290,7 +298,6 @@ impl Default for AdminConfig { fn default() -> Self { Self { session_ttl_minutes: 1440, - session_cleanup_interval_secs: default_session_cleanup_interval(), default_username: default_admin_username(), default_password: default_admin_password(), } @@ -321,10 +328,6 @@ pub enum AdminConfigError { WeakDefaultPassword, } -fn default_session_cleanup_interval() -> u64 { - 3600 // 默认每小时清理一次过期会话 -} - fn default_admin_username() -> String { "admin".to_string() } @@ -354,3 +357,30 @@ impl Default for LoggingConfig { } } } + +const CONFIG_FILE_ENV: &str = "CPR_CONFIG_FILE"; + +impl AppConfig { + /// 从 `CPR_CONFIG_FILE` 或当前目录 `config.yaml` 加载配置文件。 + pub fn load() -> Result { + let config_file = env::var(CONFIG_FILE_ENV) + .ok() + .filter(|value| !value.trim().is_empty()); + match config_file { + Some(path) => load_file(path), + None => Self::load_from_dir("."), + } + } + + /// 从指定目录加载配置文件。 + pub fn load_from_dir(config_dir: impl AsRef) -> Result { + load_file(config_dir.as_ref().join("config.yaml")) + } +} + +fn load_file(config_file: impl AsRef) -> Result { + ::config::Config::builder() + .add_source(::config::File::from(config_file.as_ref()).required(true)) + .build()? + .try_deserialize() +} diff --git a/backend/src/bootstrap/import_sqlite.rs b/backend/src/bootstrap/import_sqlite.rs new file mode 100644 index 000000000..0d513ae7a --- /dev/null +++ b/backend/src/bootstrap/import_sqlite.rs @@ -0,0 +1,233 @@ +//! 一次性 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 }, + #[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, +) -> Result { + 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,) = 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 account_model_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, 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, +) -> Result>, ImportSqliteError> { + value + .map(|value| parse_timestamp(table, column, value)) + .transpose() +} + +fn parse_cookie_expiration(value: &str) -> Option> { + 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 { + 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, +) -> Result, ImportSqliteError> { + value + .map(|value| parse_json(table, column, &value)) + .transpose() +} + +fn normalized_dimension(value: Option) -> String { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .unwrap_or_else(|| "__unknown__".to_string()) +} diff --git a/backend/src/bootstrap/import_sqlite/core.rs b/backend/src/bootstrap/import_sqlite/core.rs new file mode 100644 index 000000000..7860d097d --- /dev/null +++ b/backend/src/bootstrap/import_sqlite/core.rs @@ -0,0 +1,491 @@ +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_model_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::("id")?) + .bind(row.try_get::("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::("id")?) + .bind(row.try_get::("name")?) + .bind(row.try_get::("prefix")?) + .bind(row.try_get::("key")?) + .bind(row.try_get::, _>("label")?) + .bind(row.try_get::("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::, _>("admin_api_key")? + .map(|key| 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::("id")?) + .bind(Json(parse_json( + "runtime_settings", + "model_aliases_json", + &aliases, + )?)) + .bind(row.try_get::("refresh_margin_seconds")?) + .bind(row.try_get::("refresh_concurrency")?) + .bind(row.try_get::("max_concurrent_per_account")?) + .bind(row.try_get::("request_interval_ms")?) + .bind(row.try_get::("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::("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::, _>("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::("id")?) + .bind(row.try_get::, _>("email")?) + .bind(row.try_get::, _>("chatgpt_account_id")?) + .bind(row.try_get::, _>("chatgpt_user_id")?) + .bind(row.try_get::, _>("label")?) + .bind(row.try_get::, _>("plan_type")?) + .bind(row.try_get::("access_token")?) + .bind(row.try_get::, _>("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::("quota_limit_reached")? == 1) + .bind(row.try_get::("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::("account_id")?) + .bind(row.try_get::("request_count")?) + .bind(row.try_get::("empty_response_count")?) + .bind(row.try_get::("input_tokens")?) + .bind(row.try_get::("output_tokens")?) + .bind(row.try_get::("cached_tokens")?) + .bind(row.try_get::("reasoning_tokens")?) + .bind(row.try_get::("total_tokens")?) + .bind(row.try_get::("image_input_tokens")?) + .bind(row.try_get::("image_output_tokens")?) + .bind(row.try_get::("image_request_count")?) + .bind(row.try_get::("image_request_failed_count")?) + .bind(row.try_get::("window_request_count")?) + .bind(row.try_get::("window_input_tokens")?) + .bind(row.try_get::("window_output_tokens")?) + .bind(row.try_get::("window_cached_tokens")?) + .bind(row.try_get::("window_image_input_tokens")?) + .bind(row.try_get::("window_image_output_tokens")?) + .bind(row.try_get::("window_image_request_count")?) + .bind(row.try_get::("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::, _>("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_model_usage( + source: &SqlitePool, + target: &mut Transaction<'_, Postgres>, + report: &mut ImportSqliteReport, +) -> Result<(), ImportSqliteError> { + let mut rows = + sqlx::query("select * from account_model_usage order by account_id, model").fetch(source); + while let Some(row) = rows.try_next().await? { + sqlx::query( + "insert into account_model_usage + (account_id, model, request_count, error_count, input_tokens, output_tokens, + cached_tokens, last_used_at) + values ($1, $2, $3, $4, $5, $6, $7, $8)", + ) + .bind(row.try_get::("account_id")?) + .bind(row.try_get::("model")?) + .bind(row.try_get::("request_count")?) + .bind(row.try_get::("error_count")?) + .bind(row.try_get::("input_tokens")?) + .bind(row.try_get::("output_tokens")?) + .bind(row.try_get::("cached_tokens")?) + .bind(parse_optional_timestamp( + "account_model_usage", + "last_used_at", + row.try_get("last_used_at")?, + )?) + .execute(&mut **target) + .await?; + report.add_imported("account_model_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::, _>("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::("id")?) + .bind(row.try_get::("account_id")?) + .bind(row.try_get::("domain")?) + .bind(row.try_get::("name")?) + .bind(row.try_get::("value")?) + .bind(row.try_get::("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::("id")?) + .bind(row.try_get::("originator")?) + .bind(row.try_get::("app_version")?) + .bind(row.try_get::("build_number")?) + .bind(row.try_get::("platform")?) + .bind(row.try_get::("arch")?) + .bind(row.try_get::("chromium_version")?) + .bind(row.try_get::("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::("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::, _>("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::("id")?) + .bind(row.try_get::("current_fingerprint_id")?) + .bind(row.try_get::("app_version")?) + .bind(row.try_get::("build_number")?) + .bind(row.try_get::, _>("chromium_version")?) + .bind(row.try_get::("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): (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)", + ) + .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); + Ok(()) +} diff --git a/backend/src/bootstrap/import_sqlite/telemetry.rs b/backend/src/bootstrap/import_sqlite/telemetry.rs new file mode 100644 index 000000000..389967e42 --- /dev/null +++ b/backend/src/bootstrap/import_sqlite/telemetry.rs @@ -0,0 +1,330 @@ +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 = row.try_get("account_id")?; + let model: Option = row.try_get("model")?; + let status_code: Option = 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::("id")?) + .bind(row.try_get::, _>("request_id")?) + .bind(row.try_get::("kind")?) + .bind(row.try_get::, _>("route")?) + .bind(account_id) + .bind(model) + .bind(requested_model) + .bind(upstream_model) + .bind(service_tier) + .bind(status_code as i32) + .bind(row.try_get::, _>("transport")?) + .bind(row.try_get::, _>("attempt_index")?) + .bind(row.try_get::, _>("response_id")?) + .bind(row.try_get::, _>("upstream_request_id")?) + .bind(row.try_get::, _>("latency_ms")?) + .bind(first_token_ms) + .bind(input_tokens) + .bind(output_tokens) + .bind(cached_tokens) + .bind(reasoning_tokens) + .bind(row.try_get::("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::("id")?) + .bind(row.try_get::, _>("request_id")?) + .bind(row.try_get::("kind")?) + .bind(row.try_get::, _>("account_id")?) + .bind(row.try_get::, _>("route")?) + .bind(row.try_get::, _>("model")?) + .bind(optional_status(row.try_get("status_code")?)) + .bind(optional_status(row.try_get("upstream_status_code")?)) + .bind(row.try_get::, _>("transport")?) + .bind(row.try_get::, _>("attempt_index")?) + .bind(row.try_get::, _>("failure_class")?) + .bind(row.try_get::, _>("response_id")?) + .bind(row.try_get::, _>("upstream_request_id")?) + .bind(row.try_get::, _>("latency_ms")?) + .bind(row.try_get::("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::("id")?) + .bind(row.try_get::, _>("request_id")?) + .bind(row.try_get::("kind")?) + .bind(row.try_get::, _>("account_id")?) + .bind(row.try_get::, _>("route")?) + .bind(row.try_get::, _>("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::, _>("transport")?) + .bind(row.try_get::, _>("attempt_index")?) + .bind(row.try_get::, _>("failure_class")?) + .bind(row.try_get::, _>("response_id")?) + .bind(row.try_get::, _>("upstream_request_id")?) + .bind(row.try_get::, _>("latency_ms")?) + .bind(row.try_get::("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::("input_tokens")?) + .bind(row.try_get::("output_tokens")?) + .bind(row.try_get::("cached_tokens")?) + .bind(row.try_get::("first_token_latency_sum")?) + .bind(row.try_get::("first_token_latency_count")?) + .bind(row.try_get::("latency_sum")?) + .bind(row.try_get::("latency_count")?) + .bind(row.try_get::("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) -> Option { + status.and_then(|status| i32::try_from(status).ok()) +} + +fn metadata_string(metadata: &Value, keys: &[&str]) -> Option { + 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 { + 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 { + 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); + } +} diff --git a/backend/src/bootstrap/mod.rs b/backend/src/bootstrap/mod.rs new file mode 100644 index 000000000..360d32a56 --- /dev/null +++ b/backend/src/bootstrap/mod.rs @@ -0,0 +1,9 @@ +//! 进程装配、共享状态与后台任务。 + +pub mod application; +pub mod config; +pub mod import_sqlite; +pub mod services; +pub mod shutdown; +pub mod state; +pub mod tasks; diff --git a/backend/src/bootstrap/services.rs b/backend/src/bootstrap/services.rs new file mode 100644 index 000000000..39d6a3088 --- /dev/null +++ b/backend/src/bootstrap/services.rs @@ -0,0 +1,422 @@ +//! 运行时服务集合。 + +use std::sync::Arc; + +use crate::{ + accounts::{ + cookies::PgCookieStore, + manage::{AccountManageService, AccountManageServiceParts}, + pool::RuntimeAccountPoolService, + pool::{AccountPoolOptions, RotationStrategy}, + refresh::{ + RedisRefreshLeaseStore, RefreshPolicy, RuntimeRefreshPolicy, RuntimeTokenRefreshService, + }, + store::{AccountStore as AccountStoreTrait, PgAccountStore}, + }, + auth::{ + service::SessionService, + store::{PgAdminUserStore, RedisAdminSessionStore}, + }, + bootstrap::config::AppConfig, + dispatch::{ + affinity::{RedisSessionAffinityStore, RuntimeSessionAffinityService}, + responses::service::{ResponseDispatchService, ResponseDispatchServiceParts}, + }, + infra::redis::RedisConnection, + keys::{manage::KeyManageService, service::KeyVerifier, store::PgClientKeyStore}, + models::{ + config::ModelConfig, + service::{ModelService, ModelServiceError}, + store::{ModelSnapshotStore, RedisModelSnapshotStore}, + }, + settings::{service::RuntimeSettingsService, SettingsSnapshot}, + telemetry::{ + account_usage::query::AccountUsageQueryService, + account_usage::store::PgAccountUsageStore, + buckets::store::PgRequestBucketStore, + ops::query::OpsQueryService, + ops::store::PgOpsErrorLogStore, + usage::query::UsageQueryService, + usage::store::{PgUsageRecordStore, DEFAULT_USAGE_RECORD_CAPTURE_BODY}, + }, + upstream::openai::{ + fingerprint::{Fingerprint, PgFingerprintStore, RuntimeFingerprint}, + token_client::{default_openai_token_client, OpenAiTokenClient, TokenClientConfig}, + transport::{ + build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient, + CodexWebSocketPool, CodexWebSocketPoolConfig, + }, + }, +}; + +/// 运行时热路径内存状态初始化错误。 +#[derive(Debug, thiserror::Error)] +pub enum RuntimeServiceInitializationError { + #[error("failed to initialize model catalog runtime cache: {0}")] + Models(#[from] ModelServiceError), +} + +// ============================================================================ +// BackgroundTaskStores +// ============================================================================ + +/// 后台任务需要的具体存储适配器集合。 +#[derive(Clone)] +pub struct BackgroundTaskStores { + pub redis: RedisConnection, + pub accounts: PgAccountStore, + pub admin_users: PgAdminUserStore, + pub admin_sessions: RedisAdminSessionStore, + pub cookies: PgCookieStore, + pub fingerprints: PgFingerprintStore, + pub session_affinity: RedisSessionAffinityStore, + pub refresh_leases: RedisRefreshLeaseStore, + pub model_snapshots: RedisModelSnapshotStore, + pub client_keys: PgClientKeyStore, + pub usage_records: PgUsageRecordStore, + pub ops_errors: PgOpsErrorLogStore, + pub account_usage: PgAccountUsageStore, + pub request_buckets: PgRequestBucketStore, +} + +// ============================================================================ +// Services struct +// ============================================================================ + +/// 运行时服务集合。 +#[derive(Clone)] +pub struct Services { + pub database: sqlx::PgPool, + pub redis: RedisConnection, + pub models: Arc, + pub accounts: Arc, + pub client_keys: Arc, + pub admin_client_keys: Arc, + pub admin_sessions: Arc, + pub settings: Arc, + pub refresh_policy: RuntimeRefreshPolicy, + pub admin_accounts: Arc, + pub usage_records: Arc, + pub ops_errors: Arc, + pub usage: Arc, + pub account_pool: Arc, + pub token_refresh: Arc>, + pub responses: Arc, + pub session_affinity: Arc, + pub codex: Arc, + pub websocket_pool: Option>, + pub fingerprint: RuntimeFingerprint, + pub installation_id: Option, + pub background_tasks: BackgroundTaskStores, + account_pool_static: AccountPoolStaticSettings, +} + +#[derive(Debug, Clone)] +struct AccountPoolStaticSettings { + skip_quota_limited: bool, + tier_priority: Vec, +} + +/// 使用记录运行选项。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct UsageRecordOptions { + pub enabled: bool, + pub capture_body: bool, +} + +impl UsageRecordOptions { + pub fn from_config(config: &AppConfig) -> Self { + Self { + enabled: config.logging.enabled, + capture_body: DEFAULT_USAGE_RECORD_CAPTURE_BODY, + } + } +} + +impl Services { + pub fn new( + config: &AppConfig, + stores: BackgroundTaskStores, + fingerprint: RuntimeFingerprint, + ) -> Self { + Self::try_new(config, stores, fingerprint) + .expect("failed to build runtime services with configured TLS transport") + } + + pub fn try_new( + config: &AppConfig, + stores: BackgroundTaskStores, + fingerprint: RuntimeFingerprint, + ) -> Result { + Self::try_with_installation_id(config, stores, fingerprint, None) + } + + pub(crate) fn try_with_installation_id( + config: &AppConfig, + stores: BackgroundTaskStores, + fingerprint: RuntimeFingerprint, + installation_id: Option, + ) -> Result { + Self::try_with_installation_id_and_usage_record_options( + config, + stores, + fingerprint, + installation_id, + UsageRecordOptions::from_config(config), + ) + } + + pub fn try_with_installation_id_and_usage_record_options( + config: &AppConfig, + stores: BackgroundTaskStores, + fingerprint: RuntimeFingerprint, + installation_id: Option, + usage_record_options: UsageRecordOptions, + ) -> Result { + let installation_id = installation_id.filter(|id| !id.trim().is_empty()); + let database = stores.accounts.pool().clone(); + let redis = stores.redis.clone(); + let account_store_trait = Arc::new(stores.accounts.clone()) as Arc; + let websocket_pool = config.ws_pool.enabled.then(|| { + Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { + enabled: true, + max_age: std::time::Duration::from_millis(config.ws_pool.max_age_ms), + max_per_account: config.ws_pool.max_per_account, + first_token_timeout: (config.ws_pool.first_token_timeout_ms > 0).then(|| { + std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms) + }), + ..CodexWebSocketPoolConfig::default() + })) + }); + let codex = { + let websocket_first_token_timeout = (config.ws_pool.first_token_timeout_ms > 0) + .then(|| std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms)); + let client = CodexBackendClient::new( + build_reqwest_client(config.tls.force_http11)?, + config.api.base_url.clone(), + fingerprint.clone(), + ) + .with_websocket_first_token_timeout(websocket_first_token_timeout); + if let Some(pool) = &websocket_pool { + Arc::new(client.with_websocket_pool(pool.clone())) + } else { + Arc::new(client) + } + }; + let usage_records = Arc::new(UsageQueryService::new( + stores.usage_records.clone(), + usage_record_options.enabled, + usage_record_options.capture_body, + )); + let ops_errors = Arc::new(OpsQueryService::new( + stores.ops_errors.clone(), + usage_record_options.capture_body, + )); + let account_pool_static = AccountPoolStaticSettings { + skip_quota_limited: config.quota.skip_exhausted, + tier_priority: config.auth.tier_priority.clone(), + }; + let account_pool = Arc::new(RuntimeAccountPoolService::new( + account_store_trait.clone(), + account_pool_options_from_settings( + &settings_snapshot_from_config(config), + &account_pool_static, + ), + config.auth.request_interval_ms, + )); + let refresh_policy = RuntimeRefreshPolicy::new(RefreshPolicy { + refresh_margin_seconds: config.auth.refresh_margin_seconds, + refresh_concurrency: config.auth.refresh_concurrency, + }); + let settings = Arc::new(RuntimeSettingsService::new( + settings_snapshot_from_config(config), + stores.accounts.pool().clone(), + )); + let admin_sessions = Arc::new(SessionService::new( + stores.admin_users.clone(), + stores.admin_sessions.clone(), + config.admin.default_username.clone(), + config.admin.session_ttl_minutes, + )); + let admin_client_keys = Arc::new(KeyManageService::new(stores.client_keys.clone())); + let client_keys = Arc::new(KeyVerifier::new(stores.client_keys.clone())); + let token_client = default_openai_token_client(token_client_config(config)); + let token_refresh = Arc::new( + RuntimeTokenRefreshService::new( + stores.accounts.clone(), + refresh_policy.clone(), + token_client.clone(), + ) + .with_refresh_lease_store(stores.refresh_leases.clone()), + ); + let upstream_client: Arc = codex.clone(); + let store: Arc = Arc::new(stores.model_snapshots.clone()); + let models = Arc::new(ModelService::new( + ModelConfig { + model_aliases: config.model_aliases.clone(), + }, + Some(store), + Some(upstream_client), + )); + let admin_accounts = Arc::new(AccountManageService::new(AccountManageServiceParts { + store: stores.accounts.clone(), + cookies: stores.cookies.clone(), + codex: codex.clone(), + models: models.clone(), + account_pool: account_pool.clone(), + token_refresher: Arc::new(token_client), + refresh_leases: stores.refresh_leases.clone(), + session_affinity: stores.session_affinity.clone(), + oauth: crate::accounts::manage::oauth::AccountOAuthService::new( + reqwest::Client::new(), + config.auth.oauth_client_id.clone(), + config.auth.oauth_token_endpoint.clone(), + ), + refresh_policy: refresh_policy.clone(), + installation_id: installation_id.clone(), + })); + let usage = Arc::new(AccountUsageQueryService::new(stores.account_usage.clone())); + let session_affinity = Arc::new(RuntimeSessionAffinityService::new( + stores.session_affinity.clone(), + )); + + let cloudflare_recovery = + crate::dispatch::cloudflare::CloudflareRecovery::new(stores.cookies.clone()); + let responses = Arc::new(ResponseDispatchService::new(ResponseDispatchServiceParts { + account_pool: account_pool.clone(), + models: models.clone(), + codex: codex.clone(), + session_affinity: session_affinity.clone(), + usage_records: usage_records.clone(), + ops_errors: ops_errors.clone(), + installation_id: installation_id.clone(), + cloudflare: cloudflare_recovery, + })); + + Ok(Self { + database, + redis, + models, + accounts: account_store_trait, + client_keys, + admin_client_keys, + admin_sessions, + settings, + refresh_policy, + admin_accounts, + usage_records, + ops_errors, + usage, + account_pool, + token_refresh, + responses, + session_affinity, + codex, + websocket_pool, + fingerprint, + installation_id, + background_tasks: stores, + account_pool_static, + }) + } + + /// 初始化请求热路径依赖的内存状态。 + pub async fn initialize_hot_path_state(&self) -> Result<(), RuntimeServiceInitializationError> { + self.models.reload_from_store().await?; + Ok(()) + } + + /// 将已持久化的设置快照同步到各运行时消费者。 + pub async fn apply_settings(&self, settings: &SettingsSnapshot) { + self.models.update_config(ModelConfig { + model_aliases: settings.model_aliases.clone(), + }); + self.account_pool + .apply_options( + account_pool_options_from_settings(settings, &self.account_pool_static), + settings.request_interval_ms, + ) + .await; + self.refresh_policy.update(RefreshPolicy { + refresh_margin_seconds: settings.refresh_margin_seconds, + refresh_concurrency: settings.refresh_concurrency, + }); + } +} + +fn token_client_config(config: &AppConfig) -> TokenClientConfig { + TokenClientConfig { + client_id: config.auth.oauth_client_id.clone(), + token_endpoint: config.auth.oauth_token_endpoint.clone(), + } +} + +pub fn settings_snapshot_from_config(config: &AppConfig) -> SettingsSnapshot { + SettingsSnapshot { + model_aliases: config.model_aliases.clone(), + refresh_margin_seconds: config.auth.refresh_margin_seconds, + refresh_concurrency: config.auth.refresh_concurrency, + max_concurrent_per_account: config.auth.max_concurrent_per_account, + request_interval_ms: config.auth.request_interval_ms, + rotation_strategy: config.auth.rotation_strategy.clone(), + } +} + +pub fn apply_settings_to_config(config: &mut AppConfig, settings: &SettingsSnapshot) { + config.model_aliases = settings.model_aliases.clone(); + config.auth.refresh_margin_seconds = settings.refresh_margin_seconds; + config.auth.refresh_concurrency = settings.refresh_concurrency; + config.auth.max_concurrent_per_account = settings.max_concurrent_per_account; + config.auth.request_interval_ms = settings.request_interval_ms; + config.auth.rotation_strategy = settings.rotation_strategy.clone(); +} + +pub fn fingerprint_from_config( + config: &crate::bootstrap::config::FingerprintConfig, +) -> Fingerprint { + Fingerprint { + originator: config.originator.clone(), + app_version: config.app_version.clone(), + build_number: config.build_number.clone(), + platform: config.platform.clone(), + arch: config.arch.clone(), + chromium_version: config.chromium_version.clone(), + user_agent_template: config.user_agent_template.clone(), + default_headers: config + .default_headers + .iter() + .map(|header| (header.name.clone(), header.value.clone())) + .collect(), + header_order: config.header_order.clone(), + updated_at: None, + } +} + +fn account_pool_options_from_settings( + settings: &SettingsSnapshot, + static_settings: &AccountPoolStaticSettings, +) -> AccountPoolOptions { + AccountPoolOptions { + max_concurrent_per_account: settings.max_concurrent_per_account, + rotation_strategy: match settings.rotation_strategy.as_str() { + "smart" => RotationStrategy::Smart, + "quota_reset_priority" => RotationStrategy::QuotaResetPriority, + "round_robin" => RotationStrategy::RoundRobin, + "sticky" => RotationStrategy::Sticky, + _ => RotationStrategy::Smart, + }, + skip_quota_limited: static_settings.skip_quota_limited, + tier_priority: static_settings.tier_priority.clone(), + ..AccountPoolOptions::default() + } +} + +pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions { + account_pool_options_from_settings( + &settings_snapshot_from_config(config), + &AccountPoolStaticSettings { + skip_quota_limited: config.quota.skip_exhausted, + tier_priority: config.auth.tier_priority.clone(), + }, + ) +} diff --git a/backend/src/runtime/shutdown.rs b/backend/src/bootstrap/shutdown.rs similarity index 100% rename from backend/src/runtime/shutdown.rs rename to backend/src/bootstrap/shutdown.rs diff --git a/backend/src/runtime/state.rs b/backend/src/bootstrap/state.rs similarity index 66% rename from backend/src/runtime/state.rs rename to backend/src/bootstrap/state.rs index 5f09d3a83..b35de5a83 100644 --- a/backend/src/runtime/state.rs +++ b/backend/src/bootstrap/state.rs @@ -1,7 +1,7 @@ //! 运行时共享状态。 -use crate::config::schema::AppConfig; -use crate::runtime::services::Services; +use crate::bootstrap::config::AppConfig; +use crate::bootstrap::services::Services; /// HTTP handler 共享的运行时状态。 #[derive(Clone)] @@ -12,9 +12,9 @@ pub struct AppState { /// 运行时配置镜像（从 AppConfig 衍生）。 #[derive(Debug, Clone)] pub struct RuntimeConfig { - pub admin: crate::config::schema::AdminConfig, - pub auth: crate::config::schema::AuthConfig, - pub quota: crate::config::schema::QuotaConfig, + pub admin: crate::bootstrap::config::AdminConfig, + pub auth: crate::bootstrap::config::AuthConfig, + pub quota: crate::bootstrap::config::QuotaConfig, } impl From for RuntimeConfig { diff --git a/backend/src/runtime/tasks/cleanup.rs b/backend/src/bootstrap/tasks/cleanup.rs similarity index 100% rename from backend/src/runtime/tasks/cleanup.rs rename to backend/src/bootstrap/tasks/cleanup.rs diff --git a/backend/src/runtime/tasks/cookie_cleanup.rs b/backend/src/bootstrap/tasks/cookie_cleanup.rs similarity index 77% rename from backend/src/runtime/tasks/cookie_cleanup.rs rename to backend/src/bootstrap/tasks/cookie_cleanup.rs index 21c245395..abcb132b1 100644 --- a/backend/src/runtime/tasks/cookie_cleanup.rs +++ b/backend/src/bootstrap/tasks/cookie_cleanup.rs @@ -4,7 +4,7 @@ use std::future::Future; use chrono::{DateTime, Utc}; -use crate::upstream::accounts::cookies::{SqliteCookieStore, SqliteCookieStoreError}; +use crate::accounts::cookies::{PgCookieStore, PgCookieStoreError}; use super::{ cleanup::{CleanupTask, CleanupTaskMessages, ExpiredCleanupStore}, @@ -22,12 +22,12 @@ const MESSAGES: CleanupTaskMessages = CleanupTaskMessages { /// Cookie 清理任务。 pub struct CookieCleanupTask { - task: CleanupTask, + task: CleanupTask, } impl CookieCleanupTask { /// 构造默认 Cookie 清理任务。 - pub fn new(store: SqliteCookieStore) -> Self { + pub fn new(store: PgCookieStore) -> Self { Self { task: CleanupTask::new(store, DEFAULT_INTERVAL_SECS, MESSAGES), } @@ -39,18 +39,18 @@ impl CookieCleanupTask { } /// 执行一次清理。 - pub async fn cleanup_once(&self) -> Result { + pub async fn cleanup_once(&self) -> Result { self.task.cleanup_once().await } /// 在指定时间点执行一次清理。 - pub async fn cleanup_once_at(&self, now: DateTime) -> Result { + pub async fn cleanup_once_at(&self, now: DateTime) -> Result { self.task.cleanup_once_at(now).await } } -impl ExpiredCleanupStore for SqliteCookieStore { - type Error = SqliteCookieStoreError; +impl ExpiredCleanupStore for PgCookieStore { + type Error = PgCookieStoreError; fn delete_expired_at( &self, diff --git a/backend/src/runtime/tasks/coordinator.rs b/backend/src/bootstrap/tasks/coordinator.rs similarity index 86% rename from backend/src/runtime/tasks/coordinator.rs rename to backend/src/bootstrap/tasks/coordinator.rs index 0bf95b1c2..2299585e0 100644 --- a/backend/src/runtime/tasks/coordinator.rs +++ b/backend/src/bootstrap/tasks/coordinator.rs @@ -12,8 +12,7 @@ use super::{ }, model_refresh::ModelRefreshTask, quota_refresh::QuotaRefreshTask, - session_affinity_cleanup::SessionAffinityCleanupTask, - session_cleanup::SessionCleanupTask, + retention_trim::RetentionTrimTask, token_refresh::TokenRefreshTask, }; @@ -31,8 +30,8 @@ impl TaskCoordinator { /// 启动所有后台任务。 pub(crate) fn start( - config: &crate::runtime::state::RuntimeConfig, - services: &crate::runtime::services::Services, + config: &crate::bootstrap::state::RuntimeConfig, + services: &crate::bootstrap::services::Services, ) -> Self { let mut coordinator = TaskCoordinator::default(); let stores = &services.background_tasks; @@ -42,18 +41,11 @@ impl TaskCoordinator { CookieCleanupTask::new(stores.cookies.clone()).start(), ); coordinator.push( - "session_cleanup", - SessionCleanupTask::new( - stores.admin_sessions.clone(), - config.admin.session_cleanup_interval_secs, - ) - .start(), - ); - coordinator.push( - "session_affinity_cleanup", - SessionAffinityCleanupTask::new( - stores.session_affinity.clone(), - config.admin.session_cleanup_interval_secs, + "retention_trim", + RetentionTrimTask::new( + stores.usage_records.clone(), + stores.ops_errors.clone(), + stores.request_buckets.clone(), ) .start(), ); @@ -126,7 +118,7 @@ pub enum SchedulerHandle { handle: JoinHandle<()>, }, /// 关闭上游 WebSocket 连接池。 - WebSocketPool(Arc), + WebSocketPool(Arc), } impl SchedulerHandle { @@ -139,7 +131,9 @@ impl SchedulerHandle { } /// 使用 WebSocket 连接池构造句柄。 - pub fn from_websocket_pool(pool: Arc) -> Self { + pub fn from_websocket_pool( + pool: Arc, + ) -> Self { Self::WebSocketPool(pool) } diff --git a/backend/src/runtime/tasks/fingerprint_update.rs b/backend/src/bootstrap/tasks/fingerprint_update.rs similarity index 93% rename from backend/src/runtime/tasks/fingerprint_update.rs rename to backend/src/bootstrap/tasks/fingerprint_update.rs index d88809ef0..452ee3ba8 100644 --- a/backend/src/runtime/tasks/fingerprint_update.rs +++ b/backend/src/bootstrap/tasks/fingerprint_update.rs @@ -4,8 +4,8 @@ use std::path::PathBuf; use tracing::{info, warn}; -use crate::upstream::fingerprint::{ - FingerprintRepository, RuntimeFingerprint, UpdateChecker, APPCAST_POLL_INTERVAL, +use crate::upstream::openai::fingerprint::{ + PgFingerprintStore, RuntimeFingerprint, UpdateChecker, APPCAST_POLL_INTERVAL, }; use super::{ @@ -21,7 +21,7 @@ pub const DEFAULT_EXTRACTED_FINGERPRINT_PATH: &str = "data/extracted-fingerprint /// 启动指纹自动更新后台任务。 pub fn start_fingerprint_update_task( - repository: FingerprintRepository, + store: PgFingerprintStore, runtime_fingerprint: RuntimeFingerprint, appcast_url: String, extracted_fingerprint_path: PathBuf, @@ -29,7 +29,7 @@ pub fn start_fingerprint_update_task( current_build: String, ) -> SchedulerHandle { let update_checker = UpdateChecker::with_client( - repository, + store, reqwest::Client::new(), appcast_url, extracted_fingerprint_path, diff --git a/backend/src/runtime/tasks/mod.rs b/backend/src/bootstrap/tasks/mod.rs similarity index 77% rename from backend/src/runtime/tasks/mod.rs rename to backend/src/bootstrap/tasks/mod.rs index 62395806f..190a7567e 100644 --- a/backend/src/runtime/tasks/mod.rs +++ b/backend/src/bootstrap/tasks/mod.rs @@ -7,6 +7,5 @@ pub mod fingerprint_update; pub mod model_refresh; mod periodic; pub mod quota_refresh; -pub mod session_affinity_cleanup; -pub mod session_cleanup; +pub mod retention_trim; pub mod token_refresh; diff --git a/backend/src/runtime/tasks/model_refresh.rs b/backend/src/bootstrap/tasks/model_refresh.rs similarity index 98% rename from backend/src/runtime/tasks/model_refresh.rs rename to backend/src/bootstrap/tasks/model_refresh.rs index c97aaf429..68ada6f1e 100644 --- a/backend/src/runtime/tasks/model_refresh.rs +++ b/backend/src/bootstrap/tasks/model_refresh.rs @@ -7,8 +7,8 @@ use tokio::sync::mpsc; use tokio::time::sleep; use tracing::{info, warn}; -use crate::upstream::accounts::pool::RuntimeAccountPoolService; -use crate::upstream::models::service::{ +use crate::accounts::pool::RuntimeAccountPoolService; +use crate::models::service::{ ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError, }; diff --git a/backend/src/runtime/tasks/periodic.rs b/backend/src/bootstrap/tasks/periodic.rs similarity index 100% rename from backend/src/runtime/tasks/periodic.rs rename to backend/src/bootstrap/tasks/periodic.rs diff --git a/backend/src/runtime/tasks/quota_refresh.rs b/backend/src/bootstrap/tasks/quota_refresh.rs similarity index 92% rename from backend/src/runtime/tasks/quota_refresh.rs rename to backend/src/bootstrap/tasks/quota_refresh.rs index ff6fdcc94..2944c4373 100644 --- a/backend/src/runtime/tasks/quota_refresh.rs +++ b/backend/src/bootstrap/tasks/quota_refresh.rs @@ -6,13 +6,13 @@ use tokio::time::Instant; use tracing::{debug, info, warn}; use crate::{ - upstream::accounts::{ - cookies::SqliteCookieStore, + accounts::{ + cookies::PgCookieStore, pool::RuntimeAccountPoolService, quota::{QuotaRefreshSummary, RuntimeQuotaRefreshService}, - store::SqliteAccountStore, + store::PgAccountStore, }, - upstream::transport::CodexBackendClient, + upstream::openai::transport::CodexBackendClient, }; use super::{ @@ -30,7 +30,7 @@ pub struct QuotaRefreshTask { impl QuotaRefreshTask { /// 使用自定义间隔构造配额刷新后台任务。 pub fn with_intervals( - store: SqliteAccountStore, + store: PgAccountStore, codex: Arc, interval_secs: u64, min_refresh_interval_secs: u64, @@ -53,7 +53,7 @@ impl QuotaRefreshTask { } /// 设置 usage 请求可复用的账号 Cookie 存储。 - pub fn with_cookie_store(mut self, cookie_store: SqliteCookieStore) -> Self { + pub fn with_cookie_store(mut self, cookie_store: PgCookieStore) -> Self { self.service = self.service.with_cookie_store(cookie_store); self } diff --git a/backend/src/bootstrap/tasks/retention_trim.rs b/backend/src/bootstrap/tasks/retention_trim.rs new file mode 100644 index 000000000..96b30605e --- /dev/null +++ b/backend/src/bootstrap/tasks/retention_trim.rs @@ -0,0 +1,70 @@ +//! PostgreSQL 增长表保留期清理任务。 + +use chrono::Utc; +use tracing::{info, warn}; + +use crate::telemetry::{ + buckets::store::PgRequestBucketStore, ops::store::PgOpsErrorLogStore, + usage::store::PgUsageRecordStore, +}; + +use super::{ + coordinator::SchedulerHandle, + periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, TaskFuture}, +}; + +const RETENTION_TRIM_INTERVAL_SECS: u64 = 60 * 60; + +pub struct RetentionTrimTask { + usage_records: PgUsageRecordStore, + ops_errors: PgOpsErrorLogStore, + buckets: PgRequestBucketStore, +} + +impl RetentionTrimTask { + pub fn new( + usage_records: PgUsageRecordStore, + ops_errors: PgOpsErrorLogStore, + buckets: PgRequestBucketStore, + ) -> Self { + Self { + usage_records, + ops_errors, + buckets, + } + } + + pub fn start(self) -> SchedulerHandle { + spawn_periodic_task( + self, + PeriodicTaskConfig::new( + RETENTION_TRIM_INTERVAL_SECS, + "PostgreSQL 保留期清理任务已启动", + "PostgreSQL 保留期清理任务已关闭", + ), + ) + } +} + +impl PeriodicTaskRunner for RetentionTrimTask { + fn tick(&mut self) -> TaskFuture<'_, ()> { + Box::pin(async move { + let now = Utc::now(); + match self.usage_records.trim_to_retention(now).await { + Ok(deleted) if deleted > 0 => info!(deleted, "已清理过期成功使用事实"), + Ok(_) => {} + Err(error) => warn!(error = %error, "清理成功使用事实失败"), + } + match self.ops_errors.trim_to_retention(now).await { + Ok(deleted) if deleted > 0 => info!(deleted, "已清理过期运维错误事实"), + Ok(_) => {} + Err(error) => warn!(error = %error, "清理运维错误事实失败"), + } + match self.buckets.trim_to_retention(now).await { + Ok(deleted) if deleted > 0 => info!(deleted, "已清理过期请求时间桶"), + Ok(_) => {} + Err(error) => warn!(error = %error, "清理请求时间桶失败"), + } + }) + } +} diff --git a/backend/src/runtime/tasks/token_refresh.rs b/backend/src/bootstrap/tasks/token_refresh.rs similarity index 89% rename from backend/src/runtime/tasks/token_refresh.rs rename to backend/src/bootstrap/tasks/token_refresh.rs index 7188d6b64..e57cd616b 100644 --- a/backend/src/runtime/tasks/token_refresh.rs +++ b/backend/src/bootstrap/tasks/token_refresh.rs @@ -2,12 +2,11 @@ use tracing::{debug, info, warn}; -use crate::upstream::accounts::{ - store::SqliteAccountStore, - token_refresh::{ - RefreshLeaseStore, RuntimeRefreshPolicy, RuntimeTokenRefreshService, TokenRefresher, - }, +use crate::accounts::{ + refresh::{RedisRefreshLeaseStore, RuntimeRefreshPolicy, RuntimeTokenRefreshService}, + store::PgAccountStore, }; +use crate::upstream::openai::token_client::TokenRefresher; use super::{ coordinator::SchedulerHandle, @@ -30,7 +29,7 @@ where C: TokenRefresher, { /// 构造默认后台任务。 - pub fn new(store: SqliteAccountStore, policy: RuntimeRefreshPolicy, client: C) -> Self { + pub fn new(store: PgAccountStore, policy: RuntimeRefreshPolicy, client: C) -> Self { Self { service: RuntimeTokenRefreshService::new(store, policy, client), interval_secs: DEFAULT_INTERVAL_SECS, @@ -46,7 +45,7 @@ where } /// 使用刷新租约存储保护账号刷新。 - pub fn with_refresh_lease_store(mut self, refresh_leases: RefreshLeaseStore) -> Self { + pub fn with_refresh_lease_store(mut self, refresh_leases: RedisRefreshLeaseStore) -> Self { self.service = self.service.with_refresh_lease_store(refresh_leases); self } diff --git a/backend/src/config/loader.rs b/backend/src/config/loader.rs deleted file mode 100644 index 154729411..000000000 --- a/backend/src/config/loader.rs +++ /dev/null @@ -1,31 +0,0 @@ -use std::{env, path::Path}; - -use super::schema::AppConfig; - -const CONFIG_FILE_ENV: &str = "CPR_CONFIG_FILE"; - -impl AppConfig { - /// 从 `CPR_CONFIG_FILE` 或当前目录 `config.yaml` 加载配置文件。 - pub fn load() -> Result { - let config_file = env::var(CONFIG_FILE_ENV) - .ok() - .filter(|value| !value.trim().is_empty()); - match config_file { - Some(path) => load_file(path), - None => Self::load_from_dir("."), - } - } - - /// 从指定目录加载配置文件。 - pub fn load_from_dir(config_dir: impl AsRef) -> Result { - let config_dir = config_dir.as_ref(); - load_file(config_dir.join("config.yaml")) - } -} - -fn load_file(config_file: impl AsRef) -> Result { - ::config::Config::builder() - .add_source(::config::File::from(config_file.as_ref()).required(true)) - .build()? - .try_deserialize() -} diff --git a/backend/src/config/mod.rs b/backend/src/config/mod.rs deleted file mode 100644 index 013acb1c6..000000000 --- a/backend/src/config/mod.rs +++ /dev/null @@ -1,10 +0,0 @@ -//! 配置加载与配置 schema。 - -#![deny(missing_docs)] - -/// 配置加载器。 -pub mod loader; -/// 配置文件 schema 与默认值。 -pub mod schema; -/// 管理端可变设置与运行时写回服务。 -pub mod settings; diff --git a/backend/src/config/settings.rs b/backend/src/config/settings.rs deleted file mode 100644 index 4c76c2d58..000000000 --- a/backend/src/config/settings.rs +++ /dev/null @@ -1,601 +0,0 @@ -//! 配置设置领域逻辑与运行时数据库服务。 - -use std::{ - collections::BTreeMap, - sync::{Arc, RwLock as StdRwLock}, -}; - -use chrono::Utc; -use serde::{Deserialize, Serialize}; -use sqlx::{sqlite::SqliteRow, Row, SqlitePool}; -use subtle::ConstantTimeEq; -use thiserror::Error; - -use crate::config::schema::AppConfig; -use crate::infra::identity::generate_admin_api_key; -use crate::upstream::accounts::{ - pool::{AccountPoolOptions, RotationStrategy, RuntimeAccountPoolService}, - token_refresh::{RefreshPolicy, RuntimeRefreshPolicy}, -}; - -const ROTATION_STRATEGIES: [&str; 4] = ["smart", "quota_reset_priority", "round_robin", "sticky"]; -const RUNTIME_SETTINGS_ID: i64 = 1; - -/// 管理端可变设置。 -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct AdminSettings { - /// 模型别名映射。 - pub model_aliases: BTreeMap, - /// 访问令牌过期前多少秒开始刷新。 - pub refresh_margin_seconds: u64, - /// 访问令牌刷新并发数。 - pub refresh_concurrency: u32, - /// 单账号最大并发请求数。 - pub max_concurrent_per_account: usize, - /// 同账号请求间隔毫秒数。 - pub request_interval_ms: u64, - /// 账号轮换策略。 - pub rotation_strategy: String, -} - -impl Default for AdminSettings { - fn default() -> Self { - Self { - model_aliases: BTreeMap::new(), - refresh_margin_seconds: 3600, - refresh_concurrency: 2, - max_concurrent_per_account: 3, - request_interval_ms: 50, - rotation_strategy: "smart".to_string(), - } - } -} - -/// 管理端设置补丁。 -#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminSettingsPatch { - /// 模型别名映射。 - pub model_aliases: Option>, - /// 访问令牌过期前多少秒开始刷新。 - pub refresh_margin_seconds: Option, - /// 访问令牌刷新并发数。 - pub refresh_concurrency: Option, - /// 单账号最大并发请求数。 - pub max_concurrent_per_account: Option, - /// 同账号请求间隔毫秒数。 - pub request_interval_ms: Option, - /// 账号轮换策略。 - pub rotation_strategy: Option, -} - -/// 设置领域服务。 -#[derive(Debug, Clone, Default)] -pub struct SettingsService; - -impl SettingsService { - /// 将管理端设置补丁应用到当前设置。 - pub fn apply_patch( - current: &mut AdminSettings, - patch: AdminSettingsPatch, - ) -> Result<(), SettingsServiceError> { - if let Some(model_aliases) = patch.model_aliases { - current.model_aliases = validate_model_aliases(model_aliases)?; - } - if let Some(refresh_margin_seconds) = patch.refresh_margin_seconds { - current.refresh_margin_seconds = - positive_u64("refreshMarginSeconds", refresh_margin_seconds)?; - } - if let Some(refresh_concurrency) = patch.refresh_concurrency { - current.refresh_concurrency = positive_u32("refreshConcurrency", refresh_concurrency)?; - } - if let Some(max_concurrent_per_account) = patch.max_concurrent_per_account { - current.max_concurrent_per_account = - positive_usize("maxConcurrentPerAccount", max_concurrent_per_account)?; - } - if let Some(request_interval_ms) = patch.request_interval_ms { - current.request_interval_ms = request_interval_ms; - } - if let Some(rotation_strategy) = patch.rotation_strategy { - current.rotation_strategy = validate_rotation_strategy(&rotation_strategy)?; - } - Ok(()) - } -} - -/// 设置领域错误。 -#[derive(Debug, Clone, PartialEq, Eq, Error)] -pub enum SettingsServiceError { - /// 字段值无效。 - #[error("invalid setting `{field}`: {message}")] - InvalidField { - /// 字段名。 - field: String, - /// 错误说明。 - message: String, - }, -} - -impl SettingsServiceError { - /// 返回无效字段名。 - pub fn field(&self) -> &str { - match self { - Self::InvalidField { field, .. } => field, - } - } -} - -fn validate_model_aliases( - aliases: BTreeMap, -) -> Result, SettingsServiceError> { - let mut normalized = BTreeMap::new(); - for (alias, target) in aliases { - let alias = non_empty_string("modelAliases", &alias)?; - let target = non_empty_string("modelAliases", &target)?; - if alias == target { - return Err(invalid_field( - "modelAliases", - "alias and target must differ", - )); - } - normalized.insert(alias, target); - } - Ok(normalized) -} - -fn validate_rotation_strategy(strategy: &str) -> Result { - let strategy = non_empty_string("rotationStrategy", strategy)?; - if ROTATION_STRATEGIES.contains(&strategy.as_str()) { - Ok(strategy) - } else { - Err(invalid_field( - "rotationStrategy", - "must be one of smart, quota_reset_priority, round_robin, sticky", - )) - } -} - -fn non_empty_string(field: &str, value: &str) -> Result { - let value = value.trim().to_string(); - if value.is_empty() { - Err(invalid_field(field, "must not be empty")) - } else { - Ok(value) - } -} - -fn positive_u64(field: &str, value: u64) -> Result { - if value == 0 { - Err(invalid_field(field, "must be greater than 0")) - } else { - Ok(value) - } -} - -fn positive_u32(field: &str, value: u32) -> Result { - if value == 0 { - Err(invalid_field(field, "must be greater than 0")) - } else { - Ok(value) - } -} - -fn positive_usize(field: &str, value: usize) -> Result { - if value == 0 { - Err(invalid_field(field, "must be greater than 0")) - } else { - Ok(value) - } -} - -fn invalid_field(field: &str, message: impl Into) -> SettingsServiceError { - SettingsServiceError::InvalidField { - field: field.to_string(), - message: message.into(), - } -} - -/// 管理员 API Key 状态。 -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -#[serde(rename_all = "camelCase")] -pub struct AdminApiKeyStatus { - /// 是否已经生成。 - pub exists: bool, - /// 脱敏后的密钥。 - pub masked_key: Option, -} - -/// 运行时设置服务。 -#[derive(Clone)] -pub struct RuntimeSettingsService { - current: Arc>>, - pool: SqlitePool, - account_pool: Arc, - refresh_policy: RuntimeRefreshPolicy, -} - -impl RuntimeSettingsService { - /// 构造运行时设置服务。 - pub fn new( - config: AppConfig, - pool: SqlitePool, - account_pool: Arc, - refresh_policy: RuntimeRefreshPolicy, - ) -> Self { - Self { - current: Arc::new(StdRwLock::new(Arc::new(config))), - pool, - account_pool, - refresh_policy, - } - } - - /// 返回当前运行时配置快照。 - pub fn current(&self) -> Arc { - self.current - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone() - } - - /// 初始化缺省运行设置，并返回数据库中的运行配置快照。 - pub async fn load_or_initialize_config( - mut config: AppConfig, - pool: &SqlitePool, - ) -> Result { - let settings = admin_settings_from_config(&config); - insert_runtime_settings_if_missing(pool, &settings).await?; - let settings = load_runtime_settings(pool).await?; - apply_admin_settings_to_config(&mut config, settings); - Ok(config) - } - - /// 应用设置补丁、写入数据库并更新运行时配置快照。 - pub async fn update( - &self, - patch: AdminSettingsPatch, - ) -> Result, RuntimeSettingsError> { - let mut next = (*self.current()).clone(); - let mut settings = admin_settings_from_config(&next); - SettingsService::apply_patch(&mut settings, patch)?; - save_runtime_settings(&self.pool, &settings).await?; - apply_admin_settings_to_config(&mut next, settings); - self.account_pool - .apply_options( - account_pool_options_from_config(&next), - next.auth.request_interval_ms, - ) - .await; - self.refresh_policy.update(RefreshPolicy { - refresh_margin_seconds: next.auth.refresh_margin_seconds, - refresh_concurrency: next.auth.refresh_concurrency, - }); - let next = Arc::new(next); - *self - .current - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) = next.clone(); - Ok(next) - } - - /// 返回管理员 API Key 状态。 - pub async fn admin_api_key_status(&self) -> Result { - admin_api_key_status(&self.pool).await - } - - /// 重新生成管理员 API Key，完整值只在本次返回。 - pub async fn regenerate_admin_api_key(&self) -> Result { - let settings = admin_settings_from_config(&self.current()); - insert_runtime_settings_if_missing(&self.pool, &settings).await?; - - let key = generate_admin_api_key(); - sqlx::query( - r" -update runtime_settings -set admin_api_key = ?, updated_at = ? -where id = ?", - ) - .bind(&key) - .bind(Utc::now().to_rfc3339()) - .bind(RUNTIME_SETTINGS_ID) - .execute(&self.pool) - .await?; - - Ok(key) - } - - /// 删除管理员 API Key。 - pub async fn delete_admin_api_key(&self) -> Result<(), RuntimeSettingsError> { - sqlx::query( - r" -update runtime_settings -set admin_api_key = null, updated_at = ? -where id = ?", - ) - .bind(Utc::now().to_rfc3339()) - .bind(RUNTIME_SETTINGS_ID) - .execute(&self.pool) - .await?; - Ok(()) - } - - /// 校验管理员 API Key。 - pub async fn verify_admin_api_key(&self, key: &str) -> Result { - if key.is_empty() { - return Ok(false); - } - - let stored = load_admin_api_key(&self.pool).await?; - Ok(stored - .as_deref() - .filter(|stored_key| !stored_key.is_empty()) - .is_some_and(|stored_key| key.as_bytes().ct_eq(stored_key.as_bytes()).into())) - } -} - -/// 运行时设置错误。 -#[derive(Debug, Error)] -pub enum RuntimeSettingsError { - /// 设置字段校验失败。 - #[error(transparent)] - InvalidField(#[from] SettingsServiceError), - /// 数据库操作失败。 - #[error("runtime settings database error: {0}")] - Database(#[from] sqlx::Error), - /// JSON 编解码失败。 - #[error("runtime settings json error: {0}")] - Json(#[from] serde_json::Error), - /// 数据库存储值非法。 - #[error("invalid stored setting `{field}`: {message}")] - StoredField { - /// 字段名。 - field: String, - /// 错误说明。 - message: String, - }, -} - -fn admin_settings_from_config(config: &AppConfig) -> AdminSettings { - AdminSettings { - model_aliases: config.model_aliases.clone(), - refresh_margin_seconds: config.auth.refresh_margin_seconds, - refresh_concurrency: config.auth.refresh_concurrency, - max_concurrent_per_account: config.auth.max_concurrent_per_account, - request_interval_ms: config.auth.request_interval_ms, - rotation_strategy: config.auth.rotation_strategy.clone(), - } -} - -fn apply_admin_settings_to_config(config: &mut AppConfig, settings: AdminSettings) { - config.model_aliases = settings.model_aliases; - config.auth.refresh_margin_seconds = settings.refresh_margin_seconds; - config.auth.refresh_concurrency = settings.refresh_concurrency; - config.auth.max_concurrent_per_account = settings.max_concurrent_per_account; - config.auth.request_interval_ms = settings.request_interval_ms; - config.auth.rotation_strategy = settings.rotation_strategy; -} - -async fn insert_runtime_settings_if_missing( - pool: &SqlitePool, - settings: &AdminSettings, -) -> Result<(), RuntimeSettingsError> { - let model_aliases_json = serde_json::to_string(&settings.model_aliases)?; - sqlx::query( - r" -insert or ignore into runtime_settings ( - id, - model_aliases_json, - refresh_margin_seconds, - refresh_concurrency, - max_concurrent_per_account, - request_interval_ms, - rotation_strategy, - updated_at -) values (?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(RUNTIME_SETTINGS_ID) - .bind(model_aliases_json) - .bind( - i64::try_from(settings.refresh_margin_seconds) - .map_err(|_| stored_field_error("refreshMarginSeconds", "out of range"))?, - ) - .bind(i64::from(settings.refresh_concurrency)) - .bind( - i64::try_from(settings.max_concurrent_per_account) - .map_err(|_| stored_field_error("maxConcurrentPerAccount", "out of range"))?, - ) - .bind( - i64::try_from(settings.request_interval_ms) - .map_err(|_| stored_field_error("requestIntervalMs", "out of range"))?, - ) - .bind(&settings.rotation_strategy) - .bind(Utc::now().to_rfc3339()) - .execute(pool) - .await?; - Ok(()) -} - -async fn save_runtime_settings( - pool: &SqlitePool, - settings: &AdminSettings, -) -> Result<(), RuntimeSettingsError> { - let model_aliases_json = serde_json::to_string(&settings.model_aliases)?; - sqlx::query( - r" -insert into runtime_settings ( - id, - model_aliases_json, - refresh_margin_seconds, - refresh_concurrency, - max_concurrent_per_account, - request_interval_ms, - rotation_strategy, - updated_at -) values (?, ?, ?, ?, ?, ?, ?, ?) -on conflict(id) do update set - model_aliases_json = excluded.model_aliases_json, - refresh_margin_seconds = excluded.refresh_margin_seconds, - refresh_concurrency = excluded.refresh_concurrency, - max_concurrent_per_account = excluded.max_concurrent_per_account, - request_interval_ms = excluded.request_interval_ms, - rotation_strategy = excluded.rotation_strategy, - updated_at = excluded.updated_at", - ) - .bind(RUNTIME_SETTINGS_ID) - .bind(model_aliases_json) - .bind( - i64::try_from(settings.refresh_margin_seconds) - .map_err(|_| stored_field_error("refreshMarginSeconds", "out of range"))?, - ) - .bind(i64::from(settings.refresh_concurrency)) - .bind( - i64::try_from(settings.max_concurrent_per_account) - .map_err(|_| stored_field_error("maxConcurrentPerAccount", "out of range"))?, - ) - .bind( - i64::try_from(settings.request_interval_ms) - .map_err(|_| stored_field_error("requestIntervalMs", "out of range"))?, - ) - .bind(&settings.rotation_strategy) - .bind(Utc::now().to_rfc3339()) - .execute(pool) - .await?; - Ok(()) -} - -async fn load_runtime_settings(pool: &SqlitePool) -> Result { - let row = sqlx::query( - r" -select - model_aliases_json, - refresh_margin_seconds, - refresh_concurrency, - max_concurrent_per_account, - request_interval_ms, - rotation_strategy -from runtime_settings -where id = ?", - ) - .bind(RUNTIME_SETTINGS_ID) - .fetch_one(pool) - .await?; - runtime_settings_from_row(&row) -} - -async fn admin_api_key_status( - pool: &SqlitePool, -) -> Result { - let key = load_admin_api_key(pool).await?; - let masked_key = key - .as_deref() - .filter(|key| !key.is_empty()) - .map(mask_api_key); - Ok(AdminApiKeyStatus { - exists: masked_key.is_some(), - masked_key, - }) -} - -async fn load_admin_api_key(pool: &SqlitePool) -> Result, RuntimeSettingsError> { - Ok(sqlx::query_scalar::<_, Option>( - "select admin_api_key from runtime_settings where id = ?", - ) - .bind(RUNTIME_SETTINGS_ID) - .fetch_optional(pool) - .await? - .flatten()) -} - -fn mask_api_key(key: &str) -> String { - if key.len() > 14 { - format!("{}...{}", &key[..10], &key[key.len() - 4..]) - } else { - key.to_string() - } -} - -fn runtime_settings_from_row(row: &SqliteRow) -> Result { - let model_aliases_json: String = row.get("model_aliases_json"); - Ok(AdminSettings { - model_aliases: serde_json::from_str(&model_aliases_json)?, - refresh_margin_seconds: positive_i64_to_u64( - "refreshMarginSeconds", - row.get("refresh_margin_seconds"), - )?, - refresh_concurrency: positive_i64_to_u32( - "refreshConcurrency", - row.get("refresh_concurrency"), - )?, - max_concurrent_per_account: positive_i64_to_usize( - "maxConcurrentPerAccount", - row.get("max_concurrent_per_account"), - )?, - request_interval_ms: non_negative_i64_to_u64( - "requestIntervalMs", - row.get("request_interval_ms"), - )?, - rotation_strategy: validate_rotation_strategy(&row.get::("rotation_strategy"))?, - }) -} - -/// 从当前运行配置生成账号池调度参数。 -pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOptions { - AccountPoolOptions { - max_concurrent_per_account: config.auth.max_concurrent_per_account, - rotation_strategy: rotation_strategy_from_config(&config.auth.rotation_strategy), - skip_quota_limited: config.quota.skip_exhausted, - tier_priority: config.auth.tier_priority.clone(), - ..AccountPoolOptions::default() - } -} - -fn rotation_strategy_from_config(strategy: &str) -> RotationStrategy { - match strategy { - "smart" => RotationStrategy::Smart, - "quota_reset_priority" => RotationStrategy::QuotaResetPriority, - "round_robin" => RotationStrategy::RoundRobin, - "sticky" => RotationStrategy::Sticky, - _ => RotationStrategy::Smart, - } -} - -fn positive_i64_to_u64(field: &'static str, value: i64) -> Result { - if value <= 0 { - return Err(stored_field_error(field, "must be greater than 0")); - } - u64::try_from(value).map_err(|_| stored_field_error(field, "out of range")) -} - -fn positive_i64_to_u32(field: &'static str, value: i64) -> Result { - if value <= 0 { - return Err(stored_field_error(field, "must be greater than 0")); - } - u32::try_from(value).map_err(|_| stored_field_error(field, "out of range")) -} - -fn positive_i64_to_usize(field: &'static str, value: i64) -> Result { - if value <= 0 { - return Err(stored_field_error(field, "must be greater than 0")); - } - usize::try_from(value).map_err(|_| stored_field_error(field, "out of range")) -} - -fn non_negative_i64_to_u64(field: &'static str, value: i64) -> Result { - if value < 0 { - return Err(stored_field_error( - field, - "must be greater than or equal to 0", - )); - } - u64::try_from(value).map_err(|_| stored_field_error(field, "out of range")) -} - -fn stored_field_error( - field: impl Into, - message: impl Into, -) -> RuntimeSettingsError { - RuntimeSettingsError::StoredField { - field: field.into(), - message: message.into(), - } -} diff --git a/backend/src/dispatch/affinity/identity.rs b/backend/src/dispatch/affinity/identity.rs new file mode 100644 index 000000000..ed41093d9 --- /dev/null +++ b/backend/src/dispatch/affinity/identity.rs @@ -0,0 +1,222 @@ +//! 会话亲和变体与 Conversation identity 构建。 + +use serde_json::Value; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +use super::types::ConversationIdentity; +use crate::upstream::openai::protocol::responses::CodexResponsesRequest; + +// ==================================================================== +// Variant identity helpers +// ==================================================================== + +/// 准备用于区分并发分支的变体身份。 +pub fn prepare_variant_identity(request: &mut CodexResponsesRequest) { + request.variant_identity = build_variant_identity(request); +} + +/// 计算请求变体哈希。 +pub fn compute_variant_hash(request: &CodexResponsesRequest) -> String { + compute_variant_hash_with_identity(request, request.variant_identity.as_deref()) +} + +fn compute_variant_hash_with_identity( + request: &CodexResponsesRequest, + identity: Option<&str>, +) -> String { + let tools_json = request.tools().map_or_else(|| "[]".to_string(), tools_json); + let mut hasher = Sha256::new(); + hasher.update(request.instructions().as_bytes()); + hasher.update(b"\0"); + hasher.update(tools_json.as_bytes()); + if let Some(identity) = identity + .map(str::trim) + .filter(|identity| !identity.is_empty()) + { + hasher.update(b"\0"); + hasher.update(identity.as_bytes()); + } + hex::encode(hasher.finalize()).chars().take(12).collect() +} + +fn tools_json(tools: &[Value]) -> String { + serde_json::to_string(tools).unwrap_or_else(|_| "[]".to_string()) +} + +fn build_variant_identity(request: &CodexResponsesRequest) -> Option { + let mut parts = Vec::with_capacity(2); + 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 let Some(anchor) = derive_stable_conversation_key(request) { + parts.push(format!("anchor:{anchor}")); + } + } + + (!parts.is_empty()).then(|| parts.join("\0")) +} + +fn non_empty_str(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +/// 计算 instructions 哈希。 +pub fn hash_instructions(instructions: Option<&str>) -> String { + let mut hasher = Sha256::new(); + hasher.update(instructions.unwrap_or_default().as_bytes()); + hex::encode(hasher.finalize()) +} + +// ==================================================================== +// Conversation identity 构建器 +// ==================================================================== + +const LEADING_SYSTEM_REMINDER_OPEN: &str = ""; +const LEADING_SYSTEM_REMINDER_CLOSE: &str = ""; + +/// 从 prompt_cache_key 和可选的 window_id 构建 conversation identity +pub fn build_conversation_identity( + prompt_cache_key: Option<&str>, + client_window_id: Option<&str>, + account_scope: &str, +) -> ConversationIdentity { + let conversation_id = prompt_cache_key + .filter(|s| !s.trim().is_empty()) + .map(|key| build_account_scoped_identity("conversation", account_scope, key)); + + let window_id = if let Some(client_win) = client_window_id.filter(|s| !s.trim().is_empty()) { + Some(build_account_scoped_identity( + "window", + account_scope, + client_win, + )) + } else { + conversation_id + .as_ref() + .map(|conv_id| format!("{}:0", conv_id)) + }; + + ConversationIdentity { + conversation_id, + window_id, + } +} + +/// 确保请求拥有上游可复用的 prompt cache key。 +pub fn ensure_prompt_cache_key(request: &mut CodexResponsesRequest) { + if let Some(existing) = request + .prompt_cache_key() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + let existing = existing.to_string(); + request.set_prompt_cache_key(Some(existing)); + return; + } + + request.set_prompt_cache_key(Some( + derive_stable_conversation_key(request).unwrap_or_else(|| Uuid::new_v4().to_string()), + )); +} + +/// 按原版 `stable-conversation-key.ts` 的规则派生稳定 conversation key。 +pub fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option { + let instructions = request + .instructions() + .chars() + .take(2000) + .collect::(); + let first_user_text = first_user_text(request.input()); + let normalized_first_user_text = normalize_conversation_anchor_text(&first_user_text); + let first_user_text = if normalized_first_user_text.is_empty() { + first_user_text + } else { + normalized_first_user_text + }; + if instructions.is_empty() && first_user_text.is_empty() { + return None; + } + + let mut hasher = Sha256::new(); + hasher.update(request.model().as_bytes()); + hasher.update(b"\0"); + hasher.update(instructions.as_bytes()); + hasher.update(b"\0"); + hasher.update(first_user_text.as_bytes()); + let hex = hex::encode(hasher.finalize()); + + Some(format!( + "{}-{}-{}-{}-{}", + &hex[0..8], + &hex[8..12], + &hex[12..16], + &hex[16..20], + &hex[20..32] + )) +} + +fn first_user_text(input: &[Value]) -> String { + for item in input { + if item.get("role").and_then(Value::as_str) != Some("user") { + continue; + } + let Some(content) = item.get("content") else { + return String::new(); + }; + if let Some(text) = content.as_str() { + return text.to_string(); + } + if let Some(parts) = content.as_array() { + return parts + .iter() + .filter(|part| part.get("type").and_then(Value::as_str) == Some("input_text")) + .filter_map(|part| part.get("text").and_then(Value::as_str)) + .collect::(); + } + return String::new(); + } + + String::new() +} + +fn normalize_conversation_anchor_text(text: &str) -> String { + let mut rest = text.trim_start(); + loop { + let lower = rest.to_ascii_lowercase(); + if !lower.starts_with(LEADING_SYSTEM_REMINDER_OPEN) { + break; + } + let Some(close_start) = lower.find(LEADING_SYSTEM_REMINDER_CLOSE) else { + break; + }; + let close_end = close_start + LEADING_SYSTEM_REMINDER_CLOSE.len(); + rest = rest[close_end..].trim_start(); + } + rest.to_string() +} + +/// 构建账号作用域的身份哈希。 +fn build_account_scoped_identity(kind: &str, account_scope: &str, client_value: &str) -> String { + let mut hasher = Sha256::new(); + hasher.update(kind.as_bytes()); + hasher.update(b"\0"); + hasher.update(account_scope.as_bytes()); + hasher.update(b"\0"); + hasher.update(client_value.as_bytes()); + + let digest = hasher.finalize(); + let hex = hex::encode(digest); + let truncated = &hex[..32]; + + let prefix = match kind { + "conversation" => "cp", + "window" => "cw", + _ => "cx", + }; + + format!("{}_{}", prefix, truncated) +} diff --git a/backend/src/dispatch/affinity/mod.rs b/backend/src/dispatch/affinity/mod.rs new file mode 100644 index 000000000..576dd89d6 --- /dev/null +++ b/backend/src/dispatch/affinity/mod.rs @@ -0,0 +1,12 @@ +//! 会话亲和身份、存储、解析与运行时服务。 + +pub mod identity; +pub mod resolve; +pub mod service; +pub mod store; +pub mod types; + +pub use identity::*; +pub use service::*; +pub use store::*; +pub use types::*; diff --git a/backend/src/proxy/dispatch/responses/affinity.rs b/backend/src/dispatch/affinity/resolve.rs similarity index 88% rename from backend/src/proxy/dispatch/responses/affinity.rs rename to backend/src/dispatch/affinity/resolve.rs index 951c40a5c..1d11afbd1 100644 --- a/backend/src/proxy/dispatch/responses/affinity.rs +++ b/backend/src/dispatch/affinity/resolve.rs @@ -4,19 +4,17 @@ use chrono::Utc; use tokio::sync::Mutex; use crate::{ - proxy::dispatch::{ + dispatch::{ + affinity::{compute_variant_hash, hash_instructions, RuntimeSessionAffinityService}, reasoning_replay::ReasoningReplayCache, - session_affinity::{ - compute_variant_hash, hash_instructions, RuntimeSessionAffinityService, - }, }, - upstream::protocol::{ + upstream::openai::protocol::{ events::TokenUsage, responses::{completed_response_metadata, CodexResponsesRequest}, }, }; -pub(super) async fn record_response_affinity( +pub(crate) async fn record_response_affinity( session_affinity: &Arc, reasoning_replay: &Arc>, request: &CodexResponsesRequest, @@ -45,7 +43,7 @@ pub(super) async fn record_response_affinity( }; let variant_hash = compute_variant_hash(request); - let entry = crate::proxy::dispatch::session_affinity::SessionAffinityEntry { + let entry = crate::dispatch::affinity::SessionAffinityEntry { account_id: account_id.to_string(), conversation_id: conversation_id.to_string(), turn_state: turn_state @@ -79,7 +77,7 @@ pub(super) async fn record_response_affinity( ); } -pub(super) async fn evict_reasoning_replay( +pub(crate) async fn evict_reasoning_replay( reasoning_replay: &Arc>, request: &CodexResponsesRequest, account_id: &str, diff --git a/backend/src/dispatch/affinity/service.rs b/backend/src/dispatch/affinity/service.rs new file mode 100644 index 000000000..254015f4c --- /dev/null +++ b/backend/src/dispatch/affinity/service.rs @@ -0,0 +1,161 @@ +//! 运行时会话亲和服务。 + +use chrono::{DateTime, Duration, Utc}; +use thiserror::Error; + +use super::{ + store::{RedisSessionAffinityStore, RedisSessionAffinityStoreError}, + types::SessionAffinityEntry, +}; + +const DEFAULT_SESSION_AFFINITY_TTL_SECS: i64 = 4 * 60 * 60; + +/// 运行时会话亲和性服务。 +#[derive(Clone)] +pub struct RuntimeSessionAffinityService { + store: RedisSessionAffinityStore, + ttl: Duration, +} + +/// 运行时会话亲和性错误。 +#[derive(Debug, Error)] +pub enum RuntimeSessionAffinityError { + #[error("session affinity store error: {0}")] + Store(#[from] RedisSessionAffinityStoreError), +} + +impl RuntimeSessionAffinityService { + pub fn new(store: RedisSessionAffinityStore) -> Self { + let ttl = Duration::seconds(DEFAULT_SESSION_AFFINITY_TTL_SECS); + Self { store, ttl } + } + + pub async fn record( + &self, + response_id: String, + entry: SessionAffinityEntry, + ) -> Result<(), RuntimeSessionAffinityError> { + Ok(self.store.upsert(&response_id, &entry, self.ttl).await?) + } + + pub async fn lookup_account(&self, response_id: &str, now: DateTime) -> Option { + self.entry(response_id, now) + .await + .map(|entry| entry.account_id) + } + + pub async fn lookup_conversation_id( + &self, + response_id: &str, + now: DateTime, + ) -> Option { + self.entry(response_id, now) + .await + .map(|entry| entry.conversation_id) + } + + pub async fn lookup_turn_state(&self, response_id: &str, now: DateTime) -> Option { + self.entry(response_id, now) + .await + .and_then(|entry| entry.turn_state) + } + + pub async fn lookup_instructions_hash( + &self, + response_id: &str, + now: DateTime, + ) -> Option { + self.entry(response_id, now) + .await + .and_then(|entry| entry.instructions_hash) + } + + pub async fn lookup_function_call_ids( + &self, + response_id: &str, + now: DateTime, + ) -> Vec { + self.entry(response_id, now) + .await + .map(|entry| entry.function_call_ids) + .unwrap_or_default() + } + + pub async fn lookup_latest_response_by_conversation( + &self, + conversation_id: &str, + max_age: Option, + variant_hash: Option<&str>, + now: DateTime, + ) -> Option { + self.latest(conversation_id, max_age, variant_hash, now) + .await + .map(|(response_id, _)| response_id) + } + + pub async fn lookup_latest_account_by_conversation( + &self, + conversation_id: &str, + max_age: Option, + variant_hash: Option<&str>, + now: DateTime, + ) -> Option { + self.latest(conversation_id, max_age, variant_hash, now) + .await + .map(|(_, entry)| entry.account_id) + } + + pub async fn forget(&self, response_id: &str) -> bool { + match self.store.forget(response_id).await { + Ok(forgotten) => forgotten, + Err(error) => { + tracing::warn!( + response_id, + error = %error, + "failed to remove Redis session affinity" + ); + false + } + } + } + + pub async fn forget_account(&self, account_id: &str) -> bool { + match self.store.forget_account(account_id).await { + Ok(forgotten) => forgotten > 0, + Err(error) => { + tracing::warn!(account_id, error = %error, "failed to remove account affinities"); + false + } + } + } + + async fn entry(&self, response_id: &str, now: DateTime) -> Option { + match self.store.get(response_id, now, self.ttl).await { + Ok(entry) => entry, + Err(error) => { + tracing::warn!(response_id, error = %error, "failed to read Redis session affinity"); + None + } + } + } + + async fn latest( + &self, + conversation_id: &str, + max_age: Option, + variant_hash: Option<&str>, + now: DateTime, + ) -> Option<(String, SessionAffinityEntry)> { + match self + .store + .latest_by_conversation(conversation_id, max_age, variant_hash, now, self.ttl) + .await + { + Ok(entry) => entry, + Err(error) => { + tracing::warn!(conversation_id, error = %error, "failed to query Redis affinity index"); + None + } + } + } +} diff --git a/backend/src/dispatch/affinity/store.rs b/backend/src/dispatch/affinity/store.rs new file mode 100644 index 000000000..2242600b8 --- /dev/null +++ b/backend/src/dispatch/affinity/store.rs @@ -0,0 +1,194 @@ +//! Redis 会话亲和存储。 + +use chrono::{DateTime, Duration, Utc}; +use redis::AsyncCommands; +use thiserror::Error; + +use crate::infra::redis::RedisConnection; + +use super::types::SessionAffinityEntry; + +/// Redis 会话亲和性存储。 +#[derive(Clone)] +pub struct RedisSessionAffinityStore { + redis: RedisConnection, +} + +/// Redis 会话亲和性存储错误。 +#[derive(Debug, Error)] +pub enum RedisSessionAffinityStoreError { + #[error("Redis session affinity operation failed: {0}")] + Redis(#[from] redis::RedisError), + #[error("invalid Redis session affinity value: {0}")] + Json(#[from] serde_json::Error), +} + +pub type RedisSessionAffinityStoreResult = Result; + +impl RedisSessionAffinityStore { + pub fn new(redis: RedisConnection) -> Self { + Self { redis } + } + + pub async fn upsert( + &self, + response_id: &str, + entry: &SessionAffinityEntry, + ttl: Duration, + ) -> RedisSessionAffinityStoreResult<()> { + let value = serde_json::to_string(entry)?; + let expires_at = entry + .created_at + .checked_add_signed(ttl) + .unwrap_or(DateTime::::MAX_UTC); + if expires_at <= Utc::now() { + return Ok(()); + } + + let response_key = self.response_key(response_id); + let conversation_key = self.conversation_key(&entry.conversation_id); + let account_key = self.account_key(&entry.account_id); + let ttl_seconds = ttl.num_seconds().max(1); + let cutoff_ms = Utc::now().timestamp_millis() - ttl.num_milliseconds(); + let mut connection = self.redis.manager(); + let mut transaction = redis::pipe(); + transaction + .atomic() + .cmd("SET") + .arg(response_key) + .arg(value) + .arg("EXAT") + .arg(expires_at.timestamp()) + .ignore() + .cmd("ZADD") + .arg(&conversation_key) + .arg(entry.created_at.timestamp_millis()) + .arg(response_id) + .ignore() + .cmd("EXPIRE") + .arg(&conversation_key) + .arg(ttl_seconds) + .ignore() + .cmd("SADD") + .arg(&account_key) + .arg(response_id) + .ignore() + .cmd("EXPIRE") + .arg(account_key) + .arg(ttl_seconds) + .ignore() + .cmd("ZREMRANGEBYSCORE") + .arg(conversation_key) + .arg("-inf") + .arg(cutoff_ms) + .ignore(); + let _: () = transaction.query_async(&mut connection).await?; + Ok(()) + } + + pub async fn get( + &self, + response_id: &str, + now: DateTime, + ttl: Duration, + ) -> RedisSessionAffinityStoreResult> { + let mut connection = self.redis.manager(); + let value: Option = connection.get(self.response_key(response_id)).await?; + let Some(value) = value else { + return Ok(None); + }; + let entry: SessionAffinityEntry = serde_json::from_str(&value)?; + if entry + .created_at + .checked_add_signed(ttl) + .is_none_or(|expires_at| expires_at <= now) + { + self.forget(response_id).await?; + return Ok(None); + } + Ok(Some(entry)) + } + + pub async fn latest_by_conversation( + &self, + conversation_id: &str, + max_age: Option, + variant_hash: Option<&str>, + now: DateTime, + ttl: Duration, + ) -> RedisSessionAffinityStoreResult> { + let conversation_key = self.conversation_key(conversation_id); + let mut connection = self.redis.manager(); + let response_ids: Vec = connection.zrevrange(&conversation_key, 0, -1).await?; + for response_id in response_ids { + let Some(entry) = self.get(&response_id, now, ttl).await? else { + let _: usize = connection.zrem(&conversation_key, &response_id).await?; + continue; + }; + if entry.conversation_id != conversation_id { + let _: usize = connection.zrem(&conversation_key, &response_id).await?; + continue; + } + if variant_hash.is_some_and(|expected| entry.variant_hash.as_deref() != Some(expected)) + { + continue; + } + if max_age.is_some_and(|max_age| now.signed_duration_since(entry.created_at) > max_age) + { + continue; + } + return Ok(Some((response_id, entry))); + } + Ok(None) + } + + pub async fn forget(&self, response_id: &str) -> RedisSessionAffinityStoreResult { + let response_key = self.response_key(response_id); + let mut connection = self.redis.manager(); + let value: Option = connection.get(&response_key).await?; + let Some(value) = value else { + return Ok(false); + }; + let entry: SessionAffinityEntry = serde_json::from_str(&value)?; + let mut transaction = redis::pipe(); + transaction + .atomic() + .cmd("DEL") + .arg(response_key) + .ignore() + .cmd("ZREM") + .arg(self.conversation_key(&entry.conversation_id)) + .arg(response_id) + .ignore() + .cmd("SREM") + .arg(self.account_key(&entry.account_id)) + .arg(response_id) + .ignore(); + let _: () = transaction.query_async(&mut connection).await?; + Ok(true) + } + + pub async fn forget_account(&self, account_id: &str) -> RedisSessionAffinityStoreResult { + let account_key = self.account_key(account_id); + let mut connection = self.redis.manager(); + let response_ids: Vec = connection.smembers(&account_key).await?; + let mut forgotten = 0u64; + for response_id in response_ids { + forgotten += u64::from(self.forget(&response_id).await?); + } + let _: usize = connection.del(account_key).await?; + Ok(forgotten) + } + + fn response_key(&self, response_id: &str) -> String { + self.redis.key(&format!("affinity:resp:{response_id}")) + } + + fn conversation_key(&self, conversation_id: &str) -> String { + self.redis.key(&format!("affinity:conv:{conversation_id}")) + } + + fn account_key(&self, account_id: &str) -> String { + self.redis.key(&format!("affinity:account:{account_id}")) + } +} diff --git a/backend/src/dispatch/affinity/types.rs b/backend/src/dispatch/affinity/types.rs new file mode 100644 index 000000000..e5e1d8ac7 --- /dev/null +++ b/backend/src/dispatch/affinity/types.rs @@ -0,0 +1,25 @@ +//! 会话亲和领域类型。 + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +/// 会话亲和性条目。 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionAffinityEntry { + pub account_id: String, + pub conversation_id: String, + pub turn_state: Option, + pub instructions_hash: Option, + pub input_tokens: Option, + pub function_call_ids: Vec, + pub variant_hash: Option, + pub created_at: DateTime, +} + +/// 从请求上下文派生的 conversation identity。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ConversationIdentity { + pub conversation_id: Option, + pub window_id: Option, +} diff --git a/backend/src/proxy/dispatch/cloudflare.rs b/backend/src/dispatch/cloudflare.rs similarity index 91% rename from backend/src/proxy/dispatch/cloudflare.rs rename to backend/src/dispatch/cloudflare.rs index b10049346..47d787f20 100644 --- a/backend/src/proxy/dispatch/cloudflare.rs +++ b/backend/src/dispatch/cloudflare.rs @@ -5,25 +5,23 @@ use std::time::Duration as StdDuration; use chrono::Utc; use crate::{ - upstream::accounts::{ - cookies::{ - CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, SqliteCookieStore, - }, - model::AccountStatus, + accounts::{ + account::AccountStatus, + cookies::{CloudflareChallengeCooldownTracker, CloudflarePathBlockTracker, PgCookieStore}, pool::RuntimeAccountPoolService, }, - upstream::transport::CodexClientError, + upstream::openai::transport::CodexClientError, }; #[derive(Clone)] pub struct CloudflareRecovery { path_block_tracker: CloudflarePathBlockTracker, challenge_tracker: CloudflareChallengeCooldownTracker, - cookie_store: SqliteCookieStore, + cookie_store: PgCookieStore, } impl CloudflareRecovery { - pub fn new(cookie_store: SqliteCookieStore) -> Self { + pub fn new(cookie_store: PgCookieStore) -> Self { Self { path_block_tracker: CloudflarePathBlockTracker::default(), challenge_tracker: CloudflareChallengeCooldownTracker::default(), @@ -113,7 +111,7 @@ impl CloudflareRecovery { } } -async fn delete_account_cookies(cookie_store: &SqliteCookieStore, account_id: &str, reason: &str) { +async fn delete_account_cookies(cookie_store: &PgCookieStore, account_id: &str, reason: &str) { if let Err(error) = cookie_store.delete_account_cookies(account_id).await { tracing::warn!( account_id, diff --git a/backend/src/proxy/dispatch/errors.rs b/backend/src/dispatch/errors.rs similarity index 98% rename from backend/src/proxy/dispatch/errors.rs rename to backend/src/dispatch/errors.rs index bd461d643..cca448354 100644 --- a/backend/src/proxy/dispatch/errors.rs +++ b/backend/src/dispatch/errors.rs @@ -4,9 +4,9 @@ use chrono::{DateTime, Duration, Utc}; use serde_json::{json, Map, Value}; use crate::{ - proxy::dispatch::exhaustion::{ExhaustedAccountKind, ExhaustedAccountRef}, - upstream::accounts::model::AccountStatus, - upstream::transport::{ + accounts::account::AccountStatus, + dispatch::exhaustion::{ExhaustedAccountKind, ExhaustedAccountRef}, + upstream::openai::transport::{ is_banned_auth_signal, CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics, }, }; diff --git a/backend/src/proxy/dispatch/exhaustion.rs b/backend/src/dispatch/exhaustion.rs similarity index 99% rename from backend/src/proxy/dispatch/exhaustion.rs rename to backend/src/dispatch/exhaustion.rs index d89a1b524..b0218186c 100644 --- a/backend/src/proxy/dispatch/exhaustion.rs +++ b/backend/src/dispatch/exhaustion.rs @@ -1,6 +1,6 @@ //! 账号耗尽状态聚合。 -use crate::upstream::accounts::model::AccountStatus; +use crate::accounts::account::AccountStatus; #[derive(Clone, Copy, Debug)] pub(crate) enum ExhaustedAccountKind { diff --git a/backend/src/proxy/dispatch/implicit_resume.rs b/backend/src/dispatch/implicit_resume.rs similarity index 98% rename from backend/src/proxy/dispatch/implicit_resume.rs rename to backend/src/dispatch/implicit_resume.rs index d71b7f29c..7d8efbde9 100644 --- a/backend/src/proxy/dispatch/implicit_resume.rs +++ b/backend/src/dispatch/implicit_resume.rs @@ -4,7 +4,7 @@ use std::collections::HashSet; use serde_json::Value; -use crate::upstream::protocol::responses::CodexResponsesRequest; +use crate::upstream::openai::protocol::responses::CodexResponsesRequest; /// 隐式续接前的可恢复请求状态。 #[derive(Debug, Clone, PartialEq)] diff --git a/backend/src/proxy/dispatch/mod.rs b/backend/src/dispatch/mod.rs similarity index 83% rename from backend/src/proxy/dispatch/mod.rs rename to backend/src/dispatch/mod.rs index b778d3ded..67404e836 100644 --- a/backend/src/proxy/dispatch/mod.rs +++ b/backend/src/dispatch/mod.rs @@ -1,11 +1,10 @@ //! 请求调度模块（账号选择、回退、恢复、会话亲和性）。 +pub mod affinity; pub mod cloudflare; pub mod errors; mod exhaustion; pub mod implicit_resume; pub mod reasoning_replay; pub mod responses; -pub mod session_affinity; pub mod upstream; -mod usage_events; diff --git a/backend/src/proxy/dispatch/reasoning_replay.rs b/backend/src/dispatch/reasoning_replay.rs similarity index 100% rename from backend/src/proxy/dispatch/reasoning_replay.rs rename to backend/src/dispatch/reasoning_replay.rs diff --git a/backend/src/proxy/dispatch/responses/errors.rs b/backend/src/dispatch/responses/errors.rs similarity index 99% rename from backend/src/proxy/dispatch/responses/errors.rs rename to backend/src/dispatch/responses/errors.rs index 3818f9b75..fc96d7299 100644 --- a/backend/src/proxy/dispatch/responses/errors.rs +++ b/backend/src/dispatch/responses/errors.rs @@ -2,14 +2,14 @@ use serde_json::Value; use thiserror::Error; use crate::{ - proxy::dispatch::{ + dispatch::{ errors::{ insert_dispatch_error_metadata, upstream_error_http_status, DispatchErrorMetadata, DispatchFailureClass, }, exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef}, }, - upstream::{ + upstream::openai::{ protocol::{responses::ResponsesSseFailure, sse::SseError}, transport::CodexClientError, }, diff --git a/backend/src/proxy/dispatch/responses/event_recording.rs b/backend/src/dispatch/responses/event_recording.rs similarity index 91% rename from backend/src/proxy/dispatch/responses/event_recording.rs rename to backend/src/dispatch/responses/event_recording.rs index 428cc9e2a..d3aa55412 100644 --- a/backend/src/proxy/dispatch/responses/event_recording.rs +++ b/backend/src/dispatch/responses/event_recording.rs @@ -3,25 +3,22 @@ use std::time::Instant; use serde_json::{json, Map, Value}; use crate::{ - admin::monitoring::{ - ops_error_model::OpsErrorLog, - ops_error_service::AdminOpsErrorLogService, - usage_record_model::{UsageRecord, UsageRecordLevel}, + dispatch::errors::{ + backend_transport_name, insert_dispatch_error_metadata, + insert_upstream_diagnostics_metadata, upstream_error_http_status, DispatchErrorMetadata, }, infra::time::elapsed_millis_i64, - proxy::dispatch::{ - errors::{ - backend_transport_name, insert_dispatch_error_metadata, - insert_upstream_diagnostics_metadata, upstream_error_http_status, - DispatchErrorMetadata, - }, - usage_events::{ + telemetry::{ + ops::query::OpsQueryService, + ops::types::OpsErrorLog, + recorder::{ enrich_event_route_metadata, enrich_usage_record_identity, event_kind as response_event_kind, reasoning_effort_from_request, record_dispatch_error_event, DispatchErrorLogRecord, }, + usage::types::UsageRecordLevel, }, - upstream::{ + upstream::openai::{ protocol::responses::{CodexResponsesRequest, ResponsesSseFailure}, transport::{ CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics, @@ -40,7 +37,7 @@ use super::{ }; pub(super) struct ResponseUpstreamErrorEventRecord<'a> { - pub(super) ops_errors: &'a AdminOpsErrorLogService, + pub(super) ops_errors: &'a OpsQueryService, pub(super) request_id: &'a str, pub(super) account_id: &'a str, pub(super) account_email: Option<&'a str>, @@ -56,7 +53,7 @@ pub(super) struct ResponseUpstreamErrorEventRecord<'a> { } pub(super) struct ResponseStreamFailureEventRecord<'a> { - pub(super) ops_errors: &'a AdminOpsErrorLogService, + pub(super) ops_errors: &'a OpsQueryService, pub(super) request_id: &'a str, pub(super) account_id: &'a str, pub(super) route: &'a str, @@ -74,8 +71,9 @@ pub(super) struct ResponseStreamFailureEventRecord<'a> { } pub(super) struct ResponseDispatchErrorEventRecord<'a> { - pub(super) ops_errors: &'a AdminOpsErrorLogService, + pub(super) ops_errors: &'a OpsQueryService, pub(super) request_id: &'a str, + pub(super) client_api_key_id: Option<&'a str>, pub(super) account_id: Option<&'a str>, pub(super) route: &'a str, pub(super) model: &'a str, @@ -87,6 +85,7 @@ pub(super) struct ResponseDispatchErrorEventRecord<'a> { } pub(super) struct ResponseDispatchErrorDetails<'a> { + pub(super) client_api_key_id: Option<&'a str>, pub(super) account_id: Option<&'a str>, pub(super) stream: bool, pub(super) compact: bool, @@ -114,6 +113,8 @@ pub(super) async fn record_response_dispatch_error_event( record_dispatch_error_event(DispatchErrorLogRecord { ops_errors: record.ops_errors, request_id: record.request_id, + client_api_key_id: record.client_api_key_id, + provider: Some("openai"), account_id: record.account_id, route: record.route, model: record.model, @@ -138,6 +139,8 @@ pub(super) async fn record_response_upstream_error_event( enrich_event_route_metadata(&mut metadata, record.route); let mut event = OpsErrorLog::new("v1.response", "v1 responses upstream request failed"); event.request_id = Some(record.request_id.to_string()); + event.client_api_key_id = record.request.client_api_key_id.clone(); + event.provider = Some("openai".to_string()); event.account_id = Some(record.account_id.to_string()); event.route = Some(record.route.to_string()); event.model = Some(record.model.to_string()); @@ -225,6 +228,8 @@ pub(super) async fn record_prefetched_response_stream_failure_event( "v1 responses stream failed", ); event.request_id = Some(record.request_id.to_string()); + event.client_api_key_id = record.request.client_api_key_id.clone(); + event.provider = Some("openai".to_string()); event.account_id = Some(record.account_id.to_string()); event.route = Some(record.route.to_string()); event.model = Some(record.model.to_string()); @@ -499,20 +504,32 @@ pub(super) async fn record_live_response_stream_event( ); if is_success_usage_event(level, status_code) { - let mut event = UsageRecord::new(response_event_kind(&context.route), level, message); - event.request_id = Some(context.request_id.clone()); - event.account_id = Some(context.account_id.clone()); - event.route = Some(context.route.clone()); - event.model = Some(context.display_model.clone()); - event.status_code = Some(status_code); - event.latency_ms = Some(elapsed_millis_i64(context.started_at)); - event.metadata = metadata; - if let Err(error) = context.usage_records.record(event).await { - tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream event"); - } + crate::telemetry::recorder::record_response_event( + crate::telemetry::usage::types::ResponseUsageRecord { + usage_records: &context.usage_records, + request_id: &context.request_id, + client_api_key_id: context.request.client_api_key_id.as_deref(), + account_id: &context.account_id, + route: &context.route, + model: &context.display_model, + requested_model: Some(&context.requested_model), + client_ip: context.client_ip.as_deref(), + client_user_agent: context.request.client_user_agent.as_deref(), + reasoning_effort: reasoning_effort_from_request(&context.request), + service_tier: context.request.service_tier(), + started_at: context.started_at, + status_code, + message, + metadata, + rate_limit_headers, + }, + ) + .await; } else { let mut event = OpsErrorLog::new(response_event_kind(&context.route), message); event.request_id = Some(context.request_id.clone()); + event.client_api_key_id = context.request.client_api_key_id.clone(); + event.provider = Some("openai".to_string()); event.account_id = Some(context.account_id.clone()); event.route = Some(context.route.clone()); event.model = Some(context.display_model.clone()); diff --git a/backend/src/proxy/dispatch/responses/live_stream.rs b/backend/src/dispatch/responses/live_stream.rs similarity index 99% rename from backend/src/proxy/dispatch/responses/live_stream.rs rename to backend/src/dispatch/responses/live_stream.rs index deb3d9e16..2ccf0176a 100644 --- a/backend/src/proxy/dispatch/responses/live_stream.rs +++ b/backend/src/dispatch/responses/live_stream.rs @@ -3,13 +3,13 @@ use std::{ task::{Context, Poll}, }; -use axum::body::Bytes; +use bytes::Bytes; use futures::{stream::Stream, StreamExt}; use serde_json::Value; use tokio::sync::{mpsc, oneshot}; use crate::{ - proxy::dispatch::responses::{ + dispatch::responses::{ errors::ResponseDispatchStreamError, service::ResponseDispatchStream, stream_lifecycle::{ @@ -17,7 +17,7 @@ use crate::{ LiveResponseStreamContext, }, }, - upstream::{ + upstream::openai::{ protocol::{ responses::{ reconvert_responses_sse_event_tuple_values, response_sse_event_is_terminal, diff --git a/backend/src/proxy/dispatch/responses/mod.rs b/backend/src/dispatch/responses/mod.rs similarity index 93% rename from backend/src/proxy/dispatch/responses/mod.rs rename to backend/src/dispatch/responses/mod.rs index 58880f4a9..ac71b5f2e 100644 --- a/backend/src/proxy/dispatch/responses/mod.rs +++ b/backend/src/dispatch/responses/mod.rs @@ -1,6 +1,5 @@ //! Responses 调度领域。 -mod affinity; pub(crate) mod errors; mod event_recording; mod live_stream; diff --git a/backend/src/proxy/dispatch/responses/prefetch.rs b/backend/src/dispatch/responses/prefetch.rs similarity index 95% rename from backend/src/proxy/dispatch/responses/prefetch.rs rename to backend/src/dispatch/responses/prefetch.rs index d102c3b24..aa870dfc4 100644 --- a/backend/src/proxy/dispatch/responses/prefetch.rs +++ b/backend/src/dispatch/responses/prefetch.rs @@ -1,7 +1,7 @@ -use axum::body::Bytes; +use bytes::Bytes; use futures::StreamExt; -use crate::upstream::{ +use crate::upstream::openai::{ protocol::sse::{sse_frame_end, SseError}, transport::CodexBackendSseStream, }; diff --git a/backend/src/dispatch/responses/service.rs b/backend/src/dispatch/responses/service.rs new file mode 100644 index 000000000..33a9f9b93 --- /dev/null +++ b/backend/src/dispatch/responses/service.rs @@ -0,0 +1,452 @@ +//! Responses 创建编排与调度服务。 +//! +//! 包含了将 OpenAI 请求调度到 Codex 上游账号的完整逻辑，包括： +//! - 响应创建（非流式 / 流式 / compact） +//! - 会话亲和性与隐式续接 +//! - reasoning replay +//! - 账号回退与错误恢复 +//! - 配额验证 + +use std::{pin::Pin, sync::Arc, time::Instant}; + +use bytes::Bytes; +use chrono::{DateTime, Duration, Utc}; +use futures::stream::Stream; +use serde_json::{json, Value}; +use tokio::sync::Mutex; + +use crate::{ + accounts::{ + account::{Account, AccountStatus}, + pool::{AccountAcquireRequest, RuntimeAccountPoolService}, + }, + dispatch::{ + affinity::resolve::{evict_reasoning_replay, record_response_affinity}, + affinity::{ + compute_variant_hash, ensure_prompt_cache_key, hash_instructions, + prepare_variant_identity, RuntimeSessionAffinityService, + }, + cloudflare::{ + cloudflare_challenge_error_message, cloudflare_path_block_error_message, + is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error, + CloudflareRecovery, + }, + errors::{ + auth_failure_account_status, backend_transport_name, is_auth_upstream_error, + is_history_recovery_upstream_error, is_model_unsupported_upstream_error, + is_quota_exhausted_upstream_error, is_rate_limit_upstream_error, + rate_limit_cooldown_until, upstream_error_body, upstream_error_http_status, + upstream_error_set_cookie_headers, + }, + exhaustion::AccountExhaustionTracker, + reasoning_replay::ReasoningReplayCache, + upstream::{ + create_compact_response_with_account_retrying_5xx, + create_response_stream_with_account_retrying_5xx, + create_response_with_account_retrying_5xx, verify_acquired_quota_if_required, + QuotaVerificationContext, QuotaVerificationDecision, + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + }, + }, + models::service::ModelService, + telemetry::{ + ops::query::OpsQueryService, + recorder::{ + reasoning_effort_from_compact_request, reasoning_effort_from_request, + record_response_event, + }, + usage::query::UsageQueryService, + usage::types::ResponseUsageRecord, + }, + upstream::openai::{ + protocol::{ + events::{extract_usage, TokenUsage}, + responses::{ + response_from_codex_sse, CodexCompactRequest, CodexResponsesRequest, + CollectedResponse, + }, + }, + transport::{ + backend_transport_for_response_request, is_banned_upstream_error, CodexBackendClient, + CodexBackendResponse, + }, + }, +}; + +use crate::dispatch::implicit_resume::{ + continuation_input_start, implicit_resume_allowed, ImplicitResumeSnapshot, +}; + +use super::{ + errors::{ResponseDispatchError, ResponseDispatchStreamError}, + event_recording::{ + insert_response_status_metadata, insert_response_trace_metadata, + insert_response_upstream_diagnostics, insert_websocket_pool_decision, + record_prefetched_response_stream_failure_event, record_response_dispatch_error_event, + record_response_upstream_error_event, ResponseDispatchErrorDetails, + ResponseDispatchErrorEventRecord, ResponseStreamFailureEventRecord, + ResponseUpstreamErrorEventRecord, + }, + live_stream::spawn_live_response_stream, + prefetch::prefetch_first_sse_chunk, + sse_failure::{ + auth_sse_failure_account_status, client_error_invalid_reasoning_replay, first_sse_failure, + is_auth_sse_failure, is_history_recovery_sse_failure, is_model_unsupported_sse_failure, + is_quota_exhausted_sse_failure, sse_failure_error_body, + sse_failure_invalid_reasoning_replay, stream_failure_http_status, + }, + stream_lifecycle::LiveResponseStreamContext, + trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, +}; + +/// OpenAI Responses 调度服务。 +#[derive(Clone)] +pub struct ResponseDispatchService { + account_pool: Arc, + models: Arc, + codex: Arc, + session_affinity: Arc, + reasoning_replay: Arc>, + usage_records: Arc, + ops_errors: Arc, + installation_id: Option, + cloudflare: CloudflareRecovery, +} + +pub(crate) struct ResponseDispatchServiceParts { + pub account_pool: Arc, + pub models: Arc, + pub codex: Arc, + pub session_affinity: Arc, + pub usage_records: Arc, + pub ops_errors: Arc, + pub installation_id: Option, + pub cloudflare: CloudflareRecovery, +} + +/// 默认 reasoning replay TTL 秒数。 +const DEFAULT_REASONING_REPLAY_TTL_SECS: i64 = 55 * 60; +const IMPLICIT_RESUME_MAX_AGE_SECS: i64 = DEFAULT_REASONING_REPLAY_TTL_SECS; + +#[derive(Debug, Clone, Default)] +struct AccountAffinityDecision { + preferred_account_id: Option, + request_history_owner: RequestHistoryOwner, +} + +#[derive(Debug, Clone, Default)] +enum RequestHistoryOwner { + #[default] + Absent, + Known(String), + Unknown, +} + +impl AccountAffinityDecision { + fn preferred_account_id(&self) -> Option<&str> { + self.preferred_account_id.as_deref() + } + + fn should_strip_request_history_for(&self, account_id: &str) -> bool { + match &self.request_history_owner { + RequestHistoryOwner::Absent => false, + RequestHistoryOwner::Known(history_account_id) => history_account_id != account_id, + RequestHistoryOwner::Unknown => true, + } + } +} + +/// Responses live SSE 响应体流。 +pub type ResponseBodyStream = + Pin> + Send + 'static>>; + +/// Responses live SSE 调度结果。 +pub struct ResponseDispatchStream { + pub body: ResponseBodyStream, +} + +impl ResponseDispatchService { + pub(crate) fn new(parts: ResponseDispatchServiceParts) -> Self { + Self { + account_pool: parts.account_pool, + models: parts.models, + codex: parts.codex, + session_affinity: parts.session_affinity, + reasoning_replay: Arc::new(Mutex::new(ReasoningReplayCache::new(Duration::seconds( + DEFAULT_REASONING_REPLAY_TTL_SECS, + )))), + usage_records: parts.usage_records, + ops_errors: parts.ops_errors, + installation_id: parts.installation_id, + cloudflare: parts.cloudflare, + } + } + + async fn prepare_response_session( + &self, + request: &mut CodexResponsesRequest, + ) -> Option { + prepare_variant_identity(request); + if let Some(previous_response_id) = request.previous_response_id().map(ToString::to_string) + { + if request.prompt_cache_key().is_none() { + let conversation_id = self + .session_affinity + .lookup_conversation_id(&previous_response_id, Utc::now()) + .await; + request.set_prompt_cache_key(conversation_id); + } + if request.turn_state.as_deref().is_none_or(str::is_empty) { + request.turn_state = self + .session_affinity + .lookup_turn_state(&previous_response_id, Utc::now()) + .await; + } + ensure_prompt_cache_key(request); + return None; + } + + ensure_prompt_cache_key(request); + self.apply_implicit_resume(request).await + } + + async fn apply_implicit_resume( + &self, + request: &mut CodexResponsesRequest, + ) -> Option { + let continuation_start = continuation_input_start(request.input()); + if continuation_start == 0 || continuation_start >= request.input().len() { + return None; + } + let conversation_id = request + .prompt_cache_key() + .map(str::trim) + .filter(|value| !value.is_empty())? + .to_string(); + let snapshot = ImplicitResumeSnapshot::capture(request); + let variant_hash = compute_variant_hash(request); + let now = Utc::now(); + let previous_response_id = self + .session_affinity + .lookup_latest_response_by_conversation( + &conversation_id, + Some(Duration::seconds(IMPLICIT_RESUME_MAX_AGE_SECS)), + Some(&variant_hash), + now, + ) + .await?; + let current_instructions_hash = hash_instructions(Some(request.instructions())); + if self + .session_affinity + .lookup_instructions_hash(&previous_response_id, now) + .await + .as_deref() + != Some(current_instructions_hash.as_str()) + { + return None; + } + let stored_function_call_ids = self + .session_affinity + .lookup_function_call_ids(&previous_response_id, now) + .await; + if !implicit_resume_allowed( + &request.input()[continuation_start..], + request.input(), + &stored_function_call_ids, + ) { + return None; + } + let account_id = self + .session_affinity + .lookup_account(&previous_response_id, now) + .await?; + let replay_items = self.reasoning_replay.lock().await.lookup( + &previous_response_id, + &account_id, + &conversation_id, + &variant_hash, + now, + ); + let continuation = request.input()[continuation_start..].to_vec(); + let mut input = replay_items; + input.extend(continuation); + + request.set_previous_response_id(Some(previous_response_id.clone())); + request.use_websocket = true; + request.force_http_sse = false; + request.set_input(input); + if let Some(turn_state) = self + .session_affinity + .lookup_turn_state(&previous_response_id, now) + .await + { + request.turn_state = Some(turn_state); + } + + Some(snapshot) + } + + async fn account_affinity_for_request( + &self, + request: &CodexResponsesRequest, + now: DateTime, + ) -> AccountAffinityDecision { + let has_previous_response_id = request.previous_response_id().is_some(); + if let Some(account_id) = self.previous_response_account_id(request, now).await { + return AccountAffinityDecision { + preferred_account_id: Some(account_id.clone()), + request_history_owner: RequestHistoryOwner::Known(account_id), + }; + } + + let conversation_account_id = self.conversation_account_id(request, now).await; + AccountAffinityDecision { + preferred_account_id: conversation_account_id, + request_history_owner: if has_previous_response_id { + RequestHistoryOwner::Unknown + } else { + RequestHistoryOwner::Absent + }, + } + } + + async fn previous_response_account_id( + &self, + request: &CodexResponsesRequest, + now: DateTime, + ) -> Option { + let previous_response_id = request.previous_response_id()?; + self.session_affinity + .lookup_account(previous_response_id, now) + .await + } + + async fn conversation_account_id( + &self, + request: &CodexResponsesRequest, + now: DateTime, + ) -> Option { + let conversation_id = request + .prompt_cache_key() + .map(str::trim) + .filter(|value| !value.is_empty())?; + let variant_hash = compute_variant_hash(request); + self.session_affinity + .lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now) + .await + } + + async fn recover_request_history( + &self, + request: &mut CodexResponsesRequest, + implicit_resume: &mut Option, + ) { + if let Some(previous_response_id) = request.previous_response_id() { + self.session_affinity.forget(previous_response_id).await; + } + restore_request_without_history(request, implicit_resume); + } + + fn strip_history_if_account_changed( + request: &mut CodexResponsesRequest, + implicit_resume: &mut Option, + account_affinity: &AccountAffinityDecision, + acquired_account_id: &str, + ) { + if request.previous_response_id().is_none() + || !account_affinity.should_strip_request_history_for(acquired_account_id) + { + return; + } + + restore_request_without_history(request, implicit_resume); + } + + async fn apply_cascading_ban_defense( + &self, + request: &mut CodexResponsesRequest, + implicit_resume: &mut Option, + preferred_account_id: Option<&str>, + acquired_account_id: &str, + explicit_previous_response_id: Option<&str>, + ) -> bool { + let Some(preferred_account_id) = + preferred_account_id.filter(|account_id| *account_id != acquired_account_id) + else { + return false; + }; + let has_history = request.previous_response_id().is_some() + || request + .turn_state + .as_deref() + .is_some_and(|value| !value.is_empty()); + if !has_history { + return false; + } + let Some(preferred_account) = self + .account_pool + .account_snapshot(preferred_account_id) + .await + else { + return false; + }; + if !matches!( + preferred_account.status, + AccountStatus::Banned | AccountStatus::Disabled + ) { + return false; + } + + let response_id_to_forget = explicit_previous_response_id + .or(request.previous_response_id()) + .map(str::to_string); + restore_request_without_history(request, implicit_resume); + if let Some(response_id) = response_id_to_forget { + self.session_affinity.forget(&response_id).await; + } + true + } + + async fn evict_reasoning_replay(&self, request: &CodexResponsesRequest, account_id: &str) { + evict_reasoning_replay(&self.reasoning_replay, request, account_id).await; + } + + async fn record_response_affinity( + &self, + request: &CodexResponsesRequest, + account_id: &str, + body: &str, + turn_state: Option, + usage: Option, + ) { + record_response_affinity( + &self.session_affinity, + &self.reasoning_replay, + request, + account_id, + body, + turn_state, + usage, + ) + .await; + } +} + +mod compact; +mod complete; +mod stream; + +fn strip_request_history(request: &mut CodexResponsesRequest) { + request.set_previous_response_id(None); + request.turn_state = None; +} + +fn restore_request_without_history( + request: &mut CodexResponsesRequest, + implicit_resume: &mut Option, +) { + if let Some(snapshot) = implicit_resume.take() { + snapshot.restore(request); + } + strip_request_history(request); +} diff --git a/backend/src/dispatch/responses/service/compact.rs b/backend/src/dispatch/responses/service/compact.rs new file mode 100644 index 000000000..1e1e4a2a7 --- /dev/null +++ b/backend/src/dispatch/responses/service/compact.rs @@ -0,0 +1,290 @@ +use super::*; + +impl ResponseDispatchService { + /// 调度 Responses compact 请求到 Codex compact 上游。 + pub async fn compact( + &self, + request_id: &str, + mut request: CodexCompactRequest, + requested_model: &str, + ) -> Result { + 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 excluded_account_ids = Vec::new(); + let mut exhausted_accounts = AccountExhaustionTracker::default(); + let mut quota_verify_attempts = 0usize; + let mut trace = ResponseDispatchTrace::default(); + + loop { + let acquire_request = AccountAcquireRequest::new(request.model(), Utc::now()) + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + let acquired = match self.account_pool.acquire_with(&acquire_request).await { + Some(acquired) => acquired, + None => { + 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, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + allow_retry_with_another_account: true, + }, + acquired, + ) + .await + { + QuotaVerificationDecision::Ready(acquired) => *acquired, + QuotaVerificationDecision::RetryWithAnotherAccount => { + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); + continue; + } + QuotaVerificationDecision::MaxAttemptsReached + | QuotaVerificationDecision::RequiredAccountUnavailable => { + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + 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, + Some(&acquired_account_id), + &error, + ) + .await; + return Err(error); + } + }; + let account = acquired.account; + 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( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + ) + .await; + self.account_pool.release(&release_account_id).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, request.model(), &usage) + .await; + } + let mut metadata = json!({ + "stream": false, + "compact": true, + "usage": usage, + }); + 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 { + usage_records: &self.usage_records, + request_id, + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: &account.id, + route: "/v1/responses/compact", + model: &display_model, + requested_model: Some(requested_model), + client_ip: request.client_ip.as_deref(), + client_user_agent: request.client_user_agent.as_deref(), + reasoning_effort: reasoning_effort_from_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; + return Ok(response.body); + } + 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; + excluded_account_ids.push(release_account_id); + } + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + let account_status = auth_failure_account_status(&error); + 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; + excluded_account_ids.push(release_account_id); + } + 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; + excluded_account_ids.push(release_account_id); + } + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = ResponseDispatchError::from_exhausted_account(exhausted); + 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); + } + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); + excluded_account_ids.push(release_account_id); + } + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) => { + 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 { + ops_errors: &self.ops_errors, + 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; + } +} diff --git a/backend/src/dispatch/responses/service/complete.rs b/backend/src/dispatch/responses/service/complete.rs new file mode 100644 index 000000000..12ca29dab --- /dev/null +++ b/backend/src/dispatch/responses/service/complete.rs @@ -0,0 +1,563 @@ +use super::*; + +impl ResponseDispatchService { + /// 调度非流式 Responses 请求到 Codex Responses 上游。 + pub async fn complete( + &self, + request_id: &str, + route: &str, + mut request: CodexResponsesRequest, + requested_model: &str, + ) -> Result { + let started_at = Instant::now(); + let catalog = self.models.catalog().await; + let display_model = catalog.resolve_model_id(requested_model); + request.set_model(display_model.clone()); + let tuple_schema = request.tuple_schema.clone(); + let image_generation_requested = request.expects_image_generation(); + let now = Utc::now(); + let explicit_previous_response_id = request.previous_response_id().map(ToString::to_string); + let mut implicit_resume = self.prepare_response_session(&mut request).await; + let account_affinity = self.account_affinity_for_request(&request, now).await; + let mut acquire_request = AccountAcquireRequest::new(request.model(), now); + if let Some(preferred_account_id) = account_affinity.preferred_account_id() { + acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); + } + let mut excluded_account_ids = Vec::new(); + let mut exhausted_accounts = AccountExhaustionTracker::default(); + let mut history_recovery_used = false; + let mut next_required_account_id: Option = None; + let mut empty_response_retries = 0u8; + let mut quota_verify_attempts = 0usize; + let mut trace = ResponseDispatchTrace::default(); + const MAX_EMPTY_RESPONSE_RETRIES: u8 = 2; + let (account, response, collected_response, attempt): ( + Account, + CodexBackendResponse, + CollectedResponse, + ResponseDispatchAttempt, + ) = loop { + let mut attempt_acquire_request = acquire_request + .clone() + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + if let Some(account_id) = next_required_account_id.take() { + attempt_acquire_request = + attempt_acquire_request.with_required_account_id(account_id); + } + let allow_quota_verify_account_retry = + attempt_acquire_request.required_account_id.is_none(); + attempt_acquire_request.now = Utc::now(); + let Some(acquired) = self + .account_pool + .acquire_with(&attempt_acquire_request) + .await + else { + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: exhausted_accounts.last_account_id(), + stream: false, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, + &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, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + allow_retry_with_another_account: allow_quota_verify_account_retry, + }, + acquired, + ) + .await + { + QuotaVerificationDecision::Ready(acquired) => *acquired, + QuotaVerificationDecision::RetryWithAnotherAccount => { + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); + continue; + } + QuotaVerificationDecision::MaxAttemptsReached + | QuotaVerificationDecision::RequiredAccountUnavailable => { + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&acquired_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, + &error, + ) + .await; + return Err(error); + } + }; + + self.apply_cascading_ban_defense( + &mut request, + &mut implicit_resume, + account_affinity.preferred_account_id(), + &acquired.account.id, + explicit_previous_response_id.as_deref(), + ) + .await; + Self::strip_history_if_account_changed( + &mut request, + &mut implicit_resume, + &account_affinity, + &acquired.account.id, + ); + self.account_pool.wait_for_request_interval(&acquired).await; + let account = acquired.account; + let release_account_id = account.id.clone(); + let attempt = trace.start_attempt(&release_account_id); + let response_result = create_response_with_account_retrying_5xx( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + started_at, + ) + .await; + self.account_pool.release(&release_account_id).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.account_pool + .sync_passive_rate_limit_headers(&account, &response.rate_limit_headers) + .await; + let collected_response = + match response_from_codex_sse(&response.body, tuple_schema.as_ref()) { + Ok(collected_response) => collected_response, + Err(error) => { + let error = ResponseDispatchError::InvalidSse(error); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + return Err(error); + } + }; + if matches!(collected_response, CollectedResponse::Empty) { + self.account_pool + .record_empty_response_attempt( + &release_account_id, + request.model(), + image_generation_requested, + ) + .await; + empty_response_retries += 1; + if empty_response_retries <= MAX_EMPTY_RESPONSE_RETRIES { + continue; + } + } + if let CollectedResponse::Failed(failure) = &collected_response { + if is_history_recovery_sse_failure(failure) && !history_recovery_used { + if sse_failure_invalid_reasoning_replay(failure) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + next_required_account_id = Some(release_account_id); + continue; + } + if is_model_unsupported_sse_failure(failure) { + let upstream_error = sse_failure_error_body(failure); + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = + ResponseDispatchError::from_exhausted_account(exhausted); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + return Err(error); + } + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); + excluded_account_ids.push(release_account_id); + continue; + } + if is_quota_exhausted_sse_failure(failure) { + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + failure.message.clone(), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + if is_auth_sse_failure(failure) { + let upstream_error = sse_failure_error_body(failure); + let account_status = auth_sse_failure_account_status(failure); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(stream_failure_http_status(failure)), + ); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + } + break (account, response, collected_response, attempt); + } + 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; + excluded_account_ids.push(release_account_id); + } + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) + if is_history_recovery_upstream_error(&error) && !history_recovery_used => + { + if client_error_invalid_reasoning_replay(&error) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + next_required_account_id = Some(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + let account_status = auth_failure_account_status(&error); + 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; + excluded_account_ids.push(release_account_id); + } + 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; + excluded_account_ids.push(release_account_id); + } + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + let upstream_error = upstream_error_body(&error); + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = ResponseDispatchError::from_exhausted_account(exhausted); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&release_account_id), + stream: false, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, + &error, + ) + .await; + return Err(error); + } + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); + excluded_account_ids.push(release_account_id); + } + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) => { + record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { + ops_errors: &self.ops_errors, + request_id, + account_id: &release_account_id, + account_email: account.email.as_deref(), + route, + model: requested_model, + started_at, + stream: false, + transport: backend_transport_for_response_request(&request), + request: &request, + error: &error, + trace: &trace, + attempt: Some(&attempt), + }) + .await; + return Err(ResponseDispatchError::Upstream(error)); + } + } + }; + + match collected_response { + CollectedResponse::Completed(body) => { + let response_id = body.get("id").and_then(Value::as_str); + self.cloudflare.reset_account_recovery(&account.id).await; + if let Some(usage) = response.usage { + self.account_pool + .record_response_usage( + &account.id, + request.model(), + usage, + image_generation_requested, + ) + .await; + } + self.record_response_affinity( + &request, + &account.id, + &response.body, + response.turn_state.clone(), + response.usage, + ) + .await; + let mut metadata = json!({ + "responseId": response_id, + "stream": false, + "transport": backend_transport_name(response.transport), + "firstTokenMs": response.first_token_ms, + "usage": response.usage, + }); + 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)); + insert_websocket_pool_decision(&mut metadata, response.websocket_pool_decision); + record_response_event(ResponseUsageRecord { + usage_records: &self.usage_records, + request_id, + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: &account.id, + route, + model: &display_model, + requested_model: Some(requested_model), + client_ip: request.client_ip.as_deref(), + client_user_agent: request.client_user_agent.as_deref(), + reasoning_effort: reasoning_effort_from_request(&request), + service_tier: request.service_tier(), + started_at, + status_code: 200, + message: "v1 responses completed", + metadata, + rate_limit_headers: &response.rate_limit_headers, + }) + .await; + Ok(body) + } + CollectedResponse::Failed(failure) => { + let error = ResponseDispatchError::Failed(failure); + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + Err(error) + } + CollectedResponse::MissingCompleted => { + let error = ResponseDispatchError::MissingCompleted; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + Err(error) + } + CollectedResponse::Empty => { + let error = ResponseDispatchError::EmptyUpstreamResponse; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: Some(&account.id), + stream: false, + compact: false, + transport: Some(backend_transport_name(response.transport)), + }, + &error, + ) + .await; + Err(error) + } + } + } + + pub(super) async fn record_response_dispatch_error( + &self, + request_id: &str, + route: &str, + requested_model: &str, + started_at: Instant, + details: ResponseDispatchErrorDetails<'_>, + error: &ResponseDispatchError, + ) { + record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { + ops_errors: &self.ops_errors, + request_id, + client_api_key_id: details.client_api_key_id, + account_id: details.account_id, + route, + model: requested_model, + started_at, + stream: details.stream, + compact: details.compact, + transport: details.transport, + error, + }) + .await; + } +} diff --git a/backend/src/dispatch/responses/service/stream.rs b/backend/src/dispatch/responses/service/stream.rs new file mode 100644 index 000000000..4013fd253 --- /dev/null +++ b/backend/src/dispatch/responses/service/stream.rs @@ -0,0 +1,603 @@ +use super::*; + +impl ResponseDispatchService { + /// 调度流式 Responses 请求到 Codex Responses 上游。 + pub async fn stream( + &self, + request_id: &str, + route: &str, + mut request: CodexResponsesRequest, + requested_model: &str, + ) -> Result { + let started_at = Instant::now(); + let catalog = self.models.catalog().await; + let display_model = catalog.resolve_model_id(requested_model); + request.set_model(display_model.clone()); + request.set_stream(true); + let tuple_schema = request.tuple_schema.clone(); + let now = Utc::now(); + let explicit_previous_response_id = request.previous_response_id().map(ToString::to_string); + let mut implicit_resume = self.prepare_response_session(&mut request).await; + let account_affinity = self.account_affinity_for_request(&request, now).await; + let mut acquire_request = AccountAcquireRequest::new(request.model(), now); + if let Some(preferred_account_id) = account_affinity.preferred_account_id() { + acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); + } + let mut excluded_account_ids = Vec::new(); + let mut exhausted_accounts = AccountExhaustionTracker::default(); + let mut history_recovery_used = false; + let mut next_required_account_id: Option = None; + let mut quota_verify_attempts = 0usize; + let mut trace = ResponseDispatchTrace::default(); + macro_rules! return_stream_dispatch_error { + ($error:expr) => {{ + let error = $error; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: exhausted_accounts.last_account_id(), + stream: true, + compact: false, + transport: Some(backend_transport_name( + backend_transport_for_response_request(&request), + )), + }, + &error, + ) + .await; + return Err(error); + }}; + ($error:expr, account_id: $account_id:expr, transport: $transport:expr) => {{ + let error = $error; + self.record_response_dispatch_error( + request_id, + route, + requested_model, + started_at, + ResponseDispatchErrorDetails { + client_api_key_id: request.client_api_key_id.as_deref(), + account_id: $account_id, + stream: true, + compact: false, + transport: $transport, + }, + &error, + ) + .await; + return Err(error); + }}; + } + loop { + let mut attempt_acquire_request = acquire_request + .clone() + .with_exclude_account_ids(excluded_account_ids.iter().cloned()); + if let Some(account_id) = next_required_account_id.take() { + attempt_acquire_request = + attempt_acquire_request.with_required_account_id(account_id); + } + let allow_quota_verify_account_retry = + attempt_acquire_request.required_account_id.is_none(); + attempt_acquire_request.now = Utc::now(); + let Some(acquired) = self + .account_pool + .acquire_with(&attempt_acquire_request) + .await + else { + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + return_stream_dispatch_error!(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, + installation_id: self.installation_id.as_deref(), + request_id, + excluded_account_ids: &mut excluded_account_ids, + verify_attempts: &mut quota_verify_attempts, + allow_retry_with_another_account: allow_quota_verify_account_retry, + }, + acquired, + ) + .await + { + QuotaVerificationDecision::Ready(acquired) => *acquired, + QuotaVerificationDecision::RetryWithAnotherAccount => { + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); + continue; + } + QuotaVerificationDecision::MaxAttemptsReached + | QuotaVerificationDecision::RequiredAccountUnavailable => { + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); + return_stream_dispatch_error!( + error, + account_id: Some(&acquired_account_id), + transport: Some(backend_transport_name(backend_transport_for_response_request( + &request + ))) + ); + } + }; + + self.apply_cascading_ban_defense( + &mut request, + &mut implicit_resume, + account_affinity.preferred_account_id(), + &acquired.account.id, + explicit_previous_response_id.as_deref(), + ) + .await; + Self::strip_history_if_account_changed( + &mut request, + &mut implicit_resume, + &account_affinity, + &acquired.account.id, + ); + + self.account_pool.wait_for_request_interval(&acquired).await; + let account = acquired.account; + let release_account_id = account.id.clone(); + let attempt = trace.start_attempt(&release_account_id); + let response_result = create_response_stream_with_account_retrying_5xx( + &self.codex, + self.installation_id.as_deref(), + &self.cloudflare, + &request, + request_id, + &account, + ) + .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) => { + let transport = response.transport; + let set_cookie_headers = response.set_cookie_headers; + let rate_limit_headers = response.rate_limit_headers; + let rate_limit_header_updates = response.rate_limit_header_updates; + let turn_state_update = response.turn_state_update; + let websocket_pool_decision = response.websocket_pool_decision; + let turn_state = response.turn_state; + let diagnostics = response.diagnostics; + self.cloudflare + .capture_set_cookie_headers(&release_account_id, &set_cookie_headers) + .await; + self.account_pool + .sync_passive_rate_limit_headers(&account, &rate_limit_headers) + .await; + let (prefetched, body) = match prefetch_first_sse_chunk(response.body).await { + Ok(prefetched) => prefetched, + Err(ResponseDispatchError::Upstream(error)) + if is_history_recovery_upstream_error(&error) + && !history_recovery_used => + { + self.account_pool.release(&release_account_id).await; + if client_error_invalid_reasoning_replay(&error) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + next_required_account_id = Some(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_rate_limit_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_quota_exhausted_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_auth_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_cloudflare_challenge_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_cloudflare_path_block_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_model_unsupported_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + let upstream_error = upstream_error_body(&error); + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { + return_stream_dispatch_error!( + ResponseDispatchError::from_exhausted_account(exhausted), + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); + excluded_account_ids.push(release_account_id); + continue; + } + Err(ResponseDispatchError::Upstream(error)) + if is_banned_upstream_error(&error) => + { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + continue; + } + Err(error) => { + self.account_pool.release(&release_account_id).await; + if let ResponseDispatchError::Upstream(upstream_error) = &error { + record_response_upstream_error_event( + ResponseUpstreamErrorEventRecord { + ops_errors: &self.ops_errors, + request_id, + account_id: &release_account_id, + account_email: account.email.as_deref(), + route, + model: requested_model, + started_at, + stream: true, + transport, + request: &request, + error: upstream_error, + trace: &trace, + attempt: Some(&attempt), + }, + ) + .await; + return Err(error); + } + return_stream_dispatch_error!( + error, + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + }; + let first_failure = match first_sse_failure(&prefetched) { + Ok(failure) => failure, + Err(error) => { + self.account_pool.release(&release_account_id).await; + return_stream_dispatch_error!( + ResponseDispatchError::InvalidSse(error), + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + }; + if let Some(failure) = first_failure { + if is_history_recovery_sse_failure(&failure) && !history_recovery_used { + self.account_pool.release(&release_account_id).await; + if sse_failure_invalid_reasoning_replay(&failure) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + next_required_account_id = Some(release_account_id); + continue; + } + if is_model_unsupported_sse_failure(&failure) { + let upstream_error = sse_failure_error_body(&failure); + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { + self.account_pool.release(&release_account_id).await; + return_stream_dispatch_error!( + ResponseDispatchError::from_exhausted_account(exhausted), + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(transport)) + ); + } + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + if is_quota_exhausted_sse_failure(&failure) { + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + failure.message.clone(), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + if is_auth_sse_failure(&failure) { + let upstream_error = sse_failure_error_body(&failure); + let account_status = auth_sse_failure_account_status(&failure); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(stream_failure_http_status(&failure)), + ); + self.account_pool + .set_status(&release_account_id, account_status) + .await; + excluded_account_ids.push(release_account_id); + self.account_pool.release(&account.id).await; + continue; + } + self.account_pool.release(&release_account_id).await; + record_prefetched_response_stream_failure_event( + ResponseStreamFailureEventRecord { + ops_errors: &self.ops_errors, + request_id, + account_id: &release_account_id, + route, + model: &display_model, + requested_model, + started_at, + transport, + request: &request, + failure: &failure, + diagnostics: &diagnostics, + rate_limit_headers: &rate_limit_headers, + prefetched: &prefetched, + trace: &trace, + attempt: &attempt, + }, + ) + .await; + return Err(ResponseDispatchError::Failed(failure.clone())); + } + + let context = LiveResponseStreamContext { + account_pool: Arc::clone(&self.account_pool), + session_affinity: Arc::clone(&self.session_affinity), + reasoning_replay: Arc::clone(&self.reasoning_replay), + usage_records: Arc::clone(&self.usage_records), + ops_errors: Arc::clone(&self.ops_errors), + cloudflare: self.cloudflare.clone(), + account_id: account.id, + account_plan_type: account.plan_type, + request_id: request_id.to_string(), + route: route.to_string(), + model: request.model().to_string(), + display_model: display_model.clone(), + requested_model: requested_model.to_string(), + client_ip: request.client_ip.clone(), + request, + tuple_schema, + transport, + rate_limit_headers, + rate_limit_header_updates, + turn_state_update, + websocket_pool_decision, + turn_state, + diagnostics, + attempt: attempt.clone(), + attempts: trace.attempts().to_vec(), + started_at, + }; + return Ok(spawn_live_response_stream(context, prefetched, body)); + } + Err(error) if is_rate_limit_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_quota_exhausted_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); + self.account_pool + .set_status(&release_account_id, AccountStatus::QuotaExhausted) + .await; + excluded_account_ids.push(release_account_id); + } + Err(error) + if is_history_recovery_upstream_error(&error) && !history_recovery_used => + { + self.account_pool.release(&release_account_id).await; + if client_error_invalid_reasoning_replay(&error) { + self.evict_reasoning_replay(&request, &release_account_id) + .await; + } + self.recover_request_history(&mut request, &mut implicit_resume) + .await; + history_recovery_used = true; + next_required_account_id = Some(release_account_id); + } + Err(error) if is_auth_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_challenge_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_cloudflare_path_block_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) if is_model_unsupported_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + let upstream_error = upstream_error_body(&error); + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { + return_stream_dispatch_error!( + ResponseDispatchError::from_exhausted_account(exhausted), + account_id: Some(&release_account_id), + transport: Some(backend_transport_name(backend_transport_for_response_request( + &request + ))) + ); + } + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); + excluded_account_ids.push(release_account_id); + } + Err(error) if is_banned_upstream_error(&error) => { + self.account_pool.release(&release_account_id).await; + 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; + excluded_account_ids.push(release_account_id); + } + Err(error) => { + self.account_pool.release(&release_account_id).await; + record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { + ops_errors: &self.ops_errors, + request_id, + account_id: &release_account_id, + account_email: account.email.as_deref(), + route, + model: requested_model, + started_at, + stream: true, + transport: backend_transport_for_response_request(&request), + request: &request, + error: &error, + trace: &trace, + attempt: Some(&attempt), + }) + .await; + return Err(ResponseDispatchError::Upstream(error)); + } + } + } + } +} diff --git a/backend/src/proxy/dispatch/responses/sse_failure.rs b/backend/src/dispatch/responses/sse_failure.rs similarity index 98% rename from backend/src/proxy/dispatch/responses/sse_failure.rs rename to backend/src/dispatch/responses/sse_failure.rs index 5296f4a13..383aa3626 100644 --- a/backend/src/proxy/dispatch/responses/sse_failure.rs +++ b/backend/src/dispatch/responses/sse_failure.rs @@ -1,12 +1,12 @@ use serde_json::Value; use crate::{ - proxy::dispatch::errors::{ + accounts::account::AccountStatus, + dispatch::errors::{ is_history_recovery_signal, is_invalid_encrypted_content_signal, is_model_unsupported_signal, }, - upstream::{ - accounts::model::AccountStatus, + upstream::openai::{ protocol::{ events::TokenUsage, responses::{response_from_codex_sse, CollectedResponse, ResponsesSseFailure}, diff --git a/backend/src/proxy/dispatch/responses/stream_lifecycle.rs b/backend/src/dispatch/responses/stream_lifecycle.rs similarity index 92% rename from backend/src/proxy/dispatch/responses/stream_lifecycle.rs rename to backend/src/dispatch/responses/stream_lifecycle.rs index f5d308c63..431541c35 100644 --- a/backend/src/proxy/dispatch/responses/stream_lifecycle.rs +++ b/backend/src/dispatch/responses/stream_lifecycle.rs @@ -4,20 +4,22 @@ use serde_json::Value; use tokio::sync::Mutex; use crate::{ - admin::monitoring::{ - ops_error_service::AdminOpsErrorLogService, usage_record_model::UsageRecordLevel, - usage_record_service::AdminUsageRecordService, + accounts::pool::RuntimeAccountPoolService, + api::client::responses::response_failed_sse_event_with_id, + dispatch::{ + affinity::{ + resolve::{evict_reasoning_replay, record_response_affinity}, + RuntimeSessionAffinityService, + }, + errors::backend_transport_name, + reasoning_replay::ReasoningReplayCache, }, infra::time::elapsed_millis_i64, - proxy::{ - dispatch::{ - errors::backend_transport_name, reasoning_replay::ReasoningReplayCache, - session_affinity::RuntimeSessionAffinityService, - }, - openai::responses::response_failed_sse_event_with_id, + telemetry::{ + ops::query::OpsQueryService, usage::query::UsageQueryService, + usage::types::UsageRecordLevel, }, - upstream::{ - accounts::pool::RuntimeAccountPoolService, + upstream::openai::{ protocol::{ events::extract_sse_usage, responses::{response_from_codex_sse, CodexResponsesRequest, CollectedResponse}, @@ -31,7 +33,6 @@ use crate::{ }; use super::{ - affinity::{evict_reasoning_replay, record_response_affinity}, event_recording::{ insert_first_token_ms, live_response_rate_limit_headers, live_response_turn_state, record_live_response_stream_event, @@ -43,14 +44,14 @@ use super::{ }, trace::ResponseDispatchAttempt, }; -use crate::proxy::dispatch::cloudflare::CloudflareRecovery; +use crate::dispatch::cloudflare::CloudflareRecovery; pub(super) struct LiveResponseStreamContext { pub(super) account_pool: Arc, pub(super) session_affinity: Arc, pub(super) reasoning_replay: Arc>, - pub(super) usage_records: Arc, - pub(super) ops_errors: Arc, + pub(super) usage_records: Arc, + pub(super) ops_errors: Arc, pub(super) cloudflare: CloudflareRecovery, pub(super) account_id: String, pub(super) account_plan_type: Option, diff --git a/backend/src/proxy/dispatch/responses/trace.rs b/backend/src/dispatch/responses/trace.rs similarity index 100% rename from backend/src/proxy/dispatch/responses/trace.rs rename to backend/src/dispatch/responses/trace.rs diff --git a/backend/src/proxy/dispatch/upstream.rs b/backend/src/dispatch/upstream.rs similarity index 98% rename from backend/src/proxy/dispatch/upstream.rs rename to backend/src/dispatch/upstream.rs index 7160a3b9a..f4a2cf91f 100644 --- a/backend/src/proxy/dispatch/upstream.rs +++ b/backend/src/dispatch/upstream.rs @@ -5,19 +5,19 @@ use std::{future::Future, time::Instant}; use serde_json::Value; use crate::{ - proxy::dispatch::{ - cloudflare::CloudflareRecovery, errors::is_retryable_upstream_5xx_error, - session_affinity::build_conversation_identity, - }, - upstream::accounts::{ - model::Account, + accounts::{ + account::Account, pool::{AcquiredAccount, RuntimeAccountPoolService}, quota::{ quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at, }, }, - upstream::{ + dispatch::{ + affinity::build_conversation_identity, cloudflare::CloudflareRecovery, + errors::is_retryable_upstream_5xx_error, + }, + upstream::openai::{ protocol::responses::{CodexCompactRequest, CodexResponsesRequest}, transport::{ CodexBackendClient, CodexBackendResponse, CodexBackendStreamingResponse, diff --git a/backend/src/http/mod.rs b/backend/src/http/mod.rs deleted file mode 100644 index 96452e1ab..000000000 --- a/backend/src/http/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -//! HTTP 中间件与路由层。 - -pub mod middleware; -pub mod router; diff --git a/backend/src/infra/database.rs b/backend/src/infra/database.rs index 4d9691f03..6ce752bcb 100644 --- a/backend/src/infra/database.rs +++ b/backend/src/infra/database.rs @@ -1,27 +1,15 @@ -//! SQLite 数据库连接与迁移初始化。 +//! PostgreSQL 连接与迁移初始化。 -use std::{fs, path::Path, str::FromStr, time::Duration}; +use std::time::Duration; -use sqlx::{sqlite::SqliteConnectOptions, SqlitePool}; +use sqlx::{postgres::PgPoolOptions, PgPool, Postgres, Transaction}; -const CURRENT_SCHEMA_VERSION: i64 = 3; -const MIGRATIONS: &[Migration] = &[ - Migration { - version: 1, - name: "initial", - sql: include_str!("migrations/0001_initial.sql"), - }, - Migration { - version: 2, - name: "rename_rotation_strategies", - sql: include_str!("migrations/0002_rename_rotation_strategies.sql"), - }, - Migration { - version: 3, - name: "ops_error_logs", - sql: include_str!("migrations/0003_ops_error_logs.sql"), - }, -]; +const CURRENT_SCHEMA_VERSION: i64 = 1; +const MIGRATIONS: &[Migration] = &[Migration { + version: 1, + name: "initial", + sql: include_str!("migrations/0001_initial.sql"), +}]; struct Migration { version: i64, @@ -29,23 +17,24 @@ struct Migration { sql: &'static str, } -/// 连接 SQLite 并应用未执行的迁移。 -pub async fn connect_sqlite(database_url: &str) -> Result { - ensure_sqlite_parent_dir(database_url)?; - let options = SqliteConnectOptions::from_str(database_url)? - .create_if_missing(true) - .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) - .foreign_keys(true) - .busy_timeout(Duration::from_secs(5)); - - let pool = SqlitePool::connect_with(options).await?; - - initialize_schema(&pool).await?; - +/// 连接 PostgreSQL 并应用未执行的迁移。 +pub async fn connect(database_url: &str) -> Result { + let pool = PgPoolOptions::new() + .max_connections(10) + .acquire_timeout(Duration::from_secs(5)) + .connect(database_url) + .await?; + migrate(&pool).await?; Ok(pool) } -async fn initialize_schema(pool: &SqlitePool) -> Result<(), sqlx::Error> { +/// 检查 PostgreSQL 连接是否可用。 +pub async fn ping(pool: &PgPool) -> Result<(), sqlx::Error> { + sqlx::query("select 1").execute(pool).await.map(|_| ()) +} + +/// 将 PostgreSQL schema 推进到当前版本。 +pub async fn migrate(pool: &PgPool) -> Result<(), sqlx::Error> { let mut tx = pool.begin().await?; ensure_migrations_table(&mut tx).await?; @@ -63,22 +52,15 @@ async fn initialize_schema(pool: &SqlitePool) -> Result<(), sqlx::Error> { record_migration(&mut tx, migration).await?; } - tx.commit().await?; - Ok(()) + tx.commit().await } -async fn ensure_migrations_table( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, -) -> Result<(), sqlx::Error> { - if table_exists(tx, "schema_migrations").await? { - return Ok(()); - } - +async fn ensure_migrations_table(tx: &mut Transaction<'_, Postgres>) -> Result<(), sqlx::Error> { sqlx::query( - "create table schema_migrations ( - version integer not null check (version > 0), + "create table if not exists schema_migrations ( + version bigint not null check (version > 0), name text not null, - applied_at text not null, + applied_at timestamptz not null default now(), primary key (version) )", ) @@ -87,24 +69,8 @@ async fn ensure_migrations_table( Ok(()) } -async fn table_exists( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, - table_name: &str, -) -> Result { - let row: Option<(i64,)> = sqlx::query_as( - "select 1 - from sqlite_master - where type = 'table' and name = ? - limit 1", - ) - .bind(table_name) - .fetch_optional(&mut **tx) - .await?; - Ok(row.is_some()) -} - async fn applied_migration_versions( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + tx: &mut Transaction<'_, Postgres>, ) -> Result, sqlx::Error> { let rows: Vec<(i64,)> = sqlx::query_as("select version from schema_migrations order by version") @@ -115,13 +81,12 @@ async fn applied_migration_versions( fn validate_applied_migrations(applied_versions: &[i64]) -> Result<(), sqlx::Error> { for version in applied_versions { - if MIGRATIONS + if !MIGRATIONS .iter() .any(|migration| migration.version == *version) { - continue; + return Err(unsupported_schema_version(*version)); } - return Err(unsupported_schema_version(*version)); } if let Some(version) = applied_versions .last() @@ -134,33 +99,32 @@ fn validate_applied_migrations(applied_versions: &[i64]) -> Result<(), sqlx::Err } async fn reject_unversioned_existing_schema( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + tx: &mut Transaction<'_, Postgres>, ) -> Result<(), sqlx::Error> { let (table_count,): (i64,) = sqlx::query_as( "select count(*) - from sqlite_master - where type = 'table' - and name <> 'schema_migrations' - and name not like 'sqlite_%'", + from information_schema.tables + where table_schema = current_schema() + and table_type = 'BASE TABLE' + and table_name <> 'schema_migrations'", ) .fetch_one(&mut **tx) .await?; if table_count > 0 { return Err(sqlx::Error::Protocol( - "unversioned sqlite schema is unsupported; recreate the database or add a migration" - .to_string(), + "unversioned PostgreSQL schema is unsupported; use an empty database".to_string(), )); } Ok(()) } async fn record_migration( - tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>, + tx: &mut Transaction<'_, Postgres>, migration: &Migration, ) -> Result<(), sqlx::Error> { sqlx::query( "insert into schema_migrations (version, name, applied_at) - values (?, ?, strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))", + values ($1, $2, now())", ) .bind(migration.version) .bind(migration.name) @@ -171,36 +135,6 @@ async fn record_migration( fn unsupported_schema_version(version: i64) -> sqlx::Error { sqlx::Error::Protocol(format!( - "unsupported sqlite schema version {version}; expected {CURRENT_SCHEMA_VERSION}" + "unsupported PostgreSQL schema version {version}; expected {CURRENT_SCHEMA_VERSION}" )) } - -fn ensure_sqlite_parent_dir(database_url: &str) -> Result<(), sqlx::Error> { - let Some(path) = sqlite_file_path(database_url) else { - return Ok(()); - }; - let Some(parent) = Path::new(path) - .parent() - .filter(|parent| !parent.as_os_str().is_empty()) - else { - return Ok(()); - }; - fs::create_dir_all(parent).map_err(sqlx::Error::Io) -} - -fn sqlite_file_path(database_url: &str) -> Option<&str> { - let path = database_url - .strip_prefix("sqlite://") - .or_else(|| database_url.strip_prefix("sqlite:"))? - .split_once('?') - .map_or_else( - || { - database_url - .strip_prefix("sqlite://") - .or_else(|| database_url.strip_prefix("sqlite:")) - .unwrap_or_default() - }, - |(path, _)| path, - ); - (!path.is_empty() && path != ":memory:").then_some(path) -} diff --git a/backend/src/infra/identity.rs b/backend/src/infra/identity.rs index b20434ba6..fdd05a88b 100644 --- a/backend/src/infra/identity.rs +++ b/backend/src/infra/identity.rs @@ -6,6 +6,7 @@ use argon2::{ }; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use rand::Rng; +use sha2::{Digest, Sha256}; use thiserror::Error; // --------------------------------------------------------------------------- @@ -73,3 +74,15 @@ pub fn generate_admin_api_key() -> String { rand::rng().fill_bytes(&mut bytes); format!("admin-{}", hex::encode(bytes)) } + +/// 生成新的管理员会话令牌。 +pub fn generate_admin_session_token() -> String { + let mut bytes = [0u8; 32]; + rand::rng().fill_bytes(&mut bytes); + format!("sess_{}", URL_SAFE_NO_PAD.encode(bytes)) +} + +/// 计算高熵凭据的 SHA-256 十六进制摘要。 +pub fn hash_credential(value: &str) -> String { + hex::encode(Sha256::digest(value.as_bytes())) +} diff --git a/backend/src/infra/migrations/0001_initial.sql b/backend/src/infra/migrations/0001_initial.sql index 130c690de..f16e15616 100644 --- a/backend/src/infra/migrations/0001_initial.sql +++ b/backend/src/infra/migrations/0001_initial.sql @@ -1,62 +1,41 @@ -pragma foreign_keys = on; - --- ============================================ --- Admin Users --- ============================================ create table admin_users ( id text primary key, password_hash text not null, - created_at text not null, - updated_at text not null + created_at timestamptz not null, + updated_at timestamptz not null ); --- ============================================ --- Admin Sessions --- ============================================ -create table admin_sessions ( - id text primary key, - user_id text not null references admin_users(id) on delete cascade, - expires_at text not null, - created_at text not null -); - -create index idx_admin_sessions_expires on admin_sessions(expires_at); - --- ============================================ --- Client API Keys --- ============================================ create table client_api_keys ( id text primary key, name text not null, prefix text not null, key text not null unique, label text, - enabled integer not null default 1 check (enabled in (0, 1)), - created_at text not null, - last_used_at text + enabled boolean not null default true, + created_at timestamptz not null, + last_used_at timestamptz ); -create index idx_client_api_keys_key_enabled on client_api_keys(key) where enabled = 1; -create index idx_client_api_keys_created_id on client_api_keys(created_at desc, id desc); +create index idx_client_api_keys_created_id + on client_api_keys(created_at desc, id desc); --- ============================================ --- Runtime Settings --- ============================================ create table runtime_settings ( - id integer primary key check (id = 1), - model_aliases_json text not null default '{}', - refresh_margin_seconds integer not null check (refresh_margin_seconds > 0), - refresh_concurrency integer not null check (refresh_concurrency > 0), - max_concurrent_per_account integer not null check (max_concurrent_per_account > 0), - request_interval_ms integer not null check (request_interval_ms >= 0), - rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')), - admin_api_key text, - updated_at text not null + id bigint primary key check (id = 1), + model_aliases_json jsonb not null default '{}', + refresh_margin_seconds bigint not null check (refresh_margin_seconds > 0), + refresh_concurrency bigint not null check (refresh_concurrency > 0), + max_concurrent_per_account bigint not null check (max_concurrent_per_account > 0), + request_interval_ms bigint not null check (request_interval_ms >= 0), + rotation_strategy text not null check ( + rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky') + ), + admin_api_key_hash text, + usage_retention_days bigint not null default 30 check (usage_retention_days > 0), + ops_error_retention_days bigint not null default 30 check (ops_error_retention_days > 0), + bucket_retention_days bigint not null default 90 check (bucket_retention_days > 0), + updated_at timestamptz not null ); --- ============================================ --- Accounts --- ============================================ create table accounts ( id text primary key, email text, @@ -66,119 +45,198 @@ create table accounts ( plan_type text, access_token text not null, refresh_token text, - access_token_expires_at text, - next_refresh_at text, - status text not null check (status in ('active', 'expired', 'quota_exhausted', 'disabled', 'banned')), - quota_json text, - quota_fetched_at text, - quota_limit_reached integer not null default 0 check (quota_limit_reached in (0, 1)), - quota_verify_required integer not null default 0 check (quota_verify_required in (0, 1)), - quota_cooldown_until text, - cloudflare_cooldown_until text, - added_at text not null, - updated_at text not null + access_token_expires_at timestamptz, + next_refresh_at timestamptz, + status text not null check ( + status in ('active', 'expired', 'quota_exhausted', 'disabled', 'banned') + ), + quota_json jsonb, + quota_fetched_at timestamptz, + quota_limit_reached boolean not null default false, + quota_verify_required boolean not null default false, + quota_cooldown_until timestamptz, + cloudflare_cooldown_until timestamptz, + added_at timestamptz not null, + updated_at timestamptz not null ); create index idx_accounts_status on accounts(status); create index idx_accounts_added_id on accounts(added_at desc, id desc); create unique index ux_accounts_chatgpt_identity -on accounts(chatgpt_account_id, coalesce(chatgpt_user_id, '')) -where chatgpt_account_id is not null; + on accounts(chatgpt_account_id, coalesce(chatgpt_user_id, '')) + where chatgpt_account_id is not null; --- ============================================ --- Account Refresh Leases --- ============================================ -create table account_refresh_leases ( - account_id text primary key references accounts(id) on delete cascade, - owner text not null, - expires_at text not null, - updated_at text not null -); - -create index idx_account_refresh_leases_expires on account_refresh_leases(expires_at); - --- ============================================ --- Account Usage Statistics --- ============================================ create table account_usage ( account_id text primary key references accounts(id) on delete cascade, - request_count integer not null default 0 check (request_count >= 0), - empty_response_count integer not null default 0 check (empty_response_count >= 0), - input_tokens integer not null default 0 check (input_tokens >= 0), - output_tokens integer not null default 0 check (output_tokens >= 0), - cached_tokens integer not null default 0 check (cached_tokens >= 0), - reasoning_tokens integer not null default 0 check (reasoning_tokens >= 0), - total_tokens integer not null default 0 check (total_tokens >= 0), - image_input_tokens integer not null default 0 check (image_input_tokens >= 0), - image_output_tokens integer not null default 0 check (image_output_tokens >= 0), - image_request_count integer not null default 0 check (image_request_count >= 0), - image_request_failed_count integer not null default 0 check (image_request_failed_count >= 0), - window_request_count integer not null default 0 check (window_request_count >= 0), - window_input_tokens integer not null default 0 check (window_input_tokens >= 0), - window_output_tokens integer not null default 0 check (window_output_tokens >= 0), - window_cached_tokens integer not null default 0 check (window_cached_tokens >= 0), - window_image_input_tokens integer not null default 0 check (window_image_input_tokens >= 0), - window_image_output_tokens integer not null default 0 check (window_image_output_tokens >= 0), - window_image_request_count integer not null default 0 check (window_image_request_count >= 0), - window_image_request_failed_count integer not null default 0 check (window_image_request_failed_count >= 0), - window_started_at text, - window_reset_at text, - limit_window_seconds integer check (limit_window_seconds is null or limit_window_seconds > 0), - last_used_at text + request_count bigint not null default 0 check (request_count >= 0), + empty_response_count bigint not null default 0 check (empty_response_count >= 0), + input_tokens bigint not null default 0 check (input_tokens >= 0), + output_tokens bigint not null default 0 check (output_tokens >= 0), + cached_tokens bigint not null default 0 check (cached_tokens >= 0), + reasoning_tokens bigint not null default 0 check (reasoning_tokens >= 0), + total_tokens bigint not null default 0 check (total_tokens >= 0), + image_input_tokens bigint not null default 0 check (image_input_tokens >= 0), + image_output_tokens bigint not null default 0 check (image_output_tokens >= 0), + image_request_count bigint not null default 0 check (image_request_count >= 0), + image_request_failed_count bigint not null default 0 check (image_request_failed_count >= 0), + window_request_count bigint not null default 0 check (window_request_count >= 0), + window_input_tokens bigint not null default 0 check (window_input_tokens >= 0), + window_output_tokens bigint not null default 0 check (window_output_tokens >= 0), + window_cached_tokens bigint not null default 0 check (window_cached_tokens >= 0), + window_image_input_tokens bigint not null default 0 check (window_image_input_tokens >= 0), + window_image_output_tokens bigint not null default 0 check (window_image_output_tokens >= 0), + window_image_request_count bigint not null default 0 check (window_image_request_count >= 0), + window_image_request_failed_count bigint not null default 0 check ( + window_image_request_failed_count >= 0 + ), + window_started_at timestamptz, + window_reset_at timestamptz, + limit_window_seconds bigint check ( + limit_window_seconds is null or limit_window_seconds > 0 + ), + last_used_at timestamptz ); create index idx_account_usage_last_used_account -on account_usage(last_used_at desc, account_id desc); + on account_usage(last_used_at desc, account_id desc); --- ============================================ --- Account Model Usage Statistics --- ============================================ create table account_model_usage ( account_id text not null references accounts(id) on delete cascade, model text not null, - request_count integer not null default 0 check (request_count >= 0), - error_count integer not null default 0 check (error_count >= 0), - input_tokens integer not null default 0 check (input_tokens >= 0), - output_tokens integer not null default 0 check (output_tokens >= 0), - cached_tokens integer not null default 0 check (cached_tokens >= 0), - last_used_at text, + request_count bigint not null default 0 check (request_count >= 0), + error_count bigint not null default 0 check (error_count >= 0), + input_tokens bigint not null default 0 check (input_tokens >= 0), + output_tokens bigint not null default 0 check (output_tokens >= 0), + cached_tokens bigint not null default 0 check (cached_tokens >= 0), + last_used_at timestamptz, primary key (account_id, model) ); create index idx_account_model_usage_last_used -on account_model_usage(last_used_at desc, account_id, model); + on account_model_usage(last_used_at desc, account_id, model); --- ============================================ --- Usage Time Buckets --- ============================================ -create table usage_time_buckets ( - bucket_start text not null, - account_id text not null default '', - model text not null default '', - service_tier text not null default '', - request_count integer not null default 0 check (request_count >= 0), - error_count integer not null default 0 check (error_count >= 0), - input_tokens integer not null default 0 check (input_tokens >= 0), - output_tokens integer not null default 0 check (output_tokens >= 0), - cached_tokens integer not null default 0 check (cached_tokens >= 0), - first_token_latency_sum integer not null default 0 check (first_token_latency_sum >= 0), - first_token_latency_count integer not null default 0 check (first_token_latency_count >= 0), - latency_sum integer not null default 0 check (latency_sum >= 0), - latency_count integer not null default 0 check (latency_count >= 0), - max_latency_ms integer not null default 0 check (max_latency_ms >= 0), - min_latency_ms integer not null default 0 check (min_latency_ms >= 0), - updated_at text not null, - primary key (bucket_start, account_id, model, service_tier) +create table usage_records ( + id text primary key, + request_id text, + client_api_key_id text, + kind text not null, + route text, + provider text not null, + account_id text not null, + model text not null, + requested_model text, + upstream_model text, + service_tier text, + status_code integer not null check (status_code between 200 and 399), + transport text, + attempt_index bigint check (attempt_index is null or attempt_index >= 0), + response_id text, + upstream_request_id text, + latency_ms bigint check (latency_ms is null or latency_ms >= 0), + first_token_ms bigint check (first_token_ms is null or first_token_ms >= 0), + input_tokens bigint check (input_tokens is null or input_tokens >= 0), + output_tokens bigint check (output_tokens is null or output_tokens >= 0), + cached_tokens bigint check (cached_tokens is null or cached_tokens >= 0), + reasoning_tokens bigint check (reasoning_tokens is null or reasoning_tokens >= 0), + message text not null, + metadata_json jsonb not null, + created_at timestamptz not null ); -create index idx_usage_time_buckets_bucket -on usage_time_buckets(bucket_start); -create index idx_usage_time_buckets_model_bucket -on usage_time_buckets(model, bucket_start); +create index idx_usage_records_created_id on usage_records(created_at desc, id desc); +create index idx_usage_records_request_id + on usage_records(request_id) where request_id is not null; +create index idx_usage_records_kind_created + on usage_records(kind, created_at desc); +create index idx_usage_records_account_created + on usage_records(account_id, created_at desc); +create index idx_usage_records_model_created + on usage_records(model, created_at desc); +create index idx_usage_records_key_created + on usage_records(client_api_key_id, created_at desc) + where client_api_key_id is not null; +create index idx_usage_records_response_id + on usage_records(response_id) where response_id is not null; +create index idx_usage_records_upstream_request_id + on usage_records(upstream_request_id) where upstream_request_id is not null; + +create table ops_error_logs ( + id text primary key, + request_id text, + client_api_key_id text, + kind text not null, + provider text, + account_id text, + route text, + model text, + status_code integer check ( + status_code is null or status_code between 100 and 599 + ), + client_status_code integer check ( + client_status_code is null or client_status_code between 100 and 599 + ), + upstream_status_code integer check ( + upstream_status_code is null or upstream_status_code between 100 and 599 + ), + transport text, + attempt_index bigint check (attempt_index is null or attempt_index >= 0), + failure_class text, + response_id text, + upstream_request_id text, + latency_ms bigint check (latency_ms is null or latency_ms >= 0), + message text not null, + metadata_json jsonb not null, + created_at timestamptz not null +); + +create index idx_ops_error_logs_created_id + on ops_error_logs(created_at desc, id desc); +create index idx_ops_error_logs_request_id + on ops_error_logs(request_id) where request_id is not null; +create index idx_ops_error_logs_key_created + on ops_error_logs(client_api_key_id, created_at desc) + where client_api_key_id is not null; +create index idx_ops_error_logs_account + on ops_error_logs(account_id, created_at desc) where account_id is not null; +create index idx_ops_error_logs_route_created + on ops_error_logs(route, created_at desc) where route is not null; +create index idx_ops_error_logs_model_created + on ops_error_logs(model, created_at desc) where model is not null; +create index idx_ops_error_logs_status_created + on ops_error_logs(status_code, created_at desc) where status_code is not null; +create index idx_ops_error_logs_transport_created + on ops_error_logs(transport, created_at desc) where transport is not null; +create index idx_ops_error_logs_failure_class + on ops_error_logs(failure_class, created_at desc) where failure_class is not null; +create index idx_ops_error_logs_response_id + on ops_error_logs(response_id) where response_id is not null; +create index idx_ops_error_logs_upstream_request_id + on ops_error_logs(upstream_request_id) where upstream_request_id is not null; + +create table request_time_buckets ( + bucket_start timestamptz not null, + provider text not null default '__unknown__', + account_id text not null default '__unknown__', + model text not null default '__unknown__', + service_tier text not null default '__unknown__', + success_count bigint not null default 0 check (success_count >= 0), + error_count bigint not null default 0 check (error_count >= 0), + input_tokens bigint not null default 0 check (input_tokens >= 0), + output_tokens bigint not null default 0 check (output_tokens >= 0), + cached_tokens bigint not null default 0 check (cached_tokens >= 0), + first_token_latency_sum bigint not null default 0 check (first_token_latency_sum >= 0), + first_token_latency_count bigint not null default 0 check (first_token_latency_count >= 0), + latency_sum bigint not null default 0 check (latency_sum >= 0), + latency_count bigint not null default 0 check (latency_count >= 0), + max_latency_ms bigint not null default 0 check (max_latency_ms >= 0), + min_latency_ms bigint check (min_latency_ms is null or min_latency_ms >= 0), + updated_at timestamptz not null, + primary key (bucket_start, provider, account_id, model, service_tier) +); + +create index idx_request_time_buckets_model + on request_time_buckets(model, bucket_start); --- ============================================ --- Account Cookies --- ============================================ create table account_cookies ( id text primary key, account_id text not null references accounts(id) on delete cascade, @@ -186,20 +244,16 @@ create table account_cookies ( name text not null, value text not null, path text not null default '/', - expires_at text, - updated_at text not null, + expires_at timestamptz, + updated_at timestamptz not null, unique(account_id, domain, name, path) ); -create index idx_account_cookies_account on account_cookies(account_id); -create index idx_account_cookies_account_domain on account_cookies(account_id, domain); +create index idx_account_cookies_account_domain + on account_cookies(account_id, domain); create index idx_account_cookies_expires -on account_cookies(expires_at) -where expires_at is not null; + on account_cookies(expires_at) where expires_at is not null; --- ============================================ --- Device Fingerprints --- ============================================ create table fingerprints ( id text primary key, originator text not null, @@ -209,11 +263,11 @@ create table fingerprints ( arch text not null, chromium_version text not null, user_agent_template text not null, - default_headers_json text not null, - header_order_json text not null, + default_headers_json jsonb not null, + header_order_json jsonb not null, source text not null, - created_at text not null, - updated_at text not null + created_at timestamptz not null, + updated_at timestamptz not null ); create table fingerprint_update_history ( @@ -223,76 +277,11 @@ create table fingerprint_update_history ( build_number text not null, chromium_version text, source text not null, - manifest_json text, - created_at text not null + manifest_json jsonb, + created_at timestamptz not null ); create index idx_fingerprint_update_history_created_id -on fingerprint_update_history(created_at desc, id desc); - --- ============================================ --- Usage Records --- ============================================ -create table usage_records ( - id text primary key, - request_id text, - kind text not null, - level text not null check (level in ('debug', 'info', 'warn', 'error')), - account_id text, - route text, - model text, - status_code integer check (status_code is null or (status_code >= 100 and status_code <= 599)), - transport text, - attempt_index integer check (attempt_index is null or attempt_index >= 0), - upstream_status_code integer check (upstream_status_code is null or (upstream_status_code >= 100 and upstream_status_code <= 599)), - failure_class text, - response_id text, - upstream_request_id text, - latency_ms integer check (latency_ms is null or latency_ms >= 0), - message text not null, - metadata_json text not null, - created_at text not null -); - -create index idx_usage_records_created_id on usage_records(created_at desc, id desc); -create index idx_usage_records_kind_created on usage_records(kind, created_at desc); -create index idx_usage_records_request_id on usage_records(request_id); -create index idx_usage_records_account on usage_records(account_id, created_at desc) where account_id is not null; -create index idx_usage_records_transport on usage_records(transport, created_at desc) where transport is not null; -create index idx_usage_records_failure_class on usage_records(failure_class, created_at desc) where failure_class is not null; -create index idx_usage_records_response_id on usage_records(response_id) where response_id is not null; -create index idx_usage_records_upstream_request_id on usage_records(upstream_request_id) where upstream_request_id is not null; -create index idx_usage_records_level_created on usage_records(level, created_at desc); -create index idx_usage_records_route_created on usage_records(route, created_at desc) where route is not null; -create index idx_usage_records_model_created on usage_records(model, created_at desc) where model is not null; -create index idx_usage_records_status_created on usage_records(status_code, created_at desc) where status_code is not null; -create index idx_usage_records_upstream_status_created on usage_records(upstream_status_code, created_at desc) where upstream_status_code is not null; - --- ============================================ --- Model Plan Snapshots --- ============================================ -create table model_plan_snapshots ( - plan_type text primary key, - models_json text not null, - fetched_at text not null -); - --- ============================================ --- Session Affinities --- ============================================ -create table session_affinities ( - response_id text primary key, - account_id text not null references accounts(id) on delete cascade, - conversation_id text not null, - turn_state text, - instructions_hash text, - input_tokens integer check (input_tokens is null or input_tokens >= 0), - function_call_ids_json text not null default '[]', - variant_hash text, - expires_at text not null, - created_at text not null -); - -create index idx_session_affinities_conversation on session_affinities(conversation_id, created_at desc); -create index idx_session_affinities_expires on session_affinities(expires_at); -create index idx_session_affinities_active_order on session_affinities(expires_at, created_at, response_id); + on fingerprint_update_history(created_at desc, id desc); +create index idx_fingerprint_update_history_fingerprint + on fingerprint_update_history(current_fingerprint_id); diff --git a/backend/src/infra/migrations/0002_rename_rotation_strategies.sql b/backend/src/infra/migrations/0002_rename_rotation_strategies.sql deleted file mode 100644 index a18de3c42..000000000 --- a/backend/src/infra/migrations/0002_rename_rotation_strategies.sql +++ /dev/null @@ -1,40 +0,0 @@ -create table runtime_settings_new ( - id integer primary key check (id = 1), - model_aliases_json text not null default '{}', - refresh_margin_seconds integer not null check (refresh_margin_seconds > 0), - refresh_concurrency integer not null check (refresh_concurrency > 0), - max_concurrent_per_account integer not null check (max_concurrent_per_account > 0), - request_interval_ms integer not null check (request_interval_ms >= 0), - rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')), - admin_api_key text, - updated_at text not null -); - -insert into runtime_settings_new ( - id, - model_aliases_json, - refresh_margin_seconds, - refresh_concurrency, - max_concurrent_per_account, - request_interval_ms, - rotation_strategy, - admin_api_key, - updated_at -) -select - id, - model_aliases_json, - refresh_margin_seconds, - refresh_concurrency, - max_concurrent_per_account, - request_interval_ms, - case rotation_strategy - when 'least_used' then 'quota_reset_priority' - else rotation_strategy - end, - admin_api_key, - updated_at -from runtime_settings; - -drop table runtime_settings; -alter table runtime_settings_new rename to runtime_settings; diff --git a/backend/src/infra/migrations/0003_ops_error_logs.sql b/backend/src/infra/migrations/0003_ops_error_logs.sql deleted file mode 100644 index 6e6ee8995..000000000 --- a/backend/src/infra/migrations/0003_ops_error_logs.sql +++ /dev/null @@ -1,31 +0,0 @@ -create table ops_error_logs ( - id text primary key, - request_id text, - kind text not null, - account_id text, - route text, - model text, - status_code integer check (status_code is null or (status_code >= 100 and status_code <= 599)), - client_status_code integer check (client_status_code is null or (client_status_code >= 100 and client_status_code <= 599)), - upstream_status_code integer check (upstream_status_code is null or (upstream_status_code >= 100 and upstream_status_code <= 599)), - transport text, - attempt_index integer check (attempt_index is null or attempt_index >= 0), - failure_class text, - response_id text, - upstream_request_id text, - latency_ms integer check (latency_ms is null or latency_ms >= 0), - message text not null, - metadata_json text not null, - created_at text not null -); - -create index idx_ops_error_logs_created_id on ops_error_logs(created_at desc, id desc); -create index idx_ops_error_logs_request_id on ops_error_logs(request_id) where request_id is not null; -create index idx_ops_error_logs_account on ops_error_logs(account_id, created_at desc) where account_id is not null; -create index idx_ops_error_logs_route_created on ops_error_logs(route, created_at desc) where route is not null; -create index idx_ops_error_logs_model_created on ops_error_logs(model, created_at desc) where model is not null; -create index idx_ops_error_logs_status_created on ops_error_logs(status_code, created_at desc) where status_code is not null; -create index idx_ops_error_logs_transport_created on ops_error_logs(transport, created_at desc) where transport is not null; -create index idx_ops_error_logs_failure_class on ops_error_logs(failure_class, created_at desc) where failure_class is not null; -create index idx_ops_error_logs_response_id on ops_error_logs(response_id) where response_id is not null; -create index idx_ops_error_logs_upstream_request_id on ops_error_logs(upstream_request_id) where upstream_request_id is not null; diff --git a/backend/src/infra/mod.rs b/backend/src/infra/mod.rs index 30a1700ba..20081a5c9 100644 --- a/backend/src/infra/mod.rs +++ b/backend/src/infra/mod.rs @@ -2,11 +2,11 @@ //! 平台基础设施层，承载身份、存储、日志和 JSON 原语。 -/// 数据库连接与 SQLite 管理。 +/// PostgreSQL 连接与迁移管理。 pub mod database; /// 展示格式化辅助。 pub mod format; -/// 身份验证原语（管理员密码、API Key 哈希）。 +/// 身份验证原语（管理员密码、凭据生成与哈希）。 pub mod identity; /// JSON 序列化和分页原语。 pub mod json; @@ -14,5 +14,6 @@ pub mod json; pub mod logging; /// 路径和安装 ID 辅助。 pub mod paths; +pub mod redis; /// 时间格式化辅助。 pub mod time; diff --git a/backend/src/infra/redis.rs b/backend/src/infra/redis.rs new file mode 100644 index 000000000..43a4e59e1 --- /dev/null +++ b/backend/src/infra/redis.rs @@ -0,0 +1,54 @@ +//! Redis 运行态连接与键前缀。 + +use redis::{aio::ConnectionManager, AsyncCommands}; + +/// Redis 运行态连接。 +#[derive(Clone)] +pub struct RedisConnection { + manager: ConnectionManager, + key_prefix: String, +} + +impl RedisConnection { + /// 建立 Redis 连接。 + pub async fn connect(url: &str, key_prefix: impl Into) -> redis::RedisResult { + let client = redis::Client::open(url)?; + let manager = client.get_connection_manager().await?; + Ok(Self { + manager, + key_prefix: normalize_key_prefix(key_prefix.into()), + }) + } + + /// 返回可独立使用的多路复用连接管理器。 + pub fn manager(&self) -> ConnectionManager { + self.manager.clone() + } + + /// 为业务键添加应用前缀。 + pub fn key(&self, suffix: &str) -> String { + format!("{}{}", self.key_prefix, suffix) + } + + /// 检查 Redis 连接是否可用。 + pub async fn ping(&self) -> redis::RedisResult<()> { + let mut connection = self.manager.clone(); + let response: String = connection.ping().await?; + if response == "PONG" { + Ok(()) + } else { + Err(redis::RedisError::from(( + redis::ErrorKind::UnexpectedReturnType, + "unexpected Redis PING response", + response, + ))) + } + } +} + +fn normalize_key_prefix(mut prefix: String) -> String { + if !prefix.ends_with(':') { + prefix.push(':'); + } + prefix +} diff --git a/backend/src/keys/manage.rs b/backend/src/keys/manage.rs new file mode 100644 index 000000000..823fb7fbf --- /dev/null +++ b/backend/src/keys/manage.rs @@ -0,0 +1,101 @@ +//! 客户端 API Key 管理服务。 + +use crate::infra::json::Page; + +use super::{ + store::PgClientKeyStore, + types::{ + parse_client_key_status, BatchDeleteClientApiKeys, KeyManageError, ManagedClientApiKey, + }, +}; + +#[derive(Clone)] +pub struct KeyManageService { + store: PgClientKeyStore, +} + +impl KeyManageService { + pub fn new(store: PgClientKeyStore) -> Self { + Self { store } + } + + pub async fn create(&self, name: &str) -> Result { + let name = name.trim(); + if name.is_empty() { + return Err(KeyManageError::EmptyName); + } + self.store + .create(name) + .await + .map(Into::into) + .map_err(|_| KeyManageError::Create) + } + + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, KeyManageError> { + let page = self + .store + .list(cursor, limit) + .await + .map_err(|_| KeyManageError::List)?; + Ok(Page { + items: page.items.into_iter().map(Into::into).collect(), + next_cursor: page.next_cursor, + }) + } + + pub async fn get(&self, key_id: &str) -> Result, KeyManageError> { + self.store + .get(key_id) + .await + .map(|key| key.map(Into::into)) + .map_err(|_| KeyManageError::List) + } + + pub async fn update_label( + &self, + key_id: &str, + label: Option, + ) -> Result, KeyManageError> { + if label + .as_ref() + .is_some_and(|value| value.chars().count() > 64) + { + return Err(KeyManageError::LabelTooLong); + } + self.store + .set_label(key_id, label) + .await + .map(|key| key.map(Into::into)) + .map_err(|_| KeyManageError::UpdateLabel) + } + + pub async fn update_status(&self, key_id: &str, status: &str) -> Result { + self.store + .set_enabled(key_id, parse_client_key_status(status)?) + .await + .map_err(|_| KeyManageError::UpdateStatus) + } + + pub async fn batch_delete( + &self, + ids: Vec, + ) -> Result { + if ids.is_empty() { + return Err(KeyManageError::EmptyIds); + } + let mut deleted = 0u32; + let mut not_found = Vec::new(); + for id in ids { + match self.store.delete(&id).await { + Ok(true) => deleted += 1, + Ok(false) => not_found.push(id), + Err(_) => return Err(KeyManageError::Delete), + } + } + Ok(BatchDeleteClientApiKeys { deleted, not_found }) + } +} diff --git a/backend/src/keys/mod.rs b/backend/src/keys/mod.rs new file mode 100644 index 000000000..ac7230300 --- /dev/null +++ b/backend/src/keys/mod.rs @@ -0,0 +1,6 @@ +//! 客户端 API Key 领域。 + +pub mod manage; +pub mod service; +pub mod store; +pub mod types; diff --git a/backend/src/keys/service.rs b/backend/src/keys/service.rs new file mode 100644 index 000000000..8b56f9a6f --- /dev/null +++ b/backend/src/keys/service.rs @@ -0,0 +1,120 @@ +//! 客户端 API Key 鉴权服务。 + +use std::{ + collections::BTreeMap, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, + time::Duration, +}; + +use chrono::{DateTime, Utc}; +use thiserror::Error; +use tokio::time::sleep; + +use super::store::{PgClientKeyStore, PgClientKeyStoreError}; + +const CLIENT_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1); + +#[derive(Debug, Error)] +pub enum ClientKeyStoreError { + #[error("client key store operation failed: {message}")] + OperationFailed { message: String }, +} + +pub type ClientKeyStoreResult = Result; + +/// 客户端 API Key 验证服务。鉴权始终走 PostgreSQL 唯一索引点查。 +#[derive(Clone)] +pub struct KeyVerifier { + store: PgClientKeyStore, + pending_last_used: Arc>>>, + flush_scheduled: Arc, + flush_delay: Duration, +} + +impl KeyVerifier { + pub fn new(store: PgClientKeyStore) -> Self { + Self { + store, + pending_last_used: Arc::new(Mutex::new(BTreeMap::new())), + flush_scheduled: Arc::new(AtomicBool::new(false)), + flush_delay: CLIENT_KEY_LAST_USED_FLUSH_DELAY, + } + } + + /// 验证明文 key；成功时返回稳定 key ID，供请求事实归因。 + pub async fn verify(&self, plaintext: &str) -> ClientKeyStoreResult> { + if !plaintext.starts_with("sk_") { + return Ok(None); + } + let key_id = self + .store + .find_enabled_id_by_key(plaintext) + .await + .map_err(map_store_error)?; + if let Some(key_id) = &key_id { + self.queue_last_used_touch(key_id.clone()); + } + Ok(key_id) + } + + fn queue_last_used_touch(&self, id: String) { + self.pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(id, Utc::now()); + self.schedule_flush(); + } + + fn schedule_flush(&self) { + if self + .flush_scheduled + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return; + } + let service = self.clone(); + tokio::spawn(async move { + sleep(service.flush_delay).await; + service.flush_pending_last_used().await; + }); + } + + pub async fn flush_pending_last_used(&self) { + let updates = self.take_pending_last_used(); + if let Err(error) = self.store.touch_last_used_batch(&updates).await { + tracing::warn!(error = %error, "failed to flush client key last_used_at batch"); + self.pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .extend(updates); + } + self.flush_scheduled.store(false, Ordering::Release); + if !self + .pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty() + { + self.schedule_flush(); + } + } + + fn take_pending_last_used(&self) -> BTreeMap> { + std::mem::take( + &mut *self + .pending_last_used + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ) + } +} + +fn map_store_error(error: PgClientKeyStoreError) -> ClientKeyStoreError { + ClientKeyStoreError::OperationFailed { + message: error.to_string(), + } +} diff --git a/backend/src/keys/store.rs b/backend/src/keys/store.rs new file mode 100644 index 000000000..e2a0ac7f9 --- /dev/null +++ b/backend/src/keys/store.rs @@ -0,0 +1,227 @@ +//! 客户端 API Key PostgreSQL 存储。 + +use std::collections::BTreeMap; + +use chrono::{DateTime, Utc}; +use sqlx::{PgPool, Row}; +use thiserror::Error; +use uuid::Uuid; + +use crate::infra::{ + identity::generate_client_api_key, + json::{decode_cursor, encode_cursor, Page}, + time::parse_rfc3339_utc, +}; + +#[derive(Debug, Error)] +pub enum PgClientKeyStoreError { + #[error("PostgreSQL client key operation failed: {0}")] + Database(#[from] sqlx::Error), + #[error("invalid client key pagination cursor")] + InvalidCursor, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StoredClientApiKey { + pub id: String, + pub name: String, + pub label: Option, + pub prefix: String, + pub key: String, + pub enabled: bool, + pub created_at: DateTime, + pub last_used_at: Option>, +} + +#[derive(Clone)] +pub struct PgClientKeyStore { + pool: PgPool, +} + +impl PgClientKeyStore { + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + pub fn pool(&self) -> &PgPool { + &self.pool + } + + pub(super) async fn find_enabled_id_by_key( + &self, + key: &str, + ) -> Result, PgClientKeyStoreError> { + Ok( + sqlx::query_scalar("select id from client_api_keys where key = $1 and enabled") + .bind(key) + .fetch_optional(&self.pool) + .await?, + ) + } + + pub(super) async fn create( + &self, + name: &str, + ) -> Result { + let generated = generate_client_api_key(); + let stored = StoredClientApiKey { + id: format!("key_{}", Uuid::new_v4().simple()), + name: name.to_string(), + label: None, + prefix: generated.prefix, + key: generated.key, + enabled: true, + created_at: Utc::now(), + last_used_at: None, + }; + sqlx::query( + r" +insert into client_api_keys + (id, name, label, prefix, key, enabled, created_at, last_used_at) +values ($1, $2, null, $3, $4, true, $5, null)", + ) + .bind(&stored.id) + .bind(&stored.name) + .bind(&stored.prefix) + .bind(&stored.key) + .bind(stored.created_at) + .execute(&self.pool) + .await?; + Ok(stored) + } + + pub async fn list( + &self, + cursor: Option, + limit: u32, + ) -> Result, PgClientKeyStoreError> { + let fetch_limit = i64::from(limit) + 1; + let rows = if let Some(cursor) = cursor { + let (created_at, id) = + decode_cursor(&cursor).ok_or(PgClientKeyStoreError::InvalidCursor)?; + let created_at = + parse_rfc3339_utc(&created_at).map_err(|_| PgClientKeyStoreError::InvalidCursor)?; + sqlx::query( + r" +select id, name, label, prefix, key, enabled, created_at, last_used_at +from client_api_keys +where created_at < $1 or (created_at = $1 and id < $2) +order by created_at desc, id desc +limit $3", + ) + .bind(created_at) + .bind(id) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + } else { + sqlx::query( + r" +select id, name, label, prefix, key, enabled, created_at, last_used_at +from client_api_keys +order by created_at desc, id desc +limit $1", + ) + .bind(fetch_limit) + .fetch_all(&self.pool) + .await? + }; + let has_next = rows.len() > limit as usize; + let items = rows + .into_iter() + .take(limit as usize) + .map(|row| Self::key_from_row(&row)) + .collect::>(); + let next_cursor = has_next.then(|| { + let key = items.last().expect("next page requires at least one key"); + encode_cursor(&key.created_at.to_rfc3339(), &key.id) + }); + Ok(Page { items, next_cursor }) + } + + pub async fn get(&self, id: &str) -> Result, PgClientKeyStoreError> { + let row = sqlx::query( + r" +select id, name, label, prefix, key, enabled, created_at, last_used_at +from client_api_keys +where id = $1", + ) + .bind(id) + .fetch_optional(&self.pool) + .await?; + Ok(row.map(|row| Self::key_from_row(&row))) + } + + pub async fn set_enabled( + &self, + id: &str, + enabled: bool, + ) -> Result { + let result = sqlx::query("update client_api_keys set enabled = $1 where id = $2") + .bind(enabled) + .bind(id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn set_label( + &self, + id: &str, + label: Option, + ) -> Result, PgClientKeyStoreError> { + let result = sqlx::query("update client_api_keys set label = $1 where id = $2") + .bind(label) + .bind(id) + .execute(&self.pool) + .await?; + if result.rows_affected() == 0 { + return Ok(None); + } + self.get(id).await + } + + pub async fn delete(&self, id: &str) -> Result { + let result = sqlx::query("delete from client_api_keys where id = $1") + .bind(id) + .execute(&self.pool) + .await?; + Ok(result.rows_affected() > 0) + } + + pub(super) async fn touch_last_used_batch( + &self, + updates: &BTreeMap>, + ) -> Result<(), PgClientKeyStoreError> { + if updates.is_empty() { + return Ok(()); + } + let ids = updates.keys().cloned().collect::>(); + let timestamps = updates.values().copied().collect::>(); + sqlx::query( + r" +update client_api_keys as keys +set last_used_at = greatest(coalesce(keys.last_used_at, touched.used_at), touched.used_at) +from unnest($1::text[], $2::timestamptz[]) as touched(id, used_at) +where keys.id = touched.id", + ) + .bind(ids) + .bind(timestamps) + .execute(&self.pool) + .await?; + Ok(()) + } + + fn key_from_row(row: &sqlx::postgres::PgRow) -> StoredClientApiKey { + StoredClientApiKey { + id: row.get("id"), + name: row.get("name"), + label: row.get("label"), + prefix: row.get("prefix"), + key: row.get("key"), + enabled: row.get("enabled"), + created_at: row.get("created_at"), + last_used_at: row.get("last_used_at"), + } + } +} diff --git a/backend/src/keys/types.rs b/backend/src/keys/types.rs new file mode 100644 index 000000000..4aa1dcf1d --- /dev/null +++ b/backend/src/keys/types.rs @@ -0,0 +1,68 @@ +//! 客户端 API Key 领域类型。 + +use thiserror::Error; + +use super::store::StoredClientApiKey; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ManagedClientApiKey { + pub id: String, + pub name: String, + pub label: Option, + pub prefix: String, + pub key: String, + pub enabled: bool, + pub created_at: String, + pub last_used_at: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BatchDeleteClientApiKeys { + pub deleted: u32, + pub not_found: Vec, +} + +#[derive(Debug, Error)] +pub enum KeyManageError { + #[error("failed to list client API keys")] + List, + #[error("failed to create client API key")] + Create, + #[error("failed to delete client API key")] + Delete, + #[error("failed to update client API key label")] + UpdateLabel, + #[error("failed to update client API key status")] + UpdateStatus, + #[error("unsupported client API key status: {0}")] + InvalidStatus(String), + #[error("client API key name is required")] + EmptyName, + #[error("client API key ids are required")] + EmptyIds, + #[error("client API key label must be 64 characters or fewer")] + LabelTooLong, +} + +pub(super) fn parse_client_key_status(status: &str) -> Result { + match status.trim().to_ascii_lowercase().as_str() { + "active" => Ok(true), + "disabled" => Ok(false), + other => Err(KeyManageError::InvalidStatus(other.to_string())), + } +} + +impl From for ManagedClientApiKey { + fn from(key: StoredClientApiKey) -> Self { + Self { + id: key.id, + name: key.name, + label: key.label, + prefix: key.prefix, + key: key.key, + enabled: key.enabled, + created_at: key.created_at.to_rfc3339(), + last_used_at: key.last_used_at.map(|value| value.to_rfc3339()), + } + } +} diff --git a/backend/src/lib.rs b/backend/src/lib.rs index 2e0f38d5f..b4b7295c0 100644 --- a/backend/src/lib.rs +++ b/backend/src/lib.rs @@ -1,8 +1,12 @@ -pub mod admin; -pub mod config; -pub mod http; +pub mod accounts; +pub mod api; +pub mod auth; +pub mod bootstrap; +pub mod dispatch; pub mod infra; -pub mod proxy; -pub mod runtime; +pub mod keys; +pub mod models; +pub mod settings; +pub mod telemetry; +pub mod update; pub mod upstream; -pub mod web; diff --git a/backend/src/main.rs b/backend/src/main.rs index efea4343f..cbcb1c11c 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -1,6 +1,53 @@ -use std::error::Error; +use std::{env, error::Error, io, path::PathBuf}; #[tokio::main] async fn main() -> Result<(), Box> { - codex_proxy_rs::runtime::bootstrap::run().await + let mut arguments = env::args_os().skip(1); + match arguments.next().as_deref().and_then(|value| value.to_str()) { + None | Some("serve") => { + reject_extra_arguments(arguments)?; + codex_proxy_rs::bootstrap::application::run().await + } + Some("import-sqlite") => { + let source = arguments + .next() + .map(PathBuf::from) + .ok_or_else(|| invalid_cli("usage: codex-proxy-rs import-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()?; + let pool = codex_proxy_rs::infra::database::connect(&config.database.url).await?; + let report = codex_proxy_rs::telemetry::rebuild::rebuild_buckets(&pool).await?; + println!( + "request_time_buckets rebuilt: cutoff={}, deleted={}, rebuilt={}", + report.cutoff.to_rfc3339(), + report.deleted_rows, + report.rebuilt_rows + ); + Ok(()) + } + Some(command) => Err(invalid_cli(&format!( + "unknown command {command:?}; expected serve, import-sqlite, or rebuild-buckets" + )) + .into()), + } +} + +fn reject_extra_arguments(mut arguments: impl Iterator) -> Result<(), io::Error> { + if arguments.next().is_some() { + return Err(invalid_cli("unexpected extra command arguments")); + } + Ok(()) +} + +fn invalid_cli(message: &str) -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, message) } diff --git a/backend/src/upstream/models/catalog.rs b/backend/src/models/catalog.rs similarity index 99% rename from backend/src/upstream/models/catalog.rs rename to backend/src/models/catalog.rs index 4448f835c..aacb97371 100644 --- a/backend/src/upstream/models/catalog.rs +++ b/backend/src/models/catalog.rs @@ -2,8 +2,9 @@ use std::collections::{BTreeMap, BTreeSet}; +use crate::upstream::openai::protocol::model_catalog::{BackendModelEntry, BackendReasoningEffort}; + use super::{ - backend_entry::{BackendModelEntry, BackendReasoningEffort}, config::ModelConfig, info::{CodexModelInfo, ReasoningEffortInfo}, snapshot::ModelPlanSnapshot, diff --git a/backend/src/upstream/models/config.rs b/backend/src/models/config.rs similarity index 100% rename from backend/src/upstream/models/config.rs rename to backend/src/models/config.rs diff --git a/backend/src/upstream/models/info.rs b/backend/src/models/info.rs similarity index 100% rename from backend/src/upstream/models/info.rs rename to backend/src/models/info.rs diff --git a/backend/src/upstream/models/mod.rs b/backend/src/models/mod.rs similarity index 69% rename from backend/src/upstream/models/mod.rs rename to backend/src/models/mod.rs index 8657e7a49..68b555550 100644 --- a/backend/src/upstream/models/mod.rs +++ b/backend/src/models/mod.rs @@ -1,9 +1,8 @@ //! 模型目录领域。 -pub mod backend_entry; pub mod catalog; pub mod config; pub mod info; pub mod service; pub mod snapshot; -pub mod snapshot_store; +pub mod store; diff --git a/backend/src/upstream/models/service.rs b/backend/src/models/service.rs similarity index 93% rename from backend/src/upstream/models/service.rs rename to backend/src/models/service.rs index 4d28cd36b..e3d0c384f 100644 --- a/backend/src/upstream/models/service.rs +++ b/backend/src/models/service.rs @@ -7,16 +7,16 @@ use std::{ use thiserror::Error; -use crate::upstream::{ - accounts::model::Account, - transport::{CodexModelCatalogClient, CodexModelCatalogRequest}, +use crate::{ + accounts::account::Account, + upstream::openai::transport::{CodexModelCatalogClient, CodexModelCatalogRequest}, }; use super::{ catalog::ModelCatalog, config::ModelConfig, snapshot::ModelPlanSnapshot, - snapshot_store::{ModelSnapshotStore, ModelSnapshotStoreError}, + store::{ModelSnapshotStore, ModelSnapshotStoreError}, }; /// 模型刷新摘要。 @@ -86,7 +86,7 @@ pub struct ModelService { config: Arc>, snapshots: Arc>>, catalog: Arc>, - snapshot_store: Option>, + store: Option>, upstream_client: Option>, } @@ -94,7 +94,7 @@ impl ModelService { /// 构造模型服务。 pub fn new( config: ModelConfig, - snapshot_store: Option>, + store: Option>, upstream_client: Option>, ) -> Self { let catalog = ModelCatalog::from_config(&config); @@ -102,7 +102,7 @@ impl ModelService { config: Arc::new(RwLock::new(config)), snapshots: Arc::new(RwLock::new(Vec::new())), catalog: Arc::new(RwLock::new(catalog)), - snapshot_store, + store, upstream_client, } } @@ -118,11 +118,11 @@ impl ModelService { /// 从快照存储加载运行时内存模型目录。 pub async fn reload_from_store(&self) -> Result<(), ModelServiceError> { - let Some(snapshot_store) = self.snapshot_store.as_ref() else { + let Some(store) = self.store.as_ref() else { self.replace_snapshots(Vec::new()); return Ok(()); }; - let snapshots = snapshot_store + let snapshots = store .list_plan_snapshots() .await .map_err(map_load_snapshots_error)?; @@ -142,8 +142,8 @@ impl ModelService { request_id: &str, installation_id: Option<&str>, ) -> Result { - let snapshot_store = self - .snapshot_store + let store = self + .store .as_ref() .ok_or(ModelServiceError::SnapshotStoreUnavailable)?; let upstream_client = self @@ -186,7 +186,7 @@ impl ModelService { let snapshot = ModelPlanSnapshot::from_backend_entries(plan_account.plan_type.clone(), entries); result.model_count += snapshot.models.len(); - snapshot_store + store .replace_plan_snapshot(&snapshot) .await .map_err(map_store_snapshot_error)?; diff --git a/backend/src/upstream/models/snapshot.rs b/backend/src/models/snapshot.rs similarity index 100% rename from backend/src/upstream/models/snapshot.rs rename to backend/src/models/snapshot.rs diff --git a/backend/src/models/store.rs b/backend/src/models/store.rs new file mode 100644 index 000000000..eb6e7b95b --- /dev/null +++ b/backend/src/models/store.rs @@ -0,0 +1,106 @@ +//! 模型计划快照存储。 + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use redis::AsyncCommands; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use crate::infra::redis::RedisConnection; + +use super::{info::CodexModelInfo, snapshot::ModelPlanSnapshot}; + +/// 模型快照存储错误。 +#[derive(Debug, Error, PartialEq, Eq)] +pub enum ModelSnapshotStoreError { + /// 底层存储失败。 + #[error("model snapshot store operation failed: {message}")] + OperationFailed { + /// 错误说明。 + message: String, + }, +} + +/// 模型快照存储结果类型。 +pub type ModelSnapshotStoreResult = Result; + +/// 模型快照存储端口。 +#[async_trait] +pub trait ModelSnapshotStore: Send + Sync + 'static { + /// 用单个计划快照替换同名快照。 + async fn replace_plan_snapshot( + &self, + snapshot: &ModelPlanSnapshot, + ) -> ModelSnapshotStoreResult<()>; + + /// 列出所有计划快照。 + async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult>; +} + +/// Redis 模型快照存储。 +#[derive(Clone)] +pub struct RedisModelSnapshotStore { + redis: RedisConnection, +} + +#[derive(Debug, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +struct StoredPlanSnapshot { + models: Vec, + fetched_at: DateTime, +} + +impl RedisModelSnapshotStore { + /// 使用给定连接池构造存储。 + pub fn new(redis: RedisConnection) -> Self { + Self { redis } + } +} + +#[async_trait] +impl ModelSnapshotStore for RedisModelSnapshotStore { + async fn replace_plan_snapshot( + &self, + snapshot: &ModelPlanSnapshot, + ) -> ModelSnapshotStoreResult<()> { + let value = serde_json::to_string(&StoredPlanSnapshot { + models: snapshot.models.clone(), + fetched_at: Utc::now(), + }) + .map_err(store_error)?; + let mut connection = self.redis.manager(); + let _: usize = connection + .hset( + self.redis.key("models:plan_snapshots"), + &snapshot.plan_type, + value, + ) + .await + .map_err(store_error)?; + Ok(()) + } + + async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult> { + let mut connection = self.redis.manager(); + let mut rows: Vec<(String, String)> = connection + .hgetall(self.redis.key("models:plan_snapshots")) + .await + .map_err(store_error)?; + rows.sort_unstable_by(|left, right| left.0.cmp(&right.0)); + let mut snapshots = Vec::with_capacity(rows.len()); + for (plan_type, value) in rows { + let stored: StoredPlanSnapshot = serde_json::from_str(&value).map_err(store_error)?; + snapshots.push(ModelPlanSnapshot { + plan_type, + models: stored.models, + }); + } + Ok(snapshots) + } +} + +fn store_error(error: impl std::fmt::Display) -> ModelSnapshotStoreError { + ModelSnapshotStoreError::OperationFailed { + message: error.to_string(), + } +} diff --git a/backend/src/proxy/dispatch/responses/service.rs b/backend/src/proxy/dispatch/responses/service.rs deleted file mode 100644 index 3fad979bf..000000000 --- a/backend/src/proxy/dispatch/responses/service.rs +++ /dev/null @@ -1,1877 +0,0 @@ -//! Responses 创建编排与调度服务。 -//! -//! 包含了将 OpenAI 请求调度到 Codex 上游账号的完整逻辑，包括： -//! - 响应创建（非流式 / 流式 / compact） -//! - 会话亲和性与隐式续接 -//! - reasoning replay -//! - 账号回退与错误恢复 -//! - 配额验证 - -use std::{pin::Pin, sync::Arc, time::Instant}; - -use axum::body::Bytes; -use chrono::{DateTime, Duration, Utc}; -use futures::stream::Stream; -use serde_json::{json, Value}; -use tokio::sync::Mutex; - -use crate::{ - admin::monitoring::{ - ops_error_service::AdminOpsErrorLogService, - usage_record_model::{ResponseUsageRecord, UsageRecordLevel}, - usage_record_service::AdminUsageRecordService, - }, - proxy::dispatch::{ - cloudflare::{ - cloudflare_challenge_error_message, cloudflare_path_block_error_message, - is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error, - CloudflareRecovery, - }, - errors::{ - auth_failure_account_status, backend_transport_name, is_auth_upstream_error, - is_history_recovery_upstream_error, is_model_unsupported_upstream_error, - is_quota_exhausted_upstream_error, is_rate_limit_upstream_error, - rate_limit_cooldown_until, upstream_error_body, upstream_error_http_status, - upstream_error_set_cookie_headers, - }, - exhaustion::AccountExhaustionTracker, - reasoning_replay::ReasoningReplayCache, - session_affinity::{ - compute_variant_hash, ensure_prompt_cache_key, hash_instructions, - prepare_variant_identity, RuntimeSessionAffinityService, - }, - upstream::{ - create_compact_response_with_account_retrying_5xx, - create_response_stream_with_account_retrying_5xx, - create_response_with_account_retrying_5xx, verify_acquired_quota_if_required, - QuotaVerificationContext, QuotaVerificationDecision, - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - }, - usage_events::{ - reasoning_effort_from_compact_request, reasoning_effort_from_request, - record_response_event, - }, - }, - upstream::accounts::{ - model::{Account, AccountStatus}, - pool::{AccountAcquireRequest, RuntimeAccountPoolService}, - }, - upstream::{ - models::service::ModelService, - protocol::{ - events::{extract_usage, TokenUsage}, - responses::{ - response_from_codex_sse, CodexCompactRequest, CodexResponsesRequest, - CollectedResponse, - }, - }, - transport::{ - backend_transport_for_response_request, is_banned_upstream_error, CodexBackendClient, - CodexBackendResponse, - }, - }, -}; - -use crate::proxy::dispatch::implicit_resume::{ - continuation_input_start, implicit_resume_allowed, ImplicitResumeSnapshot, -}; - -use super::{ - affinity::{evict_reasoning_replay, record_response_affinity}, - errors::{ResponseDispatchError, ResponseDispatchStreamError}, - event_recording::{ - insert_response_status_metadata, insert_response_trace_metadata, - insert_response_upstream_diagnostics, insert_websocket_pool_decision, - record_prefetched_response_stream_failure_event, record_response_dispatch_error_event, - record_response_upstream_error_event, ResponseDispatchErrorDetails, - ResponseDispatchErrorEventRecord, ResponseStreamFailureEventRecord, - ResponseUpstreamErrorEventRecord, - }, - live_stream::spawn_live_response_stream, - prefetch::prefetch_first_sse_chunk, - sse_failure::{ - auth_sse_failure_account_status, client_error_invalid_reasoning_replay, first_sse_failure, - is_auth_sse_failure, is_history_recovery_sse_failure, is_model_unsupported_sse_failure, - is_quota_exhausted_sse_failure, sse_failure_error_body, - sse_failure_invalid_reasoning_replay, stream_failure_http_status, - }, - stream_lifecycle::LiveResponseStreamContext, - trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, -}; - -/// OpenAI Responses 调度服务。 -#[derive(Clone)] -pub struct ResponseDispatchService { - account_pool: Arc, - models: Arc, - codex: Arc, - session_affinity: Arc, - reasoning_replay: Arc>, - usage_records: Arc, - ops_errors: Arc, - installation_id: Option, - cloudflare: CloudflareRecovery, -} - -pub(crate) struct ResponseDispatchServiceParts { - pub account_pool: Arc, - pub models: Arc, - pub codex: Arc, - pub session_affinity: Arc, - pub usage_records: Arc, - pub ops_errors: Arc, - pub installation_id: Option, - pub cloudflare: CloudflareRecovery, -} - -/// 默认 reasoning replay TTL 秒数。 -const DEFAULT_REASONING_REPLAY_TTL_SECS: i64 = 55 * 60; -const IMPLICIT_RESUME_MAX_AGE_SECS: i64 = DEFAULT_REASONING_REPLAY_TTL_SECS; - -#[derive(Debug, Clone, Default)] -struct AccountAffinityDecision { - preferred_account_id: Option, - request_history_owner: RequestHistoryOwner, -} - -#[derive(Debug, Clone, Default)] -enum RequestHistoryOwner { - #[default] - Absent, - Known(String), - Unknown, -} - -impl AccountAffinityDecision { - fn preferred_account_id(&self) -> Option<&str> { - self.preferred_account_id.as_deref() - } - - fn should_strip_request_history_for(&self, account_id: &str) -> bool { - match &self.request_history_owner { - RequestHistoryOwner::Absent => false, - RequestHistoryOwner::Known(history_account_id) => history_account_id != account_id, - RequestHistoryOwner::Unknown => true, - } - } -} - -/// Responses live SSE 响应体流。 -pub type ResponseBodyStream = - Pin> + Send + 'static>>; - -/// Responses live SSE 调度结果。 -pub struct ResponseDispatchStream { - pub body: ResponseBodyStream, -} - -impl ResponseDispatchService { - pub(crate) fn new(parts: ResponseDispatchServiceParts) -> Self { - Self { - account_pool: parts.account_pool, - models: parts.models, - codex: parts.codex, - session_affinity: parts.session_affinity, - reasoning_replay: Arc::new(Mutex::new(ReasoningReplayCache::new(Duration::seconds( - DEFAULT_REASONING_REPLAY_TTL_SECS, - )))), - usage_records: parts.usage_records, - ops_errors: parts.ops_errors, - installation_id: parts.installation_id, - cloudflare: parts.cloudflare, - } - } - - async fn prepare_response_session( - &self, - request: &mut CodexResponsesRequest, - ) -> Option { - prepare_variant_identity(request); - if let Some(previous_response_id) = request.previous_response_id().map(ToString::to_string) - { - if request.prompt_cache_key().is_none() { - let conversation_id = self - .session_affinity - .lookup_conversation_id(&previous_response_id, Utc::now()) - .await; - request.set_prompt_cache_key(conversation_id); - } - if request.turn_state.as_deref().is_none_or(str::is_empty) { - request.turn_state = self - .session_affinity - .lookup_turn_state(&previous_response_id, Utc::now()) - .await; - } - ensure_prompt_cache_key(request); - return None; - } - - ensure_prompt_cache_key(request); - self.apply_implicit_resume(request).await - } - - async fn apply_implicit_resume( - &self, - request: &mut CodexResponsesRequest, - ) -> Option { - let continuation_start = continuation_input_start(request.input()); - if continuation_start == 0 || continuation_start >= request.input().len() { - return None; - } - let conversation_id = request - .prompt_cache_key() - .map(str::trim) - .filter(|value| !value.is_empty())? - .to_string(); - let snapshot = ImplicitResumeSnapshot::capture(request); - let variant_hash = compute_variant_hash(request); - let now = Utc::now(); - let previous_response_id = self - .session_affinity - .lookup_latest_response_by_conversation( - &conversation_id, - Some(Duration::seconds(IMPLICIT_RESUME_MAX_AGE_SECS)), - Some(&variant_hash), - now, - ) - .await?; - let current_instructions_hash = hash_instructions(Some(request.instructions())); - if self - .session_affinity - .lookup_instructions_hash(&previous_response_id, now) - .await - .as_deref() - != Some(current_instructions_hash.as_str()) - { - return None; - } - let stored_function_call_ids = self - .session_affinity - .lookup_function_call_ids(&previous_response_id, now) - .await; - if !implicit_resume_allowed( - &request.input()[continuation_start..], - request.input(), - &stored_function_call_ids, - ) { - return None; - } - let account_id = self - .session_affinity - .lookup_account(&previous_response_id, now) - .await?; - let replay_items = self.reasoning_replay.lock().await.lookup( - &previous_response_id, - &account_id, - &conversation_id, - &variant_hash, - now, - ); - let continuation = request.input()[continuation_start..].to_vec(); - let mut input = replay_items; - input.extend(continuation); - - request.set_previous_response_id(Some(previous_response_id.clone())); - request.use_websocket = true; - request.force_http_sse = false; - request.set_input(input); - if let Some(turn_state) = self - .session_affinity - .lookup_turn_state(&previous_response_id, now) - .await - { - request.turn_state = Some(turn_state); - } - - Some(snapshot) - } - - async fn account_affinity_for_request( - &self, - request: &CodexResponsesRequest, - now: DateTime, - ) -> AccountAffinityDecision { - let has_previous_response_id = request.previous_response_id().is_some(); - if let Some(account_id) = self.previous_response_account_id(request, now).await { - return AccountAffinityDecision { - preferred_account_id: Some(account_id.clone()), - request_history_owner: RequestHistoryOwner::Known(account_id), - }; - } - - let conversation_account_id = self.conversation_account_id(request, now).await; - AccountAffinityDecision { - preferred_account_id: conversation_account_id, - request_history_owner: if has_previous_response_id { - RequestHistoryOwner::Unknown - } else { - RequestHistoryOwner::Absent - }, - } - } - - async fn previous_response_account_id( - &self, - request: &CodexResponsesRequest, - now: DateTime, - ) -> Option { - let previous_response_id = request.previous_response_id()?; - self.session_affinity - .lookup_account(previous_response_id, now) - .await - } - - async fn conversation_account_id( - &self, - request: &CodexResponsesRequest, - now: DateTime, - ) -> Option { - let conversation_id = request - .prompt_cache_key() - .map(str::trim) - .filter(|value| !value.is_empty())?; - let variant_hash = compute_variant_hash(request); - self.session_affinity - .lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now) - .await - } - - async fn recover_request_history( - &self, - request: &mut CodexResponsesRequest, - implicit_resume: &mut Option, - ) { - if let Some(previous_response_id) = request.previous_response_id() { - self.session_affinity.forget(previous_response_id).await; - } - restore_request_without_history(request, implicit_resume); - } - - fn strip_history_if_account_changed( - request: &mut CodexResponsesRequest, - implicit_resume: &mut Option, - account_affinity: &AccountAffinityDecision, - acquired_account_id: &str, - ) { - if request.previous_response_id().is_none() - || !account_affinity.should_strip_request_history_for(acquired_account_id) - { - return; - } - - restore_request_without_history(request, implicit_resume); - } - - async fn apply_cascading_ban_defense( - &self, - request: &mut CodexResponsesRequest, - implicit_resume: &mut Option, - preferred_account_id: Option<&str>, - acquired_account_id: &str, - explicit_previous_response_id: Option<&str>, - ) -> bool { - let Some(preferred_account_id) = - preferred_account_id.filter(|account_id| *account_id != acquired_account_id) - else { - return false; - }; - let has_history = request.previous_response_id().is_some() - || request - .turn_state - .as_deref() - .is_some_and(|value| !value.is_empty()); - if !has_history { - return false; - } - let Some(preferred_account) = self - .account_pool - .account_snapshot(preferred_account_id) - .await - else { - return false; - }; - if !matches!( - preferred_account.status, - AccountStatus::Banned | AccountStatus::Disabled - ) { - return false; - } - - let response_id_to_forget = explicit_previous_response_id - .or(request.previous_response_id()) - .map(str::to_string); - restore_request_without_history(request, implicit_resume); - if let Some(response_id) = response_id_to_forget { - self.session_affinity.forget(&response_id).await; - } - true - } - - async fn evict_reasoning_replay(&self, request: &CodexResponsesRequest, account_id: &str) { - evict_reasoning_replay(&self.reasoning_replay, request, account_id).await; - } - - async fn record_response_affinity( - &self, - request: &CodexResponsesRequest, - account_id: &str, - body: &str, - turn_state: Option, - usage: Option, - ) { - record_response_affinity( - &self.session_affinity, - &self.reasoning_replay, - request, - account_id, - body, - turn_state, - usage, - ) - .await; - } - - /// 调度非流式 Responses 请求到 Codex Responses 上游。 - pub async fn complete( - &self, - request_id: &str, - route: &str, - mut request: CodexResponsesRequest, - requested_model: &str, - ) -> Result { - let started_at = Instant::now(); - let catalog = self.models.catalog().await; - let display_model = catalog.resolve_model_id(requested_model); - request.set_model(display_model.clone()); - let tuple_schema = request.tuple_schema.clone(); - let image_generation_requested = request.expects_image_generation(); - let now = Utc::now(); - let explicit_previous_response_id = request.previous_response_id().map(ToString::to_string); - let mut implicit_resume = self.prepare_response_session(&mut request).await; - let account_affinity = self.account_affinity_for_request(&request, now).await; - let mut acquire_request = AccountAcquireRequest::new(request.model(), now); - if let Some(preferred_account_id) = account_affinity.preferred_account_id() { - acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); - } - let mut excluded_account_ids = Vec::new(); - let mut exhausted_accounts = AccountExhaustionTracker::default(); - let mut history_recovery_used = false; - let mut next_required_account_id: Option = None; - let mut empty_response_retries = 0u8; - let mut quota_verify_attempts = 0usize; - let mut trace = ResponseDispatchTrace::default(); - const MAX_EMPTY_RESPONSE_RETRIES: u8 = 2; - let (account, response, collected_response, attempt): ( - Account, - CodexBackendResponse, - CollectedResponse, - ResponseDispatchAttempt, - ) = loop { - let mut attempt_acquire_request = acquire_request - .clone() - .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - if let Some(account_id) = next_required_account_id.take() { - attempt_acquire_request = - attempt_acquire_request.with_required_account_id(account_id); - } - let allow_quota_verify_account_retry = - attempt_acquire_request.required_account_id.is_none(); - attempt_acquire_request.now = Utc::now(); - let Some(acquired) = self - .account_pool - .acquire_with(&attempt_acquire_request) - .await - else { - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: exhausted_accounts.last_account_id(), - stream: false, - compact: false, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &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, - installation_id: self.installation_id.as_deref(), - request_id, - excluded_account_ids: &mut excluded_account_ids, - verify_attempts: &mut quota_verify_attempts, - allow_retry_with_another_account: allow_quota_verify_account_retry, - }, - acquired, - ) - .await - { - QuotaVerificationDecision::Ready(acquired) => *acquired, - QuotaVerificationDecision::RetryWithAnotherAccount => { - exhausted_accounts - .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); - continue; - } - QuotaVerificationDecision::MaxAttemptsReached - | QuotaVerificationDecision::RequiredAccountUnavailable => { - exhausted_accounts.record_rate_limited( - Some(&acquired_account_id), - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - ); - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: Some(&acquired_account_id), - stream: false, - compact: false, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - } - }; - - self.apply_cascading_ban_defense( - &mut request, - &mut implicit_resume, - account_affinity.preferred_account_id(), - &acquired.account.id, - explicit_previous_response_id.as_deref(), - ) - .await; - Self::strip_history_if_account_changed( - &mut request, - &mut implicit_resume, - &account_affinity, - &acquired.account.id, - ); - self.account_pool.wait_for_request_interval(&acquired).await; - let account = acquired.account; - let release_account_id = account.id.clone(); - let attempt = trace.start_attempt(&release_account_id); - let response_result = create_response_with_account_retrying_5xx( - &self.codex, - self.installation_id.as_deref(), - &self.cloudflare, - &request, - request_id, - &account, - started_at, - ) - .await; - self.account_pool.release(&release_account_id).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.account_pool - .sync_passive_rate_limit_headers(&account, &response.rate_limit_headers) - .await; - let collected_response = - match response_from_codex_sse(&response.body, tuple_schema.as_ref()) { - Ok(collected_response) => collected_response, - Err(error) => { - let error = ResponseDispatchError::InvalidSse(error); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: Some(&release_account_id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - return Err(error); - } - }; - if matches!(collected_response, CollectedResponse::Empty) { - self.account_pool - .record_empty_response_attempt( - &release_account_id, - request.model(), - image_generation_requested, - ) - .await; - empty_response_retries += 1; - if empty_response_retries <= MAX_EMPTY_RESPONSE_RETRIES { - continue; - } - } - if let CollectedResponse::Failed(failure) = &collected_response { - if is_history_recovery_sse_failure(failure) && !history_recovery_used { - if sse_failure_invalid_reasoning_replay(failure) { - self.evict_reasoning_replay(&request, &release_account_id) - .await; - } - self.recover_request_history(&mut request, &mut implicit_resume) - .await; - history_recovery_used = true; - next_required_account_id = Some(release_account_id); - continue; - } - if is_model_unsupported_sse_failure(failure) { - let upstream_error = sse_failure_error_body(failure); - if let Some(exhausted) = exhausted_accounts - .model_unsupported_retry_exhausted(upstream_error.clone()) - { - let error = - ResponseDispatchError::from_exhausted_account(exhausted); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: Some(&release_account_id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - return Err(error); - } - exhausted_accounts.record_model_unsupported( - Some(&release_account_id), - upstream_error, - ); - excluded_account_ids.push(release_account_id); - continue; - } - if is_quota_exhausted_sse_failure(failure) { - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - failure.message.clone(), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - if is_auth_sse_failure(failure) { - let upstream_error = sse_failure_error_body(failure); - let account_status = auth_sse_failure_account_status(failure); - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - account_status, - upstream_error, - Some(stream_failure_http_status(failure)), - ); - self.account_pool - .set_status(&release_account_id, account_status) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - } - break (account, response, collected_response, attempt); - } - 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; - excluded_account_ids.push(release_account_id); - } - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) - if is_history_recovery_upstream_error(&error) && !history_recovery_used => - { - if client_error_invalid_reasoning_replay(&error) { - self.evict_reasoning_replay(&request, &release_account_id) - .await; - } - self.recover_request_history(&mut request, &mut implicit_resume) - .await; - history_recovery_used = true; - next_required_account_id = Some(release_account_id); - } - Err(error) if is_auth_upstream_error(&error) => { - let upstream_error = upstream_error_body(&error); - let account_status = auth_failure_account_status(&error); - 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; - excluded_account_ids.push(release_account_id); - } - 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; - excluded_account_ids.push(release_account_id); - } - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_model_unsupported_upstream_error(&error) => { - let upstream_error = upstream_error_body(&error); - if let Some(exhausted) = - exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) - { - let error = ResponseDispatchError::from_exhausted_account(exhausted); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: Some(&release_account_id), - stream: false, - compact: false, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - } - exhausted_accounts - .record_model_unsupported(Some(&release_account_id), upstream_error); - excluded_account_ids.push(release_account_id); - } - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) => { - record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { - ops_errors: &self.ops_errors, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: false, - transport: backend_transport_for_response_request(&request), - request: &request, - error: &error, - trace: &trace, - attempt: Some(&attempt), - }) - .await; - return Err(ResponseDispatchError::Upstream(error)); - } - } - }; - - match collected_response { - CollectedResponse::Completed(body) => { - let response_id = body.get("id").and_then(Value::as_str); - self.cloudflare.reset_account_recovery(&account.id).await; - if let Some(usage) = response.usage { - self.account_pool - .record_response_usage( - &account.id, - request.model(), - usage, - image_generation_requested, - ) - .await; - } - self.record_response_affinity( - &request, - &account.id, - &response.body, - response.turn_state.clone(), - response.usage, - ) - .await; - let mut metadata = json!({ - "responseId": response_id, - "stream": false, - "transport": backend_transport_name(response.transport), - "firstTokenMs": response.first_token_ms, - "usage": response.usage, - }); - 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)); - insert_websocket_pool_decision(&mut metadata, response.websocket_pool_decision); - record_response_event(ResponseUsageRecord { - usage_records: &self.usage_records, - request_id, - account_id: &account.id, - route, - model: &display_model, - requested_model: Some(requested_model), - client_ip: request.client_ip.as_deref(), - client_user_agent: request.client_user_agent.as_deref(), - reasoning_effort: reasoning_effort_from_request(&request), - service_tier: request.service_tier(), - started_at, - status_code: 200, - level: UsageRecordLevel::Info, - message: "v1 responses completed", - metadata, - rate_limit_headers: &response.rate_limit_headers, - }) - .await; - Ok(body) - } - CollectedResponse::Failed(failure) => { - let error = ResponseDispatchError::Failed(failure); - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: Some(&account.id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - CollectedResponse::MissingCompleted => { - let error = ResponseDispatchError::MissingCompleted; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: Some(&account.id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - CollectedResponse::Empty => { - let error = ResponseDispatchError::EmptyUpstreamResponse; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: Some(&account.id), - stream: false, - compact: false, - transport: Some(backend_transport_name(response.transport)), - }, - &error, - ) - .await; - Err(error) - } - } - } - - async fn record_response_dispatch_error( - &self, - request_id: &str, - route: &str, - requested_model: &str, - started_at: Instant, - details: ResponseDispatchErrorDetails<'_>, - error: &ResponseDispatchError, - ) { - record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { - ops_errors: &self.ops_errors, - request_id, - account_id: details.account_id, - route, - model: requested_model, - started_at, - stream: details.stream, - compact: details.compact, - transport: details.transport, - error, - }) - .await; - } - - /// 调度流式 Responses 请求到 Codex Responses 上游。 - pub async fn stream( - &self, - request_id: &str, - route: &str, - mut request: CodexResponsesRequest, - requested_model: &str, - ) -> Result { - let started_at = Instant::now(); - let catalog = self.models.catalog().await; - let display_model = catalog.resolve_model_id(requested_model); - request.set_model(display_model.clone()); - request.set_stream(true); - let tuple_schema = request.tuple_schema.clone(); - let now = Utc::now(); - let explicit_previous_response_id = request.previous_response_id().map(ToString::to_string); - let mut implicit_resume = self.prepare_response_session(&mut request).await; - let account_affinity = self.account_affinity_for_request(&request, now).await; - let mut acquire_request = AccountAcquireRequest::new(request.model(), now); - if let Some(preferred_account_id) = account_affinity.preferred_account_id() { - acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); - } - let mut excluded_account_ids = Vec::new(); - let mut exhausted_accounts = AccountExhaustionTracker::default(); - let mut history_recovery_used = false; - let mut next_required_account_id: Option = None; - let mut quota_verify_attempts = 0usize; - let mut trace = ResponseDispatchTrace::default(); - macro_rules! return_stream_dispatch_error { - ($error:expr) => {{ - let error = $error; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: exhausted_accounts.last_account_id(), - stream: true, - compact: false, - transport: Some(backend_transport_name( - backend_transport_for_response_request(&request), - )), - }, - &error, - ) - .await; - return Err(error); - }}; - ($error:expr, account_id: $account_id:expr, transport: $transport:expr) => {{ - let error = $error; - self.record_response_dispatch_error( - request_id, - route, - requested_model, - started_at, - ResponseDispatchErrorDetails { - account_id: $account_id, - stream: true, - compact: false, - transport: $transport, - }, - &error, - ) - .await; - return Err(error); - }}; - } - loop { - let mut attempt_acquire_request = acquire_request - .clone() - .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - if let Some(account_id) = next_required_account_id.take() { - attempt_acquire_request = - attempt_acquire_request.with_required_account_id(account_id); - } - let allow_quota_verify_account_retry = - attempt_acquire_request.required_account_id.is_none(); - attempt_acquire_request.now = Utc::now(); - let Some(acquired) = self - .account_pool - .acquire_with(&attempt_acquire_request) - .await - else { - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - return_stream_dispatch_error!(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, - installation_id: self.installation_id.as_deref(), - request_id, - excluded_account_ids: &mut excluded_account_ids, - verify_attempts: &mut quota_verify_attempts, - allow_retry_with_another_account: allow_quota_verify_account_retry, - }, - acquired, - ) - .await - { - QuotaVerificationDecision::Ready(acquired) => *acquired, - QuotaVerificationDecision::RetryWithAnotherAccount => { - exhausted_accounts - .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); - continue; - } - QuotaVerificationDecision::MaxAttemptsReached - | QuotaVerificationDecision::RequiredAccountUnavailable => { - exhausted_accounts.record_rate_limited( - Some(&acquired_account_id), - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - ); - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - return_stream_dispatch_error!( - error, - account_id: Some(&acquired_account_id), - transport: Some(backend_transport_name(backend_transport_for_response_request( - &request - ))) - ); - } - }; - - self.apply_cascading_ban_defense( - &mut request, - &mut implicit_resume, - account_affinity.preferred_account_id(), - &acquired.account.id, - explicit_previous_response_id.as_deref(), - ) - .await; - Self::strip_history_if_account_changed( - &mut request, - &mut implicit_resume, - &account_affinity, - &acquired.account.id, - ); - - self.account_pool.wait_for_request_interval(&acquired).await; - let account = acquired.account; - let release_account_id = account.id.clone(); - let attempt = trace.start_attempt(&release_account_id); - let response_result = create_response_stream_with_account_retrying_5xx( - &self.codex, - self.installation_id.as_deref(), - &self.cloudflare, - &request, - request_id, - &account, - ) - .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) => { - let transport = response.transport; - let set_cookie_headers = response.set_cookie_headers; - let rate_limit_headers = response.rate_limit_headers; - let rate_limit_header_updates = response.rate_limit_header_updates; - let turn_state_update = response.turn_state_update; - let websocket_pool_decision = response.websocket_pool_decision; - let turn_state = response.turn_state; - let diagnostics = response.diagnostics; - self.cloudflare - .capture_set_cookie_headers(&release_account_id, &set_cookie_headers) - .await; - self.account_pool - .sync_passive_rate_limit_headers(&account, &rate_limit_headers) - .await; - let (prefetched, body) = match prefetch_first_sse_chunk(response.body).await { - Ok(prefetched) => prefetched, - Err(ResponseDispatchError::Upstream(error)) - if is_history_recovery_upstream_error(&error) - && !history_recovery_used => - { - self.account_pool.release(&release_account_id).await; - if client_error_invalid_reasoning_replay(&error) { - self.evict_reasoning_replay(&request, &release_account_id) - .await; - } - self.recover_request_history(&mut request, &mut implicit_resume) - .await; - history_recovery_used = true; - next_required_account_id = Some(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_rate_limit_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_quota_exhausted_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - upstream_error_body(&error), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_auth_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_cloudflare_challenge_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_cloudflare_path_block_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_model_unsupported_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - let upstream_error = upstream_error_body(&error); - if let Some(exhausted) = exhausted_accounts - .model_unsupported_retry_exhausted(upstream_error.clone()) - { - return_stream_dispatch_error!( - ResponseDispatchError::from_exhausted_account(exhausted), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - exhausted_accounts.record_model_unsupported( - Some(&release_account_id), - upstream_error, - ); - excluded_account_ids.push(release_account_id); - continue; - } - Err(ResponseDispatchError::Upstream(error)) - if is_banned_upstream_error(&error) => - { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - continue; - } - Err(error) => { - self.account_pool.release(&release_account_id).await; - if let ResponseDispatchError::Upstream(upstream_error) = &error { - record_response_upstream_error_event( - ResponseUpstreamErrorEventRecord { - ops_errors: &self.ops_errors, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: true, - transport, - request: &request, - error: upstream_error, - trace: &trace, - attempt: Some(&attempt), - }, - ) - .await; - return Err(error); - } - return_stream_dispatch_error!( - error, - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - }; - let first_failure = match first_sse_failure(&prefetched) { - Ok(failure) => failure, - Err(error) => { - self.account_pool.release(&release_account_id).await; - return_stream_dispatch_error!( - ResponseDispatchError::InvalidSse(error), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - }; - if let Some(failure) = first_failure { - if is_history_recovery_sse_failure(&failure) && !history_recovery_used { - self.account_pool.release(&release_account_id).await; - if sse_failure_invalid_reasoning_replay(&failure) { - self.evict_reasoning_replay(&request, &release_account_id) - .await; - } - self.recover_request_history(&mut request, &mut implicit_resume) - .await; - history_recovery_used = true; - next_required_account_id = Some(release_account_id); - continue; - } - if is_model_unsupported_sse_failure(&failure) { - let upstream_error = sse_failure_error_body(&failure); - if let Some(exhausted) = exhausted_accounts - .model_unsupported_retry_exhausted(upstream_error.clone()) - { - self.account_pool.release(&release_account_id).await; - return_stream_dispatch_error!( - ResponseDispatchError::from_exhausted_account(exhausted), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(transport)) - ); - } - exhausted_accounts.record_model_unsupported( - Some(&release_account_id), - upstream_error, - ); - excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; - continue; - } - if is_quota_exhausted_sse_failure(&failure) { - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - failure.message.clone(), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; - continue; - } - if is_auth_sse_failure(&failure) { - let upstream_error = sse_failure_error_body(&failure); - let account_status = auth_sse_failure_account_status(&failure); - exhausted_accounts.record_auth_failure( - Some(&release_account_id), - account_status, - upstream_error, - Some(stream_failure_http_status(&failure)), - ); - self.account_pool - .set_status(&release_account_id, account_status) - .await; - excluded_account_ids.push(release_account_id); - self.account_pool.release(&account.id).await; - continue; - } - self.account_pool.release(&release_account_id).await; - record_prefetched_response_stream_failure_event( - ResponseStreamFailureEventRecord { - ops_errors: &self.ops_errors, - request_id, - account_id: &release_account_id, - route, - model: &display_model, - requested_model, - started_at, - transport, - request: &request, - failure: &failure, - diagnostics: &diagnostics, - rate_limit_headers: &rate_limit_headers, - prefetched: &prefetched, - trace: &trace, - attempt: &attempt, - }, - ) - .await; - return Err(ResponseDispatchError::Failed(failure.clone())); - } - - let context = LiveResponseStreamContext { - account_pool: Arc::clone(&self.account_pool), - session_affinity: Arc::clone(&self.session_affinity), - reasoning_replay: Arc::clone(&self.reasoning_replay), - usage_records: Arc::clone(&self.usage_records), - ops_errors: Arc::clone(&self.ops_errors), - cloudflare: self.cloudflare.clone(), - account_id: account.id, - account_plan_type: account.plan_type, - request_id: request_id.to_string(), - route: route.to_string(), - model: request.model().to_string(), - display_model: display_model.clone(), - requested_model: requested_model.to_string(), - client_ip: request.client_ip.clone(), - request, - tuple_schema, - transport, - rate_limit_headers, - rate_limit_header_updates, - turn_state_update, - websocket_pool_decision, - turn_state, - diagnostics, - attempt: attempt.clone(), - attempts: trace.attempts().to_vec(), - started_at, - }; - return Ok(spawn_live_response_stream(context, prefetched, body)); - } - Err(error) if is_rate_limit_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_quota_exhausted_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - exhausted_accounts.record_quota_exhausted( - Some(&release_account_id), - upstream_error_body(&error), - ); - self.account_pool - .set_status(&release_account_id, AccountStatus::QuotaExhausted) - .await; - excluded_account_ids.push(release_account_id); - } - Err(error) - if is_history_recovery_upstream_error(&error) && !history_recovery_used => - { - self.account_pool.release(&release_account_id).await; - if client_error_invalid_reasoning_replay(&error) { - self.evict_reasoning_replay(&request, &release_account_id) - .await; - } - self.recover_request_history(&mut request, &mut implicit_resume) - .await; - history_recovery_used = true; - next_required_account_id = Some(release_account_id); - } - Err(error) if is_auth_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_cloudflare_challenge_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_cloudflare_path_block_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_model_unsupported_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - let upstream_error = upstream_error_body(&error); - if let Some(exhausted) = - exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) - { - return_stream_dispatch_error!( - ResponseDispatchError::from_exhausted_account(exhausted), - account_id: Some(&release_account_id), - transport: Some(backend_transport_name(backend_transport_for_response_request( - &request - ))) - ); - } - exhausted_accounts - .record_model_unsupported(Some(&release_account_id), upstream_error); - excluded_account_ids.push(release_account_id); - } - Err(error) if is_banned_upstream_error(&error) => { - self.account_pool.release(&release_account_id).await; - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) => { - self.account_pool.release(&release_account_id).await; - record_response_upstream_error_event(ResponseUpstreamErrorEventRecord { - ops_errors: &self.ops_errors, - request_id, - account_id: &release_account_id, - account_email: account.email.as_deref(), - route, - model: requested_model, - started_at, - stream: true, - transport: backend_transport_for_response_request(&request), - request: &request, - error: &error, - trace: &trace, - attempt: Some(&attempt), - }) - .await; - return Err(ResponseDispatchError::Upstream(error)); - } - } - } - } - - /// 调度 Responses compact 请求到 Codex compact 上游。 - pub async fn compact( - &self, - request_id: &str, - mut request: CodexCompactRequest, - requested_model: &str, - ) -> Result { - let started_at = Instant::now(); - let catalog = self.models.catalog().await; - let display_model = catalog.resolve_model_id(requested_model); - request.set_model(display_model.clone()); - let mut excluded_account_ids = Vec::new(); - let mut exhausted_accounts = AccountExhaustionTracker::default(); - let mut quota_verify_attempts = 0usize; - let mut trace = ResponseDispatchTrace::default(); - - loop { - let acquire_request = AccountAcquireRequest::new(request.model(), Utc::now()) - .with_exclude_account_ids(excluded_account_ids.iter().cloned()); - let acquired = match self.account_pool.acquire_with(&acquire_request).await { - Some(acquired) => acquired, - None => { - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_compact_dispatch_error( - request_id, - 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, - installation_id: self.installation_id.as_deref(), - request_id, - excluded_account_ids: &mut excluded_account_ids, - verify_attempts: &mut quota_verify_attempts, - allow_retry_with_another_account: true, - }, - acquired, - ) - .await - { - QuotaVerificationDecision::Ready(acquired) => *acquired, - QuotaVerificationDecision::RetryWithAnotherAccount => { - exhausted_accounts - .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); - continue; - } - QuotaVerificationDecision::MaxAttemptsReached - | QuotaVerificationDecision::RequiredAccountUnavailable => { - exhausted_accounts.record_rate_limited( - Some(&acquired_account_id), - QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, - ); - let error = exhausted_accounts - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount); - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - Some(&acquired_account_id), - &error, - ) - .await; - return Err(error); - } - }; - let account = acquired.account; - 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( - &self.codex, - self.installation_id.as_deref(), - &self.cloudflare, - &request, - request_id, - &account, - ) - .await; - self.account_pool.release(&release_account_id).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, request.model(), &usage) - .await; - } - let mut metadata = json!({ - "stream": false, - "compact": true, - "usage": usage, - }); - 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 { - usage_records: &self.usage_records, - request_id, - account_id: &account.id, - route: "/v1/responses/compact", - model: &display_model, - requested_model: Some(requested_model), - client_ip: request.client_ip.as_deref(), - client_user_agent: request.client_user_agent.as_deref(), - reasoning_effort: reasoning_effort_from_compact_request(&request), - service_tier: None, - started_at, - status_code: 200, - level: UsageRecordLevel::Info, - message: "v1 responses compact completed", - metadata, - rate_limit_headers: &response.rate_limit_headers, - }) - .await; - return Ok(response.body); - } - 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; - excluded_account_ids.push(release_account_id); - } - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_auth_upstream_error(&error) => { - let upstream_error = upstream_error_body(&error); - let account_status = auth_failure_account_status(&error); - 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; - excluded_account_ids.push(release_account_id); - } - 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; - excluded_account_ids.push(release_account_id); - } - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) if is_model_unsupported_upstream_error(&error) => { - let upstream_error = upstream_error_body(&error); - if let Some(exhausted) = - exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) - { - let error = ResponseDispatchError::from_exhausted_account(exhausted); - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - Some(&release_account_id), - &error, - ) - .await; - return Err(error); - } - exhausted_accounts - .record_model_unsupported(Some(&release_account_id), upstream_error); - excluded_account_ids.push(release_account_id); - } - 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; - excluded_account_ids.push(release_account_id); - } - Err(error) => { - let error = ResponseDispatchError::Upstream(error); - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - Some(&release_account_id), - &error, - ) - .await; - return Err(error); - } - } - } - } - - async fn record_compact_dispatch_error( - &self, - request_id: &str, - requested_model: &str, - started_at: Instant, - account_id: Option<&str>, - error: &ResponseDispatchError, - ) { - record_response_dispatch_error_event(ResponseDispatchErrorEventRecord { - ops_errors: &self.ops_errors, - request_id, - account_id, - route: "/v1/responses/compact", - model: requested_model, - started_at, - stream: false, - compact: true, - transport: Some("http"), - error, - }) - .await; - } -} - -fn strip_request_history(request: &mut CodexResponsesRequest) { - request.set_previous_response_id(None); - request.turn_state = None; -} - -fn restore_request_without_history( - request: &mut CodexResponsesRequest, - implicit_resume: &mut Option, -) { - if let Some(snapshot) = implicit_resume.take() { - snapshot.restore(request); - } - strip_request_history(request); -} diff --git a/backend/src/proxy/dispatch/session_affinity.rs b/backend/src/proxy/dispatch/session_affinity.rs deleted file mode 100644 index a3cab6f67..000000000 --- a/backend/src/proxy/dispatch/session_affinity.rs +++ /dev/null @@ -1,683 +0,0 @@ -//! 会话亲和性策略、Conversation identity 构建、SQLite 存储与运行时会话亲和性服务。 - -use std::collections::BTreeMap; - -use chrono::{DateTime, Duration, Utc}; -use serde_json::Value; -use sha2::{Digest, Sha256}; -use sqlx::SqlitePool; -use thiserror::Error; -use uuid::Uuid; - -use crate::{ - infra::time::parse_rfc3339_utc as parse_rfc3339, - upstream::protocol::responses::CodexResponsesRequest, -}; - -// ==================================================================== -// 亲和性核心类型 -// ==================================================================== - -/// 会话亲和性条目。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct SessionAffinityEntry { - pub account_id: String, - pub conversation_id: String, - pub turn_state: Option, - pub instructions_hash: Option, - pub input_tokens: Option, - pub function_call_ids: Vec, - pub variant_hash: Option, - pub created_at: DateTime, -} - -/// 持久化的会话亲和性条目。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct StoredSessionAffinity { - pub response_id: String, - pub entry: SessionAffinityEntry, - pub expires_at: DateTime, -} - -/// 纯内存会话亲和性映射。 -#[derive(Debug, Clone)] -pub struct SessionAffinityMap { - entries: BTreeMap, - ttl: Duration, -} - -impl SessionAffinityMap { - pub fn new(ttl: Duration) -> Self { - Self { - entries: BTreeMap::new(), - ttl, - } - } - - pub fn record(&mut self, response_id: String, entry: SessionAffinityEntry) { - self.entries.insert(response_id, entry); - } - - pub fn restore(&mut self, records: Vec, now: DateTime) -> usize { - let mut restored = 0usize; - for record in records { - if self.stored_record_is_active(&record, now) { - self.entries.insert(record.response_id, record.entry); - restored = restored.saturating_add(1); - } - } - restored - } - - pub fn lookup_account(&self, response_id: &str, now: DateTime) -> Option { - self.active_entry(response_id, now) - .map(|entry| entry.account_id.clone()) - } - - pub fn lookup_conversation_id(&self, response_id: &str, now: DateTime) -> Option { - self.active_entry(response_id, now) - .map(|entry| entry.conversation_id.clone()) - } - - pub fn lookup_turn_state(&self, response_id: &str, now: DateTime) -> Option { - self.active_entry(response_id, now) - .and_then(|entry| entry.turn_state.clone()) - } - - pub fn lookup_instructions_hash( - &self, - response_id: &str, - now: DateTime, - ) -> Option { - self.active_entry(response_id, now) - .and_then(|entry| entry.instructions_hash.clone()) - } - - pub fn lookup_function_call_ids(&self, response_id: &str, now: DateTime) -> Vec { - self.active_entry(response_id, now) - .map(|entry| entry.function_call_ids.clone()) - .unwrap_or_default() - } - - pub fn lookup_latest_response_by_conversation( - &self, - conversation_id: &str, - max_age: Option, - variant_hash: Option<&str>, - now: DateTime, - ) -> Option { - self.latest_entry_by_conversation(conversation_id, max_age, variant_hash, now) - .map(|(response_id, _)| response_id.clone()) - } - - pub fn lookup_latest_account_by_conversation( - &self, - conversation_id: &str, - max_age: Option, - variant_hash: Option<&str>, - now: DateTime, - ) -> Option { - self.latest_entry_by_conversation(conversation_id, max_age, variant_hash, now) - .map(|(_, entry)| entry.account_id.clone()) - } - - pub fn forget(&mut self, response_id: &str) -> bool { - self.entries.remove(response_id).is_some() - } - - fn active_entry(&self, response_id: &str, now: DateTime) -> Option<&SessionAffinityEntry> { - self.entries - .get(response_id) - .filter(|entry| self.entry_is_active(entry, now)) - } - - fn latest_entry_by_conversation( - &self, - conversation_id: &str, - max_age: Option, - variant_hash: Option<&str>, - now: DateTime, - ) -> Option<(&String, &SessionAffinityEntry)> { - self.entries - .iter() - .filter(|(_, entry)| entry.conversation_id == conversation_id) - .filter(|(_, entry)| self.entry_is_active(entry, now)) - .filter(|(_, entry)| { - variant_hash - .is_none_or(|variant_hash| entry.variant_hash.as_deref() == Some(variant_hash)) - }) - .filter(|(_, entry)| { - max_age.is_none_or(|max_age| now.signed_duration_since(entry.created_at) <= max_age) - }) - .max_by_key(|(_, entry)| entry.created_at) - } - - fn stored_record_is_active(&self, record: &StoredSessionAffinity, now: DateTime) -> bool { - record.expires_at > now && self.entry_is_active(&record.entry, now) - } - - fn entry_is_active(&self, entry: &SessionAffinityEntry, now: DateTime) -> bool { - entry - .created_at - .checked_add_signed(self.ttl) - .is_some_and(|expires_at| expires_at > now) - } -} - -// ==================================================================== -// Variant identity helpers -// ==================================================================== - -/// 准备用于区分并发分支的变体身份。 -pub fn prepare_variant_identity(request: &mut CodexResponsesRequest) { - request.variant_identity = build_variant_identity(request); -} - -/// 计算请求变体哈希。 -pub fn compute_variant_hash(request: &CodexResponsesRequest) -> String { - compute_variant_hash_with_identity(request, request.variant_identity.as_deref()) -} - -fn compute_variant_hash_with_identity( - request: &CodexResponsesRequest, - identity: Option<&str>, -) -> String { - let tools_json = request.tools().map_or_else(|| "[]".to_string(), tools_json); - let mut hasher = Sha256::new(); - hasher.update(request.instructions().as_bytes()); - hasher.update(b"\0"); - hasher.update(tools_json.as_bytes()); - if let Some(identity) = identity - .map(str::trim) - .filter(|identity| !identity.is_empty()) - { - hasher.update(b"\0"); - hasher.update(identity.as_bytes()); - } - hex::encode(hasher.finalize()).chars().take(12).collect() -} - -fn tools_json(tools: &[Value]) -> String { - serde_json::to_string(tools).unwrap_or_else(|_| "[]".to_string()) -} - -fn build_variant_identity(request: &CodexResponsesRequest) -> Option { - let mut parts = Vec::with_capacity(2); - 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 let Some(anchor) = derive_stable_conversation_key(request) { - parts.push(format!("anchor:{anchor}")); - } - } - - (!parts.is_empty()).then(|| parts.join("\0")) -} - -fn non_empty_str(value: Option<&str>) -> Option<&str> { - value.map(str::trim).filter(|value| !value.is_empty()) -} - -/// 计算 instructions 哈希。 -pub fn hash_instructions(instructions: Option<&str>) -> String { - let mut hasher = Sha256::new(); - hasher.update(instructions.unwrap_or_default().as_bytes()); - hex::encode(hasher.finalize()) -} - -// ==================================================================== -// Conversation identity 构建器 -// ==================================================================== - -const LEADING_SYSTEM_REMINDER_OPEN: &str = ""; -const LEADING_SYSTEM_REMINDER_CLOSE: &str = ""; - -/// 从请求上下文派生的 conversation identity -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ConversationIdentity { - pub conversation_id: Option, - pub window_id: Option, -} - -/// 从 prompt_cache_key 和可选的 window_id 构建 conversation identity -pub fn build_conversation_identity( - prompt_cache_key: Option<&str>, - client_window_id: Option<&str>, - account_scope: &str, -) -> ConversationIdentity { - let conversation_id = prompt_cache_key - .filter(|s| !s.trim().is_empty()) - .map(|key| build_account_scoped_identity("conversation", account_scope, key)); - - let window_id = if let Some(client_win) = client_window_id.filter(|s| !s.trim().is_empty()) { - Some(build_account_scoped_identity( - "window", - account_scope, - client_win, - )) - } else { - conversation_id - .as_ref() - .map(|conv_id| format!("{}:0", conv_id)) - }; - - ConversationIdentity { - conversation_id, - window_id, - } -} - -/// 确保请求拥有上游可复用的 prompt cache key。 -pub fn ensure_prompt_cache_key(request: &mut CodexResponsesRequest) { - if let Some(existing) = request - .prompt_cache_key() - .map(str::trim) - .filter(|value| !value.is_empty()) - { - let existing = existing.to_string(); - request.set_prompt_cache_key(Some(existing)); - return; - } - - request.set_prompt_cache_key(Some( - derive_stable_conversation_key(request).unwrap_or_else(|| Uuid::new_v4().to_string()), - )); -} - -/// 按原版 `stable-conversation-key.ts` 的规则派生稳定 conversation key。 -pub fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option { - let instructions = request - .instructions() - .chars() - .take(2000) - .collect::(); - let first_user_text = first_user_text(request.input()); - let normalized_first_user_text = normalize_conversation_anchor_text(&first_user_text); - let first_user_text = if normalized_first_user_text.is_empty() { - first_user_text - } else { - normalized_first_user_text - }; - if instructions.is_empty() && first_user_text.is_empty() { - return None; - } - - let mut hasher = Sha256::new(); - hasher.update(request.model().as_bytes()); - hasher.update(b"\0"); - hasher.update(instructions.as_bytes()); - hasher.update(b"\0"); - hasher.update(first_user_text.as_bytes()); - let hex = hex::encode(hasher.finalize()); - - Some(format!( - "{}-{}-{}-{}-{}", - &hex[0..8], - &hex[8..12], - &hex[12..16], - &hex[16..20], - &hex[20..32] - )) -} - -fn first_user_text(input: &[Value]) -> String { - for item in input { - if item.get("role").and_then(Value::as_str) != Some("user") { - continue; - } - let Some(content) = item.get("content") else { - return String::new(); - }; - if let Some(text) = content.as_str() { - return text.to_string(); - } - if let Some(parts) = content.as_array() { - return parts - .iter() - .filter(|part| part.get("type").and_then(Value::as_str) == Some("input_text")) - .filter_map(|part| part.get("text").and_then(Value::as_str)) - .collect::(); - } - return String::new(); - } - - String::new() -} - -fn normalize_conversation_anchor_text(text: &str) -> String { - let mut rest = text.trim_start(); - loop { - let lower = rest.to_ascii_lowercase(); - if !lower.starts_with(LEADING_SYSTEM_REMINDER_OPEN) { - break; - } - let Some(close_start) = lower.find(LEADING_SYSTEM_REMINDER_CLOSE) else { - break; - }; - let close_end = close_start + LEADING_SYSTEM_REMINDER_CLOSE.len(); - rest = rest[close_end..].trim_start(); - } - rest.to_string() -} - -/// 构建账号作用域的身份哈希。 -fn build_account_scoped_identity(kind: &str, account_scope: &str, client_value: &str) -> String { - let mut hasher = Sha256::new(); - hasher.update(kind.as_bytes()); - hasher.update(b"\0"); - hasher.update(account_scope.as_bytes()); - hasher.update(b"\0"); - hasher.update(client_value.as_bytes()); - - let digest = hasher.finalize(); - let hex = hex::encode(digest); - let truncated = &hex[..32]; - - let prefix = match kind { - "conversation" => "cp", - "window" => "cw", - _ => "cx", - }; - - format!("{}_{}", prefix, truncated) -} - -// ==================================================================== -// SQLite 会话亲和性存储 -// ==================================================================== - -/// SQLite 会话亲和性存储。 -#[derive(Clone)] -pub struct SqliteSessionAffinityStore { - pool: SqlitePool, -} - -/// SQLite 会话亲和性存储错误。 -#[derive(Debug, Error)] -pub enum SqliteSessionAffinityStoreError { - #[error("database error: {0}")] - Database(#[from] sqlx::Error), - #[error("invalid session affinity timestamp: {0}")] - InvalidTimestamp(#[from] chrono::ParseError), - #[error("invalid session affinity function call ids: {0}")] - InvalidFunctionCallIds(#[from] serde_json::Error), -} - -pub type SqliteSessionAffinityStoreResult = Result; - -impl SqliteSessionAffinityStore { - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - pub async fn upsert( - &self, - response_id: &str, - entry: &SessionAffinityEntry, - ttl: Duration, - ) -> SqliteSessionAffinityStoreResult<()> { - let function_call_ids_json = serde_json::to_string(&entry.function_call_ids)?; - let expires_at = entry - .created_at - .checked_add_signed(ttl) - .unwrap_or(DateTime::::MAX_UTC); - sqlx::query( - r" -insert into session_affinities ( - response_id, - account_id, - conversation_id, - turn_state, - instructions_hash, - input_tokens, - function_call_ids_json, - variant_hash, - expires_at, - created_at -) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) -on conflict(response_id) do update set - account_id = excluded.account_id, - conversation_id = excluded.conversation_id, - turn_state = excluded.turn_state, - instructions_hash = excluded.instructions_hash, - input_tokens = excluded.input_tokens, - function_call_ids_json = excluded.function_call_ids_json, - variant_hash = excluded.variant_hash, - expires_at = excluded.expires_at, - created_at = excluded.created_at", - ) - .bind(response_id) - .bind(&entry.account_id) - .bind(&entry.conversation_id) - .bind(&entry.turn_state) - .bind(&entry.instructions_hash) - .bind(entry.input_tokens.map(u64_to_i64_saturating)) - .bind(function_call_ids_json) - .bind(&entry.variant_hash) - .bind(expires_at.to_rfc3339()) - .bind(entry.created_at.to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(()) - } - - pub async fn list_active( - &self, - now: DateTime, - ) -> SqliteSessionAffinityStoreResult> { - let rows = sqlx::query( - r" -select - response_id, - account_id, - conversation_id, - turn_state, - instructions_hash, - input_tokens, - function_call_ids_json, - variant_hash, - expires_at, - created_at -from session_affinities -where expires_at > ? -order by created_at asc, response_id asc", - ) - .bind(now.to_rfc3339()) - .fetch_all(&self.pool) - .await?; - - rows.into_iter() - .map(|row| stored_session_affinity_from_row(&row)) - .collect() - } - - pub async fn delete_expired( - &self, - now: DateTime, - ) -> SqliteSessionAffinityStoreResult { - let result = sqlx::query("delete from session_affinities where expires_at <= ?") - .bind(now.to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(result.rows_affected()) - } - - pub async fn forget(&self, response_id: &str) -> SqliteSessionAffinityStoreResult { - let result = sqlx::query("delete from session_affinities where response_id = ?") - .bind(response_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } -} - -fn stored_session_affinity_from_row( - row: &sqlx::sqlite::SqliteRow, -) -> SqliteSessionAffinityStoreResult { - use sqlx::Row as _; - - let function_call_ids_json = row.get::("function_call_ids_json"); - Ok(StoredSessionAffinity { - response_id: row.get("response_id"), - entry: SessionAffinityEntry { - account_id: row.get("account_id"), - conversation_id: row.get("conversation_id"), - turn_state: row.get("turn_state"), - instructions_hash: row.get("instructions_hash"), - input_tokens: optional_nonnegative_i64_to_u64(row.get("input_tokens")), - function_call_ids: serde_json::from_str(&function_call_ids_json)?, - variant_hash: row.get("variant_hash"), - created_at: parse_rfc3339(&row.get::("created_at"))?, - }, - expires_at: parse_rfc3339(&row.get::("expires_at"))?, - }) -} - -fn optional_nonnegative_i64_to_u64(value: Option) -> Option { - value.and_then(|value| u64::try_from(value).ok()) -} - -fn u64_to_i64_saturating(value: u64) -> i64 { - value.min(i64::MAX as u64) as i64 -} - -// ==================================================================== -// 运行时会话亲和性服务 -// ==================================================================== - -use std::sync::Arc; -use tokio::sync::RwLock; - -/// 默认会话亲和性 TTL 秒数。 -const DEFAULT_SESSION_AFFINITY_TTL_SECS: i64 = 4 * 60 * 60; - -/// 运行时会话亲和性服务。 -#[derive(Clone)] -pub struct RuntimeSessionAffinityService { - store: SqliteSessionAffinityStore, - map: Arc>, - ttl: Duration, -} - -/// 运行时会话亲和性错误。 -#[derive(Debug, Error)] -pub enum RuntimeSessionAffinityError { - #[error("session affinity store error: {0}")] - Store(#[from] SqliteSessionAffinityStoreError), -} - -impl RuntimeSessionAffinityService { - pub fn new(store: SqliteSessionAffinityStore) -> Self { - let ttl = Duration::seconds(DEFAULT_SESSION_AFFINITY_TTL_SECS); - Self { - store, - map: Arc::new(RwLock::new(SessionAffinityMap::new(ttl))), - ttl, - } - } - - pub async fn restore_from_repository( - &self, - now: DateTime, - ) -> Result { - let records = self.store.list_active(now).await?; - Ok(self.map.write().await.restore(records, now)) - } - - pub async fn record( - &self, - response_id: String, - entry: SessionAffinityEntry, - ) -> Result<(), RuntimeSessionAffinityError> { - self.store.upsert(&response_id, &entry, self.ttl).await?; - self.map.write().await.record(response_id, entry); - Ok(()) - } - - pub async fn lookup_account(&self, response_id: &str, now: DateTime) -> Option { - self.map.read().await.lookup_account(response_id, now) - } - - pub async fn lookup_conversation_id( - &self, - response_id: &str, - now: DateTime, - ) -> Option { - self.map - .read() - .await - .lookup_conversation_id(response_id, now) - } - - pub async fn lookup_turn_state(&self, response_id: &str, now: DateTime) -> Option { - self.map.read().await.lookup_turn_state(response_id, now) - } - - pub async fn lookup_instructions_hash( - &self, - response_id: &str, - now: DateTime, - ) -> Option { - self.map - .read() - .await - .lookup_instructions_hash(response_id, now) - } - - pub async fn lookup_function_call_ids( - &self, - response_id: &str, - now: DateTime, - ) -> Vec { - self.map - .read() - .await - .lookup_function_call_ids(response_id, now) - } - - pub async fn lookup_latest_response_by_conversation( - &self, - conversation_id: &str, - max_age: Option, - variant_hash: Option<&str>, - now: DateTime, - ) -> Option { - self.map - .read() - .await - .lookup_latest_response_by_conversation(conversation_id, max_age, variant_hash, now) - } - - pub async fn lookup_latest_account_by_conversation( - &self, - conversation_id: &str, - max_age: Option, - variant_hash: Option<&str>, - now: DateTime, - ) -> Option { - self.map.read().await.lookup_latest_account_by_conversation( - conversation_id, - max_age, - variant_hash, - now, - ) - } - - pub async fn forget(&self, response_id: &str) -> bool { - let in_memory = self.map.write().await.forget(response_id); - match self.store.forget(response_id).await { - Ok(persisted) => in_memory || persisted, - Err(error) => { - tracing::warn!( - response_id, - error = %error, - "failed to persist session affinity removal" - ); - in_memory - } - } - } -} diff --git a/backend/src/proxy/mod.rs b/backend/src/proxy/mod.rs deleted file mode 100644 index 7d67b71c0..000000000 --- a/backend/src/proxy/mod.rs +++ /dev/null @@ -1,11 +0,0 @@ -//! OpenAI 兼容代理入口。 -//! -//! 本模块负责： -//! - OpenAI API 协议转换（Responses / Models） -//! - 请求调度（账号选择、回退、恢复） -//! - 会话亲和性与 reasoning replay - -pub mod auth; -pub mod dispatch; -pub mod openai; -pub mod router; diff --git a/backend/src/proxy/router.rs b/backend/src/proxy/router.rs deleted file mode 100644 index e47e8ce43..000000000 --- a/backend/src/proxy/router.rs +++ /dev/null @@ -1,12 +0,0 @@ -//! Proxy 路由聚合。 - -use axum::Router; - -use crate::runtime::state::AppState; - -use super::openai; - -/// 构造 Proxy 路由。 -pub fn router() -> Router { - openai::routes::router() -} diff --git a/backend/src/runtime/mod.rs b/backend/src/runtime/mod.rs deleted file mode 100644 index 17d14e096..000000000 --- a/backend/src/runtime/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -//! 运行时层：启动编排、共享状态、服务接线和后台任务生命周期。 - -pub mod bootstrap; -pub mod services; -pub mod shutdown; -pub mod state; -pub mod tasks; diff --git a/backend/src/runtime/services.rs b/backend/src/runtime/services.rs deleted file mode 100644 index 73e6048f6..000000000 --- a/backend/src/runtime/services.rs +++ /dev/null @@ -1,322 +0,0 @@ -//! 运行时服务集合。 - -use std::sync::Arc as StdArc; - -use crate::{ - admin::monitoring::{ - account_usage_service::AdminUsageService, - account_usage_store::SqliteUsageStore, - ops_error_service::AdminOpsErrorLogService, - ops_error_store::SqliteOpsErrorLogStore, - usage_record_service::AdminUsageRecordService, - usage_record_store::{SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPTURE_BODY}, - }, - admin::{ - accounts::service::{AdminAccountService, AdminAccountServiceParts}, - auth::service::{AdminSessionService, SqliteAdminSessionStore}, - keys::service::{ - AdminClientKeyService, ClientKeyService, ClientKeyStoreError, RuntimeClientKeyStore, - SqliteClientKeyStore, - }, - }, - config::{ - schema::AppConfig, - settings::{account_pool_options_from_config, RuntimeSettingsService}, - }, - proxy::dispatch::{ - responses::service::{ResponseDispatchService, ResponseDispatchServiceParts}, - session_affinity::{RuntimeSessionAffinityService, SqliteSessionAffinityStore}, - }, - upstream::accounts::{ - cookies::SqliteCookieStore, - pool::RuntimeAccountPoolService, - store::{AccountStore as AccountStoreTrait, SqliteAccountStore}, - token_refresh::{ - RefreshLeaseStore as SqliteRefreshLeaseStore, RefreshPolicy, RuntimeRefreshPolicy, - RuntimeTokenRefreshService, - }, - }, - upstream::{ - fingerprint::{FingerprintRepository, RuntimeFingerprint}, - models::{ - config::ModelConfig, - service::{ModelService, ModelServiceError}, - snapshot_store::{ModelSnapshotStore, SqliteModelSnapshotStore}, - }, - token_client::{default_openai_token_client, OpenAiTokenClient, TokenClientConfig}, - transport::{ - build_reqwest_client, tls::CustomCaError, CodexBackendClient, CodexModelCatalogClient, - CodexWebSocketPool, CodexWebSocketPoolConfig, - }, - }, -}; - -/// 运行时热路径内存状态初始化错误。 -#[derive(Debug, thiserror::Error)] -pub enum RuntimeServiceInitializationError { - #[error("failed to initialize client key runtime cache: {0}")] - ClientKeys(#[from] ClientKeyStoreError), - #[error("failed to initialize model catalog runtime cache: {0}")] - Models(#[from] ModelServiceError), -} - -// ============================================================================ -// BackgroundTaskStores -// ============================================================================ - -/// 后台任务需要的具体存储适配器集合。 -#[derive(Clone)] -pub struct BackgroundTaskStores { - pub accounts: SqliteAccountStore, - pub admin_sessions: SqliteAdminSessionStore, - pub cookies: SqliteCookieStore, - pub fingerprints: FingerprintRepository, - pub session_affinity: SqliteSessionAffinityStore, - pub refresh_leases: SqliteRefreshLeaseStore, - pub client_keys: SqliteClientKeyStore, - pub usage_records: SqliteUsageRecordStore, -} - -// ============================================================================ -// Services struct -// ============================================================================ - -/// 运行时服务集合。 -#[derive(Clone)] -pub struct Services { - pub models: StdArc, - pub accounts: StdArc, - pub client_keys: StdArc, - pub admin_client_keys: StdArc, - pub admin_sessions: StdArc, - pub settings: StdArc, - pub refresh_policy: RuntimeRefreshPolicy, - pub admin_accounts: StdArc, - pub usage_records: StdArc, - pub ops_errors: StdArc, - pub usage: StdArc, - pub account_pool: StdArc, - pub token_refresh: StdArc>, - pub responses: StdArc, - pub session_affinity: StdArc, - pub codex: StdArc, - pub websocket_pool: Option>, - pub fingerprint: RuntimeFingerprint, - pub installation_id: Option, - pub background_tasks: BackgroundTaskStores, -} - -/// 使用记录运行选项。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct UsageRecordOptions { - pub enabled: bool, - pub capture_body: bool, -} - -impl UsageRecordOptions { - pub fn from_config(config: &AppConfig) -> Self { - Self { - enabled: config.logging.enabled, - capture_body: DEFAULT_USAGE_RECORD_CAPTURE_BODY, - } - } -} - -impl Services { - pub fn new( - config: &AppConfig, - stores: BackgroundTaskStores, - fingerprint: RuntimeFingerprint, - ) -> Self { - Self::try_new(config, stores, fingerprint) - .expect("failed to build runtime services with configured TLS transport") - } - - pub fn try_new( - config: &AppConfig, - stores: BackgroundTaskStores, - fingerprint: RuntimeFingerprint, - ) -> Result { - Self::try_with_installation_id(config, stores, fingerprint, None) - } - - pub(crate) fn try_with_installation_id( - config: &AppConfig, - stores: BackgroundTaskStores, - fingerprint: RuntimeFingerprint, - installation_id: Option, - ) -> Result { - Self::try_with_installation_id_and_usage_record_options( - config, - stores, - fingerprint, - installation_id, - UsageRecordOptions::from_config(config), - ) - } - - pub fn try_with_installation_id_and_usage_record_options( - config: &AppConfig, - stores: BackgroundTaskStores, - fingerprint: RuntimeFingerprint, - installation_id: Option, - usage_record_options: UsageRecordOptions, - ) -> Result { - let installation_id = installation_id.filter(|id| !id.trim().is_empty()); - let account_store_trait = - StdArc::new(stores.accounts.clone()) as StdArc; - let websocket_pool = config.ws_pool.enabled.then(|| { - StdArc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { - enabled: true, - max_age: std::time::Duration::from_millis(config.ws_pool.max_age_ms), - max_per_account: config.ws_pool.max_per_account, - first_token_timeout: (config.ws_pool.first_token_timeout_ms > 0).then(|| { - std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms) - }), - ..CodexWebSocketPoolConfig::default() - })) - }); - let codex = { - let websocket_first_token_timeout = (config.ws_pool.first_token_timeout_ms > 0) - .then(|| std::time::Duration::from_millis(config.ws_pool.first_token_timeout_ms)); - let client = CodexBackendClient::new( - build_reqwest_client(config.tls.force_http11)?, - config.api.base_url.clone(), - fingerprint.clone(), - ) - .with_websocket_first_token_timeout(websocket_first_token_timeout); - if let Some(pool) = &websocket_pool { - StdArc::new(client.with_websocket_pool(pool.clone())) - } else { - StdArc::new(client) - } - }; - let usage_records = StdArc::new(AdminUsageRecordService::new( - stores.usage_records.clone(), - usage_record_options.enabled, - usage_record_options.capture_body, - )); - let ops_errors = StdArc::new(AdminOpsErrorLogService::new( - SqliteOpsErrorLogStore::new(stores.usage_records.pool().clone()), - usage_record_options.capture_body, - )); - let account_pool = StdArc::new(RuntimeAccountPoolService::new( - account_store_trait.clone(), - account_pool_options_from_config(config), - config.auth.request_interval_ms, - )); - let refresh_policy = RuntimeRefreshPolicy::new(RefreshPolicy { - refresh_margin_seconds: config.auth.refresh_margin_seconds, - refresh_concurrency: config.auth.refresh_concurrency, - }); - let settings = StdArc::new(RuntimeSettingsService::new( - config.clone(), - stores.accounts.pool().clone(), - account_pool.clone(), - refresh_policy.clone(), - )); - let admin_sessions = StdArc::new(AdminSessionService::new( - stores.admin_sessions.clone(), - config.admin.default_username.clone(), - config.admin.session_ttl_minutes, - )); - let runtime_client_keys = - StdArc::new(RuntimeClientKeyStore::new(stores.client_keys.clone())); - let admin_client_keys = StdArc::new(AdminClientKeyService::new( - stores.client_keys.clone(), - runtime_client_keys.clone(), - )); - let client_keys = StdArc::new(ClientKeyService::new(runtime_client_keys)); - let token_client = default_openai_token_client(token_client_config(config)); - let token_refresh = StdArc::new( - RuntimeTokenRefreshService::new( - stores.accounts.clone(), - refresh_policy.clone(), - token_client.clone(), - ) - .with_refresh_lease_store(stores.refresh_leases.clone()), - ); - let upstream_client: StdArc = codex.clone(); - let snapshot_store: StdArc = StdArc::new( - SqliteModelSnapshotStore::new(stores.accounts.pool().clone()), - ); - let models = StdArc::new(ModelService::new( - ModelConfig { - model_aliases: config.model_aliases.clone(), - }, - Some(snapshot_store), - Some(upstream_client), - )); - let admin_accounts = StdArc::new(AdminAccountService::new(AdminAccountServiceParts { - store: stores.accounts.clone(), - cookies: stores.cookies.clone(), - codex: codex.clone(), - models: models.clone(), - account_pool: account_pool.clone(), - token_refresher: StdArc::new(token_client), - refresh_leases: stores.refresh_leases.clone(), - oauth: crate::admin::accounts::service::oauth::AccountOAuthService::new( - reqwest::Client::new(), - config.auth.oauth_client_id.clone(), - config.auth.oauth_token_endpoint.clone(), - ), - refresh_policy: refresh_policy.clone(), - installation_id: installation_id.clone(), - })); - let usage_store = SqliteUsageStore::new(stores.accounts.pool().clone()); - let usage = StdArc::new(AdminUsageService::new(usage_store)); - let session_affinity = StdArc::new(RuntimeSessionAffinityService::new( - stores.session_affinity.clone(), - )); - - let cloudflare_recovery = - crate::proxy::dispatch::cloudflare::CloudflareRecovery::new(stores.cookies.clone()); - let responses = StdArc::new(ResponseDispatchService::new(ResponseDispatchServiceParts { - account_pool: account_pool.clone(), - models: models.clone(), - codex: codex.clone(), - session_affinity: session_affinity.clone(), - usage_records: usage_records.clone(), - ops_errors: ops_errors.clone(), - installation_id: installation_id.clone(), - cloudflare: cloudflare_recovery, - })); - - Ok(Self { - models, - accounts: account_store_trait, - client_keys, - admin_client_keys, - admin_sessions, - settings, - refresh_policy, - admin_accounts, - usage_records, - ops_errors, - usage, - account_pool, - token_refresh, - responses, - session_affinity, - codex, - websocket_pool, - fingerprint, - installation_id, - background_tasks: stores, - }) - } - - /// 初始化请求热路径依赖的内存状态。 - pub async fn initialize_hot_path_state(&self) -> Result<(), RuntimeServiceInitializationError> { - self.client_keys.reload_from_store().await?; - self.models.reload_from_store().await?; - Ok(()) - } -} - -fn token_client_config(config: &AppConfig) -> TokenClientConfig { - TokenClientConfig { - client_id: config.auth.oauth_client_id.clone(), - token_endpoint: config.auth.oauth_token_endpoint.clone(), - } -} diff --git a/backend/src/runtime/tasks/session_affinity_cleanup.rs b/backend/src/runtime/tasks/session_affinity_cleanup.rs deleted file mode 100644 index ca4648fb4..000000000 --- a/backend/src/runtime/tasks/session_affinity_cleanup.rs +++ /dev/null @@ -1,65 +0,0 @@ -//! 会话亲和性清理任务。 - -use std::future::Future; - -use chrono::{DateTime, Utc}; - -use crate::proxy::dispatch::session_affinity::{ - SqliteSessionAffinityStore, SqliteSessionAffinityStoreError, -}; - -use super::{ - cleanup::{CleanupTask, CleanupTaskMessages, ExpiredCleanupStore}, - coordinator::SchedulerHandle, -}; - -const MESSAGES: CleanupTaskMessages = CleanupTaskMessages { - started: "会话亲和性清理任务已启动", - deleted: "已清理过期会话亲和性", - empty: "没有需要清理的过期会话亲和性", - failed: "清理过期会话亲和性失败", - stopped: "会话亲和性清理任务已关闭", -}; - -/// 会话亲和性清理任务。 -pub struct SessionAffinityCleanupTask { - task: CleanupTask, -} - -impl SessionAffinityCleanupTask { - /// 构造会话亲和性清理任务。 - pub fn new(store: SqliteSessionAffinityStore, interval_secs: u64) -> Self { - Self { - task: CleanupTask::new(store, interval_secs, MESSAGES), - } - } - - /// 启动后台清理任务。 - pub fn start(self) -> SchedulerHandle { - self.task.start() - } - - /// 执行一次清理。 - pub async fn cleanup_once(&self) -> Result { - self.task.cleanup_once().await - } - - /// 在指定时间点执行一次清理。 - pub async fn cleanup_once_at( - &self, - now: DateTime, - ) -> Result { - self.task.cleanup_once_at(now).await - } -} - -impl ExpiredCleanupStore for SqliteSessionAffinityStore { - type Error = SqliteSessionAffinityStoreError; - - fn delete_expired_at( - &self, - now: DateTime, - ) -> impl Future> + Send + '_ { - self.delete_expired(now) - } -} diff --git a/backend/src/runtime/tasks/session_cleanup.rs b/backend/src/runtime/tasks/session_cleanup.rs deleted file mode 100644 index 64b871da1..000000000 --- a/backend/src/runtime/tasks/session_cleanup.rs +++ /dev/null @@ -1,60 +0,0 @@ -//! 会话清理任务。 - -use std::future::Future; - -use chrono::{DateTime, Utc}; - -use crate::admin::auth::service::SqliteAdminSessionStore; - -use super::{ - cleanup::{CleanupTask, CleanupTaskMessages, ExpiredCleanupStore}, - coordinator::SchedulerHandle, -}; - -const MESSAGES: CleanupTaskMessages = CleanupTaskMessages { - started: "管理员会话清理任务已启动", - deleted: "已清理过期管理员会话", - empty: "没有需要清理的过期管理员会话", - failed: "清理过期管理员会话失败", - stopped: "管理员会话清理任务已关闭", -}; - -/// 管理员会话清理任务。 -pub struct SessionCleanupTask { - task: CleanupTask, -} - -impl SessionCleanupTask { - /// 构造管理员会话清理任务。 - pub fn new(store: SqliteAdminSessionStore, interval_secs: u64) -> Self { - Self { - task: CleanupTask::new(store, interval_secs, MESSAGES), - } - } - - /// 启动后台清理任务。 - pub fn start(self) -> SchedulerHandle { - self.task.start() - } - - /// 执行一次清理。 - pub async fn cleanup_once(&self) -> Result { - self.task.cleanup_once().await - } - - /// 在指定时间点执行一次清理。 - pub async fn cleanup_once_at(&self, now: DateTime) -> Result { - self.task.cleanup_once_at(now).await - } -} - -impl ExpiredCleanupStore for SqliteAdminSessionStore { - type Error = sqlx::Error; - - fn delete_expired_at( - &self, - now: DateTime, - ) -> impl Future> + Send + '_ { - self.cleanup_expired_sessions(now) - } -} diff --git a/backend/src/settings/mod.rs b/backend/src/settings/mod.rs new file mode 100644 index 000000000..f017e31e9 --- /dev/null +++ b/backend/src/settings/mod.rs @@ -0,0 +1,8 @@ +//! 可热更新的运行时设置。 + +pub mod service; +pub mod store; +pub mod types; + +pub use service::SettingsService; +pub use types::*; diff --git a/backend/src/settings/service.rs b/backend/src/settings/service.rs new file mode 100644 index 000000000..2ba53f578 --- /dev/null +++ b/backend/src/settings/service.rs @@ -0,0 +1,188 @@ +//! 运行时设置校验、持久化与快照广播。 + +use std::{collections::BTreeMap, sync::Arc}; + +use subtle::ConstantTimeEq; +use tokio::sync::watch; + +use crate::infra::identity::{generate_admin_api_key, hash_credential}; + +use super::{ + store::PgSettingsStore, + types::{ + ManagementApiKeyStatus, SettingsError, SettingsPatch, SettingsSnapshot, + SettingsValidationError, + }, +}; + +const ROTATION_STRATEGIES: [&str; 4] = ["smart", "quota_reset_priority", "round_robin", "sticky"]; + +/// 无状态的设置校验服务。 +#[derive(Debug, Clone, Default)] +pub struct SettingsService; + +impl SettingsService { + pub fn apply_patch( + current: &mut SettingsSnapshot, + patch: SettingsPatch, + ) -> Result<(), SettingsValidationError> { + if let Some(model_aliases) = patch.model_aliases { + current.model_aliases = validate_model_aliases(model_aliases)?; + } + if let Some(value) = patch.refresh_margin_seconds { + current.refresh_margin_seconds = positive_u64("refreshMarginSeconds", value)?; + } + if let Some(value) = patch.refresh_concurrency { + current.refresh_concurrency = positive_u32("refreshConcurrency", value)?; + } + if let Some(value) = patch.max_concurrent_per_account { + current.max_concurrent_per_account = positive_usize("maxConcurrentPerAccount", value)?; + } + if let Some(value) = patch.request_interval_ms { + current.request_interval_ms = value; + } + if let Some(value) = patch.rotation_strategy { + current.rotation_strategy = validate_rotation_strategy(&value)?; + } + Ok(()) + } +} + +/// 持久化设置的当前快照与变更广播。 +#[derive(Clone)] +pub struct RuntimeSettingsService { + store: PgSettingsStore, + sender: watch::Sender>, +} + +impl RuntimeSettingsService { + pub fn new(settings: SettingsSnapshot, pool: sqlx::PgPool) -> Self { + let (sender, _) = watch::channel(Arc::new(settings)); + Self { + store: PgSettingsStore::new(pool), + sender, + } + } + + pub async fn load_or_initialize( + defaults: SettingsSnapshot, + pool: &sqlx::PgPool, + ) -> Result { + PgSettingsStore::new(pool.clone()) + .load_or_initialize(&defaults) + .await + } + + pub fn current(&self) -> Arc { + self.sender.borrow().clone() + } + + pub fn subscribe(&self) -> watch::Receiver> { + self.sender.subscribe() + } + + pub async fn update( + &self, + patch: SettingsPatch, + ) -> Result, SettingsError> { + let mut next = (*self.current()).clone(); + SettingsService::apply_patch(&mut next, patch)?; + self.store.save(&next).await?; + let next = Arc::new(next); + self.sender.send_replace(next.clone()); + Ok(next) + } + + pub async fn admin_api_key_status(&self) -> Result { + self.store.admin_api_key_status().await + } + + pub async fn regenerate_admin_api_key(&self) -> Result { + self.store.ensure(&self.current()).await?; + let key = generate_admin_api_key(); + self.store + .set_admin_api_key_hash(&hash_credential(&key)) + .await?; + Ok(key) + } + + pub async fn delete_admin_api_key(&self) -> Result<(), SettingsError> { + self.store.clear_admin_api_key_hash().await + } + + pub async fn verify_admin_api_key(&self, key: &str) -> Result { + if key.is_empty() { + return Ok(false); + } + let stored = self.store.load_admin_api_key_hash().await?; + let key_hash = hash_credential(key); + Ok(stored + .as_deref() + .filter(|stored_hash| !stored_hash.is_empty()) + .is_some_and(|stored_hash| key_hash.as_bytes().ct_eq(stored_hash.as_bytes()).into())) + } +} + +fn validate_model_aliases( + aliases: BTreeMap, +) -> Result, SettingsValidationError> { + let mut normalized = BTreeMap::new(); + for (alias, target) in aliases { + let alias = nonempty("modelAliases", &alias)?; + let target = nonempty("modelAliases", &target)?; + if alias == target { + return Err(invalid_field( + "modelAliases", + "alias and target must differ", + )); + } + normalized.insert(alias, target); + } + Ok(normalized) +} + +fn validate_rotation_strategy(strategy: &str) -> Result { + let strategy = nonempty("rotationStrategy", strategy)?; + if ROTATION_STRATEGIES.contains(&strategy.as_str()) { + Ok(strategy) + } else { + Err(invalid_field( + "rotationStrategy", + "must be one of smart, quota_reset_priority, round_robin, sticky", + )) + } +} + +fn nonempty(field: &str, value: &str) -> Result { + let value = value.trim().to_string(); + if value.is_empty() { + Err(invalid_field(field, "must not be empty")) + } else { + Ok(value) + } +} + +fn positive_u64(field: &str, value: u64) -> Result { + (value > 0) + .then_some(value) + .ok_or_else(|| invalid_field(field, "must be greater than 0")) +} + +fn positive_u32(field: &str, value: u32) -> Result { + (value > 0) + .then_some(value) + .ok_or_else(|| invalid_field(field, "must be greater than 0")) +} + +fn positive_usize(field: &str, value: usize) -> Result { + (value > 0) + .then_some(value) + .ok_or_else(|| invalid_field(field, "must be greater than 0")) +} + +fn invalid_field(field: &str, message: impl Into) -> SettingsValidationError { + SettingsValidationError::InvalidField { + field: field.to_string(), + message: message.into(), + } +} diff --git a/backend/src/settings/store.rs b/backend/src/settings/store.rs new file mode 100644 index 000000000..a406faa79 --- /dev/null +++ b/backend/src/settings/store.rs @@ -0,0 +1,210 @@ +//! 运行时设置 PostgreSQL 存取。 + +use std::collections::BTreeMap; + +use chrono::Utc; +use sqlx::{postgres::PgRow, PgPool, Row}; + +use super::types::{ManagementApiKeyStatus, SettingsError, SettingsSnapshot}; + +const RUNTIME_SETTINGS_ID: i64 = 1; + +#[derive(Clone)] +pub struct PgSettingsStore { + pool: PgPool, +} + +impl PgSettingsStore { + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + pub async fn load_or_initialize( + &self, + defaults: &SettingsSnapshot, + ) -> Result { + self.ensure(defaults).await?; + self.load().await + } + + pub async fn save(&self, settings: &SettingsSnapshot) -> Result<(), SettingsError> { + sqlx::query( + r" +insert into runtime_settings ( + id, model_aliases_json, refresh_margin_seconds, refresh_concurrency, + max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at +) values ($1, $2, $3, $4, $5, $6, $7, $8) +on conflict(id) do update set + model_aliases_json = excluded.model_aliases_json, + refresh_margin_seconds = excluded.refresh_margin_seconds, + refresh_concurrency = excluded.refresh_concurrency, + max_concurrent_per_account = excluded.max_concurrent_per_account, + request_interval_ms = excluded.request_interval_ms, + rotation_strategy = excluded.rotation_strategy, + updated_at = excluded.updated_at", + ) + .bind(RUNTIME_SETTINGS_ID) + .bind(sqlx::types::Json(&settings.model_aliases)) + .bind(to_i64( + "refreshMarginSeconds", + settings.refresh_margin_seconds, + )?) + .bind(i64::from(settings.refresh_concurrency)) + .bind(to_i64( + "maxConcurrentPerAccount", + settings.max_concurrent_per_account, + )?) + .bind(to_i64("requestIntervalMs", settings.request_interval_ms)?) + .bind(&settings.rotation_strategy) + .bind(Utc::now()) + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn admin_api_key_status(&self) -> Result { + let key_hash = self.load_admin_api_key_hash().await?; + Ok(ManagementApiKeyStatus { + exists: key_hash.is_some_and(|hash| !hash.is_empty()), + }) + } + + pub async fn load_admin_api_key_hash(&self) -> Result, SettingsError> { + Ok(sqlx::query_scalar::<_, Option>( + "select admin_api_key_hash from runtime_settings where id = $1", + ) + .bind(RUNTIME_SETTINGS_ID) + .fetch_optional(&self.pool) + .await? + .flatten()) + } + + pub async fn set_admin_api_key_hash(&self, key_hash: &str) -> Result<(), SettingsError> { + sqlx::query( + "update runtime_settings + set admin_api_key_hash = $1, updated_at = $2 + where id = $3", + ) + .bind(key_hash) + .bind(Utc::now()) + .bind(RUNTIME_SETTINGS_ID) + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn clear_admin_api_key_hash(&self) -> Result<(), SettingsError> { + sqlx::query( + "update runtime_settings + set admin_api_key_hash = null, updated_at = $1 + where id = $2", + ) + .bind(Utc::now()) + .bind(RUNTIME_SETTINGS_ID) + .execute(&self.pool) + .await?; + Ok(()) + } + + pub async fn ensure(&self, settings: &SettingsSnapshot) -> Result<(), SettingsError> { + sqlx::query( + r" +insert into runtime_settings ( + id, model_aliases_json, refresh_margin_seconds, refresh_concurrency, + max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at +) values ($1, $2, $3, $4, $5, $6, $7, $8) +on conflict (id) do nothing", + ) + .bind(RUNTIME_SETTINGS_ID) + .bind(sqlx::types::Json(&settings.model_aliases)) + .bind(to_i64( + "refreshMarginSeconds", + settings.refresh_margin_seconds, + )?) + .bind(i64::from(settings.refresh_concurrency)) + .bind(to_i64( + "maxConcurrentPerAccount", + settings.max_concurrent_per_account, + )?) + .bind(to_i64("requestIntervalMs", settings.request_interval_ms)?) + .bind(&settings.rotation_strategy) + .bind(Utc::now()) + .execute(&self.pool) + .await?; + Ok(()) + } + + async fn load(&self) -> Result { + let row = sqlx::query( + "select model_aliases_json, refresh_margin_seconds, refresh_concurrency, + max_concurrent_per_account, request_interval_ms, rotation_strategy + from runtime_settings where id = $1", + ) + .bind(RUNTIME_SETTINGS_ID) + .fetch_one(&self.pool) + .await?; + snapshot_from_row(&row) + } +} + +fn snapshot_from_row(row: &PgRow) -> Result { + let aliases: sqlx::types::Json> = row.get("model_aliases_json"); + Ok(SettingsSnapshot { + model_aliases: aliases.0, + refresh_margin_seconds: positive_u64( + "refreshMarginSeconds", + row.get("refresh_margin_seconds"), + )?, + refresh_concurrency: positive_u32("refreshConcurrency", row.get("refresh_concurrency"))?, + max_concurrent_per_account: positive_usize( + "maxConcurrentPerAccount", + row.get("max_concurrent_per_account"), + )?, + request_interval_ms: nonnegative_u64("requestIntervalMs", row.get("request_interval_ms"))?, + rotation_strategy: row.get("rotation_strategy"), + }) +} + +fn to_i64(field: &'static str, value: impl TryInto) -> Result { + value + .try_into() + .map_err(|_| stored_field_error(field, "out of range")) +} + +fn positive_u64(field: &'static str, value: i64) -> Result { + if value <= 0 { + return Err(stored_field_error(field, "must be greater than 0")); + } + u64::try_from(value).map_err(|_| stored_field_error(field, "out of range")) +} + +fn positive_u32(field: &'static str, value: i64) -> Result { + if value <= 0 { + return Err(stored_field_error(field, "must be greater than 0")); + } + u32::try_from(value).map_err(|_| stored_field_error(field, "out of range")) +} + +fn positive_usize(field: &'static str, value: i64) -> Result { + if value <= 0 { + return Err(stored_field_error(field, "must be greater than 0")); + } + usize::try_from(value).map_err(|_| stored_field_error(field, "out of range")) +} + +fn nonnegative_u64(field: &'static str, value: i64) -> Result { + if value < 0 { + return Err(stored_field_error( + field, + "must be greater than or equal to 0", + )); + } + u64::try_from(value).map_err(|_| stored_field_error(field, "out of range")) +} + +fn stored_field_error(field: impl Into, message: impl Into) -> SettingsError { + SettingsError::StoredField { + field: field.into(), + message: message.into(), + } +} diff --git a/backend/src/settings/types.rs b/backend/src/settings/types.rs new file mode 100644 index 000000000..295a279ce --- /dev/null +++ b/backend/src/settings/types.rs @@ -0,0 +1,75 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +/// 可热更新的运行时设置快照。 +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SettingsSnapshot { + pub model_aliases: BTreeMap, + pub refresh_margin_seconds: u64, + pub refresh_concurrency: u32, + pub max_concurrent_per_account: usize, + pub request_interval_ms: u64, + pub rotation_strategy: String, +} + +impl Default for SettingsSnapshot { + fn default() -> Self { + Self { + model_aliases: BTreeMap::new(), + refresh_margin_seconds: 3600, + refresh_concurrency: 2, + max_concurrent_per_account: 3, + request_interval_ms: 50, + rotation_strategy: "smart".to_string(), + } + } +} + +/// 管理端设置补丁。 +#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SettingsPatch { + pub model_aliases: Option>, + pub refresh_margin_seconds: Option, + pub refresh_concurrency: Option, + pub max_concurrent_per_account: Option, + pub request_interval_ms: Option, + pub rotation_strategy: Option, +} + +/// 管理员 API Key 状态。 +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ManagementApiKeyStatus { + pub exists: bool, +} + +/// 设置字段校验错误。 +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum SettingsValidationError { + #[error("invalid setting `{field}`: {message}")] + InvalidField { field: String, message: String }, +} + +impl SettingsValidationError { + pub fn field(&self) -> &str { + match self { + Self::InvalidField { field, .. } => field, + } + } +} + +/// 运行时设置服务错误。 +#[derive(Debug, Error)] +pub enum SettingsError { + #[error(transparent)] + InvalidField(#[from] SettingsValidationError), + #[error("runtime settings database error: {0}")] + Database(#[from] sqlx::Error), + #[error("runtime settings json error: {0}")] + Json(#[from] serde_json::Error), + #[error("invalid stored setting `{field}`: {message}")] + StoredField { field: String, message: String }, +} diff --git a/backend/src/telemetry/account_usage/mod.rs b/backend/src/telemetry/account_usage/mod.rs new file mode 100644 index 000000000..f839156b7 --- /dev/null +++ b/backend/src/telemetry/account_usage/mod.rs @@ -0,0 +1,4 @@ +//! 账号累计用量与时间桶查询。 + +pub mod query; +pub mod store; diff --git a/backend/src/admin/monitoring/account_usage_service.rs b/backend/src/telemetry/account_usage/query.rs similarity index 73% rename from backend/src/admin/monitoring/account_usage_service.rs rename to backend/src/telemetry/account_usage/query.rs index 860f823b4..0d2cb0395 100644 --- a/backend/src/admin/monitoring/account_usage_service.rs +++ b/backend/src/telemetry/account_usage/query.rs @@ -1,23 +1,33 @@ //! 用量聚合模型、端口与策略服务。 +use std::collections::HashMap; + use chrono::{DateTime, Utc}; use thiserror::Error; -use crate::admin::monitoring::{ - account_usage_store::{SqliteUsageStore, UsageListRecord, UsageSummary, UsageTimeBucketRecord}, - billing, -}; use crate::infra::{format::nonnegative_i64_to_u64, json::Page}; +use crate::telemetry::{ + account_usage::store::{PgAccountUsageStore, UsageListRecord, UsageSummary}, + billing, + buckets::query::{ + ModelBucketUsage, ModelUsageWindow, PgRequestBucketQuery, UsageBucketTotals, + UsageBucketWindow, UsageTimeBucketRecord, + }, +}; /// 管理端用量统计服务。 #[derive(Clone)] -pub struct AdminUsageService { - store: SqliteUsageStore, +pub struct AccountUsageQueryService { + store: PgAccountUsageStore, + buckets: PgRequestBucketQuery, } -impl AdminUsageService { +impl AccountUsageQueryService { /// 构造服务。 - pub fn new(store: SqliteUsageStore) -> Self { - Self { store } + pub fn new(store: PgAccountUsageStore) -> Self { + Self { + buckets: PgRequestBucketQuery::new(store.pool().clone()), + store, + } } /// 分页列出账号用量。 @@ -25,14 +35,18 @@ impl AdminUsageService { &self, cursor: Option, limit: u32, - ) -> Result, AdminUsageError> { + ) -> Result, AccountUsageQueryError> { let page = self .store .list_usage(cursor, limit) .await - .map_err(|_| AdminUsageError::List)?; + .map_err(|_| AccountUsageQueryError::List)?; Ok(Page { - items: page.items.into_iter().map(AdminUsageRecord::from).collect(), + items: page + .items + .into_iter() + .map(AccountUsageRecord::from) + .collect(), next_cursor: page.next_cursor, }) } @@ -41,21 +55,21 @@ impl AdminUsageService { pub async fn list_by_account_ids( &self, account_ids: &[String], - ) -> Result, AdminUsageError> { + ) -> Result, AccountUsageQueryError> { self.store .list_usage_by_account_ids(account_ids) .await - .map(|records| records.into_iter().map(AdminUsageRecord::from).collect()) - .map_err(|_| AdminUsageError::List) + .map(|records| records.into_iter().map(AccountUsageRecord::from).collect()) + .map_err(|_| AccountUsageQueryError::List) } /// 汇总账号用量。 - pub async fn summary(&self) -> Result { + pub async fn summary(&self) -> Result { self.store .usage_summary() .await - .map(AdminUsageSummary::from) - .map_err(|_| AdminUsageError::Summary) + .map(AccountUsageSummary::from) + .map_err(|_| AccountUsageQueryError::Summary) } /// 列出指定时间范围内的聚合时间桶。 @@ -63,23 +77,43 @@ impl AdminUsageService { &self, start: DateTime, end: DateTime, - ) -> Result, AdminUsageError> { - self.store - .list_time_buckets(start, end) + ) -> Result, AccountUsageQueryError> { + self.buckets + .list(start, end) .await .map(|records| { records .into_iter() - .map(AdminUsageTimeBucketRecord::from) + .map(AccountUsageTimeBucket::from) .collect() }) - .map_err(|_| AdminUsageError::TimeBuckets) + .map_err(|_| AccountUsageQueryError::TimeBuckets) + } + + pub async fn usage_by_windows( + &self, + windows: &[UsageBucketWindow], + ) -> Result>, AccountUsageQueryError> { + self.buckets + .usage_by_windows(windows) + .await + .map_err(|_| AccountUsageQueryError::TimeBuckets) + } + + pub async fn model_usage_by_windows( + &self, + windows: &[ModelUsageWindow], + ) -> Result, AccountUsageQueryError> { + self.buckets + .model_usage_by_windows(windows) + .await + .map_err(|_| AccountUsageQueryError::TimeBuckets) } } /// 管理端用量统计错误。 #[derive(Debug, Error)] -pub enum AdminUsageError { +pub enum AccountUsageQueryError { #[error("failed to list account usage")] List, #[error("failed to summarize account usage")] @@ -90,7 +124,7 @@ pub enum AdminUsageError { /// 管理端账号用量记录。 #[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminUsageRecord { +pub struct AccountUsageRecord { pub account_id: String, pub email: Option, pub label: Option, @@ -118,7 +152,7 @@ pub struct AdminUsageRecord { /// 管理端账号用量汇总。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct AdminUsageSummary { +pub struct AccountUsageSummary { pub account_count: i64, pub request_count: i64, pub empty_response_count: i64, @@ -135,7 +169,7 @@ pub struct AdminUsageSummary { /// 管理端时间桶用量记录。 #[derive(Debug, Clone, PartialEq)] -pub struct AdminUsageTimeBucketRecord { +pub struct AccountUsageTimeBucket { pub bucket_start: DateTime, pub model: String, pub service_tier: Option, @@ -153,7 +187,7 @@ pub struct AdminUsageTimeBucketRecord { pub cost_usd: Option, } -impl From for AdminUsageSummary { +impl From for AccountUsageSummary { fn from(s: UsageSummary) -> Self { Self { account_count: s.account_count, @@ -172,7 +206,7 @@ impl From for AdminUsageSummary { } } -impl From for AdminUsageTimeBucketRecord { +impl From for AccountUsageTimeBucket { fn from(record: UsageTimeBucketRecord) -> Self { let cost_usd = time_bucket_cost_usd(&record); Self { @@ -217,7 +251,7 @@ fn time_bucket_cost_usd(record: &UsageTimeBucketRecord) -> Option { )) } -impl From for AdminUsageRecord { +impl From for AccountUsageRecord { fn from(usage: UsageListRecord) -> Self { Self { account_id: usage.account_id, diff --git a/backend/src/admin/monitoring/account_usage_store.rs b/backend/src/telemetry/account_usage/store.rs similarity index 55% rename from backend/src/admin/monitoring/account_usage_store.rs rename to backend/src/telemetry/account_usage/store.rs index 68978460d..d438800da 100644 --- a/backend/src/admin/monitoring/account_usage_store.rs +++ b/backend/src/telemetry/account_usage/store.rs @@ -1,14 +1,12 @@ -//! SQLite 用量存储。 +//! PostgreSQL 账号用量查询存储。 use chrono::{DateTime, Utc}; -use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool}; +use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row}; use thiserror::Error; use crate::infra::{ json::{decode_cursor, encode_cursor, Page}, - time::{ - parse_optional_rfc3339_utc as parse_optional_rfc3339, parse_rfc3339_utc as parse_rfc3339, - }, + time::parse_rfc3339_utc as parse_rfc3339, }; const LIST_USAGE_SELECT_SQL: &str = r" @@ -35,63 +33,43 @@ select au.window_started_at, au.window_reset_at, au.limit_window_seconds, - au.last_used_at + au.last_used_at, + coalesce(au.last_used_at, 'epoch'::timestamptz) as sort_last_used_at from account_usage au left join accounts a on a.id = au.account_id"; const USAGE_SUMMARY_SQL: &str = r" select count(*) as account_count, - coalesce(sum(request_count), 0) as request_count, - coalesce(sum(empty_response_count), 0) as empty_response_count, - coalesce(sum(input_tokens), 0) as input_tokens, - coalesce(sum(output_tokens), 0) as output_tokens, - coalesce(sum(cached_tokens), 0) as cached_tokens, - coalesce(sum(reasoning_tokens), 0) as reasoning_tokens, - coalesce(sum(total_tokens), 0) as total_tokens, - coalesce(sum(image_input_tokens), 0) as image_input_tokens, - coalesce(sum(image_output_tokens), 0) as image_output_tokens, - coalesce(sum(image_request_count), 0) as image_request_count, - coalesce(sum(image_request_failed_count), 0) as image_request_failed_count + coalesce(sum(request_count), 0)::bigint as request_count, + coalesce(sum(empty_response_count), 0)::bigint as empty_response_count, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + coalesce(sum(reasoning_tokens), 0)::bigint as reasoning_tokens, + coalesce(sum(total_tokens), 0)::bigint as total_tokens, + coalesce(sum(image_input_tokens), 0)::bigint as image_input_tokens, + coalesce(sum(image_output_tokens), 0)::bigint as image_output_tokens, + coalesce(sum(image_request_count), 0)::bigint as image_request_count, + coalesce(sum(image_request_failed_count), 0)::bigint as image_request_failed_count from account_usage"; -const LIST_USAGE_TIME_BUCKETS_SQL: &str = r" -select - bucket_start, - model, - service_tier, - coalesce(sum(request_count), 0) as request_count, - coalesce(sum(error_count), 0) as error_count, - coalesce(sum(input_tokens), 0) as input_tokens, - coalesce(sum(output_tokens), 0) as output_tokens, - coalesce(sum(cached_tokens), 0) as cached_tokens, - coalesce(sum(first_token_latency_sum), 0) as first_token_latency_sum, - coalesce(sum(first_token_latency_count), 0) as first_token_latency_count, - coalesce(sum(latency_sum), 0) as latency_sum, - coalesce(sum(latency_count), 0) as latency_count, - coalesce(max(max_latency_ms), 0) as max_latency_ms, - coalesce(min(nullif(min_latency_ms, 0)), 0) as min_latency_ms -from usage_time_buckets -where bucket_start >= ? and bucket_start <= ? -group by bucket_start, model, service_tier -order by bucket_start asc, model asc, service_tier asc"; - -/// SQLite 用量存储错误。 +/// PostgreSQL 用量存储错误。 #[derive(Debug, Error)] -pub enum SqliteUsageStoreError { +pub enum PgAccountUsageStoreError { /// 数据库错误。 - #[error("sqlite usage store database error: {0}")] + #[error("PostgreSQL usage store database error: {0}")] Database(#[from] sqlx::Error), /// 时间格式错误。 - #[error("sqlite usage store timestamp error: {0}")] + #[error("PostgreSQL usage store timestamp error: {0}")] Timestamp(#[from] chrono::ParseError), /// 分页游标非法。 #[error("invalid usage pagination cursor")] InvalidCursor, } -/// SQLite 用量存储结果。 -pub type SqliteUsageStoreResult = Result; +/// PostgreSQL 用量存储结果。 +pub type PgAccountUsageStoreResult = Result; /// 账号用量列表记录。 #[derive(Debug, Clone)] @@ -173,72 +151,47 @@ pub struct UsageSummary { pub image_request_failed_count: i64, } -/// 时间桶聚合用量记录。 -#[derive(Debug, Clone)] -pub struct UsageTimeBucketRecord { - /// 桶开始时间。 - pub bucket_start: DateTime, - /// 模型。 - pub model: String, - /// 服务层级。 - pub service_tier: Option, - /// 请求数。 - pub request_count: i64, - /// 错误数。 - pub error_count: i64, - /// 输入 token。 - pub input_tokens: i64, - /// 输出 token。 - pub output_tokens: i64, - /// 缓存 token。 - pub cached_tokens: i64, - /// 首 token 延迟总和。 - pub first_token_latency_sum: i64, - /// 首 token 延迟样本数。 - pub first_token_latency_count: i64, - /// 完成延迟总和。 - pub latency_sum: i64, - /// 完成延迟样本数。 - pub latency_count: i64, - /// 最大完成延迟。 - pub max_latency_ms: i64, - /// 最小完成延迟。 - pub min_latency_ms: i64, -} - -/// SQLite 用量存储。 +/// PostgreSQL 用量存储。 #[derive(Clone)] -pub struct SqliteUsageStore { - pool: SqlitePool, +pub struct PgAccountUsageStore { + pool: PgPool, } -impl SqliteUsageStore { +impl PgAccountUsageStore { /// 构造存储。 - pub fn new(pool: SqlitePool) -> Self { + pub fn new(pool: PgPool) -> Self { Self { pool } } + pub fn pool(&self) -> &PgPool { + &self.pool + } + /// 分页列出账号用量。 pub async fn list_usage( &self, cursor: Option, limit: u32, - ) -> SqliteUsageStoreResult> { + ) -> PgAccountUsageStoreResult> { let limit = limit.clamp(1, 200); let fetch_limit = i64::from(limit) + 1; - let mut builder = QueryBuilder::::new(LIST_USAGE_SELECT_SQL); + let mut builder = QueryBuilder::::new(LIST_USAGE_SELECT_SQL); if let Some(cursor) = cursor { let (last_used_at, account_id) = - decode_cursor(&cursor).ok_or(SqliteUsageStoreError::InvalidCursor)?; - builder.push("\nwhere au.last_used_at < "); - builder.push_bind(last_used_at.clone()); - builder.push("\n or (au.last_used_at = "); + decode_cursor(&cursor).ok_or(PgAccountUsageStoreError::InvalidCursor)?; + let last_used_at = parse_rfc3339(&last_used_at) + .map_err(|_| PgAccountUsageStoreError::InvalidCursor)?; + builder.push("\nwhere coalesce(au.last_used_at, 'epoch'::timestamptz) < "); + builder.push_bind(last_used_at); + builder.push("\n or (coalesce(au.last_used_at, 'epoch'::timestamptz) = "); builder.push_bind(last_used_at); builder.push(" and au.account_id < "); builder.push_bind(account_id); builder.push(")"); } - builder.push("\norder by au.last_used_at desc, au.account_id desc\nlimit "); + builder.push( + "\norder by coalesce(au.last_used_at, 'epoch'::timestamptz) desc, au.account_id desc\nlimit ", + ); builder.push_bind(fetch_limit); let rows = builder.build().fetch_all(&self.pool).await?; @@ -249,12 +202,12 @@ impl SqliteUsageStore { pub async fn list_usage_by_account_ids( &self, account_ids: &[String], - ) -> SqliteUsageStoreResult> { + ) -> PgAccountUsageStoreResult> { if account_ids.is_empty() { return Ok(Vec::new()); } - let mut builder = QueryBuilder::::new(LIST_USAGE_SELECT_SQL); + let mut builder = QueryBuilder::::new(LIST_USAGE_SELECT_SQL); builder.push("\nwhere au.account_id in ("); let mut separated = builder.separated(", "); for account_id in account_ids { @@ -267,7 +220,7 @@ impl SqliteUsageStore { } /// 汇总账号用量。 - pub async fn usage_summary(&self) -> SqliteUsageStoreResult { + pub async fn usage_summary(&self) -> PgAccountUsageStoreResult { let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?; Ok(UsageSummary { account_count: row.get("account_count"), @@ -284,23 +237,9 @@ impl SqliteUsageStore { image_request_failed_count: row.get("image_request_failed_count"), }) } - - /// 列出指定时间范围内的时间桶聚合用量。 - pub async fn list_time_buckets( - &self, - start: DateTime, - end: DateTime, - ) -> SqliteUsageStoreResult> { - let rows = sqlx::query(LIST_USAGE_TIME_BUCKETS_SQL) - .bind(start.to_rfc3339()) - .bind(end.to_rfc3339()) - .fetch_all(&self.pool) - .await?; - rows.iter().map(usage_time_bucket_from_row).collect() - } } -fn usage_list_from_row(row: &SqliteRow) -> SqliteUsageStoreResult { +fn usage_list_from_row(row: &PgRow) -> PgAccountUsageStoreResult { Ok(UsageListRecord { account_id: row.get("account_id"), email: row.get("email"), @@ -321,45 +260,19 @@ fn usage_list_from_row(row: &SqliteRow) -> SqliteUsageStoreResult, _>("window_started_at").as_deref(), - )?, - window_reset_at: parse_optional_rfc3339( - row.get::, _>("window_reset_at").as_deref(), - )?, + window_started_at: row.get("window_started_at"), + window_reset_at: row.get("window_reset_at"), limit_window_seconds: row .get::, _>("limit_window_seconds") .and_then(|value| u64::try_from(value).ok()), - last_used_at: parse_optional_rfc3339( - row.get::, _>("last_used_at").as_deref(), - )?, + last_used_at: row.get("last_used_at"), }) } -fn usage_time_bucket_from_row(row: &SqliteRow) -> SqliteUsageStoreResult { - let service_tier = row.get::("service_tier").trim().to_string(); - Ok(UsageTimeBucketRecord { - bucket_start: parse_rfc3339(row.get::<&str, _>("bucket_start"))?, - model: row.get("model"), - service_tier: (!service_tier.is_empty()).then_some(service_tier), - request_count: row.get("request_count"), - error_count: row.get("error_count"), - input_tokens: row.get("input_tokens"), - output_tokens: row.get("output_tokens"), - cached_tokens: row.get("cached_tokens"), - first_token_latency_sum: row.get("first_token_latency_sum"), - first_token_latency_count: row.get("first_token_latency_count"), - latency_sum: row.get("latency_sum"), - latency_count: row.get("latency_count"), - max_latency_ms: row.get("max_latency_ms"), - min_latency_ms: row.get("min_latency_ms"), - }) -} - -fn to_page(rows: &[SqliteRow], limit: u32) -> Page { +fn to_page(rows: &[PgRow], limit: u32) -> Page { let has_more = rows.len() > limit as usize; let mut items = Vec::with_capacity(limit as usize); - let mut last_row: Option<&SqliteRow> = None; + let mut last_row: Option<&PgRow> = None; for (i, row) in rows.iter().enumerate() { if i >= limit as usize { break; @@ -371,7 +284,9 @@ fn to_page(rows: &[SqliteRow], limit: u32) -> Page { } let next_cursor = if has_more { last_row.map(|row| { - let ts: String = row.get("last_used_at"); + let ts = row + .get::, _>("sort_last_used_at") + .to_rfc3339(); let id: String = row.get("account_id"); encode_cursor(&ts, &id) }) diff --git a/backend/src/admin/monitoring/billing.rs b/backend/src/telemetry/billing.rs similarity index 99% rename from backend/src/admin/monitoring/billing.rs rename to backend/src/telemetry/billing.rs index 8593a248c..e4e1a1444 100644 --- a/backend/src/admin/monitoring/billing.rs +++ b/backend/src/telemetry/billing.rs @@ -152,7 +152,7 @@ const MODEL_PRICING_RULES: &[ModelPricingRule] = &[ /// 单次请求计费明细。 #[derive(Debug, Clone, PartialEq)] -pub(in crate::admin) struct CostBreakdown { +pub(crate) struct CostBreakdown { /// 输入成本。 pub input_cost: f64, /// 输出成本。 @@ -174,7 +174,7 @@ pub(in crate::admin) struct CostBreakdown { } /// 按 codex2api 的 billing 引擎口径计算美元计费明细。 -pub(in crate::admin) fn calculate_cost_breakdown( +pub(crate) fn calculate_cost_breakdown( input_tokens: u64, output_tokens: u64, cached_tokens: u64, @@ -241,7 +241,7 @@ pub(in crate::admin) fn calculate_cost_breakdown( } /// 按 codex2api 的 billing 引擎口径计算美元计费。 -pub(in crate::admin) fn calculate_cost( +pub(crate) fn calculate_cost( input_tokens: u64, output_tokens: u64, cached_tokens: u64, diff --git a/backend/src/telemetry/buckets/mod.rs b/backend/src/telemetry/buckets/mod.rs new file mode 100644 index 000000000..58e6b58dd --- /dev/null +++ b/backend/src/telemetry/buckets/mod.rs @@ -0,0 +1,4 @@ +//! 请求时间桶聚合存储。 + +pub mod query; +pub mod store; diff --git a/backend/src/telemetry/buckets/query.rs b/backend/src/telemetry/buckets/query.rs new file mode 100644 index 000000000..9b21b3f64 --- /dev/null +++ b/backend/src/telemetry/buckets/query.rs @@ -0,0 +1,276 @@ +//! PostgreSQL 请求时间桶聚合查询。 + +use std::collections::{HashMap, HashSet}; + +use chrono::{DateTime, Utc}; +use sqlx::{PgPool, Postgres, QueryBuilder, Row}; +use thiserror::Error; + +use crate::infra::time::china_quarter_hour_start; + +const LIST_USAGE_TIME_BUCKETS_SQL: &str = r" +select + bucket_start, + model, + nullif(service_tier, '__unknown__') as service_tier, + coalesce(sum(success_count + error_count), 0)::bigint as request_count, + coalesce(sum(error_count), 0)::bigint as error_count, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + coalesce(sum(first_token_latency_sum), 0)::bigint as first_token_latency_sum, + coalesce(sum(first_token_latency_count), 0)::bigint as first_token_latency_count, + coalesce(sum(latency_sum), 0)::bigint as latency_sum, + coalesce(sum(latency_count), 0)::bigint as latency_count, + coalesce(max(max_latency_ms), 0) as max_latency_ms, + coalesce(min(min_latency_ms), 0) as min_latency_ms +from request_time_buckets +where bucket_start >= $1 and bucket_start <= $2 +group by bucket_start, model, service_tier +order by bucket_start asc, model asc, service_tier asc"; + +/// 请求时间桶查询错误。 +#[derive(Debug, Error)] +pub enum PgRequestBucketQueryError { + /// 数据库查询失败。 + #[error("PostgreSQL request bucket query failed: {0}")] + Database(#[from] sqlx::Error), +} + +/// 时间桶聚合用量记录。 +#[derive(Debug, Clone)] +pub struct UsageTimeBucketRecord { + pub bucket_start: DateTime, + pub model: String, + pub service_tier: Option, + pub request_count: i64, + pub error_count: i64, + pub input_tokens: i64, + pub output_tokens: i64, + pub cached_tokens: i64, + pub first_token_latency_sum: i64, + pub first_token_latency_count: i64, + pub latency_sum: i64, + pub latency_count: i64, + pub max_latency_ms: i64, + pub min_latency_ms: i64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageBucketWindow { + pub account_id: String, + pub key: String, + pub start: DateTime, + pub end: DateTime, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct UsageBucketTotals { + pub request_count: i64, + pub input_tokens: i64, + pub output_tokens: i64, + pub cached_tokens: i64, +} + +impl UsageBucketTotals { + fn add(&mut self, other: Self) { + self.request_count = self.request_count.saturating_add(other.request_count); + self.input_tokens = self.input_tokens.saturating_add(other.input_tokens); + self.output_tokens = self.output_tokens.saturating_add(other.output_tokens); + self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens); + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ModelUsageWindow { + pub account_id: String, + pub start: DateTime, + pub end: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ModelBucketUsage { + pub account_id: String, + pub model: String, + pub service_tier: Option, + pub request_count: i64, + pub error_count: i64, + pub input_tokens: i64, + pub output_tokens: i64, + pub cached_tokens: i64, + pub last_used_at: Option>, +} + +/// 请求时间桶聚合查询。 +#[derive(Clone)] +pub struct PgRequestBucketQuery { + pool: PgPool, +} + +impl PgRequestBucketQuery { + /// 构造请求时间桶查询。 + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + /// 列出指定时间范围内的时间桶聚合用量。 + pub async fn list( + &self, + start: DateTime, + end: DateTime, + ) -> Result, PgRequestBucketQueryError> { + let rows = sqlx::query(LIST_USAGE_TIME_BUCKETS_SQL) + .bind(start) + .bind(end) + .fetch_all(&self.pool) + .await?; + Ok(rows.iter().map(usage_time_bucket_from_row).collect()) + } + + pub async fn usage_by_windows( + &self, + windows: &[UsageBucketWindow], + ) -> Result>, PgRequestBucketQueryError> + { + let Some(min_start) = windows.iter().map(|window| window.start).min() else { + return Ok(HashMap::new()); + }; + let Some(max_end) = windows.iter().map(|window| window.end).max() else { + return Ok(HashMap::new()); + }; + let account_ids = windows + .iter() + .map(|window| window.account_id.as_str()) + .collect::>(); + let mut windows_by_account = HashMap::<&str, Vec<&UsageBucketWindow>>::new(); + for window in windows { + windows_by_account + .entry(window.account_id.as_str()) + .or_default() + .push(window); + } + + let mut builder = QueryBuilder::::new( + "select + account_id, + bucket_start, + coalesce(sum(success_count + error_count), 0)::bigint as request_count, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens + from request_time_buckets + where account_id in (", + ); + let mut separated = builder.separated(", "); + for account_id in &account_ids { + separated.push_bind(*account_id); + } + separated.push_unseparated(")"); + builder.push(" and bucket_start >= "); + builder.push_bind(china_quarter_hour_start(min_start)); + builder.push(" and bucket_start <= "); + builder.push_bind(max_end); + builder.push(" group by account_id, bucket_start"); + + let rows = builder.build().fetch_all(&self.pool).await?; + let mut usage_by_account = HashMap::>::new(); + for row in rows { + let account_id: String = row.get("account_id"); + let bucket_start = row.get::, _>("bucket_start"); + let bucket_usage = UsageBucketTotals { + request_count: row.get("request_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + }; + let Some(account_windows) = windows_by_account.get(account_id.as_str()) else { + continue; + }; + for window in account_windows.iter().filter(|window| { + bucket_start >= china_quarter_hour_start(window.start) && bucket_start <= window.end + }) { + usage_by_account + .entry(account_id.clone()) + .or_default() + .entry(window.key.clone()) + .or_default() + .add(bucket_usage); + } + } + Ok(usage_by_account) + } + + pub async fn model_usage_by_windows( + &self, + windows: &[ModelUsageWindow], + ) -> Result, PgRequestBucketQueryError> { + if windows.is_empty() { + return Ok(Vec::new()); + } + let mut builder = QueryBuilder::::new( + "select + account_id, + model, + service_tier, + coalesce(sum(success_count + error_count), 0)::bigint as request_count, + coalesce(sum(error_count), 0)::bigint as error_count, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + max(bucket_start) as last_used_at + from request_time_buckets + where model != '__unknown__' and (", + ); + for (index, window) in windows.iter().enumerate() { + if index > 0 { + builder.push(" or "); + } + builder.push("(account_id = "); + builder.push_bind(&window.account_id); + builder.push(" and bucket_start >= "); + builder.push_bind(china_quarter_hour_start(window.start)); + builder.push(" and bucket_start <= "); + builder.push_bind(window.end); + builder.push(")"); + } + builder.push(") group by account_id, model, service_tier"); + + let rows = builder.build().fetch_all(&self.pool).await?; + Ok(rows + .iter() + .map(|row| { + let service_tier = row.get::("service_tier"); + ModelBucketUsage { + account_id: row.get("account_id"), + model: row.get("model"), + service_tier: (service_tier != "__unknown__").then_some(service_tier), + request_count: row.get("request_count"), + error_count: row.get("error_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + last_used_at: row.get("last_used_at"), + } + }) + .collect()) + } +} + +fn usage_time_bucket_from_row(row: &sqlx::postgres::PgRow) -> UsageTimeBucketRecord { + UsageTimeBucketRecord { + bucket_start: row.get("bucket_start"), + model: row.get("model"), + service_tier: row.get("service_tier"), + request_count: row.get("request_count"), + error_count: row.get("error_count"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + first_token_latency_sum: row.get("first_token_latency_sum"), + first_token_latency_count: row.get("first_token_latency_count"), + latency_sum: row.get("latency_sum"), + latency_count: row.get("latency_count"), + max_latency_ms: row.get("max_latency_ms"), + min_latency_ms: row.get("min_latency_ms"), + } +} diff --git a/backend/src/telemetry/buckets/store.rs b/backend/src/telemetry/buckets/store.rs new file mode 100644 index 000000000..26dafa7fd --- /dev/null +++ b/backend/src/telemetry/buckets/store.rs @@ -0,0 +1,143 @@ +//! PostgreSQL 请求时间桶存储。 + +use chrono::{DateTime, Duration, Utc}; +use sqlx::{PgPool, Postgres}; +use thiserror::Error; + +use crate::{ + infra::time::china_quarter_hour_start, + telemetry::{ops::types::OpsErrorLog, usage::types::UsageRecord}, +}; + +const UNKNOWN_DIMENSION: &str = "__unknown__"; + +/// PostgreSQL 请求时间桶存储错误。 +#[derive(Debug, Error)] +pub enum PgRequestBucketStoreError { + /// 数据库操作失败。 + #[error("PostgreSQL request bucket operation failed: {0}")] + Database(#[from] sqlx::Error), + /// 保留期配置非法。 + #[error("invalid request bucket retention days: {0}")] + InvalidRetention(i64), +} + +/// PostgreSQL 请求时间桶存储。 +#[derive(Clone)] +pub struct PgRequestBucketStore { + pool: PgPool, +} + +impl PgRequestBucketStore { + /// 构造请求时间桶存储。 + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + /// 在成功事实事务内累计请求时间桶。 + pub async fn upsert_success( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + event: &UsageRecord, + ) -> Result<(), sqlx::Error> { + let latency = event.latency_ms; + let first_token = event.first_token_ms; + sqlx::query( + r" +insert into request_time_buckets ( + bucket_start, provider, account_id, model, service_tier, success_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, $2, $3, $4, $5, 1, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15) +on conflict (bucket_start, provider, account_id, model, service_tier) do update set + success_count = request_time_buckets.success_count + 1, + input_tokens = request_time_buckets.input_tokens + excluded.input_tokens, + output_tokens = request_time_buckets.output_tokens + excluded.output_tokens, + cached_tokens = request_time_buckets.cached_tokens + excluded.cached_tokens, + first_token_latency_sum = request_time_buckets.first_token_latency_sum + excluded.first_token_latency_sum, + first_token_latency_count = request_time_buckets.first_token_latency_count + excluded.first_token_latency_count, + latency_sum = request_time_buckets.latency_sum + excluded.latency_sum, + latency_count = request_time_buckets.latency_count + excluded.latency_count, + max_latency_ms = greatest(request_time_buckets.max_latency_ms, excluded.max_latency_ms), + min_latency_ms = case + when request_time_buckets.min_latency_ms is null then excluded.min_latency_ms + when excluded.min_latency_ms is null then request_time_buckets.min_latency_ms + else least(request_time_buckets.min_latency_ms, excluded.min_latency_ms) + end, + updated_at = excluded.updated_at", + ) + .bind(china_quarter_hour_start(event.created_at)) + .bind(dimension(Some(&event.provider))) + .bind(dimension(Some(&event.account_id))) + .bind(dimension(Some(&event.model))) + .bind(dimension(event.service_tier.as_deref())) + .bind(event.input_tokens.unwrap_or(0)) + .bind(event.output_tokens.unwrap_or(0)) + .bind(event.cached_tokens.unwrap_or(0)) + .bind(first_token.unwrap_or(0)) + .bind(i64::from(first_token.is_some())) + .bind(latency.unwrap_or(0)) + .bind(i64::from(latency.is_some())) + .bind(latency.unwrap_or(0)) + .bind(latency) + .bind(Utc::now()) + .execute(&mut **tx) + .await?; + Ok(()) + } + + /// 在错误事实事务内累计请求时间桶。 + pub async fn upsert_error( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + event: &OpsErrorLog, + ) -> Result<(), sqlx::Error> { + sqlx::query( + r" +insert into request_time_buckets ( + bucket_start, provider, account_id, model, service_tier, error_count, updated_at +) values ($1, $2, $3, $4, $5, 1, $6) +on conflict (bucket_start, provider, account_id, model, service_tier) do update set + error_count = request_time_buckets.error_count + 1, + updated_at = excluded.updated_at", + ) + .bind(china_quarter_hour_start(event.created_at)) + .bind(dimension(event.provider.as_deref())) + .bind(dimension(event.account_id.as_deref())) + .bind(dimension(event.model.as_deref())) + .bind(dimension(event.service_tier.as_deref())) + .bind(Utc::now()) + .execute(&mut **tx) + .await?; + Ok(()) + } + + /// 按运行时配置清理过期请求时间桶。 + pub async fn trim_to_retention( + &self, + now: DateTime, + ) -> Result { + let days: i64 = + sqlx::query_scalar("select bucket_retention_days from runtime_settings where id = 1") + .fetch_one(&self.pool) + .await?; + let duration = + Duration::try_days(days).ok_or(PgRequestBucketStoreError::InvalidRetention(days))?; + let cutoff = now + .checked_sub_signed(duration) + .ok_or(PgRequestBucketStoreError::InvalidRetention(days))?; + let result = sqlx::query("delete from request_time_buckets where bucket_start < $1") + .bind(cutoff) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } +} + +fn dimension(value: Option<&str>) -> &str { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or(UNKNOWN_DIMENSION) +} diff --git a/backend/src/admin/monitoring/diagnostics.rs b/backend/src/telemetry/diagnostics.rs similarity index 96% rename from backend/src/admin/monitoring/diagnostics.rs rename to backend/src/telemetry/diagnostics.rs index e7356fb10..f3b6a413f 100644 --- a/backend/src/admin/monitoring/diagnostics.rs +++ b/backend/src/telemetry/diagnostics.rs @@ -2,7 +2,7 @@ use serde::Serialize; -use crate::{upstream::accounts::pool::AccountCapacitySummary, upstream::fingerprint::Fingerprint}; +use crate::{accounts::pool::AccountCapacitySummary, upstream::openai::fingerprint::Fingerprint}; /// 账号池摘要。 #[derive(Debug, Default, Serialize)] diff --git a/backend/src/telemetry/mod.rs b/backend/src/telemetry/mod.rs new file mode 100644 index 000000000..2d25be1de --- /dev/null +++ b/backend/src/telemetry/mod.rs @@ -0,0 +1,10 @@ +//! 遥测事实、聚合与查询。 + +pub mod account_usage; +pub(crate) mod billing; +pub mod buckets; +pub mod diagnostics; +pub mod ops; +pub mod rebuild; +pub mod recorder; +pub mod usage; diff --git a/backend/src/telemetry/ops/mod.rs b/backend/src/telemetry/ops/mod.rs new file mode 100644 index 000000000..69e3c5e7e --- /dev/null +++ b/backend/src/telemetry/ops/mod.rs @@ -0,0 +1,5 @@ +//! 运维错误事实。 + +pub mod query; +pub mod store; +pub mod types; diff --git a/backend/src/telemetry/ops/query.rs b/backend/src/telemetry/ops/query.rs new file mode 100644 index 000000000..2581dc695 --- /dev/null +++ b/backend/src/telemetry/ops/query.rs @@ -0,0 +1,169 @@ +//! 管理端运维错误事件服务。 + +use chrono::Utc; +use thiserror::Error; + +use crate::telemetry::{ + ops::store::{OpsErrorFilter, PgOpsErrorLogStore}, + ops::types::OpsErrorLog, + usage::types::{metadata_i64, metadata_service_tier, metadata_string}, +}; + +/// 管理端运维错误事件错误。 +#[derive(Debug, Error)] +pub enum OpsQueryError { + /// 写入失败。 + #[error("failed to append ops error log")] + Append, + /// 查询失败。 + #[error("failed to list ops error logs")] + List, + /// 保留期清理失败。 + #[error("failed to trim expired ops error logs")] + Retention, +} + +/// 管理端运维错误事件服务。 +#[derive(Clone)] +pub struct OpsQueryService { + store: PgOpsErrorLogStore, + capture_body: bool, +} + +impl OpsQueryService { + /// 构造管理端运维错误事件服务。 + pub fn new(store: PgOpsErrorLogStore, capture_body: bool) -> Self { + Self { + store, + capture_body, + } + } + + /// 记录运维错误事件。 + pub async fn record(&self, mut event: OpsErrorLog) -> Result<(), OpsQueryError> { + lift_error_fact_fields(&mut event); + apply_capture_body_policy(&mut event, self.capture_body); + self.store + .append(&event) + .await + .map_err(|_| OpsQueryError::Append)?; + Ok(()) + } + + pub async fn list( + &self, + cursor: Option, + limit: u32, + filter: OpsErrorFilter, + ) -> Result, OpsQueryError> { + self.store + .list(filter, cursor, limit) + .await + .map_err(|_| OpsQueryError::List) + } + + pub async fn list_page( + &self, + page: u32, + page_size: u32, + filter: OpsErrorFilter, + ) -> Result, OpsQueryError> { + self.store + .list_page(filter, page, page_size) + .await + .map_err(|_| OpsQueryError::List) + } + + /// 周期清理超过保留期的错误事实。 + pub async fn trim_to_retention( + &self, + now: chrono::DateTime, + ) -> Result { + self.store + .trim_to_retention(now) + .await + .map_err(|_| OpsQueryError::Retention) + } +} + +fn lift_error_fact_fields(event: &mut OpsErrorLog) { + event.client_status_code = event + .client_status_code + .or_else(|| metadata_i64(&event.metadata, &["clientStatusCode", "client_status_code"])); + event.upstream_status_code = event + .upstream_status_code + .or_else(|| metadata_i64(&event.metadata, &["upstreamStatusCode", "upstreamStatus"])); + event.transport = event + .transport + .take() + .or_else(|| metadata_string(&event.metadata, &["transport"])); + event.attempt_index = event + .attempt_index + .or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"])); + event.failure_class = event + .failure_class + .take() + .or_else(|| metadata_string(&event.metadata, &["failureClass", "failure_class"])); + event.response_id = event + .response_id + .take() + .or_else(|| metadata_string(&event.metadata, &["responseId", "response_id"])); + event.upstream_request_id = event.upstream_request_id.take().or_else(|| { + metadata_string( + &event.metadata, + &[ + "upstreamRequestId", + "upstream_request_id", + "openaiRequestId", + ], + ) + }); + event.service_tier = event + .service_tier + .take() + .or_else(|| metadata_service_tier(&event.metadata).map(ToString::to_string)); + + let Some(metadata) = event.metadata.as_object_mut() else { + return; + }; + for key in [ + "clientStatusCode", + "client_status_code", + "upstreamStatusCode", + "upstreamStatus", + "transport", + "attemptIndex", + "attempt_index", + "failureClass", + "failure_class", + "responseId", + "response_id", + "upstreamRequestId", + "upstream_request_id", + "openaiRequestId", + "serviceTier", + ] { + metadata.remove(key); + } +} + +fn apply_capture_body_policy(event: &mut OpsErrorLog, capture_body: bool) { + let Some(metadata) = event.metadata.as_object_mut() else { + return; + }; + for key in [ + "body", + "rawBody", + "requestBody", + "responseBody", + "upstreamBody", + ] { + if capture_body { + if let Some(serde_json::Value::String(value)) = metadata.get_mut(key) { + value.truncate(4096); + } + } else { + metadata.remove(key); + } + } +} diff --git a/backend/src/telemetry/ops/store.rs b/backend/src/telemetry/ops/store.rs new file mode 100644 index 000000000..63e18cbdf --- /dev/null +++ b/backend/src/telemetry/ops/store.rs @@ -0,0 +1,322 @@ +//! PostgreSQL 运维错误事实存储。 + +use chrono::{DateTime, Duration, Utc}; +use sqlx::{PgPool, Postgres, QueryBuilder, Row}; +use thiserror::Error; + +use crate::{ + infra::{ + json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page}, + time::parse_rfc3339_utc, + }, + telemetry::{buckets::store::PgRequestBucketStore, ops::types::OpsErrorLog}, +}; + +const OPS_ERROR_SELECT_SQL: &str = r"select + 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 +from ops_error_logs"; + +#[derive(Debug, Error)] +pub enum PgOpsErrorLogStoreError { + #[error("PostgreSQL ops error operation failed: {0}")] + Database(#[from] sqlx::Error), + #[error("invalid ops error retention days: {0}")] + InvalidRetention(i64), + #[error("invalid ops error pagination cursor")] + InvalidCursor, +} + +pub type PgOpsErrorLogStoreResult = Result; + +#[derive(Clone)] +pub struct PgOpsErrorLogStore { + pool: PgPool, + buckets: PgRequestBucketStore, +} + +impl PgOpsErrorLogStore { + pub fn new(pool: PgPool) -> Self { + Self { + buckets: PgRequestBucketStore::new(pool.clone()), + pool, + } + } + + pub async fn append(&self, event: &OpsErrorLog) -> PgOpsErrorLogStoreResult<()> { + let mut tx = self.pool.begin().await?; + sqlx::query( + r" +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, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, + $15, $16, $17, $18, $19, $20 +)", + ) + .bind(&event.id) + .bind(&event.request_id) + .bind(&event.client_api_key_id) + .bind(&event.kind) + .bind(&event.provider) + .bind(&event.account_id) + .bind(&event.route) + .bind(&event.model) + .bind(event.status_code) + .bind(event.client_status_code) + .bind(event.upstream_status_code) + .bind(&event.transport) + .bind(event.attempt_index) + .bind(&event.failure_class) + .bind(&event.response_id) + .bind(&event.upstream_request_id) + .bind(event.latency_ms) + .bind(&event.message) + .bind(sqlx::types::Json(&event.metadata)) + .bind(event.created_at) + .execute(&mut *tx) + .await?; + self.buckets.upsert_error(&mut tx, event).await?; + tx.commit().await?; + Ok(()) + } + + pub async fn list( + &self, + filter: OpsErrorFilter, + cursor: Option, + limit: u32, + ) -> PgOpsErrorLogStoreResult> { + let limit = limit.clamp(1, 200); + let mut builder = QueryBuilder::::new(OPS_ERROR_SELECT_SQL); + push_filter(&mut builder, &filter, cursor.as_deref())?; + builder.push(" order by created_at desc, id desc limit "); + builder.push_bind(i64::from(limit) + 1); + let rows = builder.build().fetch_all(&self.pool).await?; + let has_next = rows.len() > limit as usize; + let items = rows + .into_iter() + .take(limit as usize) + .map(|row| ops_error_from_row(&row)) + .collect::>(); + let next_cursor = has_next.then(|| { + let item = items.last().expect("next page requires an ops error"); + encode_cursor(&item.created_at.to_rfc3339(), &item.id) + }); + Ok(Page { items, next_cursor }) + } + + pub async fn list_page( + &self, + filter: OpsErrorFilter, + page: u32, + page_size: u32, + ) -> PgOpsErrorLogStoreResult> { + let page_size = page_size.clamp(1, 200); + let total = count_ops_errors(&self.pool, &filter).await?; + let mut builder = QueryBuilder::::new(OPS_ERROR_SELECT_SQL); + push_filter(&mut builder, &filter, None)?; + builder.push(" order by created_at desc, id desc 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); + let items = builder + .build() + .fetch_all(&self.pool) + .await? + .iter() + .map(ops_error_from_row) + .collect(); + Ok(NumberedPage { + items, + total, + page: page.max(1), + page_size, + }) + } + + pub async fn trim_to_retention(&self, now: DateTime) -> PgOpsErrorLogStoreResult { + let days: i64 = sqlx::query_scalar( + "select ops_error_retention_days from runtime_settings where id = 1", + ) + .fetch_one(&self.pool) + .await?; + let duration = + Duration::try_days(days).ok_or(PgOpsErrorLogStoreError::InvalidRetention(days))?; + let cutoff = now + .checked_sub_signed(duration) + .ok_or(PgOpsErrorLogStoreError::InvalidRetention(days))?; + let result = sqlx::query("delete from ops_error_logs where created_at < $1") + .bind(cutoff) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } + + pub async fn clear(&self) -> PgOpsErrorLogStoreResult { + let result = sqlx::query("delete from ops_error_logs") + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } + + pub async fn count(&self) -> PgOpsErrorLogStoreResult { + let total: i64 = sqlx::query_scalar("select count(*) from ops_error_logs") + .fetch_one(&self.pool) + .await?; + Ok(total.max(0) as u64) + } +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct OpsErrorFilter { + pub kind: Option, + pub client_api_key_id: Option, + pub provider: Option, + pub request_id: Option, + pub account_id: Option, + pub route: Option, + pub model: Option, + pub status_code: Option, + pub client_status_code: Option, + pub upstream_status_code: Option, + pub transport: Option, + pub attempt_index: Option, + pub failure_class: Option, + pub response_id: Option, + pub upstream_request_id: Option, + pub search: Option, + pub start_time: Option>, + pub end_time: Option>, +} + +fn push_filter( + builder: &mut QueryBuilder, + filter: &OpsErrorFilter, + cursor: Option<&str>, +) -> PgOpsErrorLogStoreResult<()> { + builder.push(" where true"); + push_optional_text(builder, "kind", filter.kind.as_deref()); + push_optional_text( + builder, + "client_api_key_id", + filter.client_api_key_id.as_deref(), + ); + push_optional_text(builder, "provider", filter.provider.as_deref()); + push_optional_text(builder, "request_id", filter.request_id.as_deref()); + push_optional_text(builder, "account_id", filter.account_id.as_deref()); + push_optional_text(builder, "route", filter.route.as_deref()); + push_optional_text(builder, "model", filter.model.as_deref()); + push_optional_i64(builder, "status_code", filter.status_code); + push_optional_i64(builder, "client_status_code", filter.client_status_code); + push_optional_i64(builder, "upstream_status_code", filter.upstream_status_code); + push_optional_text(builder, "transport", filter.transport.as_deref()); + push_optional_i64(builder, "attempt_index", filter.attempt_index); + push_optional_text(builder, "failure_class", filter.failure_class.as_deref()); + push_optional_text(builder, "response_id", filter.response_id.as_deref()); + push_optional_text( + builder, + "upstream_request_id", + filter.upstream_request_id.as_deref(), + ); + if let Some(search) = filter + .search + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + builder.push(" and (message ilike "); + builder.push_bind(format!("%{search}%")); + builder.push(" or request_id = "); + builder.push_bind(search); + builder.push(" or response_id = "); + builder.push_bind(search); + builder.push(" or upstream_request_id = "); + builder.push_bind(search); + builder.push(")"); + } + if let Some(start_time) = filter.start_time { + builder.push(" and created_at >= "); + builder.push_bind(start_time); + } + if let Some(end_time) = filter.end_time { + builder.push(" and created_at <= "); + builder.push_bind(end_time); + } + if let Some(cursor) = cursor { + let (created_at, id) = + decode_cursor(cursor).ok_or(PgOpsErrorLogStoreError::InvalidCursor)?; + let created_at = + parse_rfc3339_utc(&created_at).map_err(|_| PgOpsErrorLogStoreError::InvalidCursor)?; + builder.push(" and (created_at < "); + builder.push_bind(created_at); + builder.push(" or (created_at = "); + builder.push_bind(created_at); + builder.push(" and id < "); + builder.push_bind(id); + builder.push("))"); + } + Ok(()) +} + +fn push_optional_text(builder: &mut QueryBuilder, column: &str, value: Option<&str>) { + if let Some(value) = value { + builder.push(" and "); + builder.push(column); + builder.push(" = "); + builder.push_bind(value); + } +} + +fn push_optional_i64(builder: &mut QueryBuilder, column: &str, value: Option) { + if let Some(value) = value { + builder.push(" and "); + builder.push(column); + builder.push(" = "); + builder.push_bind(value); + } +} + +async fn count_ops_errors(pool: &PgPool, filter: &OpsErrorFilter) -> PgOpsErrorLogStoreResult { + let mut builder = QueryBuilder::::new("select count(*) from ops_error_logs"); + push_filter(&mut builder, filter, None)?; + let total: i64 = builder.build_query_scalar().fetch_one(pool).await?; + Ok(total.max(0) as u64) +} + +fn ops_error_from_row(row: &sqlx::postgres::PgRow) -> OpsErrorLog { + OpsErrorLog { + id: row.get("id"), + request_id: row.get("request_id"), + client_api_key_id: row.get("client_api_key_id"), + kind: row.get("kind"), + provider: row.get("provider"), + account_id: row.get("account_id"), + route: row.get("route"), + model: row.get("model"), + status_code: row.get::, _>("status_code").map(i64::from), + client_status_code: row + .get::, _>("client_status_code") + .map(i64::from), + upstream_status_code: row + .get::, _>("upstream_status_code") + .map(i64::from), + transport: row.get("transport"), + attempt_index: row.get("attempt_index"), + failure_class: row.get("failure_class"), + response_id: row.get("response_id"), + upstream_request_id: row.get("upstream_request_id"), + service_tier: None, + latency_ms: row.get("latency_ms"), + message: row.get("message"), + metadata: row + .get::, _>("metadata_json") + .0, + created_at: row.get("created_at"), + } +} diff --git a/backend/src/admin/monitoring/ops_error_model.rs b/backend/src/telemetry/ops/types.rs similarity index 84% rename from backend/src/admin/monitoring/ops_error_model.rs rename to backend/src/telemetry/ops/types.rs index 06eaf9dda..4186a487e 100644 --- a/backend/src/admin/monitoring/ops_error_model.rs +++ b/backend/src/telemetry/ops/types.rs @@ -13,8 +13,12 @@ pub struct OpsErrorLog { pub id: String, /// 请求 ID。 pub request_id: Option, + /// 调用方客户端 API key 的稳定 ID。 + pub client_api_key_id: Option, /// 事件类别。 pub kind: String, + /// 已知的上游 provider；调度前失败可为空。 + pub provider: Option, /// 账号 ID。预账号调度失败时为空。 pub account_id: Option, /// HTTP 路由。 @@ -37,6 +41,9 @@ pub struct OpsErrorLog { pub response_id: Option, /// 上游请求 ID。 pub upstream_request_id: Option, + /// 仅用于错误时间桶维度，不落 ops_error_logs 明细列。 + #[serde(skip)] + pub service_tier: Option, /// 请求耗时毫秒。 pub latency_ms: Option, /// 展示消息。 @@ -54,7 +61,9 @@ impl OpsErrorLog { Self { id: format!("ops_error_{}", Uuid::new_v4().simple()), request_id: None, + client_api_key_id: None, kind: kind.into(), + provider: None, account_id: None, route: None, model: None, @@ -66,6 +75,7 @@ impl OpsErrorLog { failure_class: None, response_id: None, upstream_request_id: None, + service_tier: None, latency_ms: None, message: message.into(), metadata: Value::Object(Map::default()), diff --git a/backend/src/telemetry/rebuild.rs b/backend/src/telemetry/rebuild.rs new file mode 100644 index 000000000..f5d7bf9fb --- /dev/null +++ b/backend/src/telemetry/rebuild.rs @@ -0,0 +1,138 @@ +//! `request_time_buckets` 可重建范围重算。 + +use chrono::{DateTime, Utc}; +use sqlx::PgPool; +use thiserror::Error; + +/// 桶重算结果。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RebuildBucketsReport { + pub cutoff: DateTime, + pub deleted_rows: u64, + pub rebuilt_rows: u64, +} + +/// 桶重算错误。 +#[derive(Debug, Error)] +pub enum RebuildBucketsError { + #[error("failed to rebuild request time buckets: {0}")] + Database(#[from] sqlx::Error), +} + +/// 删除两类事实都仍在保留期内的桶，并从成功/失败事实重新聚合。 +pub async fn rebuild_buckets(pool: &PgPool) -> Result { + let mut transaction = pool.begin().await?; + sqlx::query( + "lock table usage_records, ops_error_logs, request_time_buckets + in share row exclusive mode", + ) + .execute(&mut *transaction) + .await?; + + let retention_days: i64 = sqlx::query_scalar( + "select least(usage_retention_days, ops_error_retention_days) + from runtime_settings where id = 1", + ) + .fetch_one(&mut *transaction) + .await?; + let cutoff: DateTime = sqlx::query_scalar( + "select to_timestamp( + floor(extract(epoch from now() - ($1::bigint * interval '1 day')) / 900) * 900 + )", + ) + .bind(retention_days) + .fetch_one(&mut *transaction) + .await?; + + let deleted_rows = sqlx::query("delete from request_time_buckets where bucket_start >= $1") + .bind(cutoff) + .execute(&mut *transaction) + .await? + .rows_affected(); + + let rebuilt_rows = sqlx::query( + r" +with facts as ( + select + to_timestamp(floor(extract(epoch from created_at) / 900) * 900) as bucket_start, + coalesce(nullif(trim(provider), ''), '__unknown__') as provider, + coalesce(nullif(trim(account_id), ''), '__unknown__') as account_id, + coalesce(nullif(trim(model), ''), '__unknown__') as model, + coalesce(nullif(trim(service_tier), ''), '__unknown__') as service_tier, + 1::bigint as success_count, + 0::bigint as error_count, + coalesce(input_tokens, 0)::bigint as input_tokens, + coalesce(output_tokens, 0)::bigint as output_tokens, + coalesce(cached_tokens, 0)::bigint as cached_tokens, + coalesce(first_token_ms, 0)::bigint as first_token_latency_sum, + (first_token_ms is not null)::integer::bigint as first_token_latency_count, + coalesce(latency_ms, 0)::bigint as latency_sum, + (latency_ms is not null)::integer::bigint as latency_count, + coalesce(latency_ms, 0)::bigint as max_latency_ms, + latency_ms::bigint as min_latency_ms + from usage_records + where created_at >= $1 + + union all + + select + to_timestamp(floor(extract(epoch from created_at) / 900) * 900), + coalesce(nullif(trim(provider), ''), '__unknown__'), + coalesce(nullif(trim(account_id), ''), '__unknown__'), + coalesce(nullif(trim(model), ''), '__unknown__'), + '__unknown__', + 0::bigint, + 1::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + 0::bigint, + null::bigint + from ops_error_logs + where created_at >= $1 +), aggregated as ( + select + bucket_start, provider, account_id, model, service_tier, + sum(success_count) as success_count, + sum(error_count) as error_count, + sum(input_tokens) as input_tokens, + sum(output_tokens) as output_tokens, + sum(cached_tokens) as cached_tokens, + sum(first_token_latency_sum) as first_token_latency_sum, + sum(first_token_latency_count) as first_token_latency_count, + sum(latency_sum) as latency_sum, + sum(latency_count) as latency_count, + max(max_latency_ms) as max_latency_ms, + min(min_latency_ms) as min_latency_ms + from facts + group by bucket_start, provider, account_id, model, service_tier +) +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 +) +select + 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, now() +from aggregated", + ) + .bind(cutoff) + .execute(&mut *transaction) + .await? + .rows_affected(); + + transaction.commit().await?; + Ok(RebuildBucketsReport { + cutoff, + deleted_rows, + rebuilt_rows, + }) +} diff --git a/backend/src/proxy/dispatch/usage_events.rs b/backend/src/telemetry/recorder.rs similarity index 53% rename from backend/src/proxy/dispatch/usage_events.rs rename to backend/src/telemetry/recorder.rs index f1ec60835..a2f2e2ae3 100644 --- a/backend/src/proxy/dispatch/usage_events.rs +++ b/backend/src/telemetry/recorder.rs @@ -5,18 +5,20 @@ use std::time::Instant; use serde_json::Value; use crate::{ - admin::monitoring::{ - ops_error_model::OpsErrorLog, - ops_error_service::AdminOpsErrorLogService, - usage_record_model::{ResponseUsageRecord, UsageRecord}, - }, infra::time::elapsed_millis_i64, - upstream::protocol::responses::{CodexCompactRequest, CodexResponsesRequest}, + telemetry::{ + ops::query::OpsQueryService, + ops::types::OpsErrorLog, + usage::types::{ResponseUsageRecord, UsageRecord}, + }, + upstream::openai::protocol::responses::{CodexCompactRequest, CodexResponsesRequest}, }; -pub(super) struct DispatchErrorLogRecord<'a> { - pub ops_errors: &'a AdminOpsErrorLogService, +pub(crate) struct DispatchErrorLogRecord<'a> { + pub ops_errors: &'a OpsQueryService, pub request_id: &'a str, + pub client_api_key_id: Option<&'a str>, + pub provider: Option<&'a str>, pub account_id: Option<&'a str>, pub route: &'a str, pub model: &'a str, @@ -26,12 +28,14 @@ pub(super) struct DispatchErrorLogRecord<'a> { pub metadata: Value, } -pub(super) async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) { +pub(crate) async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) { let mut metadata = record.metadata; enrich_event_route_metadata(&mut metadata, record.route); let mut event = OpsErrorLog::new(event_kind(record.route), record.message); event.request_id = Some(record.request_id.to_string()); + event.client_api_key_id = record.client_api_key_id.map(ToString::to_string); + event.provider = record.provider.map(ToString::to_string); event.account_id = record.account_id.map(ToString::to_string); event.route = Some(record.route.to_string()); event.model = Some(record.model.to_string()); @@ -48,17 +52,24 @@ pub(super) async fn record_dispatch_error_event(record: DispatchErrorLogRecord<' } } -pub(super) async fn record_response_event(record: ResponseUsageRecord<'_>) { +pub(crate) async fn record_response_event(record: ResponseUsageRecord<'_>) { let mut metadata = record.metadata; enrich_event_route_metadata(&mut metadata, record.route); - let mut event = UsageRecord::new(event_kind(record.route), record.level, record.message); + let mut event = UsageRecord::new( + event_kind(record.route), + record.message, + record.account_id, + record.model, + record.status_code, + ); event.request_id = Some(record.request_id.to_string()); - event.account_id = Some(record.account_id.to_string()); + event.client_api_key_id = record.client_api_key_id.map(ToString::to_string); event.route = Some(record.route.to_string()); - event.model = Some(record.model.to_string()); - event.status_code = Some(record.status_code); event.latency_ms = Some(elapsed_millis_i64(record.started_at)); + event.requested_model = normalized_string(record.requested_model.or(Some(record.model))); + event.upstream_model = normalized_string(Some(record.model)); + event.service_tier = normalized_string(record.service_tier); enrich_usage_record_identity( &mut metadata, record.requested_model, @@ -68,10 +79,8 @@ pub(super) async fn record_response_event(record: ResponseUsageRecord<'_>) { record.reasoning_effort, record.service_tier, ); - event.metadata = metadata; - if !record.rate_limit_headers.is_empty() { - if let Some(object) = event.metadata.as_object_mut() { + if let Some(object) = metadata.as_object_mut() { object.insert( "rateLimitHeaders".to_string(), serde_json::json!(record.rate_limit_headers), @@ -79,22 +88,88 @@ pub(super) async fn record_response_event(record: ResponseUsageRecord<'_>) { } } + lift_success_fact_fields(&mut event, &mut metadata); + event.metadata = metadata; + if let Err(error) = record.usage_records.record(event).await { tracing::warn!(account_id = %record.account_id, error = %error, "failed to record response event"); } } -pub(super) fn event_kind(route: &str) -> &'static str { +fn lift_success_fact_fields(event: &mut UsageRecord, metadata: &mut Value) { + event.transport = metadata_string(metadata, "transport"); + event.attempt_index = metadata_nonnegative_i64(metadata, "attemptIndex"); + event.response_id = metadata_string(metadata, "responseId"); + event.upstream_request_id = metadata_string(metadata, "upstreamRequestId") + .or_else(|| metadata_string(metadata, "openaiRequestId")); + event.first_token_ms = metadata_nonnegative_i64(metadata, "firstTokenMs"); + if let Some(usage) = metadata.get("usage") { + event.input_tokens = value_nonnegative_i64(usage, "inputTokens"); + event.output_tokens = value_nonnegative_i64(usage, "outputTokens"); + event.cached_tokens = value_nonnegative_i64(usage, "cachedTokens"); + event.reasoning_tokens = value_nonnegative_i64(usage, "reasoningTokens"); + } + + let Some(object) = metadata.as_object_mut() else { + return; + }; + for key in [ + "usage", + "route", + "apiKind", + "requestedModel", + "upstreamModel", + "serviceTier", + "statusCode", + "transport", + "attemptIndex", + "responseId", + "upstreamRequestId", + "openaiRequestId", + "firstTokenMs", + ] { + object.remove(key); + } +} + +fn metadata_string(metadata: &Value, field: &str) -> Option { + metadata + .get(field) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) +} + +fn metadata_nonnegative_i64(metadata: &Value, field: &str) -> Option { + value_nonnegative_i64(metadata, field) +} + +fn value_nonnegative_i64(value: &Value, field: &str) -> Option { + value + .get(field) + .and_then(Value::as_i64) + .filter(|value| *value >= 0) +} + +fn normalized_string(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToString::to_string) +} + +pub(crate) fn event_kind(route: &str) -> &'static str { let _ = route; "v1.response" } -pub(super) fn api_kind(route: &str) -> &'static str { +pub(crate) fn api_kind(route: &str) -> &'static str { let _ = route; "responses" } -pub(super) fn enrich_event_route_metadata(metadata: &mut Value, route: &str) { +pub(crate) fn enrich_event_route_metadata(metadata: &mut Value, route: &str) { let Some(object) = metadata.as_object_mut() else { return; }; @@ -106,7 +181,7 @@ pub(super) fn enrich_event_route_metadata(metadata: &mut Value, route: &str) { .or_insert_with(|| Value::String(api_kind(route).to_string())); } -pub(super) fn enrich_usage_record_identity( +pub(crate) fn enrich_usage_record_identity( metadata: &mut Value, requested_model: Option<&str>, upstream_model: &str, @@ -139,11 +214,11 @@ pub(super) fn enrich_usage_record_identity( insert_trimmed_string(object, "serviceTier", service_tier); } -pub(super) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> { +pub(crate) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> { reasoning_effort_from_value(request.reasoning()) } -pub(super) fn reasoning_effort_from_compact_request(request: &CodexCompactRequest) -> Option<&str> { +pub(crate) fn reasoning_effort_from_compact_request(request: &CodexCompactRequest) -> Option<&str> { reasoning_effort_from_value(request.reasoning()) } diff --git a/backend/src/telemetry/usage/mod.rs b/backend/src/telemetry/usage/mod.rs new file mode 100644 index 000000000..b45762b9d --- /dev/null +++ b/backend/src/telemetry/usage/mod.rs @@ -0,0 +1,5 @@ +//! 成功请求事实。 + +pub mod query; +pub mod store; +pub mod types; diff --git a/backend/src/admin/monitoring/usage_record_service.rs b/backend/src/telemetry/usage/query.rs similarity index 66% rename from backend/src/admin/monitoring/usage_record_service.rs rename to backend/src/telemetry/usage/query.rs index cd13a3a86..102b6dc7d 100644 --- a/backend/src/admin/monitoring/usage_record_service.rs +++ b/backend/src/telemetry/usage/query.rs @@ -6,26 +6,27 @@ use chrono::{DateTime, Utc}; use thiserror::Error; use crate::{ - admin::monitoring::{ + infra::json::{NumberedPage, Page}, + telemetry::{ billing, - usage_record_model::metadata_service_tier, - usage_record_model::{UsageRecord, UsageRecordLevel}, - usage_record_store::{ - SqliteUsageRecordStore, UsageRecordAccountUsage, UsageRecordBreakdown, + usage::store::{ + PgUsageRecordStore, UsageRecordAccountUsage, UsageRecordBreakdown, UsageRecordEndpointSource, UsageRecordFilter, UsageRecordModelSource, UsageRecordSummary, UsageRecordTrendPoint, }, + usage::types::UsageRecord, }, - infra::json::{NumberedPage, Page}, }; /// 管理端使用记录查询过滤器。 #[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct AdminUsageRecordFilter { +pub struct UsageQueryFilter { /// 事件类别。 pub kind: Option, - /// 事件等级。 - pub level: Option, + /// 调用方客户端 API key ID。 + pub client_api_key_id: Option, + /// 上游 provider。 + pub provider: Option, /// 请求 ID。 pub request_id: Option, /// 账号 ID。 @@ -40,10 +41,6 @@ pub struct AdminUsageRecordFilter { pub transport: Option, /// 同一请求内的上游尝试序号。 pub attempt_index: Option, - /// 上游 HTTP 状态码。 - pub upstream_status_code: Option, - /// 失败分类。 - pub failure_class: Option, /// 上游响应 ID。 pub response_id: Option, /// 上游请求 ID。 @@ -56,11 +53,12 @@ pub struct AdminUsageRecordFilter { pub end_time: Option>, } -impl From for UsageRecordFilter { - fn from(filter: AdminUsageRecordFilter) -> Self { +impl From for UsageRecordFilter { + fn from(filter: UsageQueryFilter) -> Self { Self { kind: filter.kind, - level: filter.level, + client_api_key_id: filter.client_api_key_id, + provider: filter.provider, request_id: filter.request_id, account_id: filter.account_id, route: filter.route, @@ -68,8 +66,6 @@ impl From for UsageRecordFilter { status_code: filter.status_code, transport: filter.transport, attempt_index: filter.attempt_index, - upstream_status_code: filter.upstream_status_code, - failure_class: filter.failure_class, response_id: filter.response_id, upstream_request_id: filter.upstream_request_id, search: filter.search, @@ -81,14 +77,14 @@ impl From for UsageRecordFilter { /// 清空使用记录结果。 #[derive(Debug, Clone, PartialEq, Eq)] -pub struct AdminClearUsageRecords { +pub struct ClearUsageRecords { /// 清理数量。 pub cleared: u64, } /// 管理端使用记录费用明细。 #[derive(Debug, Clone, PartialEq)] -pub(crate) struct AdminUsageRecordCostDetails { +pub(crate) struct UsageRecordCostDetails { pub input_cost: f64, pub output_cost: f64, pub cache_read_cost: f64, @@ -104,7 +100,7 @@ pub(crate) struct AdminUsageRecordCostDetails { /// 管理端使用记录错误。 #[derive(Debug, Error)] -pub enum AdminUsageRecordError { +pub enum UsageQueryError { /// 列表失败。 #[error("failed to list usage records")] List, @@ -127,23 +123,23 @@ pub enum AdminUsageRecordError { /// 管理端使用记录服务。 #[derive(Clone)] -pub struct AdminUsageRecordService { - store: SqliteUsageRecordStore, - settings: Arc>, +pub struct UsageQueryService { + store: PgUsageRecordStore, + settings: Arc>, } #[derive(Debug, Clone, Copy)] -struct AdminUsageRecordSettings { +struct UsageQuerySettings { enabled: bool, capture_body: bool, } -impl AdminUsageRecordService { +impl UsageQueryService { /// 构造管理端使用记录服务。 - pub fn new(store: SqliteUsageRecordStore, enabled: bool, capture_body: bool) -> Self { + pub fn new(store: PgUsageRecordStore, enabled: bool, capture_body: bool) -> Self { Self { store, - settings: Arc::new(tokio::sync::RwLock::new(AdminUsageRecordSettings { + settings: Arc::new(tokio::sync::RwLock::new(UsageQuerySettings { enabled, capture_body, })), @@ -155,12 +151,12 @@ impl AdminUsageRecordService { &self, cursor: Option, limit: u32, - filter: AdminUsageRecordFilter, - ) -> Result, AdminUsageRecordError> { + filter: UsageQueryFilter, + ) -> Result, UsageQueryError> { self.store .list(filter.into(), cursor, limit) .await - .map_err(|_| AdminUsageRecordError::List) + .map_err(|_| UsageQueryError::List) } /// 按页码查询日志。 @@ -168,102 +164,99 @@ impl AdminUsageRecordService { &self, page: u32, page_size: u32, - filter: AdminUsageRecordFilter, - ) -> Result, AdminUsageRecordError> { + filter: UsageQueryFilter, + ) -> Result, UsageQueryError> { self.store .list_page(filter.into(), page, page_size) .await - .map_err(|_| AdminUsageRecordError::List) + .map_err(|_| UsageQueryError::List) } /// 按 ID 读取日志。 - pub async fn get(&self, id: &str) -> Result, AdminUsageRecordError> { - self.store - .get(id) - .await - .map_err(|_| AdminUsageRecordError::Get) + pub async fn get(&self, id: &str) -> Result, UsageQueryError> { + self.store.get(id).await.map_err(|_| UsageQueryError::Get) } /// 读取使用记录关联账号的邮箱映射。 pub async fn account_email_map( &self, items: &[UsageRecord], - ) -> Result, AdminUsageRecordError> { + ) -> Result, UsageQueryError> { self.store .account_email_map(items) .await - .map_err(|_| AdminUsageRecordError::Accounts) + .map_err(|_| UsageQueryError::Accounts) } /// 汇总使用记录。 pub async fn summary( &self, - filter: AdminUsageRecordFilter, - ) -> Result { + filter: UsageQueryFilter, + ) -> Result { self.store .summary(filter.into()) .await - .map_err(|_| AdminUsageRecordError::List) + .map_err(|_| UsageQueryError::List) } /// 按账号聚合使用记录。 pub async fn account_usage( &self, - filter: AdminUsageRecordFilter, + filter: UsageQueryFilter, limit: u32, - ) -> Result, AdminUsageRecordError> { + ) -> Result, UsageQueryError> { self.store .account_usage(filter.into(), limit) .await - .map_err(|_| AdminUsageRecordError::List) + .map_err(|_| UsageQueryError::List) } /// 按模型来源聚合使用记录分布。 pub async fn model_distribution( &self, - filter: AdminUsageRecordFilter, + filter: UsageQueryFilter, source: UsageRecordModelSource, - ) -> Result, AdminUsageRecordError> { + ) -> Result, UsageQueryError> { self.store .model_distribution(filter.into(), source) .await - .map_err(|_| AdminUsageRecordError::List) + .map_err(|_| UsageQueryError::List) } /// 按端点来源聚合使用记录分布。 pub async fn endpoint_distribution( &self, - filter: AdminUsageRecordFilter, + filter: UsageQueryFilter, source: UsageRecordEndpointSource, - ) -> Result, AdminUsageRecordError> { + ) -> Result, UsageQueryError> { self.store .endpoint_distribution(filter.into(), source) .await - .map_err(|_| AdminUsageRecordError::List) + .map_err(|_| UsageQueryError::List) } /// 聚合使用记录趋势。 pub async fn trend( &self, - filter: AdminUsageRecordFilter, - ) -> Result, AdminUsageRecordError> { + filter: UsageQueryFilter, + ) -> Result, UsageQueryError> { self.store .trend(filter.into()) .await - .map_err(|_| AdminUsageRecordError::List) + .map_err(|_| UsageQueryError::List) } /// 清空日志。 - pub async fn clear(&self) -> Result { + pub async fn clear(&self) -> Result { self.store .clear() .await - .map(|cleared| AdminClearUsageRecords { cleared }) - .map_err(|_| AdminUsageRecordError::Clear) + .map(|cleared| ClearUsageRecords { cleared }) + .map_err(|_| UsageQueryError::Clear) } /// 记录使用记录。 - pub async fn record(&self, event: UsageRecord) -> Result<(), AdminUsageRecordError> { + pub async fn record(&self, event: UsageRecord) -> Result<(), UsageQueryError> { let settings = *self.settings.read().await; if !settings.enabled { return Ok(()); @@ -272,13 +265,12 @@ impl AdminUsageRecordService { tracing::warn!( usage_record_id = %event.id, request_id = event.request_id.as_deref().unwrap_or(""), - account_id = event.account_id.as_deref().unwrap_or(""), - model = event.model.as_deref().unwrap_or(""), - level = ?event.level, - status_code = ?event.status_code, - "skipped non-usage event in usage record service" + account_id = event.account_id, + model = event.model, + status_code = event.status_code, + "rejected invalid success usage fact" ); - return Ok(()); + return Err(UsageQueryError::Append); } self.append_with_settings(event, settings).await } @@ -286,34 +278,30 @@ impl AdminUsageRecordService { async fn append_with_settings( &self, mut event: UsageRecord, - settings: AdminUsageRecordSettings, - ) -> Result<(), AdminUsageRecordError> { + settings: UsageQuerySettings, + ) -> Result<(), UsageQueryError> { apply_capture_body_policy(&mut event, settings.capture_body); self.store .append(&event) .await - .map_err(|_| AdminUsageRecordError::Append)?; - self.store - .trim_to_retention(Utc::now()) - .await - .map_err(|_| AdminUsageRecordError::Retention)?; + .map_err(|_| UsageQueryError::Append)?; Ok(()) } + + /// 周期清理超过保留期的成功事实。 + pub async fn trim_to_retention(&self, now: DateTime) -> Result { + self.store + .trim_to_retention(now) + .await + .map_err(|_| UsageQueryError::Retention) + } } fn is_usage_log_event(event: &UsageRecord) -> bool { - event.level != UsageRecordLevel::Error - && event.status_code.is_none_or(|status| status < 400) - && event - .account_id - .as_deref() - .map(str::trim) - .is_some_and(|account_id| !account_id.is_empty()) - && event - .model - .as_deref() - .map(str::trim) - .is_some_and(|model| !model.is_empty()) + (200..=399).contains(&event.status_code) + && !event.provider.trim().is_empty() + && !event.account_id.trim().is_empty() + && !event.model.trim().is_empty() } fn apply_capture_body_policy(event: &mut UsageRecord, capture_body: bool) { @@ -340,16 +328,17 @@ pub(crate) fn usage_record_cost_details( input_tokens: u64, output_tokens: u64, cached_tokens: u64, -) -> Option { +) -> Option { if input_tokens == 0 && output_tokens == 0 && cached_tokens == 0 { return None; } let model = upstream_model - .or(record.model.as_deref()) + .or(record.upstream_model.as_deref()) + .or(Some(record.model.as_str())) .map(str::trim) .filter(|value| !value.is_empty())?; - let service_tier = metadata_service_tier(&record.metadata); + let service_tier = record.service_tier.as_deref(); let breakdown = billing::calculate_cost_breakdown( input_tokens, output_tokens, @@ -359,7 +348,7 @@ pub(crate) fn usage_record_cost_details( ); let original_cost = breakdown.input_cost + breakdown.output_cost + breakdown.cache_read_cost; - Some(AdminUsageRecordCostDetails { + Some(UsageRecordCostDetails { input_cost: breakdown.input_cost, output_cost: breakdown.output_cost, cache_read_cost: breakdown.cache_read_cost, diff --git a/backend/src/telemetry/usage/store.rs b/backend/src/telemetry/usage/store.rs new file mode 100644 index 000000000..2cccdf32e --- /dev/null +++ b/backend/src/telemetry/usage/store.rs @@ -0,0 +1,369 @@ +//! PostgreSQL 成功使用事实存储。 + +mod query; + +use query::{ + count_usage_records, usage_account_usage, usage_breakdown, usage_record_account_email_map, + usage_record_from_row, usage_summary, usage_trend, +}; +pub use query::{ + UsageRecordAccountUsage, UsageRecordBreakdown, UsageRecordEndpointSource, + UsageRecordModelSource, UsageRecordSummary, UsageRecordTrendPoint, +}; + +use std::collections::{BTreeMap, HashMap}; + +use chrono::{DateTime, Duration, Utc}; +use sqlx::{PgPool, Postgres, QueryBuilder, Row}; +use thiserror::Error; + +use crate::{ + infra::{ + format::optional_nonnegative_i64_to_u64, + json::{decode_cursor, encode_cursor, page_offset, NumberedPage, Page}, + time::parse_rfc3339_utc, + }, + telemetry::{billing, buckets::store::PgRequestBucketStore, usage::types::UsageRecord}, +}; + +pub const USAGE_RECORD_RETENTION_DAYS: i64 = 30; +pub const DEFAULT_USAGE_RECORD_CAPTURE_BODY: bool = false; + +const USAGE_RECORD_SELECT_SQL: &str = r" +select + 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 +from usage_records"; + +#[derive(Debug, Error)] +pub enum PgUsageRecordStoreError { + #[error("PostgreSQL usage record operation failed: {0}")] + Database(#[from] sqlx::Error), + #[error("invalid usage record pagination cursor")] + InvalidCursor, + #[error("invalid usage retention days: {0}")] + InvalidRetention(i64), +} + +pub type PgUsageRecordStoreResult = Result; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct UsageRecordFilter { + pub kind: Option, + pub client_api_key_id: Option, + pub provider: Option, + pub request_id: Option, + pub account_id: Option, + pub route: Option, + pub model: Option, + pub status_code: Option, + pub transport: Option, + pub attempt_index: Option, + pub response_id: Option, + pub upstream_request_id: Option, + pub search: Option, + pub start_time: Option>, + pub end_time: Option>, +} + +#[derive(Clone)] +pub struct PgUsageRecordStore { + pool: PgPool, + buckets: PgRequestBucketStore, +} + +impl PgUsageRecordStore { + pub fn new(pool: PgPool) -> Self { + Self { + buckets: PgRequestBucketStore::new(pool.clone()), + pool, + } + } + + pub fn pool(&self) -> &PgPool { + &self.pool + } + + pub async fn append(&self, event: &UsageRecord) -> PgUsageRecordStoreResult<()> { + let mut tx = self.pool.begin().await?; + sqlx::query( + r" +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, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, + $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25 +)", + ) + .bind(&event.id) + .bind(&event.request_id) + .bind(&event.client_api_key_id) + .bind(&event.kind) + .bind(&event.route) + .bind(&event.provider) + .bind(&event.account_id) + .bind(&event.model) + .bind(&event.requested_model) + .bind(&event.upstream_model) + .bind(&event.service_tier) + .bind(event.status_code) + .bind(&event.transport) + .bind(event.attempt_index) + .bind(&event.response_id) + .bind(&event.upstream_request_id) + .bind(event.latency_ms) + .bind(event.first_token_ms) + .bind(event.input_tokens) + .bind(event.output_tokens) + .bind(event.cached_tokens) + .bind(event.reasoning_tokens) + .bind(&event.message) + .bind(sqlx::types::Json(&event.metadata)) + .bind(event.created_at) + .execute(&mut *tx) + .await?; + self.buckets.upsert_success(&mut tx, event).await?; + tx.commit().await?; + Ok(()) + } + + pub async fn trim_to_retention(&self, now: DateTime) -> PgUsageRecordStoreResult { + let days: i64 = + sqlx::query_scalar("select usage_retention_days from runtime_settings where id = 1") + .fetch_one(&self.pool) + .await?; + let duration = + Duration::try_days(days).ok_or(PgUsageRecordStoreError::InvalidRetention(days))?; + let cutoff = now + .checked_sub_signed(duration) + .ok_or(PgUsageRecordStoreError::InvalidRetention(days))?; + let result = sqlx::query("delete from usage_records where created_at < $1") + .bind(cutoff) + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } + + pub async fn list( + &self, + filter: UsageRecordFilter, + cursor: Option, + limit: u32, + ) -> PgUsageRecordStoreResult> { + let limit = limit.clamp(1, 200); + let mut builder = QueryBuilder::::new(USAGE_RECORD_SELECT_SQL); + push_filter(&mut builder, &filter, cursor.as_deref())?; + builder.push(" order by created_at desc, id desc limit "); + builder.push_bind(i64::from(limit) + 1); + let rows = builder.build().fetch_all(&self.pool).await?; + let has_next = rows.len() > limit as usize; + let items = rows + .into_iter() + .take(limit as usize) + .map(|row| usage_record_from_row(&row)) + .collect::>(); + let next_cursor = has_next.then(|| { + let item = items.last().expect("next page requires a usage record"); + encode_cursor(&item.created_at.to_rfc3339(), &item.id) + }); + Ok(Page { items, next_cursor }) + } + + pub async fn list_page( + &self, + filter: UsageRecordFilter, + page: u32, + page_size: u32, + ) -> PgUsageRecordStoreResult> { + let page_size = page_size.clamp(1, 200); + let total = count_usage_records(&self.pool, &filter).await?; + let mut builder = QueryBuilder::::new(USAGE_RECORD_SELECT_SQL); + push_filter(&mut builder, &filter, None)?; + builder.push(" order by created_at desc, id desc 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); + let items = builder + .build() + .fetch_all(&self.pool) + .await? + .iter() + .map(usage_record_from_row) + .collect(); + Ok(NumberedPage { + items, + total, + page: page.max(1), + page_size, + }) + } + + pub async fn get(&self, id: &str) -> PgUsageRecordStoreResult> { + let mut builder = QueryBuilder::::new(USAGE_RECORD_SELECT_SQL); + builder.push(" where id = "); + builder.push_bind(id); + Ok(builder + .build() + .fetch_optional(&self.pool) + .await? + .as_ref() + .map(usage_record_from_row)) + } + + pub async fn account_email_map( + &self, + items: &[UsageRecord], + ) -> PgUsageRecordStoreResult> { + usage_record_account_email_map(&self.pool, items).await + } + + pub async fn summary( + &self, + filter: UsageRecordFilter, + ) -> PgUsageRecordStoreResult { + usage_summary(&self.pool, &filter).await + } + + pub async fn account_usage( + &self, + filter: UsageRecordFilter, + limit: u32, + ) -> PgUsageRecordStoreResult> { + usage_account_usage(&self.pool, &filter, limit).await + } + + pub async fn model_distribution( + &self, + filter: UsageRecordFilter, + source: UsageRecordModelSource, + ) -> PgUsageRecordStoreResult> { + let expression = match source { + UsageRecordModelSource::Requested => { + "coalesce(nullif(requested_model, ''), model)" + } + UsageRecordModelSource::Upstream => { + "coalesce(nullif(upstream_model, ''), model)" + } + UsageRecordModelSource::Mapping => { + "coalesce(nullif(requested_model, ''), model) || ' -> ' || coalesce(nullif(upstream_model, ''), model)" + } + }; + usage_breakdown(&self.pool, &filter, expression, 8).await + } + + pub async fn endpoint_distribution( + &self, + filter: UsageRecordFilter, + source: UsageRecordEndpointSource, + ) -> PgUsageRecordStoreResult> { + let expression = match source { + UsageRecordEndpointSource::Inbound => "coalesce(nullif(route, ''), '未知端点')", + UsageRecordEndpointSource::Upstream => "provider", + UsageRecordEndpointSource::Path => { + "coalesce(nullif(route, ''), '未知端点') || ' -> ' || provider" + } + }; + usage_breakdown(&self.pool, &filter, expression, 8).await + } + + pub async fn trend( + &self, + filter: UsageRecordFilter, + ) -> PgUsageRecordStoreResult> { + usage_trend(&self.pool, &filter).await + } + + pub async fn clear(&self) -> PgUsageRecordStoreResult { + let result = sqlx::query("delete from usage_records") + .execute(&self.pool) + .await?; + Ok(result.rows_affected()) + } +} + +fn push_filter( + builder: &mut QueryBuilder, + filter: &UsageRecordFilter, + cursor: Option<&str>, +) -> PgUsageRecordStoreResult<()> { + builder.push(" where true"); + push_optional_text(builder, "kind", filter.kind.as_deref()); + push_optional_text( + builder, + "client_api_key_id", + filter.client_api_key_id.as_deref(), + ); + push_optional_text(builder, "provider", filter.provider.as_deref()); + push_optional_text(builder, "request_id", filter.request_id.as_deref()); + push_optional_text(builder, "account_id", filter.account_id.as_deref()); + push_optional_text(builder, "route", filter.route.as_deref()); + push_optional_text(builder, "model", filter.model.as_deref()); + push_optional_i64(builder, "status_code", filter.status_code); + push_optional_text(builder, "transport", filter.transport.as_deref()); + push_optional_i64(builder, "attempt_index", filter.attempt_index); + push_optional_text(builder, "response_id", filter.response_id.as_deref()); + push_optional_text( + builder, + "upstream_request_id", + filter.upstream_request_id.as_deref(), + ); + if let Some(search) = filter + .search + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + builder.push(" and (message ilike "); + builder.push_bind(format!("%{search}%")); + builder.push(" or request_id = "); + builder.push_bind(search); + builder.push(" or response_id = "); + builder.push_bind(search); + builder.push(")"); + } + if let Some(start_time) = filter.start_time { + builder.push(" and created_at >= "); + builder.push_bind(start_time); + } + if let Some(end_time) = filter.end_time { + builder.push(" and created_at <= "); + builder.push_bind(end_time); + } + if let Some(cursor) = cursor { + let (created_at, id) = + decode_cursor(cursor).ok_or(PgUsageRecordStoreError::InvalidCursor)?; + let created_at = + parse_rfc3339_utc(&created_at).map_err(|_| PgUsageRecordStoreError::InvalidCursor)?; + builder.push(" and (created_at < "); + builder.push_bind(created_at); + builder.push(" or (created_at = "); + builder.push_bind(created_at); + builder.push(" and id < "); + builder.push_bind(id); + builder.push("))"); + } + Ok(()) +} + +fn push_optional_text(builder: &mut QueryBuilder, column: &str, value: Option<&str>) { + if let Some(value) = value { + builder.push(" and "); + builder.push(column); + builder.push(" = "); + builder.push_bind(value); + } +} + +fn push_optional_i64(builder: &mut QueryBuilder, column: &str, value: Option) { + if let Some(value) = value { + builder.push(" and "); + builder.push(column); + builder.push(" = "); + builder.push_bind(value); + } +} diff --git a/backend/src/telemetry/usage/store/query.rs b/backend/src/telemetry/usage/store/query.rs new file mode 100644 index 000000000..6f862e33d --- /dev/null +++ b/backend/src/telemetry/usage/store/query.rs @@ -0,0 +1,432 @@ +use super::*; + +pub(super) async fn count_usage_records( + pool: &PgPool, + filter: &UsageRecordFilter, +) -> PgUsageRecordStoreResult { + let mut builder = QueryBuilder::::new("select count(*) from usage_records"); + push_filter(&mut builder, filter, None)?; + let total: i64 = builder.build_query_scalar().fetch_one(pool).await?; + Ok(total.max(0) as u64) +} + +pub(super) async fn usage_summary( + pool: &PgPool, + filter: &UsageRecordFilter, +) -> PgUsageRecordStoreResult { + let mut builder = QueryBuilder::::new( + r" +select + count(*) as total_requests, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + avg(latency_ms::double precision) as average_latency_ms +from usage_records", + ); + push_filter(&mut builder, filter, None)?; + let row = builder.build().fetch_one(pool).await?; + let input_tokens = nonnegative(row.get("input_tokens")); + let output_tokens = nonnegative(row.get("output_tokens")); + let cached_tokens = nonnegative(row.get("cached_tokens")); + Ok(UsageRecordSummary { + total_requests: nonnegative(row.get("total_requests")), + input_tokens, + output_tokens, + cached_tokens, + total_tokens: input_tokens + output_tokens, + average_latency_ms: row.get("average_latency_ms"), + }) +} + +pub(super) async fn usage_account_usage( + pool: &PgPool, + filter: &UsageRecordFilter, + limit: u32, +) -> PgUsageRecordStoreResult> { + let mut builder = QueryBuilder::::new( + r" +select + account_id, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + max(created_at) as last_used_at +from usage_records", + ); + push_filter(&mut builder, filter, None)?; + builder.push(" group by account_id order by last_used_at desc, account_id limit "); + builder.push_bind(i64::from(limit.clamp(1, 50))); + Ok(builder + .build() + .fetch_all(pool) + .await? + .into_iter() + .map(|row| { + let input_tokens = nonnegative(row.get("input_tokens")); + let output_tokens = nonnegative(row.get("output_tokens")); + UsageRecordAccountUsage { + account_id: row.get("account_id"), + input_tokens, + output_tokens, + total_tokens: input_tokens + output_tokens, + last_used_at: row.get("last_used_at"), + } + }) + .collect()) +} + +pub(super) async fn usage_record_account_email_map( + pool: &PgPool, + items: &[UsageRecord], +) -> PgUsageRecordStoreResult> { + let mut account_ids = items + .iter() + .map(|item| item.account_id.clone()) + .collect::>(); + account_ids.sort_unstable(); + account_ids.dedup(); + if account_ids.is_empty() { + return Ok(HashMap::new()); + } + let mut builder = QueryBuilder::::new("select id, email from accounts where id in ("); + let mut values = builder.separated(", "); + for account_id in account_ids { + values.push_bind(account_id); + } + values.push_unseparated(")"); + Ok(builder + .build() + .fetch_all(pool) + .await? + .into_iter() + .filter_map(|row| { + row.get::, _>("email") + .filter(|email| !email.trim().is_empty()) + .map(|email| (row.get("id"), email)) + }) + .collect()) +} + +pub(super) async fn usage_breakdown( + pool: &PgPool, + filter: &UsageRecordFilter, + group_expression: &str, + limit: u32, +) -> PgUsageRecordStoreResult> { + let mut builder = QueryBuilder::::new("select "); + builder.push(group_expression); + builder.push( + r" as name, + coalesce(nullif(upstream_model, ''), model) as billing_model, + service_tier, + count(*) as request_count, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + coalesce(sum(latency_ms), 0)::bigint as latency_sum, + count(latency_ms) as latency_count +from usage_records", + ); + push_filter(&mut builder, filter, None)?; + builder.push(" group by name, billing_model, service_tier order by request_count desc limit "); + builder.push_bind(i64::from(limit.clamp(1, 50) * 8)); + + let mut grouped = BTreeMap::::new(); + for row in builder.build().fetch_all(pool).await? { + let name: String = row.get("name"); + let input_tokens = nonnegative(row.get("input_tokens")); + let output_tokens = nonnegative(row.get("output_tokens")); + let cached_tokens = nonnegative(row.get("cached_tokens")); + let request_count = nonnegative(row.get("request_count")); + let latency_sum = nonnegative(row.get("latency_sum")); + let latency_count = nonnegative(row.get("latency_count")); + let cost = usage_breakdown_cost( + input_tokens, + output_tokens, + cached_tokens, + &row.get::("billing_model"), + row.get::, _>("service_tier").as_deref(), + ); + grouped + .entry(name.clone()) + .or_insert_with(|| BreakdownAccumulator::new(name)) + .push( + request_count, + input_tokens, + output_tokens, + cached_tokens, + cost, + latency_sum, + latency_count, + ); + } + let mut items = grouped + .into_values() + .map(BreakdownAccumulator::finish) + .collect::>(); + items.sort_by(|left, right| { + right + .request_count + .cmp(&left.request_count) + .then_with(|| left.name.cmp(&right.name)) + }); + items.truncate(limit.clamp(1, 50) as usize); + Ok(items) +} + +pub(super) async fn usage_trend( + pool: &PgPool, + filter: &UsageRecordFilter, +) -> PgUsageRecordStoreResult> { + let mut builder = QueryBuilder::::new( + r" +select + to_char(created_at at time zone 'Asia/Shanghai', 'YYYY-MM-DD') as date, + coalesce(nullif(upstream_model, ''), model) as billing_model, + service_tier, + coalesce(sum(input_tokens), 0)::bigint as input_tokens, + coalesce(sum(output_tokens), 0)::bigint as output_tokens, + coalesce(sum(cached_tokens), 0)::bigint as cached_tokens, + coalesce(sum(latency_ms), 0)::bigint as latency_sum, + count(latency_ms) as latency_count +from usage_records", + ); + push_filter(&mut builder, filter, None)?; + builder.push(" group by date, billing_model, service_tier order by date asc"); + + let mut days = BTreeMap::::new(); + for row in builder.build().fetch_all(pool).await? { + let date: String = row.get("date"); + let input_tokens = nonnegative(row.get("input_tokens")); + let output_tokens = nonnegative(row.get("output_tokens")); + let cached_tokens = nonnegative(row.get("cached_tokens")); + let cost = usage_breakdown_cost( + input_tokens, + output_tokens, + cached_tokens, + &row.get::("billing_model"), + row.get::, _>("service_tier").as_deref(), + ); + days.entry(date.clone()) + .or_insert_with(|| UsageTrendAccumulator::new(date)) + .push( + input_tokens, + output_tokens, + cached_tokens, + cost, + nonnegative(row.get("latency_sum")), + nonnegative(row.get("latency_count")), + ); + } + let mut points = days + .into_values() + .map(UsageTrendAccumulator::finish) + .collect::>(); + if points.len() > 60 { + points = points.split_off(points.len() - 60); + } + Ok(points) +} + +pub(super) fn usage_record_from_row(row: &sqlx::postgres::PgRow) -> UsageRecord { + UsageRecord { + id: row.get("id"), + request_id: row.get("request_id"), + client_api_key_id: row.get("client_api_key_id"), + kind: row.get("kind"), + provider: row.get("provider"), + account_id: row.get("account_id"), + route: row.get("route"), + model: row.get("model"), + requested_model: row.get("requested_model"), + upstream_model: row.get("upstream_model"), + service_tier: row.get("service_tier"), + status_code: row.get::("status_code").into(), + transport: row.get("transport"), + attempt_index: row.get("attempt_index"), + response_id: row.get("response_id"), + upstream_request_id: row.get("upstream_request_id"), + latency_ms: row.get("latency_ms"), + first_token_ms: row.get("first_token_ms"), + input_tokens: row.get("input_tokens"), + output_tokens: row.get("output_tokens"), + cached_tokens: row.get("cached_tokens"), + reasoning_tokens: row.get("reasoning_tokens"), + message: row.get("message"), + metadata: row + .get::, _>("metadata_json") + .0, + created_at: row.get("created_at"), + } +} + +fn nonnegative(value: Option) -> u64 { + optional_nonnegative_i64_to_u64(value) +} + +fn usage_breakdown_cost( + input_tokens: u64, + output_tokens: u64, + cached_tokens: u64, + model: &str, + service_tier: Option<&str>, +) -> f64 { + billing::calculate_cost( + input_tokens, + output_tokens, + cached_tokens, + model, + service_tier, + ) +} + +struct BreakdownAccumulator { + item: UsageRecordBreakdown, + latency_sum: u64, + latency_count: u64, +} + +impl BreakdownAccumulator { + fn new(name: String) -> Self { + Self { + item: UsageRecordBreakdown { + name, + ..Default::default() + }, + latency_sum: 0, + latency_count: 0, + } + } + + #[allow(clippy::too_many_arguments)] + fn push( + &mut self, + requests: u64, + input: u64, + output: u64, + cached: u64, + cost: f64, + latency_sum: u64, + latency_count: u64, + ) { + self.item.request_count += requests; + self.item.input_tokens += input; + self.item.output_tokens += output; + self.item.cached_tokens += cached; + self.item.total_tokens += input + output; + self.item.cost += cost; + self.item.actual_cost += cost; + self.item.account_cost += cost; + self.latency_sum += latency_sum; + self.latency_count += latency_count; + } + + fn finish(mut self) -> UsageRecordBreakdown { + self.item.average_latency_ms = + (self.latency_count > 0).then(|| self.latency_sum as f64 / self.latency_count as f64); + self.item + } +} + +struct UsageTrendAccumulator { + point: UsageRecordTrendPoint, + latency_sum: u64, + latency_count: u64, +} + +impl UsageTrendAccumulator { + fn new(date: String) -> Self { + Self { + point: UsageRecordTrendPoint { + date, + ..Default::default() + }, + latency_sum: 0, + latency_count: 0, + } + } + + fn push( + &mut self, + input: u64, + output: u64, + cached: u64, + cost: f64, + latency_sum: u64, + latency_count: u64, + ) { + self.point.input_tokens += input; + self.point.output_tokens += output; + self.point.cached_tokens += cached; + self.point.total_tokens += input + output; + self.point.cost += cost; + self.point.actual_cost += cost; + self.latency_sum += latency_sum; + self.latency_count += latency_count; + } + + fn finish(mut self) -> UsageRecordTrendPoint { + self.point.average_latency_ms = + (self.latency_count > 0).then(|| self.latency_sum as f64 / self.latency_count as f64); + self.point + } +} + +#[derive(Debug, Clone, Copy, Default, PartialEq)] +pub struct UsageRecordSummary { + pub total_requests: u64, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub average_latency_ms: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UsageRecordAccountUsage { + pub account_id: String, + pub input_tokens: u64, + pub output_tokens: u64, + pub total_tokens: u64, + pub last_used_at: DateTime, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UsageRecordModelSource { + Requested, + Upstream, + Mapping, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UsageRecordEndpointSource { + Inbound, + Upstream, + Path, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct UsageRecordBreakdown { + pub name: String, + pub request_count: u64, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub cost: f64, + pub actual_cost: f64, + pub account_cost: f64, + pub average_latency_ms: Option, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct UsageRecordTrendPoint { + pub date: String, + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_tokens: u64, + pub total_tokens: u64, + pub cost: f64, + pub actual_cost: f64, + pub average_latency_ms: Option, +} diff --git a/backend/src/admin/monitoring/usage_record_model.rs b/backend/src/telemetry/usage/types.rs similarity index 65% rename from backend/src/admin/monitoring/usage_record_model.rs rename to backend/src/telemetry/usage/types.rs index 877433d29..9f6f9af3a 100644 --- a/backend/src/admin/monitoring/usage_record_model.rs +++ b/backend/src/telemetry/usage/types.rs @@ -27,32 +27,46 @@ pub struct UsageRecord { pub id: String, /// 请求 ID。 pub request_id: Option, + /// 调用方客户端 API key 的稳定 ID。 + pub client_api_key_id: Option, /// 事件类别。 pub kind: String, - /// 事件等级。 - pub level: UsageRecordLevel, + /// 上游 provider。 + pub provider: String, /// 账号 ID。 - pub account_id: Option, + pub account_id: String, /// HTTP 路由。 pub route: Option, /// 模型名。 - pub model: Option, + pub model: String, + /// 客户端请求模型。 + pub requested_model: Option, + /// 上游实际模型。 + pub upstream_model: Option, + /// 服务层级。 + pub service_tier: Option, /// HTTP 状态码。 - pub status_code: Option, + pub status_code: i64, /// 上游传输方式，例如 http_sse 或 websocket。 pub transport: Option, /// 同一请求内的上游尝试序号。 pub attempt_index: Option, - /// 上游 HTTP 状态码。 - pub upstream_status_code: Option, - /// 失败分类。 - pub failure_class: Option, /// 上游响应 ID。 pub response_id: Option, /// 上游请求 ID。 pub upstream_request_id: Option, /// 请求耗时毫秒。 pub latency_ms: Option, + /// 首 token 延迟毫秒。 + pub first_token_ms: Option, + /// 输入 token。 + pub input_tokens: Option, + /// 输出 token。 + pub output_tokens: Option, + /// 缓存 token。 + pub cached_tokens: Option, + /// reasoning token。 + pub reasoning_tokens: Option, /// 展示消息。 pub message: String, /// 结构化元数据。 @@ -66,25 +80,34 @@ impl UsageRecord { /// 构造使用记录。 pub fn new( kind: impl Into, - level: UsageRecordLevel, message: impl Into, + account_id: impl Into, + model: impl Into, + status_code: i64, ) -> Self { Self { id: format!("usage_{}", Uuid::new_v4().simple()), request_id: None, + client_api_key_id: None, kind: kind.into(), - level, - account_id: None, + provider: "openai".to_string(), + account_id: account_id.into(), route: None, - model: None, - status_code: None, + model: model.into(), + requested_model: None, + upstream_model: None, + service_tier: None, + status_code, transport: None, attempt_index: None, - upstream_status_code: None, - failure_class: None, response_id: None, upstream_request_id: None, latency_ms: None, + first_token_ms: None, + input_tokens: None, + output_tokens: None, + cached_tokens: None, + reasoning_tokens: None, message: message.into(), metadata: Value::Object(Map::default()), created_at: Utc::now(), @@ -92,7 +115,7 @@ impl UsageRecord { } } -pub(in crate::admin) fn metadata_service_tier(metadata: &Value) -> Option<&str> { +pub(crate) fn metadata_service_tier(metadata: &Value) -> Option<&str> { metadata .get("serviceTier") .and_then(Value::as_str) @@ -100,7 +123,7 @@ pub(in crate::admin) fn metadata_service_tier(metadata: &Value) -> Option<&str> .filter(|value| !value.is_empty()) } -pub(in crate::admin) fn metadata_string(metadata: &Value, keys: &[&str]) -> Option { +pub(crate) fn metadata_string(metadata: &Value, keys: &[&str]) -> Option { keys.iter().find_map(|key| { metadata .get(*key) @@ -111,19 +134,16 @@ pub(in crate::admin) fn metadata_string(metadata: &Value, keys: &[&str]) -> Opti }) } -pub(in crate::admin) fn metadata_i64(metadata: &Value, keys: &[&str]) -> Option { +pub(crate) fn metadata_i64(metadata: &Value, keys: &[&str]) -> Option { keys.iter() .find_map(|key| metadata.get(*key).and_then(Value::as_i64)) } -pub(in crate::admin) fn metadata_nonnegative_i64(metadata: &Value, keys: &[&str]) -> Option { - metadata_i64(metadata, keys).filter(|value| *value >= 0) -} - /// 响应事件记录（供 dispatch 模块使用）。 pub struct ResponseUsageRecord<'a> { - pub usage_records: &'a crate::admin::monitoring::usage_record_service::AdminUsageRecordService, + pub usage_records: &'a crate::telemetry::usage::query::UsageQueryService, pub request_id: &'a str, + pub client_api_key_id: Option<&'a str>, pub account_id: &'a str, pub route: &'a str, pub model: &'a str, @@ -134,7 +154,6 @@ pub struct ResponseUsageRecord<'a> { pub service_tier: Option<&'a str>, pub started_at: std::time::Instant, pub status_code: i64, - pub level: UsageRecordLevel, pub message: &'a str, pub metadata: serde_json::Value, pub rate_limit_headers: &'a [(String, String)], diff --git a/backend/src/admin/system/updater/archive.rs b/backend/src/update/archive.rs similarity index 93% rename from backend/src/admin/system/updater/archive.rs rename to backend/src/update/archive.rs index b402e3d58..d32952290 100644 --- a/backend/src/admin/system/updater/archive.rs +++ b/backend/src/update/archive.rs @@ -8,9 +8,9 @@ use std::{ use flate2::read::GzDecoder; use tar::Archive; -use crate::admin::response::AdminError; +use crate::update::types::UpdateError; -use super::{internal_error, internal_error_with, SystemUpdateConfig, APP_BINARY_NAME}; +use super::service::{internal_error, internal_error_with, SystemUpdateConfig, APP_BINARY_NAME}; #[derive(Debug)] pub(super) struct ExtractedRelease { @@ -18,11 +18,13 @@ pub(super) struct ExtractedRelease { web_dist_dir: Option, } -pub(super) async fn rollback_release_update(config: &SystemUpdateConfig) -> Result<(), AdminError> { +pub(super) async fn rollback_release_update( + config: &SystemUpdateConfig, +) -> Result<(), UpdateError> { let exe_path = config.executable_path()?; let backup_path = backup_path_for(&exe_path); if !backup_path.exists() { - return Err(AdminError::conflict("No binary backup found for rollback")); + return Err(UpdateError::conflict("No binary backup found for rollback")); } fs::rename(&backup_path, &exe_path).map_err(internal_error_with("Binary rollback failed"))?; @@ -41,7 +43,7 @@ pub(super) async fn rollback_release_update(config: &SystemUpdateConfig) -> Resu pub(super) fn extract_release_archive( archive_path: &Path, temp_dir: &Path, -) -> Result { +) -> Result { let file = fs::File::open(archive_path).map_err(internal_error_with("Failed to open archive"))?; let decoder = GzDecoder::new(file); @@ -67,7 +69,7 @@ pub(super) fn extract_release_archive( .map_err(internal_error_with("Invalid archive path"))? .to_path_buf(); if unsafe_archive_path(&path) { - return Err(AdminError::bad_request("Unsafe archive path")); + return Err(UpdateError::bad_request("Unsafe archive path")); } if path.file_name().is_some_and(|name| name == APP_BINARY_NAME) { @@ -92,7 +94,7 @@ pub(super) fn extract_release_archive( } if !found_binary { - return Err(AdminError::bad_request( + return Err(UpdateError::bad_request( "Release archive does not contain codex-proxy-rs", )); } @@ -113,7 +115,7 @@ pub(super) fn replace_release_files( exe_path: &Path, web_dist_dir: &Path, extracted: ExtractedRelease, -) -> Result<(), AdminError> { +) -> Result<(), UpdateError> { let web_backup = backup_path_for(web_dist_dir); let web_replaced = if let Some(new_web) = extracted.web_dist_dir { replace_dir(web_dist_dir, &web_backup, &new_web)?; @@ -148,7 +150,7 @@ pub(super) fn replace_release_files( Ok(()) } -fn replace_dir(current: &Path, backup: &Path, replacement: &Path) -> Result<(), AdminError> { +fn replace_dir(current: &Path, backup: &Path, replacement: &Path) -> Result<(), UpdateError> { if backup.exists() { fs::remove_dir_all(backup) .map_err(internal_error_with("Failed to remove old web backup"))?; diff --git a/backend/src/admin/system/updater/download.rs b/backend/src/update/download.rs similarity index 89% rename from backend/src/admin/system/updater/download.rs rename to backend/src/update/download.rs index f7e550f0e..c718cb0c5 100644 --- a/backend/src/admin/system/updater/download.rs +++ b/backend/src/update/download.rs @@ -9,9 +9,9 @@ use std::{ use sha2::{Digest, Sha256}; -use crate::admin::response::AdminError; +use crate::update::types::UpdateError; -use super::{ +use super::service::{ bad_gateway_with, bad_request_with, emit_update_event, format_bytes, internal_error_with, UpdateLogLevel, }; @@ -27,7 +27,7 @@ pub(super) async fn download_file( dest: &Path, max_size: u64, progress: Option>, -) -> Result<(), AdminError> { +) -> Result<(), UpdateError> { let client = reqwest::Client::builder() .timeout(Duration::from_secs(120)) .build() @@ -38,7 +38,7 @@ pub(super) async fn download_file( .await .map_err(bad_gateway_with("Download failed"))?; if !response.status().is_success() { - return Err(AdminError::bad_gateway(format!( + return Err(UpdateError::bad_gateway(format!( "Download failed with {}", response.status() ))); @@ -54,7 +54,9 @@ pub(super) async fn download_file( { downloaded += chunk.len() as u64; if downloaded > max_size { - return Err(AdminError::bad_request("Download exceeds max allowed size")); + return Err(UpdateError::bad_request( + "Download exceeds max allowed size", + )); } file.write_all(&chunk) .map_err(internal_error_with("Failed to write download"))?; @@ -84,7 +86,7 @@ pub(super) async fn verify_checksum( file_path: &Path, file_name: &str, checksum_url: &str, -) -> Result<(), AdminError> { +) -> Result<(), UpdateError> { let client = reqwest::Client::builder() .timeout(Duration::from_secs(30)) .build() @@ -96,7 +98,7 @@ pub(super) async fn verify_checksum( .map_err(bad_gateway_with("Checksum download failed"))?; let status = body.status(); if !status.is_success() { - return Err(AdminError::bad_gateway(format!( + return Err(UpdateError::bad_gateway(format!( "Checksum download failed with {status}" ))); } @@ -111,31 +113,31 @@ pub(super) async fn verify_checksum( let name = parts.next()?; (Path::new(name).file_name()?.to_string_lossy() == file_name).then(|| hash.to_string()) }); - let expected = expected.ok_or_else(|| AdminError::bad_gateway("Checksum not found"))?; + let expected = expected.ok_or_else(|| UpdateError::bad_gateway("Checksum not found"))?; let actual = sha256_file(file_path)?; if expected.eq_ignore_ascii_case(&actual) { Ok(()) } else { - Err(AdminError::bad_gateway("Checksum mismatch")) + Err(UpdateError::bad_gateway("Checksum mismatch")) } } pub(super) fn validate_download_url( raw_url: &str, github_api_base: &str, -) -> Result<(), AdminError> { +) -> Result<(), UpdateError> { let url = reqwest::Url::parse(raw_url).map_err(bad_request_with("Invalid download URL"))?; if local_update_test_download_allowed(&url, github_api_base) { return Ok(()); } if url.scheme() != "https" { - return Err(AdminError::bad_request( + return Err(UpdateError::bad_request( "Only HTTPS release downloads are allowed", )); } let host = url .host_str() - .ok_or_else(|| AdminError::bad_request("Download URL is missing host"))?; + .ok_or_else(|| UpdateError::bad_request("Download URL is missing host"))?; let allowed = host == "github.com" || host.ends_with(".github.com") || host == "objects.githubusercontent.com" @@ -143,7 +145,7 @@ pub(super) fn validate_download_url( if allowed { Ok(()) } else { - Err(AdminError::bad_request(format!( + Err(UpdateError::bad_request(format!( "Download host is not allowed: {host}" ))) } @@ -174,7 +176,7 @@ pub(super) fn validate_github_api_base(raw_url: &str) -> Result<(), String> { Ok(()) } -fn sha256_file(path: &Path) -> Result { +fn sha256_file(path: &Path) -> Result { let mut file = fs::File::open(path).map_err(internal_error_with("Failed to open checksum file"))?; let mut hasher = Sha256::new(); diff --git a/backend/src/update/mod.rs b/backend/src/update/mod.rs new file mode 100644 index 000000000..14ec8487d --- /dev/null +++ b/backend/src/update/mod.rs @@ -0,0 +1,8 @@ +//! 在线更新领域。 + +mod archive; +mod download; +mod release; +pub mod service; +pub mod state; +pub mod types; diff --git a/backend/src/admin/system/updater/release.rs b/backend/src/update/release.rs similarity index 96% rename from backend/src/admin/system/updater/release.rs rename to backend/src/update/release.rs index c0235c77c..d3622606c 100644 --- a/backend/src/admin/system/updater/release.rs +++ b/backend/src/update/release.rs @@ -9,11 +9,11 @@ use semver::Version; use serde::{Deserialize, Serialize}; use tokio::sync::Mutex; -use crate::admin::response::AdminError; +use crate::update::types::UpdateError; use super::{ - build_type_label, deployment_mode_label, download::validate_github_api_base, - SystemUpdateConfig, APP_BINARY_NAME, + download::validate_github_api_base, + service::{build_type_label, deployment_mode_label, SystemUpdateConfig, APP_BINARY_NAME}, }; const UPDATE_CACHE_TTL: Duration = Duration::from_secs(20 * 60); @@ -22,7 +22,7 @@ static RELEASE_CACHE: Mutex> = Mutex::const_new(None); #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -pub(super) struct UpdateInfoData { +pub(crate) struct UpdateInfoData { pub current_version: String, pub latest_version: String, pub has_update: bool, @@ -194,7 +194,7 @@ pub(super) fn update_info_from_release( pub(super) fn select_release_archive<'a>( release: &'a GitHubRelease, version: &str, -) -> Result<&'a GitHubAsset, AdminError> { +) -> Result<&'a GitHubAsset, UpdateError> { let os_aliases = platform_os_aliases(); let arch_aliases = platform_arch_aliases(); let normalized = normalize_version_tag(version); @@ -216,7 +216,7 @@ pub(super) fn select_release_archive<'a>( }) }) .ok_or_else(|| { - AdminError::conflict(format!( + UpdateError::conflict(format!( "No compatible release archive found for {}/{}", env::consts::OS, env::consts::ARCH diff --git a/backend/src/admin/system/updater/mod.rs b/backend/src/update/service.rs similarity index 68% rename from backend/src/admin/system/updater/mod.rs rename to backend/src/update/service.rs index 614e39a66..ec3fdd351 100644 --- a/backend/src/admin/system/updater/mod.rs +++ b/backend/src/update/service.rs @@ -1,52 +1,34 @@ -//! 管理端系统版本与在线更新路由。 +//! 系统版本、自更新与重启领域服务。 use std::{ - convert::Infallible, env, fmt, fs, io, path::{Path, PathBuf}, process::Command, sync::OnceLock, - time::Duration, }; -use axum::{ - extract::{Json, Query}, - http::StatusCode, - response::{ - sse::{Event, KeepAlive, Sse}, - IntoResponse, - }, -}; use chrono::Utc; -use futures::Stream; -use serde::{Deserialize, Serialize}; +use serde::Serialize; use tokio::sync::{broadcast, Mutex}; -use crate::{ - admin::auth::session::AdminAuth, - admin::response::{AdminEnvelope, AdminError, AdminResponse}, - runtime::shutdown::{request_process_restart, request_shutdown, shutdown_subscription}, +use super::{ + archive::{extract_release_archive, replace_release_files, rollback_release_update}, + download::{ + download_file, validate_download_url, validate_github_api_base, verify_checksum, + DownloadProgress, + }, + release::{ + check_latest_release, fetch_latest_release, normalize_version_tag, select_release_archive, + update_info_from_release, GitHubRelease, UpdateInfoData, + }, + state::{ + finish_operation, operation_id, read_state, set_operation_running, OperationLock, + SystemOperationKind, SystemUpdateStatusData, + }, + types::UpdateError, }; -use super::state::{ - finish_operation, operation_id, read_state, set_operation_running, OperationLock, - SystemOperationKind, -}; - -mod download; -use download::{ - download_file, validate_download_url, validate_github_api_base, verify_checksum, - DownloadProgress, -}; -mod archive; -use archive::{extract_release_archive, replace_release_files, rollback_release_update}; -mod release; -use release::{ - check_latest_release, fetch_latest_release, normalize_version_tag, select_release_archive, - update_info_from_release, GitHubRelease, UpdateInfoData, -}; - -const APP_BINARY_NAME: &str = "codex-proxy-rs"; +pub(super) const APP_BINARY_NAME: &str = "codex-proxy-rs"; const DEFAULT_WEB_DIST_DIR: &str = "/app/web/dist"; const GITHUB_API_BASE: &str = "https://api.github.com/repos"; const MAX_DOWNLOAD_SIZE: u64 = 500 * 1024 * 1024; @@ -56,21 +38,9 @@ const REPLACEMENT_START_DELAY_MS: &str = "1200"; static SYSTEM_OPERATION_LOCK: Mutex<()> = Mutex::const_new(()); static UPDATE_EVENT_SENDER: OnceLock> = OnceLock::new(); -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct UpdateDetailQuery { - refresh: Option, -} - -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] -pub(crate) struct UpdateRequest { - target_version: String, -} - #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -struct VersionData { +pub(crate) struct VersionData { version: String, git_sha: String, build_time: String, @@ -85,7 +55,7 @@ struct VersionData { #[derive(Debug, Serialize)] #[serde(rename_all = "camelCase")] -struct UpdateStartedData { +pub(crate) struct UpdateStartedData { operation_id: String, deployment_mode: String, message: String, @@ -95,7 +65,7 @@ struct UpdateStartedData { #[derive(Debug, Clone, Copy, Serialize)] #[serde(rename_all = "lowercase")] -enum UpdateLogLevel { +pub(super) enum UpdateLogLevel { Info, Success, Warning, @@ -104,7 +74,7 @@ enum UpdateLogLevel { #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] -struct SystemUpdateEvent { +pub(crate) struct SystemUpdateEvent { id: String, operation_id: Option, level: UpdateLogLevel, @@ -115,88 +85,54 @@ struct SystemUpdateEvent { at: String, } -#[derive(Debug, Clone)] -struct SystemUpdateConfig { - version: String, - git_sha: String, - build_time: String, - deployment_mode: String, - build_type: String, - update_channel: String, - update_repository: Option, - github_api_base: String, - executable_path: Option, - web_dist_dir: PathBuf, - update_state_file: PathBuf, - update_lock_file: PathBuf, - update_temp_dir: PathBuf, +impl SystemUpdateEvent { + pub(crate) fn id(&self) -> &str { + &self.id + } + + pub(crate) fn is_terminal(&self) -> bool { + self.terminal + } } -/// `GET /api/admin/system/version` -pub(crate) async fn version(_auth: AdminAuth) -> Result { +#[derive(Debug, Clone)] +pub(super) struct SystemUpdateConfig { + pub(super) version: String, + pub(super) git_sha: String, + pub(super) build_time: String, + pub(super) deployment_mode: String, + pub(super) build_type: String, + pub(super) update_channel: String, + pub(super) update_repository: Option, + pub(super) github_api_base: String, + pub(super) executable_path: Option, + pub(super) web_dist_dir: PathBuf, + pub(super) update_state_file: PathBuf, + pub(super) update_lock_file: PathBuf, + pub(super) update_temp_dir: PathBuf, +} + +pub(crate) async fn version_data() -> VersionData { let config = SystemUpdateConfig::from_env(); let update_info = check_latest_release(&config, false).await; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(config.version_data(&update_info)), - )) + config.version_data(&update_info) } -/// `GET /api/admin/system/update-detail` -pub(crate) async fn update_detail( - _auth: AdminAuth, - Query(query): Query, -) -> Result { +pub(crate) async fn update_detail(refresh: bool) -> UpdateInfoData { let config = SystemUpdateConfig::from_env(); - let info = check_latest_release(&config, query.refresh.unwrap_or(false)).await; - Ok(AdminResponse::new(StatusCode::OK, AdminEnvelope::ok(info))) + check_latest_release(&config, refresh).await } -/// `GET /api/admin/system/update-events` -pub(crate) async fn update_event_stream( - _auth: AdminAuth, -) -> Result>>, AdminError> { - let receiver = update_event_sender().subscribe(); - let shutdown = shutdown_subscription(); - let stream = futures::stream::unfold( - (receiver, shutdown, false), - |(mut receiver, mut shutdown, close_after_send)| async move { - if close_after_send { - return None; - } - - loop { - tokio::select! { - _ = shutdown.recv() => return None, - received = receiver.recv() => { - match received { - Ok(message) => { - let close_after_send = message.terminal; - let id = message.id.clone(); - let data = serde_json::to_string(&message).unwrap_or_else(|_| "{}".to_string()); - let event = Event::default().event("update").id(id).data(data); - return Some((Ok(event), (receiver, shutdown, close_after_send))); - } - Err(broadcast::error::RecvError::Lagged(_)) => continue, - Err(broadcast::error::RecvError::Closed) => return None, - } - } - } - } - }, - ); - - Ok(Sse::new(stream).keep_alive(KeepAlive::default())) +pub(crate) fn subscribe_update_events() -> broadcast::Receiver { + update_event_sender().subscribe() } -/// `POST /api/admin/system/update` pub(crate) async fn perform_update( - _auth: AdminAuth, - payload: Option>, -) -> Result { + target_version: Option, +) -> Result { let _guard = SYSTEM_OPERATION_LOCK .try_lock() - .map_err(|_| AdminError::conflict("System update already running"))?; + .map_err(|_| UpdateError::conflict("System update already running"))?; let config = SystemUpdateConfig::from_env(); if let Some(reason) = config.update_support_error() { emit_terminal_update_event( @@ -205,9 +141,9 @@ pub(crate) async fn perform_update( Some("preflight"), reason.clone(), ); - return Err(AdminError::conflict(reason)); + return Err(UpdateError::conflict(reason)); } - let confirmed_target_version = confirmed_update_target(payload)?; + let confirmed_target_version = confirmed_update_target(target_version)?; emit_update_event( UpdateLogLevel::Info, @@ -220,12 +156,12 @@ pub(crate) async fn perform_update( config .update_repository .as_deref() - .ok_or_else(|| AdminError::conflict("Update checks require CPR_UPDATE_REPOSITORY"))?, + .ok_or_else(|| UpdateError::conflict("Update checks require CPR_UPDATE_REPOSITORY"))?, ) .await .map_err(|error| { emit_terminal_update_event(UpdateLogLevel::Error, None, Some("release"), error.clone()); - AdminError::bad_gateway(error) + UpdateError::bad_gateway(error) })?; let info = update_info_from_release(&config, release.clone()); if info.latest_version != confirmed_target_version { @@ -239,7 +175,7 @@ pub(crate) async fn perform_update( Some("release"), message.clone(), ); - return Err(AdminError::conflict(message)); + return Err(UpdateError::conflict(message)); } if !info.has_update { emit_terminal_update_event( @@ -248,7 +184,7 @@ pub(crate) async fn perform_update( Some("release"), "当前版本已是最新", ); - return Err(AdminError::conflict("Already up to date")); + return Err(UpdateError::conflict("Already up to date")); } let target_version = info.latest_version.clone(); let operation_id = operation_id("update"); @@ -292,47 +228,38 @@ pub(crate) async fn perform_update( drop(lock); result?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(UpdateStartedData { - operation_id, - deployment_mode: config.deployment_mode, - message: "更新完成，请重启服务。".to_string(), - need_restart: true, - target_version, - }), - )) + Ok(UpdateStartedData { + operation_id, + deployment_mode: config.deployment_mode, + message: "更新完成，请重启服务。".to_string(), + need_restart: true, + target_version, + }) } -fn confirmed_update_target(payload: Option>) -> Result { - let Some(Json(payload)) = payload else { - return Err(AdminError::conflict("更新前需要确认目标版本")); +fn confirmed_update_target(target_version: Option) -> Result { + let Some(target_version) = target_version else { + return Err(UpdateError::conflict("更新前需要确认目标版本")); }; - let target_version = normalize_version_tag(&payload.target_version); + let target_version = normalize_version_tag(&target_version); if target_version.is_empty() { - return Err(AdminError::bad_request("目标版本不能为空")); + return Err(UpdateError::bad_request("目标版本不能为空")); } Ok(target_version) } -/// `GET /api/admin/system/update-status` -pub(crate) async fn update_status(_auth: AdminAuth) -> Result { +pub(crate) fn update_status() -> Result { let config = SystemUpdateConfig::from_env(); - let status = read_state(&config.update_state_file)?; - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(status), - )) + read_state(&config.update_state_file) } -/// `POST /api/admin/system/rollback` -pub(crate) async fn rollback(_auth: AdminAuth) -> Result { +pub(crate) async fn rollback() -> Result { let _guard = SYSTEM_OPERATION_LOCK .try_lock() - .map_err(|_| AdminError::conflict("系统操作正在执行中"))?; + .map_err(|_| UpdateError::conflict("系统操作正在执行中"))?; let config = SystemUpdateConfig::from_env(); if let Some(reason) = config.update_support_error() { - return Err(AdminError::conflict(reason)); + return Err(UpdateError::conflict(reason)); } let operation_id = operation_id("rollback"); @@ -354,56 +281,31 @@ pub(crate) async fn rollback(_auth: AdminAuth) -> Result Result { +pub(crate) fn restart_plan() -> Result { if env_string("CPR_ENABLE_SELF_RESTART").as_deref() != Some("true") { - return Err(AdminError::conflict( + return Err(UpdateError::conflict( "自重启未启用，请设置 CPR_ENABLE_SELF_RESTART=true", )); } let config = SystemUpdateConfig::from_env(); - let plan = schedule_restart(&config)?; - let message = plan.message; - - tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(500)).await; - match plan.action { - RestartAction::Exec(executable_path) => request_process_restart(executable_path), - RestartAction::Shutdown => request_shutdown(), - } - }); - - Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok(serde_json::json!({ - "message": message, - "operationId": operation_id("restart") - })), - )) + schedule_restart(&config) } -struct RestartPlan { - message: &'static str, - action: RestartAction, +pub(crate) struct RestartPlan { + pub(crate) message: &'static str, + pub(crate) action: RestartAction, } -enum RestartAction { +pub(crate) enum RestartAction { Exec(PathBuf), Shutdown, } -fn schedule_restart(config: &SystemUpdateConfig) -> Result { +fn schedule_restart(config: &SystemUpdateConfig) -> Result { if config.deployment_mode == "docker" { return Ok(RestartPlan { message: "已安排进程内重启", @@ -418,7 +320,7 @@ fn schedule_restart(config: &SystemUpdateConfig) -> Result Result<(), AdminError> { +fn spawn_replacement_process(config: &SystemUpdateConfig) -> Result<(), UpdateError> { let executable_path = config.executable_path()?; let mut command = Command::new(executable_path); command @@ -485,7 +387,7 @@ impl SystemUpdateConfig { } } - fn update_support_error(&self) -> Option { + pub(super) fn update_support_error(&self) -> Option { if self.build_type != "release" { return Some("一键更新需要正式构建包".to_string()); } @@ -498,7 +400,7 @@ impl SystemUpdateConfig { None } - fn release_cache_key(&self) -> String { + pub(super) fn release_cache_key(&self) -> String { format!( "{}|{}|{}|{}|{}", self.update_repository.as_deref().unwrap_or_default(), @@ -509,7 +411,7 @@ impl SystemUpdateConfig { ) } - fn executable_path(&self) -> Result { + pub(super) fn executable_path(&self) -> Result { if let Some(path) = &self.executable_path { return Ok(path.clone()); } @@ -519,20 +421,26 @@ impl SystemUpdateConfig { } } -fn internal_error(context: &str, error: impl fmt::Display) -> AdminError { - AdminError::internal(format!("{context}: {error}")) +pub(super) fn internal_error(context: &str, error: impl fmt::Display) -> UpdateError { + UpdateError::internal(format!("{context}: {error}")) } -fn internal_error_with(context: &'static str) -> impl FnOnce(E) -> AdminError { +pub(super) fn internal_error_with( + context: &'static str, +) -> impl FnOnce(E) -> UpdateError { move |error| internal_error(context, error) } -fn bad_request_with(context: &'static str) -> impl FnOnce(E) -> AdminError { - move |error| AdminError::bad_request(format!("{context}: {error}")) +pub(super) fn bad_request_with( + context: &'static str, +) -> impl FnOnce(E) -> UpdateError { + move |error| UpdateError::bad_request(format!("{context}: {error}")) } -fn bad_gateway_with(context: &'static str) -> impl FnOnce(E) -> AdminError { - move |error| AdminError::bad_gateway(format!("{context}: {error}")) +pub(super) fn bad_gateway_with( + context: &'static str, +) -> impl FnOnce(E) -> UpdateError { + move |error| UpdateError::bad_gateway(format!("{context}: {error}")) } async fn perform_release_update( @@ -540,7 +448,7 @@ async fn perform_release_update( release: &GitHubRelease, version: &str, operation_id: &str, -) -> Result<(), AdminError> { +) -> Result<(), UpdateError> { emit_update_event( UpdateLogLevel::Info, Some(operation_id), @@ -567,13 +475,13 @@ async fn perform_release_update( ); validate_download_url(&archive.browser_download_url, &config.github_api_base)?; if archive.size > MAX_DOWNLOAD_SIZE { - return Err(AdminError::bad_request("Release archive is too large")); + return Err(UpdateError::bad_request("Release archive is too large")); } let checksum = release .assets .iter() .find(|asset| asset.name == "checksums.txt") - .ok_or_else(|| AdminError::bad_gateway("Release checksums.txt is required"))?; + .ok_or_else(|| UpdateError::bad_gateway("Release checksums.txt is required"))?; validate_download_url(&checksum.browser_download_url, &config.github_api_base)?; emit_update_event( @@ -734,7 +642,7 @@ fn update_event_sender() -> &'static broadcast::Sender { }) } -fn emit_update_event( +pub(super) fn emit_update_event( level: UpdateLogLevel, operation_id: Option<&str>, step: Option<&str>, @@ -780,7 +688,7 @@ fn is_false(value: &bool) -> bool { !*value } -fn format_bytes(bytes: u64) -> String { +pub(super) fn format_bytes(bytes: u64) -> String { const KIB: u64 = 1024; const MIB: u64 = KIB * 1024; if bytes >= MIB { @@ -792,7 +700,7 @@ fn format_bytes(bytes: u64) -> String { format!("{bytes} B") } -fn deployment_mode_label(value: &str) -> &str { +pub(super) fn deployment_mode_label(value: &str) -> &str { match value { "docker" => "Docker", "source" => "源码部署", @@ -801,7 +709,7 @@ fn deployment_mode_label(value: &str) -> &str { } } -fn build_type_label(value: &str) -> &str { +pub(super) fn build_type_label(value: &str) -> &str { match value { "release" => "正式构建", "source" => "源码构建", diff --git a/backend/src/admin/system/state.rs b/backend/src/update/state.rs similarity index 94% rename from backend/src/admin/system/state.rs rename to backend/src/update/state.rs index 13b70b627..0081e7adb 100644 --- a/backend/src/admin/system/state.rs +++ b/backend/src/update/state.rs @@ -10,7 +10,7 @@ use std::{ use chrono::Utc; use serde::{Deserialize, Serialize}; -use crate::admin::response::AdminError; +use crate::update::types::UpdateError; #[derive(Debug, Deserialize, Serialize, Default)] #[serde(rename_all = "camelCase")] @@ -71,7 +71,7 @@ pub(crate) struct OperationLock { path: PathBuf, } -pub(crate) fn read_state(path: &Path) -> Result { +pub(crate) fn read_state(path: &Path) -> Result { if !path.exists() { return Ok(SystemUpdateStatusData::default()); } @@ -80,7 +80,7 @@ pub(crate) fn read_state(path: &Path) -> Result Result<(), AdminError> { +fn write_state(path: &Path, state: &SystemUpdateStatusData) -> Result<(), UpdateError> { if let Some(parent) = path .parent() .filter(|parent| !parent.as_os_str().is_empty()) @@ -99,7 +99,7 @@ pub(crate) fn set_operation_running( operation_id: &str, kind: SystemOperationKind, version: Option<&str>, -) -> Result<(), AdminError> { +) -> Result<(), UpdateError> { let mut state = read_state(path)?; state.operation = SystemOperationState { operation_id: Some(operation_id.to_string()), @@ -158,7 +158,7 @@ pub(crate) fn operation_id(kind: &str) -> String { } impl OperationLock { - pub(crate) fn acquire(path: &Path) -> Result { + pub(crate) fn acquire(path: &Path) -> Result { if let Some(parent) = path .parent() .filter(|parent| !parent.as_os_str().is_empty()) @@ -182,7 +182,7 @@ impl OperationLock { path: path.to_path_buf(), }); } - Err(AdminError::conflict("System update already running")) + Err(UpdateError::conflict("System update already running")) } Err(error) => Err(internal_error("Failed to create update lock", error)), } @@ -212,7 +212,7 @@ impl Drop for OperationLock { } } -fn stale_lock(path: &Path) -> Result { +fn stale_lock(path: &Path) -> Result { let metadata = fs::metadata(path).map_err(internal_error_with("Failed to read update lock"))?; let modified = metadata .modified() @@ -223,12 +223,12 @@ fn stale_lock(path: &Path) -> Result { .map_err(internal_error_with("Failed to calculate update lock age")) } -fn internal_error(context: &str, error: impl std::fmt::Display) -> AdminError { - AdminError::internal(format!("{context}: {error}")) +fn internal_error(context: &str, error: impl std::fmt::Display) -> UpdateError { + UpdateError::internal(format!("{context}: {error}")) } fn internal_error_with( context: &'static str, -) -> impl FnOnce(E) -> AdminError { +) -> impl FnOnce(E) -> UpdateError { move |error| internal_error(context, error) } diff --git a/backend/src/update/types.rs b/backend/src/update/types.rs new file mode 100644 index 000000000..7fd240713 --- /dev/null +++ b/backend/src/update/types.rs @@ -0,0 +1,33 @@ +//! 自更新领域值类型与错误。 + +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum UpdateError { + #[error("{0}")] + BadRequest(String), + #[error("{0}")] + Conflict(String), + #[error("{0}")] + BadGateway(String), + #[error("{0}")] + Internal(String), +} + +impl UpdateError { + pub fn bad_request(message: impl Into) -> Self { + Self::BadRequest(message.into()) + } + + pub fn conflict(message: impl Into) -> Self { + Self::Conflict(message.into()) + } + + pub fn bad_gateway(message: impl Into) -> Self { + Self::BadGateway(message.into()) + } + + pub fn internal(message: impl Into) -> Self { + Self::Internal(message.into()) + } +} diff --git a/backend/src/upstream/accounts/mod.rs b/backend/src/upstream/accounts/mod.rs deleted file mode 100644 index c26171219..000000000 --- a/backend/src/upstream/accounts/mod.rs +++ /dev/null @@ -1,19 +0,0 @@ -//! 账号领域模块 —— 模型、仓库、池调度、Cookie、配额、刷新与导入。 - -/// Cookie 捕获与重放策略。 -pub mod cookies; -/// 账号导入逻辑。 -pub mod importing; -/// 账号领域模型。 -pub mod model; -/// 账号池调度策略。 -pub mod pool; -/// 配额检查策略。 -pub mod quota; -/// 账号管理用例辅助。 -pub mod service; -/// SQLite 仓储适配器。 -pub mod store; -/// 刷新租约存储与 JWT claims 解码。 -pub mod token_refresh; -mod window; diff --git a/backend/src/upstream/accounts/store/mod.rs b/backend/src/upstream/accounts/store/mod.rs deleted file mode 100644 index 7952db6ca..000000000 --- a/backend/src/upstream/accounts/store/mod.rs +++ /dev/null @@ -1,1142 +0,0 @@ -//! SQLite 账号仓储适配器。 - -use async_trait::async_trait; -use chrono::{DateTime, Utc}; -use secrecy::{ExposeSecret, SecretString}; -use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool}; -use thiserror::Error; - -use crate::infra::{ - json::{decode_cursor, page_offset, NumberedPage, Page}, - time::parse_optional_rfc3339_utc as parse_optional_rfc3339, -}; -use crate::upstream::accounts::model::{ - Account, AccountModelUsageDelta, AccountStatus, AccountUsageDelta, -}; -use crate::upstream::accounts::window::should_reset_usage_window; - -mod queries; -mod rows; - -use queries::*; -use rows::{ - count_account_metadata, get_pool_account, list_pool_accounts, map_account_store_error, - metadata_from_row, optional_positive_i64_to_u64, optional_update_value, - push_account_metadata_search, quota_plan_type, quota_snapshot_from_row, sqlite_usage_delta, - status_to_db, stored_account_from_row, to_page, u64_to_i64_saturating, -}; - -// ============================================================================ -// 错误类型 -// ============================================================================ - -/// SQLite 账号仓储错误。 -#[derive(Debug, Error)] -pub enum SqliteAccountStoreError { - /// 数据库错误。 - #[error("sqlite account store database error: {0}")] - Database(#[from] sqlx::Error), - /// 时间格式错误。 - #[error("sqlite account store timestamp error: {0}")] - Timestamp(#[from] chrono::ParseError), - /// 账号状态非法。 - #[error("sqlite account store status error: {0}")] - InvalidStatus(String), - /// 分页游标非法。 - #[error("invalid account pagination cursor")] - InvalidCursor, -} - -/// SQLite 账号仓储结果。 -pub type SqliteAccountStoreResult = Result; - -// ============================================================================ -// 数据类型 -// ============================================================================ - -/// 新建账号数据。 -#[derive(Debug)] -pub struct NewAccount { - /// 账号 ID。 - pub id: String, - /// 邮箱。 - pub email: Option, - /// ChatGPT 账号 ID。 - pub account_id: Option, - /// ChatGPT 用户 ID。 - pub user_id: Option, - /// 管理标签。 - pub label: Option, - /// 计划类型。 - pub plan_type: Option, - /// access token 明文。 - pub access_token: SecretString, - /// refresh token 明文。 - pub refresh_token: Option, - /// access token 过期时间。 - pub access_token_expires_at: Option>, - /// 初始状态。 - pub status: AccountStatus, - /// 账号添加时间。 - pub added_at: Option>, -} - -/// 通过 JWT claims 更新已有账号。 -#[derive(Debug)] -pub struct AccountClaimsUpdate { - /// 展示邮箱。 - pub email: Option, - /// ChatGPT 账号 ID。 - pub account_id: Option, - /// ChatGPT 用户 ID。 - pub user_id: Option, - /// 计划类型。 - pub plan_type: Option, - /// 新 access token。 - pub access_token: SecretString, - /// 可选的新 refresh token。 - pub refresh_token: Option, - /// access token 过期时间。 - pub access_token_expires_at: Option>, - /// 下一次刷新时间。 - pub next_refresh_at: Option>, - /// 更新后的状态。 - pub status: AccountStatus, -} - -/// 通过导入数据更新已有账号。 -#[derive(Debug)] -pub struct ImportedAccountUpdate { - /// 新建账号数据。 - pub account: NewAccount, - /// 缓存配额 JSON。 - pub quota_json: Option, - /// 配额抓取时间。 - pub quota_fetched_at: Option>, - /// 是否需要执行额外配额校验。 - pub quota_verify_required: bool, -} - -/// 存储中的账号完整记录。 -#[derive(Debug, Clone)] -pub struct StoredAccount { - /// 账号 ID。 - pub id: String, - /// 邮箱。 - pub email: Option, - /// ChatGPT 账号 ID。 - pub account_id: Option, - /// ChatGPT 用户 ID。 - pub user_id: Option, - /// 管理标签。 - pub label: Option, - /// 计划类型。 - pub plan_type: Option, - /// access token 明文。 - pub access_token: SecretString, - /// refresh token 明文。 - pub refresh_token: Option, - /// access token 过期时间。 - pub access_token_expires_at: Option>, - /// 下一次刷新时间。 - pub next_refresh_at: Option>, - /// 账号状态。 - pub status: AccountStatus, - /// 添加时间。 - pub added_at: String, - /// 更新时间。 - pub updated_at: String, -} - -/// 存储中的账号元数据（不含 token）。 -#[derive(Debug, Clone)] -pub struct StoredAccountMetadata { - /// 账号 ID。 - pub id: String, - /// 邮箱。 - pub email: Option, - /// ChatGPT 账号 ID。 - pub account_id: Option, - /// ChatGPT 用户 ID。 - pub user_id: Option, - /// 管理标签。 - pub label: Option, - /// 计划类型。 - pub plan_type: Option, - /// 是否保存了 refresh token。 - pub has_refresh_token: bool, - /// access token 过期时间。 - pub access_token_expires_at: Option>, - /// 账号状态。 - pub status: AccountStatus, - /// 添加时间。 - pub added_at: String, - /// 更新时间。 - pub updated_at: String, -} - -#[derive(Debug, Clone, Default)] -pub struct AccountMetadataUpdate { - pub email: Option>, - pub account_id: Option>, - pub user_id: Option>, - pub label: Option>, - pub plan_type: Option>, - pub status: Option, -} - -impl AccountMetadataUpdate { - pub fn any(&self) -> bool { - self.email.is_some() - || self.account_id.is_some() - || self.user_id.is_some() - || self.label.is_some() - || self.plan_type.is_some() - || self.status.is_some() - } -} - -/// 账号配额快照。 -#[derive(Debug, Clone)] -pub struct AccountQuotaSnapshot { - /// 账号 ID。 - pub account_id: String, - /// 邮箱。 - pub email: Option, - /// 配额 JSON。 - pub quota_json: String, - /// 配额抓取时间。 - pub quota_fetched_at: Option>, -} - -/// 用量增量（SQLite 内部表示）。 -#[derive(Debug, Clone, Copy, Default)] -pub struct UsageDelta { - request_count: i64, - empty_response_count: i64, - input_tokens: i64, - output_tokens: i64, - cached_tokens: i64, - reasoning_tokens: i64, - total_tokens: i64, - image_input_tokens: i64, - image_output_tokens: i64, - image_request_count: i64, - image_request_failed_count: i64, - window_request_count: i64, - window_input_tokens: i64, - window_output_tokens: i64, - window_cached_tokens: i64, - window_image_input_tokens: i64, - window_image_output_tokens: i64, - window_image_request_count: i64, - window_image_request_failed_count: i64, -} - -// ============================================================================ -// AccountStore 端口 trait -// ============================================================================ - -/// 账号存储错误。 -#[derive(Debug, Error)] -pub enum AccountStoreError { - /// 底层存储失败。 - #[error("account store operation failed: {message}")] - OperationFailed { - /// 错误说明。 - message: String, - }, -} - -/// 账号存储结果类型。 -pub type AccountStoreResult = Result; - -/// 提供运行时账号列表的端口。 -#[async_trait] -pub trait AccountStore: Send + Sync + 'static { - /// 列出当前账号池可见的账号。 - async fn list_pool_accounts(&self) -> AccountStoreResult>; - - /// 读取单个账号池账号快照。 - async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult> { - let accounts = self.list_pool_accounts().await?; - Ok(accounts - .into_iter() - .find(|account| account.id == account_id)) - } - - /// 标记账号进入配额冷却期。 - async fn mark_quota_limited_until( - &self, - account_id: &str, - cooldown_until: DateTime, - ) -> AccountStoreResult; - - /// 标记账号进入 Cloudflare 冷却期。 - async fn set_cloudflare_cooldown_until( - &self, - account_id: &str, - cooldown_until: DateTime, - ) -> AccountStoreResult; - - /// 更新账号状态。 - async fn set_status(&self, account_id: &str, status: AccountStatus) - -> AccountStoreResult; - - /// 记录账号用量增量。 - async fn record_usage_delta( - &self, - account_id: &str, - usage: AccountUsageDelta, - ) -> AccountStoreResult<()>; - - /// 记录账号模型维度用量增量。 - async fn record_model_usage_delta( - &self, - account_id: &str, - model: &str, - usage: AccountModelUsageDelta, - ) -> AccountStoreResult<()>; - - /// 读取账号当前配额 JSON。 - async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult>; - - /// 应用已经验证过的账号配额快照。 - async fn apply_quota_snapshot( - &self, - account_id: &str, - quota_json: &str, - limit_reached: bool, - cooldown_until: Option>, - ) -> AccountStoreResult; - - /// 同步运行时自然刷新出来的账号状态。 - async fn sync_runtime_account_state( - &self, - account: &Account, - sync_usage_window: bool, - ) -> AccountStoreResult; - - /// 同步账号当前 rate-limit 统计窗口。 - async fn sync_rate_limit_window( - &self, - account_id: &str, - reset_at: DateTime, - limit_window_seconds: Option, - ) -> AccountStoreResult<()>; - - /// 记录账号被用于一次外部请求。 - async fn record_request(&self, account_id: &str) -> AccountStoreResult<()> { - self.record_usage_delta( - account_id, - AccountUsageDelta { - requests: 1, - ..AccountUsageDelta::default() - }, - ) - .await - } -} - -#[derive(Clone)] -pub struct SqliteAccountStore { - pool: SqlitePool, -} - -impl SqliteAccountStore { - /// 构造存储。 - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - /// 返回底层连接池。 - pub fn pool(&self) -> &SqlitePool { - &self.pool - } - - /// 插入新账号。 - pub async fn insert(&self, account: NewAccount) -> SqliteAccountStoreResult<()> { - let now = Utc::now().to_rfc3339(); - let refresh_token = account - .refresh_token - .as_ref() - .map(ExposeSecret::expose_secret); - sqlx::query(INSERT_ACCOUNT_SQL) - .bind(&account.id) - .bind(&account.email) - .bind(&account.account_id) - .bind(&account.user_id) - .bind(&account.label) - .bind(&account.plan_type) - .bind(account.access_token.expose_secret()) - .bind(refresh_token) - .bind(account.access_token_expires_at.map(|dt| dt.to_rfc3339())) - .bind(status_to_db(account.status)) - .bind( - account - .added_at - .map_or_else(|| Utc::now().to_rfc3339(), |dt| dt.to_rfc3339()), - ) - .bind(&now) - .execute(&self.pool) - .await?; - Ok(()) - } - - /// 读取单个账号。 - pub async fn get(&self, account_id: &str) -> SqliteAccountStoreResult> { - let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_ID_SQL) - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - row.as_ref().map(stored_account_from_row).transpose() - } - - /// 通过 ChatGPT 身份查找账号。 - pub async fn find_by_chatgpt_identity( - &self, - chatgpt_account_id: &str, - chatgpt_user_id: Option<&str>, - ) -> SqliteAccountStoreResult> { - let row = sqlx::query(SELECT_STORED_ACCOUNT_BY_CHATGPT_IDENTITY_SQL) - .bind(chatgpt_account_id) - .bind(chatgpt_user_id) - .bind(chatgpt_user_id) - .fetch_optional(&self.pool) - .await?; - row.as_ref().map(stored_account_from_row).transpose() - } - - /// 分页列出所有账号（含 token）。 - pub async fn list( - &self, - cursor: Option, - limit: u32, - ) -> SqliteAccountStoreResult> { - let limit = limit.clamp(1, 200); - if let Some(cursor) = cursor { - let (added_at, id) = - decode_cursor(&cursor).ok_or(SqliteAccountStoreError::InvalidCursor)?; - let rows = sqlx::query(LIST_STORED_ACCOUNTS_AFTER_CURSOR_SQL) - .bind(&added_at) - .bind(&added_at) - .bind(&id) - .bind(limit + 1) - .fetch_all(&self.pool) - .await?; - Ok(to_page( - &rows, - limit, - stored_account_from_row, - ("added_at", "id"), - )) - } else { - let rows = sqlx::query(LIST_STORED_ACCOUNTS_SQL) - .bind(limit + 1) - .fetch_all(&self.pool) - .await?; - Ok(to_page( - &rows, - limit, - stored_account_from_row, - ("added_at", "id"), - )) - } - } - - /// 分页列出账号元数据（不含 token）。 - pub async fn list_metadata( - &self, - cursor: Option, - limit: u32, - ) -> SqliteAccountStoreResult> { - let limit = limit.clamp(1, 200); - if let Some(cursor) = cursor { - let (added_at, id) = - decode_cursor(&cursor).ok_or(SqliteAccountStoreError::InvalidCursor)?; - let rows = sqlx::query(LIST_ACCOUNT_METADATA_AFTER_CURSOR_SQL) - .bind(&added_at) - .bind(&added_at) - .bind(&id) - .bind(limit + 1) - .fetch_all(&self.pool) - .await?; - Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) - } else { - let rows = sqlx::query(LIST_ACCOUNT_METADATA_SQL) - .bind(limit + 1) - .fetch_all(&self.pool) - .await?; - Ok(to_page(&rows, limit, metadata_from_row, ("added_at", "id"))) - } - } - - /// 按页码列出账号元数据（不含 token）。 - pub async fn list_metadata_page( - &self, - page: u32, - page_size: u32, - search: Option<&str>, - ) -> SqliteAccountStoreResult> { - 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 offset = page_offset(page, page_size); - - let mut builder = QueryBuilder::::new(LIST_ACCOUNT_METADATA_SELECT_SQL); - push_account_metadata_search(&mut builder, search); - builder.push(" order by added_at desc, id desc limit "); - builder.push_bind(i64::from(page_size)); - builder.push(" offset "); - builder.push_bind(offset.min(i64::MAX as u64) as i64); - - let rows = builder.build().fetch_all(&self.pool).await?; - let items = rows - .iter() - .map(metadata_from_row) - .collect::>>()?; - - Ok(NumberedPage { - items, - total, - page: page.max(1), - page_size, - }) - } - - /// 读取单个账号元数据（不含 token）。 - pub async fn get_metadata( - &self, - account_id: &str, - ) -> SqliteAccountStoreResult> { - let row = sqlx::query(SELECT_ACCOUNT_METADATA_BY_ID_SQL) - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - row.as_ref().map(metadata_from_row).transpose() - } - - /// 更新单账号元数据（不含 token）。 - pub async fn update_metadata( - &self, - account_id: &str, - update: AccountMetadataUpdate, - ) -> SqliteAccountStoreResult { - if !update.any() { - return Ok(false); - } - let now = Utc::now().to_rfc3339(); - let status = update.status.map(status_to_db); - let result = sqlx::query(UPDATE_ACCOUNT_METADATA_SQL) - .bind(update.email.is_some()) - .bind(optional_update_value(&update.email)) - .bind(update.account_id.is_some()) - .bind(optional_update_value(&update.account_id)) - .bind(update.user_id.is_some()) - .bind(optional_update_value(&update.user_id)) - .bind(update.label.is_some()) - .bind(optional_update_value(&update.label)) - .bind(update.plan_type.is_some()) - .bind(optional_update_value(&update.plan_type)) - .bind(update.status.is_some()) - .bind(status) - .bind(status) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 更新账号 claims（含 refresh token）。 - pub async fn update_from_claims( - &self, - account_id: &str, - update: AccountClaimsUpdate, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let refresh_token = update - .refresh_token - .as_ref() - .map(ExposeSecret::expose_secret); - - let result = if let Some(refresh_token) = refresh_token { - sqlx::query(UPDATE_ACCOUNT_CLAIMS_WITH_REFRESH_SQL) - .bind(&update.email) - .bind(&update.account_id) - .bind(&update.user_id) - .bind(&update.plan_type) - .bind(update.access_token.expose_secret()) - .bind(refresh_token) - .bind(update.access_token_expires_at.map(|dt| dt.to_rfc3339())) - .bind(update.next_refresh_at.map(|dt| dt.to_rfc3339())) - .bind(status_to_db(update.status)) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await? - } else { - sqlx::query(UPDATE_ACCOUNT_CLAIMS_PRESERVING_REFRESH_SQL) - .bind(&update.email) - .bind(&update.account_id) - .bind(&update.user_id) - .bind(&update.plan_type) - .bind(update.access_token.expose_secret()) - .bind(update.access_token_expires_at.map(|dt| dt.to_rfc3339())) - .bind(update.next_refresh_at.map(|dt| dt.to_rfc3339())) - .bind(status_to_db(update.status)) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await? - }; - Ok(result.rows_affected() > 0) - } - - /// 通过导入数据更新已有账号。 - pub async fn update_from_import( - &self, - update: ImportedAccountUpdate, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let refresh_token = update - .account - .refresh_token - .as_ref() - .map(ExposeSecret::expose_secret); - let quota_json = update.quota_json; - let quota_fetched_at = update.quota_fetched_at.map(|dt| dt.to_rfc3339()); - let quota_verify_required = if update.quota_verify_required { 1 } else { 0 }; - - let result = if let Some(refresh_token) = refresh_token { - sqlx::query(UPDATE_IMPORTED_ACCOUNT_WITH_REFRESH_SQL) - .bind(&update.account.email) - .bind(&update.account.account_id) - .bind(&update.account.user_id) - .bind(&update.account.label) - .bind(&update.account.plan_type) - .bind(update.account.access_token.expose_secret()) - .bind(refresh_token) - .bind( - update - .account - .access_token_expires_at - .map(|dt| dt.to_rfc3339()), - ) - .bind(status_to_db(update.account.status)) - .bind("a_json) - .bind("a_fetched_at) - .bind("a_fetched_at) - .bind(quota_verify_required) - .bind(&now) - .bind(&update.account.id) - .execute(&self.pool) - .await? - } else { - sqlx::query(UPDATE_IMPORTED_ACCOUNT_PRESERVING_REFRESH_SQL) - .bind(&update.account.email) - .bind(&update.account.account_id) - .bind(&update.account.user_id) - .bind(&update.account.label) - .bind(&update.account.plan_type) - .bind(update.account.access_token.expose_secret()) - .bind( - update - .account - .access_token_expires_at - .map(|dt| dt.to_rfc3339()), - ) - .bind(status_to_db(update.account.status)) - .bind("a_json) - .bind("a_fetched_at) - .bind("a_fetched_at) - .bind(quota_verify_required) - .bind(&now) - .bind(&update.account.id) - .execute(&self.pool) - .await? - }; - Ok(result.rows_affected() > 0) - } - - /// 设置下一次刷新时间。 - pub async fn set_next_refresh_at( - &self, - account_id: &str, - next_refresh_at: Option>, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let result = sqlx::query(SET_NEXT_REFRESH_AT_SQL) - .bind(next_refresh_at.map(|dt| dt.to_rfc3339())) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 标记账号进入配额冷却期。 - pub async fn mark_quota_limited_until( - &self, - account_id: &str, - cooldown_until: DateTime, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let result = sqlx::query(MARK_QUOTA_LIMITED_UNTIL_SQL) - .bind(cooldown_until.to_rfc3339()) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 同步运行时自然刷新出来的账号状态。 - pub async fn sync_runtime_account_state( - &self, - account: &Account, - sync_usage_window: bool, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let quota_limit_reached = if account.quota_limit_reached { 1 } else { 0 }; - let quota_verify_required = if account.quota_verify_required { 1 } else { 0 }; - let quota_cooldown_until = account.quota_cooldown_until.map(|dt| dt.to_rfc3339()); - let cloudflare_cooldown_until = account.cloudflare_cooldown_until.map(|dt| dt.to_rfc3339()); - let mut tx = self.pool.begin().await?; - let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL) - .bind(quota_limit_reached) - .bind(&now) - .bind(quota_limit_reached) - .bind(status_to_db(account.status)) - .bind(quota_limit_reached) - .bind(&now) - .bind(quota_limit_reached) - .bind(quota_verify_required) - .bind(&now) - .bind(quota_verify_required) - .bind("a_cooldown_until) - .bind(&now) - .bind("a_cooldown_until) - .bind(&cloudflare_cooldown_until) - .bind(&now) - .bind(&cloudflare_cooldown_until) - .bind(&now) - .bind(&account.id) - .execute(&mut *tx) - .await?; - if result.rows_affected() == 0 { - tx.commit().await?; - return Ok(false); - } - - if sync_usage_window { - sqlx::query(SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL) - .bind(&account.id) - .bind(u64_to_i64_saturating(account.window_request_count)) - .bind(u64_to_i64_saturating(account.window_input_tokens)) - .bind(u64_to_i64_saturating(account.window_output_tokens)) - .bind(u64_to_i64_saturating(account.window_cached_tokens)) - .bind(u64_to_i64_saturating(account.window_image_input_tokens)) - .bind(u64_to_i64_saturating(account.window_image_output_tokens)) - .bind(u64_to_i64_saturating(account.window_image_request_count)) - .bind(u64_to_i64_saturating( - account.window_image_request_failed_count, - )) - .bind(account.window_started_at.map(|dt| dt.to_rfc3339())) - .bind(account.window_reset_at.map(|dt| dt.to_rfc3339())) - .bind(account.limit_window_seconds.map(u64_to_i64_saturating)) - .bind(&now) - .execute(&mut *tx) - .await?; - } - - tx.commit().await?; - Ok(true) - } - - /// 标记账号 Cloudflare 冷却期。 - pub async fn set_cloudflare_cooldown_until( - &self, - account_id: &str, - cooldown_until: DateTime, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let result = sqlx::query(SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL) - .bind(cooldown_until.to_rfc3339()) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 更新账号状态。 - pub async fn set_status( - &self, - account_id: &str, - status: AccountStatus, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let status = status_to_db(status); - let result = sqlx::query( - r" -update accounts -set - status = case - when ? = 'active' and quota_limit_reached = 1 then 'quota_exhausted' - else ? - end, - updated_at = ? -where id = ?", - ) - .bind(status) - .bind(status) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 记录用量。 - pub async fn record_usage( - &self, - account_id: &str, - delta: UsageDelta, - ) -> SqliteAccountStoreResult<()> { - let last_used_at = Utc::now().to_rfc3339(); - sqlx::query(RECORD_USAGE_SQL) - .bind(account_id) - .bind(delta.request_count) - .bind(delta.empty_response_count) - .bind(delta.input_tokens) - .bind(delta.output_tokens) - .bind(delta.cached_tokens) - .bind(delta.reasoning_tokens) - .bind(delta.total_tokens) - .bind(delta.image_input_tokens) - .bind(delta.image_output_tokens) - .bind(delta.image_request_count) - .bind(delta.image_request_failed_count) - .bind(delta.window_request_count) - .bind(delta.window_input_tokens) - .bind(delta.window_output_tokens) - .bind(delta.window_cached_tokens) - .bind(delta.window_image_input_tokens) - .bind(delta.window_image_output_tokens) - .bind(delta.window_image_request_count) - .bind(delta.window_image_request_failed_count) - .bind(&last_used_at) - .bind(&last_used_at) - .execute(&self.pool) - .await?; - Ok(()) - } - - /// 记录模型维度用量。 - pub async fn record_model_usage( - &self, - account_id: &str, - model: &str, - delta: AccountModelUsageDelta, - ) -> SqliteAccountStoreResult<()> { - let model = model.trim(); - if model.is_empty() { - return Ok(()); - } - let last_used_at = Utc::now().to_rfc3339(); - sqlx::query(RECORD_MODEL_USAGE_SQL) - .bind(account_id) - .bind(model) - .bind(u64_to_i64_saturating(delta.requests)) - .bind(u64_to_i64_saturating(delta.errors)) - .bind(u64_to_i64_saturating(delta.input_tokens)) - .bind(u64_to_i64_saturating(delta.output_tokens)) - .bind(u64_to_i64_saturating(delta.cached_tokens)) - .bind(&last_used_at) - .execute(&self.pool) - .await?; - Ok(()) - } - - /// 读取配额快照列表。 - pub async fn list_quota_snapshots( - &self, - ) -> SqliteAccountStoreResult> { - let rows = sqlx::query(LIST_QUOTA_SNAPSHOTS_SQL) - .fetch_all(&self.pool) - .await?; - rows.iter().map(quota_snapshot_from_row).collect() - } - - /// 读取单账号配额 JSON。 - pub async fn get_quota_json( - &self, - account_id: &str, - ) -> SqliteAccountStoreResult> { - let row = sqlx::query("select quota_json from accounts where id = ?") - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - Ok(row.and_then(|row| row.get::, _>("quota_json"))) - } - - /// 更新配额 JSON。 - pub async fn update_quota_json( - &self, - account_id: &str, - quota_json: &str, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let plan_type = quota_plan_type(quota_json); - let result = sqlx::query(UPDATE_QUOTA_JSON_SQL) - .bind(quota_json) - .bind(&now) - .bind(&plan_type) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 应用配额快照。 - pub async fn apply_quota_snapshot( - &self, - account_id: &str, - quota_json: &str, - limit_reached: bool, - cooldown_until: Option>, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let cooldown = cooldown_until.map(|dt| dt.to_rfc3339()); - let plan_type = quota_plan_type(quota_json); - let result = sqlx::query(APPLY_QUOTA_SNAPSHOT_SQL) - .bind(quota_json) - .bind(&now) - .bind(&plan_type) - .bind(if limit_reached { 1 } else { 0 }) - .bind(if limit_reached { 1 } else { 0 }) - .bind(&cooldown) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 应用导入的配额状态。 - pub async fn apply_imported_quota_state( - &self, - account_id: &str, - quota_json: Option<&str>, - quota_fetched_at: Option>, - quota_verify_required: bool, - ) -> SqliteAccountStoreResult { - let now = Utc::now().to_rfc3339(); - let fetched = quota_fetched_at.map(|dt| dt.to_rfc3339()); - let plan_type = quota_json.and_then(quota_plan_type); - let result = sqlx::query(APPLY_IMPORTED_QUOTA_STATE_SQL) - .bind(quota_json) - .bind(&fetched) - .bind(&fetched) - .bind(&plan_type) - .bind(if quota_verify_required { 1 } else { 0 }) - .bind(&now) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } - - /// 同步 rate-limit 窗口（含重置）。 - pub async fn sync_rate_limit_window( - &self, - account_id: &str, - reset_at: DateTime, - limit_window_seconds: Option, - ) -> SqliteAccountStoreResult<()> { - let existing = sqlx::query(SELECT_RATE_LIMIT_WINDOW_SQL) - .bind(account_id) - .fetch_optional(&self.pool) - .await?; - - let should_reset = existing - .as_ref() - .map(|row| { - let existing_reset_at = parse_optional_rfc3339( - row.get::, _>("window_reset_at").as_deref(), - ) - .ok() - .flatten(); - let existing_limit_window_seconds = - optional_positive_i64_to_u64(row.get::, _>("limit_window_seconds")); - should_reset_usage_window( - existing_reset_at, - existing_limit_window_seconds, - reset_at, - limit_window_seconds, - ) - }) - .unwrap_or_default(); - - if should_reset { - let reset_at_str = reset_at.to_rfc3339(); - let window_started_at = Utc::now().to_rfc3339(); - sqlx::query(SYNC_RATE_LIMIT_WINDOW_RESET_SQL) - .bind(account_id) - .bind(&window_started_at) - .bind(&reset_at_str) - .bind(limit_window_seconds.map(|v| v as i64)) - .execute(&self.pool) - .await?; - } else { - sqlx::query(SYNC_RATE_LIMIT_WINDOW_SQL) - .bind(account_id) - .bind(reset_at.to_rfc3339()) - .bind(limit_window_seconds.map(|v| v as i64)) - .execute(&self.pool) - .await?; - } - Ok(()) - } - - /// 删除单个账号。 - pub async fn delete(&self, account_id: &str) -> SqliteAccountStoreResult { - let result = sqlx::query(DELETE_ACCOUNT_SQL) - .bind(account_id) - .execute(&self.pool) - .await?; - Ok(result.rows_affected() > 0) - } -} - -impl SqliteAccountStore { - /// 设置账号标签。 - pub async fn set_label( - &self, - account_id: &str, - label: Option, - ) -> Result { - let updated_at = Utc::now().to_rfc3339(); - let result = sqlx::query("UPDATE accounts SET label = ?1, updated_at = ?2 WHERE id = ?3") - .bind(&label) - .bind(updated_at) - .bind(account_id) - .execute(&self.pool) - .await - .map_err(SqliteAccountStoreError::Database)?; - Ok(result.rows_affected() > 0) - } -} - -#[async_trait] -impl AccountStore for SqliteAccountStore { - async fn list_pool_accounts(&self) -> AccountStoreResult> { - list_pool_accounts(self) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult> { - get_pool_account(self, account_id) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn mark_quota_limited_until( - &self, - account_id: &str, - cooldown_until: DateTime, - ) -> AccountStoreResult { - SqliteAccountStore::mark_quota_limited_until(self, account_id, cooldown_until) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn set_cloudflare_cooldown_until( - &self, - account_id: &str, - cooldown_until: DateTime, - ) -> AccountStoreResult { - SqliteAccountStore::set_cloudflare_cooldown_until(self, account_id, cooldown_until) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn set_status( - &self, - account_id: &str, - status: AccountStatus, - ) -> AccountStoreResult { - SqliteAccountStore::set_status(self, account_id, status) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn record_usage_delta( - &self, - account_id: &str, - usage: AccountUsageDelta, - ) -> AccountStoreResult<()> { - self.record_usage(account_id, sqlite_usage_delta(usage)) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn record_model_usage_delta( - &self, - account_id: &str, - model: &str, - usage: AccountModelUsageDelta, - ) -> AccountStoreResult<()> { - self.record_model_usage(account_id, model, usage) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult> { - SqliteAccountStore::get_quota_json(self, account_id) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn apply_quota_snapshot( - &self, - account_id: &str, - quota_json: &str, - limit_reached: bool, - cooldown_until: Option>, - ) -> AccountStoreResult { - SqliteAccountStore::apply_quota_snapshot( - self, - account_id, - quota_json, - limit_reached, - cooldown_until, - ) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn sync_runtime_account_state( - &self, - account: &Account, - sync_usage_window: bool, - ) -> AccountStoreResult { - SqliteAccountStore::sync_runtime_account_state(self, account, sync_usage_window) - .await - .map_err(|error| map_account_store_error(&error)) - } - - async fn sync_rate_limit_window( - &self, - account_id: &str, - reset_at: DateTime, - limit_window_seconds: Option, - ) -> AccountStoreResult<()> { - SqliteAccountStore::sync_rate_limit_window(self, account_id, reset_at, limit_window_seconds) - .await - .map_err(|error| map_account_store_error(&error)) - } -} diff --git a/backend/src/upstream/fingerprint.rs b/backend/src/upstream/fingerprint.rs deleted file mode 100644 index de9dd0b62..000000000 --- a/backend/src/upstream/fingerprint.rs +++ /dev/null @@ -1,576 +0,0 @@ -//! 指纹类型、更新与持久化。 - -use std::{ - path::PathBuf, - sync::{Arc, RwLock}, - time::Duration, -}; - -use chrono::{DateTime, Utc}; -use indexmap::IndexMap; -use serde::{Deserialize, Serialize}; -use sqlx::{Row, SqlitePool}; -use thiserror::Error; -use tokio::sync::Mutex; -use tracing::info; -use uuid::Uuid; - -/// 运行时当前指纹槽位。 -pub const CURRENT_FINGERPRINT_ID: &str = "current"; - -const AUTO_UPDATE_SOURCE: &str = "auto_update"; -const CONFIG_SEED_SOURCE: &str = "config_seed"; -const APPCAST_TIMEOUT: Duration = Duration::from_secs(30); -pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_hours(72); - -type AppcastFields = (Option, Option, Option); - -// --------------------------------------------------------------------------- -// Fingerprint 数据类型 -// --------------------------------------------------------------------------- - -/// 上游请求指纹。 -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub struct Fingerprint { - /// 客户端来源名。 - pub originator: String, - /// 应用版本。 - pub app_version: String, - /// 构建号。 - pub build_number: String, - /// 平台名。 - pub platform: String, - /// 架构名。 - pub arch: String, - /// Chromium 主版本。 - pub chromium_version: String, - /// User-Agent 模板。 - pub user_agent_template: String, - /// 默认请求头。 - pub default_headers: IndexMap, - /// 请求头顺序。 - pub header_order: Vec, - /// DB 最后更新时间。 - #[serde(skip_serializing_if = "Option::is_none")] - pub updated_at: Option, -} - -impl Fingerprint { - /// 从配置构造指纹。 - pub fn from_config(config: &crate::config::schema::FingerprintConfig) -> Self { - Self { - originator: config.originator.clone(), - app_version: config.app_version.clone(), - build_number: config.build_number.clone(), - platform: config.platform.clone(), - arch: config.arch.clone(), - chromium_version: config.chromium_version.clone(), - user_agent_template: config.user_agent_template.clone(), - default_headers: config - .default_headers - .iter() - .map(|h| (h.name.clone(), h.value.clone())) - .collect(), - header_order: config.header_order.clone(), - updated_at: None, - } - } - - /// 根据模板展开最终 User-Agent。 - pub fn user_agent(&self) -> String { - self.user_agent_template - .replace("{version}", &self.app_version) - .replace("{platform}", &self.platform) - .replace("{arch}", &self.arch) - } - - /// 生成 `sec-ch-ua` 头值。 - pub fn sec_ch_ua(&self) -> String { - format!( - "\"Chromium\";v=\"{}\", \"Not:A-Brand\";v=\"24\"", - self.chromium_version - ) - } -} - -/// 运行时共享指纹快照。 -#[derive(Debug, Clone)] -pub struct RuntimeFingerprint { - current: Arc>>, -} - -impl RuntimeFingerprint { - pub fn new(fingerprint: Fingerprint) -> Self { - Self { - current: Arc::new(RwLock::new(Arc::new(fingerprint))), - } - } - - pub fn current(&self) -> Arc { - self.current - .read() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone() - } - - pub fn snapshot(&self) -> Fingerprint { - self.current().as_ref().clone() - } - - pub fn replace(&self, fingerprint: Fingerprint) { - *self - .current - .write() - .unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::new(fingerprint); - } -} - -// --------------------------------------------------------------------------- -// 指纹历史记录与仓储 -// --------------------------------------------------------------------------- - -#[derive(Debug, Clone, Serialize, Deserialize)] -struct StoredHeader { - name: String, - value: String, -} - -/// 指纹自动更新状态。 -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct UpdateState { - /// 最近检查时间。 - pub last_check: DateTime, - /// 最新版本。 - pub latest_version: Option, - /// 最新构建号。 - pub latest_build: Option, - /// 下载地址。 - pub download_url: Option, - /// 是否有可用更新。 - pub update_available: bool, - /// 当前版本。 - pub current_version: String, - /// 当前构建号。 - pub current_build: String, -} - -/// Appcast 检查错误。 -#[derive(Debug, Error)] -pub enum UpdateError { - /// HTTP 请求失败。 - #[error("HTTP 请求失败: {0}")] - Http(#[from] reqwest::Error), - /// 获取 appcast 失败。 - #[error("获取 appcast 失败，状态码: {0}")] - AppcastFetch(u16), - /// 解析 appcast 失败。 - #[error("解析 appcast 失败")] - AppcastParse, - /// JSON 序列化失败。 - #[error("JSON 序列化失败: {0}")] - Json(#[from] serde_json::Error), - /// 文件操作失败。 - #[error("文件操作失败: {0}")] - Io(#[from] std::io::Error), - /// 数据库存储失败。 - #[error("数据库操作失败: {0}")] - Database(#[from] sqlx::Error), -} - -/// SQLite 指纹仓储。 -#[derive(Clone)] -pub struct FingerprintRepository { - pool: SqlitePool, -} - -impl FingerprintRepository { - /// 使用给定连接池构造仓储。 - pub fn new(pool: SqlitePool) -> Self { - Self { pool } - } - - /// 写入当前指纹默认值；如果当前槽位已存在，直接读取数据库值。 - pub async fn ensure_current_seed( - &self, - default_fingerprint: &Fingerprint, - ) -> Result { - if let Some(fingerprint) = self.load_current().await? { - return Ok(fingerprint); - } - - let now = Utc::now().to_rfc3339(); - sqlx::query( - r" - 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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - ", - ) - .bind(CURRENT_FINGERPRINT_ID) - .bind(&default_fingerprint.originator) - .bind(&default_fingerprint.app_version) - .bind(&default_fingerprint.build_number) - .bind(&default_fingerprint.platform) - .bind(&default_fingerprint.arch) - .bind(&default_fingerprint.chromium_version) - .bind(&default_fingerprint.user_agent_template) - .bind(encode_default_headers( - &default_fingerprint.default_headers, - )?) - .bind(encode_header_order(&default_fingerprint.header_order)?) - .bind(CONFIG_SEED_SOURCE) - .bind(&now) - .bind(now) - .execute(&self.pool) - .await?; - - Ok(default_fingerprint.clone()) - } - - /// 更新当前指纹中的自动更新版本字段。 - pub async fn update_current_version( - &self, - app_version: &str, - build_number: &str, - chromium_version: Option<&str>, - ) -> Result<(), sqlx::Error> { - let resolved_chromium_version = match chromium_version { - Some(version) => version.to_string(), - None => { - sqlx::query_scalar::<_, String>( - "select chromium_version from fingerprints where id = ?", - ) - .bind(CURRENT_FINGERPRINT_ID) - .fetch_one(&self.pool) - .await? - } - }; - - let result = sqlx::query( - "update fingerprints set app_version = ?, build_number = ?, chromium_version = ?, source = ?, updated_at = ? where id = ?", - ) - .bind(app_version) - .bind(build_number) - .bind(resolved_chromium_version) - .bind(AUTO_UPDATE_SOURCE) - .bind(Utc::now().to_rfc3339()) - .bind(CURRENT_FINGERPRINT_ID) - .execute(&self.pool) - .await?; - if result.rows_affected() == 0 { - return Err(sqlx::Error::RowNotFound); - } - Ok(()) - } - - /// 读取当前运行时指纹。 - pub async fn load_current(&self) -> Result, sqlx::Error> { - let row = sqlx::query( - r" - select - originator, - app_version, - build_number, - platform, - arch, - chromium_version, - user_agent_template, - default_headers_json, - header_order_json, - updated_at - from fingerprints - where id = ? - ", - ) - .bind(CURRENT_FINGERPRINT_ID) - .fetch_optional(&self.pool) - .await?; - - row.map(|row| fingerprint_from_row(&row)).transpose() - } - - /// 记录自动更新历史。 - pub async fn insert_update_history( - &self, - app_version: &str, - build_number: &str, - chromium_version: Option<&str>, - manifest_json: Option<&str>, - ) -> Result<(), sqlx::Error> { - sqlx::query( - "insert into fingerprint_update_history (id, current_fingerprint_id, app_version, build_number, chromium_version, source, manifest_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?)", - ) - .bind(Uuid::new_v4().to_string()) - .bind(CURRENT_FINGERPRINT_ID) - .bind(app_version) - .bind(build_number) - .bind(chromium_version) - .bind(AUTO_UPDATE_SOURCE) - .bind(manifest_json) - .bind(Utc::now().to_rfc3339()) - .execute(&self.pool) - .await?; - Ok(()) - } -} - -fn fingerprint_from_row(row: &sqlx::sqlite::SqliteRow) -> Result { - let default_headers_json = row.get::("default_headers_json"); - let header_order_json = row.get::("header_order_json"); - Ok(Fingerprint { - originator: row.get("originator"), - app_version: row.get("app_version"), - build_number: row.get("build_number"), - platform: row.get("platform"), - arch: row.get("arch"), - chromium_version: row.get("chromium_version"), - user_agent_template: row.get("user_agent_template"), - default_headers: decode_default_headers(&default_headers_json)?, - header_order: decode_header_order(&header_order_json)?, - updated_at: row.get("updated_at"), - }) -} - -fn encode_default_headers(headers: &IndexMap) -> Result { - let headers = headers - .iter() - .map(|(name, value)| StoredHeader { - name: name.clone(), - value: value.clone(), - }) - .collect::>(); - serde_json::to_string(&headers).map_err(json_error) -} - -fn decode_default_headers(value: &str) -> Result, sqlx::Error> { - let headers = serde_json::from_str::>(value).map_err(json_error)?; - Ok(headers - .into_iter() - .map(|header| (header.name, header.value)) - .collect()) -} - -fn encode_header_order(header_order: &[String]) -> Result { - serde_json::to_string(header_order).map_err(json_error) -} - -fn decode_header_order(value: &str) -> Result, sqlx::Error> { - serde_json::from_str(value).map_err(json_error) -} - -fn json_error(error: serde_json::Error) -> sqlx::Error { - sqlx::Error::Decode(Box::new(error)) -} - -// --------------------------------------------------------------------------- -// Appcast 指纹更新检查器 -// --------------------------------------------------------------------------- - -/// Appcast 指纹更新检查器。 -#[derive(Clone)] -pub struct UpdateChecker { - repository: FingerprintRepository, - client: reqwest::Client, - appcast_url: String, - extracted_fingerprint_path: PathBuf, - state: Arc>, -} - -struct InternalState { - current_version: String, - current_build: String, -} - -impl UpdateChecker { - /// 使用自定义 client 和路径构造检查器。 - pub fn with_client( - repository: FingerprintRepository, - client: reqwest::Client, - appcast_url: impl Into, - extracted_fingerprint_path: PathBuf, - current_version: impl Into, - current_build: impl Into, - ) -> Self { - Self { - repository, - client, - appcast_url: appcast_url.into(), - extracted_fingerprint_path, - state: Arc::new(Mutex::new(InternalState { - current_version: current_version.into(), - current_build: current_build.into(), - })), - } - } - - /// 检查 appcast 是否存在更新。 - pub async fn check_for_update(&self) -> Result { - let state = self.state.lock().await; - let current_version = state.current_version.clone(); - let current_build = state.current_build.clone(); - drop(state); - - let xml = self.fetch_appcast().await?; - let (version, build, download_url) = parse_appcast(&xml)?; - - let update_available = version - .as_ref() - .is_some_and(|value| value != ¤t_version) - || build.as_ref().is_some_and(|value| value != ¤t_build); - - Ok(UpdateState { - last_check: Utc::now(), - latest_version: version, - latest_build: build, - download_url, - update_available, - current_version, - current_build, - }) - } - - /// 检查并在需要时应用更新。 - pub async fn check_and_apply_update(&self) -> Result, UpdateError> { - let update_state = self.check_for_update().await?; - - if !update_state.update_available { - return Ok(None); - } - - let (Some(version), Some(build)) = ( - update_state.latest_version.as_deref(), - update_state.latest_build.as_deref(), - ) else { - return Ok(None); - }; - - let state = self.state.lock().await; - let current_version = state.current_version.clone(); - let current_build = state.current_build.clone(); - drop(state); - - info!( - version = %version, - build = %build, - current_version = %current_version, - current_build = %current_build, - "发现新的 fingerprint 版本" - ); - - let updated_fingerprint = self.apply_version_update(version, build).await?; - - let mut state = self.state.lock().await; - state.current_version = version.to_string(); - state.current_build = build.to_string(); - - Ok(Some(updated_fingerprint)) - } - - async fn fetch_appcast(&self) -> Result { - let response = self - .client - .get(&self.appcast_url) - .timeout(APPCAST_TIMEOUT) - .send() - .await?; - - if !response.status().is_success() { - return Err(UpdateError::AppcastFetch(response.status().as_u16())); - } - - response.text().await.map_err(UpdateError::from) - } - - async fn apply_version_update( - &self, - version: &str, - build: &str, - ) -> Result { - let chromium_version = - load_matching_chromium_version(&self.extracted_fingerprint_path, version, build); - - self.repository - .update_current_version(version, build, chromium_version.as_deref()) - .await?; - self.repository - .insert_update_history(version, build, chromium_version.as_deref(), None) - .await?; - let updated = self - .repository - .load_current() - .await? - .ok_or(sqlx::Error::RowNotFound)?; - Ok(updated) - } -} - -fn parse_appcast(xml: &str) -> Result { - let item_start = xml.find("").ok_or(UpdateError::AppcastParse)?; - let item_end = xml.find("").ok_or(UpdateError::AppcastParse)?; - let item = &xml[item_start..item_end]; - - let version = extract_sparkle_field(item, "shortVersionString"); - let build = extract_sparkle_field(item, "version"); - let download_url = item.find("url=\"").and_then(|position| { - let start = position + 5; - item[start..] - .find('"') - .map(|end| item[start..start + end].to_string()) - }); - - Ok((version, build, download_url)) -} - -fn extract_sparkle_field(item: &str, field: &str) -> Option { - let attr_pattern = format!("sparkle:{}=\"", field); - if let Some(position) = item.find(&attr_pattern) { - let start = position + attr_pattern.len(); - if let Some(end) = item[start..].find('"') { - return Some(item[start..start + end].to_string()); - } - } - - let elem_start = format!("", field); - let elem_end = format!("", field); - if let Some(start_position) = item.find(&elem_start) { - let content_start = start_position + elem_start.len(); - if let Some(end_position) = item[content_start..].find(&elem_end) { - return Some(item[content_start..content_start + end_position].to_string()); - } - } - - None -} - -fn load_matching_chromium_version( - extracted_fingerprint_path: &PathBuf, - version: &str, - build: &str, -) -> Option { - let content = std::fs::read_to_string(extracted_fingerprint_path).ok()?; - let parsed: serde_json::Value = serde_json::from_str(&content).ok()?; - - let extracted_version = parsed.get("app_version")?.as_str()?; - let extracted_build = parsed.get("build_number")?.as_str()?; - - if extracted_version == version && extracted_build == build { - parsed - .get("chromium_version")? - .as_str() - .map(ToString::to_string) - } else { - None - } -} diff --git a/backend/src/upstream/mod.rs b/backend/src/upstream/mod.rs index 35c99ca09..6cfd42a80 100644 --- a/backend/src/upstream/mod.rs +++ b/backend/src/upstream/mod.rs @@ -1,9 +1,3 @@ -//! ChatGPT/Codex 上游资源、协议与传输。 +//! 纯上游客户端。 -pub mod accounts; -pub mod fingerprint; -pub mod models; -pub mod protocol; -pub mod scheduler; -pub mod token_client; -pub mod transport; +pub mod openai; diff --git a/backend/src/upstream/models/snapshot_store.rs b/backend/src/upstream/models/snapshot_store.rs deleted file mode 100644 index 045794178..000000000 --- a/backend/src/upstream/models/snapshot_store.rs +++ /dev/null @@ -1,95 +0,0 @@ -//! 模型计划快照存储。 - -use async_trait::async_trait; -use sqlx::Row; -use thiserror::Error; - -use super::{info::CodexModelInfo, snapshot::ModelPlanSnapshot}; - -/// 模型快照存储错误。 -#[derive(Debug, Error, PartialEq, Eq)] -pub enum ModelSnapshotStoreError { - /// 底层存储失败。 - #[error("model snapshot store operation failed: {message}")] - OperationFailed { - /// 错误说明。 - message: String, - }, -} - -/// 模型快照存储结果类型。 -pub type ModelSnapshotStoreResult = Result; - -/// 模型快照存储端口。 -#[async_trait] -pub trait ModelSnapshotStore: Send + Sync + 'static { - /// 用单个计划快照替换同名快照。 - async fn replace_plan_snapshot( - &self, - snapshot: &ModelPlanSnapshot, - ) -> ModelSnapshotStoreResult<()>; - - /// 列出所有计划快照。 - async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult>; -} - -/// SQLite 模型快照存储。 -#[derive(Clone)] -pub struct SqliteModelSnapshotStore { - pool: sqlx::SqlitePool, -} - -impl SqliteModelSnapshotStore { - /// 使用给定连接池构造存储。 - pub fn new(pool: sqlx::SqlitePool) -> Self { - Self { pool } - } -} - -#[async_trait] -impl ModelSnapshotStore for SqliteModelSnapshotStore { - async fn replace_plan_snapshot( - &self, - snapshot: &ModelPlanSnapshot, - ) -> ModelSnapshotStoreResult<()> { - let now = chrono::Utc::now().to_rfc3339(); - let models_json = serde_json::to_string(&snapshot.models).map_err(|e| { - ModelSnapshotStoreError::OperationFailed { - message: e.to_string(), - } - })?; - sqlx::query( - "insert or replace into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?)", - ) - .bind(&snapshot.plan_type) - .bind(&models_json) - .bind(&now) - .execute(&self.pool) - .await - .map_err(|e| ModelSnapshotStoreError::OperationFailed { - message: e.to_string(), - })?; - Ok(()) - } - - async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult> { - let rows = sqlx::query("select plan_type, models_json from model_plan_snapshots") - .fetch_all(&self.pool) - .await - .map_err(|e| ModelSnapshotStoreError::OperationFailed { - message: e.to_string(), - })?; - let mut snapshots = Vec::new(); - for row in rows { - let plan_type: String = row.get("plan_type"); - let models_json: String = row.get("models_json"); - let models: Vec = serde_json::from_str(&models_json).map_err(|e| { - ModelSnapshotStoreError::OperationFailed { - message: e.to_string(), - } - })?; - snapshots.push(ModelPlanSnapshot { plan_type, models }); - } - Ok(snapshots) - } -} diff --git a/backend/src/upstream/openai/fingerprint/mod.rs b/backend/src/upstream/openai/fingerprint/mod.rs new file mode 100644 index 000000000..67cabb3de --- /dev/null +++ b/backend/src/upstream/openai/fingerprint/mod.rs @@ -0,0 +1,11 @@ +//! OpenAI 客户端指纹。 + +mod runtime; +mod store; +mod types; +mod updater; + +pub use runtime::*; +pub use store::*; +pub use types::*; +pub use updater::*; diff --git a/backend/src/upstream/openai/fingerprint/runtime.rs b/backend/src/upstream/openai/fingerprint/runtime.rs new file mode 100644 index 000000000..5ffc442c7 --- /dev/null +++ b/backend/src/upstream/openai/fingerprint/runtime.rs @@ -0,0 +1,40 @@ +//! OpenAI 客户端运行时指纹快照。 + +use std::sync::{Arc, RwLock}; + +use super::types::Fingerprint; + +#[derive(Debug, Clone)] +pub struct RuntimeFingerprint { + current: Arc>>, +} + +impl RuntimeFingerprint { + pub fn new(fingerprint: Fingerprint) -> Self { + Self { + current: Arc::new(RwLock::new(Arc::new(fingerprint))), + } + } + + pub fn current(&self) -> Arc { + self.current + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + + pub fn snapshot(&self) -> Fingerprint { + self.current().as_ref().clone() + } + + pub fn replace(&self, fingerprint: Fingerprint) { + *self + .current + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::new(fingerprint); + } +} + +// --------------------------------------------------------------------------- +// 指纹历史记录与仓储 +// --------------------------------------------------------------------------- diff --git a/backend/src/upstream/openai/fingerprint/store.rs b/backend/src/upstream/openai/fingerprint/store.rs new file mode 100644 index 000000000..cb3dc7bf4 --- /dev/null +++ b/backend/src/upstream/openai/fingerprint/store.rs @@ -0,0 +1,216 @@ +//! OpenAI 客户端指纹 PostgreSQL 存储。 + +use chrono::{DateTime, Utc}; +use indexmap::IndexMap; +use serde::{Deserialize, Serialize}; +use sqlx::{PgPool, Row}; +use uuid::Uuid; + +use super::types::Fingerprint; + +pub const CURRENT_FINGERPRINT_ID: &str = "current"; +const AUTO_UPDATE_SOURCE: &str = "auto_update"; +const CONFIG_SEED_SOURCE: &str = "config_seed"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct StoredHeader { + name: String, + value: String, +} + +/// 指纹自动更新状态。 + +#[derive(Clone)] +pub struct PgFingerprintStore { + pool: PgPool, +} + +impl PgFingerprintStore { + /// 使用给定连接池构造仓储。 + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + /// 写入当前指纹默认值；如果当前槽位已存在，直接读取数据库值。 + pub async fn ensure_current_seed( + &self, + default_fingerprint: &Fingerprint, + ) -> Result { + if let Some(fingerprint) = self.load_current().await? { + return Ok(fingerprint); + } + + let now = Utc::now(); + sqlx::query( + r" + 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, $12) + on conflict (id) do nothing + ", + ) + .bind(CURRENT_FINGERPRINT_ID) + .bind(&default_fingerprint.originator) + .bind(&default_fingerprint.app_version) + .bind(&default_fingerprint.build_number) + .bind(&default_fingerprint.platform) + .bind(&default_fingerprint.arch) + .bind(&default_fingerprint.chromium_version) + .bind(&default_fingerprint.user_agent_template) + .bind(sqlx::types::Json(encode_default_headers( + &default_fingerprint.default_headers, + ))) + .bind(sqlx::types::Json(&default_fingerprint.header_order)) + .bind(CONFIG_SEED_SOURCE) + .bind(now) + .execute(&self.pool) + .await?; + + self.load_current().await?.ok_or(sqlx::Error::RowNotFound) + } + + /// 更新当前指纹中的自动更新版本字段。 + pub async fn update_current_version( + &self, + app_version: &str, + build_number: &str, + chromium_version: Option<&str>, + ) -> Result<(), sqlx::Error> { + let resolved_chromium_version = match chromium_version { + Some(version) => version.to_string(), + None => { + sqlx::query_scalar::<_, String>( + "select chromium_version from fingerprints where id = $1", + ) + .bind(CURRENT_FINGERPRINT_ID) + .fetch_one(&self.pool) + .await? + } + }; + + let result = sqlx::query( + "update fingerprints set app_version = $1, build_number = $2, chromium_version = $3, source = $4, updated_at = $5 where id = $6", + ) + .bind(app_version) + .bind(build_number) + .bind(resolved_chromium_version) + .bind(AUTO_UPDATE_SOURCE) + .bind(Utc::now()) + .bind(CURRENT_FINGERPRINT_ID) + .execute(&self.pool) + .await?; + if result.rows_affected() == 0 { + return Err(sqlx::Error::RowNotFound); + } + Ok(()) + } + + /// 读取当前运行时指纹。 + pub async fn load_current(&self) -> Result, sqlx::Error> { + let row = sqlx::query( + r" + select + originator, + app_version, + build_number, + platform, + arch, + chromium_version, + user_agent_template, + default_headers_json, + header_order_json, + updated_at + from fingerprints + where id = $1 + ", + ) + .bind(CURRENT_FINGERPRINT_ID) + .fetch_optional(&self.pool) + .await?; + + row.map(|row| fingerprint_from_row(&row)).transpose() + } + + /// 记录自动更新历史。 + pub async fn insert_update_history( + &self, + app_version: &str, + build_number: &str, + chromium_version: Option<&str>, + manifest_json: Option<&str>, + ) -> Result<(), sqlx::Error> { + let manifest = manifest_json + .map(serde_json::from_str::) + .transpose() + .map_err(json_error)?; + 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(Uuid::new_v4().to_string()) + .bind(CURRENT_FINGERPRINT_ID) + .bind(app_version) + .bind(build_number) + .bind(chromium_version) + .bind(AUTO_UPDATE_SOURCE) + .bind(manifest.map(sqlx::types::Json)) + .bind(Utc::now()) + .execute(&self.pool) + .await?; + Ok(()) + } +} + +fn fingerprint_from_row(row: &sqlx::postgres::PgRow) -> Result { + let default_headers = row + .get::>, _>("default_headers_json") + .0; + let header_order = row + .get::>, _>("header_order_json") + .0; + let updated_at = row.get::, _>("updated_at"); + Ok(Fingerprint { + originator: row.get("originator"), + app_version: row.get("app_version"), + build_number: row.get("build_number"), + platform: row.get("platform"), + arch: row.get("arch"), + chromium_version: row.get("chromium_version"), + user_agent_template: row.get("user_agent_template"), + default_headers: decode_default_headers(default_headers), + header_order, + updated_at: Some(updated_at.to_rfc3339()), + }) +} + +fn encode_default_headers(headers: &IndexMap) -> Vec { + headers + .iter() + .map(|(name, value)| StoredHeader { + name: name.clone(), + value: value.clone(), + }) + .collect() +} + +fn decode_default_headers(headers: Vec) -> IndexMap { + headers + .into_iter() + .map(|header| (header.name, header.value)) + .collect() +} + +fn json_error(error: serde_json::Error) -> sqlx::Error { + sqlx::Error::Decode(Box::new(error)) +} diff --git a/backend/src/upstream/openai/fingerprint/types.rs b/backend/src/upstream/openai/fingerprint/types.rs new file mode 100644 index 000000000..acece5e4a --- /dev/null +++ b/backend/src/upstream/openai/fingerprint/types.rs @@ -0,0 +1,92 @@ +//! OpenAI 客户端指纹领域类型。 + +use chrono::{DateTime, Utc}; +use indexmap::IndexMap; +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct Fingerprint { + /// 客户端来源名。 + pub originator: String, + /// 应用版本。 + pub app_version: String, + /// 构建号。 + pub build_number: String, + /// 平台名。 + pub platform: String, + /// 架构名。 + pub arch: String, + /// Chromium 主版本。 + pub chromium_version: String, + /// User-Agent 模板。 + pub user_agent_template: String, + /// 默认请求头。 + pub default_headers: IndexMap, + /// 请求头顺序。 + pub header_order: Vec, + /// DB 最后更新时间。 + #[serde(skip_serializing_if = "Option::is_none")] + pub updated_at: Option, +} + +impl Fingerprint { + /// 根据模板展开最终 User-Agent。 + pub fn user_agent(&self) -> String { + self.user_agent_template + .replace("{version}", &self.app_version) + .replace("{platform}", &self.platform) + .replace("{arch}", &self.arch) + } + + /// 生成 `sec-ch-ua` 头值。 + pub fn sec_ch_ua(&self) -> String { + format!( + "\"Chromium\";v=\"{}\", \"Not:A-Brand\";v=\"24\"", + self.chromium_version + ) + } +} + +/// 运行时共享指纹快照。 + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UpdateState { + /// 最近检查时间。 + pub last_check: DateTime, + /// 最新版本。 + pub latest_version: Option, + /// 最新构建号。 + pub latest_build: Option, + /// 下载地址。 + pub download_url: Option, + /// 是否有可用更新。 + pub update_available: bool, + /// 当前版本。 + pub current_version: String, + /// 当前构建号。 + pub current_build: String, +} + +/// Appcast 检查错误。 +#[derive(Debug, Error)] +pub enum UpdateError { + /// HTTP 请求失败。 + #[error("HTTP 请求失败: {0}")] + Http(#[from] reqwest::Error), + /// 获取 appcast 失败。 + #[error("获取 appcast 失败，状态码: {0}")] + AppcastFetch(u16), + /// 解析 appcast 失败。 + #[error("解析 appcast 失败")] + AppcastParse, + /// JSON 序列化失败。 + #[error("JSON 序列化失败: {0}")] + Json(#[from] serde_json::Error), + /// 文件操作失败。 + #[error("文件操作失败: {0}")] + Io(#[from] std::io::Error), + /// 数据库存储失败。 + #[error("数据库操作失败: {0}")] + Database(#[from] sqlx::Error), +} diff --git a/backend/src/upstream/openai/fingerprint/updater.rs b/backend/src/upstream/openai/fingerprint/updater.rs new file mode 100644 index 000000000..bbce8dbed --- /dev/null +++ b/backend/src/upstream/openai/fingerprint/updater.rs @@ -0,0 +1,213 @@ +//! OpenAI 客户端指纹更新检查器。 + +use std::{path::PathBuf, sync::Arc, time::Duration}; + +use chrono::Utc; +use tokio::sync::Mutex; +use tracing::info; + +use super::{ + store::PgFingerprintStore, + types::{Fingerprint, UpdateError, UpdateState}, +}; + +const APPCAST_TIMEOUT: Duration = Duration::from_secs(30); +pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_hours(72); +type AppcastFields = (Option, Option, Option); + +/// Appcast 指纹更新检查器。 +#[derive(Clone)] +pub struct UpdateChecker { + store: PgFingerprintStore, + client: reqwest::Client, + appcast_url: String, + extracted_fingerprint_path: PathBuf, + state: Arc>, +} + +struct InternalState { + current_version: String, + current_build: String, +} + +impl UpdateChecker { + /// 使用自定义 client 和路径构造检查器。 + pub fn with_client( + store: PgFingerprintStore, + client: reqwest::Client, + appcast_url: impl Into, + extracted_fingerprint_path: PathBuf, + current_version: impl Into, + current_build: impl Into, + ) -> Self { + Self { + store, + client, + appcast_url: appcast_url.into(), + extracted_fingerprint_path, + state: Arc::new(Mutex::new(InternalState { + current_version: current_version.into(), + current_build: current_build.into(), + })), + } + } + + /// 检查 appcast 是否存在更新。 + pub async fn check_for_update(&self) -> Result { + let state = self.state.lock().await; + let current_version = state.current_version.clone(); + let current_build = state.current_build.clone(); + drop(state); + + let xml = self.fetch_appcast().await?; + let (version, build, download_url) = parse_appcast(&xml)?; + + let update_available = version + .as_ref() + .is_some_and(|value| value != ¤t_version) + || build.as_ref().is_some_and(|value| value != ¤t_build); + + Ok(UpdateState { + last_check: Utc::now(), + latest_version: version, + latest_build: build, + download_url, + update_available, + current_version, + current_build, + }) + } + + /// 检查并在需要时应用更新。 + pub async fn check_and_apply_update(&self) -> Result, UpdateError> { + let update_state = self.check_for_update().await?; + + if !update_state.update_available { + return Ok(None); + } + + let (Some(version), Some(build)) = ( + update_state.latest_version.as_deref(), + update_state.latest_build.as_deref(), + ) else { + return Ok(None); + }; + + let state = self.state.lock().await; + let current_version = state.current_version.clone(); + let current_build = state.current_build.clone(); + drop(state); + + info!( + version = %version, + build = %build, + current_version = %current_version, + current_build = %current_build, + "发现新的 fingerprint 版本" + ); + + let updated_fingerprint = self.apply_version_update(version, build).await?; + + let mut state = self.state.lock().await; + state.current_version = version.to_string(); + state.current_build = build.to_string(); + + Ok(Some(updated_fingerprint)) + } + + async fn fetch_appcast(&self) -> Result { + let response = self + .client + .get(&self.appcast_url) + .timeout(APPCAST_TIMEOUT) + .send() + .await?; + + if !response.status().is_success() { + return Err(UpdateError::AppcastFetch(response.status().as_u16())); + } + + response.text().await.map_err(UpdateError::from) + } + + async fn apply_version_update( + &self, + version: &str, + build: &str, + ) -> Result { + let chromium_version = + load_matching_chromium_version(&self.extracted_fingerprint_path, version, build); + + self.store + .update_current_version(version, build, chromium_version.as_deref()) + .await?; + self.store + .insert_update_history(version, build, chromium_version.as_deref(), None) + .await?; + let updated = self + .store + .load_current() + .await? + .ok_or(sqlx::Error::RowNotFound)?; + Ok(updated) + } +} + +fn parse_appcast(xml: &str) -> Result { + let item_start = xml.find("").ok_or(UpdateError::AppcastParse)?; + let item_end = xml.find("").ok_or(UpdateError::AppcastParse)?; + let item = &xml[item_start..item_end]; + + let version = extract_sparkle_field(item, "shortVersionString"); + let build = extract_sparkle_field(item, "version"); + let download_url = item.find("url=\"").and_then(|position| { + let start = position + 5; + item[start..] + .find('"') + .map(|end| item[start..start + end].to_string()) + }); + + Ok((version, build, download_url)) +} + +fn extract_sparkle_field(item: &str, field: &str) -> Option { + let attr_pattern = format!("sparkle:{}=\"", field); + if let Some(position) = item.find(&attr_pattern) { + let start = position + attr_pattern.len(); + if let Some(end) = item[start..].find('"') { + return Some(item[start..start + end].to_string()); + } + } + + let elem_start = format!("", field); + let elem_end = format!("", field); + if let Some(start_position) = item.find(&elem_start) { + let content_start = start_position + elem_start.len(); + if let Some(end_position) = item[content_start..].find(&elem_end) { + return Some(item[content_start..content_start + end_position].to_string()); + } + } + + None +} + +fn load_matching_chromium_version( + extracted_fingerprint_path: &PathBuf, + version: &str, + build: &str, +) -> Option { + let content = std::fs::read_to_string(extracted_fingerprint_path).ok()?; + let parsed: serde_json::Value = serde_json::from_str(&content).ok()?; + + let extracted_version = parsed.get("app_version")?.as_str()?; + let extracted_build = parsed.get("build_number")?.as_str()?; + + if extracted_version == version && extracted_build == build { + parsed + .get("chromium_version")? + .as_str() + .map(ToString::to_string) + } else { + None + } +} diff --git a/backend/src/upstream/openai/mod.rs b/backend/src/upstream/openai/mod.rs new file mode 100644 index 000000000..209c19ea6 --- /dev/null +++ b/backend/src/upstream/openai/mod.rs @@ -0,0 +1,6 @@ +//! OpenAI/Codex 上游客户端。 + +pub mod fingerprint; +pub mod protocol; +pub mod token_client; +pub mod transport; diff --git a/backend/src/upstream/protocol/events.rs b/backend/src/upstream/openai/protocol/events.rs similarity index 99% rename from backend/src/upstream/protocol/events.rs rename to backend/src/upstream/openai/protocol/events.rs index 96731a2f4..ff68282cb 100644 --- a/backend/src/upstream/protocol/events.rs +++ b/backend/src/upstream/openai/protocol/events.rs @@ -3,7 +3,7 @@ use std::{collections::BTreeMap, time::Duration}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -use crate::upstream::protocol::sse::{parse_sse_events, SseError}; +use crate::upstream::openai::protocol::sse::{parse_sse_events, SseError}; /// 从 Codex/OpenAI usage 结构中提取出的标准化 token 用量。 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] diff --git a/backend/src/upstream/protocol/mod.rs b/backend/src/upstream/openai/protocol/mod.rs similarity index 83% rename from backend/src/upstream/protocol/mod.rs rename to backend/src/upstream/openai/protocol/mod.rs index 57085bee1..b8dacd0f2 100644 --- a/backend/src/upstream/protocol/mod.rs +++ b/backend/src/upstream/openai/protocol/mod.rs @@ -2,6 +2,8 @@ /// Codex 事件语义与用量提取。 pub mod events; +/// 上游模型目录响应类型。 +pub mod model_catalog; /// Codex Responses / Compact 请求体。 pub mod responses; /// Codex JSON schema 处理。 diff --git a/backend/src/upstream/models/backend_entry.rs b/backend/src/upstream/openai/protocol/model_catalog.rs similarity index 100% rename from backend/src/upstream/models/backend_entry.rs rename to backend/src/upstream/openai/protocol/model_catalog.rs diff --git a/backend/src/upstream/protocol/responses.rs b/backend/src/upstream/openai/protocol/responses.rs similarity index 98% rename from backend/src/upstream/protocol/responses.rs rename to backend/src/upstream/openai/protocol/responses.rs index b2f8fa0e2..9e26ca5ff 100644 --- a/backend/src/upstream/protocol/responses.rs +++ b/backend/src/upstream/openai/protocol/responses.rs @@ -5,7 +5,7 @@ use serde_json::{json, Map, Value}; use crate::{ infra::time::elapsed_millis_i64, - upstream::protocol::{ + upstream::openai::protocol::{ schema::reconvert_tuple_values, sse::{parse_sse_events, SseError, SseEvent}, }, @@ -36,6 +36,8 @@ pub struct CodexResponsesRequest { pub client_ip: Option, /// 客户端 User-Agent，仅用于管理端使用记录展示。 pub client_user_agent: Option, + /// 已鉴权客户端 API key 的稳定 ID，仅用于事实归因。 + pub client_api_key_id: Option, /// 是否偏好 WebSocket 传输。 pub use_websocket: bool, /// 是否强制 HTTP SSE。 @@ -544,6 +546,7 @@ impl CodexResponsesRequest { variant_identity: None, client_ip: None, client_user_agent: None, + client_api_key_id: None, use_websocket: false, force_http_sse: false, turn_state: None, @@ -750,6 +753,8 @@ pub struct CodexCompactRequest { pub client_ip: Option, /// 客户端 User-Agent，仅用于管理端使用记录展示。 pub client_user_agent: Option, + /// 已鉴权客户端 API key 的稳定 ID，仅用于事实归因。 + pub client_api_key_id: Option, } impl Serialize for CodexCompactRequest { diff --git a/backend/src/upstream/protocol/schema.rs b/backend/src/upstream/openai/protocol/schema.rs similarity index 100% rename from backend/src/upstream/protocol/schema.rs rename to backend/src/upstream/openai/protocol/schema.rs diff --git a/backend/src/upstream/protocol/sse.rs b/backend/src/upstream/openai/protocol/sse.rs similarity index 100% rename from backend/src/upstream/protocol/sse.rs rename to backend/src/upstream/openai/protocol/sse.rs diff --git a/backend/src/upstream/openai/protocol/websocket.rs b/backend/src/upstream/openai/protocol/websocket.rs new file mode 100644 index 000000000..d3cf250df --- /dev/null +++ b/backend/src/upstream/openai/protocol/websocket.rs @@ -0,0 +1,423 @@ +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use crate::upstream::openai::protocol::responses::{ + http_sse_fallback_allowed, transport_for_request, CodexResponsesRequest, CodexTransport, +}; +use crate::upstream::openai::protocol::sse::encode_sse_event; + +const REDACTED_PAYLOAD_VALUE: &str = ""; + +/// WebSocket 握手审计快照。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct OpeningAuditSnapshot { + /// 请求行。 + pub request_line: String, + /// 请求头顺序。 + pub header_order: Vec, + /// 红action后的请求头。 + pub headers: Vec, +} + +/// WebSocket 握手审计请求头。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct OpeningAuditHeader { + /// 请求头名。 + pub name: String, + /// 请求头值。 + pub value: String, +} + +/// WebSocket payload 审计快照。 +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +pub struct PayloadAuditSnapshot { + /// 按构造顺序记录的顶层字段。 + pub top_level_keys: Vec, + /// 红action后的 JSON payload。 + pub body: Value, +} + +/// WebSocket audit artifact. +#[derive(Debug, Clone, Serialize, PartialEq)] +pub struct WebSocketAuditArtifact { + /// 实际选择的传输模式。 + pub transport_mode: String, + /// 当前请求是否允许 HTTP/SSE fallback。 + pub fallback_allowed: bool, + /// 打开握手快照。 + pub opening: Option, + /// 首个 `response.create` payload 快照。 + pub payload: Option, +} + +/// 从 WebSocket 错误帧推导出的上游错误分类。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ClassifiedWebSocketError { + /// 应映射给 HTTP 层的状态码。 + pub status_code: u16, +} + +#[derive(Debug, Deserialize)] +struct ResponseCompleted { + #[serde(rename = "id")] + _id: String, + #[serde(default, rename = "usage")] + _usage: Option, + #[serde(default, rename = "end_turn")] + _end_turn: Option, +} + +#[derive(Debug, Deserialize)] +struct ResponseCompletedUsage { + #[serde(rename = "input_tokens")] + _input_tokens: i64, + #[serde(default, rename = "input_tokens_details")] + _input_tokens_details: Option, + #[serde(rename = "output_tokens")] + _output_tokens: i64, + #[serde(default, rename = "output_tokens_details")] + _output_tokens_details: Option, + #[serde(rename = "total_tokens")] + _total_tokens: i64, +} + +#[derive(Debug, Deserialize)] +struct ResponseCompletedInputTokensDetails { + #[serde(rename = "cached_tokens")] + _cached_tokens: i64, +} + +#[derive(Debug, Deserialize)] +struct ResponseCompletedOutputTokensDetails { + #[serde(rename = "reasoning_tokens")] + _reasoning_tokens: i64, +} + +#[derive(Debug, Deserialize)] +struct ResponsesStreamEventShape { + #[serde(rename = "type")] + _kind: String, + #[serde(default, rename = "headers")] + _headers: Option, + #[serde(default, rename = "metadata")] + _metadata: Option, + #[serde(default, rename = "response")] + _response: Option, + #[serde(default, rename = "item")] + _item: Option, + #[serde(default, rename = "item_id")] + _item_id: Option, + #[serde(default, rename = "call_id")] + _call_id: Option, + #[serde(default, rename = "delta")] + _delta: Option, + #[serde(default, rename = "summary_index")] + _summary_index: Option, + #[serde(default, rename = "content_index")] + _content_index: Option, +} + +/// 将一条公开 WebSocket JSON 事件编码为 SSE 帧。 +pub fn websocket_event_to_sse_frame(raw: &str) -> Option { + let event = websocket_event_type(raw)?; + if is_internal_websocket_event(&event) { + return None; + } + if websocket_event_should_skip(raw) { + return None; + } + Some(encode_sse_event(&event, raw)) +} + +pub(crate) fn websocket_event_type(raw: &str) -> Option { + serde_json::from_str::(raw).ok().and_then(|value| { + value + .get("type") + .and_then(Value::as_str) + .map(ToString::to_string) + }) +} + +fn is_internal_websocket_event(event: &str) -> bool { + event == "codex.rate_limits" +} + +/// 判断 WebSocket 事件是否应在转发前剥离。 +/// +/// 仅剥离传输层内部帧：`response.metadata` 承载 turn_state，由上层提取转存到 +/// 会话状态，不下发客户端。业务事件一律原样转发，不做 schema 校验。 +/// （`codex.rate_limits` 内部事件由 `is_internal_websocket_event` 单独剥离。） +fn websocket_event_should_skip(raw: &str) -> bool { + websocket_metadata_turn_state(raw).is_some() +} + +/// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。 +pub fn websocket_metadata_turn_state(raw: &str) -> Option { + let value = serde_json::from_str::(raw).ok()?; + if value.get("type").and_then(Value::as_str) != Some("response.metadata") { + return None; + } + value + .get("headers") + .and_then(Value::as_object) + .and_then(|headers| { + headers.iter().find_map(|(name, value)| { + if name.eq_ignore_ascii_case("x-codex-turn-state") { + json_value_as_string(value) + } else { + None + } + }) + }) + .or_else(|| { + value + .pointer("/metadata/turn_state") + .or_else(|| value.get("turn_state")) + .and_then(Value::as_str) + .map(ToString::to_string) + }) +} + +fn json_value_as_string(value: &Value) -> Option { + match value { + Value::String(value) => Some(value.clone()), + Value::Array(items) => items.first().and_then(json_value_as_string), + _ => None, + } +} + +/// 校验 `response.completed` 的官方响应形状。 +pub fn websocket_response_completed_parse_error(raw: &str) -> Option { + let value = serde_json::from_str::(raw).ok()?; + if value.get("type").and_then(Value::as_str) != Some("response.completed") { + return None; + } + let response = value.get("response")?; + if response.is_null() { + return None; + } + match serde_json::from_value::(response.clone()) { + Ok(_) => None, + Err(error) => Some(format!("failed to parse ResponseCompleted: {error}")), + } +} + +/// 从 `response.incomplete` 中提取 incomplete reason。 +pub fn websocket_incomplete_response_reason(raw: &str) -> Option { + let value = serde_json::from_str::(raw).ok()?; + if value.get("type").and_then(Value::as_str) != Some("response.incomplete") { + return None; + } + value + .pointer("/response/incomplete_details/reason") + .or_else(|| value.pointer("/incomplete_details/reason")) + .or_else(|| value.get("reason")) + .and_then(Value::as_str) + .map(ToString::to_string) +} + +mod validation; +pub use validation::*; + +/// 判断 WebSocket 事件是否会结束当前响应流。 +pub fn is_terminal_websocket_event(event: &str) -> bool { + matches!(event, "response.completed" | "response.failed" | "error") +} + +/// 判断是否为「首个内容帧之前的纯生命周期帧」。 +/// +/// `response.created` / `response.in_progress` 仅表示上游已受理请求，尚未开始产出 +/// 真实内容（含 reasoning）。首 token 熔断计时不应被它们解除，否则连接一收到 +/// `response.created`（通常 1s 内）就被误判为「已开始产出」，使真正的首内容帧 +/// 可以无限拖延而不触发保护。其余任何帧（reasoning/delta/output_item 等）到达即视为 +/// 上游已开始真实产出，解除熔断——对慢思考模型友好。 +pub fn is_pre_content_lifecycle_event(event: &str) -> bool { + matches!(event, "response.created" | "response.in_progress") +} + +/// 将 WebSocket 错误帧映射为上游 HTTP 状态分类。 +pub fn classify_websocket_error_frame(raw: &str) -> Option { + let value = serde_json::from_str::(raw).ok()?; + let event_type = value.get("type").and_then(Value::as_str)?; + if event_type != "error" && event_type != "response.failed" { + return None; + } + + let code = websocket_error_code(&value); + if code.as_deref() == Some("websocket_connection_limit_reached") { + return Some(ClassifiedWebSocketError { status_code: 503 }); + } + + if event_type == "error" { + if let Some(status_code) = explicit_error_status_code(&value) { + if (200..=299).contains(&status_code) { + return None; + } + return Some(ClassifiedWebSocketError { status_code }); + } + } + + if let Some(code) = code { + if let Some(status_code) = rotatable_error_status_code(&code) { + return Some(ClassifiedWebSocketError { status_code }); + } + } + + None +} + +fn websocket_error_code(value: &Value) -> Option { + value + .pointer("/response/error/code") + .or_else(|| value.pointer("/response/error/type")) + .or_else(|| value.pointer("/error/code")) + .or_else(|| value.pointer("/error/type")) + .and_then(Value::as_str) + .map(str::to_ascii_lowercase) +} + +fn explicit_error_status_code(value: &Value) -> Option { + value + .get("status") + .or_else(|| value.get("status_code")) + .and_then(Value::as_u64) + .and_then(|status| u16::try_from(status).ok()) +} + +fn rotatable_error_status_code(code: &str) -> Option { + match code { + "usage_limit_reached" | "rate_limit_exceeded" | "rate_limit_reached" => Some(429), + "quota_exhausted" | "payment_required" => Some(402), + "unauthorized" | "token_invalid" | "token_expired" | "account_deactivated" => Some(401), + "forbidden" | "account_banned" | "banned" => Some(403), + "previous_response_not_found" => Some(400), + _ => None, + } +} + +/// 从包裹在 WebSocket `error.headers` 中的 retry-after 值提取秒数。 +pub fn retry_after_seconds_from_wrapped_error_headers(raw: &str) -> Option { + let value = serde_json::from_str::(raw).ok()?; + let header_retry_after = value + .get("headers") + .and_then(Value::as_object) + .and_then(|headers| { + headers.iter().find_map(|(name, value)| { + if name.eq_ignore_ascii_case("retry-after") { + retry_after_seconds_from_json_header_value(value) + } else { + None + } + }) + }); + header_retry_after.or_else(|| { + value + .pointer("/error/retry_after_seconds") + .or_else(|| value.get("retry_after_seconds")) + .and_then(Value::as_u64) + }) +} + +fn retry_after_seconds_from_json_header_value(value: &Value) -> Option { + let seconds = match value { + Value::String(value) => value.trim().parse::().ok()?, + Value::Number(value) => value.as_u64()?, + Value::Array(values) => values + .first() + .and_then(retry_after_seconds_from_json_header_value)?, + _ => return None, + }; + (seconds > 0).then_some(seconds) +} + +/// 生成 Responses WebSocket payload 审计快照。 +pub fn websocket_payload_audit_snapshot(request: &CodexResponsesRequest) -> PayloadAuditSnapshot { + let body = websocket_response_create_payload(request); + PayloadAuditSnapshot { + top_level_keys: websocket_payload_keys(request), + body: redact_payload_body(body), + } +} + +/// 为单次 opening 尝试构建 WebSocket 审计 artifact。 +pub fn websocket_audit_artifact_from_attempt( + request: &CodexResponsesRequest, + opening: OpeningAuditSnapshot, + payload: PayloadAuditSnapshot, +) -> WebSocketAuditArtifact { + WebSocketAuditArtifact { + transport_mode: websocket_transport_mode_name(request).to_string(), + fallback_allowed: http_sse_fallback_allowed(request), + opening: Some(opening), + payload: Some(payload), + } +} + +fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static str { + match transport_for_request(request) { + CodexTransport::HttpSse => "http_sse", + CodexTransport::WebSocketPreferred => "websocket_preferred", + CodexTransport::WebSocketRequired => "websocket_required", + } +} + +/// 生成 Responses WebSocket `response.create` payload。 +/// +/// payload = `{"type": "response.create"}` 加上原始上游 body 的全部字段 +/// （保持插入顺序，含未知字段），逐字段原样透传。 +pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Value { + let mut payload = Map::new(); + payload.insert( + "type".to_string(), + Value::String("response.create".to_string()), + ); + for (key, value) in request.body() { + payload.insert(key.clone(), value.clone()); + } + Value::Object(payload) +} + +/// 生成 Responses WebSocket `response.create` 文本帧内容。 +pub fn websocket_response_create_payload_text( + request: &CodexResponsesRequest, +) -> Result { + serde_json::to_string(&websocket_response_create_payload(request)) +} + +fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec { + let mut keys = Vec::with_capacity(request.body().len() + 1); + keys.push("type".to_string()); + keys.extend(request.body().keys().cloned()); + keys +} + +fn redact_payload_body(body: Value) -> Value { + let Value::Object(body) = body else { + return body; + }; + + Value::Object( + body.into_iter() + .map(|(key, value)| { + let value = if is_sensitive_payload_key(&key) { + Value::String(REDACTED_PAYLOAD_VALUE.to_string()) + } else { + value + }; + (key, value) + }) + .collect(), + ) +} + +fn is_sensitive_payload_key(key: &str) -> bool { + matches!( + key, + "instructions" + | "input" + | "previous_response_id" + | "prompt_cache_key" + | "client_metadata" + | "tools" + ) +} diff --git a/backend/src/upstream/protocol/websocket.rs b/backend/src/upstream/openai/protocol/websocket/validation.rs similarity index 61% rename from backend/src/upstream/protocol/websocket.rs rename to backend/src/upstream/openai/protocol/websocket/validation.rs index f4028be39..5582b0111 100644 --- a/backend/src/upstream/protocol/websocket.rs +++ b/backend/src/upstream/openai/protocol/websocket/validation.rs @@ -1,220 +1,4 @@ -use serde::{Deserialize, Serialize}; -use serde_json::{Map, Value}; - -use crate::upstream::protocol::responses::{ - http_sse_fallback_allowed, transport_for_request, CodexResponsesRequest, CodexTransport, -}; -use crate::upstream::protocol::sse::encode_sse_event; - -const REDACTED_PAYLOAD_VALUE: &str = ""; - -/// WebSocket 握手审计快照。 -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -pub struct OpeningAuditSnapshot { - /// 请求行。 - pub request_line: String, - /// 请求头顺序。 - pub header_order: Vec, - /// 红action后的请求头。 - pub headers: Vec, -} - -/// WebSocket 握手审计请求头。 -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -pub struct OpeningAuditHeader { - /// 请求头名。 - pub name: String, - /// 请求头值。 - pub value: String, -} - -/// WebSocket payload 审计快照。 -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] -pub struct PayloadAuditSnapshot { - /// 按构造顺序记录的顶层字段。 - pub top_level_keys: Vec, - /// 红action后的 JSON payload。 - pub body: Value, -} - -/// WebSocket audit artifact. -#[derive(Debug, Clone, Serialize, PartialEq)] -pub struct WebSocketAuditArtifact { - /// 实际选择的传输模式。 - pub transport_mode: String, - /// 当前请求是否允许 HTTP/SSE fallback。 - pub fallback_allowed: bool, - /// 打开握手快照。 - pub opening: Option, - /// 首个 `response.create` payload 快照。 - pub payload: Option, -} - -/// 从 WebSocket 错误帧推导出的上游错误分类。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct ClassifiedWebSocketError { - /// 应映射给 HTTP 层的状态码。 - pub status_code: u16, -} - -#[derive(Debug, Deserialize)] -struct ResponseCompleted { - #[serde(rename = "id")] - _id: String, - #[serde(default, rename = "usage")] - _usage: Option, - #[serde(default, rename = "end_turn")] - _end_turn: Option, -} - -#[derive(Debug, Deserialize)] -struct ResponseCompletedUsage { - #[serde(rename = "input_tokens")] - _input_tokens: i64, - #[serde(default, rename = "input_tokens_details")] - _input_tokens_details: Option, - #[serde(rename = "output_tokens")] - _output_tokens: i64, - #[serde(default, rename = "output_tokens_details")] - _output_tokens_details: Option, - #[serde(rename = "total_tokens")] - _total_tokens: i64, -} - -#[derive(Debug, Deserialize)] -struct ResponseCompletedInputTokensDetails { - #[serde(rename = "cached_tokens")] - _cached_tokens: i64, -} - -#[derive(Debug, Deserialize)] -struct ResponseCompletedOutputTokensDetails { - #[serde(rename = "reasoning_tokens")] - _reasoning_tokens: i64, -} - -#[derive(Debug, Deserialize)] -struct ResponsesStreamEventShape { - #[serde(rename = "type")] - _kind: String, - #[serde(default, rename = "headers")] - _headers: Option, - #[serde(default, rename = "metadata")] - _metadata: Option, - #[serde(default, rename = "response")] - _response: Option, - #[serde(default, rename = "item")] - _item: Option, - #[serde(default, rename = "item_id")] - _item_id: Option, - #[serde(default, rename = "call_id")] - _call_id: Option, - #[serde(default, rename = "delta")] - _delta: Option, - #[serde(default, rename = "summary_index")] - _summary_index: Option, - #[serde(default, rename = "content_index")] - _content_index: Option, -} - -/// 将一条公开 WebSocket JSON 事件编码为 SSE 帧。 -pub fn websocket_event_to_sse_frame(raw: &str) -> Option { - let event = websocket_event_type(raw)?; - if is_internal_websocket_event(&event) { - return None; - } - if websocket_event_should_skip(raw) { - return None; - } - Some(encode_sse_event(&event, raw)) -} - -pub(crate) fn websocket_event_type(raw: &str) -> Option { - serde_json::from_str::(raw).ok().and_then(|value| { - value - .get("type") - .and_then(Value::as_str) - .map(ToString::to_string) - }) -} - -fn is_internal_websocket_event(event: &str) -> bool { - event == "codex.rate_limits" -} - -/// 判断 WebSocket 事件是否应在转发前剥离。 -/// -/// 仅剥离传输层内部帧：`response.metadata` 承载 turn_state，由上层提取转存到 -/// 会话状态，不下发客户端。业务事件一律原样转发，不做 schema 校验。 -/// （`codex.rate_limits` 内部事件由 `is_internal_websocket_event` 单独剥离。） -fn websocket_event_should_skip(raw: &str) -> bool { - websocket_metadata_turn_state(raw).is_some() -} - -/// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。 -pub fn websocket_metadata_turn_state(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - if value.get("type").and_then(Value::as_str) != Some("response.metadata") { - return None; - } - value - .get("headers") - .and_then(Value::as_object) - .and_then(|headers| { - headers.iter().find_map(|(name, value)| { - if name.eq_ignore_ascii_case("x-codex-turn-state") { - json_value_as_string(value) - } else { - None - } - }) - }) - .or_else(|| { - value - .pointer("/metadata/turn_state") - .or_else(|| value.get("turn_state")) - .and_then(Value::as_str) - .map(ToString::to_string) - }) -} - -fn json_value_as_string(value: &Value) -> Option { - match value { - Value::String(value) => Some(value.clone()), - Value::Array(items) => items.first().and_then(json_value_as_string), - _ => None, - } -} - -/// 校验 `response.completed` 的官方响应形状。 -pub fn websocket_response_completed_parse_error(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - if value.get("type").and_then(Value::as_str) != Some("response.completed") { - return None; - } - let response = value.get("response")?; - if response.is_null() { - return None; - } - match serde_json::from_value::(response.clone()) { - Ok(_) => None, - Err(error) => Some(format!("failed to parse ResponseCompleted: {error}")), - } -} - -/// 从 `response.incomplete` 中提取 incomplete reason。 -pub fn websocket_incomplete_response_reason(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - if value.get("type").and_then(Value::as_str) != Some("response.incomplete") { - return None; - } - value - .pointer("/response/incomplete_details/reason") - .or_else(|| value.pointer("/incomplete_details/reason")) - .or_else(|| value.get("reason")) - .and_then(Value::as_str) - .map(ToString::to_string) -} +use super::*; /// 判断事件是否不符合官方流事件的基本字段类型。 pub fn websocket_event_shape_parse_error(raw: &str) -> bool { @@ -828,206 +612,3 @@ pub fn websocket_reasoning_summary_part_added_missing_summary_index(raw: &str) - value.get("type").and_then(Value::as_str) == Some("response.reasoning_summary_part.added") && value.get("summary_index").and_then(Value::as_i64).is_none() } - -/// 判断 WebSocket 事件是否会结束当前响应流。 -pub fn is_terminal_websocket_event(event: &str) -> bool { - matches!(event, "response.completed" | "response.failed" | "error") -} - -/// 判断是否为「首个内容帧之前的纯生命周期帧」。 -/// -/// `response.created` / `response.in_progress` 仅表示上游已受理请求，尚未开始产出 -/// 真实内容（含 reasoning）。首 token 熔断计时不应被它们解除，否则连接一收到 -/// `response.created`（通常 1s 内）就被误判为「已开始产出」，使真正的首内容帧 -/// 可以无限拖延而不触发保护。其余任何帧（reasoning/delta/output_item 等）到达即视为 -/// 上游已开始真实产出，解除熔断——对慢思考模型友好。 -pub fn is_pre_content_lifecycle_event(event: &str) -> bool { - matches!(event, "response.created" | "response.in_progress") -} - -/// 将 WebSocket 错误帧映射为上游 HTTP 状态分类。 -pub fn classify_websocket_error_frame(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - let event_type = value.get("type").and_then(Value::as_str)?; - if event_type != "error" && event_type != "response.failed" { - return None; - } - - let code = websocket_error_code(&value); - if code.as_deref() == Some("websocket_connection_limit_reached") { - return Some(ClassifiedWebSocketError { status_code: 503 }); - } - - if event_type == "error" { - if let Some(status_code) = explicit_error_status_code(&value) { - if (200..=299).contains(&status_code) { - return None; - } - return Some(ClassifiedWebSocketError { status_code }); - } - } - - if let Some(code) = code { - if let Some(status_code) = rotatable_error_status_code(&code) { - return Some(ClassifiedWebSocketError { status_code }); - } - } - - None -} - -fn websocket_error_code(value: &Value) -> Option { - value - .pointer("/response/error/code") - .or_else(|| value.pointer("/response/error/type")) - .or_else(|| value.pointer("/error/code")) - .or_else(|| value.pointer("/error/type")) - .and_then(Value::as_str) - .map(str::to_ascii_lowercase) -} - -fn explicit_error_status_code(value: &Value) -> Option { - value - .get("status") - .or_else(|| value.get("status_code")) - .and_then(Value::as_u64) - .and_then(|status| u16::try_from(status).ok()) -} - -fn rotatable_error_status_code(code: &str) -> Option { - match code { - "usage_limit_reached" | "rate_limit_exceeded" | "rate_limit_reached" => Some(429), - "quota_exhausted" | "payment_required" => Some(402), - "unauthorized" | "token_invalid" | "token_expired" | "account_deactivated" => Some(401), - "forbidden" | "account_banned" | "banned" => Some(403), - "previous_response_not_found" => Some(400), - _ => None, - } -} - -/// 从包裹在 WebSocket `error.headers` 中的 retry-after 值提取秒数。 -pub fn retry_after_seconds_from_wrapped_error_headers(raw: &str) -> Option { - let value = serde_json::from_str::(raw).ok()?; - let header_retry_after = value - .get("headers") - .and_then(Value::as_object) - .and_then(|headers| { - headers.iter().find_map(|(name, value)| { - if name.eq_ignore_ascii_case("retry-after") { - retry_after_seconds_from_json_header_value(value) - } else { - None - } - }) - }); - header_retry_after.or_else(|| { - value - .pointer("/error/retry_after_seconds") - .or_else(|| value.get("retry_after_seconds")) - .and_then(Value::as_u64) - }) -} - -fn retry_after_seconds_from_json_header_value(value: &Value) -> Option { - let seconds = match value { - Value::String(value) => value.trim().parse::().ok()?, - Value::Number(value) => value.as_u64()?, - Value::Array(values) => values - .first() - .and_then(retry_after_seconds_from_json_header_value)?, - _ => return None, - }; - (seconds > 0).then_some(seconds) -} - -/// 生成 Responses WebSocket payload 审计快照。 -pub fn websocket_payload_audit_snapshot(request: &CodexResponsesRequest) -> PayloadAuditSnapshot { - let body = websocket_response_create_payload(request); - PayloadAuditSnapshot { - top_level_keys: websocket_payload_keys(request), - body: redact_payload_body(body), - } -} - -/// 为单次 opening 尝试构建 WebSocket 审计 artifact。 -pub fn websocket_audit_artifact_from_attempt( - request: &CodexResponsesRequest, - opening: OpeningAuditSnapshot, - payload: PayloadAuditSnapshot, -) -> WebSocketAuditArtifact { - WebSocketAuditArtifact { - transport_mode: websocket_transport_mode_name(request).to_string(), - fallback_allowed: http_sse_fallback_allowed(request), - opening: Some(opening), - payload: Some(payload), - } -} - -fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static str { - match transport_for_request(request) { - CodexTransport::HttpSse => "http_sse", - CodexTransport::WebSocketPreferred => "websocket_preferred", - CodexTransport::WebSocketRequired => "websocket_required", - } -} - -/// 生成 Responses WebSocket `response.create` payload。 -/// -/// payload = `{"type": "response.create"}` 加上原始上游 body 的全部字段 -/// （保持插入顺序，含未知字段），逐字段原样透传。 -pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Value { - let mut payload = Map::new(); - payload.insert( - "type".to_string(), - Value::String("response.create".to_string()), - ); - for (key, value) in request.body() { - payload.insert(key.clone(), value.clone()); - } - Value::Object(payload) -} - -/// 生成 Responses WebSocket `response.create` 文本帧内容。 -pub fn websocket_response_create_payload_text( - request: &CodexResponsesRequest, -) -> Result { - serde_json::to_string(&websocket_response_create_payload(request)) -} - -fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec { - let mut keys = Vec::with_capacity(request.body().len() + 1); - keys.push("type".to_string()); - keys.extend(request.body().keys().cloned()); - keys -} - -fn redact_payload_body(body: Value) -> Value { - let Value::Object(body) = body else { - return body; - }; - - Value::Object( - body.into_iter() - .map(|(key, value)| { - let value = if is_sensitive_payload_key(&key) { - Value::String(REDACTED_PAYLOAD_VALUE.to_string()) - } else { - value - }; - (key, value) - }) - .collect(), - ) -} - -fn is_sensitive_payload_key(key: &str) -> bool { - matches!( - key, - "instructions" - | "input" - | "previous_response_id" - | "prompt_cache_key" - | "client_metadata" - | "tools" - ) -} diff --git a/backend/src/upstream/token_client.rs b/backend/src/upstream/openai/token_client.rs similarity index 82% rename from backend/src/upstream/token_client.rs rename to backend/src/upstream/openai/token_client.rs index 0dd703607..222d661c2 100644 --- a/backend/src/upstream/token_client.rs +++ b/backend/src/upstream/openai/token_client.rs @@ -4,7 +4,31 @@ use async_trait::async_trait; use reqwest::{Client, StatusCode}; use serde::Deserialize; -use crate::upstream::accounts::token_refresh::{RefreshFailure, TokenPair, TokenRefresher}; +/// 一次刷新返回的 token 对。 +#[derive(Debug, Clone)] +pub struct TokenPair { + pub access_token: String, + pub refresh_token: Option, +} + +/// OpenAI 刷新端点失败分类。 +#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)] +pub enum RefreshFailure { + #[error("refresh token is invalid or expired")] + InvalidGrant, + #[error("account is banned")] + Banned, + #[error("refresh transport failed before server processing")] + RetryableTransport, + #[error("refresh transport failed after possible server processing")] + Transport, +} + +/// 刷新令牌的上游端口。 +#[async_trait] +pub trait TokenRefresher: Send + Sync + 'static { + async fn refresh(&self, refresh_token: &str) -> Result; +} /// OpenAI token 续期客户端配置。 #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/backend/src/upstream/openai/transport/client.rs b/backend/src/upstream/openai/transport/client.rs new file mode 100644 index 000000000..8011c2c37 --- /dev/null +++ b/backend/src/upstream/openai/transport/client.rs @@ -0,0 +1,606 @@ +//! Codex HTTP/SSE 上游客户端、请求头构造、TLS 与自定义 CA。 + +use std::{ + collections::HashMap, + pin::Pin, + sync::{Arc, Mutex, OnceLock}, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::{Stream, TryStreamExt}; +use reqwest::{ + header::{ + HeaderMap, HeaderName, HeaderValue, ACCEPT, ACCEPT_ENCODING, AUTHORIZATION, CONTENT_TYPE, + COOKIE, RETRY_AFTER, USER_AGENT, + }, + Client, Response as ReqwestResponse, StatusCode, +}; +use serde_json::{map::Map, Value}; +use thiserror::Error; +use tokio_tungstenite::tungstenite::handshake::client::generate_key; +use uuid::Uuid; + +use crate::upstream::openai::fingerprint::{Fingerprint, RuntimeFingerprint}; +use crate::upstream::openai::protocol::events::{extract_sse_usage, retry_after_seconds_from_body}; +use crate::upstream::openai::protocol::model_catalog::BackendModelEntry; +use crate::upstream::openai::protocol::responses::{ + http_sse_fallback_allowed, transport_for_request, CodexCompactRequest, CodexResponsesRequest, + CodexTransport, +}; +use crate::upstream::openai::protocol::sse::SseError; +use crate::upstream::openai::protocol::websocket::{ + websocket_audit_artifact_from_attempt, websocket_payload_audit_snapshot, +}; + +use super::diagnostics::CodexUpstreamDiagnostics; +use super::endpoints::{endpoint_url, CODEX_RESPONSES_COMPACT_PATH, CODEX_RESPONSES_PATH}; +use super::headers::{ + build_ordered_codex_base_headers, insert_optional_header, insert_ordered_headers, + websocket_header_pairs, +}; +use super::response_meta; +use super::tls::{build_reqwest_client_with_custom_ca, custom_ca_env_cache_key, CustomCaError}; +use super::websocket::{ + execute_response_create_request_stream_with_pool, execute_response_create_request_with_pool, + write_websocket_audit_artifact_from_env, CodexWebSocketConnection, CodexWebSocketExchangeError, + CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketTurnStateUpdate, +}; +use super::websocket_pool::{ + CodexWebSocketPool, CodexWebSocketPoolKey, WebSocketPoolDecision, DEFAULT_FIRST_TOKEN_TIMEOUT, +}; + +// --------------------------------------------------------------------------- +// Constants +// --------------------------------------------------------------------------- + +const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 1024 * 1024; +const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str = + "x-codex-ws-stream-request-start-ms"; +type ReqwestClientCacheKey = (bool, Option); +type ReqwestClientCache = Mutex>; + +/// 构建带缓存的 reqwest Client。若 `force_http11` 为 true 则强制 HTTP/1.1。 +pub fn build_reqwest_client(force_http11: bool) -> Result { + let cache_key = (force_http11, custom_ca_env_cache_key()); + static CLIENTS: OnceLock = OnceLock::new(); + let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); + let mut clients = cache + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + if let Some(client) = clients.get(&cache_key) { + return Ok(client.clone()); + } + + let builder = Client::builder() + .use_rustls_tls() + .no_proxy() + .pool_max_idle_per_host(4) + .tcp_keepalive(Duration::from_secs(30)) + .gzip(true) + .brotli(true) + .zstd(true) + .deflate(true); + let builder = if force_http11 { + builder.http1_only() + } else { + builder + }; + let client = build_reqwest_client_with_custom_ca(builder)?; + clients.insert(cache_key, client.clone()); + Ok(client) +} + +// --------------------------------------------------------------------------- +// Error types +// --------------------------------------------------------------------------- + +/// Codex 上游 HTTP 客户端错误。 +#[derive(Debug, Error)] +pub enum CodexClientError { + /// Reqwest 传输失败。 + #[error("http transport error: {0}")] + Http(#[from] reqwest::Error), + /// 自定义 CA 构建失败。 + #[error("custom CA transport error: {0}")] + CustomCa(#[from] CustomCaError), + /// 请求头名字无效。 + #[error("invalid request header name: {0}")] + InvalidHeaderName(#[from] reqwest::header::InvalidHeaderName), + /// 请求头值无效。 + #[error("invalid request header value: {0}")] + InvalidHeaderValue(#[from] reqwest::header::InvalidHeaderValue), + /// SSE 响应解析失败。 + #[error("invalid upstream SSE response: {0}")] + InvalidSse(#[from] SseError), + /// WebSocket 请求编码失败。 + #[error("failed to encode websocket request: {0}")] + WebSocketEncode(#[source] serde_json::Error), + /// WebSocket 请求失败。 + #[error("websocket request failed: {0}")] + WebSocket(#[from] CodexWebSocketExchangeError), + /// 上游返回非成功响应。 + #[error("upstream returned status {status}: {body}")] + Upstream { + /// 上游状态码。 + status: StatusCode, + /// 上游错误体。 + body: String, + /// 推导出的重试秒数。 + retry_after_seconds: Option, + /// 上游诊断元数据。 + diagnostics: CodexUpstreamDiagnostics, + /// 上游透传的 `set-cookie` 列表。 + set_cookie_headers: Vec, + }, +} + +/// 判断上游错误正文是否表示账号已封禁或停用。 +pub fn is_banned_auth_signal(value: &str) -> bool { + let value = value.to_ascii_lowercase(); + value.contains("account_deactivated") + || value.contains("account deactivated") + || value.contains("account has been deactivated") + || value.contains("deactivated") + || value.contains("banned") +} + +/// 判断 Codex 上游错误是否表示账号已封禁或停用。 +pub fn is_banned_upstream_error(error: &CodexClientError) -> bool { + matches!( + error, + CodexClientError::Upstream { status, body, .. } + if status.as_u16() == 403 && !is_html_error_body(body) + ) +} + +fn is_html_error_body(value: &str) -> bool { + let value = value.trim_start().to_ascii_lowercase(); + value.starts_with(" = Result; + +/// Codex SSE 字节流。 +pub type CodexBackendSseStream = + Pin> + Send + 'static>>; + +/// 拉取上游模型目录时的请求上下文。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CodexModelCatalogRequest<'a> { + /// 当前账号访问令牌。 + pub access_token: &'a str, + /// 上游账号 ID。 + pub account_id: Option<&'a str>, + /// 请求 ID。 + pub request_id: &'a str, + /// Codex installation id。 + pub installation_id: Option<&'a str>, + /// 订阅计划类型。 + pub plan_type: &'a str, +} + +/// 上游模型目录客户端错误。 +#[derive(Debug, Error)] +pub enum CodexModelCatalogClientError { + /// 上游请求失败。 + #[error("model catalog request failed: {message}")] + RequestFailed { + /// 错误说明。 + message: String, + }, +} + +/// 上游模型目录客户端。 +#[async_trait] +pub trait CodexModelCatalogClient: Send + Sync + 'static { + /// 读取当前账号可见的上游模型目录。 + async fn fetch_models( + &self, + request: &CodexModelCatalogRequest<'_>, + ) -> Result, CodexModelCatalogClientError>; +} + +// --------------------------------------------------------------------------- +// Request context & response types +// --------------------------------------------------------------------------- + +/// 单次 Codex 上游请求的上下文。 +#[derive(Debug, Clone, Copy)] +pub struct CodexRequestContext<'a> { + /// Access token。 + pub access_token: &'a str, + /// ChatGPT account id。 + pub account_id: Option<&'a str>, + /// 请求 ID。 + pub request_id: &'a str, + /// x-codex-turn-state。 + pub turn_state: Option<&'a str>, + /// x-codex-turn-metadata。 + pub turn_metadata: Option<&'a str>, + /// x-codex-beta-features。 + pub beta_features: Option<&'a str>, + /// x-responsesapi-include-timing-metrics。 + pub include_timing_metrics: Option<&'a str>, + /// version。 + pub version: Option<&'a str>, + /// x-codex-window-id。 + pub codex_window_id: Option<&'a str>, + /// x-codex-parent-thread-id。 + pub parent_thread_id: Option<&'a str>, + /// cookie 头。 + pub cookie_header: Option<&'a str>, + /// x-codex-installation-id。 + pub installation_id: Option<&'a str>, + /// session_id。 + pub session_id: Option<&'a str>, +} + +/// Codex Responses 上游响应。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexBackendResponse { + /// 完整 SSE 文本。 + pub body: String, + /// 实际使用的上游传输。 + pub transport: CodexBackendTransport, + /// 从 SSE 中提取出的最终 usage。 + pub usage: Option, + /// 响应头里的最新 turn state。 + pub turn_state: Option, + /// 上游透传的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游透传的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// 首个有效上游 SSE/WebSocket 事件到达代理的耗时。 + pub first_token_ms: Option, + /// WebSocket 连接池决策。 + pub websocket_pool_decision: Option, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +/// Codex Responses 实际使用的上游传输。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexBackendTransport { + /// HTTP SSE transport. + HttpSse, + /// WebSocket transport. + WebSocket, +} + +/// Map a Codex Responses request to the concrete backend transport. +pub fn backend_transport_for_response_request( + request: &CodexResponsesRequest, +) -> CodexBackendTransport { + match transport_for_request(request) { + CodexTransport::HttpSse => CodexBackendTransport::HttpSse, + CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired => { + CodexBackendTransport::WebSocket + } + } +} + +/// Live stream rate-limit updates captured after the response headers. +pub type CodexRateLimitHeaderUpdates = CodexWebSocketRateLimitHeaderUpdates; + +/// Live stream turn-state updates captured after the response headers. +pub type CodexTurnStateUpdate = CodexWebSocketTurnStateUpdate; + +/// Codex Responses 上游 live SSE 响应。 +pub struct CodexBackendStreamingResponse { + /// 上游 SSE 字节流。 + pub body: CodexBackendSseStream, + /// 实际使用的上游传输。 + pub transport: CodexBackendTransport, + /// 响应头里的最新 turn state。 + pub turn_state: Option, + /// 上游透传的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游透传的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// live stream 期间捕获的限流头更新。 + pub rate_limit_header_updates: Option, + /// live stream 期间捕获的 turn-state 更新。 + pub turn_state_update: Option, + /// WebSocket 连接池决策。 + pub websocket_pool_decision: Option, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +/// Codex compact 端点响应。 +#[derive(Debug, Clone, PartialEq)] +pub struct CodexCompactResponse { + /// 上游返回的 JSON。 + pub body: Value, + /// 上游透传的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游透传的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +// --------------------------------------------------------------------------- +// CodexBackendClient +// --------------------------------------------------------------------------- + +/// Codex HTTP/SSE 上游客户端。 +#[derive(Clone)] +pub struct CodexBackendClient { + pub(super) client: Client, + pub(super) base_url: String, + fingerprint: RuntimeFingerprint, + websocket_pool: Option>, + websocket_first_token_timeout: Option, +} + +mod requests; + +#[async_trait] +impl CodexModelCatalogClient for CodexBackendClient { + async fn fetch_models( + &self, + request: &CodexModelCatalogRequest<'_>, + ) -> Result, CodexModelCatalogClientError> { + self.fetch_models_with_context(CodexRequestContext { + access_token: request.access_token, + account_id: request.account_id, + request_id: request.request_id, + turn_state: None, + turn_metadata: None, + beta_features: None, + include_timing_metrics: None, + version: None, + codex_window_id: None, + parent_thread_id: None, + cookie_header: None, + installation_id: request.installation_id, + session_id: None, + }) + .await + .map_err(|error| CodexModelCatalogClientError::RequestFailed { + message: error.to_string(), + }) + } +} + +// --------------------------------------------------------------------------- +// Response helpers +// --------------------------------------------------------------------------- + +fn log_websocket_pool_decision( + request_id: &str, + account_id: Option<&str>, + pool_context: Option<&WebSocketPoolLogContext>, + decision: Option, +) { + let Some(decision) = decision else { + return; + }; + let rid_short = request_id.chars().take(8).collect::(); + if let Some(reason) = decision.reason() { + tracing::info!( + request_id = %request_id, + rid = %rid_short, + account_id = account_id.unwrap_or_default(), + ws_pool = decision.kind(), + ws_pool_reason = reason, + conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()), + ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()), + "websocket pool decision" + ); + } else { + tracing::info!( + request_id = %request_id, + rid = %rid_short, + account_id = account_id.unwrap_or_default(), + ws_pool = decision.kind(), + conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()), + ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()), + "websocket pool decision" + ); + } +} + +#[derive(Debug, Clone)] +struct WebSocketPoolLogContext { + conversation_id_hash: String, + pool_key_hash: String, +} + +impl WebSocketPoolLogContext { + fn from_key(key: &CodexWebSocketPoolKey) -> Self { + Self { + conversation_id_hash: key.conversation_id_hash(), + pool_key_hash: key.stable_hash(), + } + } +} + +pub(super) fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Option { + headers + .get(RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.trim().parse::().ok()) + .filter(|seconds| *seconds > 0) + .or_else(|| body.and_then(retry_after_seconds_from_body)) +} + +pub(super) fn truncate_for_error(body: &str) -> String { + body.chars().take(200).collect() +} + +pub(super) async fn read_capped_error_body( + response: ReqwestResponse, +) -> Result { + let body = response.bytes().await?; + let len = body.len().min(MAX_UPSTREAM_ERROR_BODY_BYTES); + Ok(String::from_utf8_lossy(&body[..len]).into_owned()) +} + +// --------------------------------------------------------------------------- +// Model entry extraction +// --------------------------------------------------------------------------- + +fn extract_backend_model_entries(value: &Value) -> Vec { + let Some(models) = value + .pointer("/chat_models/models") + .or_else(|| value.get("models")) + .or_else(|| value.get("data")) + .or_else(|| value.get("categories")) + .and_then(Value::as_array) + else { + return Vec::new(); + }; + + let mut entries = Vec::new(); + for model in models { + if let Some(nested) = model.get("models").and_then(Value::as_array) { + entries.extend(nested.iter().filter_map(parse_backend_model_entry)); + } else if let Some(entry) = parse_backend_model_entry(model) { + entries.push(entry); + } + } + entries +} + +fn parse_backend_model_entry(value: &Value) -> Option { + let entry = serde_json::from_value::(value.clone()).ok()?; + (entry.slug.is_some() + || entry.id.is_some() + || entry.name.is_some() + || entry.display_name.is_some() + || entry.title.is_some()) + .then_some(entry) +} + +// --------------------------------------------------------------------------- +// Request helpers +// --------------------------------------------------------------------------- + +fn response_upstream_request( + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, +) -> CodexResponsesRequest { + let mut upstream = request.clone(); + if let Some(session_id) = context.session_id { + upstream.set_prompt_cache_key(Some(session_id.to_string())); + } + upstream.set_client_metadata(response_client_metadata(request.client_metadata(), context)); + upstream +} + +fn websocket_upstream_request(request: &CodexResponsesRequest) -> CodexResponsesRequest { + let mut request = request.clone(); + stamp_ws_stream_request_start_ms(&mut request); + request +} + +fn stamp_ws_stream_request_start_ms(request: &mut CodexResponsesRequest) { + let mut metadata = match request.client_metadata() { + Some(Value::Object(metadata)) => metadata.clone(), + _ => Map::new(), + }; + metadata.insert( + X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY.to_string(), + Value::String(now_unix_timestamp_millis().to_string()), + ); + request.set_client_metadata(Some(Value::Object(metadata))); +} + +fn now_unix_timestamp_millis() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_millis()) + .unwrap_or_default() +} + +fn response_client_metadata( + client_metadata: Option<&Value>, + context: CodexRequestContext<'_>, +) -> Option { + // 以客户端原始 client_metadata 为基础（保留 number/bool/object 等所有值类型）， + // 在其上追加代理自身的上下文字段。 + let mut metadata = match client_metadata { + Some(Value::Object(input)) => input.clone(), + _ => Map::new(), + }; + + insert_metadata_string( + &mut metadata, + "x-codex-installation-id", + context.installation_id, + ); + insert_metadata_string(&mut metadata, "session_id", context.session_id); + insert_metadata_string(&mut metadata, "thread_id", context.session_id); + insert_metadata_string(&mut metadata, "x-codex-window-id", context.codex_window_id); + insert_metadata_string( + &mut metadata, + "x-codex-turn-metadata", + context.turn_metadata, + ); + insert_metadata_string( + &mut metadata, + "x-codex-parent-thread-id", + context.parent_thread_id, + ); + + if metadata.is_empty() { + None + } else { + Some(Value::Object(metadata)) + } +} + +fn insert_metadata_string(metadata: &mut Map, key: &str, value: Option<&str>) { + if let Some(value) = value.filter(|value| !value.trim().is_empty()) { + metadata.insert(key.to_string(), Value::String(value.to_string())); + } +} + +fn openai_subagent_from_metadata(client_metadata: Option<&Value>) -> Option { + let value = client_metadata? + .as_object()? + .get("x-openai-subagent")? + .as_str()? + .trim(); + if matches!( + value, + "review" | "compact" | "memory_consolidation" | "collab_spawn" + ) { + Some(value.to_string()) + } else { + None + } +} + +// --------------------------------------------------------------------------- +// Error conversion helpers +// --------------------------------------------------------------------------- + +fn websocket_exchange_error_to_client_error( + error: CodexWebSocketExchangeError, +) -> CodexClientError { + match error { + CodexWebSocketExchangeError::Upstream(upstream) => { + let upstream = *upstream; + CodexClientError::Upstream { + status: StatusCode::from_u16(upstream.status_code) + .unwrap_or(StatusCode::BAD_GATEWAY), + body: upstream.body, + retry_after_seconds: upstream.retry_after_seconds, + diagnostics: upstream.diagnostics, + set_cookie_headers: upstream.set_cookie_headers, + } + } + error => CodexClientError::WebSocket(error), + } +} + +fn websocket_error_allows_http_fallback(error: &CodexClientError) -> bool { + !matches!(error, CodexClientError::Upstream { .. }) +} diff --git a/backend/src/upstream/transport/client.rs b/backend/src/upstream/openai/transport/client/requests.rs similarity index 54% rename from backend/src/upstream/transport/client.rs rename to backend/src/upstream/openai/transport/client/requests.rs index fab00759e..2188d815c 100644 --- a/backend/src/upstream/transport/client.rs +++ b/backend/src/upstream/openai/transport/client/requests.rs @@ -1,342 +1,4 @@ -//! Codex HTTP/SSE 上游客户端、请求头构造、TLS 与自定义 CA。 - -use std::{ - collections::HashMap, - pin::Pin, - sync::{Arc, Mutex, OnceLock}, - time::{Duration, Instant, SystemTime, UNIX_EPOCH}, -}; - -use async_trait::async_trait; -use bytes::Bytes; -use futures::{Stream, TryStreamExt}; -use reqwest::{ - header::{ - HeaderMap, HeaderName, HeaderValue, ACCEPT, ACCEPT_ENCODING, AUTHORIZATION, CONTENT_TYPE, - COOKIE, RETRY_AFTER, USER_AGENT, - }, - Client, Response as ReqwestResponse, StatusCode, -}; -use serde_json::{map::Map, Value}; -use thiserror::Error; -use tokio_tungstenite::tungstenite::handshake::client::generate_key; -use uuid::Uuid; - -use crate::upstream::fingerprint::{Fingerprint, RuntimeFingerprint}; -use crate::upstream::models::backend_entry::BackendModelEntry; -use crate::upstream::protocol::events::{extract_sse_usage, retry_after_seconds_from_body}; -use crate::upstream::protocol::responses::{ - http_sse_fallback_allowed, transport_for_request, CodexCompactRequest, CodexResponsesRequest, - CodexTransport, -}; -use crate::upstream::protocol::sse::SseError; -use crate::upstream::protocol::websocket::{ - websocket_audit_artifact_from_attempt, websocket_payload_audit_snapshot, -}; - -use super::diagnostics::CodexUpstreamDiagnostics; -use super::endpoints::{endpoint_url, CODEX_RESPONSES_COMPACT_PATH, CODEX_RESPONSES_PATH}; -use super::headers::{ - build_ordered_codex_base_headers, insert_optional_header, insert_ordered_headers, - websocket_header_pairs, -}; -use super::response_meta; -use super::tls::{build_reqwest_client_with_custom_ca, custom_ca_env_cache_key, CustomCaError}; -use super::websocket::{ - execute_response_create_request_stream_with_pool, execute_response_create_request_with_pool, - write_websocket_audit_artifact_from_env, CodexWebSocketConnection, CodexWebSocketExchangeError, - CodexWebSocketRateLimitHeaderUpdates, CodexWebSocketTurnStateUpdate, -}; -use super::websocket_pool::{ - CodexWebSocketPool, CodexWebSocketPoolKey, WebSocketPoolDecision, DEFAULT_FIRST_TOKEN_TIMEOUT, -}; - -// --------------------------------------------------------------------------- -// Constants -// --------------------------------------------------------------------------- - -const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 1024 * 1024; -const X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str = - "x-codex-ws-stream-request-start-ms"; -type ReqwestClientCacheKey = (bool, Option); -type ReqwestClientCache = Mutex>; - -/// 构建带缓存的 reqwest Client。若 `force_http11` 为 true 则强制 HTTP/1.1。 -pub fn build_reqwest_client(force_http11: bool) -> Result { - let cache_key = (force_http11, custom_ca_env_cache_key()); - static CLIENTS: OnceLock = OnceLock::new(); - let cache = CLIENTS.get_or_init(|| Mutex::new(HashMap::new())); - let mut clients = cache - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - - if let Some(client) = clients.get(&cache_key) { - return Ok(client.clone()); - } - - let builder = Client::builder() - .use_rustls_tls() - .no_proxy() - .pool_max_idle_per_host(4) - .tcp_keepalive(Duration::from_secs(30)) - .gzip(true) - .brotli(true) - .zstd(true) - .deflate(true); - let builder = if force_http11 { - builder.http1_only() - } else { - builder - }; - let client = build_reqwest_client_with_custom_ca(builder)?; - clients.insert(cache_key, client.clone()); - Ok(client) -} - -// --------------------------------------------------------------------------- -// Error types -// --------------------------------------------------------------------------- - -/// Codex 上游 HTTP 客户端错误。 -#[derive(Debug, Error)] -pub enum CodexClientError { - /// Reqwest 传输失败。 - #[error("http transport error: {0}")] - Http(#[from] reqwest::Error), - /// 自定义 CA 构建失败。 - #[error("custom CA transport error: {0}")] - CustomCa(#[from] CustomCaError), - /// 请求头名字无效。 - #[error("invalid request header name: {0}")] - InvalidHeaderName(#[from] reqwest::header::InvalidHeaderName), - /// 请求头值无效。 - #[error("invalid request header value: {0}")] - InvalidHeaderValue(#[from] reqwest::header::InvalidHeaderValue), - /// SSE 响应解析失败。 - #[error("invalid upstream SSE response: {0}")] - InvalidSse(#[from] SseError), - /// WebSocket 请求编码失败。 - #[error("failed to encode websocket request: {0}")] - WebSocketEncode(#[source] serde_json::Error), - /// WebSocket 请求失败。 - #[error("websocket request failed: {0}")] - WebSocket(#[from] CodexWebSocketExchangeError), - /// 上游返回非成功响应。 - #[error("upstream returned status {status}: {body}")] - Upstream { - /// 上游状态码。 - status: StatusCode, - /// 上游错误体。 - body: String, - /// 推导出的重试秒数。 - retry_after_seconds: Option, - /// 上游诊断元数据。 - diagnostics: CodexUpstreamDiagnostics, - /// 上游透传的 `set-cookie` 列表。 - set_cookie_headers: Vec, - }, -} - -/// 判断上游错误正文是否表示账号已封禁或停用。 -pub fn is_banned_auth_signal(value: &str) -> bool { - let value = value.to_ascii_lowercase(); - value.contains("account_deactivated") - || value.contains("account deactivated") - || value.contains("account has been deactivated") - || value.contains("deactivated") - || value.contains("banned") -} - -/// 判断 Codex 上游错误是否表示账号已封禁或停用。 -pub fn is_banned_upstream_error(error: &CodexClientError) -> bool { - matches!( - error, - CodexClientError::Upstream { status, body, .. } - if status.as_u16() == 403 && !is_html_error_body(body) - ) -} - -fn is_html_error_body(value: &str) -> bool { - let value = value.trim_start().to_ascii_lowercase(); - value.starts_with(" = Result; - -/// Codex SSE 字节流。 -pub type CodexBackendSseStream = - Pin> + Send + 'static>>; - -/// 拉取上游模型目录时的请求上下文。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct CodexModelCatalogRequest<'a> { - /// 当前账号访问令牌。 - pub access_token: &'a str, - /// 上游账号 ID。 - pub account_id: Option<&'a str>, - /// 请求 ID。 - pub request_id: &'a str, - /// Codex installation id。 - pub installation_id: Option<&'a str>, - /// 订阅计划类型。 - pub plan_type: &'a str, -} - -/// 上游模型目录客户端错误。 -#[derive(Debug, Error)] -pub enum CodexModelCatalogClientError { - /// 上游请求失败。 - #[error("model catalog request failed: {message}")] - RequestFailed { - /// 错误说明。 - message: String, - }, -} - -/// 上游模型目录客户端。 -#[async_trait] -pub trait CodexModelCatalogClient: Send + Sync + 'static { - /// 读取当前账号可见的上游模型目录。 - async fn fetch_models( - &self, - request: &CodexModelCatalogRequest<'_>, - ) -> Result, CodexModelCatalogClientError>; -} - -// --------------------------------------------------------------------------- -// Request context & response types -// --------------------------------------------------------------------------- - -/// 单次 Codex 上游请求的上下文。 -#[derive(Debug, Clone, Copy)] -pub struct CodexRequestContext<'a> { - /// Access token。 - pub access_token: &'a str, - /// ChatGPT account id。 - pub account_id: Option<&'a str>, - /// 请求 ID。 - pub request_id: &'a str, - /// x-codex-turn-state。 - pub turn_state: Option<&'a str>, - /// x-codex-turn-metadata。 - pub turn_metadata: Option<&'a str>, - /// x-codex-beta-features。 - pub beta_features: Option<&'a str>, - /// x-responsesapi-include-timing-metrics。 - pub include_timing_metrics: Option<&'a str>, - /// version。 - pub version: Option<&'a str>, - /// x-codex-window-id。 - pub codex_window_id: Option<&'a str>, - /// x-codex-parent-thread-id。 - pub parent_thread_id: Option<&'a str>, - /// cookie 头。 - pub cookie_header: Option<&'a str>, - /// x-codex-installation-id。 - pub installation_id: Option<&'a str>, - /// session_id。 - pub session_id: Option<&'a str>, -} - -/// Codex Responses 上游响应。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexBackendResponse { - /// 完整 SSE 文本。 - pub body: String, - /// 实际使用的上游传输。 - pub transport: CodexBackendTransport, - /// 从 SSE 中提取出的最终 usage。 - pub usage: Option, - /// 响应头里的最新 turn state。 - pub turn_state: Option, - /// 上游透传的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游透传的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// 首个有效上游 SSE/WebSocket 事件到达代理的耗时。 - pub first_token_ms: Option, - /// WebSocket 连接池决策。 - pub websocket_pool_decision: Option, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -/// Codex Responses 实际使用的上游传输。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum CodexBackendTransport { - /// HTTP SSE transport. - HttpSse, - /// WebSocket transport. - WebSocket, -} - -/// Map a Codex Responses request to the concrete backend transport. -pub fn backend_transport_for_response_request( - request: &CodexResponsesRequest, -) -> CodexBackendTransport { - match transport_for_request(request) { - CodexTransport::HttpSse => CodexBackendTransport::HttpSse, - CodexTransport::WebSocketPreferred | CodexTransport::WebSocketRequired => { - CodexBackendTransport::WebSocket - } - } -} - -/// Live stream rate-limit updates captured after the response headers. -pub type CodexRateLimitHeaderUpdates = CodexWebSocketRateLimitHeaderUpdates; - -/// Live stream turn-state updates captured after the response headers. -pub type CodexTurnStateUpdate = CodexWebSocketTurnStateUpdate; - -/// Codex Responses 上游 live SSE 响应。 -pub struct CodexBackendStreamingResponse { - /// 上游 SSE 字节流。 - pub body: CodexBackendSseStream, - /// 实际使用的上游传输。 - pub transport: CodexBackendTransport, - /// 响应头里的最新 turn state。 - pub turn_state: Option, - /// 上游透传的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游透传的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// live stream 期间捕获的限流头更新。 - pub rate_limit_header_updates: Option, - /// live stream 期间捕获的 turn-state 更新。 - pub turn_state_update: Option, - /// WebSocket 连接池决策。 - pub websocket_pool_decision: Option, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -/// Codex compact 端点响应。 -#[derive(Debug, Clone, PartialEq)] -pub struct CodexCompactResponse { - /// 上游返回的 JSON。 - pub body: Value, - /// 上游透传的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游透传的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -// --------------------------------------------------------------------------- -// CodexBackendClient -// --------------------------------------------------------------------------- - -/// Codex HTTP/SSE 上游客户端。 -#[derive(Clone)] -pub struct CodexBackendClient { - pub(super) client: Client, - pub(super) base_url: String, - fingerprint: RuntimeFingerprint, - websocket_pool: Option>, - websocket_first_token_timeout: Option, -} +use super::*; impl CodexBackendClient { /// 构造客户端。 @@ -798,7 +460,7 @@ impl CodexBackendClient { } /// 获取后端模型目录条目。 - async fn fetch_models_with_context( + pub(super) async fn fetch_models_with_context( &self, context: CodexRequestContext<'_>, ) -> CodexClientResult> { @@ -924,7 +586,7 @@ impl CodexBackendClient { Ok(headers) } - pub(super) fn usage_request_headers( + pub(in crate::upstream::openai::transport) fn usage_request_headers( &self, context: CodexRequestContext<'_>, ) -> CodexClientResult { @@ -984,268 +646,3 @@ impl CodexBackendClient { Ok(headers) } } - -#[async_trait] -impl CodexModelCatalogClient for CodexBackendClient { - async fn fetch_models( - &self, - request: &CodexModelCatalogRequest<'_>, - ) -> Result, CodexModelCatalogClientError> { - self.fetch_models_with_context(CodexRequestContext { - access_token: request.access_token, - account_id: request.account_id, - request_id: request.request_id, - turn_state: None, - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - version: None, - codex_window_id: None, - parent_thread_id: None, - cookie_header: None, - installation_id: request.installation_id, - session_id: None, - }) - .await - .map_err(|error| CodexModelCatalogClientError::RequestFailed { - message: error.to_string(), - }) - } -} - -// --------------------------------------------------------------------------- -// Response helpers -// --------------------------------------------------------------------------- - -fn log_websocket_pool_decision( - request_id: &str, - account_id: Option<&str>, - pool_context: Option<&WebSocketPoolLogContext>, - decision: Option, -) { - let Some(decision) = decision else { - return; - }; - let rid_short = request_id.chars().take(8).collect::(); - if let Some(reason) = decision.reason() { - tracing::info!( - request_id = %request_id, - rid = %rid_short, - account_id = account_id.unwrap_or_default(), - ws_pool = decision.kind(), - ws_pool_reason = reason, - conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()), - ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()), - "websocket pool decision" - ); - } else { - tracing::info!( - request_id = %request_id, - rid = %rid_short, - account_id = account_id.unwrap_or_default(), - ws_pool = decision.kind(), - conversation_id_hash = pool_context.map_or("", |context| context.conversation_id_hash.as_str()), - ws_pool_key_hash = pool_context.map_or("", |context| context.pool_key_hash.as_str()), - "websocket pool decision" - ); - } -} - -#[derive(Debug, Clone)] -struct WebSocketPoolLogContext { - conversation_id_hash: String, - pool_key_hash: String, -} - -impl WebSocketPoolLogContext { - fn from_key(key: &CodexWebSocketPoolKey) -> Self { - Self { - conversation_id_hash: key.conversation_id_hash(), - pool_key_hash: key.stable_hash(), - } - } -} - -pub(super) fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Option { - headers - .get(RETRY_AFTER) - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.trim().parse::().ok()) - .filter(|seconds| *seconds > 0) - .or_else(|| body.and_then(retry_after_seconds_from_body)) -} - -pub(super) fn truncate_for_error(body: &str) -> String { - body.chars().take(200).collect() -} - -pub(super) async fn read_capped_error_body( - response: ReqwestResponse, -) -> Result { - let body = response.bytes().await?; - let len = body.len().min(MAX_UPSTREAM_ERROR_BODY_BYTES); - Ok(String::from_utf8_lossy(&body[..len]).into_owned()) -} - -// --------------------------------------------------------------------------- -// Model entry extraction -// --------------------------------------------------------------------------- - -fn extract_backend_model_entries(value: &Value) -> Vec { - let Some(models) = value - .pointer("/chat_models/models") - .or_else(|| value.get("models")) - .or_else(|| value.get("data")) - .or_else(|| value.get("categories")) - .and_then(Value::as_array) - else { - return Vec::new(); - }; - - let mut entries = Vec::new(); - for model in models { - if let Some(nested) = model.get("models").and_then(Value::as_array) { - entries.extend(nested.iter().filter_map(parse_backend_model_entry)); - } else if let Some(entry) = parse_backend_model_entry(model) { - entries.push(entry); - } - } - entries -} - -fn parse_backend_model_entry(value: &Value) -> Option { - let entry = serde_json::from_value::(value.clone()).ok()?; - (entry.slug.is_some() - || entry.id.is_some() - || entry.name.is_some() - || entry.display_name.is_some() - || entry.title.is_some()) - .then_some(entry) -} - -// --------------------------------------------------------------------------- -// Request helpers -// --------------------------------------------------------------------------- - -fn response_upstream_request( - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, -) -> CodexResponsesRequest { - let mut upstream = request.clone(); - if let Some(session_id) = context.session_id { - upstream.set_prompt_cache_key(Some(session_id.to_string())); - } - upstream.set_client_metadata(response_client_metadata(request.client_metadata(), context)); - upstream -} - -fn websocket_upstream_request(request: &CodexResponsesRequest) -> CodexResponsesRequest { - let mut request = request.clone(); - stamp_ws_stream_request_start_ms(&mut request); - request -} - -fn stamp_ws_stream_request_start_ms(request: &mut CodexResponsesRequest) { - let mut metadata = match request.client_metadata() { - Some(Value::Object(metadata)) => metadata.clone(), - _ => Map::new(), - }; - metadata.insert( - X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY.to_string(), - Value::String(now_unix_timestamp_millis().to_string()), - ); - request.set_client_metadata(Some(Value::Object(metadata))); -} - -fn now_unix_timestamp_millis() -> u128 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_millis()) - .unwrap_or_default() -} - -fn response_client_metadata( - client_metadata: Option<&Value>, - context: CodexRequestContext<'_>, -) -> Option { - // 以客户端原始 client_metadata 为基础（保留 number/bool/object 等所有值类型）， - // 在其上追加代理自身的上下文字段。 - let mut metadata = match client_metadata { - Some(Value::Object(input)) => input.clone(), - _ => Map::new(), - }; - - insert_metadata_string( - &mut metadata, - "x-codex-installation-id", - context.installation_id, - ); - insert_metadata_string(&mut metadata, "session_id", context.session_id); - insert_metadata_string(&mut metadata, "thread_id", context.session_id); - insert_metadata_string(&mut metadata, "x-codex-window-id", context.codex_window_id); - insert_metadata_string( - &mut metadata, - "x-codex-turn-metadata", - context.turn_metadata, - ); - insert_metadata_string( - &mut metadata, - "x-codex-parent-thread-id", - context.parent_thread_id, - ); - - if metadata.is_empty() { - None - } else { - Some(Value::Object(metadata)) - } -} - -fn insert_metadata_string(metadata: &mut Map, key: &str, value: Option<&str>) { - if let Some(value) = value.filter(|value| !value.trim().is_empty()) { - metadata.insert(key.to_string(), Value::String(value.to_string())); - } -} - -fn openai_subagent_from_metadata(client_metadata: Option<&Value>) -> Option { - let value = client_metadata? - .as_object()? - .get("x-openai-subagent")? - .as_str()? - .trim(); - if matches!( - value, - "review" | "compact" | "memory_consolidation" | "collab_spawn" - ) { - Some(value.to_string()) - } else { - None - } -} - -// --------------------------------------------------------------------------- -// Error conversion helpers -// --------------------------------------------------------------------------- - -fn websocket_exchange_error_to_client_error( - error: CodexWebSocketExchangeError, -) -> CodexClientError { - match error { - CodexWebSocketExchangeError::Upstream(upstream) => { - let upstream = *upstream; - CodexClientError::Upstream { - status: StatusCode::from_u16(upstream.status_code) - .unwrap_or(StatusCode::BAD_GATEWAY), - body: upstream.body, - retry_after_seconds: upstream.retry_after_seconds, - diagnostics: upstream.diagnostics, - set_cookie_headers: upstream.set_cookie_headers, - } - } - error => CodexClientError::WebSocket(error), - } -} - -fn websocket_error_allows_http_fallback(error: &CodexClientError) -> bool { - !matches!(error, CodexClientError::Upstream { .. }) -} diff --git a/backend/src/upstream/transport/diagnostics.rs b/backend/src/upstream/openai/transport/diagnostics.rs similarity index 100% rename from backend/src/upstream/transport/diagnostics.rs rename to backend/src/upstream/openai/transport/diagnostics.rs diff --git a/backend/src/upstream/transport/endpoints.rs b/backend/src/upstream/openai/transport/endpoints.rs similarity index 100% rename from backend/src/upstream/transport/endpoints.rs rename to backend/src/upstream/openai/transport/endpoints.rs diff --git a/backend/src/upstream/transport/headers.rs b/backend/src/upstream/openai/transport/headers.rs similarity index 98% rename from backend/src/upstream/transport/headers.rs rename to backend/src/upstream/openai/transport/headers.rs index 99adc42e7..3faf42a17 100644 --- a/backend/src/upstream/transport/headers.rs +++ b/backend/src/upstream/openai/transport/headers.rs @@ -1,7 +1,7 @@ use indexmap::IndexMap; use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; -use crate::upstream::fingerprint::Fingerprint; +use crate::upstream::openai::fingerprint::Fingerprint; use super::client::CodexClientResult; diff --git a/backend/src/upstream/transport/mod.rs b/backend/src/upstream/openai/transport/mod.rs similarity index 100% rename from backend/src/upstream/transport/mod.rs rename to backend/src/upstream/openai/transport/mod.rs diff --git a/backend/src/upstream/transport/response_meta.rs b/backend/src/upstream/openai/transport/response_meta.rs similarity index 95% rename from backend/src/upstream/transport/response_meta.rs rename to backend/src/upstream/openai/transport/response_meta.rs index 61e23a47c..d63f7458e 100644 --- a/backend/src/upstream/transport/response_meta.rs +++ b/backend/src/upstream/openai/transport/response_meta.rs @@ -4,7 +4,7 @@ use std::time::Instant; use reqwest::header::{HeaderMap, SET_COOKIE}; -use crate::upstream::protocol::responses::update_first_response_event_ms; +use crate::upstream::openai::protocol::responses::update_first_response_event_ms; use super::diagnostics::CodexUpstreamDiagnostics; diff --git a/backend/src/upstream/transport/tls.rs b/backend/src/upstream/openai/transport/tls.rs similarity index 100% rename from backend/src/upstream/transport/tls.rs rename to backend/src/upstream/openai/transport/tls.rs diff --git a/backend/src/upstream/transport/usage.rs b/backend/src/upstream/openai/transport/usage.rs similarity index 96% rename from backend/src/upstream/transport/usage.rs rename to backend/src/upstream/openai/transport/usage.rs index e58085d08..bf4408092 100644 --- a/backend/src/upstream/transport/usage.rs +++ b/backend/src/upstream/openai/transport/usage.rs @@ -1,7 +1,7 @@ use reqwest::StatusCode; use serde_json::Value; -use crate::upstream::protocol::events::retry_after_seconds_from_body; +use crate::upstream::openai::protocol::events::retry_after_seconds_from_body; use super::{ client::{read_capped_error_body, retry_after_seconds, truncate_for_error}, diff --git a/backend/src/upstream/openai/transport/websocket.rs b/backend/src/upstream/openai/transport/websocket.rs new file mode 100644 index 000000000..a47e49fc0 --- /dev/null +++ b/backend/src/upstream/openai/transport/websocket.rs @@ -0,0 +1,729 @@ +//! Codex WebSocket 连接建立（关键函数）。 + +use std::{ + collections::VecDeque, + fmt, io, + path::{Path, PathBuf}, + pin::Pin, + sync::Arc, + time::{Duration, Instant}, +}; + +use bytes::Bytes; +use chrono::Utc; +use futures::Stream; +use thiserror::Error; +use tokio::sync::mpsc; +use tokio::{sync::Mutex, time::timeout}; +use tokio_tungstenite::{connect_async_tls_with_config, Connector}; +use tungstenite::{ + self, + extensions::{compression::deflate::DeflateConfig, ExtensionsConfig}, + handshake::client::Request as WsRequest, + http::Response as WsResponse, + protocol::WebSocketConfig, + Message, +}; + +use crate::infra::time::china_filename_timestamp_millis; +use crate::upstream::openai::protocol::events::{self, TokenUsage}; +use crate::upstream::openai::protocol::responses::CodexResponsesRequest; +use crate::upstream::openai::protocol::sse::SseError; +use crate::upstream::openai::protocol::websocket::{ + classify_websocket_error_frame, is_pre_content_lifecycle_event, is_terminal_websocket_event, + retry_after_seconds_from_wrapped_error_headers, websocket_event_to_sse_frame, + websocket_event_type, websocket_incomplete_response_reason, websocket_metadata_turn_state, + websocket_response_completed_parse_error, websocket_response_create_payload_text, + OpeningAuditHeader, OpeningAuditSnapshot, WebSocketAuditArtifact, +}; + +use super::websocket_pool::{ + CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, + PooledWebSocketConnection, WebSocketPoolAcquire, WebSocketPoolDecision, +}; +use super::websocket_pump::{PumpKeepalive, PumpedWebSocket, RawWsStream}; +use super::{diagnostics::CodexUpstreamDiagnostics, response_meta}; +use uuid::Uuid; + +const REDACTED_HEADER_VALUE: &str = ""; +const CODEX_RESPONSES_PATH: &str = "/codex/responses"; +const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits"; +const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20); +const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60); +const WEBSOCKET_STREAM_BUFFER: usize = 16; +const WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS: usize = 2; +/// WebSocket audit artifact 输出目录环境变量。 +pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR"; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +/// WebSocket 连接适配器。 +mod stream; +mod types; + +pub use stream::execute_response_create_request; +use stream::{ + audit_header_value, collect_websocket_response, is_first_token_timeout, + prefetch_first_content_stream_frames, reusable_websocket_metadata, + reused_stream_prefetch_error, stream_websocket_response, websocket_audit_file_name, + websocket_connection_metadata, WebSocketStreamPoolReturn, +}; +pub use types::*; + +impl CodexWebSocketRequest { + /// 返回连接描述。 + pub fn connection(&self) -> &CodexWebSocketConnection { + &self.connection + } + + /// 返回将要发送的首个文本帧。 + pub fn payload_text(&self) -> &str { + &self.payload_text + } +} + +impl CodexWebSocketConnection { + /// 构造待打开的 WebSocket 连接描述。 + pub fn new(endpoint: impl Into, headers: Vec<(String, String)>) -> Self { + Self { + endpoint: endpoint.into(), + headers, + } + } + + /// 构造 Responses WebSocket 连接描述。 + pub fn responses( + base_url: &str, + websocket_key: &str, + business_headers: Vec<(String, String)>, + ) -> Self { + let endpoint = responses_websocket_endpoint(base_url); + let mut headers = Vec::new(); + if let Some(host) = websocket_host_header(&endpoint) { + headers.push(("Host".to_string(), host)); + } + headers.extend([ + ("Connection".to_string(), "Upgrade".to_string()), + ("Upgrade".to_string(), "websocket".to_string()), + ("Sec-WebSocket-Version".to_string(), "13".to_string()), + ("Sec-WebSocket-Key".to_string(), websocket_key.to_string()), + ]); + headers.extend(business_headers); + headers.push(( + "sec-websocket-extensions".to_string(), + WEBSOCKET_EXTENSIONS.to_string(), + )); + Self { endpoint, headers } + } + + /// 构造 Responses WebSocket opening 与首个 `response.create` 文本帧。 + pub fn responses_create_request( + base_url: &str, + websocket_key: &str, + business_headers: Vec<(String, String)>, + request: &CodexResponsesRequest, + ) -> Result { + Ok(CodexWebSocketRequest { + connection: Self::responses(base_url, websocket_key, business_headers), + payload_text: websocket_response_create_payload_text(request)?, + }) + } + + /// 返回 WebSocket endpoint。 + pub fn endpoint(&self) -> &str { + &self.endpoint + } + + /// 返回按发送顺序保存的请求头。 + pub fn headers(&self) -> &[(String, String)] { + &self.headers + } + + /// 生成打开握手审计快照。 + pub fn opening_audit_snapshot(&self) -> OpeningAuditSnapshot { + OpeningAuditSnapshot { + request_line: request_line_for_endpoint(&self.endpoint), + header_order: self.headers.iter().map(|(name, _)| name.clone()).collect(), + headers: self + .headers + .iter() + .map(|(name, value)| OpeningAuditHeader { + name: name.clone(), + value: audit_header_value(name, value), + }) + .collect(), + } + } +} + +// --------------------------------------------------------------------------- +// WebSocket endpoint helpers +// --------------------------------------------------------------------------- + +/// 将 Codex backend base URL 转换为 Responses WebSocket endpoint。 +pub fn responses_websocket_endpoint(base_url: &str) -> String { + let endpoint = format!("{}{}", base_url.trim_end_matches('/'), CODEX_RESPONSES_PATH); + if let Some(rest) = endpoint.strip_prefix("https://") { + format!("wss://{rest}") + } else if let Some(rest) = endpoint.strip_prefix("http://") { + format!("ws://{rest}") + } else { + endpoint + } +} + +fn websocket_host_header(endpoint: &str) -> Option { + let url = reqwest::Url::parse(endpoint).ok()?; + let host = url.host_str()?; + Some(match url.port() { + Some(port) => format!("{host}:{port}"), + None => host.to_string(), + }) +} + +// --------------------------------------------------------------------------- +// Connection +// --------------------------------------------------------------------------- + +async fn connect_websocket( + connection: &CodexWebSocketConnection, +) -> Result<(RawWsStream, WsResponse>>), CodexWebSocketExchangeError> { + let request = websocket_handshake_request(connection)?; + let connector = super::tls::maybe_build_rustls_client_config_with_custom_ca() + .map_err(|error| tungstenite::Error::Io(std::io::Error::other(error)))? + .map(Connector::Rustls); + match connect_async_tls_with_config(request, Some(websocket_config()), false, connector).await { + Ok((websocket, response)) => Ok((websocket, response)), + Err(tungstenite::Error::Http(response)) => Err(websocket_opening_error(response.as_ref())), + Err(error) => Err(error.into()), + } +} + +/// 建立连接并交给后台 pump 托管；`keepalive` 决定该连接是否做主动保活。 +/// +/// 非池化（即用即弃）连接用 [`PumpKeepalive::disabled`]；池化连接用连接池派生的保活策略， +/// 让空闲期在后台完成 ping/pong 与失活检测，从而复用前可零成本判活。 +async fn connect_pumped_websocket( + connection: &CodexWebSocketConnection, + keepalive: PumpKeepalive, +) -> Result<(PumpedWebSocket, WsResponse>>), CodexWebSocketExchangeError> { + let (raw, response) = connect_websocket(connection).await?; + Ok((PumpedWebSocket::new(raw, keepalive), response)) +} + +fn websocket_handshake_request( + connection: &CodexWebSocketConnection, +) -> Result { + let mut builder = WsRequest::builder() + .method("GET") + .uri(connection.endpoint()); + for (name, value) in connection.headers() { + if name.eq_ignore_ascii_case("sec-websocket-extensions") { + continue; + } + builder = builder.header(name.as_str(), value.as_str()); + } + builder.body(()) +} + +fn websocket_config() -> WebSocketConfig { + let mut extensions = ExtensionsConfig::default(); + extensions.permessage_deflate = Some(DeflateConfig::default()); + + let mut config = WebSocketConfig::default(); + config.extensions = extensions; + config +} + +fn websocket_opening_error(response: &WsResponse>>) -> CodexWebSocketExchangeError { + let status_code = response.status().as_u16(); + let body = response + .body() + .as_ref() + .map(|body| String::from_utf8_lossy(body).into_owned()) + .unwrap_or_default(); + let retry_after_seconds = response + .headers() + .get("retry-after") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + .filter(|seconds| *seconds > 0) + .or_else(|| events::retry_after_seconds_from_body(&body)); + CodexWebSocketExchangeError::upstream( + status_code, + retry_after_seconds, + body, + response_meta::set_cookie_headers(response.headers()), + response_meta::diagnostics(Some(status_code), response.headers()), + ) +} + +fn request_line_for_endpoint(endpoint: &str) -> String { + let path = reqwest::Url::parse(endpoint) + .ok() + .map(|url| { + let mut path = url.path().to_string(); + if let Some(query) = url.query() { + path.push('?'); + path.push_str(query); + } + path + }) + .filter(|path| !path.is_empty()) + .unwrap_or_else(|| endpoint.to_string()); + format!("GET {path} HTTP/1.1") +} + +// --------------------------------------------------------------------------- +// Fresh request execution +// --------------------------------------------------------------------------- + +async fn execute_fresh_response_create_request( + request: &CodexWebSocketRequest, + started_at: Instant, + first_token_timeout: Option, +) -> Result { + let (websocket, response) = + connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?; + let first_token_started_at = Instant::now(); + websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await?; + + let metadata = websocket_connection_metadata(&response); + let (exchange, _websocket, _metadata) = collect_websocket_response( + websocket, + metadata, + false, + started_at, + first_token_started_at, + first_token_timeout, + ) + .await?; + Ok(exchange) +} + +async fn execute_fresh_response_create_request_stream( + request: &CodexWebSocketRequest, + first_token_timeout: Option, +) -> Result { + let (mut websocket, response) = + connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?; + let first_token_started_at = Instant::now(); + websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await?; + + let mut metadata = websocket_connection_metadata(&response); + let prefetched_frames = match prefetch_first_content_stream_frames( + &mut websocket, + &mut metadata, + first_token_started_at, + first_token_timeout, + ) + .await + { + Ok(prefetched_frames) => prefetched_frames, + Err(error) => { + websocket.close().await; + return Err(error); + } + }; + Ok(stream_websocket_response( + websocket, + metadata, + None, + prefetched_frames, + )) +} + +async fn execute_fresh_response_create_request_with_retries( + request: &CodexWebSocketRequest, + started_at: Instant, + first_token_timeout: Option, + attempts: usize, +) -> Result { + let attempts = attempts.max(1); + for attempt in 1..=attempts { + match execute_fresh_response_create_request(request, started_at, first_token_timeout).await + { + Ok(exchange) => return Ok(exchange), + Err(error) if attempt < attempts && is_first_token_timeout(&error) => continue, + Err(error) => return Err(error), + } + } + unreachable!("fresh websocket retry loop always returns"); +} + +async fn execute_fresh_response_create_request_stream_with_retries( + request: &CodexWebSocketRequest, + first_token_timeout: Option, + attempts: usize, +) -> Result { + let attempts = attempts.max(1); + for attempt in 1..=attempts { + match execute_fresh_response_create_request_stream(request, first_token_timeout).await { + Ok(exchange) => return Ok(exchange), + Err(error) if attempt < attempts && is_first_token_timeout(&error) => continue, + Err(error) => return Err(error), + } + } + unreachable!("fresh websocket stream retry loop always returns"); +} + +// --------------------------------------------------------------------------- +// Pool-aware execution +// --------------------------------------------------------------------------- + +pub(crate) async fn execute_response_create_request_with_pool( + request: &CodexWebSocketRequest, + pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>, + started_at: Instant, + fallback_first_token_timeout: Option, +) -> Result { + let Some((pool, key)) = pool else { + return execute_fresh_response_create_request_with_retries( + request, + started_at, + fallback_first_token_timeout, + WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, + ) + .await; + }; + + match pool.acquire(&key).await { + WebSocketPoolAcquire::Reused(connection) => { + let result = + execute_pooled_response_create_request(request, pool, key, *connection, started_at) + .await; + match result { + Err(CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { + .. + }) + | Err(CodexWebSocketExchangeError::FirstTokenTimeout { .. }) => { + let mut exchange = execute_fresh_response_create_request_with_retries( + request, + started_at, + pool.first_token_timeout(), + WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, + ) + .await?; + exchange.pool_decision = Some(WebSocketPoolDecision::retry_after_stale_reuse()); + Ok(exchange) + } + Ok(mut exchange) => { + exchange.pool_decision = Some(WebSocketPoolDecision::reuse()); + Ok(exchange) + } + Err(error) => Err(error), + } + } + WebSocketPoolAcquire::FreshReserved => { + let result = execute_fresh_pooled_response_create_request( + request, + pool, + key.clone(), + started_at, + ) + .await; + match result { + Ok(mut exchange) => { + exchange.pool_decision = Some(WebSocketPoolDecision::new()); + Ok(exchange) + } + Err(error) if is_first_token_timeout(&error) => { + pool.discard(&key).await; + let mut exchange = execute_fresh_response_create_request_with_retries( + request, + started_at, + pool.first_token_timeout(), + 1, + ) + .await?; + exchange.pool_decision = Some(WebSocketPoolDecision::new()); + Ok(exchange) + } + Err(error) => { + pool.discard(&key).await; + Err(error) + } + } + } + WebSocketPoolAcquire::Bypass(reason) => { + let mut exchange = execute_fresh_response_create_request_with_retries( + request, + started_at, + pool.first_token_timeout(), + WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, + ) + .await?; + exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason)); + Ok(exchange) + } + } +} + +pub(crate) async fn execute_response_create_request_stream_with_pool( + request: &CodexWebSocketRequest, + pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>, + fallback_first_token_timeout: Option, +) -> Result { + let Some((pool, key)) = pool else { + return execute_fresh_response_create_request_stream_with_retries( + request, + fallback_first_token_timeout, + WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, + ) + .await; + }; + + match pool.acquire(&key).await { + WebSocketPoolAcquire::Reused(connection) => { + let result = execute_pooled_response_create_request_stream( + request, + pool.clone(), + key, + *connection, + ) + .await; + match result { + Err(CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { + .. + }) + | Err(CodexWebSocketExchangeError::FirstTokenTimeout { .. }) => { + let mut exchange = execute_fresh_response_create_request_stream_with_retries( + request, + pool.first_token_timeout(), + WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, + ) + .await?; + exchange.pool_decision = Some(WebSocketPoolDecision::retry_after_stale_reuse()); + Ok(exchange) + } + Ok(mut exchange) => { + exchange.pool_decision = Some(WebSocketPoolDecision::reuse()); + Ok(exchange) + } + Err(error) => Err(error), + } + } + WebSocketPoolAcquire::FreshReserved => { + let result = execute_fresh_pooled_response_create_request_stream( + request, + pool.clone(), + key.clone(), + ) + .await; + match result { + Ok(mut exchange) => { + exchange.pool_decision = Some(WebSocketPoolDecision::new()); + Ok(exchange) + } + Err(error) if is_first_token_timeout(&error) => { + pool.discard(&key).await; + let mut exchange = execute_fresh_response_create_request_stream_with_retries( + request, + pool.first_token_timeout(), + 1, + ) + .await?; + exchange.pool_decision = Some(WebSocketPoolDecision::new()); + Ok(exchange) + } + Err(error) => { + pool.discard(&key).await; + Err(error) + } + } + } + WebSocketPoolAcquire::Bypass(reason) => { + let mut exchange = execute_fresh_response_create_request_stream_with_retries( + request, + pool.first_token_timeout(), + WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, + ) + .await?; + exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason)); + Ok(exchange) + } + } +} + +async fn execute_fresh_pooled_response_create_request( + request: &CodexWebSocketRequest, + pool: &CodexWebSocketPool, + key: CodexWebSocketPoolKey, + started_at: Instant, +) -> Result { + let (websocket, response) = + connect_pumped_websocket(request.connection(), pool.keepalive()).await?; + let first_token_started_at = Instant::now(); + websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await?; + + let now = Instant::now(); + let metadata = websocket_connection_metadata(&response); + let (exchange, websocket, metadata) = collect_websocket_response( + websocket, + metadata, + false, + started_at, + first_token_started_at, + pool.first_token_timeout(), + ) + .await?; + pool.put( + key, + PooledWebSocketConnection { + websocket, + metadata: reusable_websocket_metadata(metadata), + created_at: now, + }, + ) + .await; + Ok(exchange) +} + +async fn execute_fresh_pooled_response_create_request_stream( + request: &CodexWebSocketRequest, + pool: CodexWebSocketPool, + key: CodexWebSocketPoolKey, +) -> Result { + let (mut websocket, response) = + connect_pumped_websocket(request.connection(), pool.keepalive()).await?; + let first_token_started_at = Instant::now(); + websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await?; + + let now = Instant::now(); + let mut metadata = websocket_connection_metadata(&response); + let prefetched_frames = match prefetch_first_content_stream_frames( + &mut websocket, + &mut metadata, + first_token_started_at, + pool.first_token_timeout(), + ) + .await + { + Ok(prefetched_frames) => prefetched_frames, + Err(error) => { + websocket.close().await; + return Err(error); + } + }; + Ok(stream_websocket_response( + websocket, + metadata, + Some(WebSocketStreamPoolReturn { + pool, + key, + created_at: now, + }), + prefetched_frames, + )) +} + +async fn execute_pooled_response_create_request( + request: &CodexWebSocketRequest, + pool: &CodexWebSocketPool, + key: CodexWebSocketPoolKey, + connection: PooledWebSocketConnection, + started_at: Instant, +) -> Result { + let first_token_started_at = Instant::now(); + if let Err(error) = connection + .websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await + { + pool.discard(&key).await; + return Err( + CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { + message: error.to_string(), + }, + ); + } + + let created_at = connection.created_at; + match collect_websocket_response( + connection.websocket, + reusable_websocket_metadata(connection.metadata), + true, + started_at, + first_token_started_at, + pool.first_token_timeout(), + ) + .await + { + Ok((exchange, websocket, metadata)) => { + pool.put( + key, + PooledWebSocketConnection { + websocket, + metadata: reusable_websocket_metadata(metadata), + created_at, + }, + ) + .await; + Ok(exchange) + } + Err(error) => { + pool.discard(&key).await; + Err(error) + } + } +} + +async fn execute_pooled_response_create_request_stream( + request: &CodexWebSocketRequest, + pool: CodexWebSocketPool, + key: CodexWebSocketPoolKey, + connection: PooledWebSocketConnection, +) -> Result { + let mut websocket = connection.websocket; + let mut metadata = reusable_websocket_metadata(connection.metadata); + let created_at = connection.created_at; + let first_token_started_at = Instant::now(); + if let Err(error) = websocket + .send(Message::Text(request.payload_text().to_string().into())) + .await + { + pool.discard(&key).await; + return Err( + CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { + message: error.to_string(), + }, + ); + } + + let prefetched_frames = match prefetch_first_content_stream_frames( + &mut websocket, + &mut metadata, + first_token_started_at, + pool.first_token_timeout(), + ) + .await + { + Ok(prefetched_frames) => prefetched_frames, + Err(error) => { + pool.discard(&key).await; + websocket.close().await; + return Err(reused_stream_prefetch_error(error)); + } + }; + + Ok(stream_websocket_response( + websocket, + metadata, + Some(WebSocketStreamPoolReturn { + pool, + key, + created_at, + }), + prefetched_frames, + )) +} diff --git a/backend/src/upstream/openai/transport/websocket/stream.rs b/backend/src/upstream/openai/transport/websocket/stream.rs new file mode 100644 index 000000000..1e7e8df07 --- /dev/null +++ b/backend/src/upstream/openai/transport/websocket/stream.rs @@ -0,0 +1,564 @@ +use super::*; + +/// 执行一次 prepared Responses WebSocket 请求并聚合为 SSE。 +pub async fn execute_response_create_request( + request: &CodexWebSocketRequest, +) -> Result { + execute_response_create_request_with_pool(request, None, Instant::now(), None).await +} + +// --------------------------------------------------------------------------- +// Response collection +// --------------------------------------------------------------------------- + +pub(super) async fn collect_websocket_response( + mut websocket: PumpedWebSocket, + mut metadata: CodexWebSocketConnectionMetadata, + reused_connection: bool, + started_at: Instant, + first_token_started_at: Instant, + first_token_timeout: Option, +) -> Result< + ( + CodexWebSocketExchange, + PumpedWebSocket, + CodexWebSocketConnectionMetadata, + ), + CodexWebSocketExchangeError, +> { + let mut body = String::new(); + let mut saw_content_frame = false; + let mut first_token_ms = None; + + loop { + let receive_timeout = receive_idle_timeout_with_first_token_deadline( + saw_content_frame, + first_token_started_at, + first_token_timeout, + )?; + let message = match next_websocket_message(&mut websocket, receive_timeout).await { + Ok(message) => message, + Err(error) if !saw_content_frame => { + if let Some(timeout) = + elapsed_first_token_timeout(first_token_started_at, first_token_timeout) + { + return Err(CodexWebSocketExchangeError::FirstTokenTimeout { timeout }); + } + if reused_connection { + return Err(reused_connection_died_before_first_frame(&error)); + } + return Err(error); + } + Err(error) => return Err(error), + }; + let Some(message) = message else { + break; + }; + let text = match message { + Message::Text(text) => text, + Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent), + Message::Close(_) if reused_connection && !saw_content_frame => { + return Err( + CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { + message: "websocket closed".to_string(), + }, + ); + } + Message::Close(_) => break, + _ => continue, + }; + let raw = text.to_string(); + if let Some(classified) = classify_websocket_error_frame(&raw) { + let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) + .or_else(|| events::retry_after_seconds_from_body(&raw)); + return Err(CodexWebSocketExchangeError::upstream( + classified.status_code, + retry_after_seconds, + raw, + Vec::new(), + CodexUpstreamDiagnostics::with_status(classified.status_code), + )); + } + if let Some(reason) = websocket_incomplete_response_reason(&raw) { + return Err(CodexWebSocketExchangeError::IncompleteResponse { reason }); + } + if let Some(message) = websocket_response_completed_parse_error(&raw) { + return Err(CodexWebSocketExchangeError::InvalidCompletedResponse { message }); + } + if let Some(headers) = websocket_rate_limit_event_headers(&raw) { + metadata.rate_limit_headers.extend(headers); + continue; + } + if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { + metadata.turn_state = Some(metadata_turn_state); + continue; + } + let event = websocket_event_type(&raw); + let forwarded = if let Some(frame) = websocket_event_to_sse_frame(&raw) { + body.push_str(&frame); + if event_is_first_content(event.as_deref()) { + saw_content_frame = true; + response_meta::update_first_token_ms( + started_at, + body.as_bytes(), + &mut first_token_ms, + ); + } + true + } else { + false + }; + if forwarded && event.as_deref().is_some_and(is_terminal_websocket_event) { + let usage = events::extract_sse_usage(&body)?; + let exchange = CodexWebSocketExchange { + body, + usage, + turn_state: metadata.turn_state.clone(), + set_cookie_headers: metadata.set_cookie_headers.clone(), + rate_limit_headers: metadata.rate_limit_headers.clone(), + first_token_ms, + pool_decision: None, + diagnostics: metadata.diagnostics.clone(), + }; + return Ok((exchange, websocket, metadata)); + } + } + + if reused_connection && !saw_content_frame { + return Err( + CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { + message: "websocket closed before terminal event".to_string(), + }, + ); + } + + Err(CodexWebSocketExchangeError::ClosedBeforeTerminal) +} + +pub(super) async fn prefetch_first_content_stream_frames( + websocket: &mut PumpedWebSocket, + metadata: &mut CodexWebSocketConnectionMetadata, + first_token_started_at: Instant, + first_token_timeout: Option, +) -> Result, CodexWebSocketExchangeError> { + let mut prefetched_frames = Vec::new(); + loop { + let receive_timeout = + receive_timeout_before_first_content(first_token_started_at, first_token_timeout)?; + let message = match next_websocket_message(websocket, receive_timeout).await { + Ok(Some(message)) => message, + Ok(None) => return Err(CodexWebSocketExchangeError::ClosedBeforeTerminal), + Err(error) => { + if let Some(timeout) = + elapsed_first_token_timeout(first_token_started_at, first_token_timeout) + { + return Err(CodexWebSocketExchangeError::FirstTokenTimeout { timeout }); + } + return Err(error); + } + }; + + let raw = match message { + Message::Text(text) => text.to_string(), + Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent), + Message::Close(_) => return Err(CodexWebSocketExchangeError::ClosedBeforeTerminal), + _ => continue, + }; + + if let Some(classified) = classify_websocket_error_frame(&raw) { + let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) + .or_else(|| events::retry_after_seconds_from_body(&raw)); + return Err(CodexWebSocketExchangeError::upstream( + classified.status_code, + retry_after_seconds, + raw, + Vec::new(), + CodexUpstreamDiagnostics::with_status(classified.status_code), + )); + } + if let Some(reason) = websocket_incomplete_response_reason(&raw) { + return Err(CodexWebSocketExchangeError::IncompleteResponse { reason }); + } + if let Some(message) = websocket_response_completed_parse_error(&raw) { + return Err(CodexWebSocketExchangeError::InvalidCompletedResponse { message }); + } + if let Some(headers) = websocket_rate_limit_event_headers(&raw) { + metadata.rate_limit_headers.extend(headers); + continue; + } + if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { + metadata.turn_state = Some(metadata_turn_state); + continue; + } + + let event = websocket_event_type(&raw); + if websocket_event_to_sse_frame(&raw).is_none() { + continue; + } + let is_first_content = event_is_first_content(event.as_deref()); + prefetched_frames.push(raw); + if is_first_content { + return Ok(prefetched_frames); + } + } +} + +pub(super) fn reused_stream_prefetch_error( + error: CodexWebSocketExchangeError, +) -> CodexWebSocketExchangeError { + match error { + CodexWebSocketExchangeError::ClosedBeforeTerminal + | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } + | CodexWebSocketExchangeError::Transport(_) => { + reused_connection_died_before_first_frame(&error) + } + error => error, + } +} + +fn reused_connection_died_before_first_frame( + error: &CodexWebSocketExchangeError, +) -> CodexWebSocketExchangeError { + CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { + message: error.to_string(), + } +} + +pub(super) fn websocket_connection_metadata( + response: &WsResponse>>, +) -> CodexWebSocketConnectionMetadata { + CodexWebSocketConnectionMetadata { + turn_state: response_meta::turn_state(response.headers()), + set_cookie_headers: response_meta::set_cookie_headers(response.headers()), + rate_limit_headers: response_meta::rate_limit_headers(response.headers()), + diagnostics: response_meta::diagnostics( + Some(response.status().as_u16()), + response.headers(), + ), + } +} + +pub(super) fn reusable_websocket_metadata( + mut metadata: CodexWebSocketConnectionMetadata, +) -> CodexWebSocketConnectionMetadata { + metadata.rate_limit_headers.clear(); + metadata +} + +// --------------------------------------------------------------------------- +// Streaming +// --------------------------------------------------------------------------- + +pub(super) struct WebSocketStreamPoolReturn { + pub(super) pool: CodexWebSocketPool, + pub(super) key: CodexWebSocketPoolKey, + pub(super) created_at: Instant, +} + +pub(super) fn stream_websocket_response( + websocket: PumpedWebSocket, + metadata: CodexWebSocketConnectionMetadata, + pool_return: Option, + prefetched_frames: Vec, +) -> CodexWebSocketStreamingExchange { + let response_metadata = metadata.clone(); + let rate_limit_header_updates = Arc::new(Mutex::new(Vec::new())); + let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates); + let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone())); + let turn_state_update_for_task = Arc::clone(&turn_state_update); + let (tx, rx) = mpsc::channel(WEBSOCKET_STREAM_BUFFER); + tokio::spawn(async move { + forward_websocket_response_stream( + websocket, + metadata, + pool_return, + prefetched_frames, + rate_limit_header_updates_for_task, + turn_state_update_for_task, + tx, + ) + .await; + }); + + let body = futures::stream::unfold(rx, |mut rx| async move { + rx.recv().await.map(|item| (item, rx)) + }); + + CodexWebSocketStreamingExchange { + body: Box::pin(body), + turn_state: response_metadata.turn_state, + set_cookie_headers: response_metadata.set_cookie_headers, + rate_limit_headers: response_metadata.rate_limit_headers, + rate_limit_header_updates, + turn_state_update, + pool_decision: None, + diagnostics: response_metadata.diagnostics, + } +} + +async fn forward_websocket_response_stream( + mut websocket: PumpedWebSocket, + mut metadata: CodexWebSocketConnectionMetadata, + pool_return: Option, + prefetched_frames: Vec, + rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, + turn_state_update: CodexWebSocketTurnStateUpdate, + tx: mpsc::Sender>, +) { + let mut pool_return = pool_return; + let mut prefetched_frames: VecDeque = prefetched_frames.into(); + let mut saw_content_frame = false; + loop { + let raw = if let Some(text) = prefetched_frames.pop_front() { + text + } else { + let message = tokio::select! { + biased; + // 客户端断开下游 SSE 流：立即丢弃连接并释放池 slot， + // 不再傻等上游 idle 超时（否则同会话后续请求会一直 bypass/busy）。 + () = tx.closed() => { + discard_stream_websocket(websocket, pool_return).await; + return; + } + message = next_websocket_message( + &mut websocket, + receive_idle_timeout(saw_content_frame), + ) => message, + }; + let message = match message { + Ok(message) => message, + Err(error) => { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx.send(Err(error)).await; + return; + } + }; + let Some(message) = message else { + break; + }; + match message { + Message::Text(text) => text.to_string(), + Message::Binary(_) => { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx + .send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent)) + .await; + return; + } + Message::Close(_) => { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx + .send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)) + .await; + return; + } + _ => continue, + } + }; + if let Some(classified) = classify_websocket_error_frame(&raw) { + discard_stream_websocket(websocket, pool_return).await; + let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) + .or_else(|| events::retry_after_seconds_from_body(&raw)); + let _ = tx + .send(Err(CodexWebSocketExchangeError::upstream( + classified.status_code, + retry_after_seconds, + raw, + Vec::new(), + CodexUpstreamDiagnostics::with_status(classified.status_code), + ))) + .await; + return; + } + if let Some(reason) = websocket_incomplete_response_reason(&raw) { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx + .send(Err(CodexWebSocketExchangeError::IncompleteResponse { + reason, + })) + .await; + return; + } + if let Some(message) = websocket_response_completed_parse_error(&raw) { + discard_stream_websocket(websocket, pool_return).await; + let _ = tx + .send(Err(CodexWebSocketExchangeError::InvalidCompletedResponse { + message, + })) + .await; + return; + } + if let Some(headers) = websocket_rate_limit_event_headers(&raw) { + metadata.rate_limit_headers.extend(headers.iter().cloned()); + rate_limit_header_updates.lock().await.extend(headers); + continue; + } + if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { + metadata.turn_state = Some(metadata_turn_state); + *turn_state_update.lock().await = metadata.turn_state.clone(); + continue; + } + let event = websocket_event_type(&raw); + let Some(frame) = websocket_event_to_sse_frame(&raw) else { + continue; + }; + let terminal = event.as_deref().is_some_and(is_terminal_websocket_event); + if tx.send(Ok(Bytes::from(frame))).await.is_err() { + discard_stream_websocket(websocket, pool_return).await; + return; + } + if event_is_first_content(event.as_deref()) { + saw_content_frame = true; + } + if terminal { + finish_stream_websocket(websocket, metadata, pool_return.take()).await; + return; + } + } + + discard_stream_websocket(websocket, pool_return).await; + let _ = tx + .send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)) + .await; +} + +fn receive_idle_timeout(saw_response_frame: bool) -> Duration { + if saw_response_frame { + WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT + } else { + WEBSOCKET_RECEIVE_IDLE_TIMEOUT + } +} + +fn receive_idle_timeout_with_first_token_deadline( + saw_content_frame: bool, + first_token_started_at: Instant, + first_token_timeout: Option, +) -> Result { + if saw_content_frame { + return Ok(WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT); + } + receive_timeout_before_first_content(first_token_started_at, first_token_timeout) +} + +fn receive_timeout_before_first_content( + first_token_started_at: Instant, + first_token_timeout: Option, +) -> Result { + let Some(first_token_timeout) = first_token_timeout.filter(|timeout| !timeout.is_zero()) else { + return Ok(WEBSOCKET_RECEIVE_IDLE_TIMEOUT); + }; + let Some(remaining) = first_token_timeout.checked_sub(first_token_started_at.elapsed()) else { + return Err(CodexWebSocketExchangeError::FirstTokenTimeout { + timeout: first_token_timeout, + }); + }; + if remaining.is_zero() { + return Err(CodexWebSocketExchangeError::FirstTokenTimeout { + timeout: first_token_timeout, + }); + } + Ok(WEBSOCKET_RECEIVE_IDLE_TIMEOUT.min(remaining)) +} + +fn elapsed_first_token_timeout( + first_token_started_at: Instant, + first_token_timeout: Option, +) -> Option { + let first_token_timeout = first_token_timeout.filter(|timeout| !timeout.is_zero())?; + (first_token_started_at.elapsed() >= first_token_timeout).then_some(first_token_timeout) +} + +fn event_is_first_content(event: Option<&str>) -> bool { + match event { + Some(event) => !is_pre_content_lifecycle_event(event), + None => true, + } +} + +pub(super) fn is_first_token_timeout(error: &CodexWebSocketExchangeError) -> bool { + matches!(error, CodexWebSocketExchangeError::FirstTokenTimeout { .. }) +} + +async fn next_websocket_message( + websocket: &mut PumpedWebSocket, + receive_timeout: Duration, +) -> Result, CodexWebSocketExchangeError> { + match timeout(receive_timeout, websocket.next()).await { + Ok(message) => message.transpose().map_err(Into::into), + Err(_) => Err(CodexWebSocketExchangeError::ReceiveIdleTimeout { + timeout: receive_timeout, + }), + } +} + +async fn finish_stream_websocket( + mut websocket: PumpedWebSocket, + metadata: CodexWebSocketConnectionMetadata, + pool_return: Option, +) { + let Some(pool_return) = pool_return else { + websocket.close().await; + return; + }; + pool_return + .pool + .put( + pool_return.key, + PooledWebSocketConnection { + websocket, + metadata: reusable_websocket_metadata(metadata), + created_at: pool_return.created_at, + }, + ) + .await; +} + +async fn discard_stream_websocket( + mut websocket: PumpedWebSocket, + pool_return: Option, +) { + if let Some(pool_return) = pool_return { + pool_return.pool.discard(&pool_return.key).await; + } + websocket.close().await; +} + +// --------------------------------------------------------------------------- +// Header / metadata helpers +// --------------------------------------------------------------------------- + +fn websocket_rate_limit_event_headers(raw: &str) -> Option> { + events::parse_rate_limits_event_raw(raw) + .map(|parsed| events::rate_limits_to_header_pairs(&parsed)) +} + +pub(super) fn audit_header_value(name: &str, value: &str) -> String { + if is_sensitive_opening_header(name) { + REDACTED_HEADER_VALUE.to_string() + } else { + value.to_string() + } +} + +fn is_sensitive_opening_header(name: &str) -> bool { + matches!( + name.to_ascii_lowercase().as_str(), + "authorization" + | "chatgpt-account-id" + | "cookie" + | "session_id" + | "session-id" + | "thread-id" + | "x-client-request-id" + | "x-codex-window-id" + | "x-codex-turn-metadata" + | "x-codex-turn-state" + | "x-codex-parent-thread-id" + ) +} + +pub(super) fn websocket_audit_file_name() -> String { + let timestamp = china_filename_timestamp_millis(&Utc::now()); + format!("codex-ws-audit-{timestamp}-{}.json", Uuid::new_v4()) +} diff --git a/backend/src/upstream/openai/transport/websocket/types.rs b/backend/src/upstream/openai/transport/websocket/types.rs new file mode 100644 index 000000000..9b1e0e07a --- /dev/null +++ b/backend/src/upstream/openai/transport/websocket/types.rs @@ -0,0 +1,189 @@ +use super::*; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketConnection { + pub(super) endpoint: String, + pub(super) headers: Vec<(String, String)>, +} + +/// 显式写入 WebSocket audit artifact。 +pub async fn write_websocket_audit_artifact_for_dir( + dir: Option<&Path>, + artifact: &WebSocketAuditArtifact, +) -> io::Result> { + let Some(dir) = dir.filter(|dir| !dir.as_os_str().is_empty()) else { + return Ok(None); + }; + + tokio::fs::create_dir_all(dir).await?; + let path = dir.join(websocket_audit_file_name()); + let body = serde_json::to_vec_pretty(artifact).map_err(io::Error::other)?; + tokio::fs::write(&path, body).await?; + Ok(Some(path)) +} + +/// 按环境变量配置写入 WebSocket audit artifact。 +pub async fn write_websocket_audit_artifact_from_env( + artifact: &WebSocketAuditArtifact, +) -> io::Result> { + let Some(dir) = std::env::var_os(WS_AUDIT_DIR_ENV) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) + else { + return Ok(None); + }; + + write_websocket_audit_artifact_for_dir(Some(&dir), artifact).await +} + +/// Prepared Responses WebSocket request descriptor. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketRequest { + pub(super) connection: CodexWebSocketConnection, + pub(super) payload_text: String, +} + +/// Responses WebSocket exchange result. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketExchange { + /// 完整 SSE 文本。 + pub body: String, + /// 从 SSE 中提取的 usage。 + pub usage: Option, + /// 上游 metadata 帧中的最新 turn state。 + pub turn_state: Option, + /// 上游握手响应里的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游握手响应里的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// 首个有效上游 WebSocket 事件到达代理的耗时。 + pub first_token_ms: Option, + /// WebSocket 连接池决策。 + pub pool_decision: Option, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +/// Responses WebSocket live SSE exchange result. +pub struct CodexWebSocketStreamingExchange { + /// Live SSE bytes converted from WebSocket events. + pub body: CodexWebSocketSseStream, + /// 上游 metadata 帧中的最新 turn state。 + pub turn_state: Option, + /// 上游握手响应里的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游握手响应里的限流头。 + pub rate_limit_headers: Vec<(String, String)>, + /// 上游内部 rate-limit 事件里的动态更新。 + pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, + /// 上游内部 metadata 事件里的动态 turn state。 + pub turn_state_update: CodexWebSocketTurnStateUpdate, + /// WebSocket 连接池决策。 + pub pool_decision: Option, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +/// Responses WebSocket live SSE byte stream. +pub type CodexWebSocketSseStream = + Pin> + Send + 'static>>; +/// WebSocket live stream rate-limit header updates. +pub type CodexWebSocketRateLimitHeaderUpdates = Arc>>; +/// WebSocket live stream turn-state update. +pub type CodexWebSocketTurnStateUpdate = Arc>>; + +/// Responses WebSocket exchange error. +#[derive(Debug, Error)] +pub enum CodexWebSocketExchangeError { + /// opening request 无法构造。 + #[error("invalid websocket request: {0}")] + InvalidRequest(#[from] tungstenite::http::Error), + /// WebSocket 传输失败。 + #[error("websocket transport error: {0}")] + Transport(#[from] tungstenite::Error), + /// SSE 聚合结果无法解析。 + #[error("invalid websocket SSE response: {0}")] + InvalidSse(#[from] SseError), + /// 上游 WebSocket 错误帧。 + #[error("{0}")] + Upstream(Box), + /// 上游返回 `response.incomplete`。 + #[error("Incomplete response returned, reason: {reason}")] + IncompleteResponse { + /// incomplete_details.reason。 + reason: String, + }, + /// 上游返回无法按官方形状解析的 `response.completed`。 + #[error("{message}")] + InvalidCompletedResponse { + /// 解析失败说明。 + message: String, + }, + /// 上游在 terminal 事件前关闭。 + #[error("websocket closed before terminal event")] + ClosedBeforeTerminal, + /// 上游在指定时间内没有发送任何事件。 + #[error("websocket receive idle timeout after {timeout:?}")] + ReceiveIdleTimeout { + /// 超时时长。 + timeout: Duration, + }, + /// 上游返回非文本事件帧。 + #[error("unexpected binary websocket event")] + UnexpectedBinaryEvent, + /// 复用的池连接在首个响应帧前失效。 + #[error("reused websocket connection died before first response frame: {message}")] + ReusedConnectionDiedBeforeFirstFrame { + /// 底层失效原因。 + message: String, + }, + /// 首个内容帧在配置的绝对超时内未到达（连接落到病态上游后端）。 + #[error("websocket first content frame not received within {timeout:?}")] + FirstTokenTimeout { + /// 首 token 超时时长。 + timeout: Duration, + }, +} + +/// WebSocket 上游错误帧载荷。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CodexWebSocketUpstreamError { + /// HTTP-style upstream status code. + pub status_code: u16, + /// 推导出的重试秒数。 + pub retry_after_seconds: Option, + /// 原始错误帧。 + pub body: String, + /// 上游透传的 `set-cookie` 列表。 + pub set_cookie_headers: Vec, + /// 上游诊断元数据。 + pub diagnostics: CodexUpstreamDiagnostics, +} + +impl fmt::Display for CodexWebSocketUpstreamError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "websocket upstream error {}: {}", + self.status_code, self.body + ) + } +} + +impl CodexWebSocketExchangeError { + pub(super) fn upstream( + status_code: u16, + retry_after_seconds: Option, + body: String, + set_cookie_headers: Vec, + diagnostics: CodexUpstreamDiagnostics, + ) -> Self { + Self::Upstream(Box::new(CodexWebSocketUpstreamError { + status_code, + retry_after_seconds, + body, + set_cookie_headers, + diagnostics, + })) + } +} diff --git a/backend/src/upstream/transport/websocket_pool.rs b/backend/src/upstream/openai/transport/websocket_pool.rs similarity index 100% rename from backend/src/upstream/transport/websocket_pool.rs rename to backend/src/upstream/openai/transport/websocket_pool.rs diff --git a/backend/src/upstream/transport/websocket_pump.rs b/backend/src/upstream/openai/transport/websocket_pump.rs similarity index 100% rename from backend/src/upstream/transport/websocket_pump.rs rename to backend/src/upstream/openai/transport/websocket_pump.rs diff --git a/backend/src/upstream/transport/websocket.rs b/backend/src/upstream/transport/websocket.rs deleted file mode 100644 index f5e2e528e..000000000 --- a/backend/src/upstream/transport/websocket.rs +++ /dev/null @@ -1,1470 +0,0 @@ -//! Codex WebSocket 连接建立（关键函数）。 - -use std::{ - collections::VecDeque, - fmt, io, - path::{Path, PathBuf}, - pin::Pin, - sync::Arc, - time::{Duration, Instant}, -}; - -use bytes::Bytes; -use chrono::Utc; -use futures::Stream; -use thiserror::Error; -use tokio::sync::mpsc; -use tokio::{sync::Mutex, time::timeout}; -use tokio_tungstenite::{connect_async_tls_with_config, Connector}; -use tungstenite::{ - self, - extensions::{compression::deflate::DeflateConfig, ExtensionsConfig}, - handshake::client::Request as WsRequest, - http::Response as WsResponse, - protocol::WebSocketConfig, - Message, -}; - -use crate::infra::time::china_filename_timestamp_millis; -use crate::upstream::protocol::events::{self, TokenUsage}; -use crate::upstream::protocol::responses::CodexResponsesRequest; -use crate::upstream::protocol::sse::SseError; -use crate::upstream::protocol::websocket::{ - classify_websocket_error_frame, is_pre_content_lifecycle_event, is_terminal_websocket_event, - retry_after_seconds_from_wrapped_error_headers, websocket_event_to_sse_frame, - websocket_event_type, websocket_incomplete_response_reason, websocket_metadata_turn_state, - websocket_response_completed_parse_error, websocket_response_create_payload_text, - OpeningAuditHeader, OpeningAuditSnapshot, WebSocketAuditArtifact, -}; - -use super::websocket_pool::{ - CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, - PooledWebSocketConnection, WebSocketPoolAcquire, WebSocketPoolDecision, -}; -use super::websocket_pump::{PumpKeepalive, PumpedWebSocket, RawWsStream}; -use super::{diagnostics::CodexUpstreamDiagnostics, response_meta}; -use uuid::Uuid; - -const REDACTED_HEADER_VALUE: &str = ""; -const CODEX_RESPONSES_PATH: &str = "/codex/responses"; -const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits"; -const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20); -const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60); -const WEBSOCKET_STREAM_BUFFER: usize = 16; -const WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS: usize = 2; -/// WebSocket audit artifact 输出目录环境变量。 -pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR"; - -// --------------------------------------------------------------------------- -// Types -// --------------------------------------------------------------------------- - -/// WebSocket 连接适配器。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketConnection { - endpoint: String, - headers: Vec<(String, String)>, -} - -/// 显式写入 WebSocket audit artifact。 -pub async fn write_websocket_audit_artifact_for_dir( - dir: Option<&Path>, - artifact: &WebSocketAuditArtifact, -) -> io::Result> { - let Some(dir) = dir.filter(|dir| !dir.as_os_str().is_empty()) else { - return Ok(None); - }; - - tokio::fs::create_dir_all(dir).await?; - let path = dir.join(websocket_audit_file_name()); - let body = serde_json::to_vec_pretty(artifact).map_err(io::Error::other)?; - tokio::fs::write(&path, body).await?; - Ok(Some(path)) -} - -/// 按环境变量配置写入 WebSocket audit artifact。 -pub async fn write_websocket_audit_artifact_from_env( - artifact: &WebSocketAuditArtifact, -) -> io::Result> { - let Some(dir) = std::env::var_os(WS_AUDIT_DIR_ENV) - .filter(|value| !value.is_empty()) - .map(PathBuf::from) - else { - return Ok(None); - }; - - write_websocket_audit_artifact_for_dir(Some(&dir), artifact).await -} - -/// Prepared Responses WebSocket request descriptor. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketRequest { - connection: CodexWebSocketConnection, - payload_text: String, -} - -/// Responses WebSocket exchange result. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketExchange { - /// 完整 SSE 文本。 - pub body: String, - /// 从 SSE 中提取的 usage。 - pub usage: Option, - /// 上游 metadata 帧中的最新 turn state。 - pub turn_state: Option, - /// 上游握手响应里的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游握手响应里的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// 首个有效上游 WebSocket 事件到达代理的耗时。 - pub first_token_ms: Option, - /// WebSocket 连接池决策。 - pub pool_decision: Option, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -/// Responses WebSocket live SSE exchange result. -pub struct CodexWebSocketStreamingExchange { - /// Live SSE bytes converted from WebSocket events. - pub body: CodexWebSocketSseStream, - /// 上游 metadata 帧中的最新 turn state。 - pub turn_state: Option, - /// 上游握手响应里的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游握手响应里的限流头。 - pub rate_limit_headers: Vec<(String, String)>, - /// 上游内部 rate-limit 事件里的动态更新。 - pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, - /// 上游内部 metadata 事件里的动态 turn state。 - pub turn_state_update: CodexWebSocketTurnStateUpdate, - /// WebSocket 连接池决策。 - pub pool_decision: Option, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -/// Responses WebSocket live SSE byte stream. -pub type CodexWebSocketSseStream = - Pin> + Send + 'static>>; -/// WebSocket live stream rate-limit header updates. -pub type CodexWebSocketRateLimitHeaderUpdates = Arc>>; -/// WebSocket live stream turn-state update. -pub type CodexWebSocketTurnStateUpdate = Arc>>; - -/// Responses WebSocket exchange error. -#[derive(Debug, Error)] -pub enum CodexWebSocketExchangeError { - /// opening request 无法构造。 - #[error("invalid websocket request: {0}")] - InvalidRequest(#[from] tungstenite::http::Error), - /// WebSocket 传输失败。 - #[error("websocket transport error: {0}")] - Transport(#[from] tungstenite::Error), - /// SSE 聚合结果无法解析。 - #[error("invalid websocket SSE response: {0}")] - InvalidSse(#[from] SseError), - /// 上游 WebSocket 错误帧。 - #[error("{0}")] - Upstream(Box), - /// 上游返回 `response.incomplete`。 - #[error("Incomplete response returned, reason: {reason}")] - IncompleteResponse { - /// incomplete_details.reason。 - reason: String, - }, - /// 上游返回无法按官方形状解析的 `response.completed`。 - #[error("{message}")] - InvalidCompletedResponse { - /// 解析失败说明。 - message: String, - }, - /// 上游在 terminal 事件前关闭。 - #[error("websocket closed before terminal event")] - ClosedBeforeTerminal, - /// 上游在指定时间内没有发送任何事件。 - #[error("websocket receive idle timeout after {timeout:?}")] - ReceiveIdleTimeout { - /// 超时时长。 - timeout: Duration, - }, - /// 上游返回非文本事件帧。 - #[error("unexpected binary websocket event")] - UnexpectedBinaryEvent, - /// 复用的池连接在首个响应帧前失效。 - #[error("reused websocket connection died before first response frame: {message}")] - ReusedConnectionDiedBeforeFirstFrame { - /// 底层失效原因。 - message: String, - }, - /// 首个内容帧在配置的绝对超时内未到达（连接落到病态上游后端）。 - #[error("websocket first content frame not received within {timeout:?}")] - FirstTokenTimeout { - /// 首 token 超时时长。 - timeout: Duration, - }, -} - -/// WebSocket 上游错误帧载荷。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CodexWebSocketUpstreamError { - /// HTTP-style upstream status code. - pub status_code: u16, - /// 推导出的重试秒数。 - pub retry_after_seconds: Option, - /// 原始错误帧。 - pub body: String, - /// 上游透传的 `set-cookie` 列表。 - pub set_cookie_headers: Vec, - /// 上游诊断元数据。 - pub diagnostics: CodexUpstreamDiagnostics, -} - -impl fmt::Display for CodexWebSocketUpstreamError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!( - f, - "websocket upstream error {}: {}", - self.status_code, self.body - ) - } -} - -impl CodexWebSocketExchangeError { - fn upstream( - status_code: u16, - retry_after_seconds: Option, - body: String, - set_cookie_headers: Vec, - diagnostics: CodexUpstreamDiagnostics, - ) -> Self { - Self::Upstream(Box::new(CodexWebSocketUpstreamError { - status_code, - retry_after_seconds, - body, - set_cookie_headers, - diagnostics, - })) - } -} - -impl CodexWebSocketRequest { - /// 返回连接描述。 - pub fn connection(&self) -> &CodexWebSocketConnection { - &self.connection - } - - /// 返回将要发送的首个文本帧。 - pub fn payload_text(&self) -> &str { - &self.payload_text - } -} - -impl CodexWebSocketConnection { - /// 构造待打开的 WebSocket 连接描述。 - pub fn new(endpoint: impl Into, headers: Vec<(String, String)>) -> Self { - Self { - endpoint: endpoint.into(), - headers, - } - } - - /// 构造 Responses WebSocket 连接描述。 - pub fn responses( - base_url: &str, - websocket_key: &str, - business_headers: Vec<(String, String)>, - ) -> Self { - let endpoint = responses_websocket_endpoint(base_url); - let mut headers = Vec::new(); - if let Some(host) = websocket_host_header(&endpoint) { - headers.push(("Host".to_string(), host)); - } - headers.extend([ - ("Connection".to_string(), "Upgrade".to_string()), - ("Upgrade".to_string(), "websocket".to_string()), - ("Sec-WebSocket-Version".to_string(), "13".to_string()), - ("Sec-WebSocket-Key".to_string(), websocket_key.to_string()), - ]); - headers.extend(business_headers); - headers.push(( - "sec-websocket-extensions".to_string(), - WEBSOCKET_EXTENSIONS.to_string(), - )); - Self { endpoint, headers } - } - - /// 构造 Responses WebSocket opening 与首个 `response.create` 文本帧。 - pub fn responses_create_request( - base_url: &str, - websocket_key: &str, - business_headers: Vec<(String, String)>, - request: &CodexResponsesRequest, - ) -> Result { - Ok(CodexWebSocketRequest { - connection: Self::responses(base_url, websocket_key, business_headers), - payload_text: websocket_response_create_payload_text(request)?, - }) - } - - /// 返回 WebSocket endpoint。 - pub fn endpoint(&self) -> &str { - &self.endpoint - } - - /// 返回按发送顺序保存的请求头。 - pub fn headers(&self) -> &[(String, String)] { - &self.headers - } - - /// 生成打开握手审计快照。 - pub fn opening_audit_snapshot(&self) -> OpeningAuditSnapshot { - OpeningAuditSnapshot { - request_line: request_line_for_endpoint(&self.endpoint), - header_order: self.headers.iter().map(|(name, _)| name.clone()).collect(), - headers: self - .headers - .iter() - .map(|(name, value)| OpeningAuditHeader { - name: name.clone(), - value: audit_header_value(name, value), - }) - .collect(), - } - } -} - -// --------------------------------------------------------------------------- -// WebSocket endpoint helpers -// --------------------------------------------------------------------------- - -/// 将 Codex backend base URL 转换为 Responses WebSocket endpoint。 -pub fn responses_websocket_endpoint(base_url: &str) -> String { - let endpoint = format!("{}{}", base_url.trim_end_matches('/'), CODEX_RESPONSES_PATH); - if let Some(rest) = endpoint.strip_prefix("https://") { - format!("wss://{rest}") - } else if let Some(rest) = endpoint.strip_prefix("http://") { - format!("ws://{rest}") - } else { - endpoint - } -} - -fn websocket_host_header(endpoint: &str) -> Option { - let url = reqwest::Url::parse(endpoint).ok()?; - let host = url.host_str()?; - Some(match url.port() { - Some(port) => format!("{host}:{port}"), - None => host.to_string(), - }) -} - -// --------------------------------------------------------------------------- -// Connection -// --------------------------------------------------------------------------- - -async fn connect_websocket( - connection: &CodexWebSocketConnection, -) -> Result<(RawWsStream, WsResponse>>), CodexWebSocketExchangeError> { - let request = websocket_handshake_request(connection)?; - let connector = super::tls::maybe_build_rustls_client_config_with_custom_ca() - .map_err(|error| tungstenite::Error::Io(std::io::Error::other(error)))? - .map(Connector::Rustls); - match connect_async_tls_with_config(request, Some(websocket_config()), false, connector).await { - Ok((websocket, response)) => Ok((websocket, response)), - Err(tungstenite::Error::Http(response)) => Err(websocket_opening_error(response.as_ref())), - Err(error) => Err(error.into()), - } -} - -/// 建立连接并交给后台 pump 托管；`keepalive` 决定该连接是否做主动保活。 -/// -/// 非池化（即用即弃）连接用 [`PumpKeepalive::disabled`]；池化连接用连接池派生的保活策略， -/// 让空闲期在后台完成 ping/pong 与失活检测，从而复用前可零成本判活。 -async fn connect_pumped_websocket( - connection: &CodexWebSocketConnection, - keepalive: PumpKeepalive, -) -> Result<(PumpedWebSocket, WsResponse>>), CodexWebSocketExchangeError> { - let (raw, response) = connect_websocket(connection).await?; - Ok((PumpedWebSocket::new(raw, keepalive), response)) -} - -fn websocket_handshake_request( - connection: &CodexWebSocketConnection, -) -> Result { - let mut builder = WsRequest::builder() - .method("GET") - .uri(connection.endpoint()); - for (name, value) in connection.headers() { - if name.eq_ignore_ascii_case("sec-websocket-extensions") { - continue; - } - builder = builder.header(name.as_str(), value.as_str()); - } - builder.body(()) -} - -fn websocket_config() -> WebSocketConfig { - let mut extensions = ExtensionsConfig::default(); - extensions.permessage_deflate = Some(DeflateConfig::default()); - - let mut config = WebSocketConfig::default(); - config.extensions = extensions; - config -} - -fn websocket_opening_error(response: &WsResponse>>) -> CodexWebSocketExchangeError { - let status_code = response.status().as_u16(); - let body = response - .body() - .as_ref() - .map(|body| String::from_utf8_lossy(body).into_owned()) - .unwrap_or_default(); - let retry_after_seconds = response - .headers() - .get("retry-after") - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.parse::().ok()) - .filter(|seconds| *seconds > 0) - .or_else(|| events::retry_after_seconds_from_body(&body)); - CodexWebSocketExchangeError::upstream( - status_code, - retry_after_seconds, - body, - response_meta::set_cookie_headers(response.headers()), - response_meta::diagnostics(Some(status_code), response.headers()), - ) -} - -fn request_line_for_endpoint(endpoint: &str) -> String { - let path = reqwest::Url::parse(endpoint) - .ok() - .map(|url| { - let mut path = url.path().to_string(); - if let Some(query) = url.query() { - path.push('?'); - path.push_str(query); - } - path - }) - .filter(|path| !path.is_empty()) - .unwrap_or_else(|| endpoint.to_string()); - format!("GET {path} HTTP/1.1") -} - -// --------------------------------------------------------------------------- -// Fresh request execution -// --------------------------------------------------------------------------- - -async fn execute_fresh_response_create_request( - request: &CodexWebSocketRequest, - started_at: Instant, - first_token_timeout: Option, -) -> Result { - let (websocket, response) = - connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?; - let first_token_started_at = Instant::now(); - websocket - .send(Message::Text(request.payload_text().to_string().into())) - .await?; - - let metadata = websocket_connection_metadata(&response); - let (exchange, _websocket, _metadata) = collect_websocket_response( - websocket, - metadata, - false, - started_at, - first_token_started_at, - first_token_timeout, - ) - .await?; - Ok(exchange) -} - -async fn execute_fresh_response_create_request_stream( - request: &CodexWebSocketRequest, - first_token_timeout: Option, -) -> Result { - let (mut websocket, response) = - connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?; - let first_token_started_at = Instant::now(); - websocket - .send(Message::Text(request.payload_text().to_string().into())) - .await?; - - let mut metadata = websocket_connection_metadata(&response); - let prefetched_frames = match prefetch_first_content_stream_frames( - &mut websocket, - &mut metadata, - first_token_started_at, - first_token_timeout, - ) - .await - { - Ok(prefetched_frames) => prefetched_frames, - Err(error) => { - websocket.close().await; - return Err(error); - } - }; - Ok(stream_websocket_response( - websocket, - metadata, - None, - prefetched_frames, - )) -} - -async fn execute_fresh_response_create_request_with_retries( - request: &CodexWebSocketRequest, - started_at: Instant, - first_token_timeout: Option, - attempts: usize, -) -> Result { - let attempts = attempts.max(1); - for attempt in 1..=attempts { - match execute_fresh_response_create_request(request, started_at, first_token_timeout).await - { - Ok(exchange) => return Ok(exchange), - Err(error) if attempt < attempts && is_first_token_timeout(&error) => continue, - Err(error) => return Err(error), - } - } - unreachable!("fresh websocket retry loop always returns"); -} - -async fn execute_fresh_response_create_request_stream_with_retries( - request: &CodexWebSocketRequest, - first_token_timeout: Option, - attempts: usize, -) -> Result { - let attempts = attempts.max(1); - for attempt in 1..=attempts { - match execute_fresh_response_create_request_stream(request, first_token_timeout).await { - Ok(exchange) => return Ok(exchange), - Err(error) if attempt < attempts && is_first_token_timeout(&error) => continue, - Err(error) => return Err(error), - } - } - unreachable!("fresh websocket stream retry loop always returns"); -} - -// --------------------------------------------------------------------------- -// Pool-aware execution -// --------------------------------------------------------------------------- - -pub(crate) async fn execute_response_create_request_with_pool( - request: &CodexWebSocketRequest, - pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>, - started_at: Instant, - fallback_first_token_timeout: Option, -) -> Result { - let Some((pool, key)) = pool else { - return execute_fresh_response_create_request_with_retries( - request, - started_at, - fallback_first_token_timeout, - WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, - ) - .await; - }; - - match pool.acquire(&key).await { - WebSocketPoolAcquire::Reused(connection) => { - let result = - execute_pooled_response_create_request(request, pool, key, *connection, started_at) - .await; - match result { - Err(CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { - .. - }) - | Err(CodexWebSocketExchangeError::FirstTokenTimeout { .. }) => { - let mut exchange = execute_fresh_response_create_request_with_retries( - request, - started_at, - pool.first_token_timeout(), - WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, - ) - .await?; - exchange.pool_decision = Some(WebSocketPoolDecision::retry_after_stale_reuse()); - Ok(exchange) - } - Ok(mut exchange) => { - exchange.pool_decision = Some(WebSocketPoolDecision::reuse()); - Ok(exchange) - } - Err(error) => Err(error), - } - } - WebSocketPoolAcquire::FreshReserved => { - let result = execute_fresh_pooled_response_create_request( - request, - pool, - key.clone(), - started_at, - ) - .await; - match result { - Ok(mut exchange) => { - exchange.pool_decision = Some(WebSocketPoolDecision::new()); - Ok(exchange) - } - Err(error) if is_first_token_timeout(&error) => { - pool.discard(&key).await; - let mut exchange = execute_fresh_response_create_request_with_retries( - request, - started_at, - pool.first_token_timeout(), - 1, - ) - .await?; - exchange.pool_decision = Some(WebSocketPoolDecision::new()); - Ok(exchange) - } - Err(error) => { - pool.discard(&key).await; - Err(error) - } - } - } - WebSocketPoolAcquire::Bypass(reason) => { - let mut exchange = execute_fresh_response_create_request_with_retries( - request, - started_at, - pool.first_token_timeout(), - WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, - ) - .await?; - exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason)); - Ok(exchange) - } - } -} - -pub(crate) async fn execute_response_create_request_stream_with_pool( - request: &CodexWebSocketRequest, - pool: Option<(&CodexWebSocketPool, CodexWebSocketPoolKey)>, - fallback_first_token_timeout: Option, -) -> Result { - let Some((pool, key)) = pool else { - return execute_fresh_response_create_request_stream_with_retries( - request, - fallback_first_token_timeout, - WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, - ) - .await; - }; - - match pool.acquire(&key).await { - WebSocketPoolAcquire::Reused(connection) => { - let result = execute_pooled_response_create_request_stream( - request, - pool.clone(), - key, - *connection, - ) - .await; - match result { - Err(CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { - .. - }) - | Err(CodexWebSocketExchangeError::FirstTokenTimeout { .. }) => { - let mut exchange = execute_fresh_response_create_request_stream_with_retries( - request, - pool.first_token_timeout(), - WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, - ) - .await?; - exchange.pool_decision = Some(WebSocketPoolDecision::retry_after_stale_reuse()); - Ok(exchange) - } - Ok(mut exchange) => { - exchange.pool_decision = Some(WebSocketPoolDecision::reuse()); - Ok(exchange) - } - Err(error) => Err(error), - } - } - WebSocketPoolAcquire::FreshReserved => { - let result = execute_fresh_pooled_response_create_request_stream( - request, - pool.clone(), - key.clone(), - ) - .await; - match result { - Ok(mut exchange) => { - exchange.pool_decision = Some(WebSocketPoolDecision::new()); - Ok(exchange) - } - Err(error) if is_first_token_timeout(&error) => { - pool.discard(&key).await; - let mut exchange = execute_fresh_response_create_request_stream_with_retries( - request, - pool.first_token_timeout(), - 1, - ) - .await?; - exchange.pool_decision = Some(WebSocketPoolDecision::new()); - Ok(exchange) - } - Err(error) => { - pool.discard(&key).await; - Err(error) - } - } - } - WebSocketPoolAcquire::Bypass(reason) => { - let mut exchange = execute_fresh_response_create_request_stream_with_retries( - request, - pool.first_token_timeout(), - WEBSOCKET_FIRST_TOKEN_FRESH_RETRY_ATTEMPTS, - ) - .await?; - exchange.pool_decision = Some(WebSocketPoolDecision::bypass(reason)); - Ok(exchange) - } - } -} - -async fn execute_fresh_pooled_response_create_request( - request: &CodexWebSocketRequest, - pool: &CodexWebSocketPool, - key: CodexWebSocketPoolKey, - started_at: Instant, -) -> Result { - let (websocket, response) = - connect_pumped_websocket(request.connection(), pool.keepalive()).await?; - let first_token_started_at = Instant::now(); - websocket - .send(Message::Text(request.payload_text().to_string().into())) - .await?; - - let now = Instant::now(); - let metadata = websocket_connection_metadata(&response); - let (exchange, websocket, metadata) = collect_websocket_response( - websocket, - metadata, - false, - started_at, - first_token_started_at, - pool.first_token_timeout(), - ) - .await?; - pool.put( - key, - PooledWebSocketConnection { - websocket, - metadata: reusable_websocket_metadata(metadata), - created_at: now, - }, - ) - .await; - Ok(exchange) -} - -async fn execute_fresh_pooled_response_create_request_stream( - request: &CodexWebSocketRequest, - pool: CodexWebSocketPool, - key: CodexWebSocketPoolKey, -) -> Result { - let (mut websocket, response) = - connect_pumped_websocket(request.connection(), pool.keepalive()).await?; - let first_token_started_at = Instant::now(); - websocket - .send(Message::Text(request.payload_text().to_string().into())) - .await?; - - let now = Instant::now(); - let mut metadata = websocket_connection_metadata(&response); - let prefetched_frames = match prefetch_first_content_stream_frames( - &mut websocket, - &mut metadata, - first_token_started_at, - pool.first_token_timeout(), - ) - .await - { - Ok(prefetched_frames) => prefetched_frames, - Err(error) => { - websocket.close().await; - return Err(error); - } - }; - Ok(stream_websocket_response( - websocket, - metadata, - Some(WebSocketStreamPoolReturn { - pool, - key, - created_at: now, - }), - prefetched_frames, - )) -} - -async fn execute_pooled_response_create_request( - request: &CodexWebSocketRequest, - pool: &CodexWebSocketPool, - key: CodexWebSocketPoolKey, - connection: PooledWebSocketConnection, - started_at: Instant, -) -> Result { - let first_token_started_at = Instant::now(); - if let Err(error) = connection - .websocket - .send(Message::Text(request.payload_text().to_string().into())) - .await - { - pool.discard(&key).await; - return Err( - CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { - message: error.to_string(), - }, - ); - } - - let created_at = connection.created_at; - match collect_websocket_response( - connection.websocket, - reusable_websocket_metadata(connection.metadata), - true, - started_at, - first_token_started_at, - pool.first_token_timeout(), - ) - .await - { - Ok((exchange, websocket, metadata)) => { - pool.put( - key, - PooledWebSocketConnection { - websocket, - metadata: reusable_websocket_metadata(metadata), - created_at, - }, - ) - .await; - Ok(exchange) - } - Err(error) => { - pool.discard(&key).await; - Err(error) - } - } -} - -async fn execute_pooled_response_create_request_stream( - request: &CodexWebSocketRequest, - pool: CodexWebSocketPool, - key: CodexWebSocketPoolKey, - connection: PooledWebSocketConnection, -) -> Result { - let mut websocket = connection.websocket; - let mut metadata = reusable_websocket_metadata(connection.metadata); - let created_at = connection.created_at; - let first_token_started_at = Instant::now(); - if let Err(error) = websocket - .send(Message::Text(request.payload_text().to_string().into())) - .await - { - pool.discard(&key).await; - return Err( - CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { - message: error.to_string(), - }, - ); - } - - let prefetched_frames = match prefetch_first_content_stream_frames( - &mut websocket, - &mut metadata, - first_token_started_at, - pool.first_token_timeout(), - ) - .await - { - Ok(prefetched_frames) => prefetched_frames, - Err(error) => { - pool.discard(&key).await; - websocket.close().await; - return Err(reused_stream_prefetch_error(error)); - } - }; - - Ok(stream_websocket_response( - websocket, - metadata, - Some(WebSocketStreamPoolReturn { - pool, - key, - created_at, - }), - prefetched_frames, - )) -} - -// --------------------------------------------------------------------------- -// Public entry points -// --------------------------------------------------------------------------- - -/// 执行一次 prepared Responses WebSocket 请求并聚合为 SSE。 -pub async fn execute_response_create_request( - request: &CodexWebSocketRequest, -) -> Result { - execute_response_create_request_with_pool(request, None, Instant::now(), None).await -} - -// --------------------------------------------------------------------------- -// Response collection -// --------------------------------------------------------------------------- - -async fn collect_websocket_response( - mut websocket: PumpedWebSocket, - mut metadata: CodexWebSocketConnectionMetadata, - reused_connection: bool, - started_at: Instant, - first_token_started_at: Instant, - first_token_timeout: Option, -) -> Result< - ( - CodexWebSocketExchange, - PumpedWebSocket, - CodexWebSocketConnectionMetadata, - ), - CodexWebSocketExchangeError, -> { - let mut body = String::new(); - let mut saw_content_frame = false; - let mut first_token_ms = None; - - loop { - let receive_timeout = receive_idle_timeout_with_first_token_deadline( - saw_content_frame, - first_token_started_at, - first_token_timeout, - )?; - let message = match next_websocket_message(&mut websocket, receive_timeout).await { - Ok(message) => message, - Err(error) if !saw_content_frame => { - if let Some(timeout) = - elapsed_first_token_timeout(first_token_started_at, first_token_timeout) - { - return Err(CodexWebSocketExchangeError::FirstTokenTimeout { timeout }); - } - if reused_connection { - return Err(reused_connection_died_before_first_frame(&error)); - } - return Err(error); - } - Err(error) => return Err(error), - }; - let Some(message) = message else { - break; - }; - let text = match message { - Message::Text(text) => text, - Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent), - Message::Close(_) if reused_connection && !saw_content_frame => { - return Err( - CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { - message: "websocket closed".to_string(), - }, - ); - } - Message::Close(_) => break, - _ => continue, - }; - let raw = text.to_string(); - if let Some(classified) = classify_websocket_error_frame(&raw) { - let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) - .or_else(|| events::retry_after_seconds_from_body(&raw)); - return Err(CodexWebSocketExchangeError::upstream( - classified.status_code, - retry_after_seconds, - raw, - Vec::new(), - CodexUpstreamDiagnostics::with_status(classified.status_code), - )); - } - if let Some(reason) = websocket_incomplete_response_reason(&raw) { - return Err(CodexWebSocketExchangeError::IncompleteResponse { reason }); - } - if let Some(message) = websocket_response_completed_parse_error(&raw) { - return Err(CodexWebSocketExchangeError::InvalidCompletedResponse { message }); - } - if let Some(headers) = websocket_rate_limit_event_headers(&raw) { - metadata.rate_limit_headers.extend(headers); - continue; - } - if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { - metadata.turn_state = Some(metadata_turn_state); - continue; - } - let event = websocket_event_type(&raw); - let forwarded = if let Some(frame) = websocket_event_to_sse_frame(&raw) { - body.push_str(&frame); - if event_is_first_content(event.as_deref()) { - saw_content_frame = true; - response_meta::update_first_token_ms( - started_at, - body.as_bytes(), - &mut first_token_ms, - ); - } - true - } else { - false - }; - if forwarded && event.as_deref().is_some_and(is_terminal_websocket_event) { - let usage = events::extract_sse_usage(&body)?; - let exchange = CodexWebSocketExchange { - body, - usage, - turn_state: metadata.turn_state.clone(), - set_cookie_headers: metadata.set_cookie_headers.clone(), - rate_limit_headers: metadata.rate_limit_headers.clone(), - first_token_ms, - pool_decision: None, - diagnostics: metadata.diagnostics.clone(), - }; - return Ok((exchange, websocket, metadata)); - } - } - - if reused_connection && !saw_content_frame { - return Err( - CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { - message: "websocket closed before terminal event".to_string(), - }, - ); - } - - Err(CodexWebSocketExchangeError::ClosedBeforeTerminal) -} - -async fn prefetch_first_content_stream_frames( - websocket: &mut PumpedWebSocket, - metadata: &mut CodexWebSocketConnectionMetadata, - first_token_started_at: Instant, - first_token_timeout: Option, -) -> Result, CodexWebSocketExchangeError> { - let mut prefetched_frames = Vec::new(); - loop { - let receive_timeout = - receive_timeout_before_first_content(first_token_started_at, first_token_timeout)?; - let message = match next_websocket_message(websocket, receive_timeout).await { - Ok(Some(message)) => message, - Ok(None) => return Err(CodexWebSocketExchangeError::ClosedBeforeTerminal), - Err(error) => { - if let Some(timeout) = - elapsed_first_token_timeout(first_token_started_at, first_token_timeout) - { - return Err(CodexWebSocketExchangeError::FirstTokenTimeout { timeout }); - } - return Err(error); - } - }; - - let raw = match message { - Message::Text(text) => text.to_string(), - Message::Binary(_) => return Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent), - Message::Close(_) => return Err(CodexWebSocketExchangeError::ClosedBeforeTerminal), - _ => continue, - }; - - if let Some(classified) = classify_websocket_error_frame(&raw) { - let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) - .or_else(|| events::retry_after_seconds_from_body(&raw)); - return Err(CodexWebSocketExchangeError::upstream( - classified.status_code, - retry_after_seconds, - raw, - Vec::new(), - CodexUpstreamDiagnostics::with_status(classified.status_code), - )); - } - if let Some(reason) = websocket_incomplete_response_reason(&raw) { - return Err(CodexWebSocketExchangeError::IncompleteResponse { reason }); - } - if let Some(message) = websocket_response_completed_parse_error(&raw) { - return Err(CodexWebSocketExchangeError::InvalidCompletedResponse { message }); - } - if let Some(headers) = websocket_rate_limit_event_headers(&raw) { - metadata.rate_limit_headers.extend(headers); - continue; - } - if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { - metadata.turn_state = Some(metadata_turn_state); - continue; - } - - let event = websocket_event_type(&raw); - if websocket_event_to_sse_frame(&raw).is_none() { - continue; - } - let is_first_content = event_is_first_content(event.as_deref()); - prefetched_frames.push(raw); - if is_first_content { - return Ok(prefetched_frames); - } - } -} - -fn reused_stream_prefetch_error(error: CodexWebSocketExchangeError) -> CodexWebSocketExchangeError { - match error { - CodexWebSocketExchangeError::ClosedBeforeTerminal - | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } - | CodexWebSocketExchangeError::Transport(_) => { - reused_connection_died_before_first_frame(&error) - } - error => error, - } -} - -fn reused_connection_died_before_first_frame( - error: &CodexWebSocketExchangeError, -) -> CodexWebSocketExchangeError { - CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { - message: error.to_string(), - } -} - -fn websocket_connection_metadata( - response: &WsResponse>>, -) -> CodexWebSocketConnectionMetadata { - CodexWebSocketConnectionMetadata { - turn_state: response_meta::turn_state(response.headers()), - set_cookie_headers: response_meta::set_cookie_headers(response.headers()), - rate_limit_headers: response_meta::rate_limit_headers(response.headers()), - diagnostics: response_meta::diagnostics( - Some(response.status().as_u16()), - response.headers(), - ), - } -} - -fn reusable_websocket_metadata( - mut metadata: CodexWebSocketConnectionMetadata, -) -> CodexWebSocketConnectionMetadata { - metadata.rate_limit_headers.clear(); - metadata -} - -// --------------------------------------------------------------------------- -// Streaming -// --------------------------------------------------------------------------- - -struct WebSocketStreamPoolReturn { - pool: CodexWebSocketPool, - key: CodexWebSocketPoolKey, - created_at: Instant, -} - -fn stream_websocket_response( - websocket: PumpedWebSocket, - metadata: CodexWebSocketConnectionMetadata, - pool_return: Option, - prefetched_frames: Vec, -) -> CodexWebSocketStreamingExchange { - let response_metadata = metadata.clone(); - let rate_limit_header_updates = Arc::new(Mutex::new(Vec::new())); - let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates); - let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone())); - let turn_state_update_for_task = Arc::clone(&turn_state_update); - let (tx, rx) = mpsc::channel(WEBSOCKET_STREAM_BUFFER); - tokio::spawn(async move { - forward_websocket_response_stream( - websocket, - metadata, - pool_return, - prefetched_frames, - rate_limit_header_updates_for_task, - turn_state_update_for_task, - tx, - ) - .await; - }); - - let body = futures::stream::unfold(rx, |mut rx| async move { - rx.recv().await.map(|item| (item, rx)) - }); - - CodexWebSocketStreamingExchange { - body: Box::pin(body), - turn_state: response_metadata.turn_state, - set_cookie_headers: response_metadata.set_cookie_headers, - rate_limit_headers: response_metadata.rate_limit_headers, - rate_limit_header_updates, - turn_state_update, - pool_decision: None, - diagnostics: response_metadata.diagnostics, - } -} - -async fn forward_websocket_response_stream( - mut websocket: PumpedWebSocket, - mut metadata: CodexWebSocketConnectionMetadata, - pool_return: Option, - prefetched_frames: Vec, - rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, - turn_state_update: CodexWebSocketTurnStateUpdate, - tx: mpsc::Sender>, -) { - let mut pool_return = pool_return; - let mut prefetched_frames: VecDeque = prefetched_frames.into(); - let mut saw_content_frame = false; - loop { - let raw = if let Some(text) = prefetched_frames.pop_front() { - text - } else { - let message = tokio::select! { - biased; - // 客户端断开下游 SSE 流：立即丢弃连接并释放池 slot， - // 不再傻等上游 idle 超时（否则同会话后续请求会一直 bypass/busy）。 - () = tx.closed() => { - discard_stream_websocket(websocket, pool_return).await; - return; - } - message = next_websocket_message( - &mut websocket, - receive_idle_timeout(saw_content_frame), - ) => message, - }; - let message = match message { - Ok(message) => message, - Err(error) => { - discard_stream_websocket(websocket, pool_return).await; - let _ = tx.send(Err(error)).await; - return; - } - }; - let Some(message) = message else { - break; - }; - match message { - Message::Text(text) => text.to_string(), - Message::Binary(_) => { - discard_stream_websocket(websocket, pool_return).await; - let _ = tx - .send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent)) - .await; - return; - } - Message::Close(_) => { - discard_stream_websocket(websocket, pool_return).await; - let _ = tx - .send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)) - .await; - return; - } - _ => continue, - } - }; - if let Some(classified) = classify_websocket_error_frame(&raw) { - discard_stream_websocket(websocket, pool_return).await; - let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) - .or_else(|| events::retry_after_seconds_from_body(&raw)); - let _ = tx - .send(Err(CodexWebSocketExchangeError::upstream( - classified.status_code, - retry_after_seconds, - raw, - Vec::new(), - CodexUpstreamDiagnostics::with_status(classified.status_code), - ))) - .await; - return; - } - if let Some(reason) = websocket_incomplete_response_reason(&raw) { - discard_stream_websocket(websocket, pool_return).await; - let _ = tx - .send(Err(CodexWebSocketExchangeError::IncompleteResponse { - reason, - })) - .await; - return; - } - if let Some(message) = websocket_response_completed_parse_error(&raw) { - discard_stream_websocket(websocket, pool_return).await; - let _ = tx - .send(Err(CodexWebSocketExchangeError::InvalidCompletedResponse { - message, - })) - .await; - return; - } - if let Some(headers) = websocket_rate_limit_event_headers(&raw) { - metadata.rate_limit_headers.extend(headers.iter().cloned()); - rate_limit_header_updates.lock().await.extend(headers); - continue; - } - if let Some(metadata_turn_state) = websocket_metadata_turn_state(&raw) { - metadata.turn_state = Some(metadata_turn_state); - *turn_state_update.lock().await = metadata.turn_state.clone(); - continue; - } - let event = websocket_event_type(&raw); - let Some(frame) = websocket_event_to_sse_frame(&raw) else { - continue; - }; - let terminal = event.as_deref().is_some_and(is_terminal_websocket_event); - if tx.send(Ok(Bytes::from(frame))).await.is_err() { - discard_stream_websocket(websocket, pool_return).await; - return; - } - if event_is_first_content(event.as_deref()) { - saw_content_frame = true; - } - if terminal { - finish_stream_websocket(websocket, metadata, pool_return.take()).await; - return; - } - } - - discard_stream_websocket(websocket, pool_return).await; - let _ = tx - .send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)) - .await; -} - -fn receive_idle_timeout(saw_response_frame: bool) -> Duration { - if saw_response_frame { - WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT - } else { - WEBSOCKET_RECEIVE_IDLE_TIMEOUT - } -} - -fn receive_idle_timeout_with_first_token_deadline( - saw_content_frame: bool, - first_token_started_at: Instant, - first_token_timeout: Option, -) -> Result { - if saw_content_frame { - return Ok(WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT); - } - receive_timeout_before_first_content(first_token_started_at, first_token_timeout) -} - -fn receive_timeout_before_first_content( - first_token_started_at: Instant, - first_token_timeout: Option, -) -> Result { - let Some(first_token_timeout) = first_token_timeout.filter(|timeout| !timeout.is_zero()) else { - return Ok(WEBSOCKET_RECEIVE_IDLE_TIMEOUT); - }; - let Some(remaining) = first_token_timeout.checked_sub(first_token_started_at.elapsed()) else { - return Err(CodexWebSocketExchangeError::FirstTokenTimeout { - timeout: first_token_timeout, - }); - }; - if remaining.is_zero() { - return Err(CodexWebSocketExchangeError::FirstTokenTimeout { - timeout: first_token_timeout, - }); - } - Ok(WEBSOCKET_RECEIVE_IDLE_TIMEOUT.min(remaining)) -} - -fn elapsed_first_token_timeout( - first_token_started_at: Instant, - first_token_timeout: Option, -) -> Option { - let first_token_timeout = first_token_timeout.filter(|timeout| !timeout.is_zero())?; - (first_token_started_at.elapsed() >= first_token_timeout).then_some(first_token_timeout) -} - -fn event_is_first_content(event: Option<&str>) -> bool { - match event { - Some(event) => !is_pre_content_lifecycle_event(event), - None => true, - } -} - -fn is_first_token_timeout(error: &CodexWebSocketExchangeError) -> bool { - matches!(error, CodexWebSocketExchangeError::FirstTokenTimeout { .. }) -} - -async fn next_websocket_message( - websocket: &mut PumpedWebSocket, - receive_timeout: Duration, -) -> Result, CodexWebSocketExchangeError> { - match timeout(receive_timeout, websocket.next()).await { - Ok(message) => message.transpose().map_err(Into::into), - Err(_) => Err(CodexWebSocketExchangeError::ReceiveIdleTimeout { - timeout: receive_timeout, - }), - } -} - -async fn finish_stream_websocket( - mut websocket: PumpedWebSocket, - metadata: CodexWebSocketConnectionMetadata, - pool_return: Option, -) { - let Some(pool_return) = pool_return else { - websocket.close().await; - return; - }; - pool_return - .pool - .put( - pool_return.key, - PooledWebSocketConnection { - websocket, - metadata: reusable_websocket_metadata(metadata), - created_at: pool_return.created_at, - }, - ) - .await; -} - -async fn discard_stream_websocket( - mut websocket: PumpedWebSocket, - pool_return: Option, -) { - if let Some(pool_return) = pool_return { - pool_return.pool.discard(&pool_return.key).await; - } - websocket.close().await; -} - -// --------------------------------------------------------------------------- -// Header / metadata helpers -// --------------------------------------------------------------------------- - -fn websocket_rate_limit_event_headers(raw: &str) -> Option> { - events::parse_rate_limits_event_raw(raw) - .map(|parsed| events::rate_limits_to_header_pairs(&parsed)) -} - -fn audit_header_value(name: &str, value: &str) -> String { - if is_sensitive_opening_header(name) { - REDACTED_HEADER_VALUE.to_string() - } else { - value.to_string() - } -} - -fn is_sensitive_opening_header(name: &str) -> bool { - matches!( - name.to_ascii_lowercase().as_str(), - "authorization" - | "chatgpt-account-id" - | "cookie" - | "session_id" - | "session-id" - | "thread-id" - | "x-client-request-id" - | "x-codex-window-id" - | "x-codex-turn-metadata" - | "x-codex-turn-state" - | "x-codex-parent-thread-id" - ) -} - -fn websocket_audit_file_name() -> String { - let timestamp = china_filename_timestamp_millis(&Utc::now()); - format!("codex-ws-audit-{timestamp}-{}.json", Uuid::new_v4()) -} diff --git a/backend/src/web/mod.rs b/backend/src/web/mod.rs deleted file mode 100644 index 57bbc6d7a..000000000 --- a/backend/src/web/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -//! 前端静态资源模块 —— SPA 路由与构建产物服务。 - -pub mod assets; diff --git a/backend/tests/upstream/accounts/account_pool/mod.rs b/backend/tests/accounts/account_pool/mod.rs similarity index 90% rename from backend/tests/upstream/accounts/account_pool/mod.rs rename to backend/tests/accounts/account_pool/mod.rs index e07b16ea3..a71e17124 100644 --- a/backend/tests/upstream/accounts/account_pool/mod.rs +++ b/backend/tests/accounts/account_pool/mod.rs @@ -1,8 +1,8 @@ use std::collections::{BTreeMap, BTreeSet}; use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::upstream::accounts::{ - model::{Account, AccountStatus}, +use codex_proxy_rs::accounts::{ + account::{Account, AccountStatus}, pool::{ AccountAcquireRequest, AccountPool, AccountPoolOptions, AccountWindowUsageDelta, RotationStrategy, diff --git a/backend/tests/upstream/accounts/account_pool/quota.rs b/backend/tests/accounts/account_pool/quota.rs similarity index 100% rename from backend/tests/upstream/accounts/account_pool/quota.rs rename to backend/tests/accounts/account_pool/quota.rs diff --git a/backend/tests/upstream/accounts/account_pool/selection.rs b/backend/tests/accounts/account_pool/selection.rs similarity index 94% rename from backend/tests/upstream/accounts/account_pool/selection.rs rename to backend/tests/accounts/account_pool/selection.rs index 7e6e9af42..83200af6b 100644 --- a/backend/tests/upstream/accounts/account_pool/selection.rs +++ b/backend/tests/accounts/account_pool/selection.rs @@ -171,11 +171,10 @@ fn acquire_should_skip_accounts_with_expired_jwt_when_metadata_is_missing() { #[tokio::test] async fn runtime_account_pool_should_persist_expired_status_when_jwt_expiry_is_discovered() { - let (pool, _dir) = - crate::support::sqlite::init_test_db("runtime-pool-expired-status.sqlite").await; - let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()); + let (pool, _dir) = crate::support::storage::init_test_db("runtime-pool-expired-status").await; + let store = codex_proxy_rs::accounts::store::PgAccountStore::new(pool.clone()); store - .insert(codex_proxy_rs::upstream::accounts::store::NewAccount { + .insert(codex_proxy_rs::accounts::store::NewAccount { id: "acct_expired".to_string(), email: None, account_id: Some("chatgpt-expired".to_string()), @@ -190,17 +189,17 @@ async fn runtime_account_pool_should_persist_expired_status_when_jwt_expiry_is_d }) .await .unwrap(); - let runtime_pool = codex_proxy_rs::upstream::accounts::pool::RuntimeAccountPoolService::new( + let runtime_pool = codex_proxy_rs::accounts::pool::RuntimeAccountPoolService::new( std::sync::Arc::new(store), AccountPoolOptions::default(), 0, ); - runtime_pool.restore_from_repository().await.unwrap(); + runtime_pool.restore_from_store().await.unwrap(); let acquired = runtime_pool .acquire_with(&AccountAcquireRequest::new("gpt-5.5", Utc::now())) .await; - let status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_expired") .fetch_one(&pool) .await @@ -214,10 +213,10 @@ async fn runtime_account_pool_should_persist_expired_status_when_jwt_expiry_is_d async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_expires() { let now = fixed_time(); let (pool, _dir) = - crate::support::sqlite::init_test_db("runtime-pool-quota-window-reset.sqlite").await; - let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()); + crate::support::storage::init_test_db("runtime-pool-quota-window-reset").await; + let store = codex_proxy_rs::accounts::store::PgAccountStore::new(pool.clone()); store - .insert(codex_proxy_rs::upstream::accounts::store::NewAccount { + .insert(codex_proxy_rs::accounts::store::NewAccount { id: "acct_quota_reset".to_string(), email: None, account_id: Some("chatgpt-quota-reset".to_string()), @@ -236,15 +235,15 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex r#" update accounts set - quota_limit_reached = 1, - quota_verify_required = 0, - quota_cooldown_until = ?, - cloudflare_cooldown_until = ? - where id = ? + quota_limit_reached = true, + quota_verify_required = false, + quota_cooldown_until = $1, + cloudflare_cooldown_until = $2 + where id = $3 "#, ) - .bind((now - Duration::seconds(1)).to_rfc3339()) - .bind((now - Duration::seconds(1)).to_rfc3339()) + .bind(now - Duration::seconds(1)) + .bind(now - Duration::seconds(1)) .bind("acct_quota_reset") .execute(&pool) .await @@ -261,21 +260,21 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex window_started_at, window_reset_at, limit_window_seconds - ) values (?, 9, 7, 11, 13, 17, ?, ?, 60) + ) values ($1, 9, 7, 11, 13, 17, $2, $3, 60) "#, ) .bind("acct_quota_reset") - .bind((now - Duration::seconds(120)).to_rfc3339()) - .bind((now - Duration::seconds(1)).to_rfc3339()) + .bind(now - Duration::seconds(120)) + .bind(now - Duration::seconds(1)) .execute(&pool) .await .unwrap(); - let runtime_pool = codex_proxy_rs::upstream::accounts::pool::RuntimeAccountPoolService::new( + let runtime_pool = codex_proxy_rs::accounts::pool::RuntimeAccountPoolService::new( std::sync::Arc::new(store), AccountPoolOptions::default(), 0, ); - runtime_pool.restore_from_repository().await.unwrap(); + runtime_pool.restore_from_store().await.unwrap(); let acquired = runtime_pool .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) @@ -283,15 +282,15 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex .unwrap(); type StoredQuotaResetState = ( String, - i64, - i64, - Option, - Option, + bool, + bool, + Option>, + Option>, i64, i64, i64, i64, - Option, + Option>, Option, ); let ( @@ -322,23 +321,20 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex au.limit_window_seconds from accounts a left join account_usage au on au.account_id = a.id - where a.id = ? + where a.id = $1 "#, ) .bind("acct_quota_reset") .fetch_one(&pool) .await .unwrap(); - let persisted_reset_at = stored_window_reset_at - .as_deref() - .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) - .map(|value| value.with_timezone(&Utc)); + let persisted_reset_at = stored_window_reset_at; assert_eq!(acquired.account.id, "acct_quota_reset"); assert!(acquired.account.quota_verify_required); assert_eq!(stored_status, "active"); - assert_eq!(stored_quota_limit_reached, 0); - assert_eq!(stored_quota_verify_required, 1); + assert!(!stored_quota_limit_reached); + assert!(stored_quota_verify_required); assert!(stored_quota_cooldown_until.is_none()); assert!(stored_cloudflare_cooldown_until.is_none()); assert_eq!( @@ -353,7 +349,7 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex assert!(persisted_reset_at.is_some_and(|reset_at| reset_at > now)); assert_eq!(stored_limit_window_seconds, Some(60)); let usage_record_count: i64 = - sqlx::query_scalar("select count(*) from usage_records where account_id = ?") + sqlx::query_scalar("select count(*) from usage_records where account_id = $1") .bind("acct_quota_reset") .fetch_one(&pool) .await diff --git a/backend/tests/upstream/accounts/account_pool/usage_window.rs b/backend/tests/accounts/account_pool/usage_window.rs similarity index 100% rename from backend/tests/upstream/accounts/account_pool/usage_window.rs rename to backend/tests/accounts/account_pool/usage_window.rs diff --git a/backend/tests/upstream/accounts/account_repository.rs b/backend/tests/accounts/account_store.rs similarity index 72% rename from backend/tests/upstream/accounts/account_repository.rs rename to backend/tests/accounts/account_store.rs index 4cf888d3d..c8e81588f 100644 --- a/backend/tests/upstream/accounts/account_repository.rs +++ b/backend/tests/accounts/account_store.rs @@ -1,19 +1,18 @@ use chrono::{DateTime, Duration, Utc}; -use codex_proxy_rs::{ - infra::database::connect_sqlite, - upstream::accounts::{ - model::{AccountStatus, AccountUsageDelta}, - quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at}, - store::{AccountClaimsUpdate, AccountStore, NewAccount, SqliteAccountStore}, - }, +use codex_proxy_rs::accounts::{ + account::{AccountStatus, AccountUsageDelta}, + quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at}, + store::{AccountClaimsUpdate, AccountStore, NewAccount, PgAccountStore}, }; use secrecy::{ExposeSecret, SecretString}; use serde_json::json; +use crate::support::storage::init_test_db; + #[tokio::test] -async fn account_repository_should_store_plain_tokens_and_load_secret_wrappers() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 4).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_store_plain_tokens_and_load_secret_wrappers() { + let (pool, _dir) = account_store_parts("accounts", 4).await; + let repo = PgAccountStore::new(pool.clone()); let expires_at = Utc::now() + Duration::hours(1); repo.insert(NewAccount { @@ -33,7 +32,7 @@ async fn account_repository_should_store_plain_tokens_and_load_secret_wrappers() .unwrap(); let stored: (String, String) = - sqlx::query_as("select access_token, refresh_token from accounts where id = ?") + sqlx::query_as("select access_token, refresh_token from accounts where id = $1") .bind("acct_a") .fetch_one(&pool) .await @@ -52,9 +51,9 @@ async fn account_repository_should_store_plain_tokens_and_load_secret_wrappers() } #[tokio::test] -async fn account_repository_should_cursor_page_accounts_by_added_at_desc() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 5).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_cursor_page_accounts_by_added_at_desc() { + let (pool, _dir) = account_store_parts("accounts", 5).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; seed_repo_account(&pool, "acct_b", "2026-06-11T00:01:00Z").await; @@ -82,18 +81,18 @@ async fn account_repository_should_cursor_page_accounts_by_added_at_desc() { } #[tokio::test] -async fn account_repository_should_list_metadata_without_exposing_tokens() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 35).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_list_metadata_without_exposing_tokens() { + let (pool, _dir) = account_store_parts("accounts", 35).await; + let repo = PgAccountStore::new(pool.clone()); sqlx::query( - "insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)", + "insert into accounts (id, email, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6)", ) .bind("acct_plain") .bind("user@example.com") .bind("plain-access-token") .bind("active") - .bind("2026-06-11T00:00:00Z") - .bind("2026-06-11T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-11T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-06-11T00:00:00Z")) .execute(&pool) .await .unwrap(); @@ -106,14 +105,14 @@ async fn account_repository_should_list_metadata_without_exposing_tokens() { } #[tokio::test] -async fn account_repository_should_list_pool_accounts_with_tokens_and_plan() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 39).await; - let now = Utc::now().to_rfc3339(); +async fn account_store_should_list_pool_accounts_with_tokens_and_plan() { + let (pool, _dir) = account_store_parts("accounts", 39).await; + let now = Utc::now(); sqlx::query( "insert into accounts ( id, email, chatgpt_account_id, chatgpt_user_id, label, plan_type, access_token, refresh_token, access_token_expires_at, status, added_at, updated_at - ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)", ) .bind("acct_1") .bind("user@example.com") @@ -123,15 +122,15 @@ async fn account_repository_should_list_pool_accounts_with_tokens_and_plan() { .bind("plus") .bind("access-token") .bind(Some("refresh-token")) - .bind(Option::::None) + .bind(Option::>::None) .bind("active") - .bind(&now) - .bind(&now) + .bind(now) + .bind(now) .execute(&pool) .await .expect("insert account"); - let store = SqliteAccountStore::new(pool); + let store = PgAccountStore::new(pool); let accounts = AccountStore::list_pool_accounts(&store) .await .expect("list accounts"); @@ -143,9 +142,9 @@ async fn account_repository_should_list_pool_accounts_with_tokens_and_plan() { } #[tokio::test] -async fn account_repository_should_update_status_and_label_without_rewriting_tokens() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 6).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_update_status_and_label_without_rewriting_tokens() { + let (pool, _dir) = account_store_parts("accounts", 6).await; + let repo = PgAccountStore::new(pool.clone()); repo.insert(NewAccount { id: "acct_a".to_string(), @@ -162,7 +161,7 @@ async fn account_repository_should_update_status_and_label_without_rewriting_tok }) .await .unwrap(); - let before: (String,) = sqlx::query_as("select access_token from accounts where id = ?") + let before: (String,) = sqlx::query_as("select access_token from accounts where id = $1") .bind("acct_a") .fetch_one(&pool) .await @@ -177,27 +176,28 @@ async fn account_repository_should_update_status_and_label_without_rewriting_tok .await .unwrap()); - let after: (String,) = sqlx::query_as("select access_token from accounts where id = ?") - .bind("acct_a") - .fetch_one(&pool) - .await - .unwrap(); - let updated_at: (String,) = sqlx::query_as("select updated_at from accounts where id = ?") + let after: (String,) = sqlx::query_as("select access_token from accounts where id = $1") .bind("acct_a") .fetch_one(&pool) .await .unwrap(); + let updated_at: (DateTime,) = + sqlx::query_as("select updated_at from accounts where id = $1") + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); let loaded = repo.get("acct_a").await.unwrap().unwrap(); assert_eq!(after.0, before.0); assert_eq!(loaded.status, AccountStatus::Disabled); assert_eq!(loaded.label.as_deref(), Some("work")); - assert!(updated_at.0.parse::>().is_ok()); + assert!(updated_at.0 <= Utc::now()); } #[tokio::test] -async fn account_repository_should_not_reactivate_disabled_or_banned_accounts_from_claims_update() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 11).await; - let repo = SqliteAccountStore::new(pool); +async fn account_store_should_not_reactivate_disabled_or_banned_accounts_from_claims_update() { + let (pool, _dir) = account_store_parts("accounts", 11).await; + let repo = PgAccountStore::new(pool); seed_new_account(&repo, "acct_disabled").await; seed_new_account(&repo, "acct_banned").await; repo.set_status("acct_disabled", AccountStatus::Disabled) @@ -253,9 +253,9 @@ async fn account_repository_should_not_reactivate_disabled_or_banned_accounts_fr } #[tokio::test] -async fn account_repository_claims_update_should_preserve_optional_claims_when_missing() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 49).await; - let repo = SqliteAccountStore::new(pool); +async fn account_store_claims_update_should_preserve_optional_claims_when_missing() { + let (pool, _dir) = account_store_parts("accounts", 49).await; + let repo = PgAccountStore::new(pool); seed_plan_account(&repo, "acct_preserve_without_refresh").await; repo.update_from_claims( @@ -287,10 +287,9 @@ async fn account_repository_claims_update_should_preserve_optional_claims_when_m } #[tokio::test] -async fn account_repository_claims_update_should_preserve_optional_claims_when_refresh_token_rotates( -) { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 50).await; - let repo = SqliteAccountStore::new(pool); +async fn account_store_claims_update_should_preserve_optional_claims_when_refresh_token_rotates() { + let (pool, _dir) = account_store_parts("accounts", 50).await; + let repo = PgAccountStore::new(pool); seed_plan_account(&repo, "acct_preserve_with_refresh").await; repo.update_from_claims( @@ -322,9 +321,9 @@ async fn account_repository_claims_update_should_preserve_optional_claims_when_r } #[tokio::test] -async fn account_repository_should_accumulate_usage_counters() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 7).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_accumulate_usage_counters() { + let (pool, _dir) = account_store_parts("accounts", 7).await; + let repo = PgAccountStore::new(pool.clone()); seed_new_account(&repo, "acct_a").await; AccountStore::record_usage_delta( @@ -360,7 +359,7 @@ async fn account_repository_should_accumulate_usage_counters() { .await .unwrap(); - let usage: ( + type AccountUsageRow = ( i64, i64, i64, @@ -368,15 +367,16 @@ async fn account_repository_should_accumulate_usage_counters() { i64, i64, i64, - Option, - Option, - ) = sqlx::query_as( + Option>, + Option>, + ); + let usage: AccountUsageRow = sqlx::query_as( r" select request_count, input_tokens, output_tokens, cached_tokens, reasoning_tokens, total_tokens, empty_response_count, window_started_at, last_used_at from account_usage - where account_id = ? + where account_id = $1 ", ) .bind("acct_a") @@ -395,19 +395,19 @@ async fn account_repository_should_accumulate_usage_counters() { } #[tokio::test] -async fn account_repository_should_restore_window_usage_into_runtime_pool_accounts() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 33).await; - let repo = SqliteAccountStore::new(pool.clone()); - let window_started_at = Utc::now() - Duration::minutes(2); - let window_reset_at = Utc::now() + Duration::minutes(3); +async fn account_store_should_restore_window_usage_into_runtime_pool_accounts() { + let (pool, _dir) = account_store_parts("accounts", 33).await; + let repo = PgAccountStore::new(pool.clone()); + let window_started_at = pg_timestamp(Utc::now() - Duration::minutes(2)); + let window_reset_at = pg_timestamp(Utc::now() + Duration::minutes(3)); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; sqlx::query( - "insert into account_usage (account_id, request_count, 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) values (?, 9, 31, 9, 2, 1, 7, 11, 13, 17, 19, 5, 1, 1, ?, ?, 300)", + "insert into account_usage (account_id, request_count, 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) values ($1, 9, 31, 9, 2, 1, 7, 11, 13, 17, 19, 5, 1, 1, $2, $3, 300)", ) .bind("acct_a") - .bind(window_started_at.to_rfc3339()) - .bind(window_reset_at.to_rfc3339()) + .bind(window_started_at) + .bind(window_reset_at) .execute(&pool) .await .unwrap(); @@ -435,13 +435,13 @@ async fn account_repository_should_restore_window_usage_into_runtime_pool_accoun } #[tokio::test] -async fn account_repository_should_restore_window_from_quota_json_when_usage_window_is_missing() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 36).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_restore_window_from_quota_json_when_usage_window_is_missing() { + let (pool, _dir) = account_store_parts("accounts", 36).await; + let repo = PgAccountStore::new(pool.clone()); let reset_at = DateTime::::from_timestamp(1_806_364_800, 0).unwrap(); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; - sqlx::query("update accounts set quota_json = ? where id = ?") + sqlx::query("update accounts set quota_json = $1::jsonb where id = $2") .bind( r#"{ "snapshots": [{ @@ -482,12 +482,12 @@ async fn account_repository_should_restore_window_from_quota_json_when_usage_win } #[tokio::test] -async fn account_repository_should_restore_quota_verify_required_into_runtime_pool_accounts() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 34).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_restore_quota_verify_required_into_runtime_pool_accounts() { + let (pool, _dir) = account_store_parts("accounts", 34).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; - sqlx::query("update accounts set quota_verify_required = 1 where id = ?") + sqlx::query("update accounts set quota_verify_required = true where id = $1") .bind("acct_a") .execute(&pool) .await @@ -502,9 +502,9 @@ async fn account_repository_should_restore_quota_verify_required_into_runtime_po } #[tokio::test] -async fn account_repository_should_update_quota_json_and_fetched_at() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 36).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_update_quota_json_and_fetched_at() { + let (pool, _dir) = account_store_parts("accounts", 36).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; let updated = repo @@ -518,19 +518,20 @@ async fn account_repository_should_update_quota_json_and_fetched_at() { .update_quota_json("missing", r#"{"rate_limit":null}"#) .await .unwrap(); - let row: (String, Option, Option) = - sqlx::query_as("select quota_json, quota_fetched_at, plan_type from accounts where id = ?") - .bind("acct_a") - .fetch_one(&pool) - .await - .unwrap(); + let row: (serde_json::Value, Option>, Option) = sqlx::query_as( + "select quota_json, quota_fetched_at, plan_type from accounts where id = $1", + ) + .bind("acct_a") + .fetch_one(&pool) + .await + .unwrap(); assert_eq!( (updated, missing_updated, row.0, row.1.is_some(), row.2), ( true, false, - r#"{"plan_type":"free","rate_limit":{"limit_reached":false}}"#.to_string(), + json!({"plan_type":"free","rate_limit":{"limit_reached":false}}), true, Some("free".to_string()) ) @@ -538,9 +539,9 @@ async fn account_repository_should_update_quota_json_and_fetched_at() { } #[tokio::test] -async fn account_repository_should_map_quota_snapshot_to_account_status() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 37).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_map_quota_snapshot_to_account_status() { + let (pool, _dir) = account_store_parts("accounts", 37).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; let reset_at = Utc::now() + Duration::minutes(5); @@ -553,8 +554,8 @@ async fn account_repository_should_map_quota_snapshot_to_account_status() { ) .await .unwrap(); - let limited_status: (String, i64, Option) = sqlx::query_as( - "select status, quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let limited_status: (String, bool, Option>) = sqlx::query_as( + "select status, quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct_a") .fetch_one(&pool) @@ -563,7 +564,7 @@ async fn account_repository_should_map_quota_snapshot_to_account_status() { assert!(limited); assert_eq!(limited_status.0, "quota_exhausted"); - assert_eq!(limited_status.1, 1); + assert!(limited_status.1); assert!(limited_status.2.is_some()); let restored = repo @@ -575,8 +576,8 @@ async fn account_repository_should_map_quota_snapshot_to_account_status() { ) .await .unwrap(); - let restored_status: (String, i64, Option) = sqlx::query_as( - "select status, quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let restored_status: (String, bool, Option>) = sqlx::query_as( + "select status, quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct_a") .fetch_one(&pool) @@ -585,14 +586,14 @@ async fn account_repository_should_map_quota_snapshot_to_account_status() { assert!(restored); assert_eq!(restored_status.0, "active"); - assert_eq!(restored_status.1, 0); + assert!(!restored_status.1); assert!(restored_status.2.is_none()); } #[tokio::test] -async fn account_repository_should_mark_quota_exhausted_for_any_exhausted_window() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 47).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_mark_quota_exhausted_for_any_exhausted_window() { + let (pool, _dir) = account_store_parts("accounts", 47).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; let quota = json!({ "plan_type": "plus", @@ -617,8 +618,8 @@ async fn account_repository_should_mark_quota_exhausted_for_any_exhausted_window .apply_quota_snapshot("acct_a", "a.to_string(), limit_reached, cooldown_until) .await .unwrap(); - let status: (String, i64, Option) = sqlx::query_as( - "select status, quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let status: (String, bool, Option>) = sqlx::query_as( + "select status, quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct_a") .fetch_one(&pool) @@ -627,22 +628,22 @@ async fn account_repository_should_mark_quota_exhausted_for_any_exhausted_window assert_eq!( (updated, status.0, status.1, status.2.is_some()), - (true, "quota_exhausted".to_string(), 1, true) + (true, "quota_exhausted".to_string(), true, true) ); } #[tokio::test] -async fn account_repository_should_mark_quota_cooldown_as_quota_exhausted_status() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 48).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_mark_quota_cooldown_as_quota_exhausted_status() { + let (pool, _dir) = account_store_parts("accounts", 48).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; - let future_cooldown = Utc::now() + Duration::hours(1); + let future_cooldown = pg_timestamp(Utc::now() + Duration::hours(1)); repo.mark_quota_limited_until("acct_a", future_cooldown) .await .unwrap(); - let stored: (String, i64, Option) = sqlx::query_as( - "select status, quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let stored: (String, bool, Option>) = sqlx::query_as( + "select status, quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct_a") .fetch_one(&pool) @@ -651,24 +652,20 @@ async fn account_repository_should_mark_quota_cooldown_as_quota_exhausted_status assert_eq!( (stored.0, stored.1, stored.2), - ( - String::from("quota_exhausted"), - 1, - Some(future_cooldown.to_rfc3339()) - ) + (String::from("quota_exhausted"), true, Some(future_cooldown)) ); } #[tokio::test] -async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quota_cooldown() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 42).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_runtime_state_sync_should_not_clear_newer_future_quota_cooldown() { + let (pool, _dir) = account_store_parts("accounts", 42).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; - let future_cooldown = Utc::now() + Duration::hours(1); + let future_cooldown = pg_timestamp(Utc::now() + Duration::hours(1)); sqlx::query( - "update accounts set status = 'quota_exhausted', quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ? where id = ?", + "update accounts set status = 'quota_exhausted', quota_limit_reached = true, quota_verify_required = false, quota_cooldown_until = $1 where id = $2", ) - .bind(future_cooldown.to_rfc3339()) + .bind(future_cooldown) .bind("acct_a") .execute(&pool) .await @@ -687,8 +684,8 @@ async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quo repo.sync_runtime_account_state(&stale_account, false) .await .unwrap(); - let stored: (String, i64, i64, Option) = sqlx::query_as( - "select status, quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = ?", + let stored: (String, bool, bool, Option>) = sqlx::query_as( + "select status, quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = $1", ) .bind("acct_a") .fetch_one(&pool) @@ -696,18 +693,18 @@ async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quo .unwrap(); assert_eq!(stored.0, "quota_exhausted"); - assert_eq!(stored.1, 1); - assert_eq!(stored.2, 0); - assert_eq!(stored.3, Some(future_cooldown.to_rfc3339())); + assert!(stored.1); + assert!(!stored.2); + assert_eq!(stored.3, Some(future_cooldown)); } #[tokio::test] -async fn account_repository_runtime_state_sync_should_not_regress_newer_usage_window() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 43).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_runtime_state_sync_should_not_regress_newer_usage_window() { + let (pool, _dir) = account_store_parts("accounts", 43).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; - let now = Utc::now(); - let future_reset = now + Duration::hours(1); + let now = pg_timestamp(Utc::now()); + let future_reset = pg_timestamp(now + Duration::hours(1)); let stale_reset = now + Duration::minutes(1); sqlx::query( r#" @@ -717,11 +714,11 @@ async fn account_repository_runtime_state_sync_should_not_regress_newer_usage_wi window_input_tokens, window_reset_at, limit_window_seconds - ) values (?, 5, 17, ?, 3600) + ) values ($1, 5, 17, $2, 3600) "#, ) .bind("acct_a") - .bind(future_reset.to_rfc3339()) + .bind(future_reset) .execute(&pool) .await .unwrap(); @@ -740,8 +737,8 @@ async fn account_repository_runtime_state_sync_should_not_regress_newer_usage_wi repo.sync_runtime_account_state(&stale_account, true) .await .unwrap(); - let stored: (i64, i64, Option, Option) = sqlx::query_as( - "select window_request_count, window_input_tokens, window_reset_at, limit_window_seconds from account_usage where account_id = ?", + let stored: (i64, i64, Option>, Option) = sqlx::query_as( + "select window_request_count, window_input_tokens, window_reset_at, limit_window_seconds from account_usage where account_id = $1", ) .bind("acct_a") .fetch_one(&pool) @@ -750,14 +747,14 @@ async fn account_repository_runtime_state_sync_should_not_regress_newer_usage_wi assert_eq!(stored.0, 5); assert_eq!(stored.1, 17); - assert_eq!(stored.2, Some(future_reset.to_rfc3339())); + assert_eq!(stored.2, Some(future_reset)); assert_eq!(stored.3, Some(3600)); } #[tokio::test] -async fn account_repository_should_not_override_disabled_status_from_quota_snapshot() { - let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 38).await; - let repo = SqliteAccountStore::new(pool.clone()); +async fn account_store_should_not_override_disabled_status_from_quota_snapshot() { + let (pool, _dir) = account_store_parts("accounts", 38).await; + let repo = PgAccountStore::new(pool.clone()); seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await; repo.set_status("acct_a", AccountStatus::Disabled) .await @@ -771,7 +768,7 @@ async fn account_repository_should_not_override_disabled_status_from_quota_snaps ) .await .unwrap(); - let status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_a") .fetch_one(&pool) .await @@ -780,18 +777,11 @@ async fn account_repository_should_not_override_disabled_status_from_quota_snaps assert_eq!(status.0, "disabled"); } -async fn sqlite_account_store_parts( - db_name: &str, - _key_byte: u8, -) -> (sqlx::SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - (pool, dir) +async fn account_store_parts(db_name: &str, _key_byte: u8) -> (sqlx::PgPool, tempfile::TempDir) { + init_test_db(db_name).await } -async fn seed_new_account(repo: &SqliteAccountStore, id: &str) { +async fn seed_new_account(repo: &PgAccountStore, id: &str) { repo.insert(NewAccount { id: id.to_string(), email: None, @@ -809,7 +799,7 @@ async fn seed_new_account(repo: &SqliteAccountStore, id: &str) { .unwrap(); } -async fn seed_plan_account(repo: &SqliteAccountStore, id: &str) { +async fn seed_plan_account(repo: &PgAccountStore, id: &str) { repo.insert(NewAccount { id: id.to_string(), email: Some("existing@example.com".to_string()), @@ -827,16 +817,20 @@ async fn seed_plan_account(repo: &SqliteAccountStore, id: &str) { .unwrap(); } -async fn seed_repo_account(pool: &sqlx::SqlitePool, id: &str, added_at: &str) { +async fn seed_repo_account(pool: &sqlx::PgPool, id: &str, added_at: &str) { sqlx::query( - "insert into accounts (id, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?)", + "insert into accounts (id, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5)", ) .bind(id) .bind(format!("access-{id}")) .bind("active") - .bind(added_at) - .bind(added_at) + .bind(crate::support::storage::timestamp(added_at)) + .bind(crate::support::storage::timestamp(added_at)) .execute(pool) .await .unwrap(); } + +fn pg_timestamp(value: DateTime) -> DateTime { + DateTime::from_timestamp_micros(value.timestamp_micros()).unwrap() +} diff --git a/backend/tests/upstream/accounts/cloudflare.rs b/backend/tests/accounts/cloudflare.rs similarity index 94% rename from backend/tests/upstream/accounts/cloudflare.rs rename to backend/tests/accounts/cloudflare.rs index 634a0d6ce..21b2bae3d 100644 --- a/backend/tests/upstream/accounts/cloudflare.rs +++ b/backend/tests/accounts/cloudflare.rs @@ -1,5 +1,5 @@ use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::upstream::accounts::cookies::CloudflareChallengeCooldownTracker; +use codex_proxy_rs::accounts::cookies::CloudflareChallengeCooldownTracker; #[tokio::test] async fn challenge_cooldown_should_escalate_and_cap_delay() { diff --git a/backend/tests/upstream/accounts/cookies.rs b/backend/tests/accounts/cookies.rs similarity index 62% rename from backend/tests/upstream/accounts/cookies.rs rename to backend/tests/accounts/cookies.rs index 44ddff4d3..2a7a3b714 100644 --- a/backend/tests/upstream/accounts/cookies.rs +++ b/backend/tests/accounts/cookies.rs @@ -1,27 +1,26 @@ use chrono::{Duration, Utc}; -use codex_proxy_rs::{ - infra::database::connect_sqlite, - upstream::accounts::{ - cookies::SqliteCookieStore, - model::AccountStatus, - store::{NewAccount, SqliteAccountStore}, - }, +use codex_proxy_rs::accounts::{ + account::AccountStatus, + cookies::PgCookieStore, + store::{NewAccount, PgAccountStore}, }; use secrecy::SecretString; -use sqlx::SqlitePool; +use sqlx::PgPool; + +use crate::support::storage::init_test_db; #[tokio::test] async fn cookie_store_should_cleanup_expired_cookies() { - let (pool, _dir) = sqlite_cookie_store_parts("cookies-cleanup.sqlite").await; + let (pool, _dir) = cookie_store_parts("cookies-cleanup").await; seed_account(&pool, "acct_a").await; - let store = SqliteCookieStore::new(pool.clone()); + let store = PgCookieStore::new(pool.clone()); let now = Utc::now(); let past = now - Duration::hours(2); let future = now + Duration::hours(2); sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "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("cookie_expired") .bind("acct_a") @@ -29,14 +28,14 @@ async fn cookie_store_should_cleanup_expired_cookies() { .bind("cf_clearance") .bind("old_cipher") .bind("/") - .bind(past.to_rfc3339()) - .bind(past.to_rfc3339()) + .bind(past) + .bind(past) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "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("cookie_valid") .bind("acct_a") @@ -44,8 +43,8 @@ async fn cookie_store_should_cleanup_expired_cookies() { .bind("cf_clearance") .bind("new_cipher") .bind("/valid") - .bind(future.to_rfc3339()) - .bind(now.to_rfc3339()) + .bind(future) + .bind(now) .execute(&pool) .await .unwrap(); @@ -60,16 +59,12 @@ async fn cookie_store_should_cleanup_expired_cookies() { assert_eq!(remaining.0, 1); } -async fn sqlite_cookie_store_parts(name: &str) -> (SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - (pool, dir) +async fn cookie_store_parts(name: &str) -> (PgPool, tempfile::TempDir) { + init_test_db(name).await } -async fn seed_account(pool: &SqlitePool, id: &str) { - SqliteAccountStore::new(pool.clone()) +async fn seed_account(pool: &PgPool, id: &str) { + PgAccountStore::new(pool.clone()) .insert(NewAccount { id: id.to_string(), email: Some(format!("{id}@example.com")), diff --git a/backend/tests/upstream/accounts/mod.rs b/backend/tests/accounts/mod.rs similarity index 75% rename from backend/tests/upstream/accounts/mod.rs rename to backend/tests/accounts/mod.rs index 13df78c60..2fd305afb 100644 --- a/backend/tests/upstream/accounts/mod.rs +++ b/backend/tests/accounts/mod.rs @@ -1,7 +1,8 @@ pub(crate) mod account_pool; -pub(crate) mod account_repository; +pub(crate) mod account_store; pub(crate) mod cloudflare; pub(crate) mod cookies; pub(crate) mod quota_refresh; pub(crate) mod refresh_leases; +pub(crate) mod scheduler; pub(crate) mod token_refresh; diff --git a/backend/tests/upstream/accounts/quota_refresh.rs b/backend/tests/accounts/quota_refresh.rs similarity index 87% rename from backend/tests/upstream/accounts/quota_refresh.rs rename to backend/tests/accounts/quota_refresh.rs index a224b8f67..eb23f1fbf 100644 --- a/backend/tests/upstream/accounts/quota_refresh.rs +++ b/backend/tests/accounts/quota_refresh.rs @@ -1,15 +1,14 @@ use std::{collections::HashMap, sync::Arc, time::Duration as StdDuration}; use chrono::{Duration, Utc}; -use codex_proxy_rs::infra::database::connect_sqlite; -use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore; -use codex_proxy_rs::upstream::accounts::model::AccountStatus; -use codex_proxy_rs::upstream::accounts::quota::{ +use codex_proxy_rs::accounts::account::AccountStatus; +use codex_proxy_rs::accounts::cookies::PgCookieStore; +use codex_proxy_rs::accounts::quota::{ quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, quota_snapshot_reset_at, RuntimeQuotaRefreshService, }; -use codex_proxy_rs::upstream::accounts::store::{NewAccount, SqliteAccountStore}; -use codex_proxy_rs::upstream::transport::CodexBackendClient; +use codex_proxy_rs::accounts::store::{NewAccount, PgAccountStore}; +use codex_proxy_rs::upstream::openai::transport::CodexBackendClient; use secrecy::SecretString; use serde_json::{json, Value}; use wiremock::{ @@ -17,6 +16,8 @@ use wiremock::{ Mock, MockServer, ResponseTemplate, }; +use crate::support::storage::init_test_db; + #[test] fn quota_from_usage_should_preserve_spend_control_individual_limit() { let usage = json!({ @@ -294,13 +295,9 @@ async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_con .mount(&server) .await; - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("quota-refresh-cookie.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); - let cookies = SqliteCookieStore::new(pool.clone()); + let (pool, _guard) = init_test_db("quota-refresh-cookie").await; + let store = PgAccountStore::new(pool.clone()); + let cookies = PgCookieStore::new(pool.clone()); insert_quota_locked_account(&store, &pool, "acct-quota-cookie", "access-token-cookie").await; cookies .capture_set_cookie( @@ -363,12 +360,8 @@ async fn quota_refresh_service_should_fetch_usage_for_quota_locked_accounts_and_ .mount(&server) .await; - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("quota-refresh.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); + let (pool, _guard) = init_test_db("quota-refresh").await; + let store = PgAccountStore::new(pool.clone()); store .insert(NewAccount { id: "acct-quota".to_string(), @@ -385,7 +378,7 @@ async fn quota_refresh_service_should_fetch_usage_for_quota_locked_accounts_and_ }) .await .expect("account should be inserted"); - sqlx::query("update accounts set quota_limit_reached = 1 where id = ?") + sqlx::query("update accounts set quota_limit_reached = true where id = $1") .bind("acct-quota") .execute(&pool) .await @@ -447,19 +440,15 @@ async fn quota_refresh_service_should_refresh_quota_exhausted_accounts() { .mount(&server) .await; - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("quota-refresh-exhausted.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); + let (pool, _guard) = init_test_db("quota-refresh-exhausted").await; + let store = PgAccountStore::new(pool.clone()); insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await; store .set_status("acct-quota", AccountStatus::QuotaExhausted) .await .expect("status should update"); sqlx::query( - "update accounts set quota_limit_reached = 0, quota_verify_required = 0 where id = ?", + "update accounts set quota_limit_reached = false, quota_verify_required = false where id = $1", ) .bind("acct-quota") .execute(&pool) @@ -477,29 +466,23 @@ async fn quota_refresh_service_should_refresh_quota_exhausted_accounts() { .refresh_locked_accounts(&mut last_refreshed) .await .expect("quota refresh should succeed"); - let stored: (String, i64) = - sqlx::query_as("select status, quota_limit_reached from accounts where id = ?") + let stored: (String, bool) = + sqlx::query_as("select status, quota_limit_reached from accounts where id = $1") .bind("acct-quota") .fetch_one(&pool) .await .unwrap(); assert_eq!(summary.refreshed, 1); - assert_eq!(stored, ("active".to_string(), 0)); + assert_eq!(stored, ("active".to_string(), false)); } #[tokio::test] async fn quota_refresh_service_should_skip_recent_locked_account_before_cooldown_expires() { let server = MockServer::start().await; - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir - .path() - .join("quota-refresh-recent-future-cooldown.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); + let (pool, _guard) = init_test_db("quota-refresh-recent-future-cooldown").await; + let store = PgAccountStore::new(pool.clone()); insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await; set_quota_cooldown_until(&pool, "acct-quota", Utc::now() + Duration::minutes(5)).await; let codex = CodexBackendClient::new( @@ -533,14 +516,8 @@ async fn quota_refresh_service_should_skip_recent_locked_account_before_cooldown async fn quota_refresh_service_should_skip_recent_locked_account_inside_cooldown_grace() { let server = MockServer::start().await; - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir - .path() - .join("quota-refresh-recent-cooldown-grace.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); + let (pool, _guard) = init_test_db("quota-refresh-recent-cooldown-grace").await; + let store = PgAccountStore::new(pool.clone()); insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await; set_quota_cooldown_until(&pool, "acct-quota", Utc::now() - Duration::seconds(20)).await; let codex = CodexBackendClient::new( @@ -591,14 +568,8 @@ async fn quota_refresh_service_should_bypass_recent_skip_after_cooldown_grace() .mount(&server) .await; - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir - .path() - .join("quota-refresh-recent-expired-cooldown.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); + let (pool, _guard) = init_test_db("quota-refresh-recent-expired-cooldown").await; + let store = PgAccountStore::new(pool.clone()); insert_quota_locked_account(&store, &pool, "acct-quota", "access-token").await; set_quota_cooldown_until(&pool, "acct-quota", Utc::now() - Duration::seconds(31)).await; let codex = CodexBackendClient::new( @@ -647,12 +618,8 @@ async fn quota_refresh_service_should_stagger_multiple_locked_account_requests() .mount(&server) .await; - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("quota-refresh-spacing.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); + let (pool, _guard) = init_test_db("quota-refresh-spacing").await; + let store = PgAccountStore::new(pool.clone()); insert_quota_locked_account(&store, &pool, "acct-quota-1", "access-token-1").await; insert_quota_locked_account(&store, &pool, "acct-quota-2", "access-token-2").await; let codex = CodexBackendClient::new( @@ -684,8 +651,8 @@ async fn quota_refresh_service_should_stagger_multiple_locked_account_requests() } async fn insert_quota_locked_account( - store: &SqliteAccountStore, - pool: &sqlx::SqlitePool, + store: &PgAccountStore, + pool: &sqlx::PgPool, id: &str, access_token: &str, ) { @@ -705,7 +672,7 @@ async fn insert_quota_locked_account( }) .await .expect("account should be inserted"); - sqlx::query("update accounts set quota_limit_reached = 1 where id = ?") + sqlx::query("update accounts set quota_limit_reached = true where id = $1") .bind(id) .execute(pool) .await @@ -713,12 +680,12 @@ async fn insert_quota_locked_account( } async fn set_quota_cooldown_until( - pool: &sqlx::SqlitePool, + pool: &sqlx::PgPool, id: &str, cooldown_until: chrono::DateTime, ) { - sqlx::query("update accounts set quota_cooldown_until = ? where id = ?") - .bind(cooldown_until.to_rfc3339()) + sqlx::query("update accounts set quota_cooldown_until = $1 where id = $2") + .bind(cooldown_until) .bind(id) .execute(pool) .await diff --git a/backend/tests/accounts/refresh_leases.rs b/backend/tests/accounts/refresh_leases.rs new file mode 100644 index 000000000..6a8567048 --- /dev/null +++ b/backend/tests/accounts/refresh_leases.rs @@ -0,0 +1,54 @@ +use chrono::{Duration, Utc}; +use codex_proxy_rs::accounts::refresh::RedisRefreshLeaseStore; + +use crate::support::storage::create_test_redis; + +#[tokio::test] +async fn redis_refresh_lease_acquire_renew_and_compare_release_are_atomic() { + let leases = RedisRefreshLeaseStore::new(create_test_redis("refresh-lease-atomic").await); + let now = Utc::now(); + + assert!(leases + .try_acquire("acct-lease", "owner-a", now + Duration::minutes(5), now) + .await + .unwrap()); + assert!(!leases + .try_acquire("acct-lease", "owner-b", now + Duration::minutes(5), now) + .await + .unwrap()); + assert!(leases + .try_acquire("acct-lease", "owner-a", now + Duration::minutes(10), now,) + .await + .unwrap()); + assert!(!leases.release("acct-lease", "owner-b").await.unwrap()); + assert!(leases.release("acct-lease", "owner-a").await.unwrap()); + assert!(leases + .try_acquire("acct-lease", "owner-b", now + Duration::minutes(5), now) + .await + .unwrap()); +} + +#[tokio::test] +async fn redis_refresh_lease_lists_only_existing_ttl_keys() { + let leases = RedisRefreshLeaseStore::new(create_test_redis("refresh-lease-list").await); + let now = Utc::now(); + assert!(leases + .try_acquire("acct-active", "owner", now + Duration::minutes(5), now) + .await + .unwrap()); + assert!(!leases + .try_acquire("acct-expired", "owner", now - Duration::seconds(1), now) + .await + .unwrap()); + + let ids = vec![ + "acct-active".to_string(), + "acct-expired".to_string(), + "acct-missing".to_string(), + ]; + let active = leases.active_account_ids(&ids, now).await.unwrap(); + assert_eq!( + active, + std::collections::HashSet::from(["acct-active".to_string()]) + ); +} diff --git a/backend/tests/upstream/scheduler/feedback.rs b/backend/tests/accounts/scheduler/feedback.rs similarity index 95% rename from backend/tests/upstream/scheduler/feedback.rs rename to backend/tests/accounts/scheduler/feedback.rs index 7f8266e2a..22bd5b2a0 100644 --- a/backend/tests/upstream/scheduler/feedback.rs +++ b/backend/tests/accounts/scheduler/feedback.rs @@ -1,4 +1,4 @@ -use codex_proxy_rs::upstream::scheduler::{FeedbackSample, FeedbackStats}; +use codex_proxy_rs::accounts::scheduler::{FeedbackSample, FeedbackStats}; #[test] fn sample_of_unknown_account_is_empty() { diff --git a/backend/tests/upstream/scheduler/mod.rs b/backend/tests/accounts/scheduler/mod.rs similarity index 100% rename from backend/tests/upstream/scheduler/mod.rs rename to backend/tests/accounts/scheduler/mod.rs diff --git a/backend/tests/upstream/scheduler/score.rs b/backend/tests/accounts/scheduler/score.rs similarity index 97% rename from backend/tests/upstream/scheduler/score.rs rename to backend/tests/accounts/scheduler/score.rs index bf54d7bee..58ad70b8d 100644 --- a/backend/tests/upstream/scheduler/score.rs +++ b/backend/tests/accounts/scheduler/score.rs @@ -1,6 +1,6 @@ use chrono::Utc; -use codex_proxy_rs::upstream::accounts::model::AccountStatus; -use codex_proxy_rs::upstream::scheduler::{rank_candidates, FeedbackStats, ScoreWeights}; +use codex_proxy_rs::accounts::account::AccountStatus; +use codex_proxy_rs::accounts::scheduler::{rank_candidates, FeedbackStats, ScoreWeights}; use crate::support::accounts::test_account; diff --git a/backend/tests/upstream/accounts/token_refresh/failures.rs b/backend/tests/accounts/token_refresh/failures.rs similarity index 85% rename from backend/tests/upstream/accounts/token_refresh/failures.rs rename to backend/tests/accounts/token_refresh/failures.rs index 5aa8ac6ed..39dbf24c0 100644 --- a/backend/tests/upstream/accounts/token_refresh/failures.rs +++ b/backend/tests/accounts/token_refresh/failures.rs @@ -2,12 +2,8 @@ use super::*; #[tokio::test] async fn token_refresh_task_should_keep_business_status_before_refresher_call() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-status.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-status").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 9, 0, 0).unwrap(); let new_expires_at = now + Duration::hours(1); let new_access_token = test_jwt(new_expires_at.timestamp()); @@ -39,7 +35,7 @@ async fn token_refresh_task_should_keep_business_status_before_refresher_call() refresh_token: None, }), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -70,12 +66,8 @@ async fn token_refresh_task_should_keep_business_status_before_refresher_call() #[tokio::test] async fn token_refresh_task_should_skip_expired_account_after_restart() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-recovery.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-recovery").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 10, 0, 0).unwrap(); let old_access_token = test_jwt((now + Duration::hours(1)).timestamp()); store @@ -95,7 +87,7 @@ async fn token_refresh_task_should_skip_expired_account_after_restart() { .await .expect("account should be inserted"); let refresher = CountingTokenRefresher::default(); - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -122,12 +114,8 @@ async fn token_refresh_task_should_skip_expired_account_after_restart() { #[tokio::test] async fn token_refresh_task_should_retry_transport_failure_before_recovery() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-transport-failure.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-transport-failure").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 11, 0, 0).unwrap(); let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); store @@ -153,7 +141,7 @@ async fn token_refresh_task_should_retry_transport_failure_before_recovery() { observed_statuses: observed_statuses.clone(), response: Err(RefreshFailure::Transport), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -186,12 +174,8 @@ async fn token_refresh_task_should_retry_transport_failure_before_recovery() { #[tokio::test] async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-delayed-recovery.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-delayed-recovery").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 15, 0, 0).unwrap(); let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); @@ -226,7 +210,7 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() { observed_statuses: observed_statuses.clone(), responses: Arc::new(Mutex::new(VecDeque::from(responses))), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -277,12 +261,8 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() { #[tokio::test] async fn token_refresh_task_should_not_reuse_stale_refresh_token_after_retryable_failure() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-stale-rt.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-stale-rt").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 16, 0, 0).unwrap(); let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); let rotated_access_token = test_jwt((now + Duration::hours(1)).timestamp()); @@ -309,7 +289,7 @@ async fn token_refresh_task_should_not_reuse_stale_refresh_token_after_retryable calls: calls.clone(), access_token: rotated_access_token.clone(), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -346,12 +326,8 @@ async fn token_refresh_task_should_not_reuse_stale_refresh_token_after_retryable #[tokio::test] async fn token_refresh_task_should_confirm_invalid_grant_before_expiring_account() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-permanent-failure.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-permanent-failure").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 13, 0, 0).unwrap(); let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); store @@ -382,7 +358,7 @@ async fn token_refresh_task_should_confirm_invalid_grant_before_expiring_account Err(RefreshFailure::InvalidGrant), ]))), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, diff --git a/backend/tests/upstream/accounts/token_refresh/mod.rs b/backend/tests/accounts/token_refresh/mod.rs similarity index 90% rename from backend/tests/upstream/accounts/token_refresh/mod.rs rename to backend/tests/accounts/token_refresh/mod.rs index 2c7793e1c..54adc7652 100644 --- a/backend/tests/upstream/accounts/token_refresh/mod.rs +++ b/backend/tests/accounts/token_refresh/mod.rs @@ -9,13 +9,11 @@ use std::{ use async_trait::async_trait; use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::infra::database::connect_sqlite; -use codex_proxy_rs::upstream::accounts::model::AccountStatus; -use codex_proxy_rs::upstream::accounts::store::SqliteAccountStore; -use codex_proxy_rs::upstream::accounts::store::{AccountClaimsUpdate, NewAccount}; -use codex_proxy_rs::upstream::accounts::token_refresh::{ - RefreshFailure, RefreshLeaseStore, RefreshPolicy, TokenPair, TokenRefresher, -}; +use codex_proxy_rs::accounts::account::AccountStatus; +use codex_proxy_rs::accounts::refresh::{RedisRefreshLeaseStore, RefreshPolicy}; +use codex_proxy_rs::accounts::store::PgAccountStore; +use codex_proxy_rs::accounts::store::{AccountClaimsUpdate, NewAccount}; +use codex_proxy_rs::upstream::openai::token_client::{RefreshFailure, TokenPair, TokenRefresher}; use secrecy::{ExposeSecret, SecretString}; use serde_json::json; use tokio::{ @@ -23,6 +21,8 @@ use tokio::{ time::{sleep, timeout}, }; +use crate::support::storage::{create_test_redis, init_test_db}; + mod failures; mod scheduling; mod success; @@ -91,7 +91,7 @@ impl TokenRefresher for NotifyingTokenRefresher { #[derive(Clone)] struct StoreObservingTokenRefresher { - store: SqliteAccountStore, + store: PgAccountStore, account_id: String, observed_statuses: Arc>>, response: Result, @@ -114,7 +114,7 @@ impl TokenRefresher for StoreObservingTokenRefresher { #[derive(Clone)] struct SequenceTokenRefresher { - store: SqliteAccountStore, + store: PgAccountStore, account_id: String, observed_statuses: Arc>>, responses: Arc>>>, @@ -141,7 +141,7 @@ impl TokenRefresher for SequenceTokenRefresher { #[derive(Clone)] struct RefreshTokenRotatingRefresher { - store: SqliteAccountStore, + store: PgAccountStore, account_id: String, calls: Arc, access_token: String, diff --git a/backend/tests/upstream/accounts/token_refresh/scheduling.rs b/backend/tests/accounts/token_refresh/scheduling.rs similarity index 86% rename from backend/tests/upstream/accounts/token_refresh/scheduling.rs rename to backend/tests/accounts/token_refresh/scheduling.rs index 21473c04d..f4800bef4 100644 --- a/backend/tests/upstream/accounts/token_refresh/scheduling.rs +++ b/backend/tests/accounts/token_refresh/scheduling.rs @@ -2,13 +2,9 @@ use super::*; #[tokio::test] async fn token_refresh_task_should_skip_account_when_refresh_lease_is_held() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-lease.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); - let leases = RefreshLeaseStore::new(pool.clone()); + let (pool, _guard) = init_test_db("token-refresh-lease").await; + let store = PgAccountStore::new(pool.clone()); + let leases = RedisRefreshLeaseStore::new(create_test_redis("token-refresh-lease").await); let now = Utc.with_ymd_and_hms(2026, 6, 19, 8, 0, 0).unwrap(); let old_access_token = test_jwt((now + Duration::seconds(30)).timestamp()); store @@ -37,7 +33,7 @@ async fn token_refresh_task_should_skip_account_when_refresh_lease_is_held() { .await .expect("external owner should acquire lease")); let refresher = CountingTokenRefresher::default(); - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -69,12 +65,8 @@ async fn token_refresh_task_should_skip_account_when_refresh_lease_is_held() { #[tokio::test] async fn token_refresh_task_should_skip_duplicate_in_flight_refresh() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-in-flight.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-in-flight").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 14, 0, 0).unwrap(); let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); store @@ -105,7 +97,7 @@ async fn token_refresh_task_should_skip_duplicate_in_flight_refresh() { }), }; let task = Arc::new( - codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -157,12 +149,8 @@ async fn token_refresh_task_should_skip_duplicate_in_flight_refresh() { #[tokio::test] async fn token_refresh_task_should_schedule_future_per_account_timer_without_refreshing_early() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-future-timer.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-future-timer").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 16, 0, 0).unwrap(); let old_access_token = test_jwt((now + Duration::minutes(10)).timestamp()); store @@ -182,7 +170,7 @@ async fn token_refresh_task_should_schedule_future_per_account_timer_without_ref .await .expect("account should be inserted"); let refresher = CountingTokenRefresher::default(); - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -216,12 +204,8 @@ async fn token_refresh_task_should_schedule_future_per_account_timer_without_ref #[tokio::test] async fn token_refresh_task_should_cancel_future_timer_without_refreshing_expired_account() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-unauthorized-now.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-unauthorized-now").await; + let store = PgAccountStore::new(pool); let now = Utc::now(); let old_access_token = test_jwt((now + Duration::minutes(10)).timestamp()); let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); @@ -251,7 +235,7 @@ async fn token_refresh_task_should_cancel_future_timer_without_refreshing_expire refresh_token: None, }), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -290,12 +274,8 @@ async fn token_refresh_task_should_cancel_future_timer_without_refreshing_expire #[tokio::test] async fn token_refresh_task_should_clear_stale_refresh_time_without_refresh_token() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-no-rt-clear-next.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-no-rt-clear-next").await; + let store = PgAccountStore::new(pool); let now = Utc::now(); let old_access_token = test_jwt((now + Duration::minutes(10)).timestamp()); store @@ -322,7 +302,7 @@ async fn token_refresh_task_should_clear_stale_refresh_time_without_refresh_toke .await .expect("stale future refresh time should persist"); let refresher = CountingTokenRefresher::default(); - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -372,12 +352,8 @@ async fn token_refresh_task_should_clear_stale_refresh_time_without_refresh_toke #[tokio::test] async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-per-account-timer.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-per-account-timer").await; + let store = PgAccountStore::new(pool); let now = Utc::now() + Duration::minutes(5); let old_access_token = test_jwt((now + Duration::seconds(1)).timestamp()); let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); @@ -405,7 +381,7 @@ async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() { refresh_token: None, }), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 1, @@ -443,14 +419,8 @@ async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() { #[tokio::test] async fn token_refresh_task_should_fire_quota_exhausted_timer_without_clearing_status() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir - .path() - .join("token-refresh-quota-exhausted-timer.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-quota-exhausted-timer").await; + let store = PgAccountStore::new(pool); let now = Utc::now(); let old_expires_at = now + Duration::milliseconds(50); let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); @@ -480,7 +450,7 @@ async fn token_refresh_task_should_fire_quota_exhausted_timer_without_clearing_s refresh_token: None, }), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 0, @@ -529,12 +499,8 @@ async fn token_refresh_task_should_fire_quota_exhausted_timer_without_clearing_s #[tokio::test] async fn token_refresh_task_should_reschedule_next_timer_after_scheduled_refresh() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-next-timer.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-next-timer").await; + let store = PgAccountStore::new(pool); let now = Utc::now(); let old_expires_at = now + Duration::milliseconds(50); let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); @@ -562,7 +528,7 @@ async fn token_refresh_task_should_reschedule_next_timer_after_scheduled_refresh refresh_token: None, }), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 0, diff --git a/backend/tests/upstream/accounts/token_refresh/success.rs b/backend/tests/accounts/token_refresh/success.rs similarity index 83% rename from backend/tests/upstream/accounts/token_refresh/success.rs rename to backend/tests/accounts/token_refresh/success.rs index dfee3fc87..80052022b 100644 --- a/backend/tests/upstream/accounts/token_refresh/success.rs +++ b/backend/tests/accounts/token_refresh/success.rs @@ -2,12 +2,8 @@ use super::*; #[tokio::test] async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refresh_token() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 18, 8, 0, 0).unwrap(); let new_expires_at = now + Duration::hours(1); let new_access_token = test_jwt(new_expires_at.timestamp()); @@ -35,7 +31,7 @@ async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refre refresh_token: None, }))), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -77,12 +73,8 @@ async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refre #[tokio::test] async fn token_refresh_task_should_persist_rotated_refresh_token_when_returned() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-rotated.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-rotated").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 18, 9, 0, 0).unwrap(); let new_expires_at = now + Duration::hours(1); let new_access_token = test_jwt(new_expires_at.timestamp()); @@ -110,7 +102,7 @@ async fn token_refresh_task_should_persist_rotated_refresh_token_when_returned() refresh_token: Some("refresh-new".to_string()), }))), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -142,12 +134,8 @@ async fn token_refresh_task_should_persist_rotated_refresh_token_when_returned() #[tokio::test] async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clearing_status() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-quota-exhausted.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-quota-exhausted").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 11, 0, 0).unwrap(); let new_expires_at = now + Duration::hours(1); let new_access_token = test_jwt(new_expires_at.timestamp()); @@ -181,7 +169,7 @@ async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clear refresh_token: None, }), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -214,12 +202,8 @@ async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clear #[tokio::test] async fn token_refresh_task_should_not_clear_quota_limit_when_exhausted_token_refreshes() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-exhausted-quota.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool.clone()); + let (pool, _guard) = init_test_db("token-refresh-exhausted-quota").await; + let store = PgAccountStore::new(pool.clone()); let now = Utc.with_ymd_and_hms(2026, 6, 19, 11, 30, 0).unwrap(); let cooldown_until = now + Duration::hours(1); let new_expires_at = now + Duration::hours(2); @@ -245,9 +229,9 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_exhausted_token_re .await .expect("account should be inserted"); sqlx::query( - "update accounts set quota_limit_reached = 1, quota_cooldown_until = ? where id = ?", + "update accounts set quota_limit_reached = true, quota_cooldown_until = $1 where id = $2", ) - .bind(cooldown_until.to_rfc3339()) + .bind(cooldown_until) .bind("acct-exhausted-quota") .execute(&pool) .await @@ -258,7 +242,7 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_exhausted_token_re refresh_token: None, }))), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, @@ -276,8 +260,8 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_exhausted_token_re .await .expect("account should load") .expect("account should exist"); - let quota_state: (i64, Option) = sqlx::query_as( - "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let quota_state: (bool, Option>) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct-exhausted-quota") .fetch_one(&pool) @@ -290,18 +274,14 @@ async fn token_refresh_task_should_not_clear_quota_limit_when_exhausted_token_re stored.access_token.expose_secret(), new_access_token.as_str() ); - assert_eq!(quota_state.0, 1); - assert_eq!(quota_state.1, Some(cooldown_until.to_rfc3339())); + assert!(quota_state.0); + assert_eq!(quota_state.1, Some(cooldown_until)); } #[tokio::test] async fn token_refresh_task_should_persist_success_after_transient_transport_retry() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("token-refresh-transient-success.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAccountStore::new(pool); + let (pool, _guard) = init_test_db("token-refresh-transient-success").await; + let store = PgAccountStore::new(pool); let now = Utc.with_ymd_and_hms(2026, 6, 19, 12, 0, 0).unwrap(); let new_access_token = test_jwt((now + Duration::hours(1)).timestamp()); store @@ -335,7 +315,7 @@ async fn token_refresh_task_should_persist_success_after_transient_transport_ret }), ]))), }; - let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new( + let task = codex_proxy_rs::accounts::refresh::RuntimeTokenRefreshService::new( store.clone(), RefreshPolicy { refresh_margin_seconds: 300, diff --git a/backend/tests/admin/auth/session_store.rs b/backend/tests/admin/auth/session_store.rs deleted file mode 100644 index 73d870e15..000000000 --- a/backend/tests/admin/auth/session_store.rs +++ /dev/null @@ -1,84 +0,0 @@ -use chrono::{Duration, Utc}; -use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - infra::{database::connect_sqlite, identity::hash_admin_password}, -}; - -#[tokio::test] -async fn admin_session_store_should_create_and_load_default_admin_once() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("admin-auth.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAdminSessionStore::new(pool); - let password_hash = hash_admin_password("admin-password").expect("password hash"); - - let created = store - .ensure_default_admin(&password_hash) - .await - .expect("default admin should be created"); - let skipped = store - .ensure_default_admin("different-hash") - .await - .expect("second default admin should be skipped"); - let admin = store - .load_first_admin() - .await - .expect("default admin should load") - .expect("default admin should exist"); - - assert!(created); - assert!(!skipped); - assert_eq!(admin.password_hash, password_hash); -} - -#[tokio::test] -async fn admin_session_store_should_create_validate_and_cleanup_sessions() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("admin-sessions.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAdminSessionStore::new(pool); - let password_hash = hash_admin_password("admin-password").expect("password hash"); - store - .ensure_default_admin(&password_hash) - .await - .expect("default admin should be created"); - let admin = store - .load_first_admin() - .await - .expect("admin should load") - .expect("admin should exist"); - let now = Utc::now(); - - store - .create_session("sess_future", &admin.id, now + Duration::minutes(5)) - .await - .expect("future session should be inserted"); - store - .create_session("sess_expired", &admin.id, now - Duration::minutes(5)) - .await - .expect("expired session should be inserted"); - - assert!(store - .validate_session("sess_future") - .await - .expect("future session should validate")); - assert!(!store - .validate_session("sess_expired") - .await - .expect("expired session should not validate")); - - let deleted = store - .cleanup_expired_sessions(now) - .await - .expect("expired session should be deleted"); - - assert_eq!(deleted, 1); - assert!(store - .validate_session("sess_future") - .await - .expect("future session should remain")); -} diff --git a/backend/tests/admin/keys/mod.rs b/backend/tests/admin/keys/mod.rs deleted file mode 100644 index bdca9c5dc..000000000 --- a/backend/tests/admin/keys/mod.rs +++ /dev/null @@ -1,80 +0,0 @@ -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::usage_record_store::SqliteUsageRecordStore, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::token_refresh::RefreshLeaseStore, - upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, - upstream::fingerprint::FingerprintRepository, -}; -use serde_json::json; -use tower::util::ServiceExt; - -use crate::support::{ - admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint, - http::response_json, -}; - -mod authorization; -mod lifecycle; -mod store; - -async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = std::sync::Arc::new(Services::new( - &config, - stores, - runtime_fingerprint(fingerprint), - )); - let state = AppState { - services: (*services).clone(), - }; - ( - codex_proxy_rs::http::router::router().with_state(state), - dir, - ) -} - -async fn create_admin_client_key(app: &axum::Router, name: &str) -> (String, String) { - let response = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/keys") - .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::from(json!({ "name": name }).to_string())) - .unwrap(), - ) - .await - .unwrap(); - let body = response_json(response).await; - ( - body["data"]["id"].as_str().unwrap().to_string(), - body["data"]["key"].as_str().unwrap().to_string(), - ) -} diff --git a/backend/tests/admin/keys/store.rs b/backend/tests/admin/keys/store.rs deleted file mode 100644 index e267306a0..000000000 --- a/backend/tests/admin/keys/store.rs +++ /dev/null @@ -1,115 +0,0 @@ -use codex_proxy_rs::{ - admin::keys::service::{ClientKeyStore, RuntimeClientKeyStore, SqliteClientKeyStore}, - infra::database::connect_sqlite, -}; - -#[tokio::test] -async fn client_key_store_should_create_list_disable_enable_and_delete_keys() { - let (store, _dir) = client_key_store("client-keys.sqlite").await; - - let created = store.create("cursor").await.unwrap(); - assert_eq!(created.name, "cursor"); - assert!(created.key.starts_with("sk_")); - - let key: (String,) = sqlx::query_as("select key from client_api_keys where id = ?") - .bind(&created.id) - .fetch_one(store.pool()) - .await - .unwrap(); - assert_eq!(key.0, created.key); - - let first_page = store.list(None, 10).await.unwrap(); - assert_eq!(first_page.items.len(), 1); - assert_eq!(first_page.items[0].name, "cursor"); - assert_eq!(first_page.items[0].key, created.key); - assert!(first_page.items[0].last_used_at.is_none()); - - assert!(store.set_enabled(&created.id, false).await.unwrap()); - assert!(store.set_enabled(&created.id, true).await.unwrap()); - assert!(store.delete(&created.id).await.unwrap()); -} - -#[tokio::test] -async fn runtime_client_key_store_should_verify_from_memory_and_defer_last_used_flush() { - let (store, _dir) = client_key_store("client-keys-runtime.sqlite").await; - let created = store.create("runtime").await.unwrap(); - let runtime = RuntimeClientKeyStore::new(store.clone()); - runtime.reload_from_store().await.unwrap(); - - assert!(runtime.verify_and_touch(&created.key).await.unwrap()); - - let before_flush = store.get(&created.id).await.unwrap().unwrap(); - assert!(before_flush.last_used_at.is_none()); - - runtime.flush_pending_last_used().await; - - let after_flush = store.get(&created.id).await.unwrap().unwrap(); - assert!(after_flush.last_used_at.is_some()); -} - -#[tokio::test] -async fn runtime_client_key_store_should_not_accept_disabled_keys_after_reload() { - let (store, _dir) = client_key_store("client-keys-runtime-disabled.sqlite").await; - let created = store.create("runtime-disabled").await.unwrap(); - assert!(store.set_enabled(&created.id, false).await.unwrap()); - let runtime = RuntimeClientKeyStore::new(store); - runtime.reload_from_store().await.unwrap(); - - assert!(!runtime.verify_and_touch(&created.key).await.unwrap()); -} - -#[tokio::test] -async fn client_key_store_should_page_keys_by_created_at_desc() { - let (store, _dir) = client_key_store("client-keys-page.sqlite").await; - - let key_a = store.create("a").await.unwrap(); - let key_b = store.create("b").await.unwrap(); - let key_c = store.create("c").await.unwrap(); - sqlx::query("update client_api_keys set created_at = ? where id = ?") - .bind("2026-06-11T00:00:00Z") - .bind(&key_a.id) - .execute(store.pool()) - .await - .unwrap(); - sqlx::query("update client_api_keys set created_at = ? where id = ?") - .bind("2026-06-11T00:01:00Z") - .bind(&key_b.id) - .execute(store.pool()) - .await - .unwrap(); - sqlx::query("update client_api_keys set created_at = ? where id = ?") - .bind("2026-06-11T00:02:00Z") - .bind(&key_c.id) - .execute(store.pool()) - .await - .unwrap(); - - let first_page = store.list(None, 2).await.unwrap(); - assert_eq!( - first_page - .items - .iter() - .map(|key| key.name.as_str()) - .collect::>(), - ["c", "b"] - ); - - let second_page = store.list(first_page.next_cursor, 2).await.unwrap(); - assert_eq!( - second_page - .items - .iter() - .map(|key| key.name.as_str()) - .collect::>(), - ["a"] - ); -} - -async fn client_key_store(db_name: &str) -> (SqliteClientKeyStore, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .unwrap(); - (SqliteClientKeyStore::new(pool), dir) -} diff --git a/backend/tests/admin/monitoring/usage_record_routes.rs b/backend/tests/admin/monitoring/usage_record_routes.rs deleted file mode 100644 index add23396d..000000000 --- a/backend/tests/admin/monitoring/usage_record_routes.rs +++ /dev/null @@ -1,935 +0,0 @@ -use axum::{ - body::Body, - http::{Request, StatusCode}, -}; -use chrono::{Duration, SecondsFormat, Utc}; -use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::{ - usage_record_model::{UsageRecord, UsageRecordLevel}, - usage_record_store::SqliteUsageRecordStore, - }, - config::schema::AppConfig, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::token_refresh::RefreshLeaseStore, - upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, - upstream::fingerprint::FingerprintRepository, -}; -use serde_json::json; -use sqlx::SqlitePool; -use tower::util::ServiceExt; - -use crate::support::{ - admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint, - http::response_json, -}; - -#[tokio::test] -async fn admin_usage_records_should_require_admin_session_cookie() { - let (app, _store, _dir) = admin_usage_records_test_app("admin-usage-records-auth.sqlite").await; - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records") - .header("x-request-id", "req_usage_records_auth") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::UNAUTHORIZED); - assert_eq!(response_json(response).await["code"], 40101); -} - -#[tokio::test] -async fn admin_usage_records_should_cursor_page_events_and_include_request_id() { - let (app, store, _dir) = - admin_usage_records_test_app("admin-usage-records-cursor.sqlite").await; - let now = Utc::now(); - let mut older = UsageRecord::new("request", UsageRecordLevel::Info, "older"); - older.id = "usage_older".to_string(); - older.account_id = Some("acct_usage".to_string()); - older.model = Some("gpt-5.5".to_string()); - older.status_code = Some(200); - older.created_at = now; - store.append(&older).await.unwrap(); - let mut newer = UsageRecord::new("request", UsageRecordLevel::Info, "newer"); - newer.id = "usage_newer".to_string(); - newer.account_id = Some("acct_usage".to_string()); - newer.model = Some("gpt-5.5".to_string()); - newer.status_code = Some(200); - newer.created_at = now + Duration::seconds(1); - store.append(&newer).await.unwrap(); - - let response = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records?limit=1") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_records_cursor") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let body = response_json(response).await; - assert_eq!(body["code"], 200); - assert_eq!(body["data"]["items"].as_array().unwrap().len(), 1); -} - -#[tokio::test] -async fn admin_usage_records_should_return_numbered_page_metadata() { - let (app, store, _dir) = - admin_usage_records_test_app("admin-usage-records-numbered.sqlite").await; - let now = Utc::now(); - for (id, message, offset) in [ - ("usage_old", "older timeout", 0), - ("usage_new", "newer timeout", 1), - ] { - let mut event = UsageRecord::new("request", UsageRecordLevel::Info, message); - event.id = id.to_string(); - event.account_id = Some("acct_usage".to_string()); - event.model = Some("gpt-5.5".to_string()); - event.status_code = Some(200); - event.created_at = now + Duration::seconds(offset); - store.append(&event).await.unwrap(); - } - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records?page=1&pageSize=1&search=timeout") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_records_numbered") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::OK); - let body = response_json(response).await; - assert_eq!(body["data"]["items"].as_array().unwrap().len(), 1); - assert_eq!(body["data"]["items"][0]["id"], "usage_new"); - assert_eq!(body["data"]["page"]["page"], 1); - assert_eq!(body["data"]["page"]["pageSize"], 1); - assert_eq!(body["data"]["page"]["total"], 2); - assert_eq!(body["data"]["page"]["totalPages"], 2); -} - -#[tokio::test] -async fn admin_usage_records_should_return_table_display_fields() { - let (app, store, pool, _dir) = - admin_usage_records_test_app_with_pool("admin-usage-records-display-fields.sqlite").await; - let now = Utc::now().to_rfc3339(); - sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)") - .bind("acct_display") - .bind("display@example.com") - .bind("access-token") - .bind("disabled") - .bind(&now) - .bind(&now) - .execute(&pool) - .await - .unwrap(); - - let mut event = UsageRecord::new("request", UsageRecordLevel::Info, "display fields"); - event.id = "usage_display_fields".to_string(); - event.account_id = Some("acct_display".to_string()); - event.model = Some("gpt-5.5".to_string()); - event.latency_ms = Some(12_345); - event.metadata = json!({ - "requestedModel": "client-alias", - "upstreamModel": "gpt-5.5", - "clientIp": "203.0.113.8", - "userAgent": "codex-tui/0.142.2 (Ubuntu 24.4.0; aarch64) xterm-256color", - "reasoningEffort": "xhigh", - "firstTokenMs": 342, - "usage": { - "inputTokens": 1000, - "outputTokens": 200, - "cachedTokens": 100, - "reasoningTokens": 12, - "totalTokens": 1200 - } - }); - store.append(&event).await.unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records?page=1&pageSize=20") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_records_display_fields") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let body = response_json(response).await; - - assert_eq!(body["code"], 200); - let row = &body["data"]["items"][0]; - assert_eq!(row["accountEmail"], "display@example.com"); - assert_eq!(row["requestedModel"], "client-alias"); - assert_eq!(row["upstreamModel"], "gpt-5.5"); - assert_eq!(row["clientIp"], "203.0.113.8"); - assert_eq!( - row["userAgent"], - "codex-tui/0.142.2 (Ubuntu 24.4.0; aarch64) xterm-256color" - ); - assert_eq!(row["reasoningEffort"], "xhigh"); - assert_eq!(row["tokenDetails"]["inputTokens"], 1000); - assert_eq!(row["tokenDetails"]["outputTokens"], 200); - assert_eq!(row["tokenDetails"]["cachedTokens"], 100); - assert_eq!(row["tokenDetails"]["reasoningTokens"], 12); - assert_eq!(row["tokenDetails"]["totalTokens"], 1200); - assert_eq!(row["tokenDetails"]["inputTokensDisplay"], "1,000"); - assert_eq!(row["costDetails"]["inputCostDisplay"], "$0.00"); - assert_eq!(row["costDetails"]["outputCostDisplay"], "$0.01"); - assert_eq!(row["costDetails"]["cacheReadCostDisplay"], "$0.00"); - assert_eq!(row["costDetails"]["totalCostDisplay"], "$0.01"); - assert_eq!( - row["costDetails"]["inputPriceDisplay"], - "$5.0000 / 1M Token" - ); - assert_eq!(row["firstTokenLatencyMs"], 342); - assert_eq!(row["firstTokenLatencyMsDisplay"], "342 ms"); - assert_eq!(row["latencyMsDisplay"], "12.3 s"); -} - -#[tokio::test] -async fn admin_usage_records_should_use_record_account_email_snapshot_when_account_is_missing() { - let (app, store, _dir) = - admin_usage_records_test_app("admin-usage-records-email-snapshot.sqlite").await; - let mut event = UsageRecord::new("v1.response", UsageRecordLevel::Info, "completed"); - event.id = "usage_email_snapshot".to_string(); - event.account_id = Some("acct_deleted".to_string()); - event.model = Some("gpt-5.5".to_string()); - event.status_code = Some(200); - event.metadata = json!({ - "accountEmail": "deleted@example.com" - }); - store.append(&event).await.unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records?page=1&pageSize=20") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let body = response_json(response).await; - - assert_eq!( - body["data"]["items"][0]["accountEmail"], - "deleted@example.com" - ); -} - -#[tokio::test] -async fn admin_usage_records_summary_should_aggregate_filtered_usage() { - let (app, store, _dir) = - admin_usage_records_test_app("admin-usage-records-summary.sqlite").await; - let mut success = UsageRecord::new("request", UsageRecordLevel::Info, "summary success"); - success.id = "usage_summary_success".to_string(); - success.account_id = Some("acct_summary".to_string()); - success.model = Some("gpt-5.5".to_string()); - success.status_code = Some(200); - success.latency_ms = Some(1200); - success.metadata = json!({ - "usage": { - "inputTokens": 100, - "outputTokens": 40, - "cachedTokens": 20 - } - }); - store.append(&success).await.unwrap(); - - let mut failure = UsageRecord::new("request", UsageRecordLevel::Error, "summary failure"); - failure.id = "usage_summary_failure".to_string(); - failure.account_id = Some("acct_summary".to_string()); - failure.model = Some("gpt-5.5-mini".to_string()); - failure.status_code = Some(429); - failure.latency_ms = Some(600); - failure.metadata = json!({ - "usage": { - "inputTokens": 7, - "outputTokens": 3 - } - }); - store.append(&failure).await.unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/summary?search=summary") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_records_summary") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let body = response_json(response).await; - - assert_eq!(body["code"], 200); - assert_eq!(body["data"]["totalRequests"], "2"); - assert_eq!(body["data"]["errorRequests"], "1"); - assert_eq!(body["data"]["inputTokens"], "107"); - assert_eq!(body["data"]["outputTokens"], "43"); - assert_eq!(body["data"]["cachedTokens"], "20"); - assert_eq!(body["data"]["totalTokens"], "150"); - assert_eq!(body["data"]["averageLatencyMs"], "900 ms"); -} - -#[tokio::test] -async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimensions() { - let (app, store, _pool, _dir) = - admin_usage_records_test_app_with_pool("admin-usage-records-insights.sqlite").await; - let mut first = UsageRecord::new("request", UsageRecordLevel::Info, "insight success"); - first.id = "usage_insight_success".to_string(); - first.account_id = Some("acct_insight".to_string()); - first.route = Some("/v1/responses".to_string()); - first.model = Some("gpt-5.5".to_string()); - first.status_code = Some(200); - first.latency_ms = Some(240); - first.metadata = json!({ - "stream": true, - "requestedModel": "client-alias", - "upstreamModel": "gpt-5.5", - "upstreamEndpoint": "/backend/responses", - "endpointPath": "/v1/responses", - "usage": { - "inputTokens": 30, - "outputTokens": 12, - "cachedTokens": 6 - } - }); - store.append(&first).await.unwrap(); - - let mut second = UsageRecord::new("request", UsageRecordLevel::Error, "insight failure"); - second.id = "usage_insight_failure".to_string(); - second.account_id = Some("acct_insight".to_string()); - second.route = Some("/v1/responses/review".to_string()); - second.model = Some("gpt-5.5-mini".to_string()); - second.status_code = Some(500); - second.latency_ms = Some(120); - second.metadata = json!({ - "stream": false, - "requestedModel": "gpt-5.5-mini", - "upstreamModel": "gpt-5.5-mini", - "upstreamEndpoint": "/backend/responses/review", - "endpointPath": "/v1/responses/review", - "usage": { - "inputTokens": 5, - "outputTokens": 1 - } - }); - store.append(&second).await.unwrap(); - - let requested_models = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/models?source=requested&search=insight") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let requested_body = response_json(requested_models).await; - assert_eq!(requested_body["code"], 200); - assert_eq!(requested_body["data"][0]["name"], "client-alias"); - assert_eq!(requested_body["data"][0]["totalTokens"], "42"); - assert_eq!(requested_body["data"][0]["totalTokensValue"], 42); - assert_eq!(requested_body["data"][0]["totalTokensTotal"], "48"); - assert_eq!(requested_body["data"][0]["totalTokensTotalValue"], 48); - assert_eq!(requested_body["data"][0]["actualCost"], "$0.00"); - - let upstream_models = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/models?source=upstream&search=insight") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!( - response_json(upstream_models).await["data"][0]["name"], - "gpt-5.5" - ); - - let mappings = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/models?source=mapping&search=insight") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let mapping_body = response_json(mappings).await; - assert_eq!(mapping_body["data"][0]["name"], "client-alias -> gpt-5.5"); - assert_eq!( - mapping_body["data"][1]["name"], - "gpt-5.5-mini -> gpt-5.5-mini" - ); - - let endpoints = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/endpoints?source=path&search=insight") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let endpoint_body = response_json(endpoints).await; - assert_eq!(endpoint_body["data"].as_array().unwrap().len(), 2); - let endpoint_paths = endpoint_body["data"] - .as_array() - .unwrap() - .iter() - .map(|item| item["name"].as_str().unwrap()) - .collect::>(); - assert!(endpoint_paths.contains(&"/v1/responses -> /backend/responses")); - assert!(endpoint_paths.contains(&"/v1/responses/review -> /backend/responses/review")); - - let token_trend = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/token-trend?search=insight") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let token_body = response_json(token_trend).await; - assert_eq!(token_body["data"][0]["inputTokens"], "35"); - assert_eq!(token_body["data"][0]["inputTokensValue"], 35); - assert_eq!(token_body["data"][0]["outputTokens"], "13"); - assert_eq!(token_body["data"][0]["cacheCreationTokens"], "0"); - assert_eq!(token_body["data"][0]["cachedTokens"], "6"); - assert_eq!(token_body["data"][0]["actualCost"], "$0.00"); - assert_eq!(token_body["data"][0]["cost"], "$0.00"); - - let latency_trend = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/latency-trend?search=insight") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!( - response_json(latency_trend).await["data"][0]["averageLatencyMs"], - "180 ms" - ); -} - -#[tokio::test] -async fn admin_usage_record_trends_should_bucket_by_china_calendar_day() { - let (app, store, _pool, _dir) = - admin_usage_records_test_app_with_pool("admin-usage-records-china-day-trend.sqlite").await; - let mut june_29 = UsageRecord::new("request", UsageRecordLevel::Info, "china-day trend"); - june_29.id = "usage_china_day_29".to_string(); - june_29.account_id = Some("acct_trend".to_string()); - june_29.model = Some("gpt-5.5".to_string()); - june_29.status_code = Some(200); - june_29.created_at = chrono::DateTime::parse_from_rfc3339("2026-06-29T15:30:00Z") - .unwrap() - .with_timezone(&Utc); - june_29.latency_ms = Some(100); - june_29.metadata = json!({ - "usage": { - "inputTokens": 10, - "outputTokens": 1 - } - }); - store.append(&june_29).await.unwrap(); - - let mut june_30 = UsageRecord::new("request", UsageRecordLevel::Info, "china-day trend"); - june_30.id = "usage_china_day_30".to_string(); - june_30.account_id = Some("acct_trend".to_string()); - june_30.model = Some("gpt-5.5".to_string()); - june_30.status_code = Some(200); - june_30.created_at = chrono::DateTime::parse_from_rfc3339("2026-06-29T16:30:00Z") - .unwrap() - .with_timezone(&Utc); - june_30.latency_ms = Some(300); - june_30.metadata = json!({ - "usage": { - "inputTokens": 20, - "outputTokens": 2 - } - }); - store.append(&june_30).await.unwrap(); - - let token_trend = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/token-trend?search=china-day") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let token_body = response_json(token_trend).await; - assert_eq!(token_body["data"][0]["date"], "2026-06-29"); - assert_eq!(token_body["data"][1]["date"], "2026-06-30"); - assert_eq!(token_body["data"][1]["inputTokens"], "20"); - assert_eq!(token_body["data"][1]["inputTokensValue"], 20); - - let latency_trend = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/latency-trend?search=china-day") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let latency_body = response_json(latency_trend).await; - assert_eq!(latency_body["data"][1]["date"], "2026-06-30"); - assert_eq!(latency_body["data"][1]["averageLatencyMs"], "300 ms"); -} - -#[tokio::test] -async fn admin_usage_records_insight_card_endpoints_should_return_source_specific_data() { - let (app, store, _pool, _dir) = - admin_usage_records_test_app_with_pool("admin-usage-records-insight-cards.sqlite").await; - let mut first = UsageRecord::new("request", UsageRecordLevel::Info, "card success"); - first.id = "usage_card_success".to_string(); - first.account_id = Some("acct_card".to_string()); - first.route = Some("/v1/responses".to_string()); - first.model = Some("gpt-5.5".to_string()); - first.status_code = Some(200); - first.latency_ms = Some(240); - first.metadata = json!({ - "stream": true, - "requestedModel": "client-alias", - "upstreamModel": "gpt-5.5", - "upstreamEndpoint": "/backend/responses", - "usage": { - "inputTokens": 30, - "outputTokens": 12, - "cachedTokens": 6 - } - }); - store.append(&first).await.unwrap(); - - let mut second = UsageRecord::new("request", UsageRecordLevel::Info, "card mini"); - second.id = "usage_card_mini".to_string(); - second.account_id = Some("acct_card".to_string()); - second.route = Some("/v1/responses/compact".to_string()); - second.model = Some("gpt-5.5-mini".to_string()); - second.status_code = Some(200); - second.latency_ms = Some(120); - second.metadata = json!({ - "stream": false, - "requestedModel": "gpt-5.5-mini", - "upstreamModel": "gpt-5.5-mini", - "upstreamEndpoint": "/backend/responses/compact", - "usage": { - "inputTokens": 5, - "outputTokens": 1 - } - }); - store.append(&second).await.unwrap(); - - let upstream_models = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/models?source=upstream&search=card") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let upstream_body = response_json(upstream_models).await; - assert_eq!(upstream_body["code"], 200); - assert_eq!(upstream_body["data"][0]["name"], "gpt-5.5"); - assert_eq!(upstream_body["data"][0]["totalTokens"], "42"); - assert_eq!(upstream_body["data"][0]["totalTokensValue"], 42); - - let mappings = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/models?source=mapping&search=card") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!( - response_json(mappings).await["data"][0]["name"], - "client-alias -> gpt-5.5" - ); - - let endpoints = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/endpoints?source=path&search=card") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let endpoint_body = response_json(endpoints).await; - let endpoint_paths = endpoint_body["data"] - .as_array() - .unwrap() - .iter() - .map(|item| item["name"].as_str().unwrap()) - .collect::>(); - assert!(endpoint_paths.contains(&"/v1/responses -> /backend/responses")); - assert!(endpoint_paths.contains(&"/v1/responses/compact -> /backend/responses/compact")); - - let token_trend = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/token-trend?search=card") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let token_body = response_json(token_trend).await; - assert_eq!(token_body["data"][0]["inputTokens"], "35"); - assert_eq!(token_body["data"][0]["cachedTokens"], "6"); - - let latency_trend = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/latency-trend?search=card") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!( - response_json(latency_trend).await["data"][0]["averageLatencyMs"], - "180 ms" - ); - - let invalid_source = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/insights/models?source=actual") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(invalid_source.status(), StatusCode::BAD_REQUEST); - assert_eq!(response_json(invalid_source).await["code"], 40003); -} - -#[tokio::test] -async fn admin_usage_records_should_filter_summary_and_table_by_time_range() { - let (app, store, _dir) = - admin_usage_records_test_app("admin-usage-records-time-range.sqlite").await; - let now = Utc::now(); - let mut older = UsageRecord::new("request", UsageRecordLevel::Info, "outside range"); - older.id = "usage_outside_range".to_string(); - older.account_id = Some("acct_range".to_string()); - older.model = Some("gpt-5.5".to_string()); - older.status_code = Some(200); - older.created_at = now - Duration::days(3); - store.append(&older).await.unwrap(); - - let mut current = UsageRecord::new("request", UsageRecordLevel::Info, "inside range"); - current.id = "usage_inside_range".to_string(); - current.account_id = Some("acct_range".to_string()); - current.model = Some("gpt-5.5".to_string()); - current.status_code = Some(200); - current.created_at = now; - store.append(¤t).await.unwrap(); - - let start_time = (now - Duration::hours(1)).to_rfc3339_opts(SecondsFormat::Secs, true); - let end_time = (now + Duration::hours(1)).to_rfc3339_opts(SecondsFormat::Secs, true); - let list_uri = format!( - "/api/admin/usage/records?page=1&pageSize=20&startTime={start_time}&endTime={end_time}" - ); - let list = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri(list_uri) - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let list_body = response_json(list).await; - assert_eq!(list_body["data"]["page"]["total"], 1); - assert_eq!(list_body["data"]["items"][0]["id"], "usage_inside_range"); - - let summary_uri = - format!("/api/admin/usage/records/summary?startTime={start_time}&endTime={end_time}"); - let summary = app - .oneshot( - Request::builder() - .method("GET") - .uri(summary_uri) - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response_json(summary).await["data"]["totalRequests"], "1"); -} - -#[tokio::test] -async fn admin_usage_records_should_filter_and_cursor_page_events() { - let (app, store, _dir) = admin_usage_records_test_app("admin-usage-records.sqlite").await; - let mut matching = UsageRecord::new("request", UsageRecordLevel::Info, "upstream timeout"); - matching.id = "usage_matching".to_string(); - matching.account_id = Some("acct_filter".to_string()); - matching.route = Some("/v1/responses".to_string()); - matching.model = Some("gpt-5.5".to_string()); - matching.status_code = Some(200); - store.append(&matching).await.unwrap(); - let mut non_matching_message = UsageRecord::new("request", UsageRecordLevel::Info, "completed"); - non_matching_message.account_id = Some("acct_filter".to_string()); - non_matching_message.model = Some("gpt-5.5".to_string()); - non_matching_message.status_code = Some(200); - store.append(&non_matching_message).await.unwrap(); - store - .append(&UsageRecord::new( - "account", - UsageRecordLevel::Error, - "upstream timeout", - )) - .await - .unwrap(); - - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records?kind=request&level=info&search=timeout&limit=1") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_records_filter") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - let body = response_json(response).await; - assert_eq!(body["data"]["items"][0]["id"], "usage_matching"); -} - -#[tokio::test] -async fn admin_usage_records_should_reject_unsupported_level_filter() { - let (app, _store, _dir) = - admin_usage_records_test_app("admin-usage-records-invalid-level.sqlite").await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records?level=verbose") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_records_invalid_level") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::BAD_REQUEST); -} - -#[tokio::test] -async fn admin_usage_records_should_return_detail_and_clear_events() { - let (app, store, _dir) = admin_usage_records_test_app("admin-usage-records-state.sqlite").await; - let mut event = UsageRecord::new("request", UsageRecordLevel::Warn, "detail"); - event.id = "usage_detail".to_string(); - event.request_id = Some("req_upstream".to_string()); - store.append(&event).await.unwrap(); - - let detail = app - .clone() - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/detail?id=usage_detail") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(detail.status(), StatusCode::OK); - - let cleared = app - .clone() - .oneshot( - Request::builder() - .method("POST") - .uri("/api/admin/usage/records/delete") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response_json(cleared).await["data"]["cleared"], 1); - - let empty = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records?limit=50") - .header("cookie", "cpr_admin_session=session_1") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!( - response_json(empty).await["data"]["items"] - .as_array() - .unwrap() - .len(), - 0 - ); -} - -#[tokio::test] -async fn admin_usage_records_detail_should_return_not_found_for_missing_event() { - let (app, _store, _dir) = - admin_usage_records_test_app("admin-usage-records-detail-missing.sqlite").await; - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/api/admin/usage/records/detail?id=missing") - .header("cookie", "cpr_admin_session=session_1") - .header("x-request-id", "req_usage_records_missing") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::NOT_FOUND); -} - -async fn admin_usage_records_test_app( - db_name: &str, -) -> (axum::Router, SqliteUsageRecordStore, tempfile::TempDir) { - let (app, store, _pool, dir) = admin_usage_records_test_app_with_pool(db_name).await; - (app, store, dir) -} - -async fn admin_usage_records_test_app_with_pool( - db_name: &str, -) -> ( - axum::Router, - SqliteUsageRecordStore, - SqlitePool, - tempfile::TempDir, -) { - admin_usage_records_test_app_with_config(db_name, |_| {}).await -} - -async fn admin_usage_records_test_app_with_config( - db_name: &str, - configure: impl FnOnce(&mut AppConfig), -) -> ( - axum::Router, - SqliteUsageRecordStore, - SqlitePool, - tempfile::TempDir, -) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let mut config = test_config(url); - configure(&mut config); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; - let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = std::sync::Arc::new(Services::new( - &config, - stores, - runtime_fingerprint(fingerprint), - )); - let state = AppState { - services: (*services).clone(), - }; - ( - codex_proxy_rs::http::router::router().with_state(state), - SqliteUsageRecordStore::new(pool.clone()), - pool, - dir, - ) -} diff --git a/backend/tests/admin/monitoring/usage_record_store.rs b/backend/tests/admin/monitoring/usage_record_store.rs deleted file mode 100644 index b8ae9213c..000000000 --- a/backend/tests/admin/monitoring/usage_record_store.rs +++ /dev/null @@ -1,157 +0,0 @@ -use chrono::{Duration, Utc}; -use codex_proxy_rs::{ - admin::monitoring::{ - usage_record_model::{UsageRecord, UsageRecordLevel}, - usage_record_store::{SqliteUsageRecordStore, UsageRecordFilter}, - }, - infra::database::connect_sqlite, -}; -use serde_json::json; - -#[tokio::test] -async fn usage_record_store_should_filter_and_cursor_page_events() { - let (store, _dir) = usage_record_store("usage-records.sqlite").await; - let mut matching = UsageRecord::new("request", UsageRecordLevel::Error, "upstream timeout"); - matching.id = "usage_matching".to_string(); - matching.request_id = Some("req_1".to_string()); - matching.route = Some("/v1/responses".to_string()); - store.append(&matching).await.unwrap(); - store - .append(&UsageRecord::new( - "request", - UsageRecordLevel::Info, - "upstream timeout", - )) - .await - .unwrap(); - store - .append(&UsageRecord::new( - "account", - UsageRecordLevel::Error, - "upstream timeout", - )) - .await - .unwrap(); - - let page = store - .list( - UsageRecordFilter { - kind: Some("request".to_string()), - level: Some(UsageRecordLevel::Error), - search: Some("timeout".to_string()), - ..UsageRecordFilter::default() - }, - None, - 1, - ) - .await - .unwrap(); - - assert_eq!(page.items.len(), 1); - assert_eq!(page.items[0].id, "usage_matching"); - assert_eq!(page.items[0].request_id.as_deref(), Some("req_1")); - assert_eq!(page.items[0].route.as_deref(), Some("/v1/responses")); -} - -#[tokio::test] -async fn usage_record_store_should_promote_diagnostic_metadata_fields() { - let (store, _dir) = usage_record_store("usage-records-diagnostics.sqlite").await; - let mut event = UsageRecord::new("request", UsageRecordLevel::Error, "upstream failed"); - event.id = "usage_diagnostic".to_string(); - event.metadata = json!({ - "transport": "websocket", - "attemptIndex": 2, - "upstreamStatus": 429, - "failureClass": "rate_limited", - "responseId": "resp_1", - "upstreamRequestId": "req_upstream_1" - }); - - store.append(&event).await.unwrap(); - let page = store - .list( - UsageRecordFilter { - transport: Some("websocket".to_string()), - failure_class: Some("rate_limited".to_string()), - upstream_status_code: Some(429), - ..UsageRecordFilter::default() - }, - None, - 10, - ) - .await - .unwrap(); - - assert_eq!(page.items.len(), 1); - assert_eq!(page.items[0].id, "usage_diagnostic"); - assert_eq!(page.items[0].attempt_index, Some(2)); - assert_eq!(page.items[0].response_id.as_deref(), Some("resp_1")); - assert_eq!( - page.items[0].upstream_request_id.as_deref(), - Some("req_upstream_1") - ); -} - -#[tokio::test] -async fn usage_record_store_should_not_promote_cf_ray_as_upstream_request_id() { - let (store, _dir) = usage_record_store("usage-records-cf-ray.sqlite").await; - let mut event = UsageRecord::new("request", UsageRecordLevel::Error, "upstream failed"); - event.id = "usage_cf_ray".to_string(); - event.metadata = json!({ - "upstreamStatus": 403, - "cfRay": "ray-1" - }); - - store.append(&event).await.unwrap(); - let loaded = store.get("usage_cf_ray").await.unwrap().unwrap(); - - assert_eq!(loaded.upstream_request_id, None); -} - -#[tokio::test] -async fn usage_record_store_should_get_and_clear_events() { - let (store, _dir) = usage_record_store("usage-records-clear.sqlite").await; - let mut event = UsageRecord::new("request", UsageRecordLevel::Warn, "detail"); - event.id = "usage_detail".to_string(); - store.append(&event).await.unwrap(); - - let loaded = store.get("usage_detail").await.unwrap().unwrap(); - assert_eq!(loaded.id, "usage_detail"); - assert_eq!(store.clear().await.unwrap(), 1); - let page = store - .list(UsageRecordFilter::default(), None, 10) - .await - .unwrap(); - assert!(page.items.is_empty()); -} - -#[tokio::test] -async fn usage_record_store_should_trim_records_older_than_thirty_days() { - let (store, _dir) = usage_record_store("usage-records-retention.sqlite").await; - let now = Utc::now(); - let mut expired = UsageRecord::new("v1.response", UsageRecordLevel::Info, "expired"); - expired.id = "usage_expired".to_string(); - expired.route = Some("/v1/responses".to_string()); - expired.created_at = now - Duration::days(31); - store.append(&expired).await.unwrap(); - - let mut retained = UsageRecord::new("v1.response", UsageRecordLevel::Info, "retained"); - retained.id = "usage_retained".to_string(); - retained.route = Some("/v1/responses".to_string()); - retained.created_at = now - Duration::days(29); - store.append(&retained).await.unwrap(); - - assert_eq!(store.trim_to_retention(now).await.unwrap(), 1); - - assert!(store.get("usage_expired").await.unwrap().is_none()); - assert!(store.get("usage_retained").await.unwrap().is_some()); -} - -async fn usage_record_store(db_name: &str) -> (SqliteUsageRecordStore, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .unwrap(); - (SqliteUsageRecordStore::new(pool), dir) -} diff --git a/backend/tests/admin/accounts/exporting.rs b/backend/tests/api/admin/accounts/exporting.rs similarity index 93% rename from backend/tests/admin/accounts/exporting.rs rename to backend/tests/api/admin/accounts/exporting.rs index b3c2b8f73..b19bec3a0 100644 --- a/backend/tests/admin/accounts/exporting.rs +++ b/backend/tests/api/admin/accounts/exporting.rs @@ -4,8 +4,7 @@ const EXPORT_CONFIRM: &str = "confirm=export_sensitive_accounts"; #[tokio::test] async fn admin_accounts_export_should_reject_empty_ids() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-export-all.sqlite", 131).await; + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-accounts-export-all", 131).await; seed_account( &pool, NewAccount { @@ -79,7 +78,7 @@ async fn admin_accounts_export_should_reject_empty_ids() { #[tokio::test] async fn admin_accounts_export_should_filter_by_ids() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-export-filter.sqlite", 132).await; + admin_accounts_test_app("admin-accounts-export-filter", 132).await; seed_account( &pool, NewAccount { @@ -154,7 +153,7 @@ async fn admin_accounts_export_should_filter_by_ids() { #[tokio::test] async fn admin_accounts_export_should_fail_when_selected_id_is_missing() { let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-accounts-export-missing.sqlite", 133).await; + admin_accounts_test_app("admin-accounts-export-missing", 133).await; let response = app .oneshot( @@ -176,7 +175,7 @@ async fn admin_accounts_export_should_fail_when_selected_id_is_missing() { #[tokio::test] async fn admin_accounts_export_payload_should_import_directly() { let (source_app, _source_state, source_pool, _source_dir) = - admin_accounts_test_app("admin-accounts-export-source.sqlite", 134).await; + admin_accounts_test_app("admin-accounts-export-source", 134).await; seed_account( &source_pool, NewAccount { @@ -221,7 +220,7 @@ async fn admin_accounts_export_payload_should_import_directly() { let payload = exported["data"].clone(); let (target_app, _target_state, target_pool, _target_dir) = - admin_accounts_test_app("admin-accounts-export-target.sqlite", 135).await; + admin_accounts_test_app("admin-accounts-export-target", 135).await; let response = target_app .oneshot( Request::builder() @@ -236,7 +235,7 @@ async fn admin_accounts_export_payload_should_import_directly() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(target_pool) + let stored = PgAccountStore::new(target_pool) .get("acct_export_roundtrip") .await .unwrap() @@ -253,7 +252,7 @@ async fn admin_accounts_export_payload_should_import_directly() { #[tokio::test] async fn admin_accounts_export_should_require_explicit_sensitive_export_confirmation() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-export-confirm.sqlite", 136).await; + admin_accounts_test_app("admin-accounts-export-confirm", 136).await; seed_account( &pool, NewAccount { @@ -309,7 +308,7 @@ async fn admin_accounts_export_should_require_explicit_sensitive_export_confirma #[tokio::test] async fn admin_accounts_export_should_not_write_usage_record_on_success() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-export-usage.sqlite", 137).await; + admin_accounts_test_app("admin-accounts-export-usage", 137).await; seed_account( &pool, NewAccount { @@ -358,8 +357,8 @@ async fn admin_accounts_export_should_not_write_usage_record_on_success() { ); } -async fn usage_record_count_by_request_id(pool: &SqlitePool, request_id: &str) -> i64 { - sqlx::query_scalar("select count(*) from usage_records where request_id = ?") +async fn usage_record_count_by_request_id(pool: &PgPool, request_id: &str) -> i64 { + sqlx::query_scalar("select count(*) from usage_records where request_id = $1") .bind(request_id) .fetch_one(pool) .await diff --git a/backend/tests/admin/accounts/importing.rs b/backend/tests/api/admin/accounts/importing.rs similarity index 88% rename from backend/tests/admin/accounts/importing.rs rename to backend/tests/api/admin/accounts/importing.rs index 3f105fd18..829f2fa77 100644 --- a/backend/tests/admin/accounts/importing.rs +++ b/backend/tests/api/admin/accounts/importing.rs @@ -7,28 +7,17 @@ use wiremock::{ #[tokio::test] async fn admin_accounts_import_should_store_cpr_account_tokens() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-accounts-import.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; + let (pool, _dir) = init_test_db("admin-accounts-import").await; + let redis = create_test_redis("admin-accounts-import").await; + seed_admin_session(&pool, &redis, "session_1").await; let access_token = test_jwt( "chatgpt_import", Some("user_import"), Some("import@example.com"), Some("team"), ); - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new( &config, @@ -38,7 +27,7 @@ async fn admin_accounts_import_should_store_cpr_account_tokens() { let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); let response = app .oneshot( @@ -66,7 +55,7 @@ async fn admin_accounts_import_should_store_cpr_account_tokens() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get("acct_import") .await .unwrap() @@ -107,7 +96,7 @@ async fn admin_accounts_import_should_store_cpr_access_token_from_at_field() { .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-accounts-import-cpr-at.sqlite", + "admin-accounts-import-cpr-at", 129, format!("{}/backend-api", server.uri()), ) @@ -142,7 +131,7 @@ async fn admin_accounts_import_should_store_cpr_access_token_from_at_field() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get("acct_import_cpr_at") .await .unwrap() @@ -183,7 +172,7 @@ async fn admin_accounts_import_should_read_sub2api_backup_access_token_only_acco .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-accounts-import-sub2api-backup-at.sqlite", + "admin-accounts-import-sub2api-backup-at", 127, format!("{}/backend-api", server.uri()), ) @@ -231,7 +220,7 @@ async fn admin_accounts_import_should_read_sub2api_backup_access_token_only_acco .unwrap(); let status = response.status(); let body = response_json(response).await; - let store = SqliteAccountStore::new(pool); + let store = PgAccountStore::new(pool); let stored_metadata = store .list_metadata_page(1, 10, None) .await @@ -277,7 +266,7 @@ async fn admin_accounts_import_should_read_cliproxyapi_codex_auth_files() { .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-accounts-import-cliproxyapi.sqlite", + "admin-accounts-import-cliproxyapi", 128, format!("{}/backend-api", server.uri()), ) @@ -316,7 +305,7 @@ async fn admin_accounts_import_should_read_cliproxyapi_codex_auth_files() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .list_metadata_page(1, 10, None) .await .unwrap() @@ -379,7 +368,7 @@ async fn admin_accounts_import_should_fetch_wham_usage_for_current_openai_token_ .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-accounts-import-current-openai-token.sqlite", + "admin-accounts-import-current-openai-token", 118, format!("{}/backend-api", server.uri()), ) @@ -426,7 +415,7 @@ async fn admin_accounts_import_should_fetch_wham_usage_for_current_openai_token_ .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get(CURRENT_IMPORT_ACCOUNT_ID) .await .unwrap() @@ -481,7 +470,7 @@ async fn admin_accounts_import_should_complete_chatgpt_account_id_from_refresh_t .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url_and_oauth_token_endpoint( - "admin-accounts-import-rt-complete.sqlite", + "admin-accounts-import-rt-complete", 125, format!("{}/backend-api", server.uri()), format!("{}/oauth/token", server.uri()), @@ -508,7 +497,7 @@ async fn admin_accounts_import_should_complete_chatgpt_account_id_from_refresh_t .await .unwrap(); let body = response_json(response).await; - let store = SqliteAccountStore::new(pool); + let store = PgAccountStore::new(pool); let stored = store .list_metadata_page(1, 10, None) .await @@ -557,7 +546,7 @@ async fn admin_accounts_import_from_refresh_token_should_not_store_consumed_refr .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( - "admin-accounts-import-rt-no-rotation.sqlite", + "admin-accounts-import-rt-no-rotation", 126, format!("{}/oauth/token", server.uri()), ) @@ -582,7 +571,7 @@ async fn admin_accounts_import_from_refresh_token_should_not_store_consumed_refr .await .unwrap(); let body = response_json(response).await; - let store = SqliteAccountStore::new(pool); + let store = PgAccountStore::new(pool); let stored = store .list_metadata_page(1, 10, None) .await @@ -625,23 +614,12 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q .mount(&server) .await; - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-accounts-import-supplement.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let mut config = test_config(url); + let (pool, _dir) = init_test_db("admin-accounts-import-supplement").await; + let redis = create_test_redis("admin-accounts-import-supplement").await; + seed_admin_session(&pool, &redis, "session_1").await; + let mut config = test_config(test_database_url()); config.api.base_url = format!("{}/backend-api", server.uri()); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new( &config, @@ -651,7 +629,7 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); let access_token = test_jwt_with_exp( None, Some("supplement-user"), @@ -689,18 +667,18 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool.clone()) + let stored = PgAccountStore::new(pool.clone()) .get("acct_import_supplement") .await .unwrap() .unwrap(); - let quota_json: Option = - sqlx::query_scalar("select quota_json from accounts where id = ?") + let quota_json: Option = + sqlx::query_scalar("select quota_json from accounts where id = $1") .bind("acct_import_supplement") .fetch_one(&pool) .await .unwrap(); - let quota: Value = serde_json::from_str("a_json.unwrap()).unwrap(); + let quota = quota_json.unwrap(); assert_eq!(status, StatusCode::OK); assert_eq!(body["data"]["imported"], 1); @@ -749,7 +727,7 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q #[tokio::test] async fn admin_accounts_import_should_update_existing_sub2api_account_and_clear_stale_quota_lock() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-import-update-existing.sqlite", 116).await; + admin_accounts_test_app("admin-accounts-import-update-existing", 116).await; let new_access_token = test_jwt( "chatgpt-import-update", Some("user-import-update"), @@ -773,8 +751,8 @@ async fn admin_accounts_import_should_update_existing_sub2api_account_and_clear_ }, ) .await; - sqlx::query("update accounts set quota_limit_reached = 1, quota_cooldown_until = ?, quota_verify_required = 1 where id = ?") - .bind("2026-06-25T00:00:00Z").bind("acct_import_update").execute(&pool).await.unwrap(); + sqlx::query("update accounts set quota_limit_reached = true, quota_cooldown_until = $1, quota_verify_required = true where id = $2") + .bind(crate::support::storage::timestamp("2026-06-25T00:00:00Z")).bind("acct_import_update").execute(&pool).await.unwrap(); let response = app .clone() @@ -827,13 +805,18 @@ async fn admin_accounts_import_should_update_existing_sub2api_account_and_clear_ .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool.clone()) + let stored = PgAccountStore::new(pool.clone()) .get("acct_import_update") .await .unwrap() .unwrap(); - let quota_state: (i64, Option, i64, Option) = sqlx::query_as( - "select quota_limit_reached, quota_cooldown_until, quota_verify_required, quota_json from accounts where id = ?", + let quota_state: ( + bool, + Option>, + bool, + Option, + ) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until, quota_verify_required, quota_json from accounts where id = $1", ).bind("acct_import_update").fetch_one(&pool).await.unwrap(); assert_eq!(status, StatusCode::OK); @@ -841,27 +824,16 @@ async fn admin_accounts_import_should_update_existing_sub2api_account_and_clear_ assert_eq!(body["data"]["sourceFormat"], "sub2api"); assert_eq!(stored.email.as_deref(), Some("new@example.com")); assert_eq!(stored.status, AccountStatus::Active); - assert_eq!(quota_state.0, 0); + assert!(!quota_state.0); assert!(quota_state.1.is_none()); } #[tokio::test] async fn admin_accounts_import_should_require_admin_session_cookie() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-accounts-import-auth.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let (pool, _dir) = init_test_db("admin-accounts-import-auth").await; + let redis = create_test_redis("admin-accounts-import-auth").await; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new( &config, @@ -871,7 +843,7 @@ async fn admin_accounts_import_should_require_admin_session_cookie() { let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); let response = app .oneshot( @@ -892,7 +864,7 @@ async fn admin_accounts_import_should_require_admin_session_cookie() { #[tokio::test] async fn admin_accounts_import_should_store_plain_tokens_and_list_sanitized_accounts() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-import-sanitized.sqlite", 113).await; + admin_accounts_test_app("admin-accounts-import-sanitized", 113).await; let access_token = test_jwt( "chatgpt-account", Some("chatgpt-user"), @@ -919,7 +891,7 @@ async fn admin_accounts_import_should_store_plain_tokens_and_list_sanitized_acco assert_eq!(body["data"]["imported"], 1); let stored: (String, String) = - sqlx::query_as("select access_token, refresh_token from accounts where id = ?") + sqlx::query_as("select access_token, refresh_token from accounts where id = $1") .bind("acct_imported_sanitized") .fetch_one(&pool) .await diff --git a/backend/tests/admin/accounts/lifecycle.rs b/backend/tests/api/admin/accounts/lifecycle.rs similarity index 82% rename from backend/tests/admin/accounts/lifecycle.rs rename to backend/tests/api/admin/accounts/lifecycle.rs index 9a8e6d58a..fb1fe4489 100644 --- a/backend/tests/admin/accounts/lifecycle.rs +++ b/backend/tests/api/admin/accounts/lifecycle.rs @@ -2,26 +2,15 @@ use super::*; #[tokio::test] async fn admin_accounts_lifecycle_should_update_and_delete_accounts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-accounts-lifecycle.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)") + let (pool, _dir) = init_test_db("admin-accounts-lifecycle").await; + let redis = create_test_redis("admin-accounts-lifecycle").await; + seed_admin_session(&pool, &redis, "session_1").await; + sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6)") .bind("acct_lifecycle").bind("life@example.com").bind("cipher").bind("active") - .bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z").execute(&pool).await.unwrap(); + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")).bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")).execute(&pool).await.unwrap(); seed_account_related_rows(&pool, "acct_lifecycle").await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new( &config, @@ -31,7 +20,27 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() { let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + state + .services + .background_tasks + .session_affinity + .upsert( + "resp_lifecycle", + &codex_proxy_rs::dispatch::affinity::SessionAffinityEntry { + account_id: "acct_lifecycle".to_string(), + conversation_id: "conv_lifecycle".to_string(), + turn_state: None, + instructions_hash: None, + input_tokens: None, + function_call_ids: Vec::new(), + variant_hash: None, + created_at: chrono::Utc::now(), + }, + chrono::Duration::hours(1), + ) + .await + .unwrap(); + let app = codex_proxy_rs::api::router::router().with_state(state.clone()); let label = app .clone() @@ -106,99 +115,79 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() { assert_eq!(deleted.status(), StatusCode::OK); assert_eq!(response_json(deleted).await["data"]["deleted"], 1); assert_account_related_rows_deleted(&pool, "acct_lifecycle").await; + assert!(state + .services + .background_tasks + .session_affinity + .get( + "resp_lifecycle", + chrono::Utc::now(), + chrono::Duration::hours(1), + ) + .await + .unwrap() + .is_none()); } -async fn seed_account_related_rows(pool: &SqlitePool, account_id: &str) { - sqlx::query("insert into account_usage (account_id, request_count) values (?, 1)") +async fn seed_account_related_rows(pool: &PgPool, account_id: &str) { + sqlx::query("insert into account_usage (account_id, request_count) values ($1, 1)") .bind(account_id) .execute(pool) .await .unwrap(); sqlx::query( - "insert into account_model_usage (account_id, model, request_count) values (?, ?, 1)", + "insert into account_model_usage (account_id, model, request_count) values ($1, $2, 1)", ) .bind(account_id) .bind("gpt-5.5") .execute(pool) .await .unwrap(); - sqlx::query("insert into account_cookies (id, account_id, domain, name, value, updated_at) values (?, ?, ?, ?, ?, ?)") + sqlx::query("insert into account_cookies (id, account_id, domain, name, value, updated_at) values ($1, $2, $3, $4, $5, $6)") .bind("cookie_lifecycle") .bind(account_id) .bind("chatgpt.com") .bind("cf_clearance") .bind("cipher") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query("insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?)") - .bind(account_id) - .bind("test") - .bind("2026-06-18T00:10:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query("insert into session_affinities (response_id, account_id, conversation_id, expires_at, created_at) values (?, ?, ?, ?, ?)") - .bind("resp_lifecycle") - .bind(account_id) - .bind("conv_lifecycle") - .bind("2026-06-18T01:00:00Z") - .bind("2026-06-18T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) .execute(pool) .await .unwrap(); } -async fn assert_account_related_rows_deleted(pool: &SqlitePool, account_id: &str) { - let accounts: i64 = sqlx::query_scalar("select count(*) from accounts where id = ?") +async fn assert_account_related_rows_deleted(pool: &PgPool, account_id: &str) { + let accounts: i64 = sqlx::query_scalar("select count(*) from accounts where id = $1") .bind(account_id) .fetch_one(pool) .await .unwrap(); - let usage: i64 = sqlx::query_scalar("select count(*) from account_usage where account_id = ?") + let usage: i64 = sqlx::query_scalar("select count(*) from account_usage where account_id = $1") .bind(account_id) .fetch_one(pool) .await .unwrap(); let model_usage: i64 = - sqlx::query_scalar("select count(*) from account_model_usage where account_id = ?") + sqlx::query_scalar("select count(*) from account_model_usage where account_id = $1") .bind(account_id) .fetch_one(pool) .await .unwrap(); let cookies: i64 = - sqlx::query_scalar("select count(*) from account_cookies where account_id = ?") + sqlx::query_scalar("select count(*) from account_cookies where account_id = $1") .bind(account_id) .fetch_one(pool) .await .unwrap(); - let leases: i64 = - sqlx::query_scalar("select count(*) from account_refresh_leases where account_id = ?") - .bind(account_id) - .fetch_one(pool) - .await - .unwrap(); - let affinities: i64 = - sqlx::query_scalar("select count(*) from session_affinities where account_id = ?") - .bind(account_id) - .fetch_one(pool) - .await - .unwrap(); - assert_eq!(accounts, 0); assert_eq!(usage, 0); assert_eq!(model_usage, 0); assert_eq!(cookies, 0); - assert_eq!(leases, 0); - assert_eq!(affinities, 0); } #[tokio::test] async fn admin_account_status_update_should_update_proxy_account_pool() { let (app, state, pool, _dir) = - admin_accounts_test_app("admin-account-status-runtime-pool.sqlite", 121).await; + admin_accounts_test_app("admin-account-status-runtime-pool", 121).await; seed_account( &pool, NewAccount { @@ -227,7 +216,7 @@ async fn admin_account_status_update_should_update_proxy_account_pool() { state .services .account_pool - .restore_from_repository() + .restore_from_store() .await .unwrap(); assert_eq!( @@ -270,7 +259,7 @@ async fn admin_account_status_update_should_update_proxy_account_pool() { #[tokio::test] async fn admin_account_status_update_should_reject_non_manual_statuses() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-account-status-reject-banned.sqlite", 125).await; + admin_accounts_test_app("admin-account-status-reject-banned", 125).await; seed_account( &pool, NewAccount { @@ -311,7 +300,7 @@ async fn admin_account_status_update_should_reject_non_manual_statuses() { ) .await .unwrap(); - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get("acct_reject_status") .await .unwrap() @@ -333,7 +322,7 @@ async fn admin_account_refresh_should_mark_account_expired_when_refresh_token_re .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( - "admin-account-refresh-rt-reused.sqlite", + "admin-account-refresh-rt-reused", 124, format!("{}/oauth/token", server.uri()), ) @@ -380,7 +369,7 @@ async fn admin_account_refresh_should_mark_account_expired_when_refresh_token_re .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get("acct_refresh_rt_reused") .await .unwrap() @@ -392,6 +381,87 @@ async fn admin_account_refresh_should_mark_account_expired_when_refresh_token_re assert_eq!(stored.status, AccountStatus::Expired); } +#[tokio::test] +async fn admin_account_refresh_should_preserve_quota_exhausted_status_on_success() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/oauth/token")) + .respond_with(wiremock::ResponseTemplate::new(200).set_body_json(json!({ + "access_token": test_jwt( + "chatgpt-quota-refresh-success", + Some("user-quota-refresh-success"), + Some("quota-refresh-success@example.com"), + Some("plus"), + ), + "refresh_token": "refresh-quota-success-rotated" + }))) + .expect(1) + .mount(&server) + .await; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( + "admin-account-refresh-quota-preserved", + 124, + format!("{}/oauth/token", server.uri()), + ) + .await; + seed_account( + &pool, + NewAccount { + id: "acct_refresh_quota_success".to_string(), + email: Some("quota-refresh-success@example.com".to_string()), + account_id: Some("chatgpt-quota-refresh-success".to_string()), + user_id: Some("user-quota-refresh-success".to_string()), + label: None, + plan_type: None, + access_token: SecretString::new( + test_jwt( + "chatgpt-quota-refresh-success", + Some("user-quota-refresh-success"), + Some("quota-refresh-success@example.com"), + Some("plus"), + ) + .into(), + ), + refresh_token: Some(SecretString::new( + "refresh-quota-success".to_string().into(), + )), + access_token_expires_at: None, + status: AccountStatus::QuotaExhausted, + added_at: None, + }, + ) + .await; + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/refresh") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from( + json!({"id": "acct_refresh_quota_success"}).to_string(), + )) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(response).await; + let stored = PgAccountStore::new(pool) + .get("acct_refresh_quota_success") + .await + .unwrap() + .unwrap(); + + assert_eq!(body["data"]["result"], "alive"); + assert_eq!(body["data"]["previousStatus"], "quota_exhausted"); + assert_eq!(stored.status, AccountStatus::QuotaExhausted); + assert_eq!( + stored.refresh_token.unwrap().expose_secret(), + "refresh-quota-success-rotated" + ); +} + #[tokio::test] async fn admin_account_refresh_should_skip_disabled_account_without_consuming_refresh_token() { let server = wiremock::MockServer::start().await; @@ -410,7 +480,7 @@ async fn admin_account_refresh_should_skip_disabled_account_without_consuming_re .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( - "admin-account-refresh-disabled-skip.sqlite", + "admin-account-refresh-disabled-skip", 124, format!("{}/oauth/token", server.uri()), ) @@ -457,7 +527,7 @@ async fn admin_account_refresh_should_skip_disabled_account_without_consuming_re .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get("acct_refresh_disabled") .await .unwrap() @@ -491,7 +561,7 @@ async fn admin_account_refresh_should_skip_expired_account_without_consuming_ref .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( - "admin-account-refresh-expired-skip.sqlite", + "admin-account-refresh-expired-skip", 124, format!("{}/oauth/token", server.uri()), ) @@ -538,7 +608,7 @@ async fn admin_account_refresh_should_skip_expired_account_without_consuming_ref .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get("acct_refresh_expired") .await .unwrap() @@ -557,7 +627,7 @@ async fn admin_account_refresh_should_skip_expired_account_without_consuming_ref #[tokio::test] async fn admin_account_refresh_should_skip_account_without_refresh_token() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-account-refresh-no-rt-skip.sqlite", 124).await; + admin_accounts_test_app("admin-account-refresh-no-rt-skip", 124).await; seed_account( &pool, NewAccount { @@ -598,7 +668,7 @@ async fn admin_account_refresh_should_skip_account_without_refresh_token() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get("acct_refresh_no_rt") .await .unwrap() @@ -619,7 +689,7 @@ async fn admin_account_refresh_should_not_change_status_for_non_permanent_refres .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( - "admin-account-refresh-quota-temporary.sqlite", + "admin-account-refresh-quota-temporary", 124, format!("{}/oauth/token", server.uri()), ) @@ -664,7 +734,7 @@ async fn admin_account_refresh_should_not_change_status_for_non_permanent_refres .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get("acct_refresh_quota") .await .unwrap() @@ -693,7 +763,7 @@ async fn admin_accounts_health_check_should_return_alive_dead_and_skipped_summar .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( - "admin-accounts-health-check.sqlite", + "admin-accounts-health-check", 124, format!("{}/oauth/token", server.uri()), ) @@ -790,7 +860,7 @@ async fn admin_accounts_health_check_should_return_alive_dead_and_skipped_summar .unwrap(); let status = response.status(); let body = response_json(response).await; - let stored_alive = SqliteAccountStore::new(pool) + let stored_alive = PgAccountStore::new(pool) .get("acct_health_alive") .await .unwrap() @@ -810,28 +880,17 @@ async fn admin_accounts_health_check_should_return_alive_dead_and_skipped_summar #[tokio::test] async fn admin_account_create_should_derive_claims_and_store_plain_tokens() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-account-create.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; + let (pool, _dir) = init_test_db("admin-account-create").await; + let redis = create_test_redis("admin-account-create").await; + seed_admin_session(&pool, &redis, "session_1").await; let token = test_jwt( "jwt-account", Some("jwt-user"), Some("jwt@example.com"), Some("team"), ); - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new( &config, @@ -841,7 +900,7 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() { let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state.clone()); + let app = codex_proxy_rs::api::router::router().with_state(state.clone()); let response = app.oneshot( Request::builder().method("POST").uri("/api/admin/accounts") @@ -863,22 +922,11 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() { #[tokio::test] async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_unbound_tokens() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-account-create-invalid.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let (pool, _dir) = init_test_db("admin-account-create-invalid").await; + let redis = create_test_redis("admin-account-create-invalid").await; + seed_admin_session(&pool, &redis, "session_1").await; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new( &config, @@ -888,7 +936,7 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); let cases = [ ("missing tokens", json!({})), @@ -922,7 +970,7 @@ async fn admin_account_manual_create_from_refresh_token_should_not_store_consume .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( - "admin-account-create-rt-no-rotation.sqlite", + "admin-account-create-rt-no-rotation", 126, format!("{}/oauth/token", server.uri()), ) @@ -931,7 +979,7 @@ async fn admin_account_manual_create_from_refresh_token_should_not_store_consume let response = post_admin_account(&app, json!({"refreshToken": "rt-create-source"})).await; let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get(body["data"]["id"].as_str().unwrap()) .await .unwrap() @@ -946,7 +994,7 @@ async fn admin_account_manual_create_from_refresh_token_should_not_store_consume async fn admin_account_manual_create_should_accept_current_openai_token_without_chatgpt_account_id() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-account-create-current-openai-token.sqlite", 119).await; + admin_accounts_test_app("admin-account-create-current-openai-token", 119).await; let token = test_jwt_with_exp( None, Some("current-user"), @@ -962,7 +1010,7 @@ async fn admin_account_manual_create_should_accept_current_openai_token_without_ .await; let status = response.status(); let body = response_json(response).await; - let stored = SqliteAccountStore::new(pool) + let stored = PgAccountStore::new(pool) .get(body["data"]["id"].as_str().unwrap()) .await .unwrap() diff --git a/backend/tests/admin/accounts/list.rs b/backend/tests/api/admin/accounts/list.rs similarity index 75% rename from backend/tests/admin/accounts/list.rs rename to backend/tests/api/admin/accounts/list.rs index b1d40c1b0..56cfa499b 100644 --- a/backend/tests/admin/accounts/list.rs +++ b/backend/tests/api/admin/accounts/list.rs @@ -2,30 +2,19 @@ use super::*; #[tokio::test] async fn admin_accounts_list_should_not_expose_account_tokens() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)") + let (pool, _dir) = init_test_db("admin-accounts-list").await; + let redis = create_test_redis("admin-accounts-list").await; + seed_admin_session(&pool, &redis, "session_1").await; + sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6)") .bind("acct_plain").bind("user@example.com").bind("plain-access-token") - .bind("active").bind("2026-06-18T00:00:00Z").bind("2026-06-18T00:00:00Z") + .bind("active").bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")).bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) .execute(&pool).await.unwrap(); - sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?)") + sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7)") .bind("acct_refresh").bind("refresh@example.com").bind("refresh-access-token").bind("refresh-secret-token") - .bind("active").bind("2026-06-19T00:00:00Z").bind("2026-06-19T00:00:00Z") + .bind("active").bind(crate::support::storage::timestamp("2026-06-19T00:00:00Z")).bind(crate::support::storage::timestamp("2026-06-19T00:00:00Z")) .execute(&pool).await.unwrap(); - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new( &config, @@ -35,7 +24,7 @@ async fn admin_accounts_list_should_not_expose_account_tokens() { let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); let response = app .oneshot( @@ -75,27 +64,32 @@ async fn admin_accounts_list_should_not_expose_account_tokens() { #[tokio::test] async fn admin_accounts_list_should_derive_refreshing_display_status_from_active_lease() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-refreshing-display.sqlite", 72).await; - sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?)") + let (app, state, pool, _dir) = + admin_accounts_test_app("admin-accounts-refreshing-display", 72).await; + sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7)") .bind("acct_refreshing") .bind("refreshing@example.com") .bind("access-token") .bind("refresh-token") .bind("active") - .bind("2026-07-09T00:00:00Z") - .bind("2026-07-09T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-07-09T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-07-09T00:00:00Z")) .execute(&pool) .await .unwrap(); - sqlx::query("insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?)") - .bind("acct_refreshing") - .bind("test-owner") - .bind("2099-01-01T00:00:00Z") - .bind("2026-07-09T00:00:00Z") - .execute(&pool) + let now = chrono::Utc::now(); + assert!(state + .services + .background_tasks + .refresh_leases + .try_acquire( + "acct_refreshing", + "test-owner", + now + chrono::Duration::hours(1), + now, + ) .await - .unwrap(); + .unwrap()); let response = app .oneshot( @@ -120,27 +114,32 @@ async fn admin_accounts_list_should_derive_refreshing_display_status_from_active #[tokio::test] async fn admin_accounts_list_should_ignore_refresh_lease_for_expired_account() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-expired-lease.sqlite", 72).await; - sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?)") + let (app, state, pool, _dir) = + admin_accounts_test_app("admin-accounts-expired-lease", 72).await; + sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7)") .bind("acct_expired") .bind("expired@example.com") .bind("access-token") .bind("refresh-token") .bind("expired") - .bind("2026-07-09T00:00:00Z") - .bind("2026-07-09T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-07-09T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-07-09T00:00:00Z")) .execute(&pool) .await .unwrap(); - sqlx::query("insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?)") - .bind("acct_expired") - .bind("stale-owner") - .bind("2099-01-01T00:00:00Z") - .bind("2026-07-09T00:00:00Z") - .execute(&pool) + let now = chrono::Utc::now(); + assert!(state + .services + .background_tasks + .refresh_leases + .try_acquire( + "acct_expired", + "stale-owner", + now + chrono::Duration::hours(1), + now, + ) .await - .unwrap(); + .unwrap()); let response = app .oneshot( @@ -165,8 +164,7 @@ async fn admin_accounts_list_should_ignore_refresh_lease_for_expired_account() { #[tokio::test] async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-stats.sqlite", 68).await; + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-accounts-stats", 68).await; seed_usage_account( &pool, UsageAccountSeed { @@ -190,7 +188,7 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { }, ) .await; - sqlx::query("update accounts set quota_json = ?, quota_fetched_at = ? where id = ?") + sqlx::query("update accounts set quota_json = $1::jsonb, quota_fetched_at = $2 where id = $3") .bind( json!({ "plan_type": "free", @@ -249,13 +247,13 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { }) .to_string(), ) - .bind("2026-06-23T08:51:09Z") + .bind(crate::support::storage::timestamp("2026-06-23T08:51:09Z")) .bind("acct_stats") .execute(&pool) .await .unwrap(); sqlx::query( - "insert into account_model_usage (account_id, model, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", + "insert into account_model_usage (account_id, model, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values ($1, $2, $3, $4, $5, $6, $7)", ) .bind("acct_stats") .bind("gpt-5.5") @@ -263,12 +261,12 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { .bind(3_500_100) .bind(13_920) .bind(3_400_010) - .bind("2026-06-23T08:51:13Z") + .bind(crate::support::storage::timestamp("2026-06-23T08:51:13Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into account_model_usage (account_id, model, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)", + "insert into account_model_usage (account_id, model, request_count, input_tokens, output_tokens, cached_tokens, last_used_at) values ($1, $2, $3, $4, $5, $6, $7)", ) .bind("acct_stats") .bind("gpt-5") @@ -276,63 +274,63 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { .bind(200) .bind(30) .bind(20) - .bind("2026-06-23T08:52:13Z") + .bind(crate::support::storage::timestamp("2026-06-23T08:52:13Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into request_time_buckets (bucket_start, provider, account_id, model, service_tier, success_count, input_tokens, output_tokens, cached_tokens, updated_at) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2026-06-23T08:45:00+00:00") + .bind(crate::support::storage::timestamp("2026-06-23T08:45:00+00:00")) .bind("acct_stats") .bind("gpt-5") .bind(3) .bind(60_000) .bind(1_500) .bind(20_000) - .bind("2026-06-23T08:45:00Z") + .bind(crate::support::storage::timestamp("2026-06-23T08:45:00Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into request_time_buckets (bucket_start, provider, account_id, model, service_tier, success_count, input_tokens, output_tokens, cached_tokens, updated_at) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2026-06-23T09:00:00+00:00") + .bind(crate::support::storage::timestamp("2026-06-23T09:00:00+00:00")) .bind("acct_stats") .bind("gpt-5.5") .bind(4) .bind(60_000) .bind(2_500) .bind(20_000) - .bind("2026-06-23T09:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-23T09:00:00Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into request_time_buckets (bucket_start, provider, account_id, model, service_tier, success_count, input_tokens, output_tokens, cached_tokens, updated_at) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2026-06-22T10:00:00+00:00") + .bind(crate::support::storage::timestamp("2026-06-22T10:00:00+00:00")) .bind("acct_stats") .bind("gpt-five-hour-window") .bind(1) .bind(2_000) .bind(500) .bind(100) - .bind("2026-06-22T10:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-22T10:00:00Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into request_time_buckets (bucket_start, provider, account_id, model, service_tier, success_count, input_tokens, output_tokens, cached_tokens, updated_at) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2026-06-19T23:45:00+00:00") + .bind(crate::support::storage::timestamp("2026-06-19T23:45:00+00:00")) .bind("acct_stats") .bind("gpt-outside-window") .bind(99) .bind(900_000) .bind(90_000) .bind(9_000) - .bind("2026-06-19T23:45:00Z") + .bind(crate::support::storage::timestamp("2026-06-19T23:45:00Z")) .execute(&pool) .await .unwrap(); @@ -393,7 +391,7 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() { #[tokio::test] async fn admin_accounts_list_should_use_largest_quota_reset_interval_for_usage_stats() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-largest-quota-window.sqlite", 71).await; + admin_accounts_test_app("admin-accounts-largest-quota-window", 71).await; seed_usage_account( &pool, UsageAccountSeed { @@ -417,7 +415,7 @@ async fn admin_accounts_list_should_use_largest_quota_reset_interval_for_usage_s }, ) .await; - sqlx::query("update accounts set quota_json = ?, quota_fetched_at = ? where id = ?") + sqlx::query("update accounts set quota_json = $1::jsonb, quota_fetched_at = $2 where id = $3") .bind( json!({ "plan_type": "plus", @@ -446,36 +444,36 @@ async fn admin_accounts_list_should_use_largest_quota_reset_interval_for_usage_s }) .to_string(), ) - .bind("2026-06-30T01:10:00Z") + .bind(crate::support::storage::timestamp("2026-06-30T01:10:00Z")) .bind("acct_custom_window") .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into request_time_buckets (bucket_start, provider, account_id, model, service_tier, success_count, input_tokens, output_tokens, cached_tokens, updated_at) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2026-06-30T01:00:00+00:00") + .bind(crate::support::storage::timestamp("2026-06-30T01:00:00+00:00")) .bind("acct_custom_window") .bind("gpt-custom-14d") .bind(5) .bind(10_000) .bind(2_000) .bind(1_000) - .bind("2026-06-30T01:10:00Z") + .bind(crate::support::storage::timestamp("2026-06-30T01:10:00Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into request_time_buckets (bucket_start, provider, account_id, model, service_tier, success_count, input_tokens, output_tokens, cached_tokens, updated_at) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2026-06-29T23:45:00+00:00") + .bind(crate::support::storage::timestamp("2026-06-29T23:45:00+00:00")) .bind("acct_custom_window") .bind("gpt-before-14d") .bind(99) .bind(900_000) .bind(90_000) .bind(9_000) - .bind("2026-06-29T23:59:00Z") + .bind(crate::support::storage::timestamp("2026-06-29T23:59:00Z")) .execute(&pool) .await .unwrap(); @@ -507,13 +505,13 @@ async fn admin_accounts_list_should_use_largest_quota_reset_interval_for_usage_s #[tokio::test] async fn admin_accounts_list_should_require_current_window_start_for_model_stats() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-model-window-start.sqlite", 69).await; - sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, 'active', ?, ?)") + admin_accounts_test_app("admin-accounts-model-window-start", 69).await; + sqlx::query("insert into accounts (id, email, access_token, status, added_at, updated_at) values ($1, $2, $3, 'active', $4, $5)") .bind("acct_no_window_start") .bind("stats@example.com") .bind("access-token") - .bind("2026-06-11T00:00:00Z") - .bind("2026-06-11T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-11T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-06-11T00:00:00Z")) .execute(&pool) .await .unwrap(); @@ -528,7 +526,7 @@ async fn admin_accounts_list_should_require_current_window_start_for_model_stats window_reset_at, limit_window_seconds, last_used_at - ) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) values ($1, $2, $3, $4, $5, $6, $7, $8, $9)", ) .bind("acct_no_window_start") .bind(3) @@ -536,23 +534,23 @@ async fn admin_accounts_list_should_require_current_window_start_for_model_stats .bind(120) .bind(40) .bind(20) - .bind("2026-06-27T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-27T00:00:00Z")) .bind(604_800) - .bind("2026-06-23T08:50:13Z") + .bind(crate::support::storage::timestamp("2026-06-23T08:50:13Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into request_time_buckets (bucket_start, provider, account_id, model, service_tier, success_count, input_tokens, output_tokens, cached_tokens, updated_at) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2026-06-23T08:45:00+00:00") + .bind(crate::support::storage::timestamp("2026-06-23T08:45:00+00:00")) .bind("acct_no_window_start") .bind("gpt-5.5") .bind(3) .bind(120) .bind(40) .bind(20) - .bind("2026-06-23T08:45:00Z") + .bind(crate::support::storage::timestamp("2026-06-23T08:45:00Z")) .execute(&pool) .await .unwrap(); @@ -579,7 +577,7 @@ async fn admin_accounts_list_should_require_current_window_start_for_model_stats #[tokio::test] async fn admin_accounts_list_should_include_bucket_containing_window_start_for_model_stats() { let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-model-window-bucket-start.sqlite", 70).await; + admin_accounts_test_app("admin-accounts-model-window-bucket-start", 70).await; seed_usage_account( &pool, UsageAccountSeed { @@ -604,16 +602,16 @@ async fn admin_accounts_list_should_include_bucket_containing_window_start_for_m ) .await; sqlx::query( - "insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into request_time_buckets (bucket_start, provider, account_id, model, service_tier, success_count, input_tokens, output_tokens, cached_tokens, updated_at) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2026-06-29T12:45:00+00:00") + .bind(crate::support::storage::timestamp("2026-06-29T12:45:00+00:00")) .bind("acct_bucket_start") .bind("gpt-5.5") .bind(2) .bind(120) .bind(40) .bind(20) - .bind("2026-06-29T12:59:42.044220420+00:00") + .bind(crate::support::storage::timestamp("2026-06-29T12:59:42.044220420+00:00")) .execute(&pool) .await .unwrap(); @@ -640,8 +638,7 @@ async fn admin_accounts_list_should_include_bucket_containing_window_start_for_m #[tokio::test] async fn admin_accounts_list_should_return_numbered_page_with_search_total() { - let (app, _state, pool, _dir) = - admin_accounts_test_app("admin-accounts-numbered.sqlite", 67).await; + let (app, _state, pool, _dir) = admin_accounts_test_app("admin-accounts-numbered", 67).await; for (id, email, label, added_at) in [ ( "acct_prod_new", @@ -662,14 +659,14 @@ async fn admin_accounts_list_should_return_numbered_page_with_search_total() { "2026-06-18T00:00:00Z", ), ] { - sqlx::query("insert into accounts (id, email, label, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?)") + sqlx::query("insert into accounts (id, email, label, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7)") .bind(id) .bind(email) .bind(label) .bind("cipher") .bind("active") - .bind(added_at) - .bind(added_at) + .bind(crate::support::storage::timestamp(added_at)) + .bind(crate::support::storage::timestamp(added_at)) .execute(&pool) .await .unwrap(); diff --git a/backend/tests/admin/accounts/mod.rs b/backend/tests/api/admin/accounts/mod.rs similarity index 65% rename from backend/tests/admin/accounts/mod.rs rename to backend/tests/api/admin/accounts/mod.rs index 27259b274..9d5ce1956 100644 --- a/backend/tests/admin/accounts/mod.rs +++ b/backend/tests/api/admin/accounts/mod.rs @@ -5,30 +5,26 @@ use axum::{ http::{Request, StatusCode}, }; use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::usage_record_store::SqliteUsageRecordStore, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::{ - cookies::SqliteCookieStore, - model::AccountStatus, - store::{NewAccount, SqliteAccountStore}, - token_refresh::RefreshLeaseStore, + accounts::{ + account::AccountStatus, + cookies::PgCookieStore, + store::{NewAccount, PgAccountStore}, }, - upstream::fingerprint::FingerprintRepository, + bootstrap::services::Services, + bootstrap::state::AppState, }; use secrecy::{ExposeSecret, SecretString}; use serde_json::{json, Value}; -use sqlx::SqlitePool; +use sqlx::PgPool; use tower::util::ServiceExt; use crate::support::jwt::unsigned_jwt; use crate::support::{ - admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint, + admin::seed_admin_session, + config::test_config, + fingerprint::runtime_fingerprint, http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; mod exporting; @@ -59,12 +55,12 @@ struct UsageAccountSeed<'a> { last_used_at: &'a str, } -async fn seed_usage_account(pool: &SqlitePool, seed: UsageAccountSeed<'_>) { - sqlx::query("insert into accounts (id, email, label, plan_type, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, 'active', ?, ?)") +async fn seed_usage_account(pool: &PgPool, seed: UsageAccountSeed<'_>) { + sqlx::query("insert into accounts (id, email, label, plan_type, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, 'active', $6, $7)") .bind(seed.id).bind(seed.email).bind(seed.label).bind(seed.plan_type).bind("access-token") - .bind("2026-06-11T00:00:00Z").bind("2026-06-11T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-11T00:00:00Z")).bind(crate::support::storage::timestamp("2026-06-11T00:00:00Z")) .execute(pool).await.unwrap(); - sqlx::query("insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds, last_used_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)") + sqlx::query("insert into account_usage (account_id, request_count, empty_response_count, input_tokens, output_tokens, cached_tokens, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, 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)") .bind(seed.id) .bind(seed.request_count) .bind(seed.empty_response_count) @@ -75,10 +71,10 @@ async fn seed_usage_account(pool: &SqlitePool, seed: UsageAccountSeed<'_>) { .bind(seed.window_input_tokens) .bind(seed.window_output_tokens) .bind(seed.window_cached_tokens) - .bind(seed.window_started_at) - .bind(seed.window_reset_at) + .bind(crate::support::storage::timestamp(seed.window_started_at)) + .bind(crate::support::storage::timestamp(seed.window_reset_at)) .bind(seed.limit_window_seconds) - .bind(seed.last_used_at) + .bind(crate::support::storage::timestamp(seed.last_used_at)) .execute(pool).await.unwrap(); } @@ -100,7 +96,7 @@ async fn post_admin_account(app: &axum::Router, payload: Value) -> axum::respons async fn admin_accounts_test_app( db_name: &str, key_byte: u8, -) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { +) -> (axum::Router, AppState, PgPool, tempfile::TempDir) { admin_accounts_test_app_with_api_base_url( db_name, key_byte, @@ -113,7 +109,7 @@ async fn admin_accounts_test_app_with_api_base_url( db_name: &str, key_byte: u8, api_base_url: String, -) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { +) -> (axum::Router, AppState, PgPool, tempfile::TempDir) { admin_accounts_test_app_with_overrides(db_name, key_byte, api_base_url, None).await } @@ -121,7 +117,7 @@ async fn admin_accounts_test_app_with_oauth_token_endpoint( db_name: &str, key_byte: u8, oauth_token_endpoint: String, -) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { +) -> (axum::Router, AppState, PgPool, tempfile::TempDir) { admin_accounts_test_app_with_overrides( db_name, key_byte, @@ -136,7 +132,7 @@ async fn admin_accounts_test_app_with_api_base_url_and_oauth_token_endpoint( key_byte: u8, api_base_url: String, oauth_token_endpoint: String, -) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { +) -> (axum::Router, AppState, PgPool, tempfile::TempDir) { admin_accounts_test_app_with_overrides( db_name, key_byte, @@ -151,27 +147,16 @@ async fn admin_accounts_test_app_with_overrides( _key_byte: u8, api_base_url: String, oauth_token_endpoint: Option, -) -> (axum::Router, AppState, SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let mut config = test_config(url); +) -> (axum::Router, AppState, PgPool, tempfile::TempDir) { + let (pool, dir) = init_test_db(db_name).await; + let redis = create_test_redis(db_name).await; + seed_admin_session(&pool, &redis, "session_1").await; + let mut config = test_config(test_database_url()); config.api.base_url = api_base_url; if let Some(oauth_token_endpoint) = oauth_token_endpoint { config.auth.oauth_token_endpoint = oauth_token_endpoint; } - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new( &config, @@ -181,12 +166,12 @@ async fn admin_accounts_test_app_with_overrides( let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state.clone()); + let app = codex_proxy_rs::api::router::router().with_state(state.clone()); (app, state, pool, dir) } -async fn seed_account(pool: &SqlitePool, account: NewAccount) { - SqliteAccountStore::new(pool.clone()) +async fn seed_account(pool: &PgPool, account: NewAccount) { + PgAccountStore::new(pool.clone()) .insert(account) .await .unwrap(); diff --git a/backend/tests/admin/accounts/oauth.rs b/backend/tests/api/admin/accounts/oauth.rs similarity index 96% rename from backend/tests/admin/accounts/oauth.rs rename to backend/tests/api/admin/accounts/oauth.rs index ecbcaf2cf..498070077 100644 --- a/backend/tests/admin/accounts/oauth.rs +++ b/backend/tests/api/admin/accounts/oauth.rs @@ -7,7 +7,7 @@ use wiremock::{ #[tokio::test] async fn admin_account_oauth_authorize_should_return_pkce_auth_url() { let (app, _state, _pool, _dir) = - admin_accounts_test_app("admin-account-oauth-authorize.sqlite", 94).await; + admin_accounts_test_app("admin-account-oauth-authorize", 94).await; let response = post_oauth_authorize(&app, "req_oauth_authorize").await; @@ -50,7 +50,7 @@ async fn admin_account_oauth_authorize_should_return_pkce_auth_url() { async fn admin_account_oauth_exchange_should_reject_callback_without_state() { let server = MockServer::start().await; let (app, _state, _pool, _dir) = admin_accounts_test_app_with_oauth_token_endpoint( - "admin-account-oauth-missing-state.sqlite", + "admin-account-oauth-missing-state", 95, format!("{}/oauth/token", server.uri()), ) @@ -128,7 +128,7 @@ async fn admin_account_oauth_exchange_should_import_account_tokens() { .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url_and_oauth_token_endpoint( - "admin-account-oauth-exchange.sqlite", + "admin-account-oauth-exchange", 96, server.uri(), format!("{}/oauth/token", server.uri()), @@ -168,7 +168,7 @@ async fn admin_account_oauth_exchange_should_import_account_tokens() { let body = response_json(response).await; assert_eq!(body["data"]["imported"], 1); let stored: (String, String, String, String) = sqlx::query_as( - "select email, chatgpt_account_id, access_token, refresh_token from accounts where chatgpt_account_id = ?", + "select email, chatgpt_account_id, access_token, refresh_token from accounts where chatgpt_account_id = $1", ) .bind("chatgpt-oauth") .fetch_one(&pool) diff --git a/backend/tests/admin/accounts/quota.rs b/backend/tests/api/admin/accounts/quota.rs similarity index 82% rename from backend/tests/admin/accounts/quota.rs rename to backend/tests/api/admin/accounts/quota.rs index 4f2141bca..c6d6b1ee0 100644 --- a/backend/tests/admin/accounts/quota.rs +++ b/backend/tests/api/admin/accounts/quota.rs @@ -29,7 +29,7 @@ async fn admin_account_quota_should_send_usage_cookie() { .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-quota-cookie.sqlite", + "admin-account-quota-cookie", 88, format!("{}/backend-api", server.uri()), ) @@ -63,7 +63,7 @@ async fn admin_account_quota_should_send_usage_cookie() { window_reset_at, limit_window_seconds, last_used_at - ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)", ) .bind("acct_quota_cookie") .bind(9) @@ -71,60 +71,68 @@ async fn admin_account_quota_should_send_usage_cookie() { .bind(9_000) .bind(900) .bind(90) - .bind("2027-01-01T00:00:00Z") - .bind("2027-01-01T05:00:00Z") + .bind(crate::support::storage::timestamp("2027-01-01T00:00:00Z")) + .bind(crate::support::storage::timestamp("2027-01-01T05:00:00Z")) .bind(18_000) - .bind("2027-01-15T04:00:00Z") + .bind(crate::support::storage::timestamp("2027-01-15T04:00:00Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets ( + "insert into request_time_buckets ( bucket_start, + provider, account_id, model, - request_count, + service_tier, + success_count, input_tokens, output_tokens, cached_tokens, updated_at - ) values (?, ?, ?, ?, ?, ?, ?, ?)", + ) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2027-01-15T04:00:00+00:00") + .bind(crate::support::storage::timestamp( + "2027-01-15T04:00:00+00:00", + )) .bind("acct_quota_cookie") .bind("gpt-5") .bind(1) .bind(1_500) .bind(500) .bind(300) - .bind("2027-01-15T04:00:00Z") + .bind(crate::support::storage::timestamp("2027-01-15T04:00:00Z")) .execute(&pool) .await .unwrap(); sqlx::query( - "insert into usage_time_buckets ( + "insert into request_time_buckets ( bucket_start, + provider, account_id, model, - request_count, + service_tier, + success_count, input_tokens, output_tokens, cached_tokens, updated_at - ) values (?, ?, ?, ?, ?, ?, ?, ?)", + ) values ($1, 'openai', $2, $3, '__unknown__', $4, $5, $6, $7, $8)", ) - .bind("2027-01-15T09:00:00+00:00") + .bind(crate::support::storage::timestamp( + "2027-01-15T09:00:00+00:00", + )) .bind("acct_quota_cookie") .bind("gpt-5") .bind(1) .bind(1_500) .bind(500) .bind(300) - .bind("2027-01-15T09:00:00Z") + .bind(crate::support::storage::timestamp("2027-01-15T09:00:00Z")) .execute(&pool) .await .unwrap(); - SqliteCookieStore::new(pool.clone()) + PgCookieStore::new(pool.clone()) .capture_set_cookie( "acct_quota_cookie", "cf_clearance=admin-quota; Domain=.chatgpt.com; Path=/", diff --git a/backend/tests/admin/accounts/testing.rs b/backend/tests/api/admin/accounts/testing.rs similarity index 75% rename from backend/tests/admin/accounts/testing.rs rename to backend/tests/api/admin/accounts/testing.rs index 5f2fde6ae..6962e25b6 100644 --- a/backend/tests/admin/accounts/testing.rs +++ b/backend/tests/api/admin/accounts/testing.rs @@ -1,12 +1,18 @@ use super::*; use axum::body::to_bytes; -use codex_proxy_rs::upstream::protocol::sse::{encode_sse_event, parse_sse_events}; +use codex_proxy_rs::{ + models::{snapshot::ModelPlanSnapshot, store::ModelSnapshotStore}, + upstream::openai::protocol::{ + model_catalog::BackendModelEntry, + sse::{encode_sse_event, parse_sse_events}, + }, +}; use wiremock::{ matchers::{header, method, path}, Mock, MockServer, ResponseTemplate, }; -async fn seed_test_account(pool: &SqlitePool) { +async fn seed_test_account(pool: &PgPool) { seed_account( pool, NewAccount { @@ -35,16 +41,16 @@ fn test_events(body: &str) -> Vec { } #[tokio::test] -async fn account_models_should_return_database_snapshot_models() { +async fn account_models_should_return_redis_snapshot_models() { let server = MockServer::start().await; let (app, state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-models.sqlite", + "admin-account-models", 91, format!("{}/backend-api", server.uri()), ) .await; seed_test_account(&pool).await; - seed_model_snapshot(&pool, "plus").await; + seed_model_snapshot(&state, "plus").await; state .services .models @@ -67,7 +73,10 @@ async fn account_models_should_return_database_snapshot_models() { let requests = server.received_requests().await.unwrap(); assert_eq!(response.status(), StatusCode::OK); - assert!(requests.is_empty(), "model list should be loaded from DB"); + assert!( + requests.is_empty(), + "model list should be loaded from Redis" + ); let body = response_json(response).await; assert_eq!(body["data"]["models"][0]["id"], "gpt-5.5"); assert_eq!(body["data"]["models"][0]["label"], "GPT 5.5"); @@ -100,7 +109,7 @@ async fn account_models_should_fetch_missing_plan_snapshot_with_account_token() .mount(&server) .await; let (app, state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-models-unfetched-plan.sqlite", + "admin-account-models-unfetched-plan", 93, format!("{}/backend-api", server.uri()), ) @@ -122,7 +131,7 @@ async fn account_models_should_fetch_missing_plan_snapshot_with_account_token() }, ) .await; - seed_model_snapshot(&pool, "plus").await; + seed_model_snapshot(&state, "plus").await; state .services .models @@ -149,20 +158,21 @@ async fn account_models_should_fetch_missing_plan_snapshot_with_account_token() let body = response_json(response).await; assert_eq!(body["data"]["models"][0]["id"], "gpt-plan-b-live"); assert_eq!(body["data"]["models"][0]["label"], "GPT Plan B Live"); - let stored_plan_b: i64 = - sqlx::query_scalar("select count(*) from model_plan_snapshots where plan_type = ?") - .bind("plan-b") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(stored_plan_b, 1); + let stored = state + .services + .background_tasks + .model_snapshots + .list_plan_snapshots() + .await + .unwrap(); + assert!(stored.iter().any(|snapshot| snapshot.plan_type == "plan-b")); } #[tokio::test] async fn account_models_should_return_models_for_account_plan_snapshot() { let server = MockServer::start().await; let (app, state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-models-matching-plan.sqlite", + "admin-account-models-matching-plan", 94, format!("{}/backend-api", server.uri()), ) @@ -184,8 +194,8 @@ async fn account_models_should_return_models_for_account_plan_snapshot() { }, ) .await; - seed_model_snapshot(&pool, "plan-a").await; - seed_single_model_snapshot(&pool, "plan-b", "gpt-plan-b-only", "GPT Plan B Only").await; + seed_model_snapshot(&state, "plan-a").await; + seed_single_model_snapshot(&state, "plan-b", "gpt-plan-b-only", "GPT Plan B Only").await; state .services .models @@ -208,85 +218,62 @@ async fn account_models_should_return_models_for_account_plan_snapshot() { let requests = server.received_requests().await.unwrap(); assert_eq!(response.status(), StatusCode::OK); - assert!(requests.is_empty(), "model list should be loaded from DB"); + assert!( + requests.is_empty(), + "model list should be loaded from Redis" + ); let body = response_json(response).await; assert_eq!(body["data"]["models"].as_array().unwrap().len(), 1); assert_eq!(body["data"]["models"][0]["id"], "gpt-plan-b-only"); assert_eq!(body["data"]["models"][0]["label"], "GPT Plan B Only"); } -async fn seed_model_snapshot(pool: &SqlitePool, plan_type: &str) { - let models_json = json!([ - { - "id": "gpt-5.5", - "displayName": "GPT 5.5", - "description": "Test model", - "isDefault": false, - "supportedReasoningEfforts": [{"reasoningEffort": "medium", "description": "medium"}], - "defaultReasoningEffort": "medium", - "inputModalities": ["text"], - "outputModalities": ["text"], - "supportsPersonality": false, - "upgrade": null, - "source": "test" - }, - { - "id": "gpt-5.4", - "displayName": "GPT 5.4", - "description": "Test model", - "isDefault": false, - "supportedReasoningEfforts": [{"reasoningEffort": "medium", "description": "medium"}], - "defaultReasoningEffort": "medium", - "inputModalities": ["text"], - "outputModalities": ["text"], - "supportsPersonality": false, - "upgrade": null, - "source": "test" - } - ]) - .to_string(); - sqlx::query( - "insert or replace into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?)", - ) - .bind(plan_type) - .bind(models_json) - .bind("2026-06-29T00:00:00Z") - .execute(pool) - .await - .unwrap(); +async fn seed_model_snapshot(state: &AppState, plan_type: &str) { + let snapshot = ModelPlanSnapshot::from_backend_entries( + plan_type, + vec![ + BackendModelEntry { + id: Some("gpt-5.5".to_string()), + name: Some("GPT 5.5".to_string()), + ..BackendModelEntry::default() + }, + BackendModelEntry { + id: Some("gpt-5.4".to_string()), + name: Some("GPT 5.4".to_string()), + ..BackendModelEntry::default() + }, + ], + ); + state + .services + .background_tasks + .model_snapshots + .replace_plan_snapshot(&snapshot) + .await + .unwrap(); } async fn seed_single_model_snapshot( - pool: &SqlitePool, + state: &AppState, plan_type: &str, model_id: &str, display_name: &str, ) { - let models_json = json!([ - { - "id": model_id, - "displayName": display_name, - "description": "Test model", - "isDefault": false, - "supportedReasoningEfforts": [{"reasoningEffort": "medium", "description": "medium"}], - "defaultReasoningEffort": "medium", - "inputModalities": ["text"], - "outputModalities": ["text"], - "supportsPersonality": false, - "upgrade": null, - "source": "test" - } - ]) - .to_string(); - sqlx::query( - "insert or replace into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?)", - ) - .bind(plan_type) - .bind(models_json) - .bind("2026-06-29T00:00:00Z") - .execute(pool) - .await - .unwrap(); + let snapshot = ModelPlanSnapshot::from_backend_entries( + plan_type, + vec![BackendModelEntry { + id: Some(model_id.to_string()), + name: Some(display_name.to_string()), + ..BackendModelEntry::default() + }], + ); + state + .services + .background_tasks + .model_snapshots + .replace_plan_snapshot(&snapshot) + .await + .unwrap(); } #[tokio::test] @@ -314,13 +301,13 @@ async fn account_test_stream_should_translate_upstream_responses_sse() { .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-test-stream.sqlite", + "admin-account-test-stream", 92, format!("{}/backend-api", server.uri()), ) .await; seed_test_account(&pool).await; - let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()); + let store = codex_proxy_rs::accounts::store::PgAccountStore::new(pool.clone()); store .set_status("acct_test", AccountStatus::QuotaExhausted) .await @@ -386,6 +373,58 @@ async fn account_test_stream_should_translate_upstream_responses_sse() { ); } +#[tokio::test] +async fn account_test_stream_should_preserve_manually_disabled_status_on_success() { + let server = MockServer::start().await; + let upstream_sse = encode_sse_event("", &json!({ "type": "response.completed" }).to_string()); + Mock::given(method("POST")) + .and(path("/backend-api/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(upstream_sse), + ) + .mount(&server) + .await; + let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( + "admin-account-test-disabled-preserved", + 94, + format!("{}/backend-api", server.uri()), + ) + .await; + seed_test_account(&pool).await; + let store = codex_proxy_rs::accounts::store::PgAccountStore::new(pool.clone()); + store + .set_status("acct_test", AccountStatus::Disabled) + .await + .expect("pre-test status should persist"); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/api/admin/accounts/test?id=acct_test&modelId=gpt-5.5") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let body = String::from_utf8( + to_bytes(response.into_body(), usize::MAX) + .await + .unwrap() + .to_vec(), + ) + .unwrap(); + + assert_eq!(test_events(&body).last().unwrap()["type"], "test_complete"); + assert_eq!( + store.get("acct_test").await.unwrap().unwrap().status, + AccountStatus::Disabled + ); +} + #[tokio::test] async fn account_test_stream_should_mark_expired_after_auth_failure() { let server = MockServer::start().await; @@ -400,7 +439,7 @@ async fn account_test_stream_should_mark_expired_after_auth_failure() { .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-test-expired.sqlite", + "admin-account-test-expired", 95, format!("{}/backend-api", server.uri()), ) @@ -426,7 +465,7 @@ async fn account_test_stream_should_mark_expired_after_auth_failure() { ) .unwrap(); let events = test_events(&body); - let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool); + let store = codex_proxy_rs::accounts::store::PgAccountStore::new(pool); let stored = store.get("acct_test").await.unwrap().unwrap(); assert_eq!(events.last().unwrap()["type"], "error"); @@ -460,7 +499,7 @@ async fn account_test_stream_should_mark_quota_exhausted_after_failed_sse() { .mount(&server) .await; let (app, _state, pool, _dir) = admin_accounts_test_app_with_api_base_url( - "admin-account-test-quota.sqlite", + "admin-account-test-quota", 96, format!("{}/backend-api", server.uri()), ) @@ -486,7 +525,7 @@ async fn account_test_stream_should_mark_quota_exhausted_after_failed_sse() { ) .unwrap(); let events = test_events(&body); - let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool); + let store = codex_proxy_rs::accounts::store::PgAccountStore::new(pool); let stored = store.get("acct_test").await.unwrap().unwrap(); assert_eq!(events.last().unwrap()["type"], "error"); diff --git a/backend/tests/admin/auth/mod.rs b/backend/tests/api/admin/auth/mod.rs similarity index 100% rename from backend/tests/admin/auth/mod.rs rename to backend/tests/api/admin/auth/mod.rs diff --git a/backend/tests/admin/auth/password.rs b/backend/tests/api/admin/auth/password.rs similarity index 100% rename from backend/tests/admin/auth/password.rs rename to backend/tests/api/admin/auth/password.rs diff --git a/backend/tests/admin/auth/session.rs b/backend/tests/api/admin/auth/session.rs similarity index 70% rename from backend/tests/admin/auth/session.rs rename to backend/tests/api/admin/auth/session.rs index 187d1e34b..99a97e18d 100644 --- a/backend/tests/admin/auth/session.rs +++ b/backend/tests/api/admin/auth/session.rs @@ -3,15 +3,20 @@ use axum::{ http::{Request, StatusCode}, }; use chrono::Utc; -use codex_proxy_rs::infra::{database::connect_sqlite, identity::hash_admin_password}; -use sqlx::SqlitePool; +use codex_proxy_rs::infra::identity::hash_admin_password; +use sqlx::PgPool; use tower::util::ServiceExt; -use crate::support::{config::test_config, fingerprint::runtime_fingerprint, http::response_json}; +use crate::support::{ + config::test_config, + fingerprint::runtime_fingerprint, + http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, +}; #[tokio::test] async fn admin_login_should_issue_http_only_session_cookie() { - let (app, _dir, _pool) = admin_login_test_app("admin-login.sqlite", "correct-password").await; + let (app, _dir, _pool) = admin_login_test_app("admin-login", "correct-password").await; let response = app .clone() @@ -60,7 +65,7 @@ async fn admin_login_should_issue_http_only_session_cookie() { #[tokio::test] async fn admin_login_should_reject_client_api_key_as_password_or_authorization() { let (app, _dir, _pool) = - admin_login_test_app("admin-login-rejects-client-key.sqlite", "correct-password").await; + admin_login_test_app("admin-login-rejects-client-key", "correct-password").await; let response = app .clone() @@ -87,8 +92,7 @@ async fn admin_login_should_reject_client_api_key_as_password_or_authorization() #[tokio::test] async fn admin_login_should_throttle_repeated_failures_from_same_source() { - let (app, _dir, _pool) = - admin_login_test_app("admin-login-throttle.sqlite", "correct-password").await; + let (app, _dir, _pool) = admin_login_test_app("admin-login-throttle", "correct-password").await; for _ in 0..5 { let response = post_login(&app, "wrong-password").await; @@ -106,7 +110,7 @@ async fn admin_login_should_throttle_repeated_failures_from_same_source() { #[tokio::test] async fn admin_login_should_not_write_usage_record_on_success() { let (app, _dir, pool) = - admin_login_test_app("admin-login-success-usage.sqlite", "correct-password").await; + admin_login_test_app("admin-login-success-usage", "correct-password").await; let response = post_login_with_request_id(&app, "correct-password", "req_login_usage_ok").await; assert_eq!(response.status(), StatusCode::OK); @@ -120,7 +124,7 @@ async fn admin_login_should_not_write_usage_record_on_success() { #[tokio::test] async fn admin_login_should_not_write_usage_record_on_invalid_credentials() { let (app, _dir, pool) = - admin_login_test_app("admin-login-invalid-usage.sqlite", "correct-password").await; + admin_login_test_app("admin-login-invalid-usage", "correct-password").await; let response = app .clone() @@ -148,7 +152,7 @@ async fn admin_login_should_not_write_usage_record_on_invalid_credentials() { #[tokio::test] async fn admin_login_should_not_write_usage_record_when_throttled() { let (app, _dir, pool) = - admin_login_test_app("admin-login-throttled-usage.sqlite", "correct-password").await; + admin_login_test_app("admin-login-throttled-usage", "correct-password").await; for index in 0..5 { let response = @@ -169,46 +173,23 @@ async fn admin_login_should_not_write_usage_record_when_throttled() { async fn admin_login_test_app( db_name: &str, password: &str, -) -> (axum::Router, tempfile::TempDir, SqlitePool) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); +) -> (axum::Router, tempfile::TempDir, PgPool) { + let (pool, dir) = init_test_db(db_name).await; + let redis = create_test_redis(db_name).await; seed_admin_user(&pool, password).await; - let config = test_config(url); - let stores = codex_proxy_rs::runtime::services::BackgroundTaskStores { - accounts: codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()), - admin_sessions: codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore::new( - pool.clone(), - ), - cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new(pool.clone()), - fingerprints: codex_proxy_rs::upstream::fingerprint::FingerprintRepository::new( - pool.clone(), - ), - session_affinity: - codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new( - pool.clone(), - ), - refresh_leases: codex_proxy_rs::upstream::accounts::token_refresh::RefreshLeaseStore::new( - pool.clone(), - ), - client_keys: codex_proxy_rs::admin::keys::service::SqliteClientKeyStore::new(pool.clone()), - usage_records: - codex_proxy_rs::admin::monitoring::usage_record_store::SqliteUsageRecordStore::new( - pool.clone(), - ), - }; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = std::sync::Arc::new(codex_proxy_rs::runtime::services::Services::new( + let services = std::sync::Arc::new(codex_proxy_rs::bootstrap::services::Services::new( &config, stores, runtime_fingerprint(fingerprint), )); - let state = codex_proxy_rs::runtime::state::AppState { + let state = codex_proxy_rs::bootstrap::state::AppState { services: (*services).clone(), }; ( - codex_proxy_rs::http::router::router().with_state(state), + codex_proxy_rs::api::router::router().with_state(state), dir, pool, ) @@ -237,24 +218,23 @@ async fn post_login_with_request_id( .unwrap() } -async fn usage_record_count_by_request_id(pool: &SqlitePool, request_id: &str) -> i64 { - sqlx::query_scalar("select count(*) from usage_records where request_id = ?") +async fn usage_record_count_by_request_id(pool: &PgPool, request_id: &str) -> i64 { + sqlx::query_scalar("select count(*) from usage_records where request_id = $1") .bind(request_id) .fetch_one(pool) .await .unwrap() } -async fn seed_admin_user(pool: &SqlitePool, password: &str) { - let now = Utc::now().to_rfc3339(); +async fn seed_admin_user(pool: &PgPool, password: &str) { + let now = Utc::now(); let hash = hash_admin_password(password).unwrap(); sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + "insert into admin_users (id, password_hash, created_at, updated_at) values ($1, $2, $3, $3)", ) .bind("admin_1") .bind(hash) - .bind(&now) - .bind(&now) + .bind(now) .execute(pool) .await .unwrap(); diff --git a/backend/tests/api/admin/auth/session_store.rs b/backend/tests/api/admin/auth/session_store.rs new file mode 100644 index 000000000..d0b5faa61 --- /dev/null +++ b/backend/tests/api/admin/auth/session_store.rs @@ -0,0 +1,63 @@ +use chrono::Duration; +use codex_proxy_rs::{ + auth::store::{PgAdminUserStore, RedisAdminSessionStore}, + infra::identity::hash_admin_password, +}; + +use crate::support::storage::{create_test_redis, init_test_db}; + +#[tokio::test] +async fn postgres_admin_user_store_creates_default_admin_once() { + let (pool, _guard) = init_test_db("admin-user-store").await; + let store = PgAdminUserStore::new(pool); + let password_hash = hash_admin_password("admin-password").expect("password hash"); + + let created = store + .ensure_default_admin(&password_hash) + .await + .expect("default admin should be created"); + let skipped = store + .ensure_default_admin("different-hash") + .await + .expect("second default admin should be skipped"); + let admin = store + .load_first_admin() + .await + .expect("default admin should load") + .expect("default admin should exist"); + + assert!(created); + assert!(!skipped); + assert_eq!(admin.password_hash, password_hash); +} + +#[tokio::test] +async fn redis_admin_session_store_hashes_tokens_and_uses_ttl() { + let redis = create_test_redis("admin-session-store").await; + let store = RedisAdminSessionStore::new(redis.clone()); + store + .create_session("session-secret", "admin_1", Duration::minutes(5)) + .await + .unwrap(); + + assert!(store.validate_session("session-secret").await.unwrap()); + assert!(!store.validate_session("wrong-secret").await.unwrap()); + + let mut connection = redis.manager(); + let keys: Vec = redis::cmd("KEYS") + .arg(redis.key("admin:session:*")) + .query_async(&mut connection) + .await + .unwrap(); + assert_eq!(keys.len(), 1); + assert!(!keys[0].contains("session-secret")); + let ttl: i64 = redis::cmd("TTL") + .arg(&keys[0]) + .query_async(&mut connection) + .await + .unwrap(); + assert!((1..=300).contains(&ttl)); + + assert!(store.delete_session("session-secret").await.unwrap()); + assert!(!store.validate_session("session-secret").await.unwrap()); +} diff --git a/backend/tests/admin/keys/authorization.rs b/backend/tests/api/admin/keys/authorization.rs similarity index 97% rename from backend/tests/admin/keys/authorization.rs rename to backend/tests/api/admin/keys/authorization.rs index 6efc12444..45690bf7c 100644 --- a/backend/tests/admin/keys/authorization.rs +++ b/backend/tests/api/admin/keys/authorization.rs @@ -2,7 +2,7 @@ use super::*; #[tokio::test] async fn admin_client_keys_route_should_create_list_and_authorize_v1_requests() { - let (app, _dir) = admin_client_key_test_app("admin-client-key-create.sqlite").await; + let (app, _dir) = admin_client_key_test_app("admin-client-key-create").await; let rejected = app .clone() @@ -55,6 +55,7 @@ async fn admin_client_keys_route_should_create_list_and_authorize_v1_requests() let list_body = response_json(list_response).await; assert_eq!(list_body["data"]["items"][0]["name"], "cursor"); assert_eq!(list_body["data"]["items"][0]["key"], key); + assert!(key.starts_with(list_body["data"]["items"][0]["prefix"].as_str().unwrap())); let models_response = app .oneshot( @@ -72,7 +73,7 @@ async fn admin_client_keys_route_should_create_list_and_authorize_v1_requests() #[tokio::test] async fn admin_client_keys_status_should_disable_and_enable_authorization() { - let (app, _dir) = admin_client_key_test_app("admin-client-key-status.sqlite").await; + let (app, _dir) = admin_client_key_test_app("admin-client-key-status").await; let (key_id, key) = create_admin_client_key(&app, "status-key").await; let disabled = app @@ -141,7 +142,7 @@ async fn admin_client_keys_status_should_disable_and_enable_authorization() { #[tokio::test] async fn admin_client_keys_delete_should_remove_authorization() { - let (app, _dir) = admin_client_key_test_app("admin-client-key-delete.sqlite").await; + let (app, _dir) = admin_client_key_test_app("admin-client-key-delete").await; let (key_id, key) = create_admin_client_key(&app, "delete-key").await; let deleted = app diff --git a/backend/tests/admin/keys/lifecycle.rs b/backend/tests/api/admin/keys/lifecycle.rs similarity index 98% rename from backend/tests/admin/keys/lifecycle.rs rename to backend/tests/api/admin/keys/lifecycle.rs index fa9ea4a1e..2647cf35b 100644 --- a/backend/tests/admin/keys/lifecycle.rs +++ b/backend/tests/api/admin/keys/lifecycle.rs @@ -2,7 +2,7 @@ use super::*; #[tokio::test] async fn admin_client_keys_label_should_update_clear_and_validate_label() { - let (app, _dir) = admin_client_key_test_app("admin-client-key-label.sqlite").await; + let (app, _dir) = admin_client_key_test_app("admin-client-key-label").await; let (key_id, _key) = create_admin_client_key(&app, "label-key").await; let renamed = app @@ -59,7 +59,7 @@ async fn admin_client_keys_label_should_update_clear_and_validate_label() { #[tokio::test] async fn admin_client_keys_batch_delete_should_remove_found_keys_and_report_missing_ids() { - let (app, _dir) = admin_client_key_test_app("admin-client-key-delete.sqlite").await; + let (app, _dir) = admin_client_key_test_app("admin-client-key-delete").await; let (key_a, api_key_a) = create_admin_client_key(&app, "batch-a").await; let (key_b, api_key_b) = create_admin_client_key(&app, "batch-b").await; diff --git a/backend/tests/api/admin/keys/mod.rs b/backend/tests/api/admin/keys/mod.rs new file mode 100644 index 000000000..1acf331e2 --- /dev/null +++ b/backend/tests/api/admin/keys/mod.rs @@ -0,0 +1,58 @@ +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use codex_proxy_rs::{bootstrap::services::Services, bootstrap::state::AppState}; +use serde_json::json; +use tower::util::ServiceExt; + +use crate::support::{ + admin::seed_admin_session, + config::test_config, + fingerprint::runtime_fingerprint, + http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, +}; + +mod authorization; +mod lifecycle; +mod store; + +async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, tempfile::TempDir) { + let (pool, dir) = init_test_db(db_name).await; + let redis = create_test_redis(db_name).await; + seed_admin_session(&pool, &redis, "session_1").await; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); + let fingerprint = crate::support::fingerprint::test_fingerprint(); + let services = std::sync::Arc::new(Services::new( + &config, + stores, + runtime_fingerprint(fingerprint), + )); + let state = AppState { + services: (*services).clone(), + }; + (codex_proxy_rs::api::router::router().with_state(state), dir) +} + +async fn create_admin_client_key(app: &axum::Router, name: &str) -> (String, String) { + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/keys") + .header("content-type", "application/json") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::from(json!({ "name": name }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + let body = response_json(response).await; + ( + body["data"]["id"].as_str().unwrap().to_string(), + body["data"]["key"].as_str().unwrap().to_string(), + ) +} diff --git a/backend/tests/api/admin/keys/store.rs b/backend/tests/api/admin/keys/store.rs new file mode 100644 index 000000000..bf1ca69d5 --- /dev/null +++ b/backend/tests/api/admin/keys/store.rs @@ -0,0 +1,124 @@ +use codex_proxy_rs::keys::{ + manage::KeyManageService, service::KeyVerifier, store::PgClientKeyStore, +}; + +use crate::support::storage::init_test_db; + +#[tokio::test] +async fn client_key_store_should_create_list_disable_enable_and_delete_keys() { + let (store, _dir) = client_key_store("client-keys").await; + + let created = KeyManageService::new(store.clone()) + .create("cursor") + .await + .unwrap(); + assert_eq!(created.name, "cursor"); + assert!(created.key.starts_with("sk_")); + + let stored_key: (String,) = sqlx::query_as("select key from client_api_keys where id = $1") + .bind(&created.id) + .fetch_one(store.pool()) + .await + .unwrap(); + assert_eq!(stored_key.0, created.key); + + let first_page = store.list(None, 10).await.unwrap(); + assert_eq!(first_page.items.len(), 1); + assert_eq!(first_page.items[0].name, "cursor"); + assert_eq!(first_page.items[0].prefix, created.prefix); + assert_eq!(first_page.items[0].key, created.key); + assert!(first_page.items[0].last_used_at.is_none()); + + assert!(store.set_enabled(&created.id, false).await.unwrap()); + assert!(store.set_enabled(&created.id, true).await.unwrap()); + assert!(store.delete(&created.id).await.unwrap()); +} + +#[tokio::test] +async fn client_key_service_should_verify_by_unique_key_and_defer_last_used_flush() { + let (store, _dir) = client_key_store("client-keys-runtime").await; + let created = KeyManageService::new(store.clone()) + .create("runtime") + .await + .unwrap(); + let runtime = KeyVerifier::new(store.clone()); + + assert_eq!( + runtime.verify(&created.key).await.unwrap().as_deref(), + Some(created.id.as_str()) + ); + + let before_flush = store.get(&created.id).await.unwrap().unwrap(); + assert!(before_flush.last_used_at.is_none()); + + runtime.flush_pending_last_used().await; + + let after_flush = store.get(&created.id).await.unwrap().unwrap(); + assert!(after_flush.last_used_at.is_some()); +} + +#[tokio::test] +async fn client_key_service_should_not_accept_disabled_keys() { + let (store, _dir) = client_key_store("client-keys-runtime-disabled").await; + let created = KeyManageService::new(store.clone()) + .create("runtime-disabled") + .await + .unwrap(); + assert!(store.set_enabled(&created.id, false).await.unwrap()); + let runtime = KeyVerifier::new(store); + + assert!(runtime.verify(&created.key).await.unwrap().is_none()); +} + +#[tokio::test] +async fn client_key_store_should_page_keys_by_created_at_desc() { + let (store, _dir) = client_key_store("client-keys-page").await; + + let admin = KeyManageService::new(store.clone()); + let key_a = admin.create("a").await.unwrap(); + let key_b = admin.create("b").await.unwrap(); + let key_c = admin.create("c").await.unwrap(); + sqlx::query("update client_api_keys set created_at = $1 where id = $2") + .bind(crate::support::storage::timestamp("2026-06-11T00:00:00Z")) + .bind(&key_a.id) + .execute(store.pool()) + .await + .unwrap(); + sqlx::query("update client_api_keys set created_at = $1 where id = $2") + .bind(crate::support::storage::timestamp("2026-06-11T00:01:00Z")) + .bind(&key_b.id) + .execute(store.pool()) + .await + .unwrap(); + sqlx::query("update client_api_keys set created_at = $1 where id = $2") + .bind(crate::support::storage::timestamp("2026-06-11T00:02:00Z")) + .bind(&key_c.id) + .execute(store.pool()) + .await + .unwrap(); + + let first_page = store.list(None, 2).await.unwrap(); + assert_eq!( + first_page + .items + .iter() + .map(|key| key.name.as_str()) + .collect::>(), + ["c", "b"] + ); + + let second_page = store.list(first_page.next_cursor, 2).await.unwrap(); + assert_eq!( + second_page + .items + .iter() + .map(|key| key.name.as_str()) + .collect::>(), + ["a"] + ); +} + +async fn client_key_store(db_name: &str) -> (PgClientKeyStore, tempfile::TempDir) { + let (pool, dir) = init_test_db(db_name).await; + (PgClientKeyStore::new(pool), dir) +} diff --git a/backend/tests/admin/mod.rs b/backend/tests/api/admin/mod.rs similarity index 100% rename from backend/tests/admin/mod.rs rename to backend/tests/api/admin/mod.rs diff --git a/backend/tests/admin/monitoring/dashboard.rs b/backend/tests/api/admin/monitoring/dashboard.rs similarity index 82% rename from backend/tests/admin/monitoring/dashboard.rs rename to backend/tests/api/admin/monitoring/dashboard.rs index a0345e826..84dc21a5a 100644 --- a/backend/tests/admin/monitoring/dashboard.rs +++ b/backend/tests/api/admin/monitoring/dashboard.rs @@ -4,42 +4,26 @@ use axum::{ }; use chrono::{Duration, Utc}; use codex_proxy_rs::{ - admin::{ - auth::service::SqliteAdminSessionStore, - keys::service::SqliteClientKeyStore, - monitoring::{ - usage_record_model::{UsageRecord, UsageRecordLevel}, - usage_record_store::SqliteUsageRecordStore, - }, - }, - infra::{ - database::connect_sqlite, - time::{china_datetime, china_day_start, china_hour}, - }, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::{ - services::{BackgroundTaskStores, Services}, - state::AppState, - }, - upstream::{ - accounts::{ - cookies::SqliteCookieStore, store::SqliteAccountStore, token_refresh::RefreshLeaseStore, - }, - fingerprint::{Fingerprint, FingerprintRepository}, - }, + bootstrap::{services::Services, state::AppState}, + infra::time::{china_datetime, china_day_start, china_hour}, + telemetry::{usage::store::PgUsageRecordStore, usage::types::UsageRecord}, + upstream::openai::fingerprint::Fingerprint, }; use serde_json::{json, Value}; -use sqlx::SqlitePool; +use sqlx::PgPool; use tower::util::ServiceExt; use crate::support::{ - admin::seed_admin_session, fingerprint::runtime_fingerprint, http::response_json, + admin::seed_admin_session, + fingerprint::runtime_fingerprint, + http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; #[tokio::test] async fn dashboard_summary_should_render_fingerprint_updated_at_as_beijing_time() { let (app, _store, _pool, _dir) = dashboard_test_app( - "dashboard-fingerprint-updated.sqlite", + "dashboard-fingerprint-updated", crate::support::fingerprint::test_fingerprint_with_updated_at(Some( "2026-06-23T11:36:46.965574590+00:00", )), @@ -57,7 +41,7 @@ async fn dashboard_summary_should_render_fingerprint_updated_at_as_beijing_time( #[tokio::test] async fn dashboard_summary_should_render_dash_when_fingerprint_updated_at_is_missing() { let (app, _store, _pool, _dir) = dashboard_test_app( - "dashboard-fingerprint-missing.sqlite", + "dashboard-fingerprint-missing", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -71,7 +55,7 @@ async fn dashboard_summary_should_render_dash_when_fingerprint_updated_at_is_mis async fn dashboard_summary_should_calculate_today_traffic_and_average_first_token_latency_from_time_buckets( ) { let (app, store, _pool, _dir) = dashboard_test_app( - "dashboard-traffic-latency.sqlite", + "dashboard-traffic-latency", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -86,10 +70,10 @@ async fn dashboard_summary_should_calculate_today_traffic_and_average_first_toke let mut first = usage_record_with_tokens(today_record_at, 10); first.latency_ms = Some(1_000); - first.metadata["firstTokenMs"] = json!(200u64); + first.first_token_ms = Some(200); let mut second = usage_record_with_tokens(today_record_at, 20); second.latency_ms = Some(2_000); - second.metadata["firstTokenMs"] = json!(400u64); + second.first_token_ms = Some(400); store.append(&first).await.unwrap(); store.append(&second).await.unwrap(); store @@ -126,7 +110,7 @@ async fn dashboard_summary_should_calculate_today_traffic_and_average_first_toke #[tokio::test] async fn dashboard_summary_should_keep_trend_after_usage_records_are_cleared() { let (app, store, _pool, _dir) = dashboard_test_app( - "dashboard-time-buckets-survive-usage-record-clear.sqlite", + "dashboard-time-buckets-survive-usage-record-clear", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -154,7 +138,7 @@ async fn dashboard_summary_should_keep_trend_after_usage_records_are_cleared() { #[tokio::test] async fn dashboard_trend_should_bucket_usage_by_china_hour() { let (app, store, _pool, _dir) = dashboard_test_app( - "dashboard-trend-china-hour.sqlite", + "dashboard-trend-china-hour", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -186,7 +170,7 @@ async fn dashboard_trend_should_bucket_usage_by_china_hour() { #[tokio::test] async fn dashboard_trend_should_only_include_china_today() { let (app, store, _pool, _dir) = dashboard_test_app( - "dashboard-trend-china-today.sqlite", + "dashboard-trend-china-today", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -218,19 +202,19 @@ async fn dashboard_trend_should_only_include_china_today() { #[tokio::test] async fn dashboard_latency_trend_should_use_first_token_latency() { let (app, store, _pool, _dir) = dashboard_test_app( - "dashboard-trend-first-token-latency.sqlite", + "dashboard-trend-first-token-latency", crate::support::fingerprint::test_fingerprint(), ) .await; let now = Utc::now(); let mut slow_completion = usage_record_with_tokens(now, 10); slow_completion.latency_ms = Some(10_000); - slow_completion.metadata["firstTokenMs"] = json!(100_i64); - slow_completion.metadata["serviceTier"] = json!("default"); + slow_completion.first_token_ms = Some(100); + slow_completion.service_tier = Some("default".to_string()); let mut fast_completion = usage_record_with_tokens(now, 10); fast_completion.latency_ms = Some(1_000); - fast_completion.metadata["firstTokenMs"] = json!(500_i64); - fast_completion.metadata["serviceTier"] = json!("priority"); + fast_completion.first_token_ms = Some(500); + fast_completion.service_tier = Some("priority".to_string()); store.append(&slow_completion).await.unwrap(); store.append(&fast_completion).await.unwrap(); @@ -258,14 +242,14 @@ async fn dashboard_latency_trend_should_use_first_token_latency() { #[tokio::test] async fn dashboard_summary_should_return_requested_trend_kind() { let (app, store, _pool, _dir) = dashboard_test_app( - "dashboard-summary-requested-trend-kind.sqlite", + "dashboard-summary-requested-trend-kind", crate::support::fingerprint::test_fingerprint(), ) .await; let now = Utc::now(); let mut record = usage_record_with_tokens(now - Duration::seconds(1), 10); record.latency_ms = Some(10_000); - record.metadata["firstTokenMs"] = json!(250_i64); + record.first_token_ms = Some(250); store.append(&record).await.unwrap(); let body = dashboard_summary_with_kind(app, "latency").await; @@ -284,7 +268,7 @@ async fn dashboard_summary_should_return_requested_trend_kind() { #[tokio::test] async fn dashboard_summary_should_calculate_priority_cost_from_event_service_tier() { let cost = dashboard_summary_today_cost_for_usage_record( - "dashboard-cost-priority.sqlite", + "dashboard-cost-priority", "gpt-5.5", json!({ "usage": { @@ -303,7 +287,7 @@ async fn dashboard_summary_should_calculate_priority_cost_from_event_service_tie #[tokio::test] async fn dashboard_summary_should_charge_cached_input_at_cache_read_price() { let cost = dashboard_summary_today_cost_for_usage_record( - "dashboard-cost-cached.sqlite", + "dashboard-cost-cached", "gpt-5.5", json!({ "usage": { @@ -321,7 +305,7 @@ async fn dashboard_summary_should_charge_cached_input_at_cache_read_price() { #[tokio::test] async fn dashboard_summary_should_apply_long_context_prices() { let cost = dashboard_summary_today_cost_for_usage_record( - "dashboard-cost-long-context.sqlite", + "dashboard-cost-long-context", "gpt-5.5", json!({ "usage": { @@ -339,7 +323,7 @@ async fn dashboard_summary_should_apply_long_context_prices() { #[tokio::test] async fn dashboard_summary_should_normalize_codex_model_names_for_cost() { let cost = dashboard_summary_today_cost_for_usage_record( - "dashboard-cost-codex-model.sqlite", + "dashboard-cost-codex-model", "codex-auto-review", json!({ "usage": { @@ -357,7 +341,7 @@ async fn dashboard_summary_should_normalize_codex_model_names_for_cost() { #[tokio::test] async fn dashboard_summary_should_not_price_usage_summary_without_model_dimension() { let (app, _store, pool, _dir) = dashboard_test_app( - "dashboard-cost-summary-fallback.sqlite", + "dashboard-cost-summary-fallback", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -371,7 +355,7 @@ async fn dashboard_summary_should_not_price_usage_summary_without_model_dimensio #[tokio::test] async fn dashboard_summary_should_return_backend_formatted_time_fields() { let (app, store, _pool, _dir) = dashboard_test_app( - "dashboard-backend-formatted-time.sqlite", + "dashboard-backend-formatted-time", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -393,7 +377,7 @@ async fn dashboard_summary_should_return_backend_formatted_time_fields() { #[tokio::test] async fn dashboard_summary_should_only_return_today_account_usage_and_usage_records() { let (app, store, pool, _dir) = dashboard_test_app( - "dashboard-summary-today-records.sqlite", + "dashboard-summary-today-records", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -442,7 +426,7 @@ async fn dashboard_summary_should_only_return_today_account_usage_and_usage_reco #[tokio::test] async fn dashboard_summary_should_order_account_usage_by_recent_usage() { let (app, store, pool, _dir) = dashboard_test_app( - "dashboard-account-usage-recent-order.sqlite", + "dashboard-account-usage-recent-order", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -484,7 +468,7 @@ async fn dashboard_summary_should_order_account_usage_by_recent_usage() { #[tokio::test] async fn dashboard_summary_should_use_five_hour_quota_window_for_account_usage() { let (app, store, pool, _dir) = dashboard_test_app( - "dashboard-account-usage-five-hour-quota.sqlite", + "dashboard-account-usage-five-hour-quota", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -535,7 +519,7 @@ async fn dashboard_summary_should_use_five_hour_quota_window_for_account_usage() #[tokio::test] async fn dashboard_summary_should_fall_back_to_weekly_quota_window_when_five_hour_missing() { let (app, store, pool, _dir) = dashboard_test_app( - "dashboard-account-usage-weekly-quota.sqlite", + "dashboard-account-usage-weekly-quota", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -586,7 +570,7 @@ async fn dashboard_summary_should_fall_back_to_weekly_quota_window_when_five_hou #[tokio::test] async fn dashboard_summary_should_return_today_health_timeline() { let (app, store, _pool, _dir) = dashboard_test_app( - "dashboard-health-timeline.sqlite", + "dashboard-health-timeline", crate::support::fingerprint::test_fingerprint(), ) .await; @@ -618,10 +602,24 @@ async fn dashboard_summary_today_cost_for_usage_record( ) -> String { let (app, store, _pool, _dir) = dashboard_test_app(db_name, crate::support::fingerprint::test_fingerprint()).await; - store - .append(&usage_record(Utc::now(), model, metadata)) - .await - .unwrap(); + let mut record = usage_record(Utc::now(), model, metadata.clone()); + record.input_tokens = metadata + .pointer("/usage/inputTokens") + .and_then(Value::as_u64) + .map(|value| value as i64); + record.output_tokens = metadata + .pointer("/usage/outputTokens") + .and_then(Value::as_u64) + .map(|value| value as i64); + record.cached_tokens = metadata + .pointer("/usage/cachedTokens") + .and_then(Value::as_u64) + .map(|value| value as i64); + record.service_tier = metadata + .get("serviceTier") + .and_then(Value::as_str) + .map(str::to_string); + store.append(&record).await.unwrap(); let body = dashboard_summary(app).await; body["data"]["cards"]["tokens"]["todayCostUsd"] @@ -633,28 +631,12 @@ async fn dashboard_summary_today_cost_for_usage_record( async fn dashboard_test_app( db_name: &str, fingerprint: Fingerprint, -) -> ( - axum::Router, - SqliteUsageRecordStore, - SqlitePool, - tempfile::TempDir, -) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let config = crate::support::config::test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; +) -> (axum::Router, PgUsageRecordStore, PgPool, tempfile::TempDir) { + let (pool, dir) = init_test_db(db_name).await; + let redis = create_test_redis(db_name).await; + seed_admin_session(&pool, &redis, "session_1").await; + let config = crate::support::config::test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let services = std::sync::Arc::new(Services::new( &config, stores, @@ -664,8 +646,8 @@ async fn dashboard_test_app( services: (*services).clone(), }; ( - codex_proxy_rs::http::router::router().with_state(state), - SqliteUsageRecordStore::new(pool.clone()), + codex_proxy_rs::api::router::router().with_state(state), + PgUsageRecordStore::new(pool.clone()), pool, dir, ) @@ -713,46 +695,46 @@ async fn dashboard_trend(app: axum::Router, kind: &str) -> Value { response_json(response).await } -async fn seed_usage_summary(pool: &SqlitePool, input_tokens: u64, output_tokens: u64) { +async fn seed_usage_summary(pool: &PgPool, input_tokens: u64, output_tokens: u64) { seed_usage_summary_with_last_used(pool, input_tokens, output_tokens, "2026-06-18T00:00:00Z") .await; } async fn seed_usage_summary_with_last_used( - pool: &SqlitePool, + pool: &PgPool, input_tokens: u64, output_tokens: u64, last_used_at: &str, ) { sqlx::query( - "insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)", + "insert into accounts (id, email, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6)", ) .bind("acct_usage") .bind("acct-usage@example.com") .bind("cipher") .bind("active") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) .execute(pool) .await .unwrap(); sqlx::query( - "insert into account_usage (account_id, request_count, input_tokens, output_tokens, total_tokens, last_used_at) values (?, ?, ?, ?, ?, ?)", + "insert into account_usage (account_id, request_count, input_tokens, output_tokens, total_tokens, last_used_at) values ($1, $2, $3, $4, $5, $6)", ) .bind("acct_usage") .bind(1_i64) .bind(input_tokens as i64) .bind(output_tokens as i64) .bind((input_tokens + output_tokens) as i64) - .bind(last_used_at) + .bind(crate::support::storage::timestamp(last_used_at)) .execute(pool) .await .unwrap(); } -async fn seed_dashboard_account(pool: &SqlitePool, account_id: &str, email: &str) { +async fn seed_dashboard_account(pool: &PgPool, account_id: &str, email: &str) { sqlx::query( - "insert into accounts (id, email, chatgpt_account_id, plan_type, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)", + "insert into accounts (id, email, chatgpt_account_id, plan_type, access_token, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7, $8)", ) .bind(account_id) .bind(email) @@ -760,17 +742,17 @@ async fn seed_dashboard_account(pool: &SqlitePool, account_id: &str, email: &str .bind("K12") .bind("cipher") .bind("active") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) .execute(pool) .await .unwrap(); } -async fn set_account_quota_json(pool: &SqlitePool, account_id: &str, quota_json: Value) { - sqlx::query("update accounts set quota_json = ?, quota_fetched_at = ? where id = ?") +async fn set_account_quota_json(pool: &PgPool, account_id: &str, quota_json: Value) { + sqlx::query("update accounts set quota_json = $1::jsonb, quota_fetched_at = $2 where id = $3") .bind(quota_json.to_string()) - .bind("2026-06-30T08:01:00Z") + .bind(crate::support::storage::timestamp("2026-06-30T08:01:00Z")) .bind(account_id) .execute(pool) .await @@ -797,17 +779,11 @@ fn usage_record_with_token_parts( output_tokens: u64, cached_tokens: u64, ) -> UsageRecord { - usage_record( - created_at, - "gpt-5.5", - json!({ - "usage": { - "inputTokens": input_tokens, - "outputTokens": output_tokens, - "cachedTokens": cached_tokens - } - }), - ) + let mut record = usage_record(created_at, "gpt-5.5", json!({})); + record.input_tokens = Some(input_tokens as i64); + record.output_tokens = Some(output_tokens as i64); + record.cached_tokens = Some(cached_tokens as i64); + record } fn usage_record_with_account_tokens( @@ -816,14 +792,12 @@ fn usage_record_with_account_tokens( total_tokens: u64, ) -> UsageRecord { let mut record = usage_record_with_tokens(created_at, total_tokens); - record.account_id = Some(account_id.to_string()); + record.account_id = account_id.to_string(); record } fn usage_record(created_at: chrono::DateTime, model: &str, metadata: Value) -> UsageRecord { - let mut record = UsageRecord::new("v1.response", UsageRecordLevel::Info, "completed"); - record.model = Some(model.to_string()); - record.status_code = Some(200); + let mut record = UsageRecord::new("v1.response", "completed", "acct_dashboard", model, 200); record.created_at = created_at; record.metadata = metadata; record diff --git a/backend/tests/admin/monitoring/mod.rs b/backend/tests/api/admin/monitoring/mod.rs similarity index 75% rename from backend/tests/admin/monitoring/mod.rs rename to backend/tests/api/admin/monitoring/mod.rs index 23a82b62d..22de2ac66 100644 --- a/backend/tests/admin/monitoring/mod.rs +++ b/backend/tests/api/admin/monitoring/mod.rs @@ -1,3 +1,4 @@ pub(crate) mod dashboard; +pub(crate) mod ops_error_store; pub(crate) mod usage_record_routes; pub(crate) mod usage_record_store; diff --git a/backend/tests/api/admin/monitoring/ops_error_store.rs b/backend/tests/api/admin/monitoring/ops_error_store.rs new file mode 100644 index 000000000..d00a4cd9a --- /dev/null +++ b/backend/tests/api/admin/monitoring/ops_error_store.rs @@ -0,0 +1,61 @@ +use codex_proxy_rs::telemetry::{ + ops::store::{OpsErrorFilter, PgOpsErrorLogStore}, + ops::types::OpsErrorLog, +}; + +use crate::support::storage::init_test_db; + +#[tokio::test] +async fn ops_store_filters_diagnostic_dimensions() { + let (pool, _guard) = init_test_db("ops-error-filter").await; + let store = PgOpsErrorLogStore::new(pool); + let mut error = OpsErrorLog::new("upstream", "rate limited"); + error.id = "ops_error_matching".to_string(); + error.request_id = Some("req_error".to_string()); + error.client_api_key_id = Some("key_1".to_string()); + error.failure_class = Some("rate_limited".to_string()); + error.upstream_status_code = Some(429); + error.transport = Some("http_sse".to_string()); + store.append(&error).await.unwrap(); + + let page = store + .list_page( + OpsErrorFilter { + client_api_key_id: Some("key_1".to_string()), + failure_class: Some("rate_limited".to_string()), + upstream_status_code: Some(429), + search: Some("req_error".to_string()), + ..OpsErrorFilter::default() + }, + 1, + 20, + ) + .await + .unwrap(); + assert_eq!(page.total, 1); + assert_eq!(page.items[0].id, "ops_error_matching"); +} + +#[tokio::test] +async fn ops_error_and_bucket_are_one_transaction() { + let (pool, _guard) = init_test_db("ops-error-transaction").await; + let store = PgOpsErrorLogStore::new(pool.clone()); + sqlx::raw_sql( + "create function reject_error_bucket() returns trigger language plpgsql as $$ + begin raise exception 'bucket rejected'; end $$; + create trigger reject_error_bucket before insert or update on request_time_buckets + for each row execute function reject_error_bucket();", + ) + .execute(&pool) + .await + .unwrap(); + let mut error = OpsErrorLog::new("upstream", "must rollback"); + error.id = "ops_error_rolled_back".to_string(); + assert!(store.append(&error).await.is_err()); + let count: i64 = sqlx::query_scalar("select count(*) from ops_error_logs where id = $1") + .bind(&error.id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(count, 0); +} diff --git a/backend/tests/api/admin/monitoring/usage_record_routes.rs b/backend/tests/api/admin/monitoring/usage_record_routes.rs new file mode 100644 index 000000000..856c0ed4c --- /dev/null +++ b/backend/tests/api/admin/monitoring/usage_record_routes.rs @@ -0,0 +1,165 @@ +use axum::{ + body::Body, + http::{Request, StatusCode}, +}; +use codex_proxy_rs::{ + bootstrap::{services::Services, state::AppState}, + telemetry::{ + ops::store::PgOpsErrorLogStore, ops::types::OpsErrorLog, usage::store::PgUsageRecordStore, + usage::types::UsageRecord, + }, +}; +use serde_json::{json, Value}; +use tower::util::ServiceExt; + +use crate::support::{ + admin::seed_admin_session, + config::test_config, + fingerprint::runtime_fingerprint, + http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, +}; + +#[tokio::test] +async fn usage_and_ops_routes_require_admin_auth() { + let (app, _, _, _guard) = monitoring_test_app("monitoring-auth", false).await; + for uri in ["/api/admin/usage/records", "/api/admin/ops/errors"] { + let response = app + .clone() + .oneshot(Request::builder().uri(uri).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + } +} + +#[tokio::test] +async fn usage_route_returns_success_facts_without_legacy_level() { + let (app, usage, _, _guard) = monitoring_test_app("monitoring-success", true).await; + let mut record = success_record("completed"); + record.id = "usage_route_success".to_string(); + record.client_api_key_id = Some("key_42".to_string()); + record.input_tokens = Some(12); + record.output_tokens = Some(4); + record.latency_ms = Some(0); + usage.append(&record).await.unwrap(); + + let response = app + .oneshot( + Request::builder() + .uri("/api/admin/usage/records?page=1&pageSize=20&clientApiKeyId=key_42") + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + let item = &body["data"]["items"][0]; + assert_eq!(item["id"], "usage_route_success"); + assert_eq!(item["clientApiKeyId"], "key_42"); + assert_eq!(item["statusCode"], 200); + assert!(item.get("level").is_none()); +} + +#[tokio::test] +async fn usage_summary_and_trend_expose_only_real_success_metrics() { + let (app, usage, _, _guard) = monitoring_test_app("monitoring-summary", true).await; + let mut record = success_record("summary"); + record.input_tokens = Some(100); + record.output_tokens = Some(20); + record.cached_tokens = Some(30); + record.latency_ms = Some(10); + usage.append(&record).await.unwrap(); + + let summary = admin_get(&app, "/api/admin/usage/records/summary").await; + assert_eq!(summary["data"]["totalRequests"], "1"); + assert_eq!(summary["data"]["totalTokens"], "120"); + assert!(summary["data"].get("errorRequests").is_none()); + assert!(summary["data"].get("errorRate").is_none()); + + let trend = admin_get(&app, "/api/admin/usage/records/insights/token-trend").await; + assert_eq!(trend["data"][0]["inputTokensValue"], 100); + assert!(trend["data"][0].get("cacheCreationTokens").is_none()); +} + +#[tokio::test] +async fn ops_errors_route_filters_failure_facts_independently() { + let (app, usage, ops, _guard) = monitoring_test_app("monitoring-errors", true).await; + usage.append(&success_record("success only")).await.unwrap(); + let mut error = OpsErrorLog::new("upstream", "rate limited"); + error.id = "ops_route_error".to_string(); + error.request_id = Some("req_ops".to_string()); + error.client_api_key_id = Some("key_ops".to_string()); + error.provider = Some("openai".to_string()); + error.failure_class = Some("rate_limited".to_string()); + error.status_code = Some(429); + error.upstream_status_code = Some(429); + error.metadata = json!({"retryAfter": 5}); + ops.append(&error).await.unwrap(); + + let body = admin_get( + &app, + "/api/admin/ops/errors?page=1&pageSize=20&failureClass=rate_limited&search=req_ops", + ) + .await; + assert_eq!(body["data"]["page"]["total"], 1); + assert_eq!(body["data"]["items"][0]["id"], "ops_route_error"); + assert_eq!(body["data"]["items"][0]["upstreamStatusCode"], 429); + + let success = admin_get(&app, "/api/admin/usage/records?page=1&pageSize=20").await; + assert_eq!(success["data"]["page"]["total"], 1); + assert_eq!(success["data"]["items"][0]["message"], "success only"); +} + +async fn monitoring_test_app( + label: &str, + authenticated: bool, +) -> ( + axum::Router, + PgUsageRecordStore, + PgOpsErrorLogStore, + tempfile::TempDir, +) { + let (pool, guard) = init_test_db(label).await; + let redis = create_test_redis(label).await; + if authenticated { + seed_admin_session(&pool, &redis, "session_1").await; + } + let mut config = test_config(test_database_url()); + config.logging.enabled = true; + let stores = background_task_stores(pool.clone(), redis); + let services = Services::new( + &config, + stores, + runtime_fingerprint(crate::support::fingerprint::test_fingerprint()), + ); + let app = codex_proxy_rs::api::router::router().with_state(AppState { services }); + ( + app, + PgUsageRecordStore::new(pool.clone()), + PgOpsErrorLogStore::new(pool), + guard, + ) +} + +async fn admin_get(app: &axum::Router, uri: &str) -> Value { + let response = app + .clone() + .oneshot( + Request::builder() + .uri(uri) + .header("cookie", "cpr_admin_session=session_1") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + response_json(response).await +} + +fn success_record(message: &str) -> UsageRecord { + UsageRecord::new("request", message, "acct_usage", "gpt-5", 200) +} diff --git a/backend/tests/api/admin/monitoring/usage_record_store.rs b/backend/tests/api/admin/monitoring/usage_record_store.rs new file mode 100644 index 000000000..96fe81c6a --- /dev/null +++ b/backend/tests/api/admin/monitoring/usage_record_store.rs @@ -0,0 +1,115 @@ +use chrono::{Duration, Utc}; +use codex_proxy_rs::telemetry::{ + usage::store::{PgUsageRecordStore, UsageRecordFilter}, + usage::types::UsageRecord, +}; + +use crate::support::storage::init_test_db; + +#[tokio::test] +async fn usage_store_filters_and_cursor_pages_success_facts() { + let (pool, _guard) = init_test_db("usage-record-filter").await; + let store = PgUsageRecordStore::new(pool); + let mut matching = success_record("matching"); + matching.id = "usage_matching".to_string(); + matching.request_id = Some("req_1".to_string()); + matching.route = Some("/v1/responses".to_string()); + matching.client_api_key_id = Some("key_1".to_string()); + store.append(&matching).await.unwrap(); + store.append(&success_record("other")).await.unwrap(); + + let page = store + .list( + UsageRecordFilter { + kind: Some("request".to_string()), + client_api_key_id: Some("key_1".to_string()), + search: Some("matching".to_string()), + ..UsageRecordFilter::default() + }, + None, + 1, + ) + .await + .unwrap(); + + assert_eq!(page.items.len(), 1); + assert_eq!(page.items[0].id, "usage_matching"); + assert_eq!(page.items[0].request_id.as_deref(), Some("req_1")); +} + +#[tokio::test] +async fn usage_record_and_bucket_are_one_transaction_and_keep_zero_ms_samples() { + let (pool, _guard) = init_test_db("usage-record-transaction").await; + let store = PgUsageRecordStore::new(pool.clone()); + let mut record = success_record("zero latency"); + record.id = "usage_zero_ms".to_string(); + record.latency_ms = Some(0); + record.first_token_ms = Some(0); + record.input_tokens = Some(11); + record.output_tokens = Some(7); + store.append(&record).await.unwrap(); + + let bucket: (i64, i64, i64, i64, i64) = sqlx::query_as( + "select success_count, input_tokens, output_tokens, latency_count, + first_token_latency_count + from request_time_buckets where account_id = $1 and model = $2", + ) + .bind(&record.account_id) + .bind(&record.model) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(bucket, (1, 11, 7, 1, 1)); + + sqlx::raw_sql( + "create function reject_test_bucket() returns trigger language plpgsql as $$ + begin raise exception 'bucket rejected'; end $$; + create trigger reject_test_bucket before insert or update on request_time_buckets + for each row execute function reject_test_bucket();", + ) + .execute(&pool) + .await + .unwrap(); + let mut rejected = success_record("must rollback"); + rejected.id = "usage_rolled_back".to_string(); + assert!(store.append(&rejected).await.is_err()); + let count: i64 = sqlx::query_scalar("select count(*) from usage_records where id = $1") + .bind(&rejected.id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(count, 0); +} + +#[tokio::test] +async fn usage_store_trims_by_runtime_retention() { + let (pool, _guard) = init_test_db("usage-record-retention").await; + sqlx::query( + "insert into runtime_settings ( + id, refresh_margin_seconds, refresh_concurrency, + max_concurrent_per_account, request_interval_ms, + rotation_strategy, updated_at + ) values (1, 3600, 2, 1, 0, 'smart', now())", + ) + .execute(&pool) + .await + .unwrap(); + let store = PgUsageRecordStore::new(pool); + let now = Utc::now(); + let mut expired = success_record("expired"); + expired.id = "usage_expired".to_string(); + expired.created_at = now - Duration::days(31); + store.append(&expired).await.unwrap(); + let mut retained = success_record("retained"); + retained.id = "usage_retained".to_string(); + retained.created_at = now - Duration::days(29); + store.append(&retained).await.unwrap(); + + assert_eq!(store.trim_to_retention(now).await.unwrap(), 1); + assert!(store.get("usage_expired").await.unwrap().is_none()); + assert!(store.get("usage_retained").await.unwrap().is_some()); +} + +fn success_record(message: &str) -> UsageRecord { + UsageRecord::new("request", message, "acct_usage", "gpt-5", 200) +} diff --git a/backend/tests/admin/response/mod.rs b/backend/tests/api/admin/response/mod.rs similarity index 98% rename from backend/tests/admin/response/mod.rs rename to backend/tests/api/admin/response/mod.rs index ac4838976..b33f987f3 100644 --- a/backend/tests/admin/response/mod.rs +++ b/backend/tests/api/admin/response/mod.rs @@ -3,7 +3,7 @@ use axum::{ http::StatusCode, response::{IntoResponse, Response}, }; -use codex_proxy_rs::admin::response::{ +use codex_proxy_rs::api::admin::response::{ AdminEnvelope, AdminError, AdminPageEnvelope, AdminResponse, }; use codex_proxy_rs::infra::json::{NumberedPage, Page}; diff --git a/backend/tests/admin/settings/mod.rs b/backend/tests/api/admin/settings/mod.rs similarity index 81% rename from backend/tests/admin/settings/mod.rs rename to backend/tests/api/admin/settings/mod.rs index daac741e7..fc2e9fe1d 100644 --- a/backend/tests/admin/settings/mod.rs +++ b/backend/tests/api/admin/settings/mod.rs @@ -1,19 +1,19 @@ pub(crate) mod routes; -use codex_proxy_rs::config::settings::{AdminSettings, AdminSettingsPatch, SettingsService}; +use codex_proxy_rs::settings::{SettingsPatch, SettingsService, SettingsSnapshot}; #[test] fn settings_service_should_apply_retained_settings_patch() { - let mut current = AdminSettings { + let mut current = SettingsSnapshot { max_concurrent_per_account: 4, request_interval_ms: 50, rotation_strategy: "smart".to_string(), - ..AdminSettings::default() + ..SettingsSnapshot::default() }; SettingsService::apply_patch( &mut current, - AdminSettingsPatch { + SettingsPatch { model_aliases: Some( [ (" claude-sonnet ".to_string(), " gpt-5.5 ".to_string()), @@ -41,12 +41,12 @@ fn settings_service_should_apply_retained_settings_patch() { #[test] fn settings_service_should_reject_invalid_patch_values() { - let mut current = AdminSettings::default(); + let mut current = SettingsSnapshot::default(); let error = SettingsService::apply_patch( &mut current, - AdminSettingsPatch { + SettingsPatch { rotation_strategy: Some("random".to_string()), - ..AdminSettingsPatch::default() + ..SettingsPatch::default() }, ) .expect_err("invalid rotation strategy should be rejected"); @@ -56,12 +56,12 @@ fn settings_service_should_reject_invalid_patch_values() { #[test] fn settings_service_should_reject_invalid_model_aliases() { - let mut current = AdminSettings::default(); + let mut current = SettingsSnapshot::default(); let error = SettingsService::apply_patch( &mut current, - AdminSettingsPatch { + SettingsPatch { model_aliases: Some([("gpt-5.5".to_string(), "gpt-5.5".to_string())].into()), - ..AdminSettingsPatch::default() + ..SettingsPatch::default() }, ) .expect_err("self-referencing model alias should be rejected"); diff --git a/backend/tests/admin/settings/routes.rs b/backend/tests/api/admin/settings/routes.rs similarity index 74% rename from backend/tests/admin/settings/routes.rs rename to backend/tests/api/admin/settings/routes.rs index 4a6db9b7a..b10c06d9b 100644 --- a/backend/tests/admin/settings/routes.rs +++ b/backend/tests/api/admin/settings/routes.rs @@ -9,32 +9,29 @@ use axum::{ }; use chrono::Utc; use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::usage_record_store::SqliteUsageRecordStore, - config::{schema::AppConfig, settings::RuntimeSettingsService}, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::{ - cookies::SqliteCookieStore, pool::AccountAcquireRequest, store::SqliteAccountStore, - token_refresh::RefreshLeaseStore, + accounts::pool::AccountAcquireRequest, + bootstrap::{ + config::AppConfig, + services::{apply_settings_to_config, settings_snapshot_from_config, Services}, + state::AppState, }, - upstream::fingerprint::FingerprintRepository, + settings::service::RuntimeSettingsService, }; use serde_json::json; -use sqlx::SqlitePool; +use sqlx::PgPool; use tower::util::ServiceExt; use crate::support::{ - admin::seed_admin_session, config::test_config as base_test_config, - fingerprint::runtime_fingerprint, http::response_json, + admin::seed_admin_session, + config::test_config as base_test_config, + fingerprint::runtime_fingerprint, + http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; #[tokio::test] async fn admin_settings_should_require_admin_session_cookie() { - let (app, _dir) = admin_settings_test_app("admin-settings-auth.sqlite").await; + let (app, _dir) = admin_settings_test_app("admin-settings-auth").await; let response = app .oneshot( Request::builder() @@ -52,7 +49,7 @@ async fn admin_settings_should_require_admin_session_cookie() { #[tokio::test] async fn admin_settings_should_return_runtime_fields() { - let (app, _dir) = admin_settings_test_app("admin-settings-get.sqlite").await; + let (app, _dir) = admin_settings_test_app("admin-settings-get").await; let response = app .oneshot( Request::builder() @@ -74,7 +71,7 @@ async fn admin_settings_should_return_runtime_fields() { #[tokio::test] async fn admin_settings_admin_api_key_status_should_be_empty_by_default() { - let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-empty.sqlite").await; + let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-empty").await; let response = app .oneshot( Request::builder() @@ -94,8 +91,8 @@ async fn admin_settings_admin_api_key_status_should_be_empty_by_default() { } #[tokio::test] -async fn admin_settings_admin_api_key_should_regenerate_and_return_masked_status() { - let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-regenerate.sqlite").await; +async fn admin_settings_admin_api_key_should_regenerate_and_return_exists_only() { + let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-regenerate").await; let response = app .clone() .oneshot( @@ -130,15 +127,12 @@ async fn admin_settings_admin_api_key_should_regenerate_and_return_masked_status assert_eq!(response.status(), StatusCode::OK); let body = response_json(response).await; assert_eq!(body["data"]["exists"], true); - assert_eq!( - body["data"]["maskedKey"], - format!("{}...{}", &key[..10], &key[key.len() - 4..]) - ); + assert!(body["data"]["maskedKey"].is_null()); } #[tokio::test] async fn admin_settings_should_accept_valid_admin_api_key_header() { - let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-auth.sqlite").await; + let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-auth").await; let response = app .clone() .oneshot( @@ -174,7 +168,7 @@ async fn admin_settings_should_accept_valid_admin_api_key_header() { #[tokio::test] async fn admin_settings_should_reject_invalid_admin_api_key_header() { - let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-invalid.sqlite").await; + let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-invalid").await; let response = app .oneshot( Request::builder() @@ -193,7 +187,7 @@ async fn admin_settings_should_reject_invalid_admin_api_key_header() { #[tokio::test] async fn admin_settings_admin_api_key_should_delete_and_revoke_access() { - let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-delete.sqlite").await; + let (app, _dir) = admin_settings_test_app("admin-settings-admin-key-delete").await; let response = app .clone() .oneshot( @@ -243,7 +237,7 @@ async fn admin_settings_admin_api_key_should_delete_and_revoke_access() { #[tokio::test] async fn admin_settings_update_should_require_admin_session_cookie() { - let (app, _dir) = admin_settings_test_app("admin-settings-update-auth.sqlite").await; + let (app, _dir) = admin_settings_test_app("admin-settings-update-auth").await; let response = app .oneshot( Request::builder() @@ -261,29 +255,18 @@ async fn admin_settings_update_should_require_admin_session_cookie() { #[tokio::test] async fn admin_settings_update_should_persist_runtime_settings_to_database() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-settings-update.sqlite"); - let url = format!("sqlite://{}", db.display()); - let config = test_config(url); - let pool = connect_sqlite(&config.database.url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let (pool, dir) = init_test_db("admin-settings-update").await; + let redis = create_test_redis("admin-settings-update").await; + let config = test_config(test_database_url()); + seed_admin_session(&pool, &redis, "session_1").await; + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Services::new(&config, stores, runtime_fingerprint(fingerprint)); let services = std::sync::Arc::new(services); let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); let response = app .clone() @@ -334,13 +317,13 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() { "round_robin" ); - let row: (String, i64, i64, i64, i64, String) = sqlx::query_as( + let row: (serde_json::Value, i64, i64, i64, i64, String) = sqlx::query_as( "select model_aliases_json, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy from runtime_settings where id = 1", ) .fetch_one(&pool) .await .unwrap(); - let aliases: BTreeMap = serde_json::from_str(&row.0).unwrap(); + let aliases: BTreeMap = serde_json::from_value(row.0).unwrap(); assert_eq!(aliases["gpt-5.2"], "gpt-5.5"); assert_eq!(aliases["claude-sonnet"], "gpt-5.5"); assert_eq!(row.1, 900); @@ -350,9 +333,12 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() { assert_eq!(row.5, "round_robin"); assert!(!dir.path().join("config.yaml").exists()); - let restarted_config = RuntimeSettingsService::load_or_initialize_config(config.clone(), &pool) - .await - .unwrap(); + let restarted_settings = + RuntimeSettingsService::load_or_initialize(settings_snapshot_from_config(&config), &pool) + .await + .unwrap(); + let mut restarted_config = config.clone(); + apply_settings_to_config(&mut restarted_config, &restarted_settings); assert_eq!(restarted_config.model_aliases["gpt-5.2"], "gpt-5.5"); assert_eq!(restarted_config.auth.refresh_margin_seconds, 900); assert_eq!(restarted_config.auth.refresh_concurrency, 4); @@ -364,39 +350,24 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() { #[tokio::test] async fn admin_settings_update_should_apply_runtime_services() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-settings-runtime.sqlite"); - let url = format!("sqlite://{}", db.display()); - let config = test_config(url); - let pool = connect_sqlite(&config.database.url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; + let (pool, _dir) = init_test_db("admin-settings-runtime").await; + let redis = create_test_redis("admin-settings-runtime").await; + let config = test_config(test_database_url()); + seed_admin_session(&pool, &redis, "session_1").await; seed_schedulable_account(&pool, "acct_runtime_a").await; seed_schedulable_account(&pool, "acct_runtime_b").await; - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new( &config, stores, runtime_fingerprint(fingerprint), )); - services - .account_pool - .restore_from_repository() - .await - .unwrap(); + services.account_pool.restore_from_store().await.unwrap(); let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); let response = app .oneshot( @@ -468,22 +439,11 @@ async fn admin_settings_update_should_apply_runtime_services() { #[tokio::test] async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("admin-settings-update-invalid.sqlite"); - let url = format!("sqlite://{}", db.display()); - let config = test_config(url); - let pool = connect_sqlite(&config.database.url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let (pool, _dir) = init_test_db("admin-settings-update-invalid").await; + let redis = create_test_redis("admin-settings-update-invalid").await; + let config = test_config(test_database_url()); + seed_admin_session(&pool, &redis, "session_1").await; + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new( &config, @@ -493,7 +453,7 @@ async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() { let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); let response = app .oneshot( @@ -514,22 +474,11 @@ async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() { } async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let (pool, dir) = init_test_db(db_name).await; + let redis = create_test_redis(db_name).await; + seed_admin_session(&pool, &redis, "session_1").await; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new( &config, @@ -539,13 +488,10 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp let state = AppState { services: (*services).clone(), }; - ( - codex_proxy_rs::http::router::router().with_state(state), - dir, - ) + (codex_proxy_rs::api::router::router().with_state(state), dir) } -async fn seed_schedulable_account(pool: &SqlitePool, account_id: &str) { +async fn seed_schedulable_account(pool: &PgPool, account_id: &str) { sqlx::query( r" insert into accounts ( @@ -555,13 +501,13 @@ insert into accounts ( status, added_at, updated_at -) values (?, ?, ?, 'active', ?, ?)", +) values ($1, $2, $3, 'active', $4, $5)", ) .bind(account_id) .bind("access-token") - .bind("2099-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") + .bind(crate::support::storage::timestamp("2099-01-01T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) .execute(pool) .await .unwrap(); diff --git a/backend/tests/admin/system/mod.rs b/backend/tests/api/admin/system/mod.rs similarity index 94% rename from backend/tests/admin/system/mod.rs rename to backend/tests/api/admin/system/mod.rs index 29d528989..711ed01bf 100644 --- a/backend/tests/admin/system/mod.rs +++ b/backend/tests/api/admin/system/mod.rs @@ -9,23 +9,10 @@ use axum::{ body::{to_bytes, Body}, http::{header::CONTENT_TYPE, Request, StatusCode}, }; -use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::usage_record_store::SqliteUsageRecordStore, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::{ - cookies::SqliteCookieStore, store::SqliteAccountStore, token_refresh::RefreshLeaseStore, - }, - upstream::fingerprint::FingerprintRepository, -}; +use codex_proxy_rs::{bootstrap::services::Services, bootstrap::state::AppState}; use flate2::{write::GzEncoder, Compression}; use serde_json::{json, Value}; use sha2::{Digest, Sha256}; -use sqlx::SqlitePool; use tar::{Builder, EntryType, Header}; use tokio::sync::Mutex; use tower::util::ServiceExt; @@ -54,7 +41,7 @@ async fn version_should_return_backend_build_metadata() { set_system_update_env(repository, &github.uri()); std::env::set_var("CPR_GIT_SHA", "version-test-sha"); std::env::set_var("CPR_BUILD_TIME", "2026-07-01T00:00:00Z"); - let (app, _dir) = admin_system_test_app("system-version.sqlite").await; + let (app, _dir) = admin_system_test_app("system-version").await; let response = app .oneshot( @@ -91,7 +78,7 @@ async fn update_detail_should_use_cached_release_when_not_refreshed() { let success_server = MockServer::start().await; mount_latest_release(&success_server, repository, "0.2.0").await; set_system_update_env(repository, &success_server.uri()); - let (app, _dir) = admin_system_test_app("system-cache.sqlite").await; + let (app, _dir) = admin_system_test_app("system-cache").await; let initial = get_update_detail(&app, true, "req_system_cache_initial").await; let failing_server = MockServer::start().await; @@ -112,7 +99,7 @@ async fn update_detail_should_fallback_to_cache_when_refresh_fetch_fails() { let success_server = MockServer::start().await; mount_latest_release(&success_server, repository, "0.3.0").await; set_system_update_env(repository, &success_server.uri()); - let (app, _dir) = admin_system_test_app("system-cache-forced.sqlite").await; + let (app, _dir) = admin_system_test_app("system-cache-forced").await; let initial = get_update_detail(&app, true, "req_system_cache_forced_initial").await; let failing_server = MockServer::start().await; @@ -140,7 +127,7 @@ async fn update_detail_should_report_source_build_as_unsupported() { mount_latest_release(&github, repository, "0.2.0").await; set_system_update_env(repository, &github.uri()); std::env::set_var("CPR_BUILD_TYPE", "source"); - let (app, _dir) = admin_system_test_app("system-source-build.sqlite").await; + let (app, _dir) = admin_system_test_app("system-source-build").await; let info = get_update_detail(&app, true, "req_system_source_build").await; @@ -162,7 +149,7 @@ async fn update_detail_should_reject_untrusted_github_api_base() { "zyycn/codex-proxy-rs-untrusted-api-base-route", "https://mirror.example", ); - let (app, _dir) = admin_system_test_app("system-untrusted-api-base.sqlite").await; + let (app, _dir) = admin_system_test_app("system-untrusted-api-base").await; let info = get_update_detail(&app, true, "req_system_untrusted_api_base").await; @@ -178,7 +165,7 @@ async fn restart_should_request_process_restart_inside_docker() { let _guard = SYSTEM_ENV_LOCK.lock().await; set_system_update_env("zyycn/codex-proxy-rs-restart-route", "http://127.0.0.1:9"); std::env::set_var("CPR_ENABLE_SELF_RESTART", "true"); - let (app, _dir) = admin_system_test_app("system-restart.sqlite").await; + let (app, _dir) = admin_system_test_app("system-restart").await; let response = app .oneshot( @@ -195,7 +182,7 @@ async fn restart_should_request_process_restart_inside_docker() { let body = response_json(response).await; let shutdown = tokio::time::timeout( Duration::from_secs(2), - codex_proxy_rs::runtime::shutdown::shutdown_signal(), + codex_proxy_rs::bootstrap::shutdown::shutdown_signal(), ) .await; @@ -222,7 +209,7 @@ async fn restart_should_spawn_replacement_before_shutdown_outside_docker() { std::env::set_var("CPR_DEPLOYMENT_MODE", "binary"); std::env::set_var("CPR_ENABLE_SELF_RESTART", "true"); std::env::set_var("CPR_UPDATE_EXE_PATH", replacement); - let (app, _dir) = admin_system_test_app("system-binary-restart.sqlite").await; + let (app, _dir) = admin_system_test_app("system-binary-restart").await; let response = app .oneshot( @@ -239,7 +226,7 @@ async fn restart_should_spawn_replacement_before_shutdown_outside_docker() { let body = response_json(response).await; let shutdown = tokio::time::timeout( Duration::from_secs(2), - codex_proxy_rs::runtime::shutdown::shutdown_signal(), + codex_proxy_rs::bootstrap::shutdown::shutdown_signal(), ) .await; @@ -259,7 +246,7 @@ async fn restart_should_not_shutdown_when_replacement_spawn_fails() { "CPR_UPDATE_EXE_PATH", "/tmp/codex-proxy-rs-missing-restart-binary", ); - let (app, _dir) = admin_system_test_app("system-restart-spawn-failure.sqlite").await; + let (app, _dir) = admin_system_test_app("system-restart-spawn-failure").await; let response = app .oneshot( @@ -277,7 +264,7 @@ async fn restart_should_not_shutdown_when_replacement_spawn_fails() { let body = response_json(response).await; let shutdown = tokio::time::timeout( Duration::from_millis(100), - codex_proxy_rs::runtime::shutdown::shutdown_signal(), + codex_proxy_rs::bootstrap::shutdown::shutdown_signal(), ) .await; let message = body["message"].as_str().unwrap_or_default(); @@ -303,7 +290,7 @@ async fn update_should_replace_local_release_files_with_latest_asset() { mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update").await; let response = app .clone() @@ -355,7 +342,7 @@ async fn update_should_reject_when_confirmed_target_differs_from_remote_latest() mount_latest_release(&latest_github, repository, "0.5.0").await; set_system_update_env(repository, &initial_github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-target-changed.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-target-changed").await; let initial = get_update_detail(&app, true, "req_system_update_target_initial").await; set_system_update_env(repository, &latest_github.uri()); @@ -399,7 +386,7 @@ async fn update_should_reject_untrusted_github_api_base() { "https://mirror.example", ); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-untrusted-api-base.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-untrusted-api-base").await; let response = app .oneshot( @@ -445,7 +432,7 @@ async fn update_should_fail_when_release_checksum_is_missing() { .await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-missing-checksum.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-missing-checksum").await; assert_update_failure_preserves_files( &app, @@ -479,7 +466,7 @@ async fn update_should_fail_when_release_checksum_mismatches() { .await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-bad-checksum.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-bad-checksum").await; assert_update_failure_preserves_files( &app, @@ -513,7 +500,7 @@ async fn update_should_reject_release_archive_from_untrusted_host() { .await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-untrusted-download.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-untrusted-download").await; assert_update_failure_preserves_files( &app, @@ -548,7 +535,7 @@ async fn update_should_reject_insecure_release_archive_even_when_env_flag_is_set set_system_update_env(repository, &github.uri()); std::env::set_var("CPR_UPDATE_ALLOW_INSECURE_DOWNLOADS", "true"); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-insecure-download.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-insecure-download").await; assert_update_failure_preserves_files( &app, @@ -576,7 +563,7 @@ async fn update_should_reject_release_archive_with_unsafe_path() { .await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-unsafe-archive.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-unsafe-archive").await; assert_update_failure_preserves_files( &app, @@ -604,7 +591,7 @@ async fn update_should_restore_web_assets_when_binary_backup_fails() { mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-binary-backup-failure.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-binary-backup-failure").await; let response = app .clone() @@ -663,7 +650,7 @@ async fn update_should_remove_stale_file_lock_and_continue() { mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-update-stale-lock.sqlite").await; + let (app, _dir) = admin_system_test_app("system-update-stale-lock").await; let response = app .clone() @@ -708,7 +695,7 @@ async fn update_should_replace_web_assets_across_filesystems() { mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-cross-device-update.sqlite").await; + let (app, _dir) = admin_system_test_app("system-cross-device-update").await; let response = app .clone() @@ -767,7 +754,7 @@ async fn update_status_should_read_local_update_state() { .to_string(), ) .unwrap(); - let (app, _dir) = admin_system_test_app("system-status.sqlite").await; + let (app, _dir) = admin_system_test_app("system-status").await; let response = app .oneshot( @@ -800,7 +787,7 @@ async fn update_events_should_open_authenticated_sse_stream() { &deploy.path().join("codex-proxy-rs"), &deploy.path().join("web/dist"), ); - let (app, _dir) = admin_system_test_app("system-events.sqlite").await; + let (app, _dir) = admin_system_test_app("system-events").await; let response = app .oneshot( @@ -836,7 +823,7 @@ async fn update_events_should_close_on_shutdown() { &deploy.path().join("codex-proxy-rs"), &deploy.path().join("web/dist"), ); - let (app, _dir) = admin_system_test_app("system-events-shutdown.sqlite").await; + let (app, _dir) = admin_system_test_app("system-events-shutdown").await; let response = app .oneshot( @@ -850,7 +837,7 @@ async fn update_events_should_close_on_shutdown() { ) .await .unwrap(); - codex_proxy_rs::runtime::shutdown::request_shutdown(); + codex_proxy_rs::bootstrap::shutdown::request_shutdown(); let body = tokio::time::timeout(Duration::from_secs(2), to_bytes(response.into_body(), 1024)) .await @@ -874,7 +861,7 @@ async fn update_events_should_close_after_terminal_update_log() { mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await; set_system_update_env(repository, &github.uri()); set_system_update_paths(deploy.path(), &exe_path, &web_dist); - let (app, _dir) = admin_system_test_app("system-events-terminal.sqlite").await; + let (app, _dir) = admin_system_test_app("system-events-terminal").await; let sse_response = app .clone() @@ -1277,13 +1264,11 @@ async fn wait_for_operation_status(app: &axum::Router, expected: &str) -> Value } async fn admin_system_test_app(db_name: &str) -> (axum::Router, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - seed_admin_session(&pool, "session_1").await; - let config = test_config(url); - let stores = stores(pool); + let (pool, dir) = crate::support::storage::init_test_db(db_name).await; + let redis = crate::support::storage::create_test_redis(db_name).await; + seed_admin_session(&pool, &redis, "session_1").await; + let config = test_config(crate::support::storage::test_database_url()); + let stores = crate::support::storage::background_task_stores(pool, redis); let services = Arc::new(Services::new( &config, stores, @@ -1292,19 +1277,6 @@ async fn admin_system_test_app(db_name: &str) -> (axum::Router, tempfile::TempDi let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); (app, dir) } - -fn stores(pool: SqlitePool) -> BackgroundTaskStores { - BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool), - } -} diff --git a/backend/tests/proxy/auth.rs b/backend/tests/api/client/auth.rs similarity index 94% rename from backend/tests/proxy/auth.rs rename to backend/tests/api/client/auth.rs index 7346c20a7..e2a9069fe 100644 --- a/backend/tests/proxy/auth.rs +++ b/backend/tests/api/client/auth.rs @@ -1,5 +1,5 @@ use axum::http::{header::AUTHORIZATION, HeaderMap, HeaderValue}; -use codex_proxy_rs::proxy::auth::{bearer_client_api_key, ClientApiKeyAuthError}; +use codex_proxy_rs::api::client::auth::{bearer_client_api_key, ClientApiKeyAuthError}; #[test] fn bearer_client_api_key_should_report_missing_authorization() { diff --git a/backend/tests/proxy/keys/auth.rs b/backend/tests/api/client/key_auth.rs similarity index 100% rename from backend/tests/proxy/keys/auth.rs rename to backend/tests/api/client/key_auth.rs diff --git a/backend/tests/api/client/mod.rs b/backend/tests/api/client/mod.rs new file mode 100644 index 000000000..c8a1717a8 --- /dev/null +++ b/backend/tests/api/client/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod auth; +pub(crate) mod key_auth; +pub(crate) mod openai; diff --git a/backend/tests/proxy/openai/mod.rs b/backend/tests/api/client/openai/mod.rs similarity index 100% rename from backend/tests/proxy/openai/mod.rs rename to backend/tests/api/client/openai/mod.rs diff --git a/backend/tests/proxy/openai/models_routes.rs b/backend/tests/api/client/openai/models_routes.rs similarity index 65% rename from backend/tests/proxy/openai/models_routes.rs rename to backend/tests/api/client/openai/models_routes.rs index 0b3965ea6..246643af4 100644 --- a/backend/tests/proxy/openai/models_routes.rs +++ b/backend/tests/api/client/openai/models_routes.rs @@ -5,28 +5,20 @@ use axum::{ http::{Request, StatusCode}, Router, }; -use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::usage_record_store::SqliteUsageRecordStore, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::token_refresh::RefreshLeaseStore, - upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, - upstream::fingerprint::FingerprintRepository, -}; +use codex_proxy_rs::{bootstrap::services::Services, bootstrap::state::AppState}; use tower::util::ServiceExt; use crate::support::{ - client_keys::insert_client_api_key, config::test_config, fingerprint::runtime_fingerprint, + client_keys::insert_client_api_key, + config::test_config, + fingerprint::runtime_fingerprint, http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; #[tokio::test] async fn models_route_should_reject_unknown_client_api_key() { - let (app, _key, _dir) = test_app("openai-models-auth.sqlite", false).await; + let (app, _key, _dir) = test_app("openai-models-auth", false).await; let response = app .oneshot( @@ -44,7 +36,7 @@ async fn models_route_should_reject_unknown_client_api_key() { #[tokio::test] async fn models_route_should_return_uniform_error_for_malformed_authorization() { - let (app, _key, _dir) = test_app("openai-models-auth-malformed.sqlite", false).await; + let (app, _key, _dir) = test_app("openai-models-auth-malformed", false).await; let response = app .oneshot( @@ -67,7 +59,7 @@ async fn models_route_should_return_uniform_error_for_malformed_authorization() #[tokio::test] async fn models_route_should_accept_stored_client_api_key() { - let (app, plaintext, _dir) = test_app("openai-models-auth-valid.sqlite", true).await; + let (app, plaintext, _dir) = test_app("openai-models-auth-valid", true).await; let plaintext = plaintext.expect("client api key should be seeded"); let response = app @@ -94,7 +86,7 @@ async fn models_route_should_accept_stored_client_api_key() { async fn model_detail_route_should_accept_configured_alias() { let mut aliases = BTreeMap::new(); aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string()); - let (app, plaintext, _dir) = test_app_with_aliases("openai-models-alias.sqlite", aliases).await; + let (app, plaintext, _dir) = test_app_with_aliases("openai-models-alias", aliases).await; let response = app .oneshot( @@ -116,31 +108,25 @@ async fn test_app( db_name: &str, seed_client_key: bool, ) -> (Router, Option, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); + let (pool, dir) = init_test_db(db_name).await; let plaintext = if seed_client_key { Some(insert_client_api_key(&pool).await) } else { None }; - let config = test_config(url); - build_test_app(pool, config, plaintext, dir).await + let config = test_config(test_database_url()); + build_test_app(pool, config, plaintext, dir, db_name).await } async fn test_app_with_aliases( db_name: &str, model_aliases: BTreeMap, ) -> (Router, String, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); + let (pool, dir) = init_test_db(db_name).await; let plaintext = insert_client_api_key(&pool).await; - let mut config = test_config(url); + let mut config = test_config(test_database_url()); config.model_aliases = model_aliases; - let (app, plaintext, dir) = build_test_app(pool, config, Some(plaintext), dir).await; + let (app, plaintext, dir) = build_test_app(pool, config, Some(plaintext), dir, db_name).await; ( app, plaintext.expect("client api key should be seeded"), @@ -149,21 +135,15 @@ async fn test_app_with_aliases( } async fn build_test_app( - pool: sqlx::SqlitePool, - config: codex_proxy_rs::config::schema::AppConfig, + pool: sqlx::PgPool, + config: codex_proxy_rs::bootstrap::config::AppConfig, plaintext: Option, dir: tempfile::TempDir, + label: &str, ) -> (Router, Option, tempfile::TempDir) { - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + let redis = create_test_redis(label).await; + seed_model_snapshot(&redis).await; + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new( &config, @@ -177,7 +157,34 @@ async fn build_test_app( let state = AppState { services: (*services).clone(), }; - let app = codex_proxy_rs::http::router::router().with_state(state); + let app = codex_proxy_rs::api::router::router().with_state(state); (app, plaintext, dir) } + +async fn seed_model_snapshot(redis: &codex_proxy_rs::infra::redis::RedisConnection) { + let mut connection = redis.manager(); + let value = serde_json::json!({ + "models": [{ + "id": "gpt-5.5", + "displayName": "GPT 5.5", + "description": "Test model", + "isDefault": false, + "supportedReasoningEfforts": [{"reasoningEffort": "medium", "description": "medium"}], + "defaultReasoningEffort": "medium", + "inputModalities": ["text"], + "outputModalities": ["text"], + "supportsPersonality": false, + "upgrade": null, + "source": "test" + }], + "fetchedAt": chrono::Utc::now() + }); + let _: usize = redis::cmd("HSET") + .arg(redis.key("models:plan_snapshots")) + .arg("plus") + .arg(value.to_string()) + .query_async(&mut connection) + .await + .unwrap(); +} diff --git a/backend/tests/proxy/openai/responses_routes.rs b/backend/tests/api/client/openai/responses_routes.rs similarity index 67% rename from backend/tests/proxy/openai/responses_routes.rs rename to backend/tests/api/client/openai/responses_routes.rs index 318f9ae18..e017040a3 100644 --- a/backend/tests/proxy/openai/responses_routes.rs +++ b/backend/tests/api/client/openai/responses_routes.rs @@ -2,24 +2,15 @@ use axum::{ body::{to_bytes, Body}, http::{Request, StatusCode}, }; -use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::usage_record_store::SqliteUsageRecordStore, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::SqliteSessionAffinityStore, - runtime::services::{BackgroundTaskStores, Services}, - runtime::state::AppState, - upstream::accounts::token_refresh::RefreshLeaseStore, - upstream::accounts::{cookies::SqliteCookieStore, store::SqliteAccountStore}, - upstream::fingerprint::FingerprintRepository, -}; -use sqlx::SqlitePool; +use codex_proxy_rs::{bootstrap::services::Services, bootstrap::state::AppState}; use tower::util::ServiceExt; use crate::support::{ - client_keys::insert_client_api_key, config::test_config, fingerprint::runtime_fingerprint, + client_keys::insert_client_api_key, + config::test_config, + fingerprint::runtime_fingerprint, http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; #[tokio::test] @@ -137,23 +128,12 @@ async fn responses_route_should_default_omitted_stream_to_sse() { } async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("openai-responses-routes.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); + let (pool, dir) = init_test_db("openai-responses-routes").await; + let redis = create_test_redis("openai-responses-routes").await; let plaintext = insert_client_api_key(&pool).await; - insert_model_snapshot(&pool).await; - let config = test_config(url); - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool.clone()), - }; + insert_model_snapshot(&redis).await; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new( &config, @@ -168,22 +148,35 @@ async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::Temp services: (*services).clone(), }; ( - codex_proxy_rs::http::router::router().with_state(state), + codex_proxy_rs::api::router::router().with_state(state), plaintext, dir, ) } -async fn insert_model_snapshot(pool: &SqlitePool) { - sqlx::query( - r"insert or replace into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?)", - ) - .bind("plus") - .bind( - r#"[{"id":"gpt-5.5","displayName":"GPT 5.5","description":"Test model","isDefault":false,"supportedReasoningEfforts":[{"reasoningEffort":"medium","description":"medium"}],"defaultReasoningEffort":"medium","inputModalities":["text"],"outputModalities":["text"],"supportsPersonality":false,"upgrade":null,"source":"test"}]"#, - ) - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); +async fn insert_model_snapshot(redis: &codex_proxy_rs::infra::redis::RedisConnection) { + let mut connection = redis.manager(); + let value = serde_json::json!({ + "models": [{ + "id":"gpt-5.5", + "displayName":"GPT 5.5", + "description":"Test model", + "isDefault":false, + "supportedReasoningEfforts":[{"reasoningEffort":"medium","description":"medium"}], + "defaultReasoningEffort":"medium", + "inputModalities":["text"], + "outputModalities":["text"], + "supportsPersonality":false, + "upgrade":null, + "source":"test" + }], + "fetchedAt": chrono::Utc::now() + }); + let _: usize = redis::cmd("HSET") + .arg(redis.key("models:plan_snapshots")) + .arg("plus") + .arg(value.to_string()) + .query_async(&mut connection) + .await + .unwrap(); } diff --git a/backend/tests/http/mod.rs b/backend/tests/api/http/mod.rs similarity index 100% rename from backend/tests/http/mod.rs rename to backend/tests/api/http/mod.rs diff --git a/backend/tests/http/trace_middleware/mod.rs b/backend/tests/api/http/trace_middleware/mod.rs similarity index 99% rename from backend/tests/http/trace_middleware/mod.rs rename to backend/tests/api/http/trace_middleware/mod.rs index f0511eed4..495eafacf 100644 --- a/backend/tests/http/trace_middleware/mod.rs +++ b/backend/tests/api/http/trace_middleware/mod.rs @@ -13,7 +13,7 @@ use axum::{ routing::get, Extension, Router, }; -use codex_proxy_rs::http::middleware::{ +use codex_proxy_rs::api::middleware::{ request_id::{attach_request_id, ClientIp}, trace::http_trace_layer, }; diff --git a/backend/tests/http/web_assets/mod.rs b/backend/tests/api/http/web_assets/mod.rs similarity index 71% rename from backend/tests/http/web_assets/mod.rs rename to backend/tests/api/http/web_assets/mod.rs index bda99353b..0b5072810 100644 --- a/backend/tests/http/web_assets/mod.rs +++ b/backend/tests/api/http/web_assets/mod.rs @@ -5,13 +5,16 @@ use axum::{ http::{header, Request, StatusCode}, Router, }; -use codex_proxy_rs::{config::schema::AppConfig, infra::database::connect_sqlite}; +use codex_proxy_rs::bootstrap::config::AppConfig; use serde_json::Value; -use sqlx::SqlitePool; +use sqlx::PgPool; use tower::ServiceExt; use crate::support::{ - config::test_config as base_test_config, fingerprint::runtime_fingerprint, http::response_json, + config::test_config as base_test_config, + fingerprint::runtime_fingerprint, + http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; #[tokio::test] @@ -31,9 +34,7 @@ async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes .expect("favicon should be written"); std::fs::write(dist.join("assets").join("app.js"), "window.__asset = true;") .expect("asset should be written"); - let db = dir.path().join("assets-routes.sqlite"); - let database_url = format!("sqlite://{}", db.display()); - let (app, _pool) = router_with_dist_and_database(&dist, &database_url).await; + let (app, _pool) = router_with_dist_and_database(&dist, "assets-routes").await; let index = app.clone().oneshot(request("/")).await.expect("index"); let route_fallback = app @@ -98,9 +99,7 @@ async fn server_router_should_return_json_404_for_unknown_api_routes() { "
Codex Proxy Admin
", ) .expect("index should be written"); - let db = dir.path().join("unknown-api-routes.sqlite"); - let database_url = format!("sqlite://{}", db.display()); - let (app, _pool) = router_with_dist_and_database(&dist, &database_url).await; + let (app, _pool) = router_with_dist_and_database(&dist, "unknown-api-routes").await; let admin = app .clone() @@ -117,14 +116,12 @@ async fn server_router_should_return_json_404_for_unknown_api_routes() { } #[tokio::test] -async fn healthz_should_report_unavailable_when_sqlite_is_closed() { +async fn healthz_should_report_unavailable_when_postgres_is_closed() { let dir = tempfile::tempdir().expect("temp dir"); let dist = dir.path().join("dist"); std::fs::create_dir_all(&dist).expect("dist dir should be created"); std::fs::write(dist.join("index.html"), "").expect("index should be written"); - let db = dir.path().join("healthz.sqlite"); - let database_url = format!("sqlite://{}", db.display()); - let (app, pool) = router_with_dist_and_database(&dist, &database_url).await; + let (app, pool) = router_with_dist_and_database(&dist, "healthz-postgres").await; pool.close().await; let response = app @@ -173,42 +170,22 @@ fn assert_no_static_policy_headers(response: &axum::response::Response) { assert!(!response.headers().contains_key("referrer-policy")); } -async fn router_with_dist_and_database(dist: &Path, database_url: &str) -> (Router, SqlitePool) { - let pool = connect_sqlite(database_url).await.expect("sqlite pool"); - let config = test_config(database_url.to_string()); - let stores = codex_proxy_rs::runtime::services::BackgroundTaskStores { - accounts: codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone()), - admin_sessions: codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore::new( - pool.clone(), - ), - cookies: codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore::new(pool.clone()), - fingerprints: codex_proxy_rs::upstream::fingerprint::FingerprintRepository::new( - pool.clone(), - ), - session_affinity: - codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new( - pool.clone(), - ), - refresh_leases: codex_proxy_rs::upstream::accounts::token_refresh::RefreshLeaseStore::new( - pool.clone(), - ), - client_keys: codex_proxy_rs::admin::keys::service::SqliteClientKeyStore::new(pool.clone()), - usage_records: - codex_proxy_rs::admin::monitoring::usage_record_store::SqliteUsageRecordStore::new( - pool.clone(), - ), - }; +async fn router_with_dist_and_database(dist: &Path, label: &str) -> (Router, PgPool) { + let (pool, _guard) = init_test_db(label).await; + let redis = create_test_redis(label).await; + let config = test_config(test_database_url()); + let stores = background_task_stores(pool.clone(), redis); let fingerprint = crate::support::fingerprint::test_fingerprint(); - let services = std::sync::Arc::new(codex_proxy_rs::runtime::services::Services::new( + let services = std::sync::Arc::new(codex_proxy_rs::bootstrap::services::Services::new( &config, stores, runtime_fingerprint(fingerprint), )); - let state = codex_proxy_rs::runtime::state::AppState { + let state = codex_proxy_rs::bootstrap::state::AppState { services: (*services).clone(), }; ( - codex_proxy_rs::http::router::router_with_assets(dist).with_state(state), + codex_proxy_rs::api::router::router_with_assets(dist).with_state(state), pool, ) } diff --git a/backend/tests/api/mod.rs b/backend/tests/api/mod.rs new file mode 100644 index 000000000..e9410c3dd --- /dev/null +++ b/backend/tests/api/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod admin; +pub(crate) mod client; +pub(crate) mod http; diff --git a/backend/tests/runtime/account_pool_restore/mod.rs b/backend/tests/bootstrap/account_pool_restore/mod.rs similarity index 73% rename from backend/tests/runtime/account_pool_restore/mod.rs rename to backend/tests/bootstrap/account_pool_restore/mod.rs index 8fac6ae60..8dd0c8107 100644 --- a/backend/tests/runtime/account_pool_restore/mod.rs +++ b/backend/tests/bootstrap/account_pool_restore/mod.rs @@ -1,21 +1,21 @@ use std::sync::Arc; use chrono::Utc; -use codex_proxy_rs::upstream::accounts::{ +use codex_proxy_rs::accounts::{ pool::{ AccountAcquireRequest, AccountPoolOptions, RotationStrategy, RuntimeAccountPoolService, }, - store::{AccountStore, SqliteAccountStore}, + store::{AccountStore, PgAccountStore}, }; -use crate::support::sqlite::init_test_db; +use crate::support::storage::init_test_db; #[tokio::test] async fn runtime_account_pool_should_restore_capacity_and_clear_runtime_state() { - let (pool, _dir) = init_test_db("runtime-account-pool.sqlite").await; + let (pool, _dir) = init_test_db("runtime-account-pool").await; insert_account(&pool, "acct_pool").await; - let store = SqliteAccountStore::new(pool); + let store = PgAccountStore::new(pool); let service = RuntimeAccountPoolService::new( Arc::new(store) as Arc, AccountPoolOptions { @@ -26,7 +26,7 @@ async fn runtime_account_pool_should_restore_capacity_and_clear_runtime_state() ); let restored = service - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore"); let capacity = service.capacity_summary_now().await; @@ -36,7 +36,7 @@ async fn runtime_account_pool_should_restore_capacity_and_clear_runtime_state() .expect("active account should be acquired"); let used_capacity = service.capacity_summary_now().await; let restored_again = service - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore again"); let restored_capacity = service.capacity_summary_now().await; @@ -53,12 +53,12 @@ async fn runtime_account_pool_should_restore_capacity_and_clear_runtime_state() } #[tokio::test] -async fn runtime_account_pool_should_restore_accounts_in_insert_order() { - let (pool, _dir) = init_test_db("runtime-account-pool-order.sqlite").await; +async fn runtime_account_pool_should_restore_accounts_by_added_at_and_id() { + let (pool, _dir) = init_test_db("runtime-account-pool-order").await; insert_account(&pool, "acct_c").await; insert_account(&pool, "acct_a").await; - let store = SqliteAccountStore::new(pool); + let store = PgAccountStore::new(pool); let service = RuntimeAccountPoolService::new( Arc::new(store) as Arc, AccountPoolOptions { @@ -69,7 +69,7 @@ async fn runtime_account_pool_should_restore_accounts_in_insert_order() { ); service - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore"); let acquired = service @@ -77,22 +77,22 @@ async fn runtime_account_pool_should_restore_accounts_in_insert_order() { .await .expect("active account should be acquired"); - assert_eq!(acquired.account.id, "acct_c"); + assert_eq!(acquired.account.id, "acct_a"); } -async fn insert_account(pool: &sqlx::SqlitePool, id: &str) { +async fn insert_account(pool: &sqlx::PgPool, id: &str) { sqlx::query( - "insert into accounts (id, email, chatgpt_account_id, chatgpt_user_id, access_token, access_token_expires_at, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + "insert into accounts (id, email, chatgpt_account_id, chatgpt_user_id, access_token, access_token_expires_at, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7, $8, $9)", ) .bind(id) .bind(format!("{id}@example.com")) .bind(format!("chatgpt-{id}")) .bind(format!("user-{id}")) .bind(format!("access-{id}")) - .bind("2999-01-01T00:00:00Z") + .bind(crate::support::storage::timestamp("2999-01-01T00:00:00Z")) .bind("active") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) + .bind(crate::support::storage::timestamp("2026-06-18T00:00:00Z")) .execute(pool) .await .expect("account should be inserted"); diff --git a/backend/tests/config/mod.rs b/backend/tests/bootstrap/config/mod.rs similarity index 91% rename from backend/tests/config/mod.rs rename to backend/tests/bootstrap/config/mod.rs index e17ffadb7..25be2324d 100644 --- a/backend/tests/config/mod.rs +++ b/backend/tests/bootstrap/config/mod.rs @@ -1,6 +1,6 @@ -use codex_proxy_rs::config::schema::LoggingConfig; -use codex_proxy_rs::config::schema::{AdminConfig, AdminConfigError, AppConfig}; -use codex_proxy_rs::config::settings::account_pool_options_from_config; +use codex_proxy_rs::bootstrap::config::LoggingConfig; +use codex_proxy_rs::bootstrap::config::{AdminConfig, AdminConfigError, AppConfig}; +use codex_proxy_rs::bootstrap::services::account_pool_options_from_config; use serde::de::DeserializeOwned; use std::fs; @@ -11,7 +11,9 @@ server: api: base_url: https://chatgpt.com/backend-api database: - url: sqlite://.runtime/data/codex-proxy-rs.sqlite + url: postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy +redis: + url: redis://127.0.0.1:6379 quota: refresh_interval_minutes: 5 skip_exhausted: true @@ -104,8 +106,9 @@ fn default_config_keeps_runtime_artifacts_under_runtime_directory() { assert_eq!( cfg.database.url, - "sqlite://.runtime/data/codex-proxy-rs.sqlite" + "postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy" ); + assert_eq!(cfg.redis.url, "redis://127.0.0.1:6379"); assert_eq!(cfg.logging.directory, ".runtime/logs"); } @@ -153,7 +156,9 @@ api: base_url: https://chatgpt.com/backend-api unexpected: {} database: - url: sqlite://.runtime/data/codex-proxy-rs.sqlite + url: postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy +redis: + url: redis://127.0.0.1:6379 ", ) .unwrap_err(); @@ -176,7 +181,9 @@ server: api: base_url: https://chatgpt.com/backend-api database: - url: sqlite://.runtime/data/codex-proxy-rs.sqlite + url: postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy +redis: + url: redis://127.0.0.1:6379 quota: skip_exhausted: false tls: @@ -259,8 +266,9 @@ ws_pool: assert_eq!(cfg.server.port, 8080); assert_eq!( cfg.database.url, - "sqlite://.runtime/data/codex-proxy-rs.sqlite" + "postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy" ); + assert_eq!(cfg.redis.url, "redis://127.0.0.1:6379"); assert!(cfg.model_aliases.is_empty()); assert_eq!(cfg.auth.max_concurrent_per_account, 3); assert_eq!(cfg.quota.refresh_interval_minutes, 5); diff --git a/backend/tests/bootstrap/import_sqlite.rs b/backend/tests/bootstrap/import_sqlite.rs new file mode 100644 index 000000000..30fad6f00 --- /dev/null +++ b/backend/tests/bootstrap/import_sqlite.rs @@ -0,0 +1,256 @@ +use std::path::PathBuf; + +use codex_proxy_rs::bootstrap::import_sqlite::{import_sqlite, ImportSqliteError}; +use sqlx::{sqlite::SqliteConnectOptions, sqlite::SqlitePoolOptions, Row}; + +use crate::support::storage::init_test_db; + +#[tokio::test] +async fn import_sqlite_should_split_facts_preserve_retrievable_keys_and_report_discards() { + let (_source_guard, source_path) = sqlite_v3_source("success").await; + seed_representative_v3_data(&source_path).await; + let (target, _target_guard) = init_test_db("import-sqlite-success").await; + + let report = import_sqlite(&target, &source_path).await.unwrap(); + + assert_eq!(report.imported_rows("client_api_keys"), 1); + assert_eq!(report.imported_rows("accounts"), 2); + assert_eq!(report.normalized_rows("accounts.refreshing_to_expired"), 1); + assert_eq!(report.imported_rows("usage_records"), 1); + assert_eq!(report.imported_rows("ops_error_logs"), 2); + assert_eq!(report.discarded_rows("usage_record_noise"), 1); + assert_eq!(report.discarded_rows("admin_sessions"), 1); + assert_eq!(report.discarded_rows("session_affinities"), 1); + assert_eq!(report.discarded_rows("account_refresh_leases"), 1); + assert_eq!(report.discarded_rows("model_plan_snapshots"), 1); + assert_eq!(report.cookie_expiration_parse_failures(), 1); + + let key: String = sqlx::query_scalar("select key from client_api_keys where id = $1") + .bind("key-1") + .fetch_one(&target) + .await + .unwrap(); + assert_eq!(key, "sk-old-secret"); + + let refreshing_account: (String, Option>) = + sqlx::query_as("select status, next_refresh_at from accounts where id = $1") + .bind("acct-refreshing") + .fetch_one(&target) + .await + .unwrap(); + assert_eq!(refreshing_account.0, "expired"); + assert!(refreshing_account.1.is_none()); + + let success = sqlx::query( + "select requested_model, upstream_model, service_tier, first_token_ms, + input_tokens, output_tokens, cached_tokens, reasoning_tokens, metadata_json + from usage_records where id = $1", + ) + .bind("usage-success") + .fetch_one(&target) + .await + .unwrap(); + assert_eq!( + success + .get::, _>("requested_model") + .as_deref(), + Some("alias") + ); + assert_eq!( + success + .get::, _>("upstream_model") + .as_deref(), + Some("gpt-5.5") + ); + assert_eq!( + success.get::, _>("service_tier").as_deref(), + Some("priority") + ); + assert_eq!(success.get::, _>("first_token_ms"), Some(0)); + assert_eq!(success.get::, _>("input_tokens"), Some(12)); + assert_eq!(success.get::, _>("output_tokens"), Some(5)); + assert_eq!(success.get::, _>("cached_tokens"), Some(3)); + assert_eq!(success.get::, _>("reasoning_tokens"), Some(2)); + let metadata: serde_json::Value = success.get("metadata_json"); + assert_eq!(metadata["clientIp"], "127.0.0.1"); + assert!(metadata.get("usage").is_none()); + + let bucket: (String, String, String, i64, i64, Option) = sqlx::query_as( + "select provider, account_id, model, success_count, error_count, min_latency_ms + from request_time_buckets", + ) + .fetch_one(&target) + .await + .unwrap(); + assert_eq!(bucket.0, "openai"); + assert_eq!(bucket.1, "__unknown__"); + assert_eq!(bucket.2, "__unknown__"); + assert_eq!(bucket.3, 2); + assert_eq!(bucket.4, 1); + assert_eq!(bucket.5, None); +} + +#[tokio::test] +async fn import_sqlite_should_rollback_all_target_writes_on_invalid_source_value() { + let (_source_guard, source_path) = sqlite_v3_source("rollback").await; + let source = open_sqlite(&source_path, false).await; + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) + values ('admin', 'hash', 'not-a-time', '2026-01-01T00:00:00Z')", + ) + .execute(&source) + .await + .unwrap(); + source.close().await; + let (target, _target_guard) = init_test_db("import-sqlite-rollback").await; + + let error = import_sqlite(&target, &source_path).await.unwrap_err(); + + assert!(matches!(error, ImportSqliteError::InvalidTimestamp { .. })); + let rows: i64 = sqlx::query_scalar("select count(*) from admin_users") + .fetch_one(&target) + .await + .unwrap(); + assert_eq!(rows, 0); +} + +#[tokio::test] +async fn import_sqlite_should_reject_unknown_source_version() { + let (_source_guard, source_path) = sqlite_v3_source("version").await; + let source = open_sqlite(&source_path, false).await; + sqlx::query("delete from schema_migrations where version = 3") + .execute(&source) + .await + .unwrap(); + source.close().await; + let (target, _target_guard) = init_test_db("import-sqlite-version").await; + + let error = import_sqlite(&target, &source_path).await.unwrap_err(); + + assert!(matches!( + error, + ImportSqliteError::UnsupportedSourceVersion { actual: Some(2) } + )); +} + +#[tokio::test] +async fn import_sqlite_should_reject_nonempty_target() { + let (_source_guard, source_path) = sqlite_v3_source("nonempty").await; + let (target, _target_guard) = init_test_db("import-sqlite-nonempty").await; + sqlx::query( + "insert into admin_users (id, password_hash, created_at, updated_at) + values ($1, $2, now(), now())", + ) + .bind("existing") + .bind("hash") + .execute(&target) + .await + .unwrap(); + + let error = import_sqlite(&target, &source_path).await.unwrap_err(); + + assert!(matches!(error, ImportSqliteError::TargetNotEmpty)); +} + +async fn sqlite_v3_source(label: &str) -> (tempfile::TempDir, PathBuf) { + let guard = tempfile::tempdir().unwrap(); + let path = guard.path().join(format!("{label}.sqlite")); + let source = open_sqlite(&path, true).await; + sqlx::raw_sql(include_str!("../fixtures/sqlite_v3.sql")) + .execute(&source) + .await + .unwrap(); + source.close().await; + (guard, path) +} + +async fn open_sqlite(path: &PathBuf, create: bool) -> sqlx::SqlitePool { + SqlitePoolOptions::new() + .max_connections(1) + .connect_with( + SqliteConnectOptions::new() + .filename(path) + .create_if_missing(create), + ) + .await + .unwrap() +} + +async fn seed_representative_v3_data(path: &PathBuf) { + let source = open_sqlite(path, false).await; + sqlx::raw_sql( + r#" +insert into admin_users values + ('admin', 'argon-hash', '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z'); +insert into admin_sessions values + ('session', 'admin', '2026-01-02T00:00:00Z', '2026-01-01T00:00:00Z'); +insert into client_api_keys values + ('key-1', 'legacy', 'sk-old', 'sk-old-secret', null, 1, + '2026-01-01T00:00:00Z', null); +insert into runtime_settings values + (1, '{}', 3600, 2, 3, 50, 'smart', 'admin-machine-secret', + '2026-01-01T00:00:00Z'); +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 + ('acct-1', 'a@example.com', 'chatgpt-1', 'user-1', null, 'plus', + 'access', 'refresh', '2026-02-01T00:00:00Z', '2026-01-31T23:00:00Z', + 'active', '{"remaining":1}', '2026-01-01T00:00:00Z', 0, 0, null, null, + '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z'), + ('acct-refreshing', 'refreshing@example.com', 'chatgpt-refreshing', 'user-refreshing', + null, 'plus', 'access-refreshing', 'refresh-refreshing', '2026-02-01T00:00:00Z', + '2026-01-31T23:00:00Z', 'refreshing', null, null, 0, 0, null, null, + '2026-01-01T00:00:00Z', '2026-01-01T00:00:00Z'); +insert into account_cookies values + ('cookie-1', 'acct-1', '.chatgpt.com', 'session', 'secret', '/', + 'not-a-cookie-date', '2026-01-01T00:00:00Z'); +insert into account_refresh_leases values + ('acct-1', 'owner', '2026-01-01T00:05:00Z', '2026-01-01T00:00:00Z'); +insert into model_plan_snapshots values + ('plus', '[]', '2026-01-01T00:00:00Z'); +insert into session_affinities values + ('response-1', 'acct-1', 'conversation-1', null, null, null, '[]', null, + '2026-01-02T00:00:00Z', '2026-01-01T00:00:00Z'); +insert into usage_records + (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) +values + ('usage-success', 'req-1', 'v1.response', 'info', 'acct-1', '/v1/responses', + 'gpt-5.5', 200, 'http', 0, 200, null, 'resp-1', 'upstream-1', 0, + 'ok', + '{"requestedModel":"alias","upstreamModel":"gpt-5.5","serviceTier":"priority","firstTokenMs":0,"usage":{"inputTokens":12,"outputTokens":5,"cachedTokens":3,"reasoningTokens":2},"clientIp":"127.0.0.1"}', + '2026-01-01T00:01:00Z'), + ('usage-error', 'req-2', 'v1.response', 'error', 'acct-1', '/v1/responses', + 'gpt-5.5', 502, 'http', 1, 500, 'upstream', null, null, 10, + 'failed', '{}', '2026-01-01T00:02:00Z'), + ('usage-noise', null, 'diagnostic', 'info', null, null, null, null, null, + null, null, null, null, null, null, 'noise', '{}', '2026-01-01T00:03:00Z'); +insert into ops_error_logs + (id, request_id, kind, account_id, route, model, status_code, + client_status_code, upstream_status_code, transport, attempt_index, + failure_class, response_id, upstream_request_id, latency_ms, message, + metadata_json, created_at) +values + ('ops-1', 'req-3', 'v1.response', 'acct-1', '/v1/responses', 'gpt-5.5', + 429, 429, 429, 'http', 0, 'quota', null, null, 5, 'limited', '{}', + '2026-01-01T00:04:00Z'); +insert into usage_time_buckets + (bucket_start, account_id, model, service_tier, request_count, error_count, + input_tokens, output_tokens, cached_tokens, first_token_latency_sum, + first_token_latency_count, latency_sum, latency_count, max_latency_ms, + min_latency_ms, updated_at) +values + ('2026-01-01T00:00:00Z', '', '', '', 3, 1, 12, 5, 3, 0, 1, 0, 1, 0, 0, + '2026-01-01T00:15:00Z'); +"#, + ) + .execute(&source) + .await + .unwrap(); + source.close().await; +} diff --git a/backend/tests/runtime/mod.rs b/backend/tests/bootstrap/mod.rs similarity index 61% rename from backend/tests/runtime/mod.rs rename to backend/tests/bootstrap/mod.rs index 90cc4e238..0ce8f5300 100644 --- a/backend/tests/runtime/mod.rs +++ b/backend/tests/bootstrap/mod.rs @@ -1,3 +1,5 @@ pub(crate) mod account_pool_restore; +pub(crate) mod config; +pub(crate) mod import_sqlite; pub(crate) mod services; pub(crate) mod tasks; diff --git a/backend/tests/bootstrap/services.rs b/backend/tests/bootstrap/services.rs new file mode 100644 index 000000000..980be50b1 --- /dev/null +++ b/backend/tests/bootstrap/services.rs @@ -0,0 +1,81 @@ +use std::process::Command; + +use crate::support::{ + config::test_config, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, +}; +use codex_proxy_rs::{ + bootstrap::services::Services, upstream::openai::transport::tls::CODEX_CA_CERT_ENV, +}; + +#[tokio::test] +async fn services_try_new_should_use_configured_tls_transport_builder() { + const CASE_ENV: &str = "CODEX_PROXY_TEST_SERVICES_TLS_CASE"; + + if std::env::var(CASE_ENV).as_deref() == Ok("invalid_ca") { + let (pool, _temp_dir) = init_test_db("services-tls").await; + let config = test_config(test_database_url()); + let redis = create_test_redis("services-tls").await; + let stores = background_task_stores(pool, redis); + + let Err(error) = Services::try_new( + &config, + stores, + crate::support::fingerprint::runtime_test_fingerprint(), + ) else { + panic!("invalid custom CA should fail service transport construction"); + }; + assert!(error + .to_string() + .contains("Failed to read CA certificate file")); + return; + } + + let current_exe = std::env::current_exe().expect("current test binary path"); + let output = Command::new(current_exe) + .arg("--exact") + .arg("bootstrap::services::services_try_new_should_use_configured_tls_transport_builder") + .arg("--nocapture") + .env(CASE_ENV, "invalid_ca") + .env(CODEX_CA_CERT_ENV, "/tmp/codex-proxy-rs-missing-ca.pem") + .output() + .expect("run isolated services TLS test case"); + + assert!( + output.status.success(), + "isolated services TLS test failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); +} + +#[tokio::test] +async fn services_try_new_should_expose_websocket_pool_when_enabled() { + let (pool, _temp_dir) = init_test_db("services-ws-pool-enabled").await; + let config = test_config(test_database_url()); + let redis = create_test_redis("services-ws-pool-enabled").await; + let services = Services::try_new( + &config, + background_task_stores(pool, redis), + crate::support::fingerprint::runtime_test_fingerprint(), + ) + .expect("services should build"); + + assert!(services.websocket_pool.is_some()); +} + +#[tokio::test] +async fn services_try_new_should_not_expose_websocket_pool_when_disabled() { + let (pool, _temp_dir) = init_test_db("services-ws-pool-disabled").await; + let mut config = test_config(test_database_url()); + let redis = create_test_redis("services-ws-pool-disabled").await; + config.ws_pool.enabled = false; + let services = Services::try_new( + &config, + background_task_stores(pool, redis), + crate::support::fingerprint::runtime_test_fingerprint(), + ) + .expect("services should build"); + + assert!(services.websocket_pool.is_none()); +} diff --git a/backend/tests/bootstrap/tasks/cleanup.rs b/backend/tests/bootstrap/tasks/cleanup.rs new file mode 100644 index 000000000..aa17a4abb --- /dev/null +++ b/backend/tests/bootstrap/tasks/cleanup.rs @@ -0,0 +1,34 @@ +use super::*; + +#[tokio::test] +async fn cookie_cleanup_task_should_delete_only_expired_cookie_rows() { + let (pool, _guard) = init_test_db("cookie-cleanup").await; + let store = PgCookieStore::new(pool.clone()); + insert_account(&pool, "acct-cookie").await; + let now = Utc::now(); + insert_cookie( + &pool, + "expired-cookie", + "acct-cookie", + "old", + now - Duration::minutes(1), + ) + .await; + insert_cookie( + &pool, + "active-cookie", + "acct-cookie", + "fresh", + now + Duration::minutes(10), + ) + .await; + + let deleted = codex_proxy_rs::bootstrap::tasks::cookie_cleanup::CookieCleanupTask::new(store) + .cleanup_once_at(now) + .await + .expect("cookie cleanup should succeed"); + + let remaining = cookie_ids(&pool).await; + + assert_eq!((deleted, remaining), (1, vec!["active-cookie".to_string()])); +} diff --git a/backend/tests/runtime/tasks/coordinator.rs b/backend/tests/bootstrap/tasks/coordinator.rs similarity index 76% rename from backend/tests/runtime/tasks/coordinator.rs rename to backend/tests/bootstrap/tasks/coordinator.rs index eaba6cd25..69d0df54f 100644 --- a/backend/tests/runtime/tasks/coordinator.rs +++ b/backend/tests/bootstrap/tasks/coordinator.rs @@ -1,13 +1,13 @@ use std::{sync::Arc, time::Duration}; use codex_proxy_rs::{ - runtime::tasks::coordinator::SchedulerHandle, - upstream::transport::{CodexWebSocketPool, CodexWebSocketPoolConfig}, + bootstrap::tasks::coordinator::SchedulerHandle, + upstream::openai::transport::{CodexWebSocketPool, CodexWebSocketPoolConfig}, }; #[tokio::test] async fn task_coordinator_should_shutdown_without_panicking() { - let coordinator = codex_proxy_rs::runtime::tasks::coordinator::TaskCoordinator::default(); + let coordinator = codex_proxy_rs::bootstrap::tasks::coordinator::TaskCoordinator::default(); coordinator.shutdown().await; } diff --git a/backend/tests/runtime/tasks/fingerprint.rs b/backend/tests/bootstrap/tasks/fingerprint.rs similarity index 64% rename from backend/tests/runtime/tasks/fingerprint.rs rename to backend/tests/bootstrap/tasks/fingerprint.rs index 3f59066b4..d7b84c136 100644 --- a/backend/tests/runtime/tasks/fingerprint.rs +++ b/backend/tests/bootstrap/tasks/fingerprint.rs @@ -3,30 +3,28 @@ use super::*; #[tokio::test] async fn fingerprint_update_task_should_start_background_checker() { let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("fingerprints.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool.clone()); + let (pool, _guard) = init_test_db("fingerprint-task-start").await; + let repo = PgFingerprintStore::new(pool.clone()); repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) .await .expect("seed current fingerprint"); let runtime_fingerprint = runtime_fingerprint_from_repo(&repo).await; - let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task( - repo, - runtime_fingerprint, - "https://example.invalid/appcast.xml".to_string(), - dir.path().join("extracted-fingerprint.json"), - "26.800.1".to_string(), - "6001".to_string(), - ); + let handle = + codex_proxy_rs::bootstrap::tasks::fingerprint_update::start_fingerprint_update_task( + repo, + runtime_fingerprint, + "https://example.invalid/appcast.xml".to_string(), + dir.path().join("extracted-fingerprint.json"), + "26.800.1".to_string(), + "6001".to_string(), + ); handle.shutdown().await; } #[tokio::test] -async fn fingerprint_update_task_should_apply_initial_appcast_update_to_repository() { +async fn fingerprint_update_task_should_apply_initial_appcast_update_to_store() { let server = MockServer::start().await; mount_appcast( &server, @@ -43,11 +41,8 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito .await; let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("fingerprints-initial-update.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool.clone()); + let (pool, _guard) = init_test_db("fingerprint-task-update").await; + let repo = PgFingerprintStore::new(pool.clone()); repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) .await .expect("seed current fingerprint"); @@ -59,14 +54,15 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito ) .expect("extracted fingerprint should be written"); - let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task( - repo.clone(), - runtime_fingerprint.clone(), - format!("{}/appcast.xml", server.uri()), - extracted_path, - "26.800.1".to_string(), - "6001".to_string(), - ); + let handle = + codex_proxy_rs::bootstrap::tasks::fingerprint_update::start_fingerprint_update_task( + repo.clone(), + runtime_fingerprint.clone(), + format!("{}/appcast.xml", server.uri()), + extracted_path, + "26.800.1".to_string(), + "6001".to_string(), + ); let stored = wait_for_current_fingerprint_version(&repo).await; handle.shutdown().await; @@ -108,23 +104,21 @@ async fn fingerprint_update_task_should_not_persist_when_appcast_matches_current .await; let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("fingerprints-no-update.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool.clone()); + let (pool, _guard) = init_test_db("fingerprint-task-no-update").await; + let repo = PgFingerprintStore::new(pool.clone()); repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) .await .expect("seed current fingerprint"); let runtime_fingerprint = runtime_fingerprint_from_repo(&repo).await; - let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task( - repo.clone(), - runtime_fingerprint, - format!("{}/appcast.xml", server.uri()), - dir.path().join("missing-extracted-fingerprint.json"), - "26.900.1".to_string(), - "7001".to_string(), - ); + let handle = + codex_proxy_rs::bootstrap::tasks::fingerprint_update::start_fingerprint_update_task( + repo.clone(), + runtime_fingerprint, + format!("{}/appcast.xml", server.uri()), + dir.path().join("missing-extracted-fingerprint.json"), + "26.900.1".to_string(), + "7001".to_string(), + ); wait_for_appcast_requests(&server, 1).await; handle.shutdown().await; diff --git a/backend/tests/runtime/tasks/mod.rs b/backend/tests/bootstrap/tasks/mod.rs similarity index 59% rename from backend/tests/runtime/tasks/mod.rs rename to backend/tests/bootstrap/tasks/mod.rs index d41cc3cee..b1e1ca30f 100644 --- a/backend/tests/runtime/tasks/mod.rs +++ b/backend/tests/bootstrap/tasks/mod.rs @@ -2,22 +2,20 @@ use std::{sync::Arc, time::Duration as StdDuration}; use async_trait::async_trait; use chrono::{Duration, Utc}; -use codex_proxy_rs::admin::auth::service::SqliteAdminSessionStore; -use codex_proxy_rs::infra::database::connect_sqlite; -use codex_proxy_rs::proxy::dispatch::session_affinity::SessionAffinityEntry; -use codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore; -use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore; -use codex_proxy_rs::upstream::accounts::model::{Account, AccountStatus}; -use codex_proxy_rs::upstream::accounts::model::{AccountModelUsageDelta, AccountUsageDelta}; -use codex_proxy_rs::upstream::accounts::store::{AccountStore, AccountStoreResult}; -use codex_proxy_rs::upstream::fingerprint::{FingerprintRepository, RuntimeFingerprint}; -use codex_proxy_rs::upstream::models::config::ModelConfig; -use codex_proxy_rs::upstream::models::service::ModelService; +use codex_proxy_rs::accounts::account::{Account, AccountStatus}; +use codex_proxy_rs::accounts::account::{AccountModelUsageDelta, AccountUsageDelta}; +use codex_proxy_rs::accounts::cookies::PgCookieStore; +use codex_proxy_rs::accounts::store::{AccountStore, AccountStoreResult}; +use codex_proxy_rs::models::config::ModelConfig; +use codex_proxy_rs::models::service::ModelService; +use codex_proxy_rs::upstream::openai::fingerprint::{PgFingerprintStore, RuntimeFingerprint}; use wiremock::{ matchers::{method, path}, Mock, MockServer, ResponseTemplate, }; +use crate::support::storage::init_test_db; + mod cleanup; mod coordinator; mod fingerprint; @@ -116,8 +114,8 @@ async fn mount_appcast(server: &MockServer, body: &str) { } async fn wait_for_current_fingerprint_version( - repo: &FingerprintRepository, -) -> codex_proxy_rs::upstream::fingerprint::Fingerprint { + repo: &PgFingerprintStore, +) -> codex_proxy_rs::upstream::openai::fingerprint::Fingerprint { let deadline = tokio::time::Instant::now() + StdDuration::from_secs(2); loop { if let Some(stored) = repo @@ -137,7 +135,7 @@ async fn wait_for_current_fingerprint_version( } } -async fn runtime_fingerprint_from_repo(repo: &FingerprintRepository) -> RuntimeFingerprint { +async fn runtime_fingerprint_from_repo(repo: &PgFingerprintStore) -> RuntimeFingerprint { RuntimeFingerprint::new( repo.load_current() .await @@ -168,106 +166,42 @@ async fn wait_for_appcast_requests(server: &MockServer, expected_count: usize) { } } -async fn insert_account(pool: &sqlx::SqlitePool, account_id: &str) { - let now = Utc::now().to_rfc3339(); +async fn insert_account(pool: &sqlx::PgPool, account_id: &str) { + let now = Utc::now(); sqlx::query( - "insert into accounts (id, access_token, status, added_at, updated_at) values (?, ?, 'active', ?, ?)", + "insert into accounts (id, access_token, status, added_at, updated_at) values ($1, $2, 'active', $3, $3)", ) .bind(account_id) .bind("access-token") - .bind(&now) - .bind(&now) + .bind(now) .execute(pool) .await .expect("account should be inserted"); } async fn insert_cookie( - pool: &sqlx::SqlitePool, + pool: &sqlx::PgPool, id: &str, account_id: &str, name: &str, expires_at: chrono::DateTime, ) { sqlx::query( - "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values (?, ?, 'chatgpt.com', ?, 'cipher', '/', ?, ?)", + "insert into account_cookies (id, account_id, domain, name, value, path, expires_at, updated_at) values ($1, $2, 'chatgpt.com', $3, 'cipher', '/', $4, $5)", ) .bind(id) .bind(account_id) .bind(name) - .bind(expires_at.to_rfc3339()) - .bind(Utc::now().to_rfc3339()) + .bind(expires_at) + .bind(Utc::now()) .execute(pool) .await .expect("cookie should be inserted"); } -async fn insert_admin_user(pool: &sqlx::SqlitePool, user_id: &str) { - let now = Utc::now().to_rfc3339(); - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, 'hash', ?, ?)", - ) - .bind(user_id) - .bind(&now) - .bind(&now) - .execute(pool) - .await - .expect("admin user should be inserted"); -} - -async fn insert_admin_session( - pool: &sqlx::SqlitePool, - session_id: &str, - user_id: &str, - expires_at: chrono::DateTime, -) { - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind(user_id) - .bind(expires_at.to_rfc3339()) - .bind(Utc::now().to_rfc3339()) - .execute(pool) - .await - .expect("admin session should be inserted"); -} - -async fn cookie_ids(pool: &sqlx::SqlitePool) -> Vec { +async fn cookie_ids(pool: &sqlx::PgPool) -> Vec { sqlx::query_scalar::<_, String>("select id from account_cookies order by id") .fetch_all(pool) .await .expect("cookie ids should load") } - -async fn admin_session_ids(pool: &sqlx::SqlitePool) -> Vec { - sqlx::query_scalar::<_, String>("select id from admin_sessions order by id") - .fetch_all(pool) - .await - .expect("session ids should load") -} - -async fn session_affinity_response_ids(pool: &sqlx::SqlitePool) -> Vec { - sqlx::query_scalar::<_, String>( - "select response_id from session_affinities order by response_id", - ) - .fetch_all(pool) - .await - .expect("session affinity response ids should load") -} - -fn session_affinity_entry( - account_id: &str, - created_at: chrono::DateTime, -) -> SessionAffinityEntry { - SessionAffinityEntry { - account_id: account_id.to_string(), - conversation_id: format!("conversation-{account_id}"), - turn_state: Some(format!("turn-{account_id}")), - instructions_hash: None, - input_tokens: Some(1), - function_call_ids: Vec::new(), - variant_hash: None, - created_at, - } -} diff --git a/backend/tests/runtime/tasks/model_refresh.rs b/backend/tests/bootstrap/tasks/model_refresh.rs similarity index 91% rename from backend/tests/runtime/tasks/model_refresh.rs rename to backend/tests/bootstrap/tasks/model_refresh.rs index 03038c911..e2c587e24 100644 --- a/backend/tests/runtime/tasks/model_refresh.rs +++ b/backend/tests/bootstrap/tasks/model_refresh.rs @@ -5,14 +5,18 @@ use std::{ sync::atomic::{AtomicUsize, Ordering}, }; -use codex_proxy_rs::upstream::{ +use codex_proxy_rs::{ accounts::pool::{AccountAcquireRequest, AccountPoolOptions, RuntimeAccountPoolService}, models::{ - backend_entry::BackendModelEntry, snapshot::ModelPlanSnapshot, - snapshot_store::{ModelSnapshotStore, ModelSnapshotStoreResult}, + store::{ModelSnapshotStore, ModelSnapshotStoreResult}, + }, + upstream::openai::{ + protocol::model_catalog::BackendModelEntry, + transport::{ + CodexModelCatalogClient, CodexModelCatalogClientError, CodexModelCatalogRequest, + }, }, - transport::{CodexModelCatalogClient, CodexModelCatalogClientError, CodexModelCatalogRequest}, }; #[tokio::test] @@ -25,7 +29,7 @@ async fn model_refresh_task_should_start_and_shutdown() { 0, )); - let handle = codex_proxy_rs::runtime::tasks::model_refresh::ModelRefreshTask::new( + let handle = codex_proxy_rs::bootstrap::tasks::model_refresh::ModelRefreshTask::new( model_service, account_pool, ) @@ -36,7 +40,7 @@ async fn model_refresh_task_should_start_and_shutdown() { #[tokio::test] async fn model_refresh_task_should_sync_model_plan_allowlist_to_account_pool() { - let snapshot_store = Arc::new(InMemorySnapshotStore::default()); + let store = Arc::new(InMemorySnapshotStore::default()); let upstream = Arc::new(FakeModelCatalogClient { models_by_plan: BTreeMap::from([ ("plus".to_string(), backend_models("gpt-plus")), @@ -45,7 +49,7 @@ async fn model_refresh_task_should_sync_model_plan_allowlist_to_account_pool() { }); let model_service = Arc::new(ModelService::new( empty_model_config(), - Some(snapshot_store), + Some(store), Some(upstream), )); let account_store = Arc::new(PlanAccountStore { @@ -62,11 +66,11 @@ async fn model_refresh_task_should_sync_model_plan_allowlist_to_account_pool() { 0, )); account_pool - .restore_from_repository() + .restore_from_store() .await .expect("accounts should restore"); - codex_proxy_rs::runtime::tasks::model_refresh::ModelRefreshTask::new( + codex_proxy_rs::bootstrap::tasks::model_refresh::ModelRefreshTask::new( model_service, account_pool.clone(), ) diff --git a/backend/tests/dispatch/affinity.rs b/backend/tests/dispatch/affinity.rs new file mode 100644 index 000000000..2f3324eca --- /dev/null +++ b/backend/tests/dispatch/affinity.rs @@ -0,0 +1,125 @@ +use chrono::{Duration, Utc}; +use codex_proxy_rs::dispatch::affinity::{RedisSessionAffinityStore, SessionAffinityEntry}; + +use crate::support::storage::create_test_redis; + +#[tokio::test] +async fn redis_affinity_store_upserts_reads_and_forgets_entries() { + let store = RedisSessionAffinityStore::new(create_test_redis("affinity-basic").await); + let now = Utc::now(); + let entry = affinity("acct_1", "conv_1", Some("variant_1"), now); + store + .upsert("resp_1", &entry, Duration::hours(4)) + .await + .unwrap(); + + let loaded = store + .get("resp_1", now + Duration::minutes(1), Duration::hours(4)) + .await + .unwrap() + .unwrap(); + assert_eq!(loaded, entry); + assert!(store.forget("resp_1").await.unwrap()); + assert!(store + .get("resp_1", now, Duration::hours(4)) + .await + .unwrap() + .is_none()); +} + +#[tokio::test] +async fn redis_affinity_store_returns_latest_matching_conversation_variant() { + let store = RedisSessionAffinityStore::new(create_test_redis("affinity-latest").await); + let now = Utc::now(); + store + .upsert( + "resp_old", + &affinity( + "acct_old", + "conv_shared", + Some("variant_a"), + now - Duration::minutes(2), + ), + Duration::hours(4), + ) + .await + .unwrap(); + store + .upsert( + "resp_new", + &affinity( + "acct_new", + "conv_shared", + Some("variant_a"), + now - Duration::minutes(1), + ), + Duration::hours(4), + ) + .await + .unwrap(); + store + .upsert( + "resp_other", + &affinity("acct_other", "conv_shared", Some("variant_b"), now), + Duration::hours(4), + ) + .await + .unwrap(); + + let (response_id, entry) = store + .latest_by_conversation( + "conv_shared", + None, + Some("variant_a"), + now, + Duration::hours(4), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(response_id, "resp_new"); + assert_eq!(entry.account_id, "acct_new"); +} + +#[tokio::test] +async fn redis_affinity_store_forgets_all_entries_for_account() { + let store = RedisSessionAffinityStore::new(create_test_redis("affinity-account").await); + let now = Utc::now(); + for response_id in ["resp_1", "resp_2"] { + store + .upsert( + response_id, + &affinity("acct_delete", response_id, None, now), + Duration::hours(4), + ) + .await + .unwrap(); + } + + assert_eq!(store.forget_account("acct_delete").await.unwrap(), 2); + for response_id in ["resp_1", "resp_2"] { + assert!(store + .get(response_id, now, Duration::hours(4)) + .await + .unwrap() + .is_none()); + } +} + +fn affinity( + account_id: &str, + conversation_id: &str, + variant_hash: Option<&str>, + created_at: chrono::DateTime, +) -> SessionAffinityEntry { + SessionAffinityEntry { + account_id: account_id.to_string(), + conversation_id: conversation_id.to_string(), + turn_state: Some(format!("turn_{conversation_id}")), + instructions_hash: Some("instructions".to_string()), + input_tokens: Some(7), + function_call_ids: vec!["call_1".to_string()], + variant_hash: variant_hash.map(ToString::to_string), + created_at, + } +} diff --git a/backend/tests/proxy/dispatch/chat_upstream/compact_routes.rs b/backend/tests/dispatch/chat_upstream/compact_routes.rs similarity index 96% rename from backend/tests/proxy/dispatch/chat_upstream/compact_routes.rs rename to backend/tests/dispatch/chat_upstream/compact_routes.rs index acd7046bb..cb27f8f79 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/compact_routes.rs +++ b/backend/tests/dispatch/chat_upstream/compact_routes.rs @@ -220,12 +220,10 @@ async fn responses_compact_should_post_json_to_codex_compact_upstream() { assert_eq!(event.account_id.as_deref(), Some("acct_chat")); assert_eq!(event.route.as_deref(), Some("/v1/responses/compact")); assert_eq!(event.status_code, Some(200)); - assert_eq!(metadata["route"], "/v1/responses/compact"); - assert_eq!(metadata["apiKind"], "responses"); assert_eq!(metadata["stream"], false); assert_eq!(metadata["compact"], true); - assert_eq!(metadata["usage"]["inputTokens"], 11); - assert_eq!(metadata["usage"]["outputTokens"], 4); + assert_eq!(event.input_tokens, Some(11)); + assert_eq!(event.output_tokens, Some(4)); assert_rate_limit_header(&metadata, "x-ratelimit-limit-requests", "33"); } @@ -284,13 +282,11 @@ async fn responses_compact_should_return_rate_limit_error_when_fallback_is_exhau assert_eq!(event.route.as_deref(), Some("/v1/responses/compact")); assert_eq!(event.status_code, Some(429)); assert_eq!(event.level, "error"); - assert_eq!(metadata["route"], "/v1/responses/compact"); - assert_eq!(metadata["apiKind"], "responses"); assert_eq!(metadata["stream"], false); assert_eq!(metadata["compact"], true); - assert_eq!(metadata["transport"], "http"); + assert_eq!(event.transport.as_deref(), Some("http")); assert_eq!(metadata["failed"], true); - assert_eq!(metadata["failureClass"], "rate_limited"); + assert_eq!(event.failure_class.as_deref(), Some("rate_limited")); assert_eq!(metadata["exhaustedCount"], 1); assert!(metadata["upstreamError"] .as_str() @@ -348,11 +344,10 @@ async fn responses_compact_should_preserve_upstream_client_error_status() { assert_eq!(event.route.as_deref(), Some("/v1/responses/compact")); assert_eq!(event.status_code, Some(400)); assert_eq!(event.level, "error"); - assert_eq!(metadata["route"], "/v1/responses/compact"); assert_eq!(metadata["compact"], true); assert_eq!(metadata["failed"], true); - assert_eq!(metadata["failureClass"], "upstream"); - assert_eq!(metadata["upstreamStatus"], 400); + assert_eq!(event.failure_class.as_deref(), Some("upstream")); + assert_eq!(event.upstream_status_code, Some(400)); assert!(metadata["error"] .as_str() .is_some_and(|value| value.contains("invalid_encrypted_content"))); diff --git a/backend/tests/proxy/dispatch/chat_upstream/mod.rs b/backend/tests/dispatch/chat_upstream/mod.rs similarity index 80% rename from backend/tests/proxy/dispatch/chat_upstream/mod.rs rename to backend/tests/dispatch/chat_upstream/mod.rs index 601784ae9..65f1cfb09 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/mod.rs +++ b/backend/tests/dispatch/chat_upstream/mod.rs @@ -10,33 +10,24 @@ use axum::{ body::{to_bytes, Body, Bytes}, http::{Request, StatusCode}, }; -use chrono::{DateTime, Duration, Utc}; +use chrono::{Duration, Utc}; use codex_proxy_rs::{ - admin::auth::service::SqliteAdminSessionStore, - admin::keys::service::SqliteClientKeyStore, - admin::monitoring::{ - usage_record_model::UsageRecordLevel, - usage_record_store::{SqliteUsageRecordStore, UsageRecordFilter}, - }, - config::schema::AppConfig, - http::router, - infra::database::connect_sqlite, - proxy::dispatch::session_affinity::{ + accounts::{account::AccountStatus, cookies::PgCookieStore, pool::AccountAcquireRequest}, + api::router, + bootstrap::config::AppConfig, + bootstrap::services::{Services, UsageRecordOptions}, + bootstrap::state::AppState, + dispatch::affinity::{ build_conversation_identity, compute_variant_hash, prepare_variant_identity, - SessionAffinityEntry, SqliteSessionAffinityStore, + SessionAffinityEntry, }, - runtime::services::{BackgroundTaskStores, Services, UsageRecordOptions}, - runtime::state::AppState, - upstream::accounts::{ - cookies::SqliteCookieStore, model::AccountStatus, pool::AccountAcquireRequest, - store::SqliteAccountStore, token_refresh::RefreshLeaseStore, - }, - upstream::fingerprint::FingerprintRepository, - upstream::protocol::responses::CodexResponsesRequest, + infra::identity::hash_credential, + telemetry::usage::store::{PgUsageRecordStore, UsageRecordFilter}, + upstream::openai::protocol::responses::CodexResponsesRequest, }; use futures::{SinkExt, StreamExt}; use serde_json::{json, Value}; -use sqlx::SqlitePool; +use sqlx::PgPool; use tokio::{ io::{AsyncReadExt, AsyncWriteExt}, net::{TcpListener, TcpStream}, @@ -62,8 +53,10 @@ use wiremock::{ }; use crate::support::{ - client_keys::insert_client_api_key, config::test_config as base_test_config, + client_keys::insert_client_api_key, + config::test_config as base_test_config, http::response_json, + storage::{background_task_stores, create_test_redis, init_test_db, test_database_url}, }; const TEST_INSTALLATION_ID: &str = "b4f9d503-07b1-457b-a0da-87e6836b1c43"; @@ -113,126 +106,126 @@ where } const RESPONSES_SUCCESS_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/success.sse"), + include_str!("../../fixtures/responses/http_sse/success.sse"), "\n" ); const RESPONSES_COMPLETED_USAGE_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/completed_usage.sse"), + include_str!("../../fixtures/responses/http_sse/completed_usage.sse"), "\n" ); const RESPONSES_COMPLETED_IMAGE_USAGE_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/completed_image_usage.sse"), + include_str!("../../fixtures/responses/http_sse/completed_image_usage.sse"), "\n" ); const RESPONSES_EMPTY_COMPLETED_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/empty_completed.sse"), + include_str!("../../fixtures/responses/http_sse/empty_completed.sse"), "\n" ); const RESPONSES_TEXT_DELTAS_COMPLETED_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/text_deltas_completed.sse"), + include_str!("../../fixtures/responses/http_sse/text_deltas_completed.sse"), "\n" ); const RESPONSES_DONE_ITEM_COMPLETED_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/done_item_completed.sse"), + include_str!("../../fixtures/responses/http_sse/done_item_completed.sse"), "\n" ); const RESPONSES_TUPLE_OBJECT_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/tuple_object.sse"), + include_str!("../../fixtures/responses/http_sse/tuple_object.sse"), "\n" ); const RESPONSES_AFTER_5XX_RETRY_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/after_5xx_retry.sse"), + include_str!("../../fixtures/responses/http_sse/after_5xx_retry.sse"), "\n" ); const RESPONSES_AFTER_402_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/after_402.sse"), + include_str!("../../fixtures/responses/http_sse/after_402.sse"), "\n" ); const RESPONSES_FAILED_QUOTA_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/failed_quota.sse"), + include_str!("../../fixtures/responses/http_sse/failed_quota.sse"), "\n" ); const RESPONSES_FAILED_AUTH_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/failed_auth.sse"), + include_str!("../../fixtures/responses/http_sse/failed_auth.sse"), "\n" ); const RESPONSES_AFTER_401_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/after_401.sse"), + include_str!("../../fixtures/responses/http_sse/after_401.sse"), "\n" ); const RESPONSES_AFTER_403_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/after_403.sse"), + include_str!("../../fixtures/responses/http_sse/after_403.sse"), "\n" ); const RESPONSES_STREAM_USAGE_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/stream_usage.sse"), + include_str!("../../fixtures/responses/http_sse/stream_usage.sse"), "\n" ); const RESPONSES_INCOMPLETE_USAGE_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"), + include_str!("../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"), "\n" ); const RESPONSES_STREAM_AFTER_429_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/stream_after_429.sse"), + include_str!("../../fixtures/responses/http_sse/stream_after_429.sse"), "\n" ); const RESPONSES_STREAM_AFTER_5XX_RETRY_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/stream_after_5xx_retry.sse"), + include_str!("../../fixtures/responses/http_sse/stream_after_5xx_retry.sse"), "\n" ); const RESPONSES_STREAM_AFTER_402_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/stream_after_402.sse"), + include_str!("../../fixtures/responses/http_sse/stream_after_402.sse"), "\n" ); const RESPONSES_STREAM_AFTER_401_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/stream_after_401.sse"), + include_str!("../../fixtures/responses/http_sse/stream_after_401.sse"), "\n" ); const RESPONSES_STREAM_AFTER_403_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/stream_after_403.sse"), + include_str!("../../fixtures/responses/http_sse/stream_after_403.sse"), "\n" ); const RESPONSES_AFTER_CLOUDFLARE_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/after_cloudflare.sse"), + include_str!("../../fixtures/responses/http_sse/after_cloudflare.sse"), "\n" ); const RESPONSES_STREAM_AFTER_CLOUDFLARE_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/stream_after_cloudflare.sse"), + include_str!("../../fixtures/responses/http_sse/stream_after_cloudflare.sse"), "\n" ); const RESPONSES_AFTER_MODEL_UNSUPPORTED_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/after_model_unsupported.sse"), + include_str!("../../fixtures/responses/http_sse/after_model_unsupported.sse"), "\n" ); const RESPONSES_STREAM_AFTER_MODEL_UNSUPPORTED_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/stream_after_model_unsupported.sse"), + include_str!("../../fixtures/responses/http_sse/stream_after_model_unsupported.sse"), "\n" ); const RESPONSES_FAILED_MODEL_UNSUPPORTED_SSE: &str = concat!( - include_str!("../../../fixtures/responses/http_sse/failed_model_unsupported.sse"), + include_str!("../../fixtures/responses/http_sse/failed_model_unsupported.sse"), "\n" ); const WEBSOCKET_COMPLETED_WITH_REASONING_REPLAY: &str = - include_str!("../../../fixtures/responses/websocket/completed_with_reasoning_replay.json"); + include_str!("../../fixtures/responses/websocket/completed_with_reasoning_replay.json"); const WEBSOCKET_HISTORY_RATE_LIMITED: &str = - include_str!("../../../fixtures/responses/websocket/history_rate_limited.json"); + include_str!("../../fixtures/responses/websocket/history_rate_limited.json"); const WEBSOCKET_RATE_LIMITED: &str = - include_str!("../../../fixtures/responses/websocket/rate_limited.json"); + include_str!("../../fixtures/responses/websocket/rate_limited.json"); const WEBSOCKET_TOKEN_REVOKED: &str = - include_str!("../../../fixtures/responses/websocket/token_revoked.json"); + include_str!("../../fixtures/responses/websocket/token_revoked.json"); const WEBSOCKET_FIRST_ACCOUNT_LIMITED: &str = - include_str!("../../../fixtures/responses/websocket/first_account_limited.json"); + include_str!("../../fixtures/responses/websocket/first_account_limited.json"); const WEBSOCKET_SECOND_ACCOUNT_LIMITED: &str = - include_str!("../../../fixtures/responses/websocket/second_account_limited.json"); + include_str!("../../fixtures/responses/websocket/second_account_limited.json"); const WEBSOCKET_INVALID_ENCRYPTED_CONTENT: &str = - include_str!("../../../fixtures/responses/websocket/invalid_encrypted_content.json"); + include_str!("../../fixtures/responses/websocket/invalid_encrypted_content.json"); const WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND: &str = - include_str!("../../../fixtures/responses/websocket/previous_response_not_found.json"); + include_str!("../../fixtures/responses/websocket/previous_response_not_found.json"); const WEBSOCKET_UNANSWERED_FUNCTION_CALL: &str = - include_str!("../../../fixtures/responses/websocket/unanswered_function_call.json"); + include_str!("../../fixtures/responses/websocket/unanswered_function_call.json"); const REASONING_REPLAY_REQUEST_GOLDEN: &str = - include_str!("../../../fixtures/responses/golden/reasoning_replay_request.json"); + include_str!("../../fixtures/responses/golden/reasoning_replay_request.json"); mod compact_routes; mod responses_http; @@ -242,7 +235,7 @@ mod usage_logging; async fn test_app_state_with_pool_and_installation_id( config: &AppConfig, - pool: SqlitePool, + pool: PgPool, installation_id: String, ) -> AppState { test_app_state_with_pool_installation_id_and_usage_record_options( @@ -256,20 +249,23 @@ async fn test_app_state_with_pool_and_installation_id( async fn test_app_state_with_pool_installation_id_and_usage_record_options( config: AppConfig, - pool: SqlitePool, + pool: PgPool, installation_id: String, usage_record_options: UsageRecordOptions, ) -> AppState { - let stores = BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: SqliteSessionAffinityStore::new(pool.clone()), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool), - }; + let redis = create_test_redis("chat-upstream").await; + test_app_state_with_storage(config, pool, redis, installation_id, usage_record_options).await +} + +async fn test_app_state_with_storage( + config: AppConfig, + pool: PgPool, + redis: codex_proxy_rs::infra::redis::RedisConnection, + installation_id: String, + usage_record_options: UsageRecordOptions, +) -> AppState { + insert_model_snapshot(&redis).await; + let stores = background_task_stores(pool, redis); let services = Services::try_with_installation_id_and_usage_record_options( &config, stores, @@ -293,14 +289,11 @@ async fn test_app_with_account_and_installation_id( base_url: String, installation_id: String, ) -> (axum::Router, String, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("openai-chat-upstream.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); + let (pool, dir) = init_test_db("openai-chat-upstream").await; let api_key = insert_client_api_key(&pool).await; insert_account(&pool).await; let state = test_app_state_with_pool_and_installation_id( - &test_config(url, base_url), + &test_config(test_database_url(), base_url), pool, installation_id, ) @@ -308,7 +301,7 @@ async fn test_app_with_account_and_installation_id( state .services .account_pool - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore"); (router::router().with_state(state), api_key, dir) @@ -316,28 +309,24 @@ async fn test_app_with_account_and_installation_id( async fn test_app_with_account_and_pool( base_url: String, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { +) -> (axum::Router, String, PgPool, tempfile::TempDir) { test_app_with_account_pool_config(base_url, |_| {}).await } async fn test_app_with_account_pool_and_affinity( base_url: String, response_id: &str, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { +) -> (axum::Router, String, PgPool, tempfile::TempDir) { test_app_with_account_pool_config_and_affinity(base_url, |_| {}, Some(response_id)).await } async fn test_app_without_accounts( base_url: String, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("openai-responses-no-accounts.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); +) -> (axum::Router, String, PgPool, tempfile::TempDir) { + let (pool, dir) = init_test_db("openai-responses-no-accounts").await; let api_key = insert_client_api_key(&pool).await; - insert_model_snapshot(&pool).await; let state = test_app_state_with_pool_and_installation_id( - &test_config(url, base_url), + &test_config(test_database_url(), base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), ) @@ -345,7 +334,7 @@ async fn test_app_without_accounts( state .services .account_pool - .restore_from_repository() + .restore_from_store() .await .expect("empty account pool should restore"); (router::router().with_state(state), api_key, pool, dir) @@ -353,7 +342,7 @@ async fn test_app_without_accounts( async fn test_app_with_account_pool_and_logging( base_url: String, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { +) -> (axum::Router, String, PgPool, tempfile::TempDir) { test_app_with_account_pool_config(base_url, |config| { config.logging.enabled = true; }) @@ -362,15 +351,12 @@ async fn test_app_with_account_pool_and_logging( async fn test_app_with_account_pool_and_disabled_logging( base_url: String, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("openai-record-affinity.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); +) -> (axum::Router, String, PgPool, tempfile::TempDir) { + let (pool, dir) = init_test_db("openai-record-disabled-logging").await; let api_key = insert_client_api_key(&pool).await; insert_account(&pool).await; let state = test_app_state_with_pool_and_installation_id( - &test_config(url, base_url), + &test_config(test_database_url(), base_url), pool.clone(), TEST_INSTALLATION_ID.to_string(), ) @@ -378,7 +364,7 @@ async fn test_app_with_account_pool_and_disabled_logging( state .services .account_pool - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore"); (router::router().with_state(state), api_key, pool, dir) @@ -386,14 +372,11 @@ async fn test_app_with_account_pool_and_disabled_logging( async fn test_app_with_account_pool_and_logging_capture_body( base_url: String, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("openai-record-affinity.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); +) -> (axum::Router, String, PgPool, tempfile::TempDir) { + let (pool, dir) = init_test_db("openai-record-capture-body").await; let api_key = insert_client_api_key(&pool).await; insert_account(&pool).await; - let mut config = test_config(url, base_url); + let mut config = test_config(test_database_url(), base_url); config.logging.enabled = true; let usage_record_options = UsageRecordOptions { enabled: true, @@ -409,7 +392,7 @@ async fn test_app_with_account_pool_and_logging_capture_body( state .services .account_pool - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore"); (router::router().with_state(state), api_key, pool, dir) @@ -418,7 +401,7 @@ async fn test_app_with_account_pool_and_logging_capture_body( async fn test_app_with_account_pool_config( base_url: String, configure: impl FnOnce(&mut AppConfig), -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { +) -> (axum::Router, String, PgPool, tempfile::TempDir) { test_app_with_account_pool_config_and_affinity(base_url, configure, None).await } @@ -426,17 +409,11 @@ async fn test_app_with_account_pool_config_and_affinity( base_url: String, configure: impl FnOnce(&mut AppConfig), affinity_response_id: Option<&str>, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("openai-record-affinity.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); +) -> (axum::Router, String, PgPool, tempfile::TempDir) { + let (pool, dir) = init_test_db("openai-record-affinity").await; let api_key = insert_client_api_key(&pool).await; insert_account(&pool).await; - if let Some(response_id) = affinity_response_id { - insert_session_affinity(&pool, response_id, "acct_chat").await; - } - let mut config = test_config(url, base_url); + let mut config = test_config(test_database_url(), base_url); configure(&mut config); let state = test_app_state_with_pool_and_installation_id( &config, @@ -444,35 +421,41 @@ async fn test_app_with_account_pool_config_and_affinity( TEST_INSTALLATION_ID.to_string(), ) .await; + if let Some(response_id) = affinity_response_id { + insert_session_affinity( + &state.services.background_tasks.session_affinity, + response_id, + "acct_chat", + ) + .await; + } state .services .account_pool - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore"); - if affinity_response_id.is_some() { - state - .services - .session_affinity - .restore_from_repository(Utc::now()) - .await - .expect("session affinity should restore"); - } (router::router().with_state(state), api_key, pool, dir) } async fn test_app_with_two_accounts_and_affinity( base_url: String, ) -> (axum::Router, String, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("openai-responses-affinity.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); + let (pool, dir) = init_test_db("openai-responses-affinity").await; let api_key = insert_client_api_key(&pool).await; insert_named_account(&pool, "acct_a", "access-default", "chatgpt-default").await; insert_named_account(&pool, "acct_z", "access-affinity", "chatgpt-affinity").await; let now = Utc::now(); - SqliteSessionAffinityStore::new(pool.clone()) + let state = test_app_state_with_pool_and_installation_id( + &test_config(test_database_url(), base_url), + pool, + TEST_INSTALLATION_ID.to_string(), + ) + .await; + state + .services + .background_tasks + .session_affinity .upsert( "resp_previous", &SessionAffinityEntry { @@ -489,37 +472,23 @@ async fn test_app_with_two_accounts_and_affinity( ) .await .unwrap(); - let state = test_app_state_with_pool_and_installation_id( - &test_config(url, base_url), - pool, - TEST_INSTALLATION_ID.to_string(), - ) - .await; state .services .account_pool - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore"); - state - .services - .session_affinity - .restore_from_repository(now + Duration::minutes(1)) - .await - .expect("session affinity should restore"); (router::router().with_state(state), api_key, dir) } async fn test_app_with_two_accounts_and_affinity_status( base_url: String, affinity_account_status: &str, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join(format!( - "openai-responses-affinity-{affinity_account_status}.sqlite" - )); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); +) -> (axum::Router, AppState, String, PgPool, tempfile::TempDir) { + let (pool, dir) = init_test_db(&format!( + "openai-responses-affinity-{affinity_account_status}" + )) + .await; let api_key = insert_client_api_key(&pool).await; insert_named_account(&pool, "acct_primary", "access-primary", "chatgpt-primary").await; insert_named_account( @@ -529,14 +498,23 @@ async fn test_app_with_two_accounts_and_affinity_status( "chatgpt-affinity", ) .await; - sqlx::query("update accounts set status = ? where id = ?") + sqlx::query("update accounts set status = $1 where id = $2") .bind(affinity_account_status) .bind("acct_affinity") .execute(&pool) .await .unwrap(); let now = Utc::now(); - SqliteSessionAffinityStore::new(pool.clone()) + let state = test_app_state_with_pool_and_installation_id( + &test_config(test_database_url(), base_url), + pool.clone(), + TEST_INSTALLATION_ID.to_string(), + ) + .await; + state + .services + .background_tasks + .session_affinity .upsert( "resp_affinity_risk", &SessionAffinityEntry { @@ -553,92 +531,10 @@ async fn test_app_with_two_accounts_and_affinity_status( ) .await .unwrap(); - let state = test_app_state_with_pool_and_installation_id( - &test_config(url, base_url), - pool.clone(), - TEST_INSTALLATION_ID.to_string(), - ) - .await; state .services .account_pool - .restore_from_repository() - .await - .expect("account pool should restore"); - state - .services - .session_affinity - .restore_from_repository(now + Duration::minutes(1)) - .await - .expect("session affinity should restore"); - (router::router().with_state(state), api_key, pool, dir) -} - -async fn test_app_with_two_accounts( - base_url: String, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { - let (app, _state, api_key, pool, dir) = test_app_with_two_accounts_and_state(base_url).await; - (app, api_key, pool, dir) -} - -async fn test_app_with_two_accounts_and_logging( - base_url: String, -) -> (axum::Router, String, SqlitePool, tempfile::TempDir) { - let (app, _state, api_key, pool, dir) = - test_app_with_two_accounts_and_state_config(base_url, |config| { - config.logging.enabled = true; - }) - .await; - (app, api_key, pool, dir) -} - -async fn test_app_with_two_accounts_and_state( - base_url: String, -) -> ( - axum::Router, - AppState, - String, - SqlitePool, - tempfile::TempDir, -) { - test_app_with_two_accounts_and_state_config(base_url, |_| {}).await -} - -async fn test_app_with_two_accounts_and_state_config( - base_url: String, - configure: impl FnOnce(&mut AppConfig), -) -> ( - axum::Router, - AppState, - String, - SqlitePool, - tempfile::TempDir, -) { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("openai-responses-fallback.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let api_key = insert_client_api_key(&pool).await; - insert_named_account(&pool, "acct_primary", "access-primary", "chatgpt-primary").await; - insert_named_account( - &pool, - "acct_secondary", - "access-secondary", - "chatgpt-secondary", - ) - .await; - let mut config = test_config(url, base_url); - configure(&mut config); - let state = test_app_state_with_pool_and_installation_id( - &config, - pool.clone(), - TEST_INSTALLATION_ID.to_string(), - ) - .await; - state - .services - .account_pool - .restore_from_repository() + .restore_from_store() .await .expect("account pool should restore"); ( @@ -650,24 +546,79 @@ async fn test_app_with_two_accounts_and_state_config( ) } -async fn seed_openai_admin_session(pool: &SqlitePool, session_id: &str) { - sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", +async fn test_app_with_two_accounts( + base_url: String, +) -> (axum::Router, String, PgPool, tempfile::TempDir) { + let (app, _state, api_key, pool, dir) = test_app_with_two_accounts_and_state(base_url).await; + (app, api_key, pool, dir) +} + +async fn test_app_with_two_accounts_and_logging( + base_url: String, +) -> (axum::Router, String, PgPool, tempfile::TempDir) { + let (app, _state, api_key, pool, dir) = + test_app_with_two_accounts_and_state_config(base_url, |config| { + config.logging.enabled = true; + }) + .await; + (app, api_key, pool, dir) +} + +async fn test_app_with_two_accounts_and_state( + base_url: String, +) -> (axum::Router, AppState, String, PgPool, tempfile::TempDir) { + test_app_with_two_accounts_and_state_config(base_url, |_| {}).await +} + +async fn test_app_with_two_accounts_and_state_config( + base_url: String, + configure: impl FnOnce(&mut AppConfig), +) -> (axum::Router, AppState, String, PgPool, tempfile::TempDir) { + let (pool, dir) = init_test_db("openai-responses-fallback").await; + let api_key = insert_client_api_key(&pool).await; + insert_named_account(&pool, "acct_primary", "access-primary", "chatgpt-primary").await; + insert_named_account( + &pool, + "acct_secondary", + "access-secondary", + "chatgpt-secondary", ) - .bind("admin_openai") - .bind("hash") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", + .await; + let mut config = test_config(test_database_url(), base_url); + configure(&mut config); + let state = test_app_state_with_pool_and_installation_id( + &config, + pool.clone(), + TEST_INSTALLATION_ID.to_string(), ) - .bind(session_id) - .bind("admin_openai") - .bind("2999-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") + .await; + state + .services + .account_pool + .restore_from_store() + .await + .expect("account pool should restore"); + ( + router::router().with_state(state.clone()), + state, + api_key, + pool, + dir, + ) +} + +async fn seed_openai_admin_key(pool: &PgPool, key: &str) { + sqlx::query( + "insert into runtime_settings ( + id, refresh_margin_seconds, refresh_concurrency, + max_concurrent_per_account, request_interval_ms, + rotation_strategy, admin_api_key_hash, updated_at + ) values (1, 3600, 2, 3, 0, 'smart', $1, now()) + on conflict (id) do update + set admin_api_key_hash = excluded.admin_api_key_hash, + updated_at = excluded.updated_at", + ) + .bind(hash_credential(key)) .execute(pool) .await .unwrap(); @@ -681,7 +632,7 @@ async fn update_admin_account_status(app: &axum::Router, account_id: &str, statu .method("POST") .uri("/api/admin/accounts/update") .header("content-type", "application/json") - .header("cookie", "cpr_admin_session=session_status_cycle") + .header("x-api-key", "admin-status-cycle") .body(Body::from( json!({"id": account_id, "status": status}).to_string(), )) @@ -692,27 +643,27 @@ async fn update_admin_account_status(app: &axum::Router, account_id: &str, statu assert_eq!(response.status(), StatusCode::OK); } -async fn insert_account(pool: &SqlitePool) { - insert_model_snapshot(pool).await; +async fn insert_account(pool: &PgPool) { + let now = Utc::now(); sqlx::query( - "insert into accounts (id, email, chatgpt_account_id, chatgpt_user_id, access_token, access_token_expires_at, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + "insert into accounts (id, email, chatgpt_account_id, chatgpt_user_id, access_token, access_token_expires_at, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7, $8, $9)", ) .bind("acct_chat") .bind("user@example.com") .bind("chatgpt-account") .bind("chatgpt-user") .bind("access-secret") - .bind("2100-01-01T00:00:00Z") + .bind(now + Duration::days(36500)) .bind("active") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") + .bind(now) + .bind(now) .execute(pool) .await .unwrap(); } -async fn insert_model_snapshot(pool: &SqlitePool) { - let models_json = json!([ +async fn insert_model_snapshot(redis: &codex_proxy_rs::infra::redis::RedisConnection) { + let models = json!([ { "id": "gpt-5.5", "displayName": "GPT 5.5", @@ -730,21 +681,24 @@ async fn insert_model_snapshot(pool: &SqlitePool) { "upgrade": null, "source": "test" } - ]) - .to_string(); - sqlx::query( - "insert or replace into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?)", - ) - .bind("plus") - .bind(models_json) - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); + ]); + let value = json!({"models": models, "fetchedAt": Utc::now()}); + let mut connection = redis.manager(); + let _: usize = redis::cmd("HSET") + .arg(redis.key("models:plan_snapshots")) + .arg("plus") + .arg(value.to_string()) + .query_async(&mut connection) + .await + .unwrap(); } -async fn insert_session_affinity(pool: &SqlitePool, response_id: &str, account_id: &str) { - SqliteSessionAffinityStore::new(pool.clone()) +async fn insert_session_affinity( + store: &codex_proxy_rs::dispatch::affinity::RedisSessionAffinityStore, + response_id: &str, + account_id: &str, +) { + store .upsert( response_id, &SessionAffinityEntry { @@ -764,24 +718,24 @@ async fn insert_session_affinity(pool: &SqlitePool, response_id: &str, account_i } async fn insert_named_account( - pool: &SqlitePool, + pool: &PgPool, id: &str, access_token: &str, chatgpt_account_id: &str, ) { - insert_model_snapshot(pool).await; + let now = Utc::now(); sqlx::query( - "insert into accounts (id, email, chatgpt_account_id, chatgpt_user_id, access_token, access_token_expires_at, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)", + "insert into accounts (id, email, chatgpt_account_id, chatgpt_user_id, access_token, access_token_expires_at, status, added_at, updated_at) values ($1, $2, $3, $4, $5, $6, $7, $8, $9)", ) .bind(id) .bind(format!("{id}@example.com")) .bind(chatgpt_account_id) .bind(format!("user-{id}")) .bind(access_token) - .bind("2100-01-01T00:00:00Z") + .bind(now + Duration::days(36500)) .bind("active") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") + .bind(now) + .bind(now) .execute(pool) .await .unwrap(); @@ -1407,39 +1361,24 @@ async fn response_text(response: axum::response::Response) -> String { String::from_utf8(bytes.to_vec()).unwrap() } -async fn wait_for_session_affinity_turn_state( - pool: &SqlitePool, - response_id: &str, -) -> Option { - timeout(StdDuration::from_secs(1), async { - loop { - if let Some(turn_state) = sqlx::query_scalar::<_, Option>( - "select turn_state from session_affinities where response_id = ?", - ) - .bind(response_id) - .fetch_optional(pool) - .await - .unwrap() - .flatten() - { - return Some(turn_state); - } - tokio::time::sleep(StdDuration::from_millis(10)).await; - } - }) - .await - .ok() - .flatten() -} - struct ResponseUsageRecordSnapshot { level: String, request_id: Option, account_id: Option, route: Option, model: Option, + requested_model: Option, + upstream_model: Option, + transport: Option, status_code: Option, + client_status_code: Option, + upstream_status_code: Option, + failure_class: Option, + attempt_index: Option, response_id: Option, + first_token_ms: Option, + input_tokens: Option, + output_tokens: Option, metadata_json: String, } @@ -1449,7 +1388,12 @@ struct OpsErrorLogSnapshot { account_id: Option, route: Option, model: Option, - status_code: Option, + transport: Option, + status_code: Option, + client_status_code: Option, + upstream_status_code: Option, + failure_class: Option, + attempt_index: Option, response_id: Option, metadata_json: String, } @@ -1462,22 +1406,34 @@ impl From for ResponseUsageRecordSnapshot { account_id: snapshot.account_id, route: snapshot.route, model: snapshot.model, - status_code: snapshot.status_code, + requested_model: None, + upstream_model: None, + transport: snapshot.transport, + status_code: snapshot.status_code.map(i64::from), + client_status_code: snapshot.client_status_code.map(i64::from), + upstream_status_code: snapshot.upstream_status_code.map(i64::from), + failure_class: snapshot.failure_class, + attempt_index: snapshot.attempt_index, response_id: snapshot.response_id, + first_token_ms: None, + input_tokens: None, + output_tokens: None, metadata_json: snapshot.metadata_json, } } } -async fn latest_response_usage_record(pool: &SqlitePool) -> ResponseUsageRecordSnapshot { +async fn latest_response_usage_record(pool: &PgPool) -> ResponseUsageRecordSnapshot { latest_usage_record(pool, "v1.response").await } -async fn latest_response_ops_error_log(pool: &SqlitePool) -> ResponseUsageRecordSnapshot { +async fn latest_response_ops_error_log(pool: &PgPool) -> ResponseUsageRecordSnapshot { let row = sqlx::query_as::<_, OpsErrorLogSnapshot>( - "select request_id, account_id, route, model, status_code, response_id, metadata_json + "select request_id, account_id, route, model, transport, status_code, + client_status_code, upstream_status_code, failure_class, attempt_index, response_id, + metadata_json::text as metadata_json from ops_error_logs - where kind = ? + where kind = $1 order by created_at desc, id desc limit 1", ) @@ -1489,8 +1445,8 @@ async fn latest_response_ops_error_log(pool: &SqlitePool) -> ResponseUsageRecord .into() } -async fn latest_usage_record(pool: &SqlitePool, kind: &str) -> ResponseUsageRecordSnapshot { - let page = SqliteUsageRecordStore::new(pool.clone()) +async fn latest_usage_record(pool: &PgPool, kind: &str) -> ResponseUsageRecordSnapshot { + let page = PgUsageRecordStore::new(pool.clone()) .list( UsageRecordFilter { kind: Some(kind.to_string()), @@ -1507,19 +1463,29 @@ async fn latest_usage_record(pool: &SqlitePool, kind: &str) -> ResponseUsageReco .next() .unwrap_or_else(|| panic!("expected a {kind} usage record")); ResponseUsageRecordSnapshot { - level: usage_record_level_name(event.level).to_string(), + level: "info".to_string(), request_id: event.request_id, - account_id: event.account_id, + account_id: Some(event.account_id), route: event.route, - model: event.model, - status_code: event.status_code, + model: Some(event.model), + requested_model: event.requested_model, + upstream_model: event.upstream_model, + transport: event.transport, + status_code: Some(event.status_code), + client_status_code: None, + upstream_status_code: None, + failure_class: None, + attempt_index: event.attempt_index, response_id: event.response_id, + first_token_ms: event.first_token_ms, + input_tokens: event.input_tokens, + output_tokens: event.output_tokens, metadata_json: event.metadata.to_string(), } } -async fn response_usage_record_count(pool: &SqlitePool) -> i64 { - let (count,): (i64,) = sqlx::query_as("select count(*) from usage_records where kind = ?") +async fn response_usage_record_count(pool: &PgPool) -> i64 { + let (count,): (i64,) = sqlx::query_as("select count(*) from usage_records where kind = $1") .bind("v1.response") .fetch_one(pool) .await @@ -1527,15 +1493,6 @@ async fn response_usage_record_count(pool: &SqlitePool) -> i64 { count } -fn usage_record_level_name(level: UsageRecordLevel) -> &'static str { - match level { - UsageRecordLevel::Debug => "debug", - UsageRecordLevel::Info => "info", - UsageRecordLevel::Warn => "warn", - UsageRecordLevel::Error => "error", - } -} - fn assert_rate_limit_header(metadata: &Value, name: &str, value: &str) { let headers = metadata["rateLimitHeaders"] .as_array() diff --git a/backend/tests/proxy/dispatch/chat_upstream/responses_http.rs b/backend/tests/dispatch/chat_upstream/responses_http.rs similarity index 97% rename from backend/tests/proxy/dispatch/chat_upstream/responses_http.rs rename to backend/tests/dispatch/chat_upstream/responses_http.rs index f53192e2c..f4ad92fc5 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/responses_http.rs +++ b/backend/tests/dispatch/chat_upstream/responses_http.rs @@ -298,7 +298,7 @@ async fn responses_non_stream_should_record_websocket_transport_metadata() { let event = latest_response_usage_record(&pool).await; let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); assert_eq!(metadata["stream"], false); - assert_eq!(metadata["transport"], "websocket"); + assert_eq!(event.transport.as_deref(), Some("websocket")); assert_eq!(metadata["websocketPool"]["kind"], "new"); } @@ -865,7 +865,7 @@ async fn responses_should_scope_upstream_cookie_by_codex_response_path() { .await; let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone()); + let cookie_store = PgCookieStore::new(pool.clone()); cookie_store .capture_set_cookie( "acct_chat", @@ -922,7 +922,7 @@ async fn responses_stream_should_close_http_sse_upstream_when_client_disconnects write_chunked_http_sse_headers(&mut socket).await; write_http_chunk( &mut socket, - include_bytes!("../../../fixtures/responses/http_sse/text_delta_hello.sse"), + include_bytes!("../../fixtures/responses/http_sse/text_delta_hello.sse"), ) .await; socket.flush().await.unwrap(); @@ -978,8 +978,8 @@ async fn responses_stream_should_close_http_sse_upstream_when_client_disconnects #[tokio::test] async fn responses_stream_should_forward_first_chunk_before_upstream_completes() { let (base_url, first_chunk_sent, finish_upstream) = spawn_chunked_sse_upstream( - include_str!("../../../fixtures/responses/http_sse/live_stream_hello_delta.sse"), - include_str!("../../../fixtures/responses/http_sse/live_stream_completed.sse"), + include_str!("../../fixtures/responses/http_sse/live_stream_hello_delta.sse"), + include_str!("../../fixtures/responses/http_sse/live_stream_completed.sse"), ) .await; @@ -1038,7 +1038,7 @@ async fn responses_stream_should_forward_first_chunk_before_upstream_completes() } let rest = String::from_utf8(rest).unwrap(); let usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_chat") .fetch_one(&pool) @@ -1058,7 +1058,7 @@ async fn responses_stream_should_emit_failed_event_after_upstream_read_error_onc ) { let (base_url, first_chunk_sent, close_upstream) = spawn_chunked_sse_upstream_then_abrupt_close(include_str!( - "../../../fixtures/responses/http_sse/partial_transport_failure.sse" + "../../fixtures/responses/http_sse/partial_transport_failure.sse" )) .await; @@ -1111,7 +1111,7 @@ async fn responses_stream_should_emit_failed_event_after_upstream_read_error_onc #[tokio::test] async fn responses_stream_should_emit_failed_event_when_upstream_closes_without_completed() { let (base_url, first_chunk_sent, close_upstream) = spawn_chunked_sse_upstream_then_clean_close( - include_str!("../../../fixtures/responses/http_sse/partial_clean_close.sse"), + include_str!("../../fixtures/responses/http_sse/partial_clean_close.sse"), ) .await; @@ -1204,7 +1204,7 @@ async fn responses_should_prefer_session_affinity_account_for_previous_response( async fn responses_should_strip_banned_affinity_history_when_switching_to_active_account() { let (base_url, upstream) = spawn_single_websocket_completed_upstream("resp_after_banned_affinity").await; - let (app, api_key, pool, _dir) = + let (app, state, api_key, _pool, _dir) = test_app_with_two_accounts_and_affinity_status(base_url, "banned").await; let response = app .oneshot( @@ -1227,12 +1227,13 @@ async fn responses_should_strip_banned_affinity_history_when_switching_to_active .await .unwrap(); let captured = upstream.await.unwrap(); - let affinity_count: (i64,) = - sqlx::query_as("select count(*) from session_affinities where response_id = ?") - .bind("resp_affinity_risk") - .fetch_one(&pool) - .await - .unwrap(); + let affinity = state + .services + .background_tasks + .session_affinity + .get("resp_affinity_risk", Utc::now(), chrono::Duration::hours(4)) + .await + .unwrap(); assert_eq!(response.status(), StatusCode::OK); assert_eq!( @@ -1245,14 +1246,14 @@ async fn responses_should_strip_banned_affinity_history_when_switching_to_active ); assert!(captured_header(&captured.headers, "x-codex-turn-state").is_none()); assert!(captured.payload.get("previous_response_id").is_none()); - assert_eq!(affinity_count.0, 0); + assert!(affinity.is_none()); } #[tokio::test] async fn responses_should_strip_quota_history_when_switching_to_active_account() { let (base_url, upstream) = spawn_single_websocket_completed_upstream("resp_after_quota_affinity").await; - let (app, api_key, pool, _dir) = + let (app, state, api_key, _pool, _dir) = test_app_with_two_accounts_and_affinity_status(base_url, "quota_exhausted").await; let response = app .oneshot( @@ -1275,12 +1276,13 @@ async fn responses_should_strip_quota_history_when_switching_to_active_account() .await .unwrap(); let captured = upstream.await.unwrap(); - let affinity_count: (i64,) = - sqlx::query_as("select count(*) from session_affinities where response_id = ?") - .bind("resp_affinity_risk") - .fetch_one(&pool) - .await - .unwrap(); + let affinity = state + .services + .background_tasks + .session_affinity + .get("resp_affinity_risk", Utc::now(), chrono::Duration::hours(4)) + .await + .unwrap(); assert_eq!(response.status(), StatusCode::OK); assert_eq!( @@ -1289,5 +1291,5 @@ async fn responses_should_strip_quota_history_when_switching_to_active_account() ); assert!(captured_header(&captured.headers, "x-codex-turn-state").is_none()); assert!(captured.payload.get("previous_response_id").is_none()); - assert_eq!(affinity_count.0, 1); + assert!(affinity.is_some()); } diff --git a/backend/tests/proxy/dispatch/chat_upstream/responses_recovery.rs b/backend/tests/dispatch/chat_upstream/responses_recovery.rs similarity index 97% rename from backend/tests/proxy/dispatch/chat_upstream/responses_recovery.rs rename to backend/tests/dispatch/chat_upstream/responses_recovery.rs index ebb55035e..caf7e45d3 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/responses_recovery.rs +++ b/backend/tests/dispatch/chat_upstream/responses_recovery.rs @@ -55,7 +55,7 @@ async fn responses_should_classify_sse_cyber_policy_failure_as_bad_request() { ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") .set_body_string(include_str!( - "../../../fixtures/responses/http_sse/failed_cyber_policy.sse" + "../../fixtures/responses/http_sse/failed_cyber_policy.sse" )), ) .mount(&server) @@ -138,7 +138,7 @@ async fn responses_should_return_no_available_accounts_error_when_no_accounts_ar .await .unwrap(); let (bucket_requests, bucket_errors): (i64, i64) = sqlx::query_as( - "select coalesce(sum(request_count), 0), coalesce(sum(error_count), 0) from usage_time_buckets", + "select coalesce(sum(success_count + error_count), 0)::bigint, coalesce(sum(error_count), 0)::bigint from request_time_buckets", ) .fetch_one(&pool) .await @@ -210,14 +210,14 @@ async fn responses_should_fallback_to_next_account_after_rate_limit() { assert_eq!(status, StatusCode::OK); assert_eq!(body["id"], "resp_response_1"); let primary_usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let secondary_usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_secondary") .fetch_one(&pool) @@ -274,7 +274,7 @@ async fn responses_should_mark_expired_after_401_and_fallback() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await @@ -334,7 +334,7 @@ async fn responses_should_recover_auth_failure_from_sse_failed_event() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await @@ -382,7 +382,7 @@ async fn responses_should_return_auth_error_when_401_fallback_is_exhausted() { let status = response.status(); let body = response_json(response).await; let message = body["error"]["message"].as_str().unwrap_or_default(); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -430,7 +430,7 @@ async fn responses_should_mark_banned_when_401_says_account_deactivated() { ) .await .unwrap(); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -487,7 +487,7 @@ async fn responses_should_mark_banned_after_403_and_fallback() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await @@ -512,7 +512,7 @@ async fn responses_should_return_cloudflare_challenge_error_when_403_fallback_is .await; let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone()); + let cookie_store = PgCookieStore::new(pool.clone()); cookie_store .capture_set_cookie("acct_chat", "cf_clearance=old; Domain=.chatgpt.com; Path=/") .await @@ -540,8 +540,8 @@ async fn responses_should_return_cloudflare_challenge_error_when_403_fallback_is let status = response.status(); let body = response_json(response).await; let message = body["error"]["message"].as_str().unwrap_or_default(); - let primary_state: (String, Option) = - sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = ?") + let primary_state: (String, Option>) = + sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -588,7 +588,7 @@ async fn responses_should_cool_down_cloudflare_403_and_fallback() { .await; let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone()); + let cookie_store = PgCookieStore::new(pool.clone()); cookie_store .capture_set_cookie( "acct_primary", @@ -625,15 +625,13 @@ async fn responses_should_cool_down_cloudflare_403_and_fallback() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let primary_state: (String, Option) = - sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = ?") + let primary_state: (String, Option>) = + sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); - let cooldown_until = chrono::DateTime::parse_from_rfc3339(primary_state.1.as_deref().unwrap()) - .unwrap() - .with_timezone(&Utc); + let cooldown_until = primary_state.1.unwrap(); assert_eq!(status, StatusCode::OK); assert_eq!(body["id"], "resp_after_cloudflare"); @@ -680,7 +678,7 @@ async fn responses_should_clear_cookies_after_cloudflare_path_block_404_and_fall .await; let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone()); + let cookie_store = PgCookieStore::new(pool.clone()); cookie_store .capture_set_cookie( "acct_primary", @@ -717,7 +715,7 @@ async fn responses_should_clear_cookies_after_cloudflare_path_block_404_and_fall .unwrap(); let status = response.status(); let body = response_json(response).await; - let primary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let primary_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await @@ -755,7 +753,7 @@ async fn responses_should_disable_account_after_three_cloudflare_path_block_404s .await; let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone()); + let cookie_store = PgCookieStore::new(pool.clone()); cookie_store .capture_set_cookie("acct_chat", "cf_clearance=old; Domain=.chatgpt.com; Path=/") .await @@ -789,7 +787,7 @@ async fn responses_should_disable_account_after_three_cloudflare_path_block_404s assert_eq!(status, StatusCode::BAD_GATEWAY); assert!(message.contains("cloudflare-path-block")); } - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -855,7 +853,7 @@ async fn responses_should_fallback_after_http_model_unsupported() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let primary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let primary_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await @@ -1026,7 +1024,7 @@ async fn responses_stream_should_return_model_unsupported_error_when_fallback_is .unwrap_or_default() .to_string(); let body = response_text(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -1393,22 +1391,22 @@ async fn responses_stream_should_fallback_to_next_account_after_rate_limit() { .unwrap_or_default() .to_string(); let body = response_text(response).await; - let primary_quota_state: (i64, Option) = sqlx::query_as( - "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let primary_quota_state: (bool, Option>) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let primary_usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let secondary_usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_secondary") .fetch_one(&pool) @@ -1418,7 +1416,7 @@ async fn responses_stream_should_fallback_to_next_account_after_rate_limit() { assert_eq!(status, StatusCode::OK); assert!(content_type.starts_with("text/event-stream")); assert!(body.contains("resp_stream_after_429")); - assert_eq!(primary_quota_state.0, 1); + assert!(primary_quota_state.0); assert!(primary_quota_state.1.is_some()); assert_eq!(primary_usage, (1, 0, 0)); assert_eq!(secondary_usage, (1, 4, 1)); @@ -1477,13 +1475,13 @@ async fn responses_stream_should_mark_expired_after_401_and_fallback() { .unwrap_or_default() .to_string(); let body = response_text(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let secondary_usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_secondary") .fetch_one(&pool) @@ -1539,7 +1537,7 @@ async fn responses_stream_should_return_auth_error_when_401_fallback_is_exhauste .unwrap_or_default() .to_string(); let body = response_text(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -1607,7 +1605,7 @@ async fn responses_stream_should_mark_banned_after_403_and_fallback() { .unwrap_or_default() .to_string(); let body = response_text(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await @@ -1645,7 +1643,7 @@ async fn responses_stream_should_cool_down_cloudflare_403_and_fallback() { .await; let (app, api_key, pool, _dir) = test_app_with_two_accounts(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone()); + let cookie_store = PgCookieStore::new(pool.clone()); cookie_store .capture_set_cookie( "acct_primary", @@ -1688,15 +1686,13 @@ async fn responses_stream_should_cool_down_cloudflare_403_and_fallback() { .unwrap_or_default() .to_string(); let body = response_text(response).await; - let primary_state: (String, Option) = - sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = ?") + let primary_state: (String, Option>) = + sqlx::query_as("select status, cloudflare_cooldown_until from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); - let cooldown_until = chrono::DateTime::parse_from_rfc3339(primary_state.1.as_deref().unwrap()) - .unwrap() - .with_timezone(&Utc); + let cooldown_until = primary_state.1.unwrap(); assert_eq!(status, StatusCode::OK); assert!(content_type.starts_with("text/event-stream")); @@ -1734,7 +1730,7 @@ async fn responses_stream_should_return_cloudflare_path_block_error_when_404_fal .await; let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; - let cookie_store = SqliteCookieStore::new(pool.clone()); + let cookie_store = PgCookieStore::new(pool.clone()); cookie_store .capture_set_cookie("acct_chat", "cf_clearance=old; Domain=.chatgpt.com; Path=/") .await @@ -1962,13 +1958,13 @@ async fn responses_should_mark_quota_exhausted_after_402_and_fallback() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let failed_usage: Option<(i64,)> = - sqlx::query_as("select request_count from account_usage where account_id = ?") + sqlx::query_as("select request_count from account_usage where account_id = $1") .bind("acct_primary") .fetch_optional(&pool) .await @@ -2033,13 +2029,13 @@ async fn responses_stream_should_mark_quota_exhausted_after_402_and_fallback() { .unwrap_or_default() .to_string(); let body = response_text(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let secondary_usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_secondary") .fetch_one(&pool) @@ -2089,7 +2085,7 @@ async fn responses_should_return_quota_error_when_402_fallback_is_exhausted() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -2141,8 +2137,8 @@ async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted .unwrap(); let status = response.status(); let body = response_json(response).await; - let quota_state: (i64, Option) = sqlx::query_as( - "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let quota_state: (bool, Option>) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct_chat") .fetch_one(&pool) @@ -2155,7 +2151,7 @@ async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted assert!(message.contains("rate limited")); assert_eq!(body["error"]["type"], "rate_limit_error"); assert_eq!(body["error"]["code"], "rate_limit_exceeded"); - assert_eq!(quota_state.0, 1); + assert!(quota_state.0); assert!(quota_state.1.is_some()); let event = latest_response_ops_error_log(&pool).await; let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); @@ -2168,8 +2164,8 @@ async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted assert_eq!(event.route.as_deref(), Some("/v1/responses")); assert_eq!(event.status_code, Some(429)); assert_eq!(metadata["stream"], false); - assert_eq!(metadata["transport"], "http_sse"); - assert_eq!(metadata["failureClass"], "rate_limited"); + assert_eq!(event.transport.as_deref(), Some("http_sse")); + assert_eq!(event.failure_class.as_deref(), Some("rate_limited")); assert_eq!(metadata["exhaustedCount"], 1); assert!(metadata["upstreamError"] .as_str() @@ -2224,8 +2220,8 @@ async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_ex .unwrap_or_default() .to_string(); let body = response_text(response).await; - let quota_state: (i64, Option) = sqlx::query_as( - "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let quota_state: (bool, Option>) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct_chat") .fetch_one(&pool) @@ -2238,7 +2234,7 @@ async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_ex assert!(body.contains("rate_limit_exceeded")); assert!(body.contains("All accounts exhausted (1 rate-limited)")); assert!(body.contains("rate limited")); - assert_eq!(quota_state.0, 1); + assert!(quota_state.0); assert!(quota_state.1.is_some()); assert!(!response_request_id.is_empty()); let event = latest_response_ops_error_log(&pool).await; @@ -2252,8 +2248,8 @@ async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_ex assert_eq!(event.route.as_deref(), Some("/v1/responses")); assert_eq!(event.status_code, Some(429)); assert_eq!(metadata["stream"], true); - assert_eq!(metadata["transport"], "http_sse"); - assert_eq!(metadata["failureClass"], "rate_limited"); + assert_eq!(event.transport.as_deref(), Some("http_sse")); + assert_eq!(event.failure_class.as_deref(), Some("rate_limited")); assert_eq!(metadata["exhaustedCount"], 1); assert!(metadata["upstreamError"] .as_str() @@ -2302,7 +2298,7 @@ async fn responses_stream_should_return_quota_error_when_402_fallback_is_exhaust .unwrap_or_default() .to_string(); let body = response_text(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -2366,7 +2362,7 @@ async fn responses_should_classify_sse_quota_failure_and_fallback() { .unwrap(); let status = response.status(); let body = response_json(response).await; - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await diff --git a/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs b/backend/tests/dispatch/chat_upstream/responses_websocket.rs similarity index 96% rename from backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs rename to backend/tests/dispatch/chat_upstream/responses_websocket.rs index fbce8df3a..a08cf43bb 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/responses_websocket.rs +++ b/backend/tests/dispatch/chat_upstream/responses_websocket.rs @@ -188,21 +188,21 @@ async fn responses_websocket_stream_first_error_429_should_retry_fallback_accoun .to_string(); let body = response_text(response).await; let (first_payload, second_payload) = upstream.await.unwrap(); - let primary_quota_state: (i64, Option) = sqlx::query_as( - "select quota_limit_reached, quota_cooldown_until from accounts where id = ?", + let primary_quota_state: (bool, Option>) = sqlx::query_as( + "select quota_limit_reached, quota_cooldown_until from accounts where id = $1", ) .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let primary_usage: (i64,) = - sqlx::query_as("select request_count from account_usage where account_id = ?") + sqlx::query_as("select request_count from account_usage where account_id = $1") .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let secondary_usage: (i64,) = - sqlx::query_as("select request_count from account_usage where account_id = ?") + sqlx::query_as("select request_count from account_usage where account_id = $1") .bind("acct_secondary") .fetch_one(&pool) .await @@ -213,7 +213,7 @@ async fn responses_websocket_stream_first_error_429_should_retry_fallback_accoun assert!(body.contains("resp_after_ws_stream_rate_limit")); assert_eq!(first_payload["type"], "response.create"); assert_eq!(second_payload["type"], "response.create"); - assert_eq!(primary_quota_state.0, 1); + assert!(primary_quota_state.0); assert!(primary_quota_state.1.is_some()); assert_eq!(primary_usage.0, 1); assert_eq!(secondary_usage.0, 1); @@ -622,7 +622,7 @@ async fn responses_websocket_stream_should_record_metadata_turn_state_for_contin let second_headers = second_headers.lock().unwrap().clone(); (first_payload, second_payload, second_headers) }); - let (app, api_key, pool, _dir) = test_app_with_account_pool_config(base_url, |config| { + let (app, api_key, _pool, _dir) = test_app_with_account_pool_config(base_url, |config| { config.ws_pool.enabled = false; }) .await; @@ -652,11 +652,6 @@ async fn responses_websocket_stream_should_record_metadata_turn_state_for_contin let first_body = response_text(first_response).await; assert!(first_body.contains("\"id\":\"resp_metadata_turn_state\"")); assert!(!first_body.contains("response.metadata")); - assert_eq!( - wait_for_session_affinity_turn_state(&pool, "resp_metadata_turn_state").await, - Some("turn-from-metadata".to_string()) - ); - let second_response = app .oneshot( Request::builder() @@ -834,7 +829,7 @@ async fn responses_websocket_should_strip_implicit_affinity_when_switching_accou (first_payload, second_payload, captured_second_headers) }); - let (app, state, api_key, pool, _dir) = test_app_with_two_accounts_and_state(base_url).await; + let (app, state, api_key, _pool, _dir) = test_app_with_two_accounts_and_state(base_url).await; let first_response = app .clone() @@ -883,12 +878,13 @@ async fn responses_websocket_should_strip_implicit_affinity_when_switching_accou .await .contains(&format!("\"id\":\"{second_response_id}\""))); let (first_payload, second_payload, second_headers) = upstream.await.unwrap(); - let affinity_count: (i64,) = - sqlx::query_as("select count(*) from session_affinities where response_id = ?") - .bind(&first_response_id) - .fetch_one(&pool) - .await - .unwrap(); + let affinity = state + .services + .background_tasks + .session_affinity + .get(&first_response_id, Utc::now(), chrono::Duration::hours(4)) + .await + .unwrap(); assert!(first_payload["prompt_cache_key"].as_str().is_some()); assert_eq!( @@ -898,7 +894,7 @@ async fn responses_websocket_should_strip_implicit_affinity_when_switching_accou assert!(captured_header(&second_headers, "x-codex-turn-state").is_none()); assert!(second_payload.get("previous_response_id").is_none()); assert_eq!(second_payload["input"].as_array().unwrap().len(), 3); - assert_eq!(affinity_count.0, 0); + assert!(affinity.is_none()); } #[tokio::test] @@ -974,7 +970,7 @@ async fn responses_websocket_should_not_implicitly_resume_unmatched_function_cal } #[tokio::test] -async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_restore() { +async fn responses_websocket_should_implicitly_resume_after_redis_affinity_restore() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let base_url = format!("http://{}", listener.local_addr().unwrap()); let server_base_url = base_url.clone(); @@ -991,7 +987,26 @@ async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_rest .await; (first_payload, second_payload) }); - let (app, api_key, pool, dir) = test_app_with_account_pool_config(base_url, |_| {}).await; + let (pool, _guard) = init_test_db("websocket-affinity-restart").await; + let redis = create_test_redis("websocket-affinity-restart").await; + let api_key = insert_client_api_key(&pool).await; + insert_account(&pool).await; + let config = test_config(test_database_url(), base_url); + let initial_state = test_app_state_with_storage( + config.clone(), + pool.clone(), + redis.clone(), + TEST_INSTALLATION_ID.to_string(), + UsageRecordOptions::from_config(&config), + ) + .await; + initial_state + .services + .account_pool + .restore_from_store() + .await + .unwrap(); + let app = router::router().with_state(initial_state); let first_response = app .oneshot(responses_json_request( @@ -1007,27 +1022,20 @@ async fn responses_websocket_should_implicitly_resume_after_sqlite_affinity_rest .unwrap(); assert_eq!(first_response.status(), StatusCode::OK); - let db = dir.path().join("openai-record-affinity.sqlite"); - let restored_state = test_app_state_with_pool_and_installation_id( - &test_config(format!("sqlite://{}", db.display()), server_base_url), + let restored_config = test_config(test_database_url(), server_base_url); + let restored_state = test_app_state_with_storage( + restored_config.clone(), pool.clone(), + redis, TEST_INSTALLATION_ID.to_string(), + UsageRecordOptions::from_config(&restored_config), ) .await; assert_eq!( restored_state .services .account_pool - .restore_from_repository() - .await - .unwrap(), - 1 - ); - assert_eq!( - restored_state - .services - .session_affinity - .restore_from_repository(Utc::now()) + .restore_from_store() .await .unwrap(), 1 @@ -1116,7 +1124,7 @@ async fn responses_websocket_pool_should_be_evicted_after_admin_account_status_c } }); let (app, api_key, pool, _dir) = test_app_with_account_and_pool(base_url).await; - seed_openai_admin_session(&pool, "session_status_cycle").await; + seed_openai_admin_key(&pool, "admin-status-cycle").await; let first_response = app .clone() @@ -1474,7 +1482,7 @@ async fn responses_websocket_should_route_previous_response_id_to_recorded_accou .await; (first_payload, second_payload) }); - let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await; + let (app, state, api_key, _pool, _dir) = test_app_with_two_accounts_and_state(base_url).await; let first_response = app .clone() @@ -1495,18 +1503,22 @@ async fn responses_websocket_should_route_previous_response_id_to_recorded_accou assert!(response_text(first_response) .await .contains("\"id\":\"resp_affinity_first\"")); - let stored_affinity: (String, String, String, Option, String) = sqlx::query_as( - "select account_id, conversation_id, function_call_ids_json, input_tokens, expires_at from session_affinities where response_id = ?", - ) - .bind("resp_affinity_first") - .fetch_one(&pool) - .await - .unwrap(); - assert_eq!(stored_affinity.0, "acct_primary"); - assert!(!stored_affinity.1.is_empty()); - assert_eq!(stored_affinity.2, "[]"); - assert_eq!(stored_affinity.3, Some(3)); - assert!(!stored_affinity.4.is_empty()); + let stored_affinity = state + .services + .background_tasks + .session_affinity + .get( + "resp_affinity_first", + Utc::now(), + chrono::Duration::hours(4), + ) + .await + .unwrap() + .unwrap(); + assert_eq!(stored_affinity.account_id, "acct_primary"); + assert!(!stored_affinity.conversation_id.is_empty()); + assert!(stored_affinity.function_call_ids.is_empty()); + assert_eq!(stored_affinity.input_tokens, Some(3)); let second_response = app .oneshot(responses_json_request( @@ -1821,7 +1833,7 @@ async fn responses_websocket_previous_response_id_should_retry_fallback_account_ let body = response_text(response).await; let fallback_payload = upstream.await.unwrap(); let secondary_usage: (i64,) = - sqlx::query_as("select request_count from account_usage where account_id = ?") + sqlx::query_as("select request_count from account_usage where account_id = $1") .bind("acct_secondary") .fetch_one(&pool) .await @@ -1873,7 +1885,7 @@ async fn responses_websocket_non_stream_previous_response_id_should_retry_fallba let body = response_json(response).await; let fallback_payload = upstream.await.unwrap(); let secondary_usage: (i64,) = - sqlx::query_as("select request_count from account_usage where account_id = ?") + sqlx::query_as("select request_count from account_usage where account_id = $1") .bind("acct_secondary") .fetch_one(&pool) .await @@ -1926,13 +1938,13 @@ async fn responses_websocket_without_history_should_mark_expired_after_fallback_ .unwrap(); let body = response_text(response).await; upstream.await.unwrap(); - let secondary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let secondary_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_secondary") .fetch_one(&pool) .await .unwrap(); let primary_usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_primary") .fetch_one(&pool) @@ -1990,13 +2002,13 @@ async fn responses_websocket_without_history_should_return_rate_limit_stream_err let body = response_text(response).await; upstream.await.unwrap(); let primary_usage: (i64,) = - sqlx::query_as("select request_count from account_usage where account_id = ?") + sqlx::query_as("select request_count from account_usage where account_id = $1") .bind("acct_primary") .fetch_one(&pool) .await .unwrap(); let secondary_usage: (i64,) = - sqlx::query_as("select request_count from account_usage where account_id = ?") + sqlx::query_as("select request_count from account_usage where account_id = $1") .bind("acct_secondary") .fetch_one(&pool) .await @@ -2065,7 +2077,7 @@ async fn responses_websocket_response_failed_quota_should_retry_fallback_account .unwrap(); let body = response_json(response).await; upstream.await.unwrap(); - let primary_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let primary_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_primary") .fetch_one(&pool) .await @@ -2107,7 +2119,7 @@ async fn responses_websocket_without_history_should_return_quota_stream_error_wh .unwrap(); let body = response_text(response).await; upstream.await.unwrap(); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -2157,7 +2169,7 @@ async fn responses_websocket_without_history_should_return_model_unsupported_str .unwrap(); let body = response_text(response).await; upstream.await.unwrap(); - let account_status: (String,) = sqlx::query_as("select status from accounts where id = ?") + let account_status: (String,) = sqlx::query_as("select status from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await @@ -2362,24 +2374,22 @@ async fn responses_stream_with_previous_response_id_should_record_websocket_audi assert_eq!(captured.payload["previous_response_id"], Value::Null); assert_eq!(event.level, "info"); assert_eq!(metadata["stream"], true); - assert_eq!(metadata["transport"], "websocket"); - assert_eq!(metadata["usage"]["inputTokens"], 3); - assert_eq!(metadata["usage"]["outputTokens"], 1); + assert_eq!(event.transport.as_deref(), Some("websocket")); + assert_eq!(event.input_tokens, Some(3)); + assert_eq!(event.output_tokens, Some(1)); assert!( - metadata["firstTokenMs"] - .as_i64() - .is_some_and(|value| value > 0), + event.first_token_ms.is_some_and(|value| value > 0), "websocket stream usage metadata should include first token latency: {metadata:?}", ); assert_rate_limit_header(&metadata, "x-codex-primary-used-percent", "44"); assert_rate_limit_header(&metadata, "x-codex-primary-window-minutes", "5"); - let stored_quota: (Option,) = - sqlx::query_as("select quota_json from accounts where id = ?") + let stored_quota: (Option,) = + sqlx::query_as("select quota_json from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await .unwrap(); - let quota: Value = serde_json::from_str(stored_quota.0.as_deref().unwrap()).unwrap(); + let quota = stored_quota.0.unwrap(); assert_eq!( quota["snapshots"][0]["primary"]["used_percent"].as_f64(), Some(44.0) diff --git a/backend/tests/proxy/dispatch/chat_upstream/usage_logging.rs b/backend/tests/dispatch/chat_upstream/usage_logging.rs similarity index 92% rename from backend/tests/proxy/dispatch/chat_upstream/usage_logging.rs rename to backend/tests/dispatch/chat_upstream/usage_logging.rs index 93ba315fc..864720e9a 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/usage_logging.rs +++ b/backend/tests/dispatch/chat_upstream/usage_logging.rs @@ -50,7 +50,7 @@ async fn responses_should_use_imported_account_record_usage_cookie_and_usage_rec assert_eq!(body["id"], "resp_usage"); assert_eq!(body["usage"]["input_tokens"], 7); let usage: (i64, i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens, cached_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens, cached_tokens from account_usage where account_id = $1", ) .bind("acct_chat") .fetch_one(&pool) @@ -58,7 +58,7 @@ async fn responses_should_use_imported_account_record_usage_cookie_and_usage_rec .unwrap(); assert_eq!(usage, (1, 7, 4, 2)); let model_usage: (i64, i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens, cached_tokens from account_model_usage where account_id = ? and model = ?", + "select request_count, input_tokens, output_tokens, cached_tokens from account_model_usage where account_id = $1 and model = $2", ) .bind("acct_chat") .bind("gpt-5.5") @@ -66,7 +66,7 @@ async fn responses_should_use_imported_account_record_usage_cookie_and_usage_rec .await .unwrap(); assert_eq!(model_usage, (1, 7, 4, 2)); - let cookie_header = SqliteCookieStore::new(pool.clone()) + let cookie_header = PgCookieStore::new(pool.clone()) .cookie_header("acct_chat", "chatgpt.com") .await .unwrap(); @@ -81,9 +81,9 @@ async fn responses_should_use_imported_account_record_usage_cookie_and_usage_rec assert_eq!(event.account_id.as_deref(), Some("acct_chat")); assert_eq!(event.status_code, Some(200)); assert_eq!(event.response_id.as_deref(), Some("resp_usage")); - assert_eq!(metadata["responseId"], "resp_usage"); + assert_eq!(event.response_id.as_deref(), Some("resp_usage")); assert_eq!(metadata["stream"], false); - assert_eq!(metadata["usage"]["inputTokens"], 7); + assert_eq!(event.input_tokens, Some(7)); } #[tokio::test] @@ -134,9 +134,8 @@ async fn responses_usage_record_should_store_resolved_upstream_model_after_alias let event = latest_response_usage_record(&pool).await; assert_eq!(event.request_id.as_deref(), Some("req_alias_usage_model")); assert_eq!(event.model.as_deref(), Some("gpt-5.5")); - let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); - assert_eq!(metadata["requestedModel"], "client-alias"); - assert_eq!(metadata["upstreamModel"], "gpt-5.5"); + assert_eq!(event.requested_model.as_deref(), Some("client-alias")); + assert_eq!(event.upstream_model.as_deref(), Some("gpt-5.5")); let received = server.received_requests().await.unwrap(); let upstream_body: Value = serde_json::from_slice(&received[0].body).unwrap(); @@ -225,7 +224,7 @@ async fn responses_should_record_image_generation_usage_when_tool_succeeds() { assert_eq!(response.status(), StatusCode::OK); let usage: (i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens, cached_tokens, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens, cached_tokens, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count from account_usage where account_id = $1", ) .bind("acct_chat") .fetch_one(&pool) @@ -273,7 +272,7 @@ async fn responses_should_record_failed_image_generation_attempt_when_tool_has_n assert_eq!(response.status(), StatusCode::OK); let usage: (i64, i64, i64, i64) = sqlx::query_as( - "select image_request_count, image_request_failed_count, window_image_request_count, window_image_request_failed_count from account_usage where account_id = ?", + "select image_request_count, image_request_failed_count, window_image_request_count, window_image_request_failed_count from account_usage where account_id = $1", ) .bind("acct_chat") .fetch_one(&pool) @@ -302,7 +301,7 @@ async fn responses_should_passively_cache_rate_limit_headers() { let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; sqlx::query( - r#"update accounts set quota_json = '{"credits":{"has_credits":true,"unlimited":false,"balance":12}}', quota_verify_required = 1 where id = ?"#, + r#"update accounts set quota_json = '{"credits":{"has_credits":true,"unlimited":false,"balance":12}}', quota_verify_required = true where id = $1"#, ) .bind("acct_chat") .execute(&pool) @@ -332,25 +331,39 @@ async fn responses_should_passively_cache_rate_limit_headers() { .unwrap(); assert_eq!(response.status(), StatusCode::OK); - let stored: (String, String, i64, Option, i64) = sqlx::query_as( - "select quota_json, status, quota_limit_reached, quota_cooldown_until, quota_verify_required from accounts where id = ?", + let stored: ( + Value, + String, + bool, + Option>, + bool, + ) = sqlx::query_as( + "select quota_json, status, quota_limit_reached, quota_cooldown_until, quota_verify_required from accounts where id = $1", ) .bind("acct_chat") .fetch_one(&pool) .await .unwrap(); - let quota: Value = serde_json::from_str(&stored.0).unwrap(); + let quota = stored.0; assert_eq!(quota["snapshots"][0]["source"], "core"); assert_eq!(quota["snapshots"][0]["primary"]["limit_reached"], true); assert_eq!(quota["snapshots"][0]["primary"]["reset_at"], reset_at); assert_eq!(quota["credits"]["balance"], 12); assert_eq!(stored.1, "quota_exhausted"); - assert_eq!(stored.2, 1); + assert!(stored.2); assert!(stored.3.is_some()); - assert_eq!(stored.4, 0); - let window: (i64, i64, i64, i64, Option, Option, Option) = - sqlx::query_as( - "select window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds from account_usage where account_id = ?", + assert!(!stored.4); + type UsageWindowRow = ( + i64, + i64, + i64, + i64, + Option>, + Option>, + Option, + ); + let window: UsageWindowRow = sqlx::query_as( + "select window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_started_at, window_reset_at, limit_window_seconds from account_usage where account_id = $1", ) .bind("acct_chat") .fetch_one(&pool) @@ -361,12 +374,7 @@ async fn responses_should_passively_cache_rate_limit_headers() { assert_eq!(window.2, 4); assert_eq!(window.3, 2); assert!(window.4.is_some()); - assert_eq!( - DateTime::parse_from_rfc3339(window.5.as_deref().unwrap()) - .unwrap() - .timestamp(), - reset_at - ); + assert_eq!(window.5.unwrap().timestamp(), reset_at); assert_eq!(window.6, Some(300)); let second_response = app @@ -430,7 +438,7 @@ async fn responses_should_record_empty_response_attempts() { assert_eq!(response.status(), StatusCode::BAD_GATEWAY); let usage: (i64, i64, i64, i64) = sqlx::query_as( - "select request_count, empty_response_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, empty_response_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_chat") .fetch_one(&pool) @@ -484,14 +492,14 @@ async fn responses_stream_should_proxy_sse_and_record_usage() { .to_string(); let body = response_text(response).await; let usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_usage where account_id = ?", + "select request_count, input_tokens, output_tokens from account_usage where account_id = $1", ) .bind("acct_chat") .fetch_one(&pool) .await .unwrap(); let model_usage: (i64, i64, i64) = sqlx::query_as( - "select request_count, input_tokens, output_tokens from account_model_usage where account_id = ? and model = ?", + "select request_count, input_tokens, output_tokens from account_model_usage where account_id = $1 and model = $2", ) .bind("acct_chat") .bind("gpt-5.5") @@ -561,19 +569,17 @@ async fn responses_stream_should_record_usage_record_after_completed_stream() { assert_eq!(event.account_id.as_deref(), Some("acct_chat")); assert_eq!(event.status_code, Some(200)); assert_eq!(metadata["stream"], true); - assert_eq!(metadata["transport"], "http_sse"); + assert_eq!(event.transport.as_deref(), Some("http_sse")); assert_eq!(metadata["completed"], true); - assert_eq!(metadata["responseId"], "resp_stream_usage"); - assert_eq!(metadata["usage"]["inputTokens"], 3); - assert_eq!(metadata["usage"]["outputTokens"], 5); + assert_eq!(event.response_id.as_deref(), Some("resp_stream_usage")); + assert_eq!(event.input_tokens, Some(3)); + assert_eq!(event.output_tokens, Some(5)); assert!( - metadata["firstTokenMs"] - .as_i64() - .is_some_and(|value| value > 0), + event.first_token_ms.is_some_and(|value| value > 0), "stream usage metadata should include first token latency: {metadata:?}", ); let first_token_bucket: (i64, i64) = sqlx::query_as( - "select first_token_latency_sum, first_token_latency_count from usage_time_buckets limit 1", + "select first_token_latency_sum, first_token_latency_count from request_time_buckets limit 1", ) .fetch_one(&pool) .await @@ -643,7 +649,7 @@ async fn responses_stream_should_preserve_body_metadata_when_capture_body_enable async fn responses_stream_should_record_usage_record_after_late_disconnect() { let (base_url, first_chunk_sent, close_upstream) = spawn_chunked_sse_upstream_then_clean_close_with_headers( - include_str!("../../../fixtures/responses/http_sse/partial_logged_disconnect.sse"), + include_str!("../../fixtures/responses/http_sse/partial_logged_disconnect.sse"), &[ ("retry-after", "11"), ("x-codex-primary-used-percent", "88"), @@ -696,13 +702,13 @@ async fn responses_stream_should_record_usage_record_after_late_disconnect() { assert_eq!(metadata["synthetic"], true); assert_rate_limit_header(&metadata, "retry-after", "11"); assert_rate_limit_header(&metadata, "x-codex-primary-used-percent", "88"); - let stored_quota: (Option,) = - sqlx::query_as("select quota_json from accounts where id = ?") + let stored_quota: (Option,) = + sqlx::query_as("select quota_json from accounts where id = $1") .bind("acct_chat") .fetch_one(&pool) .await .unwrap(); - let quota: Value = serde_json::from_str(stored_quota.0.as_deref().unwrap()).unwrap(); + let quota = stored_quota.0.unwrap(); assert_eq!( quota["snapshots"][0]["primary"]["used_percent"].as_f64(), Some(88.0) @@ -713,7 +719,7 @@ async fn responses_stream_should_record_usage_record_after_late_disconnect() { async fn responses_stream_should_record_failure_detail_after_transport_error() { let (base_url, first_chunk_sent, close_upstream) = spawn_chunked_sse_upstream_then_abrupt_close(include_str!( - "../../../fixtures/responses/http_sse/partial_transport_failure.sse" + "../../fixtures/responses/http_sse/partial_transport_failure.sse" )) .await; @@ -757,7 +763,7 @@ async fn responses_stream_should_record_failure_detail_after_transport_error() { #[tokio::test] async fn responses_stream_should_not_record_first_token_when_metadata_only_prefix_fails() { let (base_url, first_chunk_sent, close_upstream) = spawn_chunked_sse_upstream_then_clean_close( - include_str!("../../../fixtures/responses/http_sse/metadata_only_prefix.sse"), + include_str!("../../fixtures/responses/http_sse/metadata_only_prefix.sse"), ) .await; @@ -791,7 +797,7 @@ async fn responses_stream_should_not_record_first_token_when_metadata_only_prefi let event = latest_response_ops_error_log(&pool).await; let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap(); let first_token_bucket: (i64, i64) = sqlx::query_as( - "select coalesce(sum(first_token_latency_sum), 0), coalesce(sum(first_token_latency_count), 0) from usage_time_buckets", + "select coalesce(sum(first_token_latency_sum), 0)::bigint, coalesce(sum(first_token_latency_count), 0)::bigint from request_time_buckets", ) .fetch_one(&pool) .await @@ -846,7 +852,7 @@ async fn responses_should_record_request_count_when_5xx_retries_are_exhausted() let status = response.status(); let authorizations = received_authorizations(&server).await; let failed_usage: Option<(i64,)> = - sqlx::query_as("select request_count from account_usage where account_id = ?") + sqlx::query_as("select request_count from account_usage where account_id = $1") .bind("acct_chat") .fetch_optional(&pool) .await @@ -914,11 +920,9 @@ async fn responses_should_record_attempt_trace_after_retrying_another_account() assert_eq!(event.request_id.as_deref(), Some("req_attempt_trace_retry")); assert_eq!(event.account_id.as_deref(), Some("acct_secondary")); assert_eq!(event.status_code, Some(500)); - assert_eq!(metadata["eventStatusCode"], 500); - assert_eq!(metadata["clientStatusCode"], 502); - assert_eq!(metadata["upstreamStatusCode"], 500); - assert_eq!(metadata["attemptIndex"], 1); - assert_eq!(metadata["attemptAccountId"], "acct_secondary"); + assert_eq!(event.client_status_code, Some(502)); + assert_eq!(event.upstream_status_code, Some(500)); + assert_eq!(event.attempt_index, Some(1)); assert_eq!(metadata["attemptCount"], 2); assert_eq!(metadata["attempts"][0]["index"], 0); assert_eq!(metadata["attempts"][0]["accountId"], "acct_primary"); @@ -938,7 +942,7 @@ async fn responses_should_record_attempt_trace_after_retrying_another_account() } #[tokio::test] -async fn responses_should_record_session_affinity_for_completed_response() { +async fn responses_should_complete_when_session_affinity_recording_is_enabled() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/codex/responses")) @@ -952,7 +956,7 @@ async fn responses_should_record_session_affinity_for_completed_response() { .mount(&server) .await; - let (app, api_key, pool, _dir) = test_app_with_account_and_pool(server.uri()).await; + let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(server.uri()).await; let response = app .oneshot( Request::builder() @@ -973,13 +977,5 @@ async fn responses_should_record_session_affinity_for_completed_response() { ) .await .unwrap(); - let records = SqliteSessionAffinityStore::new(pool) - .list_active(Utc::now()) - .await - .unwrap(); - assert_eq!(response.status(), StatusCode::OK); - assert_eq!(records.len(), 1); - assert_eq!(records[0].response_id, "resp_response_1"); - assert_eq!(records[0].entry.account_id, "acct_chat"); } diff --git a/backend/tests/dispatch/mod.rs b/backend/tests/dispatch/mod.rs new file mode 100644 index 000000000..74cc09489 --- /dev/null +++ b/backend/tests/dispatch/mod.rs @@ -0,0 +1,2 @@ +pub(crate) mod affinity; +pub(crate) mod chat_upstream; diff --git a/backend/tests/fixtures/sqlite_v3.sql b/backend/tests/fixtures/sqlite_v3.sql new file mode 100644 index 000000000..5b104610a --- /dev/null +++ b/backend/tests/fixtures/sqlite_v3.sql @@ -0,0 +1,382 @@ +create table schema_migrations ( + version integer primary key, + name text not null, + applied_at text not null +); + +pragma foreign_keys = on; + +-- ============================================ +-- Admin Users +-- ============================================ +create table admin_users ( + id text primary key, + password_hash text not null, + created_at text not null, + updated_at text not null +); + +-- ============================================ +-- Admin Sessions +-- ============================================ +create table admin_sessions ( + id text primary key, + user_id text not null references admin_users(id) on delete cascade, + expires_at text not null, + created_at text not null +); + +create index idx_admin_sessions_expires on admin_sessions(expires_at); + +-- ============================================ +-- Client API Keys +-- ============================================ +create table client_api_keys ( + id text primary key, + name text not null, + prefix text not null, + key text not null unique, + label text, + enabled integer not null default 1 check (enabled in (0, 1)), + created_at text not null, + last_used_at text +); + +create index idx_client_api_keys_key_enabled on client_api_keys(key) where enabled = 1; +create index idx_client_api_keys_created_id on client_api_keys(created_at desc, id desc); + +-- ============================================ +-- Runtime Settings +-- ============================================ +create table runtime_settings ( + id integer primary key check (id = 1), + model_aliases_json text not null default '{}', + refresh_margin_seconds integer not null check (refresh_margin_seconds > 0), + refresh_concurrency integer not null check (refresh_concurrency > 0), + max_concurrent_per_account integer not null check (max_concurrent_per_account > 0), + request_interval_ms integer not null check (request_interval_ms >= 0), + rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')), + admin_api_key text, + updated_at text not null +); + +-- ============================================ +-- Accounts +-- ============================================ +create table accounts ( + id text primary key, + email text, + chatgpt_account_id text, + chatgpt_user_id text, + label text, + plan_type text, + access_token text not null, + refresh_token text, + access_token_expires_at text, + next_refresh_at text, + status text not null check (status in ('active', 'expired', 'quota_exhausted', 'refreshing', 'disabled', 'banned')), + quota_json text, + quota_fetched_at text, + quota_limit_reached integer not null default 0 check (quota_limit_reached in (0, 1)), + quota_verify_required integer not null default 0 check (quota_verify_required in (0, 1)), + quota_cooldown_until text, + cloudflare_cooldown_until text, + added_at text not null, + updated_at text not null +); + +create index idx_accounts_status on accounts(status); +create index idx_accounts_added_id on accounts(added_at desc, id desc); +create unique index ux_accounts_chatgpt_identity +on accounts(chatgpt_account_id, coalesce(chatgpt_user_id, '')) +where chatgpt_account_id is not null; + +-- ============================================ +-- Account Refresh Leases +-- ============================================ +create table account_refresh_leases ( + account_id text primary key references accounts(id) on delete cascade, + owner text not null, + expires_at text not null, + updated_at text not null +); + +create index idx_account_refresh_leases_expires on account_refresh_leases(expires_at); + +-- ============================================ +-- Account Usage Statistics +-- ============================================ +create table account_usage ( + account_id text primary key references accounts(id) on delete cascade, + request_count integer not null default 0 check (request_count >= 0), + empty_response_count integer not null default 0 check (empty_response_count >= 0), + input_tokens integer not null default 0 check (input_tokens >= 0), + output_tokens integer not null default 0 check (output_tokens >= 0), + cached_tokens integer not null default 0 check (cached_tokens >= 0), + reasoning_tokens integer not null default 0 check (reasoning_tokens >= 0), + total_tokens integer not null default 0 check (total_tokens >= 0), + image_input_tokens integer not null default 0 check (image_input_tokens >= 0), + image_output_tokens integer not null default 0 check (image_output_tokens >= 0), + image_request_count integer not null default 0 check (image_request_count >= 0), + image_request_failed_count integer not null default 0 check (image_request_failed_count >= 0), + window_request_count integer not null default 0 check (window_request_count >= 0), + window_input_tokens integer not null default 0 check (window_input_tokens >= 0), + window_output_tokens integer not null default 0 check (window_output_tokens >= 0), + window_cached_tokens integer not null default 0 check (window_cached_tokens >= 0), + window_image_input_tokens integer not null default 0 check (window_image_input_tokens >= 0), + window_image_output_tokens integer not null default 0 check (window_image_output_tokens >= 0), + window_image_request_count integer not null default 0 check (window_image_request_count >= 0), + window_image_request_failed_count integer not null default 0 check (window_image_request_failed_count >= 0), + window_started_at text, + window_reset_at text, + limit_window_seconds integer check (limit_window_seconds is null or limit_window_seconds > 0), + last_used_at text +); + +create index idx_account_usage_last_used_account +on account_usage(last_used_at desc, account_id desc); + +-- ============================================ +-- Account Model Usage Statistics +-- ============================================ +create table account_model_usage ( + account_id text not null references accounts(id) on delete cascade, + model text not null, + request_count integer not null default 0 check (request_count >= 0), + error_count integer not null default 0 check (error_count >= 0), + input_tokens integer not null default 0 check (input_tokens >= 0), + output_tokens integer not null default 0 check (output_tokens >= 0), + cached_tokens integer not null default 0 check (cached_tokens >= 0), + last_used_at text, + primary key (account_id, model) +); + +create index idx_account_model_usage_last_used +on account_model_usage(last_used_at desc, account_id, model); + +-- ============================================ +-- Usage Time Buckets +-- ============================================ +create table usage_time_buckets ( + bucket_start text not null, + account_id text not null default '', + model text not null default '', + service_tier text not null default '', + request_count integer not null default 0 check (request_count >= 0), + error_count integer not null default 0 check (error_count >= 0), + input_tokens integer not null default 0 check (input_tokens >= 0), + output_tokens integer not null default 0 check (output_tokens >= 0), + cached_tokens integer not null default 0 check (cached_tokens >= 0), + first_token_latency_sum integer not null default 0 check (first_token_latency_sum >= 0), + first_token_latency_count integer not null default 0 check (first_token_latency_count >= 0), + latency_sum integer not null default 0 check (latency_sum >= 0), + latency_count integer not null default 0 check (latency_count >= 0), + max_latency_ms integer not null default 0 check (max_latency_ms >= 0), + min_latency_ms integer not null default 0 check (min_latency_ms >= 0), + updated_at text not null, + primary key (bucket_start, account_id, model, service_tier) +); + +create index idx_usage_time_buckets_bucket +on usage_time_buckets(bucket_start); +create index idx_usage_time_buckets_model_bucket +on usage_time_buckets(model, bucket_start); + +-- ============================================ +-- Account Cookies +-- ============================================ +create table account_cookies ( + id text primary key, + account_id text not null references accounts(id) on delete cascade, + domain text not null, + name text not null, + value text not null, + path text not null default '/', + expires_at text, + updated_at text not null, + unique(account_id, domain, name, path) +); + +create index idx_account_cookies_account on account_cookies(account_id); +create index idx_account_cookies_account_domain on account_cookies(account_id, domain); +create index idx_account_cookies_expires +on account_cookies(expires_at) +where expires_at is not null; + +-- ============================================ +-- Device Fingerprints +-- ============================================ +create table fingerprints ( + id text primary key, + originator text not null, + app_version text not null, + build_number text not null, + platform text not null, + arch text not null, + chromium_version text not null, + user_agent_template text not null, + default_headers_json text not null, + header_order_json text not null, + source text not null, + created_at text not null, + updated_at text not null +); + +create table fingerprint_update_history ( + id text primary key, + current_fingerprint_id text not null references fingerprints(id) on delete cascade, + app_version text not null, + build_number text not null, + chromium_version text, + source text not null, + manifest_json text, + created_at text not null +); + +create index idx_fingerprint_update_history_created_id +on fingerprint_update_history(created_at desc, id desc); + +-- ============================================ +-- Usage Records +-- ============================================ +create table usage_records ( + id text primary key, + request_id text, + kind text not null, + level text not null check (level in ('debug', 'info', 'warn', 'error')), + account_id text, + route text, + model text, + status_code integer check (status_code is null or (status_code >= 100 and status_code <= 599)), + transport text, + attempt_index integer check (attempt_index is null or attempt_index >= 0), + upstream_status_code integer check (upstream_status_code is null or (upstream_status_code >= 100 and upstream_status_code <= 599)), + failure_class text, + response_id text, + upstream_request_id text, + latency_ms integer check (latency_ms is null or latency_ms >= 0), + message text not null, + metadata_json text not null, + created_at text not null +); + +create index idx_usage_records_created_id on usage_records(created_at desc, id desc); +create index idx_usage_records_kind_created on usage_records(kind, created_at desc); +create index idx_usage_records_request_id on usage_records(request_id); +create index idx_usage_records_account on usage_records(account_id, created_at desc) where account_id is not null; +create index idx_usage_records_transport on usage_records(transport, created_at desc) where transport is not null; +create index idx_usage_records_failure_class on usage_records(failure_class, created_at desc) where failure_class is not null; +create index idx_usage_records_response_id on usage_records(response_id) where response_id is not null; +create index idx_usage_records_upstream_request_id on usage_records(upstream_request_id) where upstream_request_id is not null; +create index idx_usage_records_level_created on usage_records(level, created_at desc); +create index idx_usage_records_route_created on usage_records(route, created_at desc) where route is not null; +create index idx_usage_records_model_created on usage_records(model, created_at desc) where model is not null; +create index idx_usage_records_status_created on usage_records(status_code, created_at desc) where status_code is not null; +create index idx_usage_records_upstream_status_created on usage_records(upstream_status_code, created_at desc) where upstream_status_code is not null; + +-- ============================================ +-- Model Plan Snapshots +-- ============================================ +create table model_plan_snapshots ( + plan_type text primary key, + models_json text not null, + fetched_at text not null +); + +-- ============================================ +-- Session Affinities +-- ============================================ +create table session_affinities ( + response_id text primary key, + account_id text not null references accounts(id) on delete cascade, + conversation_id text not null, + turn_state text, + instructions_hash text, + input_tokens integer check (input_tokens is null or input_tokens >= 0), + function_call_ids_json text not null default '[]', + variant_hash text, + expires_at text not null, + created_at text not null +); + +create index idx_session_affinities_conversation on session_affinities(conversation_id, created_at desc); +create index idx_session_affinities_expires on session_affinities(expires_at); +create index idx_session_affinities_active_order on session_affinities(expires_at, created_at, response_id); + +create table runtime_settings_new ( + id integer primary key check (id = 1), + model_aliases_json text not null default '{}', + refresh_margin_seconds integer not null check (refresh_margin_seconds > 0), + refresh_concurrency integer not null check (refresh_concurrency > 0), + max_concurrent_per_account integer not null check (max_concurrent_per_account > 0), + request_interval_ms integer not null check (request_interval_ms >= 0), + rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')), + admin_api_key text, + updated_at text not null +); + +insert into runtime_settings_new ( + id, + model_aliases_json, + refresh_margin_seconds, + refresh_concurrency, + max_concurrent_per_account, + request_interval_ms, + rotation_strategy, + admin_api_key, + updated_at +) +select + id, + model_aliases_json, + refresh_margin_seconds, + refresh_concurrency, + max_concurrent_per_account, + request_interval_ms, + case rotation_strategy + when 'least_used' then 'quota_reset_priority' + else rotation_strategy + end, + admin_api_key, + updated_at +from runtime_settings; + +drop table runtime_settings; +alter table runtime_settings_new rename to runtime_settings; + +create table ops_error_logs ( + id text primary key, + request_id text, + kind text not null, + account_id text, + route text, + model text, + status_code integer check (status_code is null or (status_code >= 100 and status_code <= 599)), + client_status_code integer check (client_status_code is null or (client_status_code >= 100 and client_status_code <= 599)), + upstream_status_code integer check (upstream_status_code is null or (upstream_status_code >= 100 and upstream_status_code <= 599)), + transport text, + attempt_index integer check (attempt_index is null or attempt_index >= 0), + failure_class text, + response_id text, + upstream_request_id text, + latency_ms integer check (latency_ms is null or latency_ms >= 0), + message text not null, + metadata_json text not null, + created_at text not null +); + +create index idx_ops_error_logs_created_id on ops_error_logs(created_at desc, id desc); +create index idx_ops_error_logs_request_id on ops_error_logs(request_id) where request_id is not null; +create index idx_ops_error_logs_account on ops_error_logs(account_id, created_at desc) where account_id is not null; +create index idx_ops_error_logs_route_created on ops_error_logs(route, created_at desc) where route is not null; +create index idx_ops_error_logs_model_created on ops_error_logs(model, created_at desc) where model is not null; +create index idx_ops_error_logs_status_created on ops_error_logs(status_code, created_at desc) where status_code is not null; +create index idx_ops_error_logs_transport_created on ops_error_logs(transport, created_at desc) where transport is not null; +create index idx_ops_error_logs_failure_class on ops_error_logs(failure_class, created_at desc) where failure_class is not null; +create index idx_ops_error_logs_response_id on ops_error_logs(response_id) where response_id is not null; +create index idx_ops_error_logs_upstream_request_id on ops_error_logs(upstream_request_id) where upstream_request_id is not null; + +insert into schema_migrations (version, name, applied_at) values + (1, 'initial', '2026-01-01T00:00:00Z'), + (2, 'rename_rotation_strategies', '2026-01-01T00:00:00Z'), + (3, 'ops_error_logs', '2026-01-01T00:00:00Z'); diff --git a/backend/tests/infra/storage_schema/mod.rs b/backend/tests/infra/storage_schema/mod.rs index 0fa574dbc..411c1763b 100644 --- a/backend/tests/infra/storage_schema/mod.rs +++ b/backend/tests/infra/storage_schema/mod.rs @@ -1,708 +1,193 @@ -use std::path::Path; +use std::collections::BTreeSet; -use codex_proxy_rs::infra::database::connect_sqlite; +use crate::support::storage::init_test_db; #[tokio::test] -async fn sqlite_schema_should_create_current_tables() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("test.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from sqlite_master where type = 'table' and name in ('schema_migrations', 'admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'ops_error_logs', 'model_plan_snapshots', 'session_affinities') order by name", +async fn postgres_schema_is_terminal_v1_without_runtime_tables() { + let (pool, _guard) = init_test_db("schema-terminal-v1").await; + let tables = sqlx::query_scalar::<_, String>( + "select table_name + from information_schema.tables + where table_schema = current_schema() and table_type = 'BASE TABLE' + order by table_name", ) .fetch_all(&pool) .await .unwrap(); assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "account_cookies", - "account_model_usage", - "account_refresh_leases", - "account_usage", - "accounts", - "admin_sessions", - "admin_users", - "client_api_keys", - "fingerprint_update_history", - "fingerprints", - "model_plan_snapshots", - "ops_error_logs", - "runtime_settings", - "schema_migrations", - "session_affinities", - "usage_records", - "usage_time_buckets", - ] + tables.into_iter().collect::>(), + BTreeSet::from_iter( + [ + "account_cookies", + "account_model_usage", + "account_usage", + "accounts", + "admin_users", + "client_api_keys", + "fingerprint_update_history", + "fingerprints", + "ops_error_logs", + "request_time_buckets", + "runtime_settings", + "schema_migrations", + "usage_records", + ] + .map(str::to_string) + ) ); + + let versions = sqlx::query_as::<_, (i64, String)>( + "select version, name from schema_migrations order by version", + ) + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(versions, vec![(1, "initial".to_string())]); } #[tokio::test] -async fn sqlite_schema_should_record_applied_migrations() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("schema-migrations.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); +async fn postgres_schema_keeps_retrievable_client_keys_and_hashes_admin_credentials() { + let (pool, _guard) = init_test_db("schema-credential-storage").await; + let key_columns = columns(&pool, "client_api_keys").await; + assert!(key_columns.contains("key")); + assert!(!key_columns.contains("key_hash")); + let settings_columns = columns(&pool, "runtime_settings").await; + assert!(settings_columns.contains("admin_api_key_hash")); + assert!(!settings_columns.contains("admin_api_key")); - let rows: Vec<(i64, String)> = - sqlx::query_as("select version, name from schema_migrations order by version") - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows, - [ - (1, "initial".to_string()), - (2, "rename_rotation_strategies".to_string()), - (3, "ops_error_logs".to_string()), - ] - ); + let key_unique: bool = sqlx::query_scalar( + "select exists ( + select 1 from pg_indexes + where schemaname = current_schema() + and tablename = 'client_api_keys' + and indexdef ilike '%unique%key%' + )", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert!(key_unique); } #[tokio::test] -async fn sqlite_schema_should_reject_refreshing_account_status() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("account-status.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); +async fn postgres_schema_separates_success_and_error_facts() { + let (pool, _guard) = init_test_db("schema-fact-split").await; + let usage_columns = columns(&pool, "usage_records").await; + for column in [ + "client_api_key_id", + "provider", + "requested_model", + "upstream_model", + "service_tier", + "first_token_ms", + "input_tokens", + "output_tokens", + "cached_tokens", + "reasoning_tokens", + ] { + assert!(usage_columns.contains(column), "missing {column}"); + } + assert!(!usage_columns.contains("level")); + assert!(!usage_columns.contains("failure_class")); - let error = sqlx::query( + let error_columns = columns(&pool, "ops_error_logs").await; + for column in [ + "client_api_key_id", + "provider", + "client_status_code", + "upstream_status_code", + "failure_class", + ] { + assert!(error_columns.contains(column), "missing {column}"); + } + + assert!(sqlx::query( + "insert into usage_records ( + id, kind, provider, account_id, model, status_code, message, metadata_json, created_at + ) values ('bad', 'request', 'openai', 'acct', 'gpt-5', 500, 'bad', '{}'::jsonb, now())", + ) + .execute(&pool) + .await + .is_err()); + assert!(sqlx::query( + "insert into ops_error_logs (id, kind, status_code, message, metadata_json, created_at) + values ('bad', 'request', 700, 'bad', '{}'::jsonb, now())", + ) + .execute(&pool) + .await + .is_err()); +} + +#[tokio::test] +async fn postgres_schema_has_required_fact_and_bucket_indexes() { + let (pool, _guard) = init_test_db("schema-indexes").await; + let indexes = sqlx::query_scalar::<_, String>( + "select indexname from pg_indexes where schemaname = current_schema()", + ) + .fetch_all(&pool) + .await + .unwrap() + .into_iter() + .collect::>(); + for index in [ + "idx_usage_records_created_id", + "idx_usage_records_key_created", + "idx_usage_records_response_id", + "idx_ops_error_logs_created_id", + "idx_ops_error_logs_key_created", + "idx_ops_error_logs_failure_class", + "idx_request_time_buckets_model", + ] { + assert!(indexes.contains(index), "missing {index}"); + } +} + +#[tokio::test] +async fn postgres_schema_enforces_account_and_counter_constraints() { + let (pool, _guard) = init_test_db("schema-constraints").await; + assert!(sqlx::query( "insert into accounts (id, access_token, status, added_at, updated_at) - values ('acct_refreshing', 'access', 'refreshing', '2026-07-09T00:00:00Z', '2026-07-09T00:00:00Z')", + values ('acct_bad', 'token', 'refreshing', now(), now())", ) .execute(&pool) .await - .unwrap_err(); - - assert!(error.to_string().contains("CHECK constraint failed")); -} - -#[test] -fn initial_migration_should_not_use_idempotent_business_schema_ddl() { - let migration_path = Path::new(env!("CARGO_MANIFEST_DIR")) - .join("src") - .join("infra") - .join("migrations") - .join("0001_initial.sql"); - let migration = std::fs::read_to_string(migration_path).unwrap(); - - assert!(!migration.to_lowercase().contains("if not exists")); -} - -#[tokio::test] -async fn connect_sqlite_should_reject_unversioned_existing_schema() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("unversioned.sqlite"); - let url = format!("sqlite://{}", db.display()); - - std::fs::File::create(&db).unwrap(); - let pool = sqlx::SqlitePool::connect(&url).await.unwrap(); - sqlx::query("create table admin_users (id text primary key)") - .execute(&pool) - .await - .unwrap(); - pool.close().await; - - let err = match connect_sqlite(&url).await { - Ok(pool) => { - pool.close().await; - panic!("unversioned sqlite schema should be rejected"); - } - Err(err) => err, - }; - - assert!(err.to_string().contains("unversioned sqlite schema")); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_runtime_settings_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("settings.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('runtime_settings') where name in ('id', 'model_aliases_json', 'refresh_margin_seconds', 'refresh_concurrency', 'max_concurrent_per_account', 'request_interval_ms', 'rotation_strategy', 'admin_api_key', 'updated_at') order by name", + .is_err()); + assert!(sqlx::query( + "insert into request_time_buckets ( + bucket_start, provider, account_id, model, service_tier, success_count + ) values (now(), 'openai', 'acct', 'gpt-5', '__unknown__', -1)", ) - .fetch_all(&pool) + .execute(&pool) .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "admin_api_key", - "id", - "max_concurrent_per_account", - "model_aliases_json", - "refresh_concurrency", - "refresh_margin_seconds", - "request_interval_ms", - "rotation_strategy", - "updated_at", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_runtime_query_indexes() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("indexes.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from sqlite_master where type = 'index' and name in ('idx_accounts_added_id', 'idx_account_usage_last_used_account', 'idx_account_model_usage_last_used', 'idx_client_api_keys_created_id', 'idx_fingerprint_update_history_created_id') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "idx_account_model_usage_last_used", - "idx_account_usage_last_used_account", - "idx_accounts_added_id", - "idx_client_api_keys_created_id", - "idx_fingerprint_update_history_created_id", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_diagnostic_filter_indexes() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("diagnostic-indexes.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from sqlite_master where type = 'index' and name in ('idx_account_cookies_expires', 'idx_usage_records_level_created', 'idx_usage_records_model_created', 'idx_usage_records_route_created', 'idx_usage_records_status_created', 'idx_usage_records_upstream_status_created', 'idx_ops_error_logs_created_id', 'idx_ops_error_logs_status_created', 'idx_ops_error_logs_failure_class', 'idx_session_affinities_active_order', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "idx_account_cookies_expires", - "idx_ops_error_logs_created_id", - "idx_ops_error_logs_failure_class", - "idx_ops_error_logs_status_created", - "idx_session_affinities_active_order", - "idx_usage_records_level_created", - "idx_usage_records_model_created", - "idx_usage_records_route_created", - "idx_usage_records_status_created", - "idx_usage_records_upstream_status_created", - "idx_usage_time_buckets_bucket", - "idx_usage_time_buckets_model_bucket", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_account_model_usage_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("model-usage.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('account_model_usage') where name in ('account_id', 'model', 'request_count', 'error_count', 'input_tokens', 'output_tokens', 'cached_tokens', 'last_used_at') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "account_id", - "cached_tokens", - "error_count", - "input_tokens", - "last_used_at", - "model", - "output_tokens", - "request_count", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_usage_time_bucket_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("usage-time-buckets.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('usage_time_buckets') where name in ('bucket_start', 'account_id', 'model', 'service_tier', 'request_count', 'error_count', 'input_tokens', 'output_tokens', 'cached_tokens', 'first_token_latency_sum', 'first_token_latency_count', 'latency_sum', 'latency_count', 'max_latency_ms', 'min_latency_ms', 'updated_at') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "account_id", - "bucket_start", - "cached_tokens", - "error_count", - "first_token_latency_count", - "first_token_latency_sum", - "input_tokens", - "latency_count", - "latency_sum", - "max_latency_ms", - "min_latency_ms", - "model", - "output_tokens", - "request_count", - "service_tier", - "updated_at", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_plain_account_cookie_value_column() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("cookies.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let value_column: Option<(String,)> = sqlx::query_as( - "select name from pragma_table_info('account_cookies') where name = 'value'", - ) - .fetch_optional(&pool) - .await - .unwrap(); - - assert_eq!(value_column, Some(("value".to_string(),))); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_account_usage_window_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("test.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('account_usage') where name in ('window_request_count', 'window_input_tokens', 'window_output_tokens', 'window_cached_tokens', 'window_started_at', 'window_reset_at', 'limit_window_seconds') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "limit_window_seconds", - "window_cached_tokens", - "window_input_tokens", - "window_output_tokens", - "window_request_count", - "window_reset_at", - "window_started_at", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_image_generation_usage_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("test.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('account_usage') where name in ('image_input_tokens', 'image_output_tokens', 'image_request_count', 'image_request_failed_count', 'window_image_input_tokens', 'window_image_output_tokens', 'window_image_request_count', 'window_image_request_failed_count') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "image_input_tokens", - "image_output_tokens", - "image_request_count", - "image_request_failed_count", - "window_image_input_tokens", - "window_image_output_tokens", - "window_image_request_count", - "window_image_request_failed_count", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_token_total_usage_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("test.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('account_usage') where name in ('reasoning_tokens', 'total_tokens') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - ["reasoning_tokens", "total_tokens"] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_quota_verify_required_flag() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("test.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let row: Option<(String,)> = sqlx::query_as( - "select name from pragma_table_info('accounts') where name = 'quota_verify_required'", - ) - .fetch_optional(&pool) - .await - .unwrap(); - - assert_eq!(row, Some(("quota_verify_required".to_string(),))); -} - -#[tokio::test] -async fn sqlite_schema_should_enforce_unique_chatgpt_identity() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("account-identity.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - + .is_err()); sqlx::query( - "insert into accounts (id, chatgpt_account_id, chatgpt_user_id, access_token, status, added_at, updated_at) values ('acct_a', 'chatgpt-account', 'chatgpt-user', 'access-token-a', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", + "insert into runtime_settings ( + id, refresh_margin_seconds, refresh_concurrency, + max_concurrent_per_account, request_interval_ms, + rotation_strategy, updated_at + ) values (1, 3600, 2, 1, 0, 'smart', now())", ) .execute(&pool) .await .unwrap(); - let result = sqlx::query( - "insert into accounts (id, chatgpt_account_id, chatgpt_user_id, access_token, status, added_at, updated_at) values ('acct_b', 'chatgpt-account', 'chatgpt-user', 'access-token-b', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - - assert!(result.is_err()); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_session_affinity_function_call_ids() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("test.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let row: Option<(String,)> = sqlx::query_as( - "select name from pragma_table_info('session_affinities') where name = 'function_call_ids_json'", - ) - .fetch_optional(&pool) - .await - .unwrap(); - - assert_eq!(row, Some(("function_call_ids_json".to_string(),))); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_complete_current_fingerprint_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("fingerprints.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('fingerprints') where name in ('originator', 'app_version', 'build_number', 'platform', 'arch', 'chromium_version', 'user_agent_template', 'default_headers_json', 'header_order_json', 'source', 'created_at', 'updated_at') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "app_version", - "arch", - "build_number", - "chromium_version", - "created_at", - "default_headers_json", - "header_order_json", - "originator", - "platform", - "source", - "updated_at", - "user_agent_template", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_fingerprint_history_table() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("fingerprint-history.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let row: Option<(String,)> = sqlx::query_as( - "select name from sqlite_master where type = 'table' and name = 'fingerprint_update_history'", - ) - .fetch_optional(&pool) - .await - .unwrap(); - - assert_eq!(row, Some(("fingerprint_update_history".to_string(),))); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_structured_event_diagnostic_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("event-diagnostics.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('usage_records') where name in ('transport', 'attempt_index', 'upstream_status_code', 'failure_class', 'response_id', 'upstream_request_id') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "attempt_index", - "failure_class", - "response_id", - "transport", - "upstream_request_id", - "upstream_status_code", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_persist_ops_error_log_columns() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("ops-error-columns.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let rows: Vec<(String,)> = sqlx::query_as( - "select name from pragma_table_info('ops_error_logs') where name in ('request_id', 'account_id', 'route', 'model', 'status_code', 'client_status_code', 'upstream_status_code', 'transport', 'attempt_index', 'failure_class', 'response_id', 'upstream_request_id', 'latency_ms', 'metadata_json') order by name", - ) - .fetch_all(&pool) - .await - .unwrap(); - - assert_eq!( - rows.into_iter().map(|row| row.0).collect::>(), - [ - "account_id", - "attempt_index", - "client_status_code", - "failure_class", - "latency_ms", - "metadata_json", - "model", - "request_id", - "response_id", - "route", - "status_code", - "transport", - "upstream_request_id", - "upstream_status_code", - ] - ); -} - -#[tokio::test] -async fn sqlite_schema_should_reject_invalid_account_status() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("status-check.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let result = sqlx::query( - "insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_bad', 'access-token', 'unknown', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - - assert!(result.is_err()); -} - -#[tokio::test] -async fn sqlite_schema_should_reject_invalid_usage_record_level() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("event-level-check.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let result = sqlx::query( - "insert into usage_records (id, kind, level, message, metadata_json, created_at) values ('usage_bad', 'request', 'fatal', 'bad', '{}', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - - assert!(result.is_err()); -} - -#[tokio::test] -async fn sqlite_schema_should_reject_invalid_ops_error_status_code() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("ops-error-status-check.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let result = sqlx::query( - "insert into ops_error_logs (id, kind, status_code, message, metadata_json, created_at) values ('ops_bad', 'request', 99, 'bad', '{}', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - - assert!(result.is_err()); -} - -#[tokio::test] -async fn sqlite_schema_should_reject_non_boolean_flags() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("boolean-check.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let api_key_result = sqlx::query( - "insert into client_api_keys (id, name, prefix, key, enabled, created_at) values ('key_bad', 'bad', 'sk_x', 'sk_test', 2, '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - let account_result = sqlx::query( - "insert into accounts (id, access_token, status, quota_limit_reached, added_at, updated_at) values ('acct_bad', 'access-token', 'active', 2, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - let quota_verify_result = sqlx::query( - "insert into accounts (id, access_token, status, quota_verify_required, added_at, updated_at) values ('acct_verify_bad', 'access-token', 'active', 2, '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - assert!(api_key_result.is_err()); - assert!(account_result.is_err()); - assert!(quota_verify_result.is_err()); -} - -#[tokio::test] -async fn sqlite_schema_should_reject_invalid_runtime_settings() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("settings-check.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let invalid_id_result = sqlx::query( - "insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (2, 300, 2, 3, 50, 'smart', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - let invalid_refresh_result = sqlx::query( - "insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (1, 0, 2, 3, 50, 'smart', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - let invalid_strategy_result = sqlx::query( - "insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (1, 300, 2, 3, 50, 'random', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - - assert!(invalid_id_result.is_err()); - assert!(invalid_refresh_result.is_err()); - assert!(invalid_strategy_result.is_err()); -} - -#[tokio::test] -async fn sqlite_schema_should_reject_negative_usage_time_bucket_counts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("usage-time-bucket-check.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - let result = sqlx::query( - "insert into usage_time_buckets (bucket_start, request_count, updated_at) values ('2026-06-14T00:00:00Z', -1, '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await; - - assert!(result.is_err()); -} - -#[tokio::test] -async fn sqlite_schema_should_reject_negative_account_usage_counts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("usage-check.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - sqlx::query( - "insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_a', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", - ) - .execute(&pool) - .await - .unwrap(); - let result = - sqlx::query("insert into account_usage (account_id, request_count) values ('acct_a', -1)") + assert!( + sqlx::query("update runtime_settings set usage_retention_days = 0 where id = 1",) .execute(&pool) - .await; - - assert!(result.is_err()); + .await + .is_err() + ); } -#[tokio::test] -async fn sqlite_schema_should_reject_negative_account_model_usage_counts() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("model-usage-check.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - - sqlx::query( - "insert into accounts (id, access_token, status, added_at, updated_at) values ('acct_a', 'access-token', 'active', '2026-06-14T00:00:00Z', '2026-06-14T00:00:00Z')", +async fn columns(pool: &sqlx::PgPool, table: &str) -> BTreeSet { + sqlx::query_scalar::<_, String>( + "select column_name + from information_schema.columns + where table_schema = current_schema() and table_name = $1", ) - .execute(&pool) + .bind(table) + .fetch_all(pool) .await - .unwrap(); - let result = sqlx::query( - "insert into account_model_usage (account_id, model, request_count) values ('acct_a', 'gpt-5.5', -1)", - ) - .execute(&pool) - .await; - - assert!(result.is_err()); -} - -#[tokio::test] -async fn connect_sqlite_creates_parent_directory_for_file_database() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("nested").join("startup.sqlite"); - let url = format!("sqlite://{}", db.display()); - - let pool = connect_sqlite(&url).await.unwrap(); - pool.close().await; - - assert!(db.exists()); + .unwrap() + .into_iter() + .collect() } diff --git a/backend/tests/main.rs b/backend/tests/main.rs index 3d2e535be..67a5380f8 100644 --- a/backend/tests/main.rs +++ b/backend/tests/main.rs @@ -1,8 +1,9 @@ -mod admin; -mod config; -mod http; +mod accounts; +mod api; +mod bootstrap; +mod dispatch; mod infra; -mod proxy; -mod runtime; +mod models; mod support; +mod telemetry; mod upstream; diff --git a/backend/tests/upstream/models/catalog.rs b/backend/tests/models/catalog.rs similarity index 88% rename from backend/tests/upstream/models/catalog.rs rename to backend/tests/models/catalog.rs index bd8940be2..1c668fdb8 100644 --- a/backend/tests/upstream/models/catalog.rs +++ b/backend/tests/models/catalog.rs @@ -8,20 +8,18 @@ use std::{ use async_trait::async_trait; -use codex_proxy_rs::upstream::accounts::model::{Account, AccountStatus}; -use codex_proxy_rs::upstream::models::backend_entry::{ - BackendModelEntry, BackendReasoningEffort, BackendTruncationPolicy, -}; -use codex_proxy_rs::upstream::models::catalog::ModelCatalog; -use codex_proxy_rs::upstream::models::config::ModelConfig; -use codex_proxy_rs::upstream::models::service::{ +use codex_proxy_rs::accounts::account::{Account, AccountStatus}; +use codex_proxy_rs::models::catalog::ModelCatalog; +use codex_proxy_rs::models::config::ModelConfig; +use codex_proxy_rs::models::service::{ ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError, }; -use codex_proxy_rs::upstream::models::snapshot::ModelPlanSnapshot; -use codex_proxy_rs::upstream::models::snapshot_store::{ - ModelSnapshotStore, ModelSnapshotStoreResult, +use codex_proxy_rs::models::snapshot::ModelPlanSnapshot; +use codex_proxy_rs::models::store::{ModelSnapshotStore, ModelSnapshotStoreResult}; +use codex_proxy_rs::upstream::openai::protocol::model_catalog::{ + BackendModelEntry, BackendReasoningEffort, BackendTruncationPolicy, }; -use codex_proxy_rs::upstream::transport::{ +use codex_proxy_rs::upstream::openai::transport::{ CodexModelCatalogClient, CodexModelCatalogClientError, CodexModelCatalogRequest, }; @@ -104,7 +102,7 @@ fn model_catalog_should_merge_backend_snapshots_and_build_plan_allowlist() { } #[tokio::test] -async fn model_service_should_return_builtin_catalog_when_snapshot_store_is_missing() { +async fn model_service_should_return_builtin_catalog_when_store_is_missing() { let service = ModelService::new(test_model_config(), None, None); let catalog = service.catalog().await; @@ -116,15 +114,15 @@ async fn model_service_should_return_builtin_catalog_when_snapshot_store_is_miss #[tokio::test] async fn model_service_catalog_should_use_loaded_memory_snapshot() { - let snapshot_store = Arc::new(InMemorySnapshotStore::default()); - snapshot_store + let store = Arc::new(InMemorySnapshotStore::default()); + store .replace_plan_snapshot(&ModelPlanSnapshot::from_backend_entries( "plus", backend_models("gpt-6"), )) .await .unwrap(); - let service = ModelService::new(test_model_config(), Some(snapshot_store.clone()), None); + let service = ModelService::new(test_model_config(), Some(store.clone()), None); service .reload_from_store() @@ -135,19 +133,19 @@ async fn model_service_catalog_should_use_loaded_memory_snapshot() { assert!(first.is_recognized_model_name("gpt-6")); assert!(second.is_recognized_model_name("gpt-6")); - assert_eq!(snapshot_store.list_calls.load(Ordering::SeqCst), 1); + assert_eq!(store.list_calls.load(Ordering::SeqCst), 1); } #[tokio::test] async fn model_service_should_refresh_plan_accounts_and_build_routing() { - let snapshot_store = Arc::new(InMemorySnapshotStore::default()); + let store = Arc::new(InMemorySnapshotStore::default()); let upstream = Arc::new(FakeModelCatalogClient::with_models(BTreeMap::from([ ("plus".to_string(), backend_models("gpt-6")), ("team".to_string(), backend_models("gpt-7")), ]))); let service = ModelService::new( test_model_config(), - Some(snapshot_store.clone()), + Some(store.clone()), Some(upstream.clone()), ); @@ -173,7 +171,7 @@ async fn model_service_should_refresh_plan_accounts_and_build_routing() { ); assert_eq!(upstream.recorded_plan_types().await, vec!["plus", "team"]); - let snapshots = snapshot_store.list_plan_snapshots().await.unwrap(); + let snapshots = store.list_plan_snapshots().await.unwrap(); assert_eq!(snapshots.len(), 2); assert_eq!(snapshots[0].plan_type, "plus"); assert_eq!(snapshots[1].plan_type, "team"); @@ -196,9 +194,9 @@ async fn model_service_should_refresh_plan_accounts_and_build_routing() { #[tokio::test] async fn model_service_should_return_no_accounts_when_no_plan_accounts_exist() { - let snapshot_store = Arc::new(InMemorySnapshotStore::default()); + let store = Arc::new(InMemorySnapshotStore::default()); let upstream = Arc::new(FakeModelCatalogClient::default()); - let service = ModelService::new(test_model_config(), Some(snapshot_store), Some(upstream)); + let service = ModelService::new(test_model_config(), Some(store), Some(upstream)); let error = service .refresh_backend_models_with_installation_id(&[], "req-model-refresh", None) diff --git a/backend/tests/upstream/models/mod.rs b/backend/tests/models/mod.rs similarity index 100% rename from backend/tests/upstream/models/mod.rs rename to backend/tests/models/mod.rs diff --git a/backend/tests/upstream/models/store.rs b/backend/tests/models/store.rs similarity index 54% rename from backend/tests/upstream/models/store.rs rename to backend/tests/models/store.rs index a2a507886..d3f411752 100644 --- a/backend/tests/upstream/models/store.rs +++ b/backend/tests/models/store.rs @@ -1,18 +1,17 @@ -use codex_proxy_rs::upstream::models::{ - backend_entry::BackendModelEntry, - snapshot::ModelPlanSnapshot, - snapshot_store::{ModelSnapshotStore, SqliteModelSnapshotStore}, +use codex_proxy_rs::{ + models::{ + snapshot::ModelPlanSnapshot, + store::{ModelSnapshotStore, RedisModelSnapshotStore}, + }, + upstream::openai::protocol::model_catalog::BackendModelEntry, }; +use crate::support::storage::create_test_redis; + #[tokio::test] -async fn model_snapshot_repository_should_replace_and_load_plan_snapshots() { - let dir = tempfile::tempdir().unwrap(); - let db = dir.path().join("models.sqlite"); - let pool = - codex_proxy_rs::infra::database::connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .unwrap(); - let repo = SqliteModelSnapshotStore::new(pool); +async fn model_snapshot_store_should_replace_and_load_plan_snapshots() { + let redis = create_test_redis("model-snapshots").await; + let repo = RedisModelSnapshotStore::new(redis); let snapshot = ModelPlanSnapshot::from_backend_entries( "team", vec![BackendModelEntry { diff --git a/backend/tests/proxy/dispatch/mod.rs b/backend/tests/proxy/dispatch/mod.rs deleted file mode 100644 index 50837bb5d..000000000 --- a/backend/tests/proxy/dispatch/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub(crate) mod chat_upstream; -pub(crate) mod session_affinity; -pub(crate) mod session_affinity_integration; diff --git a/backend/tests/proxy/dispatch/session_affinity.rs b/backend/tests/proxy/dispatch/session_affinity.rs deleted file mode 100644 index 0cec956a9..000000000 --- a/backend/tests/proxy/dispatch/session_affinity.rs +++ /dev/null @@ -1,276 +0,0 @@ -use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::infra::database::connect_sqlite; -use codex_proxy_rs::proxy::dispatch::session_affinity::{ - SessionAffinityEntry, SqliteSessionAffinityStore, -}; - -#[tokio::test] -async fn session_affinity_store_should_upsert_and_list_active_records() { - let dir = tempfile::tempdir().expect("tempdir"); - let db = dir.path().join("session-affinity.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - insert_account(&pool, "acct_1").await; - let store = SqliteSessionAffinityStore::new(pool); - let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); - - store - .upsert( - "resp_1", - &SessionAffinityEntry { - account_id: "acct_1".to_string(), - conversation_id: "conv_1".to_string(), - turn_state: Some("turn_1".to_string()), - instructions_hash: Some("hash_1".to_string()), - input_tokens: Some(128), - function_call_ids: vec!["call_a".to_string(), "call_b".to_string()], - variant_hash: Some("variant_1".to_string()), - created_at: now, - }, - Duration::hours(4), - ) - .await - .expect("affinity should be stored"); - - let records = store - .list_active(now + Duration::minutes(1)) - .await - .expect("active affinities should load"); - - assert_eq!(records.len(), 1); - assert_eq!(records[0].response_id, "resp_1"); - assert_eq!(records[0].entry.account_id, "acct_1"); - assert_eq!(records[0].entry.function_call_ids, vec!["call_a", "call_b"]); -} - -#[tokio::test] -async fn session_affinity_store_should_delete_expired_records() { - let dir = tempfile::tempdir().expect("tempdir"); - let db = dir.path().join("session-affinity-cleanup.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - insert_account(&pool, "acct_1").await; - insert_account(&pool, "acct_2").await; - let store = SqliteSessionAffinityStore::new(pool); - let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); - - store - .upsert( - "resp_expired", - &SessionAffinityEntry { - account_id: "acct_1".to_string(), - conversation_id: "conv_1".to_string(), - turn_state: None, - instructions_hash: None, - input_tokens: None, - function_call_ids: Vec::new(), - variant_hash: None, - created_at: now - Duration::hours(5), - }, - Duration::hours(1), - ) - .await - .expect("expired affinity should be stored"); - store - .upsert( - "resp_active", - &SessionAffinityEntry { - account_id: "acct_2".to_string(), - conversation_id: "conv_2".to_string(), - turn_state: None, - instructions_hash: None, - input_tokens: None, - function_call_ids: Vec::new(), - variant_hash: None, - created_at: now, - }, - Duration::hours(1), - ) - .await - .expect("active affinity should be stored"); - - let deleted = store - .delete_expired(now) - .await - .expect("expired affinities should be deleted"); - let records = store - .list_active(now) - .await - .expect("active affinities should load"); - - assert_eq!(deleted, 1); - assert_eq!(records.len(), 1); - assert_eq!(records[0].response_id, "resp_active"); -} - -#[tokio::test] -async fn session_affinity_store_should_forget_record_by_response_id() { - let dir = tempfile::tempdir().expect("tempdir"); - let db = dir.path().join("session-affinity-forget.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - insert_account(&pool, "acct_1").await; - let store = SqliteSessionAffinityStore::new(pool); - let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); - - store - .upsert( - "resp_forget", - &SessionAffinityEntry { - account_id: "acct_1".to_string(), - conversation_id: "conv_forget".to_string(), - turn_state: Some("turn_forget".to_string()), - instructions_hash: None, - input_tokens: None, - function_call_ids: Vec::new(), - variant_hash: None, - created_at: now, - }, - Duration::hours(1), - ) - .await - .expect("affinity should be stored"); - - let forgotten = store - .forget("resp_forget") - .await - .expect("affinity should be forgotten"); - let records = store - .list_active(now) - .await - .expect("active affinities should load"); - - assert!(forgotten); - assert!(records.is_empty()); -} - -async fn insert_account(pool: &sqlx::SqlitePool, id: &str) { - sqlx::query( - "insert into accounts (id, access_token, status, added_at, updated_at) values (?, 'access-token', 'active', '2026-06-18T12:00:00Z', '2026-06-18T12:00:00Z')", - ) - .bind(id) - .execute(pool) - .await - .expect("account should be inserted"); -} - -// --------------------------------------------------------------------------- -// In-memory SessionAffinityMap tests -// --------------------------------------------------------------------------- - -use codex_proxy_rs::proxy::dispatch::session_affinity::{ - SessionAffinityMap, StoredSessionAffinity, -}; - -#[test] -fn session_affinity_map_should_restore_active_records_and_skip_expired_records() { - let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); - let mut map = SessionAffinityMap::new(Duration::hours(4)); - - let restored = map.restore( - vec![ - StoredSessionAffinity { - response_id: "resp_active".to_string(), - entry: SessionAffinityEntry { - account_id: "acct_active".to_string(), - conversation_id: "conv_active".to_string(), - turn_state: Some("turn_active".to_string()), - instructions_hash: Some("hash_active".to_string()), - input_tokens: Some(42), - function_call_ids: vec!["call_1".to_string()], - variant_hash: Some("variant_a".to_string()), - created_at: now - Duration::minutes(30), - }, - expires_at: now + Duration::hours(1), - }, - StoredSessionAffinity { - response_id: "resp_expired".to_string(), - entry: SessionAffinityEntry { - account_id: "acct_expired".to_string(), - conversation_id: "conv_expired".to_string(), - turn_state: None, - instructions_hash: None, - input_tokens: None, - function_call_ids: Vec::new(), - variant_hash: None, - created_at: now - Duration::hours(5), - }, - expires_at: now + Duration::hours(1), - }, - ], - now, - ); - - assert_eq!(restored, 1); - assert_eq!( - map.lookup_account("resp_active", now), - Some("acct_active".to_string()) - ); - assert_eq!(map.lookup_account("resp_expired", now), None); -} - -#[test] -fn session_affinity_map_should_find_latest_response_by_conversation_and_variant() { - let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); - let mut map = SessionAffinityMap::new(Duration::hours(4)); - - map.record( - "resp_old".to_string(), - SessionAffinityEntry { - account_id: "acct_a".to_string(), - conversation_id: "conv".to_string(), - turn_state: Some("turn_old".to_string()), - instructions_hash: None, - input_tokens: None, - function_call_ids: Vec::new(), - variant_hash: Some("variant_a".to_string()), - created_at: now - Duration::minutes(20), - }, - ); - map.record( - "resp_new".to_string(), - SessionAffinityEntry { - account_id: "acct_b".to_string(), - conversation_id: "conv".to_string(), - turn_state: Some("turn_new".to_string()), - instructions_hash: None, - input_tokens: None, - function_call_ids: Vec::new(), - variant_hash: Some("variant_a".to_string()), - created_at: now - Duration::minutes(5), - }, - ); - map.record( - "resp_other_variant".to_string(), - SessionAffinityEntry { - account_id: "acct_c".to_string(), - conversation_id: "conv".to_string(), - turn_state: Some("turn_other".to_string()), - instructions_hash: None, - input_tokens: None, - function_call_ids: Vec::new(), - variant_hash: Some("variant_b".to_string()), - created_at: now - Duration::minutes(1), - }, - ); - - assert_eq!( - map.lookup_latest_response_by_conversation("conv", None, Some("variant_a"), now), - Some("resp_new".to_string()) - ); - assert_eq!( - map.lookup_latest_account_by_conversation("conv", None, Some("variant_a"), now), - Some("acct_b".to_string()) - ); - assert_eq!( - map.lookup_latest_account_by_conversation("conv", None, Some("variant_b"), now), - Some("acct_c".to_string()) - ); - assert_eq!( - map.lookup_turn_state("resp_new", now), - Some("turn_new".to_string()) - ); -} diff --git a/backend/tests/proxy/dispatch/session_affinity_integration.rs b/backend/tests/proxy/dispatch/session_affinity_integration.rs deleted file mode 100644 index 55bb48503..000000000 --- a/backend/tests/proxy/dispatch/session_affinity_integration.rs +++ /dev/null @@ -1,53 +0,0 @@ -use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::infra::database::connect_sqlite; -use codex_proxy_rs::proxy::dispatch::session_affinity::{ - SessionAffinityEntry, SqliteSessionAffinityStore, -}; - -#[tokio::test] -async fn session_affinity_store_should_restore_active_records() { - let dir = tempfile::tempdir().expect("tempdir"); - let db = dir.path().join("runtime-session-affinity.sqlite"); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.expect("sqlite pool"); - let store = SqliteSessionAffinityStore::new(pool.clone()); - let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap(); - - insert_account(&pool, "acct_restore").await; - store - .upsert( - "resp_restore", - &SessionAffinityEntry { - account_id: "acct_restore".to_string(), - conversation_id: "conv_restore".to_string(), - turn_state: Some("turn_restore".to_string()), - instructions_hash: Some("hash_restore".to_string()), - input_tokens: Some(7), - function_call_ids: vec!["call_restore".to_string()], - variant_hash: Some("variant_restore".to_string()), - created_at: now, - }, - Duration::hours(4), - ) - .await - .expect("affinity should be stored"); - - let records = store - .list_active(now + Duration::minutes(1)) - .await - .expect("active affinities should load"); - - assert_eq!(records.len(), 1); - assert_eq!(records[0].response_id, "resp_restore"); - assert_eq!(records[0].entry.account_id, "acct_restore"); -} - -async fn insert_account(pool: &sqlx::SqlitePool, id: &str) { - sqlx::query( - "insert into accounts (id, access_token, status, added_at, updated_at) values (?, 'access-token', 'active', '2026-06-18T12:00:00Z', '2026-06-18T12:00:00Z')", - ) - .bind(id) - .execute(pool) - .await - .expect("account should be inserted"); -} diff --git a/backend/tests/proxy/keys/mod.rs b/backend/tests/proxy/keys/mod.rs deleted file mode 100644 index 6f8ff8b9d..000000000 --- a/backend/tests/proxy/keys/mod.rs +++ /dev/null @@ -1 +0,0 @@ -pub(crate) mod auth; diff --git a/backend/tests/proxy/mod.rs b/backend/tests/proxy/mod.rs deleted file mode 100644 index 62d3d5566..000000000 --- a/backend/tests/proxy/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -pub(crate) mod auth; -pub(crate) mod dispatch; -pub(crate) mod keys; -pub(crate) mod openai; diff --git a/backend/tests/runtime/services.rs b/backend/tests/runtime/services.rs deleted file mode 100644 index d5ba0bb40..000000000 --- a/backend/tests/runtime/services.rs +++ /dev/null @@ -1,119 +0,0 @@ -use std::process::Command; - -use codex_proxy_rs::{ - admin::{ - auth::service::SqliteAdminSessionStore, keys::service::SqliteClientKeyStore, - monitoring::usage_record_store::SqliteUsageRecordStore, - }, - runtime::services::{BackgroundTaskStores, Services}, - upstream::{ - accounts::{ - cookies::SqliteCookieStore, store::SqliteAccountStore, token_refresh::RefreshLeaseStore, - }, - fingerprint::FingerprintRepository, - transport::tls::CODEX_CA_CERT_ENV, - }, -}; -use sqlx::SqlitePool; - -use crate::support::{config::test_config, sqlite::init_test_db}; - -fn stores(pool: SqlitePool) -> BackgroundTaskStores { - BackgroundTaskStores { - accounts: SqliteAccountStore::new(pool.clone()), - admin_sessions: SqliteAdminSessionStore::new(pool.clone()), - cookies: SqliteCookieStore::new(pool.clone()), - fingerprints: FingerprintRepository::new(pool.clone()), - session_affinity: - codex_proxy_rs::proxy::dispatch::session_affinity::SqliteSessionAffinityStore::new( - pool.clone(), - ), - refresh_leases: RefreshLeaseStore::new(pool.clone()), - client_keys: SqliteClientKeyStore::new(pool.clone()), - usage_records: SqliteUsageRecordStore::new(pool), - } -} - -#[tokio::test] -async fn services_try_new_should_use_configured_tls_transport_builder() { - const CASE_ENV: &str = "CODEX_PROXY_TEST_SERVICES_TLS_CASE"; - - if std::env::var(CASE_ENV).as_deref() == Ok("invalid_ca") { - let (pool, temp_dir) = init_test_db("services-tls.sqlite").await; - let config = test_config(format!( - "sqlite://{}", - temp_dir.path().join("services-tls.sqlite").display() - )); - let stores = stores(pool); - - let Err(error) = Services::try_new( - &config, - stores, - crate::support::fingerprint::runtime_test_fingerprint(), - ) else { - panic!("invalid custom CA should fail service transport construction"); - }; - assert!(error - .to_string() - .contains("Failed to read CA certificate file")); - return; - } - - let current_exe = std::env::current_exe().expect("current test binary path"); - let output = Command::new(current_exe) - .arg("--exact") - .arg("runtime::services::services_try_new_should_use_configured_tls_transport_builder") - .arg("--nocapture") - .env(CASE_ENV, "invalid_ca") - .env(CODEX_CA_CERT_ENV, "/tmp/codex-proxy-rs-missing-ca.pem") - .output() - .expect("run isolated services TLS test case"); - - assert!( - output.status.success(), - "isolated services TLS test failed\nstdout:\n{}\nstderr:\n{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); -} - -#[tokio::test] -async fn services_try_new_should_expose_websocket_pool_when_enabled() { - let (pool, temp_dir) = init_test_db("services-ws-pool-enabled.sqlite").await; - let config = test_config(format!( - "sqlite://{}", - temp_dir - .path() - .join("services-ws-pool-enabled.sqlite") - .display() - )); - let services = Services::try_new( - &config, - stores(pool), - crate::support::fingerprint::runtime_test_fingerprint(), - ) - .expect("services should build"); - - assert!(services.websocket_pool.is_some()); -} - -#[tokio::test] -async fn services_try_new_should_not_expose_websocket_pool_when_disabled() { - let (pool, temp_dir) = init_test_db("services-ws-pool-disabled.sqlite").await; - let mut config = test_config(format!( - "sqlite://{}", - temp_dir - .path() - .join("services-ws-pool-disabled.sqlite") - .display() - )); - config.ws_pool.enabled = false; - let services = Services::try_new( - &config, - stores(pool), - crate::support::fingerprint::runtime_test_fingerprint(), - ) - .expect("services should build"); - - assert!(services.websocket_pool.is_none()); -} diff --git a/backend/tests/runtime/tasks/cleanup.rs b/backend/tests/runtime/tasks/cleanup.rs deleted file mode 100644 index 4f8978d6e..000000000 --- a/backend/tests/runtime/tasks/cleanup.rs +++ /dev/null @@ -1,121 +0,0 @@ -use super::*; - -#[tokio::test] -async fn cookie_cleanup_task_should_delete_only_expired_cookie_rows() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("cookies.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteCookieStore::new(pool.clone()); - insert_account(&pool, "acct-cookie").await; - let now = Utc::now(); - insert_cookie( - &pool, - "expired-cookie", - "acct-cookie", - "old", - now - Duration::minutes(1), - ) - .await; - insert_cookie( - &pool, - "active-cookie", - "acct-cookie", - "fresh", - now + Duration::minutes(10), - ) - .await; - - let deleted = codex_proxy_rs::runtime::tasks::cookie_cleanup::CookieCleanupTask::new(store) - .cleanup_once_at(now) - .await - .expect("cookie cleanup should succeed"); - - let remaining = cookie_ids(&pool).await; - - assert_eq!((deleted, remaining), (1, vec!["active-cookie".to_string()])); -} - -#[tokio::test] -async fn session_cleanup_task_should_delete_only_expired_sessions() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("sessions.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteAdminSessionStore::new(pool.clone()); - let now = Utc::now(); - insert_admin_user(&pool, "admin").await; - insert_admin_session( - &pool, - "expired-session", - "admin", - now - Duration::minutes(1), - ) - .await; - insert_admin_session( - &pool, - "active-session", - "admin", - now + Duration::minutes(10), - ) - .await; - - let deleted = - codex_proxy_rs::runtime::tasks::session_cleanup::SessionCleanupTask::new(store, 3600) - .cleanup_once_at(now) - .await - .expect("session cleanup should succeed"); - - let remaining = admin_session_ids(&pool).await; - - assert_eq!( - (deleted, remaining), - (1, vec!["active-session".to_string()]) - ); -} - -#[tokio::test] -async fn session_affinity_cleanup_task_should_delete_only_expired_affinities() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("session-affinity-cleanup.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let store = SqliteSessionAffinityStore::new(pool.clone()); - let now = Utc::now(); - insert_account(&pool, "expired-account").await; - insert_account(&pool, "active-account").await; - store - .upsert( - "expired-response", - &session_affinity_entry("expired-account", now - Duration::hours(2)), - Duration::hours(1), - ) - .await - .expect("expired affinity should be inserted"); - store - .upsert( - "active-response", - &session_affinity_entry("active-account", now), - Duration::hours(1), - ) - .await - .expect("active affinity should be inserted"); - - let deleted = - codex_proxy_rs::runtime::tasks::session_affinity_cleanup::SessionAffinityCleanupTask::new( - store, 3600, - ) - .cleanup_once_at(now) - .await - .expect("session affinity cleanup should succeed"); - - let remaining = session_affinity_response_ids(&pool).await; - - assert_eq!( - (deleted, remaining), - (1, vec!["active-response".to_string()]) - ); -} diff --git a/backend/tests/support/accounts.rs b/backend/tests/support/accounts.rs index 72460533c..95caef765 100644 --- a/backend/tests/support/accounts.rs +++ b/backend/tests/support/accounts.rs @@ -1,5 +1,5 @@ use chrono::{Duration, Utc}; -use codex_proxy_rs::upstream::accounts::model::{Account, AccountStatus}; +use codex_proxy_rs::accounts::account::{Account, AccountStatus}; pub(crate) fn test_account(id: &str, status: AccountStatus) -> Account { Account { diff --git a/backend/tests/support/admin.rs b/backend/tests/support/admin.rs index f7e536588..f51dba364 100644 --- a/backend/tests/support/admin.rs +++ b/backend/tests/support/admin.rs @@ -1,25 +1,20 @@ -use sqlx::SqlitePool; +use chrono::Duration; +use codex_proxy_rs::{auth::store::RedisAdminSessionStore, infra::redis::RedisConnection}; +use sqlx::PgPool; -pub(crate) async fn seed_admin_session(pool: &SqlitePool, session_id: &str) { +pub(crate) async fn seed_admin_session(pool: &PgPool, redis: &RedisConnection, session_id: &str) { sqlx::query( - "insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)", + "insert into admin_users (id, password_hash, created_at, updated_at) + values ($1, $2, now(), now()) on conflict (id) do nothing", ) .bind("admin_1") .bind("hash") - .bind("2026-06-18T00:00:00Z") - .bind("2026-06-18T00:00:00Z") .execute(pool) .await .unwrap(); - sqlx::query( - "insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)", - ) - .bind(session_id) - .bind("admin_1") - .bind("2999-01-01T00:00:00Z") - .bind("2026-06-18T00:00:00Z") - .execute(pool) - .await - .unwrap(); + RedisAdminSessionStore::new(redis.clone()) + .create_session(session_id, "admin_1", Duration::days(1)) + .await + .unwrap(); } diff --git a/backend/tests/support/client_keys.rs b/backend/tests/support/client_keys.rs index 1044513ed..a195d51c9 100644 --- a/backend/tests/support/client_keys.rs +++ b/backend/tests/support/client_keys.rs @@ -1,8 +1,8 @@ -use codex_proxy_rs::admin::keys::service::SqliteClientKeyStore; -use sqlx::SqlitePool; +use codex_proxy_rs::keys::{manage::KeyManageService, store::PgClientKeyStore}; +use sqlx::PgPool; -pub(crate) async fn insert_client_api_key(pool: &SqlitePool) -> String { - SqliteClientKeyStore::new(pool.clone()) +pub(crate) async fn insert_client_api_key(pool: &PgPool) -> String { + KeyManageService::new(PgClientKeyStore::new(pool.clone())) .create("test") .await .unwrap() diff --git a/backend/tests/support/config.rs b/backend/tests/support/config.rs index cd479bb95..77c156cb0 100644 --- a/backend/tests/support/config.rs +++ b/backend/tests/support/config.rs @@ -1,6 +1,6 @@ -use codex_proxy_rs::config::schema::{ +use codex_proxy_rs::bootstrap::config::{ AdminConfig, ApiConfig, AppConfig, AuthConfig, DatabaseConfig, FingerprintConfig, - LoggingConfig, QuotaConfig, ServerConfig, TlsConfig, WebSocketPoolConfig, + LoggingConfig, QuotaConfig, RedisConfig, ServerConfig, TlsConfig, WebSocketPoolConfig, }; pub(crate) fn test_config(database_url: String) -> AppConfig { @@ -29,6 +29,9 @@ pub(crate) fn test_config(database_url: String) -> AppConfig { skip_exhausted: true, }, database: DatabaseConfig { url: database_url }, + redis: RedisConfig { + url: crate::support::storage::test_redis_url(), + }, tls: TlsConfig { force_http11: false, }, @@ -36,7 +39,6 @@ pub(crate) fn test_config(database_url: String) -> AppConfig { fingerprint: FingerprintConfig::default(), admin: AdminConfig { session_ttl_minutes: 1440, - session_cleanup_interval_secs: 3600, default_username: "admin".to_string(), default_password: "test-admin-password".to_string(), }, diff --git a/backend/tests/support/fingerprint.rs b/backend/tests/support/fingerprint.rs index 038da36a1..a9c1eb0a5 100644 --- a/backend/tests/support/fingerprint.rs +++ b/backend/tests/support/fingerprint.rs @@ -1,10 +1,10 @@ use codex_proxy_rs::{ - config::schema::FingerprintConfig, - upstream::fingerprint::{Fingerprint, RuntimeFingerprint}, + bootstrap::{config::FingerprintConfig, services::fingerprint_from_config}, + upstream::openai::fingerprint::{Fingerprint, RuntimeFingerprint}, }; pub(crate) fn test_fingerprint() -> Fingerprint { - Fingerprint::from_config(&FingerprintConfig::default()) + fingerprint_from_config(&FingerprintConfig::default()) } pub(crate) fn test_fingerprint_with_updated_at(updated_at: Option<&str>) -> Fingerprint { diff --git a/backend/tests/support/mod.rs b/backend/tests/support/mod.rs index d44fcc735..059cce885 100644 --- a/backend/tests/support/mod.rs +++ b/backend/tests/support/mod.rs @@ -6,4 +6,4 @@ pub(crate) mod config; pub(crate) mod fingerprint; pub(crate) mod http; pub(crate) mod jwt; -pub(crate) mod sqlite; +pub(crate) mod storage; diff --git a/backend/tests/support/sqlite.rs b/backend/tests/support/sqlite.rs deleted file mode 100644 index c6735c29e..000000000 --- a/backend/tests/support/sqlite.rs +++ /dev/null @@ -1,10 +0,0 @@ -use codex_proxy_rs::infra::database::connect_sqlite; -use sqlx::SqlitePool; - -pub(crate) async fn init_test_db(db_name: &str) -> (SqlitePool, tempfile::TempDir) { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join(db_name); - let url = format!("sqlite://{}", db.display()); - let pool = connect_sqlite(&url).await.unwrap(); - (pool, dir) -} diff --git a/backend/tests/support/storage.rs b/backend/tests/support/storage.rs new file mode 100644 index 000000000..f0327835f --- /dev/null +++ b/backend/tests/support/storage.rs @@ -0,0 +1,150 @@ +use std::env; + +use chrono::{DateTime, Utc}; +use codex_proxy_rs::{ + accounts::{cookies::PgCookieStore, refresh::RedisRefreshLeaseStore, store::PgAccountStore}, + auth::store::{PgAdminUserStore, RedisAdminSessionStore}, + bootstrap::services::BackgroundTaskStores, + dispatch::affinity::RedisSessionAffinityStore, + infra::{database::migrate, redis::RedisConnection}, + keys::store::PgClientKeyStore, + models::store::RedisModelSnapshotStore, + telemetry::{ + account_usage::store::PgAccountUsageStore, buckets::store::PgRequestBucketStore, + ops::store::PgOpsErrorLogStore, usage::store::PgUsageRecordStore, + }, + upstream::openai::fingerprint::PgFingerprintStore, +}; +use sqlx::{postgres::PgPoolOptions, PgPool}; +use uuid::Uuid; + +const DEFAULT_TEST_DATABASE_URL: &str = + "postgres://codex_proxy:codex_proxy@127.0.0.1:5432/codex_proxy"; +const DEFAULT_TEST_REDIS_URL: &str = "redis://127.0.0.1:6379"; +static DATABASE_INIT_LIMIT: tokio::sync::Semaphore = tokio::sync::Semaphore::const_new(2); + +pub(crate) async fn init_test_db(label: &str) -> (PgPool, tempfile::TempDir) { + let guard = tempfile::tempdir().expect("test database guard"); + (create_test_database(label).await, guard) +} + +pub(crate) async fn create_test_database(label: &str) -> PgPool { + let _permit = DATABASE_INIT_LIMIT.acquire().await.unwrap(); + let database_name = test_database_name(label); + let base_url = test_database_url(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&base_url) + .await + .unwrap_or_else(|error| panic!("connect CPR_TEST_DATABASE_URL: {error}")); + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + r#"create database "{database_name}""# + ))) + .execute(&admin) + .await + .unwrap_or_else(|error| panic!("create test database {database_name}: {error}")); + admin.close().await; + + let pool = PgPoolOptions::new() + .max_connections(1) + .connect(&database_url_for(&base_url, &database_name)) + .await + .unwrap_or_else(|error| panic!("connect test database {database_name}: {error}")); + sqlx::query("set max_parallel_workers_per_gather = 0") + .execute(&pool) + .await + .unwrap(); + sqlx::query("set max_parallel_maintenance_workers = 0") + .execute(&pool) + .await + .unwrap(); + migrate(&pool) + .await + .unwrap_or_else(|error| panic!("initialize test database {database_name}: {error}")); + pool +} + +pub(crate) async fn create_test_redis(label: &str) -> RedisConnection { + let label = sanitize_label(label); + let prefix = format!("cpr:test:{label}:{}", Uuid::new_v4().simple()); + RedisConnection::connect(&test_redis_url(), prefix) + .await + .unwrap_or_else(|error| panic!("connect CPR_TEST_REDIS_URL: {error}")) +} + +pub(crate) fn background_task_stores(pool: PgPool, redis: RedisConnection) -> BackgroundTaskStores { + BackgroundTaskStores { + redis: redis.clone(), + accounts: PgAccountStore::new(pool.clone()), + admin_users: PgAdminUserStore::new(pool.clone()), + admin_sessions: RedisAdminSessionStore::new(redis.clone()), + cookies: PgCookieStore::new(pool.clone()), + fingerprints: PgFingerprintStore::new(pool.clone()), + session_affinity: RedisSessionAffinityStore::new(redis.clone()), + refresh_leases: RedisRefreshLeaseStore::new(redis.clone()), + model_snapshots: RedisModelSnapshotStore::new(redis), + client_keys: PgClientKeyStore::new(pool.clone()), + usage_records: PgUsageRecordStore::new(pool.clone()), + ops_errors: PgOpsErrorLogStore::new(pool.clone()), + account_usage: PgAccountUsageStore::new(pool.clone()), + request_buckets: PgRequestBucketStore::new(pool), + } +} + +pub(crate) fn test_database_url() -> String { + env::var("CPR_TEST_DATABASE_URL").unwrap_or_else(|_| DEFAULT_TEST_DATABASE_URL.to_string()) +} + +pub(crate) fn test_redis_url() -> String { + env::var("CPR_TEST_REDIS_URL").unwrap_or_else(|_| DEFAULT_TEST_REDIS_URL.to_string()) +} + +pub(crate) fn timestamp(value: &str) -> DateTime { + DateTime::parse_from_rfc3339(value) + .unwrap_or_else(|error| panic!("invalid test timestamp {value}: {error}")) + .to_utc() +} + +fn test_database_name(label: &str) -> String { + let label = sanitize_label(label); + let suffix = Uuid::new_v4().simple().to_string(); + format!( + "cpr_test_{}_{}", + label.chars().take(24).collect::(), + &suffix[..12] + ) +} + +fn sanitize_label(label: &str) -> String { + let label = label + .chars() + .map(|character| { + if character.is_ascii_alphanumeric() { + character.to_ascii_lowercase() + } else { + '_' + } + }) + .collect::(); + let label = label.trim_matches('_'); + if label.is_empty() { + "case".to_string() + } else { + label.to_string() + } +} + +fn database_url_for(base_url: &str, database_name: &str) -> String { + let (without_query, query) = base_url + .split_once('?') + .map_or((base_url, None), |(url, query)| (url, Some(query))); + let slash = without_query + .rfind('/') + .unwrap_or_else(|| panic!("CPR_TEST_DATABASE_URL must include a database name")); + let mut url = format!("{}/{database_name}", &without_query[..slash]); + if let Some(query) = query { + url.push('?'); + url.push_str(query); + } + url +} diff --git a/backend/tests/telemetry/mod.rs b/backend/tests/telemetry/mod.rs new file mode 100644 index 000000000..c24290dac --- /dev/null +++ b/backend/tests/telemetry/mod.rs @@ -0,0 +1 @@ +pub(crate) mod rebuild; diff --git a/backend/tests/telemetry/rebuild.rs b/backend/tests/telemetry/rebuild.rs new file mode 100644 index 000000000..75e5caafc --- /dev/null +++ b/backend/tests/telemetry/rebuild.rs @@ -0,0 +1,97 @@ +use chrono::{Duration, Utc}; +use codex_proxy_rs::telemetry::{ + ops::{store::PgOpsErrorLogStore, types::OpsErrorLog}, + rebuild::rebuild_buckets, + usage::{store::PgUsageRecordStore, types::UsageRecord}, +}; + +use crate::support::storage::init_test_db; + +#[tokio::test] +async fn rebuild_buckets_should_replace_only_fully_reconstructible_range() { + let (pool, _guard) = init_test_db("rebuild-buckets").await; + seed_runtime_settings(&pool, 30, 10).await; + + let now = Utc::now(); + let mut success = UsageRecord::new("v1.response", "ok", "acct-1", "gpt-5.5", 200); + success.created_at = now; + success.latency_ms = Some(0); + success.first_token_ms = Some(0); + success.input_tokens = Some(12); + success.output_tokens = Some(5); + success.cached_tokens = Some(3); + PgUsageRecordStore::new(pool.clone()) + .append(&success) + .await + .unwrap(); + + let mut failure = OpsErrorLog::new("v1.response", "failed"); + failure.created_at = now; + failure.provider = Some("openai".to_string()); + failure.account_id = Some("acct-1".to_string()); + failure.model = Some("gpt-5.5".to_string()); + failure.status_code = Some(502); + PgOpsErrorLogStore::new(pool.clone()) + .append(&failure) + .await + .unwrap(); + + sqlx::query( + "update request_time_buckets + set success_count = 99, error_count = 99, input_tokens = 999", + ) + .execute(&pool) + .await + .unwrap(); + + let protected_bucket = now - Duration::days(20); + sqlx::query( + "insert into request_time_buckets + (bucket_start, provider, account_id, model, service_tier, + success_count, error_count, updated_at) + values ($1, 'openai', 'old-account', 'old-model', '__unknown__', 7, 2, now())", + ) + .bind(protected_bucket) + .execute(&pool) + .await + .unwrap(); + + let report = rebuild_buckets(&pool).await.unwrap(); + + assert_eq!(report.deleted_rows, 1); + assert_eq!(report.rebuilt_rows, 1); + let rebuilt: (i64, i64, i64, i64, i64, i64, i64, Option) = sqlx::query_as( + "select success_count, error_count, input_tokens, output_tokens, cached_tokens, + first_token_latency_count, latency_count, min_latency_ms + from request_time_buckets where account_id = 'acct-1'", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(rebuilt, (1, 1, 12, 5, 3, 1, 1, Some(0))); + + let preserved: (i64, i64) = sqlx::query_as( + "select success_count, error_count + from request_time_buckets where account_id = 'old-account'", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(preserved, (7, 2)); + assert!(report.cutoff > now - Duration::days(11)); +} + +async fn seed_runtime_settings(pool: &sqlx::PgPool, usage_days: i64, ops_days: i64) { + sqlx::query( + "insert into runtime_settings + (id, model_aliases_json, refresh_margin_seconds, refresh_concurrency, + max_concurrent_per_account, request_interval_ms, rotation_strategy, + usage_retention_days, ops_error_retention_days, bucket_retention_days, updated_at) + values (1, '{}', 3600, 2, 3, 50, 'smart', $1, $2, 90, now())", + ) + .bind(usage_days) + .bind(ops_days) + .execute(pool) + .await + .unwrap(); +} diff --git a/backend/tests/upstream/accounts/refresh_leases.rs b/backend/tests/upstream/accounts/refresh_leases.rs deleted file mode 100644 index 31a3966bc..000000000 --- a/backend/tests/upstream/accounts/refresh_leases.rs +++ /dev/null @@ -1,133 +0,0 @@ -use chrono::{Duration, TimeZone, Utc}; -use codex_proxy_rs::{ - infra::database::connect_sqlite, - upstream::accounts::{ - model::AccountStatus, - store::{NewAccount, SqliteAccountStore}, - token_refresh::RefreshLeaseStore, - }, -}; -use secrecy::SecretString; - -#[tokio::test] -async fn refresh_lease_store_should_acquire_release_and_respect_expiry() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("refresh-leases.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let accounts = SqliteAccountStore::new(pool.clone()); - accounts - .insert(NewAccount { - id: "acct-lease".to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access".to_string().into()), - refresh_token: Some(SecretString::new("refresh".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - added_at: None, - }) - .await - .expect("account should be inserted"); - let leases = RefreshLeaseStore::new(pool); - let now = Utc.with_ymd_and_hms(2026, 6, 19, 8, 0, 0).unwrap(); - let first_expires_at = now + Duration::minutes(5); - let second_expires_at = now + Duration::minutes(10); - - assert!(leases - .try_acquire("acct-lease", "owner-a", first_expires_at, now) - .await - .expect("first owner should acquire")); - assert!(!leases - .try_acquire("acct-lease", "owner-b", second_expires_at, now) - .await - .expect("second owner should be blocked")); - assert!(leases - .try_acquire( - "acct-lease", - "owner-b", - second_expires_at, - first_expires_at + Duration::seconds(1), - ) - .await - .expect("second owner should acquire expired lease")); - assert!(!leases - .release("acct-lease", "owner-a") - .await - .expect("wrong owner should not release")); - assert!(leases - .release("acct-lease", "owner-b") - .await - .expect("current owner should release")); - assert!(leases - .try_acquire("acct-lease", "owner-c", second_expires_at, now) - .await - .expect("lease should be acquirable after release")); -} - -#[tokio::test] -async fn refresh_lease_store_should_list_active_account_ids() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("refresh-lease-active-ids.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let accounts = SqliteAccountStore::new(pool.clone()); - for id in ["acct-active-lease", "acct-expired-lease"] { - accounts - .insert(NewAccount { - id: id.to_string(), - email: None, - account_id: None, - user_id: None, - label: None, - plan_type: None, - access_token: SecretString::new("access".to_string().into()), - refresh_token: Some(SecretString::new("refresh".to_string().into())), - access_token_expires_at: None, - status: AccountStatus::Active, - added_at: None, - }) - .await - .expect("account should be inserted"); - } - let leases = RefreshLeaseStore::new(pool); - let now = Utc.with_ymd_and_hms(2026, 7, 9, 8, 0, 0).unwrap(); - leases - .try_acquire( - "acct-active-lease", - "owner-a", - now + Duration::minutes(5), - now, - ) - .await - .expect("active lease should be acquired"); - leases - .try_acquire( - "acct-expired-lease", - "owner-b", - now - Duration::minutes(5), - now - Duration::minutes(10), - ) - .await - .expect("expired lease should be acquired"); - - let account_ids = vec![ - "acct-active-lease".to_string(), - "acct-expired-lease".to_string(), - "acct-missing".to_string(), - ]; - let active = leases - .active_account_ids(&account_ids, now) - .await - .expect("active leases should load"); - - assert_eq!( - active, - std::collections::HashSet::from(["acct-active-lease".to_string()]) - ); -} diff --git a/backend/tests/upstream/mod.rs b/backend/tests/upstream/mod.rs index 9e415e77b..1e30e941c 100644 --- a/backend/tests/upstream/mod.rs +++ b/backend/tests/upstream/mod.rs @@ -1,7 +1 @@ -pub(crate) mod accounts; -pub(crate) mod fingerprint; -pub(crate) mod models; -pub(crate) mod protocol; -pub(crate) mod scheduler; -pub(crate) mod token_client; -pub(crate) mod transport; +pub(crate) mod openai; diff --git a/backend/tests/upstream/fingerprint/mod.rs b/backend/tests/upstream/openai/fingerprint/mod.rs similarity index 80% rename from backend/tests/upstream/fingerprint/mod.rs rename to backend/tests/upstream/openai/fingerprint/mod.rs index f9bf94d67..6449ba000 100644 --- a/backend/tests/upstream/fingerprint/mod.rs +++ b/backend/tests/upstream/openai/fingerprint/mod.rs @@ -1,14 +1,11 @@ -use codex_proxy_rs::infra::database::connect_sqlite; -use codex_proxy_rs::upstream::fingerprint::{FingerprintRepository, UpdateChecker}; +use codex_proxy_rs::upstream::openai::fingerprint::{PgFingerprintStore, UpdateChecker}; + +use crate::support::storage::init_test_db; #[tokio::test] -async fn fingerprint_repository_should_update_current_record() { - let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("fingerprints.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool.clone()); +async fn fingerprint_store_should_update_current_record() { + let (pool, _guard) = init_test_db("fingerprint-store").await; + let repo = PgFingerprintStore::new(pool.clone()); let default_fingerprint = crate::support::fingerprint::test_fingerprint(); repo.ensure_current_seed(&default_fingerprint) @@ -58,11 +55,8 @@ async fn update_checker_should_report_available_update_from_appcast() { .await; let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("fingerprints-check.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool); + let (pool, _guard) = init_test_db("fingerprint-update-check").await; + let repo = PgFingerprintStore::new(pool); repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) .await .expect("seed current fingerprint"); @@ -84,7 +78,7 @@ async fn update_checker_should_report_available_update_from_appcast() { } #[tokio::test] -async fn update_checker_should_apply_available_update_to_repository() { +async fn update_checker_should_apply_available_update_to_store() { let server = wiremock::MockServer::start().await; wiremock::Mock::given(wiremock::matchers::method("GET")) .and(wiremock::matchers::path("/appcast.xml")) @@ -104,11 +98,8 @@ async fn update_checker_should_apply_available_update_to_repository() { .await; let dir = tempfile::tempdir().expect("temp dir"); - let db = dir.path().join("fingerprints.sqlite"); - let pool = connect_sqlite(&format!("sqlite://{}", db.display())) - .await - .expect("sqlite pool"); - let repo = FingerprintRepository::new(pool); + let (pool, _guard) = init_test_db("fingerprint-update-apply").await; + let repo = PgFingerprintStore::new(pool); repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint()) .await .expect("seed current fingerprint"); diff --git a/backend/tests/upstream/openai/mod.rs b/backend/tests/upstream/openai/mod.rs new file mode 100644 index 000000000..c0d61d017 --- /dev/null +++ b/backend/tests/upstream/openai/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod fingerprint; +pub(crate) mod protocol; +pub(crate) mod token_client; +pub(crate) mod transport; diff --git a/backend/tests/upstream/protocol/codex_websocket.rs b/backend/tests/upstream/openai/protocol/codex_websocket.rs similarity index 99% rename from backend/tests/upstream/protocol/codex_websocket.rs rename to backend/tests/upstream/openai/protocol/codex_websocket.rs index 02d893820..d54e24c89 100644 --- a/backend/tests/upstream/protocol/codex_websocket.rs +++ b/backend/tests/upstream/openai/protocol/codex_websocket.rs @@ -64,7 +64,7 @@ fn codex_websocket_response_create_payload_text_should_preserve_canonical_field_ }))); let payload = - codex_proxy_rs::upstream::protocol::websocket::websocket_response_create_payload_text( + codex_proxy_rs::upstream::openai::protocol::websocket::websocket_response_create_payload_text( &request, ) .expect("payload should serialize"); @@ -91,7 +91,7 @@ fn codex_websocket_response_create_payload_text_should_include_empty_instruction let request = CodexResponsesRequest::new_http_sse("gpt-5.5", "", Vec::new()); let payload = - codex_proxy_rs::upstream::protocol::websocket::websocket_response_create_payload_text( + codex_proxy_rs::upstream::openai::protocol::websocket::websocket_response_create_payload_text( &request, ) .expect("payload should serialize"); diff --git a/backend/tests/upstream/protocol/mod.rs b/backend/tests/upstream/openai/protocol/mod.rs similarity index 89% rename from backend/tests/upstream/protocol/mod.rs rename to backend/tests/upstream/openai/protocol/mod.rs index 5614329a8..1e5f5728e 100644 --- a/backend/tests/upstream/protocol/mod.rs +++ b/backend/tests/upstream/openai/protocol/mod.rs @@ -1,18 +1,18 @@ use axum::http::HeaderMap; -use codex_proxy_rs::proxy::openai::responses::{ +use codex_proxy_rs::api::client::responses::{ build_codex_request, build_compact_request, response_failed_sse_event, }; -use codex_proxy_rs::upstream::protocol::events::{ +use codex_proxy_rs::upstream::openai::protocol::events::{ extract_sse_usage, extract_usage, parse_rate_limit_headers, parse_rate_limits_event, rate_limit_quota, retry_after_seconds_from_body, RateLimitWindow, TokenUsage, }; -use codex_proxy_rs::upstream::protocol::responses::{ +use codex_proxy_rs::upstream::openai::protocol::responses::{ http_sse_fallback_allowed, response_body_has_first_event, response_from_codex_sse, response_sse_event_is_terminal, transport_for_request, CodexResponsesRequest, CodexTransport, CollectedResponse, }; -use codex_proxy_rs::upstream::protocol::sse::parse_sse_events; -use codex_proxy_rs::upstream::protocol::websocket::{ +use codex_proxy_rs::upstream::openai::protocol::sse::parse_sse_events; +use codex_proxy_rs::upstream::openai::protocol::websocket::{ classify_websocket_error_frame, is_terminal_websocket_event, retry_after_seconds_from_wrapped_error_headers, websocket_agent_message_output_item_event_invalid_required_fields, diff --git a/backend/tests/upstream/protocol/openai_responses.rs b/backend/tests/upstream/openai/protocol/openai_responses.rs similarity index 98% rename from backend/tests/upstream/protocol/openai_responses.rs rename to backend/tests/upstream/openai/protocol/openai_responses.rs index d983e5aa6..7f26da1db 100644 --- a/backend/tests/upstream/protocol/openai_responses.rs +++ b/backend/tests/upstream/openai/protocol/openai_responses.rs @@ -70,7 +70,7 @@ fn codex_responses_transport_flags_should_not_serialize_to_upstream_json() { #[test] fn response_from_codex_sse_should_collect_completed_response_with_deltas() { - let body = include_str!("../../fixtures/responses/http_sse/text_delta_completed_usage.sse"); + let body = include_str!("../../../fixtures/responses/http_sse/text_delta_completed_usage.sse"); let response = response_from_codex_sse(body, None).expect("conversion should succeed"); @@ -101,7 +101,7 @@ fn response_from_codex_sse_should_collect_completed_response_with_deltas() { #[test] fn response_from_codex_sse_should_collect_incomplete_response() { - let body = include_str!("../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"); + let body = include_str!("../../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"); let response = response_from_codex_sse(body, None).expect("conversion should succeed"); @@ -135,7 +135,7 @@ fn response_from_codex_sse_should_collect_incomplete_response() { #[test] fn response_sse_event_is_terminal_should_match_incomplete_response() { - let body = include_str!("../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"); + let body = include_str!("../../../fixtures/responses/http_sse/chat_delta_incomplete_usage.sse"); let events = parse_sse_events(body).expect("sse should parse"); assert!(events.iter().any(response_sse_event_is_terminal)); diff --git a/backend/tests/upstream/protocol/response_first_event.rs b/backend/tests/upstream/openai/protocol/response_first_event.rs similarity index 58% rename from backend/tests/upstream/protocol/response_first_event.rs rename to backend/tests/upstream/openai/protocol/response_first_event.rs index db612c008..d32472568 100644 --- a/backend/tests/upstream/protocol/response_first_event.rs +++ b/backend/tests/upstream/openai/protocol/response_first_event.rs @@ -2,14 +2,14 @@ use super::*; #[test] fn first_event_detection_should_match_text_delta() { - let body = include_bytes!("../../fixtures/responses/http_sse/text_delta_hello.sse"); + let body = include_bytes!("../../../fixtures/responses/http_sse/text_delta_hello.sse"); assert!(response_body_has_first_event(body)); } #[test] fn first_event_detection_should_match_reasoning_delta() { - let body = include_bytes!("../../fixtures/responses/http_sse/reasoning_delta.sse"); + let body = include_bytes!("../../../fixtures/responses/http_sse/reasoning_delta.sse"); assert!(response_body_has_first_event(body)); } @@ -17,21 +17,21 @@ fn first_event_detection_should_match_reasoning_delta() { #[test] fn first_event_detection_should_match_function_call_arguments_delta() { let body = - include_bytes!("../../fixtures/responses/http_sse/function_call_arguments_delta.sse"); + include_bytes!("../../../fixtures/responses/http_sse/function_call_arguments_delta.sse"); assert!(response_body_has_first_event(body)); } #[test] fn first_event_detection_should_match_type_when_event_name_is_missing() { - let body = include_bytes!("../../fixtures/responses/http_sse/data_type_text_delta.sse"); + let body = include_bytes!("../../../fixtures/responses/http_sse/data_type_text_delta.sse"); assert!(response_body_has_first_event(body)); } #[test] fn first_event_detection_should_ignore_response_created() { - let body = include_bytes!("../../fixtures/responses/http_sse/response_created.sse"); + let body = include_bytes!("../../../fixtures/responses/http_sse/response_created.sse"); assert!(!response_body_has_first_event(body)); } @@ -39,7 +39,7 @@ fn first_event_detection_should_ignore_response_created() { #[test] fn first_event_detection_should_ignore_metadata_and_rate_limit_events() { let body = include_bytes!( - "../../fixtures/responses/http_sse/metadata_rate_limits_no_first_output.sse" + "../../../fixtures/responses/http_sse/metadata_rate_limits_no_first_output.sse" ); assert!(!response_body_has_first_event(body)); @@ -47,28 +47,28 @@ fn first_event_detection_should_ignore_metadata_and_rate_limit_events() { #[test] fn first_event_detection_should_ignore_empty_delta() { - let body = include_bytes!("../../fixtures/responses/http_sse/empty_text_delta.sse"); + let body = include_bytes!("../../../fixtures/responses/http_sse/empty_text_delta.sse"); assert!(!response_body_has_first_event(body)); } #[test] fn first_event_detection_should_ignore_done_frame() { - let body = include_bytes!("../../fixtures/responses/http_sse/done_only.sse"); + let body = include_bytes!("../../../fixtures/responses/http_sse/done_only.sse"); assert!(!response_body_has_first_event(body)); } #[test] fn first_event_detection_should_ignore_incomplete_frames() { - let body = include_bytes!("../../fixtures/responses/http_sse/incomplete_text_delta.sse"); + let body = include_bytes!("../../../fixtures/responses/http_sse/incomplete_text_delta.sse"); assert!(!response_body_has_first_event(body)); } #[test] fn first_event_detection_should_match_crlf_frames() { - let body = include_bytes!("../../fixtures/responses/http_sse/text_delta_crlf.sse"); + let body = include_bytes!("../../../fixtures/responses/http_sse/text_delta_crlf.sse"); assert!(response_body_has_first_event(body)); } diff --git a/backend/tests/upstream/protocol/usage_rate_limits.rs b/backend/tests/upstream/openai/protocol/usage_rate_limits.rs similarity index 96% rename from backend/tests/upstream/protocol/usage_rate_limits.rs rename to backend/tests/upstream/openai/protocol/usage_rate_limits.rs index 9ff396030..b477392a4 100644 --- a/backend/tests/upstream/protocol/usage_rate_limits.rs +++ b/backend/tests/upstream/openai/protocol/usage_rate_limits.rs @@ -120,7 +120,7 @@ fn extract_usage_should_read_openai_usage_shape() { #[test] fn extract_sse_usage_should_prefer_completed_response_usage() { - let body = include_str!("../../fixtures/responses/http_sse/created_completed_usage.sse"); + let body = include_str!("../../../fixtures/responses/http_sse/created_completed_usage.sse"); let usage = extract_sse_usage(body) .expect("usage extraction should succeed") @@ -142,7 +142,7 @@ fn extract_sse_usage_should_prefer_completed_response_usage() { #[test] fn extract_sse_usage_should_read_completed_image_generation_tokens() { - let body = include_str!("../../fixtures/responses/http_sse/completed_image_tool_usage.sse"); + let body = include_str!("../../../fixtures/responses/http_sse/completed_image_tool_usage.sse"); let usage = extract_sse_usage(body) .expect("usage extraction should succeed") @@ -368,9 +368,10 @@ fn rate_limit_quota_should_block_when_window_is_exhausted_even_if_flag_is_false( async fn refresh_scheduler_should_refresh_before_expiry_and_preserve_refresh_token() { use async_trait::async_trait; use chrono::{Duration as ChronoDuration, Utc}; - use codex_proxy_rs::upstream::accounts::model::AccountStatus; - use codex_proxy_rs::upstream::accounts::token_refresh::{ - RefreshFailure, RefreshPolicy, RefreshScheduler, TokenPair, TokenRefresher, + use codex_proxy_rs::accounts::account::AccountStatus; + use codex_proxy_rs::{ + accounts::refresh::{RefreshPolicy, RefreshScheduler}, + upstream::openai::token_client::{RefreshFailure, TokenPair, TokenRefresher}, }; #[derive(Clone)] diff --git a/backend/tests/upstream/token_client.rs b/backend/tests/upstream/openai/token_client.rs similarity index 96% rename from backend/tests/upstream/token_client.rs rename to backend/tests/upstream/openai/token_client.rs index 4026c4aa0..2659e7c00 100644 --- a/backend/tests/upstream/token_client.rs +++ b/backend/tests/upstream/openai/token_client.rs @@ -3,9 +3,8 @@ use wiremock::{ Mock, MockServer, ResponseTemplate, }; -use codex_proxy_rs::{ - upstream::accounts::token_refresh::{RefreshFailure, TokenRefresher}, - upstream::token_client::{OpenAiTokenClient, TokenClientConfig}, +use codex_proxy_rs::upstream::openai::token_client::{ + OpenAiTokenClient, RefreshFailure, TokenClientConfig, TokenRefresher, }; #[tokio::test] diff --git a/backend/tests/upstream/transport/client.rs b/backend/tests/upstream/openai/transport/client.rs similarity index 95% rename from backend/tests/upstream/transport/client.rs rename to backend/tests/upstream/openai/transport/client.rs index f3001ce4f..9c24b690c 100644 --- a/backend/tests/upstream/transport/client.rs +++ b/backend/tests/upstream/openai/transport/client.rs @@ -4,8 +4,8 @@ use std::sync::{ }; use super::{accept_codex_test_websocket, completed_websocket_response, request_context}; -use codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest; -use codex_proxy_rs::upstream::transport::{CodexBackendClient, CodexWebSocketPool}; +use codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest; +use codex_proxy_rs::upstream::openai::transport::{CodexBackendClient, CodexWebSocketPool}; use futures::{SinkExt, StreamExt}; use tokio::net::TcpListener; use tokio_tungstenite::tungstenite::Message; diff --git a/backend/tests/upstream/transport/diagnostics.rs b/backend/tests/upstream/openai/transport/diagnostics.rs similarity index 92% rename from backend/tests/upstream/transport/diagnostics.rs rename to backend/tests/upstream/openai/transport/diagnostics.rs index 8db2265fd..b93691aad 100644 --- a/backend/tests/upstream/transport/diagnostics.rs +++ b/backend/tests/upstream/openai/transport/diagnostics.rs @@ -1,4 +1,4 @@ -use codex_proxy_rs::upstream::transport::CodexUpstreamDiagnostics; +use codex_proxy_rs::upstream::openai::transport::CodexUpstreamDiagnostics; use reqwest::header::{HeaderMap, HeaderValue}; #[test] diff --git a/backend/tests/upstream/transport/headers.rs b/backend/tests/upstream/openai/transport/headers.rs similarity index 94% rename from backend/tests/upstream/transport/headers.rs rename to backend/tests/upstream/openai/transport/headers.rs index fe032d7c1..406989b14 100644 --- a/backend/tests/upstream/transport/headers.rs +++ b/backend/tests/upstream/openai/transport/headers.rs @@ -1,7 +1,7 @@ use super::*; -use codex_proxy_rs::upstream::fingerprint::RuntimeFingerprint; -use codex_proxy_rs::upstream::transport::build_ordered_codex_headers; -use codex_proxy_rs::upstream::transport::websocket::CodexWebSocketConnection; +use codex_proxy_rs::upstream::openai::fingerprint::RuntimeFingerprint; +use codex_proxy_rs::upstream::openai::transport::build_ordered_codex_headers; +use codex_proxy_rs::upstream::openai::transport::websocket::CodexWebSocketConnection; use serde_json::Value; #[test] @@ -188,7 +188,7 @@ async fn websocket_execute_response_create_request_should_capture_handshake_head websocket.close(None).await.unwrap(); }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -259,7 +259,7 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an payload }); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -272,7 +272,7 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an "ignored_non_string": 42 }))); let pool = Arc::new( - codex_proxy_rs::upstream::transport::websocket_pool::CodexWebSocketPool::new( + codex_proxy_rs::upstream::openai::transport::websocket_pool::CodexWebSocketPool::new( 8, std::time::Duration::from_mins(1), ), @@ -384,7 +384,7 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an #[tokio::test] async fn codex_backend_client_should_send_desktop_headers_and_capture_response_metadata() { let server = wiremock::MockServer::start().await; - let sse_body = include_str!("../../fixtures/responses/http_sse/completed_usage_basic.sse"); + let sse_body = include_str!("../../../fixtures/responses/http_sse/completed_usage_basic.sse"); wiremock::Mock::given(wiremock::matchers::method("POST")) .and(wiremock::matchers::path("/codex/responses")) .and(wiremock::matchers::header( @@ -424,7 +424,7 @@ async fn codex_backend_client_should_send_desktop_headers_and_capture_response_m crate::support::fingerprint::runtime_test_fingerprint(), ); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "", Vec::new(), @@ -455,15 +455,17 @@ async fn codex_backend_client_should_send_desktop_headers_and_capture_response_m assert_eq!( response.usage, - Some(codex_proxy_rs::upstream::protocol::events::TokenUsage { - input_tokens: 2, - output_tokens: 3, - cached_tokens: 1, - reasoning_tokens: 0, - image_input_tokens: 0, - image_output_tokens: 0, - total_tokens: 5, - }) + Some( + codex_proxy_rs::upstream::openai::protocol::events::TokenUsage { + input_tokens: 2, + output_tokens: 3, + cached_tokens: 1, + reasoning_tokens: 0, + image_input_tokens: 0, + image_output_tokens: 0, + total_tokens: 5, + } + ) ); assert_eq!(response.turn_state.as_deref(), Some("turn_2")); assert_eq!( @@ -475,7 +477,7 @@ async fn codex_backend_client_should_send_desktop_headers_and_capture_response_m #[tokio::test] async fn codex_backend_client_should_use_latest_runtime_fingerprint_for_new_requests() { let server = wiremock::MockServer::start().await; - let sse_body = include_str!("../../fixtures/responses/http_sse/completed_usage_basic.sse"); + let sse_body = include_str!("../../../fixtures/responses/http_sse/completed_usage_basic.sse"); wiremock::Mock::given(wiremock::matchers::method("POST")) .and(wiremock::matchers::path("/codex/responses")) .respond_with( @@ -495,7 +497,7 @@ async fn codex_backend_client_should_use_latest_runtime_fingerprint_for_new_requ runtime_fingerprint.clone(), ); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "", Vec::new(), @@ -642,7 +644,9 @@ async fn codex_backend_client_usage_should_use_wham_usage_headers() { #[tokio::test] async fn codex_backend_client_models_should_use_original_auxiliary_headers() { - use codex_proxy_rs::upstream::transport::{CodexModelCatalogClient, CodexModelCatalogRequest}; + use codex_proxy_rs::upstream::openai::transport::{ + CodexModelCatalogClient, CodexModelCatalogRequest, + }; let server = wiremock::MockServer::start().await; wiremock::Mock::given(wiremock::matchers::method("GET")) @@ -709,7 +713,7 @@ async fn codex_backend_client_should_send_http_sse_headers_in_fingerprint_order( }); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "", Vec::new(), @@ -804,7 +808,7 @@ async fn codex_backend_client_should_send_compact_headers_in_fingerprint_order() client .create_compact_response( - &codex_proxy_rs::upstream::protocol::responses::CodexCompactRequest { + &codex_proxy_rs::upstream::openai::protocol::responses::CodexCompactRequest { body: json!({ "model": "gpt-5.5", "input": [], @@ -815,6 +819,7 @@ async fn codex_backend_client_should_send_compact_headers_in_fingerprint_order() .unwrap(), client_ip: None, client_user_agent: None, + client_api_key_id: None, }, CodexRequestContext { access_token: "access-token", diff --git a/backend/tests/upstream/transport/http_client.rs b/backend/tests/upstream/openai/transport/http_client.rs similarity index 87% rename from backend/tests/upstream/transport/http_client.rs rename to backend/tests/upstream/openai/transport/http_client.rs index 40bceba03..d62baf5ec 100644 --- a/backend/tests/upstream/transport/http_client.rs +++ b/backend/tests/upstream/openai/transport/http_client.rs @@ -3,14 +3,14 @@ use super::*; #[test] fn endpoints_should_join_backend_paths() { assert_eq!( - codex_proxy_rs::upstream::transport::endpoint_url( + codex_proxy_rs::upstream::openai::transport::endpoint_url( "https://api.example.com/", "/codex/responses" ), "https://api.example.com/codex/responses" ); assert_eq!( - codex_proxy_rs::upstream::transport::endpoint_request_path( + codex_proxy_rs::upstream::openai::transport::endpoint_request_path( "https://api.example.com/backend-api", "/codex/usage" ), @@ -29,17 +29,17 @@ fn custom_ca_should_report_environment_cache_key_consistently() { "unset" => None, "ssl_cert_file" => Some(format!( "{}={SSL_CERT_PATH}", - codex_proxy_rs::upstream::transport::tls::SSL_CERT_FILE_ENV + codex_proxy_rs::upstream::openai::transport::tls::SSL_CERT_FILE_ENV )), "codex_ca_priority" => Some(format!( "{}={CODEX_CA_PATH}", - codex_proxy_rs::upstream::transport::tls::CODEX_CA_CERT_ENV + codex_proxy_rs::upstream::openai::transport::tls::CODEX_CA_CERT_ENV )), _ => panic!("unknown custom CA cache key test case: {case}"), }; assert_eq!( - codex_proxy_rs::upstream::transport::tls::custom_ca_env_cache_key(), + codex_proxy_rs::upstream::openai::transport::tls::custom_ca_env_cache_key(), expected ); return; @@ -52,18 +52,18 @@ fn custom_ca_should_report_environment_cache_key_consistently() { "ssl_cert_file", None, Some(( - codex_proxy_rs::upstream::transport::tls::SSL_CERT_FILE_ENV, + codex_proxy_rs::upstream::openai::transport::tls::SSL_CERT_FILE_ENV, SSL_CERT_PATH, )), ), ( "codex_ca_priority", Some(( - codex_proxy_rs::upstream::transport::tls::CODEX_CA_CERT_ENV, + codex_proxy_rs::upstream::openai::transport::tls::CODEX_CA_CERT_ENV, CODEX_CA_PATH, )), Some(( - codex_proxy_rs::upstream::transport::tls::SSL_CERT_FILE_ENV, + codex_proxy_rs::upstream::openai::transport::tls::SSL_CERT_FILE_ENV, SSL_CERT_PATH, )), ), @@ -76,8 +76,8 @@ fn custom_ca_should_report_environment_cache_key_consistently() { .arg("custom_ca_should_report_environment_cache_key_consistently") .arg("--nocapture") .env(CASE_ENV, case) - .env_remove(codex_proxy_rs::upstream::transport::tls::CODEX_CA_CERT_ENV) - .env_remove(codex_proxy_rs::upstream::transport::tls::SSL_CERT_FILE_ENV); + .env_remove(codex_proxy_rs::upstream::openai::transport::tls::CODEX_CA_CERT_ENV) + .env_remove(codex_proxy_rs::upstream::openai::transport::tls::SSL_CERT_FILE_ENV); if let Some((key, value)) = codex_ca { command.env(key, value); } @@ -110,7 +110,7 @@ async fn codex_backend_client_should_cap_non_success_error_body_at_one_mib() { crate::support::fingerprint::runtime_test_fingerprint(), ); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "", Vec::new(), @@ -164,7 +164,7 @@ async fn codex_backend_client_should_parse_retry_after_from_rate_limit_error_bod crate::support::fingerprint::runtime_test_fingerprint(), ); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "", Vec::new(), diff --git a/backend/tests/upstream/transport/mod.rs b/backend/tests/upstream/openai/transport/mod.rs similarity index 94% rename from backend/tests/upstream/transport/mod.rs rename to backend/tests/upstream/openai/transport/mod.rs index 3662207c8..5fdff49f1 100644 --- a/backend/tests/upstream/transport/mod.rs +++ b/backend/tests/upstream/openai/transport/mod.rs @@ -7,7 +7,7 @@ use std::{ time::Duration, }; -use codex_proxy_rs::upstream::{ +use codex_proxy_rs::upstream::openai::{ transport::websocket::{ execute_response_create_request, responses_websocket_endpoint, CodexWebSocketConnection, CodexWebSocketExchangeError, CodexWebSocketRequest, @@ -145,7 +145,7 @@ pub(crate) async fn write_empty_http_response(stream: &mut TcpStream) { } pub(crate) async fn write_completed_sse_response(stream: &mut TcpStream) { - let body = include_str!("../../fixtures/responses/http_sse/completed_order_usage.sse"); + let body = include_str!("../../../fixtures/responses/http_sse/completed_order_usage.sse"); stream .write_all( format!( @@ -181,7 +181,7 @@ pub(crate) fn request_context<'a>( pub(crate) fn prepared_websocket_request(base_url: &str) -> CodexWebSocketRequest { let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -200,9 +200,9 @@ pub(crate) fn prepared_websocket_request(base_url: &str) -> CodexWebSocketReques pub(crate) fn pooled_websocket_request( conversation_id: &str, -) -> codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest { +) -> codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest { let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), diff --git a/backend/tests/upstream/transport/websocket.rs b/backend/tests/upstream/openai/transport/websocket.rs similarity index 96% rename from backend/tests/upstream/transport/websocket.rs rename to backend/tests/upstream/openai/transport/websocket.rs index 94965e658..e1b6d03fb 100644 --- a/backend/tests/upstream/transport/websocket.rs +++ b/backend/tests/upstream/openai/transport/websocket.rs @@ -1,5 +1,5 @@ use super::*; -use codex_proxy_rs::upstream::{ +use codex_proxy_rs::upstream::openai::{ protocol::websocket::{OpeningAuditSnapshot, WebSocketAuditArtifact}, transport::websocket::write_websocket_audit_artifact_for_dir, }; @@ -134,7 +134,7 @@ async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_ ) .with_websocket_pool(pool); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -159,7 +159,7 @@ async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_ #[test] fn websocket_connection_should_prepare_response_create_payload_text() { let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", vec![json!({ @@ -195,7 +195,7 @@ fn websocket_connection_should_prepare_response_create_payload_text() { #[test] fn websocket_connection_should_prepare_capture_payload_with_canonical_field_order() { let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "private capture instructions", vec![json!({ @@ -269,7 +269,7 @@ async fn websocket_execute_response_create_request_should_collect_completed_sse( payload }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -324,7 +324,7 @@ async fn websocket_execute_response_create_request_should_surface_response_faile websocket.close(None).await.unwrap(); }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -578,7 +578,7 @@ async fn websocket_execute_response_create_request_should_forward_typed_events_w websocket.close(None).await.unwrap(); }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -781,7 +781,7 @@ async fn websocket_execute_response_create_request_should_surface_incomplete_res websocket.close(None).await.unwrap(); }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -838,7 +838,7 @@ async fn websocket_execute_response_create_request_should_reject_invalid_complet websocket.close(None).await.unwrap(); }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -887,7 +887,7 @@ async fn websocket_execute_response_create_request_should_forward_completed_with websocket.close(None).await.unwrap(); }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -940,7 +940,7 @@ async fn websocket_execute_response_create_request_should_pass_through_unclassif websocket.close(None).await.unwrap(); }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1040,7 +1040,7 @@ async fn websocket_execute_response_create_request_should_reply_to_server_ping_b websocket.close(None).await.unwrap(); }); let request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1081,7 +1081,7 @@ async fn codex_backend_client_stream_should_reject_binary_websocket_event() { websocket.close(None).await.unwrap(); }); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1214,7 +1214,7 @@ async fn codex_backend_client_stream_should_error_when_websocket_closes_before_t websocket.close(None).await.unwrap(); }); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1284,7 +1284,7 @@ async fn codex_backend_client_stream_should_wait_for_terminal_after_active_webso websocket.close(None).await.unwrap(); }); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1348,7 +1348,7 @@ async fn codex_backend_client_stream_should_timeout_when_active_websocket_stalls futures::future::pending::<()>().await; }); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1444,7 +1444,7 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_ payload }); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), diff --git a/backend/tests/upstream/transport/websocket_pool.rs b/backend/tests/upstream/openai/transport/websocket_pool.rs similarity index 98% rename from backend/tests/upstream/transport/websocket_pool.rs rename to backend/tests/upstream/openai/transport/websocket_pool.rs index a700765ff..d522014f7 100644 --- a/backend/tests/upstream/transport/websocket_pool.rs +++ b/backend/tests/upstream/openai/transport/websocket_pool.rs @@ -43,7 +43,7 @@ async fn codex_backend_client_should_reuse_pooled_websocket_for_same_account_and ) .with_websocket_pool(Arc::clone(&pool)); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -121,7 +121,7 @@ async fn codex_backend_client_should_open_fresh_socket_when_idle_pooled_websocke ) .with_websocket_pool(Arc::clone(&pool)); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -458,7 +458,7 @@ async fn codex_backend_client_should_not_reuse_pooled_websocket_across_local_acc ) .with_websocket_pool(Arc::clone(&pool)); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -982,7 +982,7 @@ async fn codex_backend_client_should_keep_idle_pooled_websocket_alive_across_rep ) .with_websocket_pool(Arc::clone(&pool)); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1088,7 +1088,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_when_account_is ) .with_websocket_pool(Arc::clone(&pool)); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1192,7 +1192,7 @@ async fn codex_backend_client_should_stop_reusing_pooled_websockets_after_shutdo ) .with_websocket_pool(Arc::clone(&pool)); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1296,7 +1296,7 @@ async fn codex_backend_client_should_close_idle_pooled_websocket_after_liveness_ ) .with_websocket_pool(Arc::clone(&pool)); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), @@ -1391,7 +1391,7 @@ async fn codex_backend_client_should_discard_pooled_websocket_after_upstream_err ) .with_websocket_pool(pool); let mut request = - codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + codex_proxy_rs::upstream::openai::protocol::responses::CodexResponsesRequest::new_http_sse( "gpt-5.5", "be brief", Vec::new(), diff --git a/deploy/config.example.yaml b/deploy/config.example.yaml index 70b69b528..aacbb4672 100644 --- a/deploy/config.example.yaml +++ b/deploy/config.example.yaml @@ -9,8 +9,12 @@ api: base_url: 'https://chatgpt.com/backend-api' database: - # Docker 部署默认写入 /app/data，由 Compose 挂载到仓库 .runtime/data。 - url: 'sqlite:///app/data/codex-proxy-rs.sqlite' + # PostgreSQL 持久层。 + url: 'postgres://codex_proxy:codex_proxy@postgres:5432/codex_proxy' + +redis: + # Redis 仅承载可丢弃的运行态与缓存。 + url: 'redis://redis:6379/' quota: # 账号额度刷新间隔，单位分钟。 @@ -80,8 +84,6 @@ fingerprint: admin: # 管理端登录会话有效期，单位分钟。 session_ttl_minutes: 1440 - # 过期会话清理间隔，单位秒。 - session_cleanup_interval_secs: 3600 # 首次启动创建的默认管理员账号。default_password 必须改成至少 12 位的非默认密码。 default_username: admin@cpr.local default_password: '' diff --git a/deploy/docker-compose.yml b/deploy/docker-compose.yml index ee19fde8e..9b793ed5e 100644 --- a/deploy/docker-compose.yml +++ b/deploy/docker-compose.yml @@ -1,6 +1,41 @@ name: codex-proxy-rs services: + postgres: + image: postgres:18-alpine + restart: unless-stopped + environment: + POSTGRES_DB: codex_proxy + POSTGRES_USER: codex_proxy + POSTGRES_PASSWORD: ${CPR_POSTGRES_PASSWORD:-codex_proxy} + volumes: + - postgres-data:/var/lib/postgresql + ports: + - "127.0.0.1:5432:5432" + healthcheck: + test: ["CMD-SHELL", "pg_isready -U codex_proxy -d codex_proxy"] + interval: 5s + timeout: 5s + retries: 10 + security_opt: + - no-new-privileges:true + + redis: + image: redis:8-alpine + restart: unless-stopped + command: ["redis-server", "--appendonly", "yes"] + volumes: + - redis-data:/data + ports: + - "127.0.0.1:6379:6379" + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 5s + timeout: 5s + retries: 10 + security_opt: + - no-new-privileges:true + codex-proxy-rs: image: ${CPR_IMAGE:-ghcr.io/zyycn/codex-proxy-rs:latest} build: @@ -17,6 +52,11 @@ services: user: "10001:10001" init: true stop_grace_period: 30s + depends_on: + postgres: + condition: service_healthy + redis: + condition: service_healthy pids_limit: 256 cpus: "2.0" mem_limit: 1g @@ -50,3 +90,7 @@ services: CPR_UPDATE_LOCK_FILE: /app/data/update.lock CPR_WEB_DIST_DIR: /app/web/dist CPR_ENABLE_SELF_RESTART: "true" + +volumes: + postgres-data: + redis-data: diff --git a/docs/architecture.md b/docs/architecture.md index 8dbba0960..607b0b900 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -116,7 +116,7 @@ backend/src/ │ ├── mod.rs store.rs query.rs ├── keys/ # L1 客户端 key 域（原 admin/keys/service.rs 拆分） │ ├── mod.rs types.rs store.rs service.rs manage.rs -│ │ # 鉴权 = SHA-256 → PG 点查（database.md §4.3），无内存明文缓存文件 +│ │ # 鉴权 = 完整 key → PG unique 点查（database.md §4.3），无进程内鉴权缓存 ├── auth/ # L1 管理端认证域（原 admin/auth/service.rs） │ ├── mod.rs types.rs store.rs service.rs # store = PG admin_users + Redis 会话键（§4B.1） ├── settings/ # L1 运行时设置域（原 config/settings.rs） @@ -155,7 +155,8 @@ backend/src/ │ ├── errors.rs # 原 dispatch/errors.rs + responses/errors.rs 合并 │ ├── recording.rs # 原 responses/event_recording.rs，调 telemetry::recorder │ ├── affinity/ # 原 session_affinity.rs（683 行三合一）+ responses/affinity.rs -│ │ ├── mod.rs types.rs store.rs service.rs resolve.rs # store = Redis 键 + 二级索引（§4B.3），无内存映射副本 +│ │ ├── mod.rs types.rs identity.rs store.rs service.rs resolve.rs +│ │ │ # store = Redis 键 + 二级索引（§4B.3），无内存映射副本 │ ├── recovery/ │ │ ├── mod.rs auth.rs cloudflare.rs exhaustion.rs │ │ ├── implicit_resume.rs reasoning_replay.rs @@ -233,7 +234,7 @@ backend/src/ | `Admin*` | **仅** `api/admin` 下的 DTO/提取器 | `AdminAccountPayload` | | 无前缀 | 领域服务与实体 | `Recorder`、`AccountScheduler`、`SettingsService` | -由此更名（现状 → 终态）：`Sqlite*Store` 全系 → `Pg*Store` 或 `Redis*Store`（按 database.md 的归属，引擎切换 PR0 完成）；`AdminUsageRecordService` → `telemetry::usage::UsageQueryService`（读）+ `telemetry::Recorder`（写）；`AdminOpsErrorLogService` → `telemetry::ops` 同拆；`AdminUsageService`（实际管 account_usage）→ `telemetry::account_usage::AccountUsageQueryService`；`AdminSessionService` → `auth::SessionService`；`AdminClientKeyService` → `keys::KeyManageService`；`ClientKeyService` → `keys::KeyVerifier`（SHA-256 → PG 点查，`RuntimeClientKeyStore` 内存明文表随哈希化消亡）。**类型名与文件名对不上的历史包袱全部在搬家时清算，不留别名 re-export。** +由此更名（现状 → 终态）：`Sqlite*Store` 全系 → `Pg*Store` 或 `Redis*Store`（按 database.md 的归属，引擎切换 PR0 完成）；`AdminUsageRecordService` → `telemetry::usage::UsageQueryService`（读）+ `telemetry::Recorder`（写）；`AdminOpsErrorLogService` → `telemetry::ops` 同拆；`AdminUsageService`（实际管 account_usage）→ `telemetry::account_usage::AccountUsageQueryService`；`AdminSessionService` → `auth::SessionService`；`AdminClientKeyService` → `keys::KeyManageService`；`ClientKeyService` → `keys::KeyVerifier`（完整 key → PG unique 点查，旧 `RuntimeClientKeyStore` 内存鉴权表消亡）。**类型名与文件名对不上的历史包袱全部在搬家时清算，不留别名 re-export。** 禁止无意义别名导入：`use std::sync::Arc as StdArc` 这类一律 `use std::sync::Arc`。 @@ -255,7 +256,7 @@ axum/tower 类型只允许出现在 `api` 与 `bootstrap`。领域函数签名 ### 5.4 文件行数预算 -- 软上限 **400 行**（超过即在 review 中说明理由）；硬上限 **800 行**（超过必须拆，CI 警告，§9）。 +- 软上限 **400 行**（超过即在 review 中说明理由）；硬上限 **800 行**（超过必须拆，CI 直接阻断，§9）。 - 现状超硬线清单（11 个，全部有既定拆法）： | 文件（现路径） | 行数 | 拆法 | @@ -364,7 +365,7 @@ axum/tower 类型只允许出现在 `api` 与 `bootstrap`。领域函数签名 | `admin/monitoring/dashboard.rs` | `telemetry/dashboard.rs` + `api/admin/dashboard_routes.rs` † | | `admin/monitoring/diagnostics.rs` | `telemetry/diagnostics.rs` | | `admin/monitoring/usage_record_routes.rs` | `api/admin/usage_routes.rs` | -| `admin/keys/service.rs` | `keys/{types,store,service,manage}.rs` †（明文内存缓存消亡，database.md §4.3） | +| `admin/keys/service.rs` | `keys/{types,store,service,manage}.rs` †（进程内鉴权缓存消亡，database.md §4.3） | | `admin/keys/routes.rs` | `api/admin/keys_routes.rs` | | `admin/auth/service.rs` | `auth/{types,store,service}.rs` †（会话改 Redis，database.md §4B.1） | | `admin/auth/session.rs` | `api/admin/session.rs` | @@ -392,19 +393,19 @@ axum/tower 类型只允许出现在 `api` 与 `bootstrap`。领域函数签名 --- -## 8. 执行计划（PR0 + 六个顺序搬家 PR，每个独立编译绿、测试绿） +## 8. 执行记录（PR0 + 六个顺序搬家阶段） -每个 PR 的固定验收：`cargo fmt --check`、`cargo clippy -- -D warnings`、`cargo test`、`backend/scripts/check_architecture.sh`（§9，随 PR 逐步启用对应规则）、`backend/tests` 镜像目录同步搬移。**PR0 是唯一的行为变更 PR**（存储引擎与 schema 语义一次到位）；PR1 起全部为零行为搬家（更名与移动，不改逻辑）。 +以下阶段已于 **2026-07-10** 全部完成。固定验收：`cargo fmt --check`、`cargo clippy -- -D warnings`、`cargo test`、`backend/scripts/check_architecture.sh`、`backend/tests` 镜像目录同步搬移。**PR0 是唯一的迁移行为变更阶段**（存储引擎与 schema 语义一次到位）；PR1 起为搬家、更名与拆分。 | PR | 内容 | 关键动作 | | --- | --- | --- | -| **PR0** | **存储终态切换**（database.md §7 全部 12 条，在现有目录结构内完成） | Cargo：sqlx 加 postgres、新增 redis；`infra/database.rs` 改 PG + 0001 终态基线、新增 `infra/redis.rs`；全部 store 换 PG 方言 / 运行态四键空间改 Redis；成功/失败事实拆表 + provider 维度 + 两条写入事务；key/session/admin-key 哈希化 + `client_api_key_id` 归因贯通；Dashboard 口径换源；trim 周期任务化；`import-sqlite` / `rebuild-buckets` 子命令；tests support 换 PG+Redis;deploy compose 加双服务。验收对齐 database.md §6.4 | +| **PR0** | **存储终态切换**（database.md §7 全部 12 条，在现有目录结构内完成） | Cargo：sqlx 加 postgres、新增 redis；`infra/database.rs` 改 PG + 0001 终态基线、新增 `infra/redis.rs`；全部 store 换 PG 方言 / 运行态四键空间改 Redis；成功/失败事实拆表 + provider 维度 + 两条写入事务；client key 可取回、session/admin-key 哈希化 + `client_api_key_id` 归因贯通；Dashboard 口径换源；trim 周期任务化；`import-sqlite` / `rebuild-buckets` 子命令；tests support 换 PG+Redis;deploy compose 加双服务。验收对齐 database.md §6.4 | | **PR1** | `telemetry` 域成立 | monitoring 全部迁出 admin；`usage_events.rs` → `recorder.rs`；Admin* 遥测类型按 §4.3 更名；dispatch/api 改依赖 telemetry；启用"禁止 `crate::admin::monitoring`"检查 | | **PR2** | `keys` / `auth` / `settings` / `update` 域成立 | 四个域从 admin/config 拆出；settings 改 watch 广播（§2.3）；config 只剩 schema+loader | | **PR3** | `accounts` / `models` 域成立 | `upstream/accounts`、`upstream/scheduler`、`admin/accounts/service` 三处合并为 `accounts`；`upstream` 缩为纯客户端并落成 `openai/` 子目录；动名词文件更名；Repository→Store | | **PR4** | `dispatch` 域成立 | proxy/dispatch 整体迁出并按目标树重组；affinity store 独立成文件；`proxy` 只剩 api/client 的料 | | **PR5** | `api` / `bootstrap` 壳成立 | http + web + proxy 入口 + admin routes 合并为 `api`；runtime + config 装配部分改名 `bootstrap`；删除 `crate::{admin,proxy,runtime,web,http,config}` 六个旧根模块；启用全部依赖检查 | -| **PR6** | 行数预算清账 | §5.4 清单内 11 个文件拆到硬线以下；启用行数预算 CI 警告 | +| **PR6** | 行数预算清账 | §5.4 清单内 11 个文件拆到硬线以下；启用行数预算 CI 阻断 | 顺序依据：存储切换触及所有 store、无法按域切片，必须整体先行（PR0）——先切引擎再搬家，每个 store 文件只被"重写"一次、"移动"一次，不叠加；PR1–4 自底向上消除旧模块的存在必要；PR5 收壳；PR6 纯内部拆分无跨域影响。原"PR1 与数据库改造互锁"的约定由 PR0 替代，PR2 中的哈希化条目亦已并入 PR0。 @@ -437,12 +438,21 @@ deny "axum outside api" '\baxum(::|\b)' \ src/infra src/upstream src/telemetry src/keys src/auth src/settings src/update \ src/accounts src/models src/dispatch # SQL 与 Redis 命令只在 store / query -deny "sqlx outside store/query" 'sqlx::query' \ +deny "sqlx outside store/query" '\bsqlx(::|\b)' \ $(find src/api -name '*.rs') deny "redis outside store" '\bredis::' \ $(find src/api -name '*.rs') # 旧模块根不得复活 deny "legacy module path" 'crate::(admin|proxy|runtime|web|http|config)::' src +# 领域类型与命名纪律 +deny "Admin-prefixed domain type" '\bAdmin[A-Z][A-Za-z0-9]*\b' \ + src/infra src/upstream src/telemetry src/keys src/auth src/settings src/update \ + src/accounts src/models src/dispatch +deny "Repository naming" '\b[A-Za-z0-9]*Repository\b' src +deny "persistence repository naming" '\b[A-Za-z0-9_]*repository[A-Za-z0-9_]*\b' \ + src/accounts src/auth src/dispatch src/keys src/models src/settings src/telemetry \ + src/upstream/openai/fingerprint +deny "Arc alias" '\bArc[[:space:]]+as[[:space:]]+' src # 禁用文件名 banned=$(find src -name 'importing.rs' -o -name 'exporting.rs' -o -name 'testing.rs' \ -o -name 'util.rs' -o -name 'helper.rs' -o -name 'common.rs' -o -name 'misc.rs') @@ -452,15 +462,15 @@ while read -r f; do lines=$(wc -l < "$f") [ "$lines" -gt 50 ] && { echo "ARCH VIOLATION: $f has $lines lines (mod.rs <= 50)" >&2; fail=1; } done < <(find src -name mod.rs) -# 行数硬上限（警告期：仅报告；PR6 后改 fail） -find src -name '*.rs' ! -name mod.rs | while read -r f; do +# 行数硬上限（PR6 已完成，违规即失败） +while read -r f; do lines=$(wc -l < "$f") - [ "$lines" -gt 800 ] && echo "WARN: $f has $lines lines (> 800)" -done + [ "$lines" -gt 800 ] && { echo "ARCH VIOLATION: $f has $lines lines (> 800)" >&2; fail=1; } +done < <(find src -name '*.rs' ! -name mod.rs) exit $fail ``` -规则随 PR 分阶段启用（PR1 起启用 legacy-path 的 monitoring 子集，PR5 后全量）；`deny` 的路径参数在过渡期只指向已完成搬家的目录。 +全部规则已启用；后续新增代码一旦恢复旧路径、越层依赖、违规命名或突破硬行数上限，CI 直接失败。 --- @@ -469,7 +479,7 @@ exit $fail - `backend/tests/` 镜像 `src/`，不在 `src/` 写测试模块。测试依赖本机（或 CI 服务容器）的 PostgreSQL 与 Redis：`CPR_TEST_DATABASE_URL` 下每测试建独立数据库、`CPR_TEST_REDIS_URL` 下每测试随机键前缀（database.md §7）。 - `docs/` 只保留权威文档：`architecture.md`（本文）、`database.md`。行为细节（调度策略、WS 池、更新流程等）以模块文档注释为准（`scheduler/mod.rs` 是范本），不再在 docs 里维护第二份会腐化的副本。 - `deploy/` 随 PR0 扩容（compose 增 postgres/redis 服务与健康检查依赖）；`release/`、`skills/` 现状合格，不动。 -- 前端联动（全部随 PR0）：keys 列表不再返回完整 key（只余 `prefix`）；admin API key 状态不再有脱敏值（哈希不可逆，只余"已启用"）；成功记录响应不再有 `level` 字段；错误明细/筛选改走 `/api/admin/ops/errors`；Dashboard "总请求/总 token" 卡片改为 account_usage 来源并标注"成功累计"。 +- 前端联动（全部随 PR0）：keys 列表继续返回完整 key，并长期提供复制与 CCSwitch 导入；admin API key 状态不再有脱敏值（哈希不可逆，只余"已启用"）；成功记录响应不再有 `level` 字段；错误明细/筛选改走 `/api/admin/ops/errors`；Dashboard "总请求/总 token" 卡片改为 account_usage 来源并标注"成功累计"。 --- diff --git a/docs/database.md b/docs/database.md index 63bc2b99b..5c3ec8e3d 100644 --- a/docs/database.md +++ b/docs/database.md @@ -48,7 +48,7 @@ | JSON | `*_json`,**jsonb** | `metadata_json` | | 布尔 | 裸形容词/过去式,**boolean** | `enabled` | | 外键/逻辑引用列 | `<被引用实体单数>_id` | `client_api_key_id` | -| 凭据哈希 | `*_hash`,明文永不落库(§2.6) | `key_hash`、`token_hash` | +| 凭据哈希 | `*_hash`,用于不可逆凭据(§2.6) | `password_hash`、`token_hash` | | 主键 | `id`,除非自然复合键或哈希键 | — | | 普通/唯一索引 | `idx_<表>_<列摘要>` / `ux_<表>_<列摘要>` | `ux_accounts_chatgpt_identity` | | 双端状态码 | 视角前缀 | `status_code` / `client_status_code` / `upstream_status_code` | @@ -76,7 +76,7 @@ ### 2.4 ID -主键一律 text 随机 ID(`infra/identity.rs` 生成):无中心分配、写入前可生成、跨表无碰撞。例外:`schema_migrations.version`(天然有序整数)、`runtime_settings.id`(单例锚点)、凭据存储以哈希为键(§2.6)。**不用 serial/identity/uuid 列类型**——ID 是业务字符串,不是数据库生成物。 +主键一律 text 随机 ID(`infra/identity.rs` 生成):无中心分配、写入前可生成、跨表无碰撞。例外:`schema_migrations.version`(天然有序整数)、`runtime_settings.id`(单例锚点)。**不用 serial/identity/uuid 列类型**——ID 是业务字符串,不是数据库生成物。 ### 2.5 NULL 与 sentinel @@ -84,14 +84,14 @@ - 聚合维度列不可空时,未知值用 `__unknown__`;预聚合汇总行(预留)用 `__all__`。 - 唯一例外:`ux_accounts_chatgpt_identity` 的索引表达式 `coalesce(chatgpt_user_id, '')`——空串只存在于索引表达式,不是存储值,不构成先例。 -### 2.6 Secret 纪律(无例外) +### 2.6 Secret 纪律与可取回边界 -凡是"出示即获得权限"的值都是凭据,**明文一律不落库(PG 与 Redis 同等适用)**: +凡是"出示即获得权限"的值都是凭据。默认只保存不可逆摘要；但客户端 API key 与上游凭据有明确的可取回功能需求，必须可逆保存。所有可逆凭据均不得进入日志、遥测或非管理端响应，数据库备份与访问权限必须按 secret 级别保护。 | 凭据 | 存储 | 哈希类型 | 理由 | | --- | --- | --- | --- | | 管理员密码 | PG `admin_users.password_hash` | 慢哈希(argon2) | 低熵人类输入,必须抗字典攻击 | -| 客户端 API key | PG `client_api_keys.key_hash` | SHA-256 | 高熵随机串无字典攻击面,快哈希保住 unique 索引查找;明文只在创建响应出现一次,`prefix` 供事后定位 | +| 客户端 API key | PG `client_api_keys.key` 明文列,标记 secret | —(功能上必须可逆) | 管理端需要长期复制并导入 CCSwitch；鉴权走 `key` unique 索引点查，列表只对已认证管理员返回完整值 | | 管理端 session token | Redis `cpr:admin:session:` | SHA-256 | 键即哈希。Redis RDB/AOF 或内存转储泄露不再等于会话劫持 | | 管理 API key | PG `runtime_settings.admin_api_key_hash` | SHA-256 | 校验 = 哈希后 `ct_eq`;哈希不可脱敏展示,状态 API 只回答"是否已启用" | | 上游 access/refresh token、Cookie value | PG 明文列,标记 secret | —(功能上必须可逆) | 不进日志、不进列表 API;静态加密的触发条件见 §8 | @@ -220,7 +220,7 @@ create table client_api_keys ( id text primary key, name text not null, prefix text not null, - key_hash text not null unique, + key text not null unique, label text, enabled boolean not null default true, created_at timestamptz not null, @@ -234,12 +234,12 @@ create index idx_client_api_keys_created_id on client_api_keys(created_at desc, | --- | --- | | `id` | 稳定标识;管理操作与事实归因(`usage_records.client_api_key_id`)都引用 id,key 轮换不影响引用 | | `name` / `label` | `name` 必填主展示,`label` 可选注记 | -| `prefix` | 明文前几位(`sk_abc1…`),明文销毁后定位 key 的唯一手段 | -| `key_hash` | `SHA-256(key)` hex,unique 索引即鉴权查找路径。明文只在创建响应出现一次,**管理端列表 API 永不回显完整 key** | +| `prefix` | 明文前几位(`sk_abc1…`),用于列表快速定位与遮罩展示 | +| `key` | 完整客户端凭据,unique 索引即鉴权查找路径。只在已认证管理端列表/创建响应返回,供长期复制与导入 CCSwitch；严禁进入日志和遥测 | | `enabled` | 软开关;刻意不建索引(低基数布尔) | | `last_used_at` | 鉴权成功后延迟批量更新(1s 去抖),不在请求关键路径强一致 | -**鉴权路径**:`SHA-256(来键)` → `where key_hash = $1 and enabled` 单点查询,**删除进程内明文鉴权表**——单一代码路径、禁用即刻生效、无缓存同步逻辑;一次 unique 索引点查的代价可忽略。每 key 配额/限流不进本表(§8)。 +**鉴权路径**:`where key = $1 and enabled` unique 索引单点查询,**删除进程内明文鉴权表**——单一代码路径、禁用即刻生效、无缓存同步逻辑。管理端取回与客户端鉴权读取同一事实列,不存在 hash/明文双写。每 key 配额/限流不进本表(§8)。 ### 4.4 runtime_settings @@ -709,7 +709,7 @@ errors: 错误分布 / 排查 ← ops_error ## 6. SQLite v3 → PG 一次性搬迁 -引擎切换,不存在就地迁移:PG 库由 0001 基线全新建立,历史数据经**导入命令**搬入。旧文档的 0004 迁移草案(数据决策 1–4 与哈希回填)全部由导入命令吸收,SQLite 侧不再演进(v3 即其终版)。 +引擎切换,不存在就地迁移:PG 库由 0001 基线全新建立,历史数据经**导入命令**搬入。旧文档的 0004 迁移草案(数据决策 1–4 与管理 API key 哈希回填)全部由导入命令吸收,SQLite 侧不再演进(v3 即其终版)。 ### 6.1 导入命令 @@ -721,16 +721,16 @@ codex-proxy-rs import-sqlite <旧库路径.sqlite> - 源库校验:`schema_migrations` 最高版本必须为 3,否则拒绝(未知谱系不猜)。 - 全程单个 PG 事务:要么全部导入,要么全部回滚。源库只读打开。 - Redis 不参与导入:v3 的运行态数据按 §6.3 显式丢弃。 -- 结束打印逐表行数与丢弃计数报告。 +- 结束打印逐表行数、规范化计数与丢弃计数报告。 ### 6.2 逐表规则 | v3 源表 | 终态去向 | 变换 | | --- | --- | --- | | `admin_users` | PG 原样 | 时间戳 text→timestamptz | -| `client_api_keys` | PG | `key` 明文 → SHA-256 落 `key_hash`;明文自此不存在于任何存储 | +| `client_api_keys` | PG | `key`、`prefix` 原值搬迁；完整 key 继续作为管理端长期复制与 CCSwitch 导入的唯一事实 | | `runtime_settings` | PG | `admin_api_key` 明文 → 哈希落 `admin_api_key_hash`;三个 retention 列取默认值 | -| `accounts` | PG 原样 | 布尔 0/1→boolean,时间戳→timestamptz | +| `accounts` | PG | 布尔 0/1→boolean,时间戳→timestamptz；旧版曾落库的瞬态 `refreshing` → `expired` 并清空 `next_refresh_at`，防止迁移后自动消费 RT；其余业务状态原样 | | `account_usage` / `account_model_usage` | PG 原样 | 同上 | | `account_cookies` | PG | `expires_at` 字符串解析为 timestamptz(RFC2822/RFC3339);解析失败 → NULL(降级为会话 Cookie)并计数报告 | | `fingerprints` / `fingerprint_update_history` | PG 原样 | JSON text→jsonb | @@ -749,6 +749,7 @@ codex-proxy-rs import-sqlite <旧库路径.sqlite> - 全部管理端会话(重新登录)、会话亲和(退化为重新调度)、刷新租约(自愈)、模型清单快照(按需重拉)——运行态按 §2.9 本就可丢弃; - v3 usage_records 中的噪音行(决策 3); +- 旧版误落库的 `accounts.status='refreshing'` 不作为业务状态保留，按 §6.2 规范化为 `expired`； - 历史桶中错误路径污染的 token 按旧口径保留展示、不追溯清洗;事实保留期内的窗口随后由 `rebuild-buckets` 用新口径重算覆盖。 ### 6.4 迁移后动作与验收 @@ -757,7 +758,8 @@ codex-proxy-rs import-sqlite <旧库路径.sqlite> | --- | --- | --- | | 1 | 服务指向 PG+Redis 启动 | `/healthz` 204(PG `select 1` + Redis `PING` 双活) | | 2 | 管理员重新登录 | 登录成功;Redis 出现 `cpr:admin:session:*`,TTL ≈ session_ttl | -| 3 | 抽查 `select count(*) from client_api_keys where length(key_hash) <> 64` = 0 | 明文 key 消失,旧 key 鉴权仍通过 | +| 3 | 抽查 `select count(*) from client_api_keys where key is null or key = ''` = 0 | 旧 key 值不变；鉴权、管理端复制与 CCSwitch 导入均可用 | +| 3a | 抽查 `select count(*) from accounts where status = 'refreshing'` = 0 | 旧瞬态状态已规范化为 `expired`，且 `next_refresh_at is null` | | 4 | 执行 `rebuild-buckets`(§7),用新口径重算事实保留期内的桶 | 保留期内桶 token 列 = usage_records 聚合值 | | 5 | 核对 Dashboard 卡片全部走 bucket 口径,生命周期数字标注"成功累计" | 同卡片无混合口径数字 | @@ -765,19 +767,19 @@ codex-proxy-rs import-sqlite <旧库路径.sqlite> ## 7. 代码改造清单(本次交付) -与导入命令同一 PR 交付,按依赖顺序: +与导入命令同一 PR 交付,按依赖顺序。**2026-07-10 已逐项核验完成**，对应自动化验收由后端集成测试、前端构建、架构检查和 Compose 配置检查共同覆盖： -1. **连接层**(`infra/database.rs` + 新增 `infra/redis.rs`):PgPool(参数见 §1)+ 迁移框架(PG 谱系 0001 起);Redis ConnectionManager + 键前缀 + `PING` 健康检查。`Cargo.toml`:sqlx 加 `postgres`(保留 `sqlite` 仅供导入命令读源库),新增 `redis` crate。 -2. **配置**(`config/schema.rs`、`deploy/config.example.yaml`):`database.url` 改 `postgres://…`,新增 `redis.url`。 -3. **鉴权归因**(`admin/keys/service.rs`、`proxy/auth.rs`):key 创建落 `key_hash`;鉴权 = SHA-256 后 PG 点查(删除进程内明文鉴权表);`authorize_client_api_key_result` 返回 key `id`,调用链把 `client_api_key_id` 装进请求上下文直至事件写入。key 列表响应删除完整 key 字段(只余 `prefix`)——**前端同步移除完整 key 展示**。 -4. **事件结构**(`usage_record_model.rs`、`ops_error_model.rs`、`usage_events.rs`、`event_recording.rs`):成功/失败事件携带 `client_api_key_id` 与 `provider`;token、requested/upstream_model、service_tier、first_token_ms 从 metadata 移到一等字段;metadata 不再写已提升字段;UsageRecord 删除 `level`(成功事实无等级)。 -5. **写入路径**(usage record store、ops error store):§5.2 的两条 PG 事务;错误路径对桶只 `error_count + 1`;`__unknown__` sentinel;min/max 延迟直接 `least()`/`greatest()`(无 0-sentinel 分支);内联 trim 移除。 -6. **查询路径**(`usage_record_store/service/routes`、`dashboard.rs`、`account_usage_store.rs`、`admin/accounts/routes.rs`):summary/分布/趋势全部走列,删除所有 `->>` 聚合与 metadata LIKE;新增 `/api/admin/ops/errors` 错误明细查询面;Dashboard 按 §5.3 改卡片来源与文案。 -7. **运行态改 Redis**(`admin/auth/service.rs` 会话、`token_refresh` 租约、`session_affinity.rs`、`models/snapshot_store.rs`):按 §4B 契约实现;删除 admin 会话/亲和清理任务与亲和重启恢复;账号删除路径挂 Redis 亲和级联(§4B.3)。 +1. **连接层**(`infra/database.rs` + `infra/redis.rs`):PgPool(参数见 §1)+ 迁移框架(PG 谱系 0001 起);Redis ConnectionManager + 键前缀 + `PING` 健康检查。`Cargo.toml`:sqlx 加 `postgres`(保留 `sqlite` 仅供导入命令读源库),新增 `redis` crate。 +2. **配置**(`bootstrap/config.rs`、`deploy/config.example.yaml`):`database.url` 改 `postgres://…`,新增 `redis.url`。 +3. **鉴权归因**(`keys/{store,service,manage}.rs`、`api/client/auth.rs`):key 创建落唯一 `key` 列；鉴权按完整 key 做 PG unique 点查(删除进程内鉴权缓存);鉴权返回 key `id`,调用链把 `client_api_key_id` 装进请求上下文直至事件写入。管理端列表持续返回完整 key，前端长期保留复制与 CCSwitch 导入入口。 +4. **事件结构**(`telemetry/{usage,ops}`、`telemetry/recorder.rs`、`dispatch/responses/event_recording.rs`):成功/失败事件携带 `client_api_key_id` 与 `provider`;token、requested/upstream_model、service_tier、first_token_ms 从 metadata 移到一等字段;metadata 不再写已提升字段;UsageRecord 删除 `level`(成功事实无等级)。 +5. **写入路径**(`telemetry/usage/store.rs`、`telemetry/ops/store.rs`、`telemetry/buckets/store.rs`):§5.2 的两条 PG 事务;错误路径对桶只 `error_count + 1`;`__unknown__` sentinel;min/max 延迟使用 nullable-safe 的 `least()`/`greatest()`;内联 trim 移除。 +6. **查询路径**(`telemetry/{usage,ops,account_usage}`、`api/admin/{usage,ops,dashboard,accounts}_routes*`):summary/分布/趋势全部走列,删除所有 `->>` 聚合与 metadata LIKE;新增 `/api/admin/ops/errors` 错误明细查询面;Dashboard 按 §5.3 改卡片来源与文案。 +7. **运行态改 Redis**(`auth/store.rs` 会话、`accounts/refresh/lease.rs`、`dispatch/affinity/store.rs`、`models/store.rs`):按 §4B 契约实现;删除 admin 会话/亲和清理任务与亲和重启恢复;账号删除路径挂 Redis 亲和级联(§4B.3)。 8. **维护命令**(`main.rs` CLI):`import-sqlite `(§6);`rebuild-buckets`——删除事实保留期内的桶并从两张事实表重算(15 分钟槽对齐在 SQL 侧 `floor(epoch/900)*900`,§2.3);保留期外只读不动。范围明确**不含** `account_usage` / `account_model_usage`(§2.10,不可重建)。 -9. **清理任务**(`runtime/tasks/`):三张增长表的周期 trim(读 `runtime_settings` 三列);cookie 清理保留;admin 会话与亲和清理任务删除(TTL 接管)。 -10. **健康检查**(`http/router.rs`):PG `select 1` + Redis `PING`,任一失败即 503。 -11. **测试**(`backend/tests/`):support 换 PG(每测试独立数据库,`CPR_TEST_DATABASE_URL` 指定服务端)+ Redis(每测试随机键前缀,`CPR_TEST_REDIS_URL`);storage 测试断言终态 DDL、导入拆分规则(成功行迁入 / error 行转 ops / 噪音行丢弃)、两条写入事务的原子性、0ms 延迟样本可入桶、租约原子性。 +9. **清理任务**(`bootstrap/tasks/`):三张增长表的周期 trim(读 `runtime_settings` 三列);cookie 清理保留;admin 会话与亲和清理任务删除(TTL 接管)。 +10. **健康检查**(`api/router.rs`):PG `select 1` + Redis `PING`,任一失败即 503。 +11. **测试**(`backend/tests/`):support 换 PG(每测试独立数据库,`CPR_TEST_DATABASE_URL` 指定服务端)+ Redis(每测试随机键前缀,`CPR_TEST_REDIS_URL`);storage 测试断言终态 DDL、导入拆分规则(成功行迁入 / error 行转 ops / 噪音行丢弃 / 旧 `refreshing` 规范化)、两条写入事务的原子性、0ms 延迟样本可入桶、租约原子性。 12. **部署**(`deploy/docker-compose.yml`):新增 postgres、redis 服务与健康检查依赖。 --- diff --git a/frontend/src/api/modules/usage.ts b/frontend/src/api/modules/usage.ts index b223f0811..765957853 100644 --- a/frontend/src/api/modules/usage.ts +++ b/frontend/src/api/modules/usage.ts @@ -8,6 +8,14 @@ export function getUsageRecords(params?: any) { }) } +export function getOpsErrors(params?: any) { + return request({ + url: '/api/admin/ops/errors', + method: 'GET', + params, + }) +} + export function getUsageRecordDetail(params: any) { return request({ url: '/api/admin/usage/records/detail', diff --git a/frontend/src/views/api-keys/components/ApiKeyActions.vue b/frontend/src/views/api-keys/components/ApiKeyActions.vue index 5c72a3fe7..d919f5aa1 100644 --- a/frontend/src/views/api-keys/components/ApiKeyActions.vue +++ b/frontend/src/views/api-keys/components/ApiKeyActions.vue @@ -1,5 +1,5 @@ - - diff --git a/frontend/src/views/api-keys/composables/useApiKeyMutations.ts b/frontend/src/views/api-keys/composables/useApiKeyMutations.ts index 9d8125974..3d3d460e5 100644 --- a/frontend/src/views/api-keys/composables/useApiKeyMutations.ts +++ b/frontend/src/views/api-keys/composables/useApiKeyMutations.ts @@ -16,6 +16,7 @@ export function useApiKeyMutations(selectedIds: Ref>) { const showSingleDeleteModal = ref(false) const showKeyModal = ref(false) const createdKey = ref('') + const createdKeyName = ref('') const pendingDeleteKey = ref(null) const creatingKeyAction = useAsyncAction() const deletingKeyAction = useAsyncAction() @@ -58,6 +59,7 @@ export function useApiKeyMutations(selectedIds: Ref>) { }) createdKey.value = result.key + createdKeyName.value = result.name showCreateModal.value = false showKeyModal.value = true createForm.value = { name: '', label: '' } @@ -135,10 +137,6 @@ export function useApiKeyMutations(selectedIds: Ref>) { } } - function maskKey(prefix: string): string { - return `${prefix}••••••••••••••••` - } - onMounted(() => { loadApiKeys() }) @@ -151,6 +149,7 @@ export function useApiKeyMutations(selectedIds: Ref>) { showSingleDeleteModal, showKeyModal, createdKey, + createdKeyName, pendingDeleteKey, creatingKey, deletingKey, @@ -163,6 +162,5 @@ export function useApiKeyMutations(selectedIds: Ref>) { handleBatchDelete, handleToggleStatus, copyToClipboard, - maskKey, } } diff --git a/frontend/src/views/api-keys/constants.ts b/frontend/src/views/api-keys/constants.ts index 35999289b..b3a9bcb51 100644 --- a/frontend/src/views/api-keys/constants.ts +++ b/frontend/src/views/api-keys/constants.ts @@ -3,7 +3,7 @@ const relaxedActionCellClass = '!px-4' export const apiKeyColumns = [ { key: 'selection', label: '', width: '48px', minWidth: '48px', align: 'center' as const }, { key: 'identity', label: '名称', minWidth: '220px', flex: 0.9, fixed: 'left' as const }, - { key: 'prefix', label: '密钥前缀', minWidth: '300px', flex: 1.35 }, + { key: 'prefix', label: '密钥前缀', minWidth: '240px', flex: 1.35 }, { key: 'enabled', label: '状态', width: '112px', minWidth: '112px', align: 'center' as const }, { key: 'createdAtDisplay', @@ -22,8 +22,8 @@ export const apiKeyColumns = [ { key: 'actions', label: '操作', - width: '184px', - minWidth: '184px', + width: '112px', + minWidth: '112px', fixed: 'right' as const, headerClass: relaxedActionCellClass, cellClass: relaxedActionCellClass, diff --git a/frontend/src/views/api-keys/index.vue b/frontend/src/views/api-keys/index.vue index 3b951337d..9d9b2249b 100644 --- a/frontend/src/views/api-keys/index.vue +++ b/frontend/src/views/api-keys/index.vue @@ -1,11 +1,11 @@ @@ -155,11 +157,7 @@ function importToCcs(apiKey: any) { @@ -189,14 +186,7 @@ function importToCcs(apiKey: any) { :saving="creatingKey" @copy="copyToClipboard" @create="handleCreate" - /> - - -import { computed, shallowRef } from 'vue' - import BaseCard from '@/components/base/BaseCard.vue' -import BaseSegmented from '@/components/base/BaseSegmented.vue' import BaseTable from '@/components/base/BaseTable/index.vue' import { usageAccountText, @@ -16,24 +13,11 @@ import UsageCostCell from '@/views/usage/components/UsageCostCell.vue' import UsageModelCell from '@/views/usage/components/UsageModelCell.vue' import UsageTokenCell from '@/views/usage/components/UsageTokenCell.vue' -const props = defineProps<{ +defineProps<{ rows: any[] }>() -const activeFilter = shallowRef('all') -const filters = [ - { label: '全部', value: 'all' }, - { label: '错误', value: 'error' }, -] const dashboardUsageRecordColumns = usageRecordColumns.filter((column) => column.key !== 'actions') - -const filteredRows = computed(() => { - if (activeFilter.value === 'error') { - return props.rows.filter((row) => row.level === 'error' || Number(row.statusCode || 0) >= 400) - } - - return props.rows -})