refactor: complete postgresql redis migration

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
Zyy
2026-07-10 19:00:14 +08:00
co-authored by Codex
parent 68e76d8138
commit 45e4b36225
381 changed files with 21571 additions and 20990 deletions
+26
View File
@@ -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
+22 -13
View File
@@ -9,7 +9,8 @@
<img src="https://img.shields.io/badge/Axum-0.8-2f7d95?style=flat-square" />
<img src="https://img.shields.io/badge/Vue-3.5-42b883?style=flat-square&logo=vuedotjs&logoColor=white" />
<img src="https://img.shields.io/badge/Vite-8-646cff?style=flat-square&logo=vite&logoColor=white" />
<img src="https://img.shields.io/badge/SQLite-local-003b57?style=flat-square&logo=sqlite" />
<img src="https://img.shields.io/badge/PostgreSQL-18-4169e1?style=flat-square&logo=postgresql&logoColor=white" />
<img src="https://img.shields.io/badge/Redis-8-ff4438?style=flat-square&logo=redis&logoColor=white" />
<img src="https://img.shields.io/badge/License-MIT-blue?style=flat-square" />
</p>
@@ -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/ 架构和维护文档
+99
View File
@@ -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"
+2 -2
View File
@@ -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"
+59
View File
@@ -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"
+5
View File
@@ -0,0 +1,5 @@
//! 账号 Cookie 存储。
mod store;
pub use store::*;
@@ -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<T> = Result<T, SqliteCookieStoreError>;
/// PostgreSQL Cookie 存储结果。
pub type PgCookieStoreResult<T> = Result<T, PgCookieStoreError>;
/// 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<Option<String>> {
) -> PgCookieStoreResult<Option<String>> {
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<Option<String>> {
) -> PgCookieStoreResult<Option<String>> {
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::<Option<String>, _>("expires_at");
if cookie_is_expired(expires_at.as_deref(), now) {
let expires_at = row.get::<Option<DateTime<Utc>>, _>("expires_at");
if expires_at.is_some_and(|expires_at| expires_at <= now) {
continue;
}
let name = row.get::<String, _>("name");
@@ -114,8 +110,8 @@ impl SqliteCookieStore {
}
/// 删除账号全部 Cookie。
pub async fn delete_account_cookies(&self, account_id: &str) -> SqliteCookieStoreResult<u64> {
let result = sqlx::query("delete from account_cookies where account_id = ?")
pub async fn delete_account_cookies(&self, account_id: &str) -> PgCookieStoreResult<u64> {
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<Utc>) -> SqliteCookieStoreResult<u64> {
pub async fn cleanup_expired(&self, now: DateTime<Utc>) -> PgCookieStoreResult<u64> {
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<String>,
expires_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -197,7 +193,7 @@ fn parse_set_cookie(raw: &str) -> Option<ParsedCookie> {
}
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<ParsedCookie> {
})
}
fn max_age_expires_at(seconds: i64) -> String {
fn max_age_expires_at(seconds: i64) -> DateTime<Utc> {
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<Utc>) -> 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<Utc>> {
DateTime::parse_from_rfc2822(value)
.or_else(|_| DateTime::parse_from_rfc3339(value))
@@ -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)]
@@ -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<String> {
self.cookies
.cookie_header_for_request(account_id, "chatgpt.com", "/codex/usage")
@@ -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<String>) -> Result<ExportedAccounts, AdminAccountError> {
impl AccountManageService {
pub async fn export(&self, ids: Vec<String>) -> Result<ExportedAccounts, AccountManageError> {
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<String>,
) -> Result<Vec<StoredAccount>, AdminAccountError> {
) -> Result<Vec<StoredAccount>, 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);
}
@@ -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<ImportedAccounts, AdminAccountError> {
let parsed = crate::upstream::accounts::importing::parse_account_import_payload(&data)
.map_err(|_| AdminAccountError::NoImportableAccounts)?;
) -> Result<ImportedAccounts, AccountManageError> {
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<ImportedAccountState, AdminAccountError> {
entry: crate::accounts::import::AccountImportEntry,
source: crate::accounts::import::AccountImportSource,
) -> Result<ImportedAccountState, AccountManageError> {
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<String>,
refresh_token: Option<String>,
) -> Result<Option<ResolvedImportTokens>, AdminAccountError> {
let refresh_token = crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
let Some(access_token) = crate::upstream::accounts::importing::normalize_nonempty(
) -> Result<Option<ResolvedImportTokens>, 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<String, AdminAccountError> {
let provided_id = crate::upstream::accounts::importing::normalize_nonempty_str(id)
.map(ToString::to_string);
) -> Result<String, AccountManageError> {
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)))
}
}
@@ -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<String>,
limit: u32,
) -> Result<Page<AdminAccountMetadata>, AdminAccountError> {
) -> Result<Page<ManagedAccount>, 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<String>,
) -> Result<NumberedPage<AdminAccountMetadata>, AdminAccountError> {
) -> Result<NumberedPage<ManagedAccount>, 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<Option<AdminAccountMetadata>, AdminAccountError> {
) -> Result<Option<ManagedAccount>, 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<String>,
refresh_token: Option<String>,
) -> Result<AdminAccountMetadata, AdminAccountError> {
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<ManagedAccount, AccountManageError> {
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<Option<AdminAccountMetadata>, AdminAccountError> {
update: AccountUpdate,
) -> Result<Option<ManagedAccount>, 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<String>,
) -> Result<BatchDeleteAccounts, AdminAccountError> {
) -> Result<BatchDeleteAccounts, AccountManageError> {
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<AdminAccountRefreshResult, AdminAccountError> {
) -> Result<AccountRefreshResult, AccountManageError> {
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<Vec<String>>,
stagger_ms: u64,
concurrency: usize,
) -> Result<AdminAccountHealthCheck, AdminAccountError> {
) -> Result<AccountHealthCheck, AccountManageError> {
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::<HashSet<_>>());
let candidate_ids = accounts
.into_iter()
@@ -437,12 +432,12 @@ impl AdminAccountService {
.into_iter()
.collect::<Result<Vec<_>, _>>()?;
Ok(AdminAccountHealthCheck { results })
Ok(AccountHealthCheck { results })
}
}
fn create_next_refresh_at(
service: &AdminAccountService,
service: &AccountManageService,
account_id: &str,
expires_at: DateTime<Utc>,
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,
}
}
+15
View File
@@ -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::*;
@@ -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<OAuthAuthorizeResult, AdminAccountError> {
async fn authorize(&self) -> Result<OAuthAuthorizeResult, AccountManageError> {
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<AccountOAuthTokens, AdminAccountError> {
) -> Result<AccountOAuthTokens, AccountManageError> {
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::<AccountOAuthTokens>(&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<AccountOAuthSession, AdminAccountError> {
async fn session(&self, session_id: &str) -> Result<AccountOAuthSession, AccountManageError> {
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<OAuthAuthorizeResult, AdminAccountError> {
impl AccountManageService {
pub async fn oauth_authorize(&self) -> Result<OAuthAuthorizeResult, AccountManageError> {
self.oauth.authorize().await
}
pub async fn oauth_exchange(
&self,
input: OAuthExchangeInput,
) -> Result<super::contracts::ImportedAccounts, AdminAccountError> {
) -> Result<super::types::ImportedAccounts, AccountManageError> {
let tokens = self.oauth.exchange(input).await?;
self.import(serde_json::json!({
"sourceFormat": "cpr",
@@ -170,21 +170,23 @@ struct AccountOAuthTokens {
refresh_token: Option<String>,
}
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<String> {
@@ -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<Box<dyn Stream<Item = Result<Bytes, Infallible>> + Send>>;
@@ -59,17 +59,17 @@ impl AccountTestOutcome {
}
}
impl AdminAccountService {
impl AccountManageService {
pub async fn account_models(
&self,
account_id: &str,
) -> Result<Vec<AccountModelOption>, AdminAccountError> {
) -> Result<Vec<AccountModelOption>, 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::<Vec<_>>();
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<AccountTestStream, AdminAccountError> {
) -> Result<AccountTestStream, 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 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<Bytes>,
) -> AccountTestOutcome {
let mut buffer = String::new();
@@ -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<serde_json::Value, AdminAccountError> {
) -> Result<serde_json::Value, AccountManageError> {
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;
@@ -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<AdminAccountQuotaWindowData>,
windows: Vec<AccountQuotaWindowData>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountQuotaWindowData {
struct AccountQuotaWindowData {
key: String,
group: String,
window_seconds: Option<u64>,
@@ -32,7 +32,7 @@ struct AdminAccountQuotaWindowData {
used_percent: Option<f64>,
used_percent_display: String,
#[serde(skip_serializing_if = "Option::is_none")]
local_usage: Option<AdminAccountQuotaWindowLocalUsageData>,
local_usage: Option<AccountQuotaWindowLocalUsageData>,
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<Utc>,
pub(crate) end: DateTime<Utc>,
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<AdminAccountQuotaWindowLocalUsage> for AdminAccountQuotaWindowLocalUsageData {
fn from(usage: AdminAccountQuotaWindowLocalUsage) -> Self {
impl From<AccountQuotaWindowLocalUsage> 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<AdminAccountQuotaWindowLocalUsage> 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<AdminAccountQuotaUsageWindow> {
pub(crate) fn usage_windows(&self) -> Vec<AccountQuotaUsageWindow> {
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<String, AdminAccountQuotaWindowLocalUsage>,
usage_by_window: &HashMap<String, AccountQuotaWindowLocalUsage>,
) {
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<AdminAccountQuotaUsageWindow> {
impl AccountQuotaWindowData {
fn usage_window(&self) -> Option<AccountQuotaUsageWindow> {
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<DateTime<Utc>>,
) -> AdminAccountQuotaData {
pub(crate) fn quota_data(quota_json: &str, fetched_at: Option<DateTime<Utc>>) -> AccountQuotaData {
let quota = serde_json::from_str::<Value>(quota_json).unwrap_or(Value::Null);
let windows = quota_windows(&quota);
AdminAccountQuotaData {
AccountQuotaData {
refreshed_at_display: china_relative_time(fetched_at, Utc::now()),
windows,
}
}
fn quota_windows(quota: &Value) -> Vec<AdminAccountQuotaWindowData> {
fn quota_windows(quota: &Value) -> Vec<AccountQuotaWindowData> {
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<AdminAccountQuotaWindowData> {
}
fn push_monthly_quota_window(
windows: &mut Vec<AdminAccountQuotaWindowData>,
windows: &mut Vec<AccountQuotaWindowData>,
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<AdminAccountQuotaWindowData>,
windows: &mut Vec<AccountQuotaWindowData>,
snapshot: &Value,
skip_core_monthly: bool,
) {
@@ -269,7 +257,7 @@ fn push_snapshot_quota_windows(
}
fn push_quota_window(
windows: &mut Vec<AdminAccountQuotaWindowData>,
windows: &mut Vec<AccountQuotaWindowData>,
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<String> {
}
fn unique_quota_window_key(
windows: &[AdminAccountQuotaWindowData],
windows: &[AccountQuotaWindowData],
source_key: &str,
role: &str,
window_seconds: Option<u64>,
@@ -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,
@@ -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<CodexBackendClient>,
pub(crate) models: StdArc<ModelService>,
pub(crate) account_pool: StdArc<RuntimeAccountPoolService>,
pub(crate) token_refresher: StdArc<dyn TokenRefresher>,
pub(crate) refresh_leases: RefreshLeaseStore,
pub struct AccountManageService {
pub store: PgAccountStore,
pub(crate) cookies: PgCookieStore,
pub(crate) codex: Arc<CodexBackendClient>,
pub(crate) models: Arc<ModelService>,
pub(crate) account_pool: Arc<RuntimeAccountPoolService>,
pub(crate) token_refresher: Arc<dyn TokenRefresher>,
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<String>,
}
pub struct AdminAccountServiceParts {
pub store: SqliteAccountStore,
pub cookies: SqliteCookieStore,
pub codex: StdArc<CodexBackendClient>,
pub models: StdArc<ModelService>,
pub account_pool: StdArc<RuntimeAccountPoolService>,
pub token_refresher: StdArc<dyn TokenRefresher>,
pub refresh_leases: RefreshLeaseStore,
pub struct AccountManageServiceParts {
pub store: PgAccountStore,
pub cookies: PgCookieStore,
pub codex: Arc<CodexBackendClient>,
pub models: Arc<ModelService>,
pub account_pool: Arc<RuntimeAccountPoolService>,
pub token_refresher: Arc<dyn TokenRefresher>,
pub refresh_leases: RedisRefreshLeaseStore,
pub session_affinity: RedisSessionAffinityStore,
pub oauth: oauth::AccountOAuthService,
pub refresh_policy: RuntimeRefreshPolicy,
pub installation_id: Option<String>,
}
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<chrono::Utc>,
) -> DateTime<chrono::Utc> {
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<contracts::RefreshedAccountTokens, AdminAccountError> {
) -> Result<RefreshedAccountTokens, AccountManageError> {
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,
@@ -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<String>,
pub account_id: Option<String>,
@@ -27,7 +27,7 @@ pub struct AdminAccountMetadata {
pub updated_at: DateTime<Utc>,
}
impl From<StoredAccountMetadata> for AdminAccountMetadata {
impl From<StoredAccountMetadata> for ManagedAccount {
fn from(m: StoredAccountMetadata) -> Self {
let added_at = m
.added_at
@@ -54,7 +54,7 @@ impl From<StoredAccountMetadata> 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<String>,
pub previous_status: AcctStatus,
pub outcome: AdminAccountRefreshOutcome,
pub outcome: AccountRefreshOutcome,
pub error: Option<String>,
pub duration_ms: i64,
}
#[derive(Debug, Clone, Default)]
pub struct AdminAccountHealthCheck {
pub results: Vec<AdminAccountRefreshResult>,
pub struct AccountHealthCheck {
pub results: Vec<AccountRefreshResult>,
}
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<Option<String>>,
pub status: Option<String>,
}
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<String>,
pub(super) refresh_token_for_existing: Option<String>,
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<String>,
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<String>,
pub(super) claims: Option<crate::upstream::accounts::token_refresh::ManualAccountClaims>,
pub(super) claims: Option<crate::accounts::refresh::ManualAccountClaims>,
}
#[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<String> {
}
pub(super) fn import_status_from_usage_error(error: &CodexClientError) -> Option<AccountStatus> {
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<AcctStatus, AdminAccountError> {
pub(super) fn parse_account_status(status: &str) -> Result<AcctStatus, AccountManageError> {
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(),
)),
}
+13
View File
@@ -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;
@@ -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<String>,
/// 模型到允许订阅计划的映射。
pub model_plan_allowlist: BTreeMap<String, Vec<String>>,
/// 已成功拉取过模型列表的订阅计划。
pub fetched_model_plan_types: BTreeSet<String>,
}
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<String>,
/// 本次请求必须使用的账号 ID,用于同请求内的同账号重试。
pub required_account_id: Option<String>,
/// 会话亲和性建议的优先账号 ID。
pub preferred_account_id: Option<String>,
/// 本次调度使用的当前时间。
pub now: DateTime<Utc>,
}
impl AccountAcquireRequest {
/// 构造账号获取请求。
pub fn new(model: impl Into<String>, now: DateTime<Utc>) -> 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<Item = impl Into<String>>,
) -> 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<String>) -> Self {
self.preferred_account_id = Some(account_id.into());
self
}
/// 设置本次调度必须使用的账号 ID。
pub fn with_required_account_id(mut self, account_id: impl Into<String>) -> Self {
self.required_account_id = Some(account_id.into());
self
}
}
/// 成功获取到的账号及调度元数据。
#[derive(Debug, Clone)]
pub struct AcquiredAccount {
/// 被选中的账号快照。
pub account: Account,
/// 同一账号上一个在途槽位的创建时间。
pub previous_slot_at: Option<DateTime<Utc>>,
}
/// 模型刷新时每个订阅计划选中的账号。
#[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<Utc>) -> bool {
}
/// 运行时账号池服务。
#[derive(Clone)]
pub struct RuntimeAccountPoolService {
pool: StdArc<tokio::sync::Mutex<AccountPool>>,
store: StdArc<dyn AccountStore>,
request_interval: StdArc<std::sync::RwLock<StdDuration>>,
}
impl RuntimeAccountPoolService {
/// 构造运行时账号池服务。
pub fn new(
store: StdArc<dyn AccountStore>,
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<String, Vec<String>>,
fetched_model_plan_types: BTreeSet<String>,
) {
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<usize, RuntimeAccountPoolError> {
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<Utc>) -> 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<Utc>) -> 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<AcquiredAccount> {
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<Utc>) -> Vec<DistinctPlanAccount> {
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<RefreshedAccountState>) {
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::<serde_json::Value>(&quota_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, &quota).await;
}
/// 读取运行时账号快照。
pub async fn account_snapshot(&self, account_id: &str) -> Option<Account> {
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, &quota_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<bool, RuntimeAccountPoolError> {
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<Utc>,
) -> 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<Utc>) {
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<DateTime<Utc>>,
limit_window_seconds: Option<u64>,
now: DateTime<Utc>,
) -> Option<DateTime<Utc>> {
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<Utc>) {
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,
};
+593
View File
@@ -0,0 +1,593 @@
use super::*;
#[derive(Clone)]
pub struct RuntimeAccountPoolService {
pool: Arc<tokio::sync::Mutex<AccountPool>>,
store: Arc<dyn AccountStore>,
request_interval: Arc<std::sync::RwLock<StdDuration>>,
}
impl RuntimeAccountPoolService {
/// 构造运行时账号池服务。
pub fn new(
store: Arc<dyn AccountStore>,
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<String, Vec<String>>,
fetched_model_plan_types: BTreeSet<String>,
) {
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<usize, RuntimeAccountPoolError> {
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<Utc>) -> 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<Utc>) -> 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<AcquiredAccount> {
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<Utc>) -> Vec<DistinctPlanAccount> {
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<RefreshedAccountState>) {
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::<serde_json::Value>(&quota_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, &quota).await;
}
/// 读取运行时账号快照。
pub async fn account_snapshot(&self, account_id: &str) -> Option<Account> {
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, &quota_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<bool, RuntimeAccountPoolError> {
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<Utc>,
) -> 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<Utc>) {
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<DateTime<Utc>>,
limit_window_seconds: Option<u64>,
now: DateTime<Utc>,
) -> Option<DateTime<Utc>> {
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<Utc>) {
if account
.cloudflare_cooldown_until
.is_some_and(|cooldown_until| now >= cooldown_until)
{
account.cloudflare_cooldown_until = None;
}
}
+100
View File
@@ -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<String>,
/// 模型到允许订阅计划的映射。
pub model_plan_allowlist: BTreeMap<String, Vec<String>>,
/// 已成功拉取过模型列表的订阅计划。
pub fetched_model_plan_types: BTreeSet<String>,
}
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<String>,
/// 本次请求必须使用的账号 ID,用于同请求内的同账号重试。
pub required_account_id: Option<String>,
/// 会话亲和性建议的优先账号 ID。
pub preferred_account_id: Option<String>,
/// 本次调度使用的当前时间。
pub now: DateTime<Utc>,
}
impl AccountAcquireRequest {
/// 构造账号获取请求。
pub fn new(model: impl Into<String>, now: DateTime<Utc>) -> 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<Item = impl Into<String>>,
) -> 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<String>) -> Self {
self.preferred_account_id = Some(account_id.into());
self
}
/// 设置本次调度必须使用的账号 ID。
pub fn with_required_account_id(mut self, account_id: impl Into<String>) -> Self {
self.required_account_id = Some(account_id.into());
self
}
}
/// 成功获取到的账号及调度元数据。
#[derive(Debug, Clone)]
pub struct AcquiredAccount {
/// 被选中的账号快照。
pub account: Account,
/// 同一账号上一个在途槽位的创建时间。
pub previous_slot_at: Option<DateTime<Utc>>,
}
/// 模型刷新时每个订阅计划选中的账号。
#[derive(Debug, Clone)]
pub struct DistinctPlanAccount {
/// 订阅计划类型。
pub plan_type: String,
/// 被选中的账号快照。
pub account: Account,
}
+7
View File
@@ -0,0 +1,7 @@
//! 账号配额解析与刷新。
mod runtime;
mod service;
pub use runtime::*;
pub use service::*;
@@ -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<CodexBackendClient>,
min_refresh_interval_secs: u64,
request_spacing: Duration,
installation_id: Option<String>,
cookie_store: Option<SqliteCookieStore>,
cookie_store: Option<PgCookieStore>,
account_pool: Option<Arc<RuntimeAccountPoolService>>,
}
impl RuntimeQuotaRefreshService {
/// 构造默认 quota refresh 服务。
pub fn new(store: SqliteAccountStore, codex: Arc<CodexBackendClient>) -> Self {
pub fn new(store: PgAccountStore, codex: Arc<CodexBackendClient>) -> 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<CodexBackendClient>,
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),
@@ -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,
};
+118
View File
@@ -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<T> = Result<T, RedisRefreshLeaseStoreError>;
/// 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<Utc>,
now: DateTime<Utc>,
) -> RedisRefreshLeaseStoreResult<bool> {
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<bool> {
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<Utc>,
) -> RedisRefreshLeaseStoreResult<HashSet<String>> {
if account_ids.is_empty() {
return Ok(HashSet::new());
}
let keys = account_ids
.iter()
.map(|account_id| self.lease_key(account_id))
.collect::<Vec<_>>();
let mut connection = self.redis.manager();
let values: Vec<Option<String>> = 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}"))
}
}
+9
View File
@@ -0,0 +1,9 @@
//! Token 刷新策略、租约与运行时服务。
mod lease;
mod runtime;
mod service;
pub use lease::*;
pub use runtime::*;
pub use service::*;
@@ -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<C>
where
C: TokenRefresher,
{
store: SqliteAccountStore,
store: PgAccountStore,
scheduler: Arc<RefreshScheduler<C>>,
policy: RuntimeRefreshPolicy,
refresh_leases: Option<RefreshLeaseStore>,
refresh_leases: Option<RedisRefreshLeaseStore>,
lease_owner: String,
retry_delays: Vec<Duration>,
in_flight: Arc<Mutex<HashSet<String>>>,
@@ -88,11 +88,7 @@ where
C: TokenRefresher,
{
/// 构造运行时 token refresh 服务。
pub fn new(
store: SqliteAccountStore,
policy: impl Into<RuntimeRefreshPolicy>,
client: C,
) -> Self {
pub fn new(store: PgAccountStore, policy: impl Into<RuntimeRefreshPolicy>, 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<Utc>,
) -> TokenRefreshServiceResult<TokenRefreshOutcome> {
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<Account>),
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<T> = Result<T, TokenRefreshServiceError>;
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<DateTime<Utc>>,
) -> TokenRefreshServiceResult<()> {
@@ -845,44 +775,3 @@ fn same_refresh_time(left: Option<DateTime<Utc>>, right: Option<DateTime<Utc>>)
_ => 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,
}
}
@@ -0,0 +1,109 @@
use super::*;
pub(super) enum TokenRefreshOutcome {
Refreshed(Box<Account>),
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<T> = Result<T, TokenRefreshServiceError>;
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,
}
}
@@ -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<String>,
}
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<TokenPair, RefreshFailure>;
}
/// 负责执行单账号刷新策略的调度器。
#[derive(Clone)]
pub struct RefreshScheduler<C> {
@@ -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<T> = Result<T, RefreshLeaseStoreError>;
/// 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<Utc>,
now: DateTime<Utc>,
) -> RefreshLeaseStoreResult<bool> {
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<bool> {
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<Utc>,
) -> RefreshLeaseStoreResult<HashSet<String>> {
if account_ids.is_empty() {
return Ok(HashSet::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(
"select account_id from account_refresh_leases where datetime(expires_at) > datetime(",
);
builder.push_bind(now.to_rfc3339());
builder.push(") and account_id in (");
let mut separated = builder.separated(", ");
for account_id in account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder
.build_query_as::<(String,)>()
.fetch_all(&self.pool)
.await?;
Ok(rows.into_iter().map(|(account_id,)| account_id).collect())
}
}
mod runtime;
pub use runtime::{
RuntimeTokenRefreshService, TokenRefreshServiceError, TokenRefreshServiceResult,
TokenRefreshSummary, TokenTimerSummary,
};
@@ -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};
@@ -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> {
@@ -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};
@@ -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};
@@ -3,7 +3,7 @@
//! 最简单的策略:忽略账号负载与用量,纯粹按候选集合的顺序用轮转游标依次选中,
//! 保证请求在候选账号间均匀轮转。
use crate::upstream::accounts::model::Account;
use crate::accounts::account::Account;
use super::SelectionInput;
@@ -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};
@@ -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};
@@ -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<Account> {
match strategy {
RotationStrategy::Smart => smart::select(input, cursor),
RotationStrategy::QuotaResetPriority => quota_reset::select(input, cursor),
RotationStrategy::RoundRobin => round_robin::select(input, cursor),
RotationStrategy::Sticky => sticky::select(input, cursor),
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),
}
}
@@ -2,7 +2,7 @@
use chrono::{DateTime, Utc};
use crate::upstream::accounts::model::Account;
use crate::accounts::account::Account;
/// 账号服务。
#[derive(Debug, Clone, Default)]
+477
View File
@@ -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<T> = Result<T, PgAccountStoreError>;
// ============================================================================
// 数据类型
// ============================================================================
/// 新建账号数据。
#[derive(Debug)]
pub struct NewAccount {
/// 账号 ID。
pub id: String,
/// 邮箱。
pub email: Option<String>,
/// ChatGPT 账号 ID。
pub account_id: Option<String>,
/// ChatGPT 用户 ID。
pub user_id: Option<String>,
/// 管理标签。
pub label: Option<String>,
/// 计划类型。
pub plan_type: Option<String>,
/// access token 明文。
pub access_token: SecretString,
/// refresh token 明文。
pub refresh_token: Option<SecretString>,
/// access token 过期时间。
pub access_token_expires_at: Option<DateTime<Utc>>,
/// 初始状态。
pub status: AccountStatus,
/// 账号添加时间。
pub added_at: Option<DateTime<Utc>>,
}
/// 通过 JWT claims 更新已有账号。
#[derive(Debug)]
pub struct AccountClaimsUpdate {
/// 展示邮箱。
pub email: Option<String>,
/// ChatGPT 账号 ID。
pub account_id: Option<String>,
/// ChatGPT 用户 ID。
pub user_id: Option<String>,
/// 计划类型。
pub plan_type: Option<String>,
/// 新 access token。
pub access_token: SecretString,
/// 可选的新 refresh token。
pub refresh_token: Option<SecretString>,
/// access token 过期时间。
pub access_token_expires_at: Option<DateTime<Utc>>,
/// 下一次刷新时间。
pub next_refresh_at: Option<DateTime<Utc>>,
/// 更新后的状态。
pub status: AccountStatus,
}
/// 通过导入数据更新已有账号。
#[derive(Debug)]
pub struct ImportedAccountUpdate {
/// 新建账号数据。
pub account: NewAccount,
/// 缓存配额 JSON。
pub quota_json: Option<String>,
/// 配额抓取时间。
pub quota_fetched_at: Option<DateTime<Utc>>,
/// 是否需要执行额外配额校验。
pub quota_verify_required: bool,
}
/// 存储中的账号完整记录。
#[derive(Debug, Clone)]
pub struct StoredAccount {
/// 账号 ID。
pub id: String,
/// 邮箱。
pub email: Option<String>,
/// ChatGPT 账号 ID。
pub account_id: Option<String>,
/// ChatGPT 用户 ID。
pub user_id: Option<String>,
/// 管理标签。
pub label: Option<String>,
/// 计划类型。
pub plan_type: Option<String>,
/// access token 明文。
pub access_token: SecretString,
/// refresh token 明文。
pub refresh_token: Option<SecretString>,
/// access token 过期时间。
pub access_token_expires_at: Option<DateTime<Utc>>,
/// 下一次刷新时间。
pub next_refresh_at: Option<DateTime<Utc>>,
/// 账号状态。
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<String>,
/// ChatGPT 账号 ID。
pub account_id: Option<String>,
/// ChatGPT 用户 ID。
pub user_id: Option<String>,
/// 管理标签。
pub label: Option<String>,
/// 计划类型。
pub plan_type: Option<String>,
/// 是否保存了 refresh token。
pub has_refresh_token: bool,
/// access token 过期时间。
pub access_token_expires_at: Option<DateTime<Utc>>,
/// 账号状态。
pub status: AccountStatus,
/// 添加时间。
pub added_at: String,
/// 更新时间。
pub updated_at: String,
}
#[derive(Debug, Clone, Default)]
pub struct AccountMetadataUpdate {
pub email: Option<Option<String>>,
pub account_id: Option<Option<String>>,
pub user_id: Option<Option<String>>,
pub label: Option<Option<String>>,
pub plan_type: Option<Option<String>>,
pub status: Option<AccountStatus>,
}
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<String>,
/// 配额 JSON。
pub quota_json: String,
/// 配额抓取时间。
pub quota_fetched_at: Option<DateTime<Utc>>,
}
/// 用量增量(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<T> = Result<T, AccountStoreError>;
/// 提供运行时账号列表的端口。
#[async_trait]
pub trait AccountStore: Send + Sync + 'static {
/// 列出当前账号池可见的账号。
async fn list_pool_accounts(&self) -> AccountStoreResult<Vec<Account>>;
/// 读取单个账号池账号快照。
async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult<Option<Account>> {
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<Utc>,
) -> AccountStoreResult<bool>;
/// 标记账号进入 Cloudflare 冷却期。
async fn set_cloudflare_cooldown_until(
&self,
account_id: &str,
cooldown_until: DateTime<Utc>,
) -> AccountStoreResult<bool>;
/// 更新账号状态。
async fn set_status(&self, account_id: &str, status: AccountStatus)
-> AccountStoreResult<bool>;
/// 记录账号用量增量。
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<Option<String>>;
/// 应用已经验证过的账号配额快照。
async fn apply_quota_snapshot(
&self,
account_id: &str,
quota_json: &str,
limit_reached: bool,
cooldown_until: Option<DateTime<Utc>>,
) -> AccountStoreResult<bool>;
/// 同步运行时自然刷新出来的账号状态。
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool>;
/// 同步账号当前 rate-limit 统计窗口。
async fn sync_rate_limit_window(
&self,
account_id: &str,
reset_at: DateTime<Utc>,
limit_window_seconds: Option<u64>,
) -> 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<String>,
) -> Result<bool, PgAccountStoreError> {
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<Vec<Account>> {
list_pool_accounts(self)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn get_pool_account(&self, account_id: &str) -> AccountStoreResult<Option<Account>> {
get_pool_account(self, account_id)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn mark_quota_limited_until(
&self,
account_id: &str,
cooldown_until: DateTime<Utc>,
) -> AccountStoreResult<bool> {
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<Utc>,
) -> AccountStoreResult<bool> {
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<bool> {
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<Option<String>> {
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<DateTime<Utc>>,
) -> AccountStoreResult<bool> {
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<bool> {
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<Utc>,
limit_window_seconds: Option<u64>,
) -> AccountStoreResult<()> {
PgAccountStore::sync_rate_limit_window(self, account_id, reset_at, limit_window_seconds)
.await
.map_err(|error| map_account_store_error(&error))
}
}
@@ -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";
@@ -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<Vec<Account>> {
store: &PgAccountStore,
) -> PgAccountStoreResult<Vec<Account>> {
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<Option<Account>> {
) -> PgAccountStoreResult<Option<Account>> {
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<Account> {
pub(super) fn pool_account_from_row(row: &PgRow) -> PgAccountStoreResult<Account> {
let quota_json = row
.get::<Option<String>, _>("quota_json")
.and_then(|quota_json| serde_json::from_str::<Value>(&quota_json).ok());
.get::<Option<sqlx::types::Json<Value>>, _>("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::<Option<String>, _>("access_token_expires_at")
.as_deref(),
)?,
next_refresh_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("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::<String, _>("status"))?,
quota_limit_reached: row.get::<i64, _>("quota_limit_reached") != 0,
quota_verify_required: row.get::<i64, _>("quota_verify_required") != 0,
quota_cooldown_until: parse_optional_rfc3339(
row.get::<Option<String>, _>("quota_cooldown_until")
.as_deref(),
)?,
cloudflare_cooldown_until: parse_optional_rfc3339(
row.get::<Option<String>, _>("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::<i64, _>("usage_request_count")),
empty_response_count: nonnegative_i64_to_u64(
row.get::<i64, _>("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::<i64, _>("usage_window_image_request_failed_count"),
),
window_started_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("usage_window_started_at")
.as_deref(),
)?,
window_reset_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("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::<Option<chrono::DateTime<chrono::Utc>>, _>("usage_window_reset_at")
.or(quota_window_reset_at),
limit_window_seconds: optional_positive_i64_to_u64(
row.get::<Option<i64>, _>("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::<chrono::DateTime<chrono::Utc>, _>("added_at")
.to_rfc3339(),
last_used_at: row
.get::<Option<chrono::DateTime<chrono::Utc>>, _>("usage_last_used_at")
.map(|value| value.to_rfc3339()),
})
}
pub(super) fn stored_account_from_row(row: &SqliteRow) -> SqliteAccountStoreResult<StoredAccount> {
pub(super) fn stored_account_from_row(row: &PgRow) -> PgAccountStoreResult<StoredAccount> {
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::<Option<String>, _>("refresh_token")
.map(|token| SecretString::new(token.into())),
access_token_expires_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("access_token_expires_at")
.as_deref(),
)?,
next_refresh_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("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::<String, _>("status"))?,
added_at: row.get("added_at"),
updated_at: row.get("updated_at"),
added_at: row
.get::<chrono::DateTime<chrono::Utc>, _>("added_at")
.to_rfc3339(),
updated_at: row
.get::<chrono::DateTime<chrono::Utc>, _>("updated_at")
.to_rfc3339(),
})
}
pub(super) fn metadata_from_row(
row: &SqliteRow,
) -> SqliteAccountStoreResult<StoredAccountMetadata> {
pub(super) fn metadata_from_row(row: &PgRow) -> PgAccountStoreResult<StoredAccountMetadata> {
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::<i64, _>("has_refresh_token") != 0,
access_token_expires_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("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::<String, _>("status"))?,
added_at: row.get("added_at"),
updated_at: row.get("updated_at"),
added_at: row
.get::<chrono::DateTime<chrono::Utc>, _>("added_at")
.to_rfc3339(),
updated_at: row
.get::<chrono::DateTime<chrono::Utc>, _>("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<String> {
}
pub(super) fn quota_snapshot_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountQuotaSnapshot> {
row: &sqlx::postgres::PgRow,
) -> PgAccountStoreResult<AccountQuotaSnapshot> {
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::<Option<String>, _>("quota_fetched_at").as_deref(),
)?,
quota_json: row
.get::<sqlx::types::Json<Value>, _>("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<String>>) -> Option<&s
value.as_ref().and_then(|value| value.as_deref())
}
pub(super) fn status_from_db(value: &str) -> SqliteAccountStoreResult<AccountStatus> {
pub(super) fn status_from_db(value: &str) -> PgAccountStoreResult<AccountStatus> {
match value {
"active" => Ok(AccountStatus::Active),
"expired" => Ok(AccountStatus::Expired),
"quota_exhausted" => Ok(AccountStatus::QuotaExhausted),
"disabled" => Ok(AccountStatus::Disabled),
"banned" => Ok(AccountStatus::Banned),
other => Err(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<i64>) -> Option<u64> {
}
pub(super) async fn count_account_metadata(
pool: &SqlitePool,
pool: &PgPool,
search: Option<&str>,
) -> SqliteAccountStoreResult<u64> {
let mut builder = QueryBuilder::<Sqlite>::new("select count(*) from accounts");
) -> PgAccountStoreResult<u64> {
let mut builder = QueryBuilder::<Postgres>::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<Sqlite>,
builder: &mut QueryBuilder<Postgres>,
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<T>(
rows: &[SqliteRow],
rows: &[PgRow],
limit: u32,
mapper: impl Fn(&SqliteRow) -> SqliteAccountStoreResult<T>,
mapper: impl Fn(&PgRow) -> PgAccountStoreResult<T>,
cursor_fields: (&str, &str),
) -> Page<T> {
let has_more = rows.len() > limit as usize;
let mut items: Vec<T> = 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<T>(
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::<chrono::DateTime<chrono::Utc>, _>(cursor_fields.0)
.to_rfc3339();
let id: String = row.get(cursor_fields.1);
crate::infra::json::encode_cursor(&ts, &id)
})
+660
View File
@@ -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<Option<StoredAccount>> {
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<Option<StoredAccount>> {
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<String>,
limit: u32,
) -> PgAccountStoreResult<Page<StoredAccount>> {
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<String>,
limit: u32,
) -> PgAccountStoreResult<Page<StoredAccountMetadata>> {
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<NumberedPage<StoredAccountMetadata>> {
let page_size = page_size.clamp(1, 200);
let search = search.map(str::trim).filter(|value| !value.is_empty());
let total = count_account_metadata(&self.pool, search).await?;
let offset = page_offset(page, page_size);
let mut builder = QueryBuilder::<Postgres>::new(LIST_ACCOUNT_METADATA_SELECT_SQL);
push_account_metadata_search(&mut builder, search);
builder.push(" order by added_at desc, id desc limit ");
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::<PgAccountStoreResult<Vec<_>>>()?;
Ok(NumberedPage {
items,
total,
page: page.max(1),
page_size,
})
}
/// 读取单个账号元数据(不含 token)。
pub async fn get_metadata(
&self,
account_id: &str,
) -> PgAccountStoreResult<Option<StoredAccountMetadata>> {
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<bool> {
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<bool> {
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<bool> {
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>(&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(&quota_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(&quota_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<DateTime<Utc>>,
) -> PgAccountStoreResult<bool> {
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<Utc>,
) -> PgAccountStoreResult<bool> {
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<bool> {
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<Utc>,
) -> PgAccountStoreResult<bool> {
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<bool> {
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<Vec<AccountQuotaSnapshot>> {
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<Option<String>> {
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::<Option<sqlx::types::Json<Value>>, _>("quota_json"))
.map(|value| value.0.to_string()))
}
/// 更新配额 JSON。
pub async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> PgAccountStoreResult<bool> {
let now = Utc::now();
let plan_type = quota_plan_type(quota_json);
let quota_json = sqlx::types::Json(serde_json::from_str::<Value>(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<DateTime<Utc>>,
) -> PgAccountStoreResult<bool> {
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::<Value>(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<DateTime<Utc>>,
quota_verify_required: bool,
) -> PgAccountStoreResult<bool> {
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::<Value>)
.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<Utc>,
limit_window_seconds: Option<u64>,
) -> 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::<Option<DateTime<Utc>>, _>("window_reset_at");
let existing_limit_window_seconds =
optional_positive_i64_to_u64(row.get::<Option<i64>, _>("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<bool> {
let result = sqlx::query(DELETE_ACCOUNT_SQL)
.bind(account_id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
}
-5
View File
@@ -1,5 +0,0 @@
//! 管理端账号 HTTP 模块。
pub(crate) mod quota_view;
pub mod routes;
pub mod service;
File diff suppressed because it is too large Load Diff
-4
View File
@@ -1,4 +0,0 @@
//! 管理端认证与本地会话模块。
pub mod service;
pub mod session;
-341
View File
@@ -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<Utc>,
}
/// 管理员会话错误。
#[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<Mutex<HashMap<String, LoginFailureState>>>,
}
#[derive(Debug, Clone)]
struct LoginFailureState {
failures: u32,
first_failed_at: DateTime<Utc>,
locked_until: Option<DateTime<Utc>>,
}
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<bool, AdminSessionError> {
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<bool, AdminSessionError> {
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<Option<AdminLoginSession>, 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<bool, AdminSessionError> {
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<String, LoginFailureState>,
now: DateTime<Utc>,
) {
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<T> = Result<T, sqlx::Error>;
/// 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<bool> {
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<Option<StoredAdminUser>> {
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<Utc>,
) -> 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<bool> {
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<Utc>,
) -> SqliteAdminSessionStoreResult<u64> {
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<bool> {
let result = sqlx::query("delete from admin_sessions where id = ?")
.bind(session_id)
.execute(&self.pool)
.await?;
Ok(result.rows_affected() > 0)
}
}
-4
View File
@@ -1,4 +0,0 @@
//! 管理端 v1 接口访问 Key 模块。
pub mod routes;
pub mod service;
-698
View File
@@ -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<T> = Result<T, ClientKeyStoreError>;
/// 提供客户端 API Key 验证能力的端口。
#[async_trait]
pub trait ClientKeyStore: Send + Sync + 'static {
/// 验证明文客户端 API Key,并在成功时记录使用时间。
async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult<bool>;
}
const CLIENT_KEY_LAST_USED_FLUSH_DELAY: Duration = Duration::from_secs(1);
// ---------------------------------------------------------------------------
// 业务服务
// ---------------------------------------------------------------------------
/// 客户端 API Key 服务。
#[derive(Clone)]
pub struct ClientKeyService {
store: Arc<RuntimeClientKeyStore>,
}
impl ClientKeyService {
/// 构造服务。
pub fn new(store: Arc<RuntimeClientKeyStore>) -> Self {
Self { store }
}
/// 验证客户端 API Key。
pub async fn verify(&self, plaintext: &str) -> ClientKeyStoreResult<bool> {
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<RuntimeClientKeyStore>,
}
impl AdminClientKeyService {
/// 构造服务。
pub fn new(store: SqliteClientKeyStore, runtime: Arc<RuntimeClientKeyStore>) -> Self {
Self { store, runtime }
}
/// 创建客户端 API Key。
pub async fn create(&self, name: &str) -> Result<AdminClientApiKey, AdminClientKeyError> {
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<String>,
limit: u32,
) -> Result<Page<AdminClientApiKey>, 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<Option<AdminClientApiKey>, 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<String>,
) -> Result<Option<AdminClientApiKey>, 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<bool, AdminClientKeyError> {
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<String>,
) -> Result<BatchDeleteClientApiKeys, AdminClientKeyError> {
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<String>,
pub prefix: String,
pub key: String,
pub enabled: bool,
pub created_at: String,
pub last_used_at: Option<String>,
}
/// 客户端 API Key 批量删除结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BatchDeleteClientApiKeys {
pub deleted: u32,
pub not_found: Vec<String>,
}
/// 管理端客户端 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<bool, AdminClientKeyError> {
match status.trim().to_ascii_lowercase().as_str() {
"active" => Ok(true),
"disabled" => Ok(false),
other => Err(AdminClientKeyError::InvalidStatus(other.to_string())),
}
}
impl From<StoredClientApiKey> 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<String>,
/// 明文 API Key 的短前缀。
pub prefix: String,
/// 管理端可复制的完整 API Key。
pub key: String,
/// 是否允许用于 `/v1` 认证。
pub enabled: bool,
/// 创建时间。
pub created_at: String,
/// 最近一次成功使用时间。
pub last_used_at: Option<String>,
}
/// 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<Vec<StoredClientApiKey>, 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<StoredClientApiKey, SqliteClientKeyStoreError> {
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<String>,
limit: u32,
) -> Result<Page<StoredClientApiKey>, 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::<Vec<_>>();
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<Option<StoredClientApiKey>, 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<bool, SqliteClientKeyStoreError> {
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<String>,
) -> Result<Option<StoredClientApiKey>, 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<bool, SqliteClientKeyStoreError> {
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<bool, SqliteClientKeyStoreError> {
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::<i64, _>("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<bool> {
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<RwLock<RuntimeClientKeyCache>>,
pending_last_used: Arc<Mutex<BTreeMap<String, String>>>,
flush_scheduled: Arc<AtomicBool>,
flush_delay: Duration,
}
#[derive(Default)]
struct RuntimeClientKeyCache {
keys_by_plaintext: HashMap<String, String>,
plaintext_by_id: HashMap<String, String>,
}
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<String> {
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<String, String> {
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(),
}
}
-19
View File
@@ -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;
-14
View File
@@ -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;
@@ -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);
}
}
@@ -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<T> = Result<T, SqliteOpsErrorLogStoreError>;
/// SQLite 运维错误事件存储。
#[derive(Clone)]
pub struct SqliteOpsErrorLogStore {
pool: SqlitePool,
}
impl SqliteOpsErrorLogStore {
/// 构造存储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 写入运维错误事件。
pub async fn append(&self, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> {
append_event(&self.pool, event).await
}
/// 按保留期清理运维错误事件。
pub async fn trim_to_retention(&self, now: DateTime<Utc>) -> SqliteOpsErrorLogStoreResult<u64> {
trim_to_retention(&self.pool, now).await
}
/// 清空运维错误事件。
pub async fn clear(&self) -> SqliteOpsErrorLogStoreResult<u64> {
let result = sqlx::query("delete from ops_error_logs")
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
/// 读取错误事件数量。
pub async fn count(&self) -> SqliteOpsErrorLogStoreResult<u64> {
let (total,): (i64,) = sqlx::query_as("select count(*) from ops_error_logs")
.fetch_one(&self.pool)
.await?;
Ok(total.max(0) as u64)
}
}
#[derive(Debug, Clone)]
struct OpsErrorSummaryFields {
client_status_code: Option<i64>,
upstream_status_code: Option<i64>,
transport: Option<String>,
attempt_index: Option<i64>,
failure_class: Option<String>,
response_id: Option<String>,
upstream_request_id: Option<String>,
}
fn ops_error_summary_fields(event: &OpsErrorLog) -> OpsErrorSummaryFields {
OpsErrorSummaryFields {
client_status_code: event.client_status_code.or_else(|| {
metadata_i64(
&event.metadata,
&["clientStatusCode", "client_status_code", "clientStatus"],
)
}),
upstream_status_code: event.upstream_status_code.or_else(|| {
metadata_i64(
&event.metadata,
&[
"upstreamStatus",
"upstreamStatusCode",
"upstream_status_code",
],
)
}),
transport: event
.transport
.clone()
.or_else(|| metadata_string(&event.metadata, &["transport"])),
attempt_index: event
.attempt_index
.or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"])),
failure_class: event
.failure_class
.clone()
.or_else(|| metadata_string(&event.metadata, &["failureClass", "failure_class"])),
response_id: event
.response_id
.clone()
.or_else(|| metadata_string(&event.metadata, &["responseId", "response_id"])),
upstream_request_id: event.upstream_request_id.clone().or_else(|| {
metadata_string(
&event.metadata,
&[
"upstreamRequestId",
"upstream_request_id",
"openaiRequestId",
],
)
}),
}
}
async fn append_event(pool: &SqlitePool, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> {
let summary = ops_error_summary_fields(event);
sqlx::query(
"insert into ops_error_logs (id, request_id, kind, account_id, route, model, status_code, client_status_code, upstream_status_code, transport, attempt_index, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(&event.id)
.bind(&event.request_id)
.bind(&event.kind)
.bind(&event.account_id)
.bind(&event.route)
.bind(&event.model)
.bind(event.status_code)
.bind(summary.client_status_code)
.bind(summary.upstream_status_code)
.bind(summary.transport)
.bind(summary.attempt_index)
.bind(summary.failure_class)
.bind(summary.response_id)
.bind(summary.upstream_request_id)
.bind(event.latency_ms)
.bind(&event.message)
.bind(event.metadata.to_string())
.bind(event.created_at.to_rfc3339())
.execute(pool)
.await?;
append_error_time_bucket(pool, event).await?;
Ok(())
}
async fn append_error_time_bucket(
pool: &SqlitePool,
event: &OpsErrorLog,
) -> SqliteOpsErrorLogStoreResult<()> {
let bucket_start = china_quarter_hour_start(event.created_at);
let input_tokens = metadata_usage_i64(&event.metadata, "inputTokens");
let output_tokens = metadata_usage_i64(&event.metadata, "outputTokens");
let cached_tokens = metadata_usage_i64(&event.metadata, "cachedTokens");
let first_token_latency = metadata_first_token_latency(&event.metadata).unwrap_or(0);
let first_token_latency_count = i64::from(first_token_latency > 0);
let latency = event.latency_ms.filter(|value| *value > 0).unwrap_or(0);
let latency_count = i64::from(latency > 0);
let now = Utc::now().to_rfc3339();
sqlx::query(
r"
insert into usage_time_buckets (
bucket_start,
account_id,
model,
service_tier,
request_count,
error_count,
input_tokens,
output_tokens,
cached_tokens,
first_token_latency_sum,
first_token_latency_count,
latency_sum,
latency_count,
max_latency_ms,
min_latency_ms,
updated_at
) values (?, ?, ?, ?, 1, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
on conflict(bucket_start, account_id, model, service_tier) do update set
request_count = request_count + excluded.request_count,
error_count = error_count + excluded.error_count,
input_tokens = input_tokens + excluded.input_tokens,
output_tokens = output_tokens + excluded.output_tokens,
cached_tokens = cached_tokens + excluded.cached_tokens,
first_token_latency_sum = first_token_latency_sum + excluded.first_token_latency_sum,
first_token_latency_count = first_token_latency_count + excluded.first_token_latency_count,
latency_sum = latency_sum + excluded.latency_sum,
latency_count = latency_count + excluded.latency_count,
max_latency_ms = max(max_latency_ms, excluded.max_latency_ms),
min_latency_ms = case
when min_latency_ms = 0 then excluded.min_latency_ms
when excluded.min_latency_ms = 0 then min_latency_ms
else min(min_latency_ms, excluded.min_latency_ms)
end,
updated_at = excluded.updated_at",
)
.bind(bucket_start.to_rfc3339())
.bind(event.account_id.as_deref().unwrap_or_default())
.bind(event.model.as_deref().unwrap_or_default())
.bind(metadata_service_tier(&event.metadata).unwrap_or_default())
.bind(input_tokens)
.bind(output_tokens)
.bind(cached_tokens)
.bind(first_token_latency)
.bind(first_token_latency_count)
.bind(latency)
.bind(latency_count)
.bind(latency)
.bind(latency)
.bind(now)
.execute(pool)
.await?;
Ok(())
}
fn metadata_usage_i64(metadata: &Value, field: &str) -> i64 {
metadata
.get("usage")
.and_then(|usage| usage.get(field))
.or_else(|| metadata.get(field))
.and_then(Value::as_i64)
.filter(|value| *value > 0)
.unwrap_or(0)
}
fn metadata_first_token_latency(metadata: &Value) -> Option<i64> {
[
"firstTokenMs",
"first_token_ms",
"firstTokenLatencyMs",
"first_token_latency_ms",
]
.into_iter()
.find_map(|field| {
metadata
.get(field)
.or_else(|| metadata.get("usage").and_then(|usage| usage.get(field)))
.and_then(Value::as_i64)
.filter(|value| *value > 0)
})
}
async fn trim_to_retention(
pool: &SqlitePool,
now: DateTime<Utc>,
) -> SqliteOpsErrorLogStoreResult<u64> {
let cutoff = now - Duration::days(USAGE_RECORD_RETENTION_DAYS);
let result = sqlx::query("delete from ops_error_logs where created_at < ?")
.bind(cutoff.to_rfc3339())
.execute(pool)
.await?;
Ok(result.rows_affected())
}
File diff suppressed because it is too large Load Diff
-3
View File
@@ -1,3 +0,0 @@
//! 管理端设置 HTTP 模块。
pub mod routes;
-5
View File
@@ -1,5 +0,0 @@
//! 系统版本与在线更新处理器。
pub mod routes;
mod state;
mod updater;
-5
View File
@@ -1,5 +0,0 @@
//! 管理端系统路由入口。
pub(crate) use super::updater::{
perform_update, restart, rollback, update_detail, update_event_stream, update_status, version,
};
+606
View File
@@ -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<String>,
limit: Option<u32>,
page: Option<u32>,
page_size: Option<u32>,
search: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct CreateAccountRequest {
token: Option<String>,
refresh_token: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct BatchDeleteAccountsRequest {
ids: Vec<String>,
}
#[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<String>,
confirm: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct AccountTestQuery {
id: String,
model_id: Option<String>,
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct AccountOAuthExchangeRequest {
session_id: String,
callback_url: Option<String>,
code: Option<String>,
state: Option<String>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountModelsData {
models: Vec<AccountModelData>,
}
#[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<String>,
account_id: Option<String>,
user_id: Option<String>,
label: Option<String>,
plan_type: Option<String>,
has_refresh_token: bool,
status: String,
display_status: String,
token_refreshing: bool,
access_token_expires_at: Option<String>,
access_token_expires_at_display: Option<String>,
added_at: String,
added_at_display: String,
updated_at: String,
updated_at_display: String,
quota: AccountQuotaData,
usage: AdminAccountUsageData,
}
impl From<ManagedAccount> 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<AccountQuotaData>,
models: Vec<AdminAccountModelUsageData>,
token_refreshing: bool,
) -> Self {
let access_token_expires_at = a.access_token_expires_at.as_ref().map(china_rfc3339);
let access_token_expires_at_display =
a.access_token_expires_at.as_ref().map(china_datetime);
let display_status = account_display_status(a.status, token_refreshing).to_string();
Self {
id: a.id,
email: a.email,
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<String>,
last_used_at_display: String,
models: Vec<AdminAccountModelUsageData>,
}
impl AdminAccountUsageData {
fn from_usage(
usage: Option<&AccountUsageRecord>,
models: Vec<AdminAccountModelUsageData>,
) -> 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<String>,
last_used_at_display: String,
}
#[derive(Debug, Clone, Default)]
struct AccountListStats {
usage_by_account: HashMap<String, AccountUsageRecord>,
quota_by_account: HashMap<String, AccountQuotaData>,
models_by_account: HashMap<String, Vec<AdminAccountModelUsageData>>,
refreshing_account_ids: HashSet<String>,
}
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<DateTime<Utc>>,
}
#[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<AdminAccountData>,
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<AdminAccountData>, 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<AdminAccountData>,
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<AdminAccountData>),
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AccountRefreshData {
id: String,
email: Option<String>,
previous_status: String,
result: &'static str,
error: Option<String>,
duration_ms: i64,
}
impl From<AccountRefreshResult> 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<String>,
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<Vec<String>>,
stagger_ms: Option<u64>,
concurrency: Option<usize>,
}
#[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<AccountRefreshData>,
}
impl From<AccountHealthCheck> 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<AppState>,
_auth: AdminAuth,
Query(params): Query<AccountsQuery>,
) -> Result<impl IntoResponse, AdminError> {
let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50));
let use_numbered_page = params.page.is_some() || params.page_size.is_some();
let quota_by_account = quota_snapshots_by_account(&state).await;
let summary = account_summary_data(&state, &quota_by_account).await;
if use_numbered_page {
return match state
.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, &quota_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, &quota_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<AppState>,
_auth: AdminAuth,
Json(payload): Json<CreateAccountRequest>,
) -> Result<impl IntoResponse, AdminError> {
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::*;
@@ -0,0 +1,312 @@
use super::*;
/// `GET /api/admin/accounts/export`
pub(crate) async fn export_accounts(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<AccountExportQuery>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountActionRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountHealthCheckRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Query(query): Query<AccountIdQuery>,
) -> Result<impl IntoResponse, AdminError> {
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(&quota_json, Some(Utc::now()));
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
let account =
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(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<AdminAccountData, AdminError> {
let Some(account) = state
.services
.admin_accounts
.get(account_id)
.await
.map_err(|error| account_error(&error))?
else {
return Err(account_not_found());
};
let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]);
let stats =
account_list_stats(state, std::slice::from_ref(&account), &quota_by_account).await?;
Ok(stats.data_for(account))
}
/// `POST /api/admin/accounts/import`
pub(crate) async fn import_accounts(
State(state): State<AppState>,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<AccountOAuthExchangeRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
Query(query): Query<AccountTestQuery>,
_auth: AdminAuth,
) -> Result<Response, AdminError> {
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<AppState>,
Query(query): Query<AccountIdQuery>,
_auth: AdminAuth,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<BatchDeleteAccountsRequest>,
) -> Result<impl IntoResponse, AdminError> {
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<AppState>,
_auth: AdminAuth,
Json(payload): Json<Value>,
) -> Result<impl IntoResponse, AdminError> {
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)),
}
}
@@ -0,0 +1,512 @@
use super::*;
pub(super) async fn account_list_stats(
state: &AppState,
accounts: &[ManagedAccount],
quota_by_account: &HashMap<String, AccountQuotaData>,
) -> Result<AccountListStats, AdminError> {
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
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, &quota_by_account);
let quota_window_usage_stats =
quota_window_local_usage_by_account(state, &all_quota_windows).await;
apply_quota_window_local_usage(accounts, &mut quota_by_account, &quota_window_usage_stats);
let selected_quota_windows = selected_quota_windows_by_account(accounts, &quota_by_account);
let usage_records = usage_records
.into_iter()
.map(|usage| {
apply_selected_quota_window_usage(
usage,
&selected_quota_windows,
&quota_window_usage_stats,
)
})
.collect::<Vec<_>>();
let models_by_account = list_current_window_model_usage(state, &usage_records).await;
let refreshing_account_ids = state
.services
.token_refresh
.refreshing_account_ids(&account_ids, Utc::now())
.await
.map_err(|error| AdminError::internal(error.to_string()))?;
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<String, AccountQuotaData>,
) -> HashMap<String, AccountQuotaUsageWindow> {
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<AccountQuotaUsageWindow> {
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<String, AccountQuotaData>,
) -> Vec<AccountQuotaWindowSelection> {
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::<Vec<_>>();
let usage_by_account = quota_window_local_usage_by_account(state, &windows).await;
let empty_usage = HashMap::new();
quota.apply_local_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage));
}
pub(super) fn apply_quota_window_local_usage(
accounts: &[ManagedAccount],
quota_by_account: &mut HashMap<String, AccountQuotaData>,
usage_by_account: &HashMap<String, HashMap<String, AccountQuotaWindowLocalUsage>>,
) {
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<String, HashMap<String, AccountQuotaWindowLocalUsage>> {
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::<Vec<_>>();
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<String, AccountQuotaUsageWindow>,
quota_window_usage_stats: &HashMap<String, HashMap<String, AccountQuotaWindowLocalUsage>>,
) -> 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<String, AccountQuotaData>,
) -> 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<String, AccountQuotaData> {
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<ManagedAccount> {
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<String, Vec<AdminAccountModelUsageData>> {
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::<Vec<_>>();
if windows.is_empty() {
return HashMap::new();
}
let windows = windows
.into_iter()
.map(|(account_id, start, end)| ModelUsageWindow {
account_id,
start,
end,
})
.collect::<Vec<_>>();
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::<Vec<_>>();
models_by_account(records)
}
pub(super) fn current_usage_window(
usage: &AccountUsageRecord,
now: DateTime<Utc>,
) -> Option<(DateTime<Utc>, DateTime<Utc>)> {
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<AccountModelUsageRecord>,
) -> HashMap<String, Vec<AdminAccountModelUsageData>> {
let mut by_account = HashMap::<String, Vec<AccountModelUsageRecord>>::new();
for record in records {
by_account
.entry(record.account_id.clone())
.or_default()
.push(record);
}
by_account
.into_iter()
.map(|(account_id, mut 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<ParsedAccountUpdate, AdminError> {
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<String> {
value
.into_iter()
.flat_map(|ids| ids.split(','))
.map(str::trim)
.filter(|id| !id.is_empty())
.map(ToString::to_string)
.collect()
}
+388
View File
@@ -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<DashboardTrendKind>,
}
#[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<DashboardAccountUsageData>,
service_statuses: Vec<DashboardServiceStatusData>,
usage_records: Vec<DashboardUsageRecordData>,
pool_summary: AccountPoolDiagnostics,
capacity_info: AccountCapacityDiagnostics,
rotation_strategy: Option<String>,
}
#[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<f64>,
yesterday_hit_rate_value: Option<f64>,
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<DashboardTrendPointData>,
summary: Vec<DashboardTrendSummaryData>,
}
#[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<String>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardAccountUsageData {
id: String,
email: String,
plan_type: Option<String>,
tokens: String,
quota_used_percent: Option<f64>,
last_used: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardUsageRecordData {
id: String,
route: Option<String>,
model: String,
status_code: i64,
transport: Option<String>,
created_at_display: String,
account_email: Option<String>,
requested_model: Option<String>,
upstream_model: Option<String>,
client_ip: Option<String>,
user_agent: Option<String>,
reasoning_effort: Option<String>,
token_details: UsageRecordTokenDetailsData,
cost_details: Option<UsageRecordCostDetailsData>,
first_token_latency_ms_display: String,
latency_ms_display: String,
}
#[derive(Debug, Serialize)]
#[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<f64> {
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<u64> {
self.first_token_latency_sum
.checked_div(self.first_token_latency_count)
}
fn max_first_token_bucket_latency(self) -> Option<u64> {
(self.max_first_token_bucket_latency > 0).then_some(self.max_first_token_bucket_latency)
}
fn min_first_token_bucket_latency(self) -> Option<u64> {
(self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency)
}
}
/// `GET /api/admin/dashboard/summary`
pub(crate) async fn dashboard_summary(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
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::<Vec<_>>();
let refreshing_account_ids = state
.services
.token_refresh
.refreshing_account_ids(&account_ids, now)
.await
.map_err(|error| AdminError::internal(error.to_string()))?;
let pool_summary = account_pool_summary(&accounts, &refreshing_account_ids);
let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default());
let quota_used_by_account =
account_quota_used_percent_by_id(&state, &account_usage_records).await;
let settings = state.services.settings.current();
let recent_usage_records = recent_events;
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,
&quota_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<Utc>) -> 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<AppState>,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
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::*;
@@ -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<DashboardTrendKind>,
}
#[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<DashboardAccountUsageData>,
service_statuses: Vec<DashboardServiceStatusData>,
usage_records: Vec<DashboardUsageRecordData>,
pool_summary: AccountPoolDiagnostics,
capacity_info: AccountCapacityDiagnostics,
rotation_strategy: Option<String>,
}
#[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<f64>,
yesterday_hit_rate_value: Option<f64>,
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<DashboardTrendPointData>,
summary: Vec<DashboardTrendSummaryData>,
}
#[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<String>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardAccountUsageData {
id: String,
email: String,
plan_type: Option<String>,
tokens: String,
quota_used_percent: Option<f64>,
last_used: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct DashboardUsageRecordData {
id: String,
level: UsageRecordLevel,
route: Option<String>,
model: Option<String>,
status_code: Option<i64>,
transport: Option<String>,
created_at_display: String,
account_email: Option<String>,
requested_model: Option<String>,
upstream_model: Option<String>,
client_ip: Option<String>,
user_agent: Option<String>,
reasoning_effort: Option<String>,
token_details: UsageRecordTokenDetailsData,
cost_details: Option<UsageRecordCostDetailsData>,
first_token_latency_ms_display: String,
latency_ms_display: String,
}
#[derive(Debug, Serialize)]
#[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<f64> {
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<u64> {
self.first_token_latency_sum
.checked_div(self.first_token_latency_count)
}
fn max_first_token_bucket_latency(self) -> Option<u64> {
(self.max_first_token_bucket_latency > 0).then_some(self.max_first_token_bucket_latency)
}
fn min_first_token_bucket_latency(self) -> Option<u64> {
(self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency)
}
}
/// `GET /api/admin/dashboard/summary`
pub(crate) async fn dashboard_summary(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
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::<Vec<_>>();
let refreshing_account_ids = state
.services
.token_refresh
.refreshing_account_ids(&account_ids, now)
.await
.map_err(|error| AdminError::internal(error.to_string()))?;
let pool_summary = account_pool_summary(&accounts, &refreshing_account_ids);
let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default());
let quota_used_by_account =
account_quota_used_percent_by_id(&state, &account_usage_records).await;
let settings = state.services.settings.current();
let recent_usage_records = recent_events;
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,
&quota_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<Utc>) -> AdminUsageRecordFilter {
AdminUsageRecordFilter {
start_time: Some(china_day_start(now)),
end_time: Some(now),
..AdminUsageRecordFilter::default()
}
}
/// `GET /api/admin/dashboard/trend?kind=usage|latency|errors`
pub(crate) async fn dashboard_trend(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
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<Utc>,
) -> DashboardCardsData {
@@ -462,7 +77,7 @@ fn dashboard_cards_at(
}
}
fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option<f64> {
pub(super) fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option<f64> {
let input_tokens = summary.input_tokens;
if input_tokens > 0 {
Some(summary.cached_tokens as f64 / input_tokens as f64)
@@ -471,14 +86,14 @@ fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option<f64> {
}
}
fn format_optional_duration_ms(value: Option<u64>) -> String {
pub(super) fn format_optional_duration_ms(value: Option<u64>) -> 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<Utc>,
) -> Vec<AdminUsageTimeBucketRecord> {
) -> Vec<AccountUsageTimeBucket> {
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<Utc>,
) -> 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<Utc>,
) -> 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<DashboardTrendSummaryData> {
@@ -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<Utc>,
end: DateTime<Utc>,
) -> UsageWindow {
@@ -775,8 +390,8 @@ fn usage_window(
window
}
fn cost_window(
records: &[AdminUsageTimeBucketRecord],
pub(super) fn cost_window(
records: &[AccountUsageTimeBucket],
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Option<f64> {
@@ -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<u64> {
pub(super) fn first_token_bucket_latency(record: &AccountUsageTimeBucket) -> Option<u64> {
let sum = nonnegative_i64_to_u64(record.first_token_latency_sum);
let count = nonnegative_i64_to_u64(record.first_token_latency_count);
sum.checked_div(count).filter(|latency| *latency > 0)
}
fn account_pool_summary(
pub(super) fn account_pool_summary(
accounts: &[Account],
refreshing_account_ids: &HashSet<String>,
) -> 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<String, f64>,
@@ -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<String, f64> {
@@ -906,7 +521,7 @@ async fn account_quota_used_percent_by_id(
quota_used_by_account
}
fn quota_used_percent(quota_json: &str) -> Option<f64> {
pub(super) fn quota_used_percent(quota_json: &str) -> Option<f64> {
let quota = serde_json::from_str::<Value>(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<f64> {
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<f64>,
@@ -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<f64> {
pub(super) fn percent_value(value: &Value) -> Option<f64> {
let percent = value
.as_f64()
.or_else(|| value.as_str().and_then(|value| value.parse::<f64>().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<UsageRecord>,
account_emails: &HashMap<String, String>,
) -> Vec<DashboardUsageRecordData> {
@@ -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<String, String>,
) -> 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<DashboardServiceStatusData> {
pub(super) fn service_status_data(state: &AppState) -> Vec<DashboardServiceStatusData> {
let current_fingerprint = state.services.fingerprint.snapshot();
let fingerprint = fingerprint_diagnostics(&current_fingerprint);
vec![
@@ -1081,7 +688,7 @@ fn service_status_data(state: &AppState) -> Vec<DashboardServiceStatusData> {
]
}
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>) -> String {
pub(super) fn format_fingerprint_updated_at(value: Option<String>) -> String {
let Some(value) = value else {
return "-".to_string();
};
@@ -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<AdminClientApiKey> for ClientApiKeyData {
fn from(k: AdminClientApiKey) -> Self {
impl From<ManagedClientApiKey> for ClientApiKeyData {
fn from(k: ManagedClientApiKey) -> Self {
Self {
id: k.id,
name: k.name,
@@ -74,12 +74,12 @@ impl From<AdminClientApiKey> 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()),
}
}
+13
View File
@@ -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;
+155
View File
@@ -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<String>,
limit: Option<u32>,
page: Option<u32>,
page_size: Option<u32>,
kind: Option<String>,
client_api_key_id: Option<String>,
provider: Option<String>,
request_id: Option<String>,
account_id: Option<String>,
route: Option<String>,
model: Option<String>,
status_code: Option<i64>,
client_status_code: Option<i64>,
upstream_status_code: Option<i64>,
transport: Option<String>,
attempt_index: Option<i64>,
failure_class: Option<String>,
response_id: Option<String>,
upstream_request_id: Option<String>,
search: Option<String>,
start_time: Option<String>,
end_time: Option<String>,
}
#[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<AppState>,
_auth: AdminAuth,
Query(query): Query<OpsErrorsQuery>,
) -> Result<impl IntoResponse, AdminError> {
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<OpsErrorFilter, AdminError> {
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<OpsErrorLog>) -> Vec<OpsErrorData> {
items
.into_iter()
.map(|error| OpsErrorData {
created_at_display: china_datetime(&error.created_at),
error,
})
.collect()
}
fn optional_datetime(
value: Option<String>,
) -> Result<Option<chrono::DateTime<chrono::Utc>>, 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<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn log_error(error: OpsQueryError) -> AdminError {
AdminError::internal(error.to_string())
}
@@ -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<AppState> {
)
.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))
@@ -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")),
};
@@ -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<AdminSettingsPatch, AdminError> {
fn parse_settings_patch(body: &[u8]) -> Result<SettingsPatch, AdminError> {
let value = serde_json::from_slice::<Value>(body)
.map_err(|_| AdminError::malformed_json("Malformed JSON body"))?;
let object = value
@@ -157,15 +153,12 @@ fn parse_settings_patch(body: &[u8]) -> Result<AdminSettingsPatch, AdminError> {
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")
}
}
+151
View File
@@ -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<bool>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct UpdateRequest {
target_version: String,
}
pub(crate) async fn version(_auth: AdminAuth) -> Result<impl IntoResponse, AdminError> {
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(service::version_data().await),
))
}
pub(crate) async fn update_detail(
_auth: AdminAuth,
Query(query): Query<UpdateDetailQuery>,
) -> Result<impl IntoResponse, AdminError> {
Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(service::update_detail(query.refresh.unwrap_or(false)).await),
))
}
pub(crate) async fn update_event_stream(
_auth: AdminAuth,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, 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<Json<UpdateRequest>>,
) -> Result<impl IntoResponse, AdminError> {
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<impl IntoResponse, AdminError> {
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<impl IntoResponse, AdminError> {
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<impl IntoResponse, AdminError> {
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),
}
}
@@ -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)]
@@ -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<u32>,
page_size: Option<u32>,
kind: Option<String>,
level: Option<String>,
client_api_key_id: Option<String>,
provider: Option<String>,
request_id: Option<String>,
account_id: Option<String>,
route: Option<String>,
@@ -55,8 +50,6 @@ pub(crate) struct UsageRecordsQuery {
status_code: Option<i64>,
transport: Option<String>,
attempt_index: Option<i64>,
upstream_status_code: Option<i64>,
failure_class: Option<String>,
response_id: Option<String>,
upstream_request_id: Option<String>,
search: Option<String>,
@@ -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<AdminUsageRecordFilter, AdminError> {
Ok(AdminUsageRecordFilter {
fn filter_from_query(query: UsageRecordsQuery) -> Result<UsageQueryFilter, AdminError> {
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<AdminUsageRecordFilter,
status_code: query.status_code,
transport: non_empty(query.transport),
attempt_index: query.attempt_index,
upstream_status_code: query.upstream_status_code,
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),
@@ -457,29 +445,14 @@ fn optional_datetime(
.map_err(|_| AdminError::invalid_time_range("Invalid time range"))
}
fn log_error(error: &AdminUsageRecordError) -> 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<String>) -> Result<Option<UsageRecordLevel>, 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<String, String>,
) -> 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<String>, Option<String>) {
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<u64> {
value
.get(key)
.and_then(|value| value.as_u64().or_else(|| value.as_i64()?.try_into().ok()))
}
impl From<UsageRecordSummary> 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<UsageRecordBreakdown>) -> Vec<UsageRe
impl From<UsageRecordTrendPoint> 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<UsageRecordTrendPoint> 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),
@@ -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<String> {
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<String, ClientApiKeyAuthError> {
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) {
@@ -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 =
@@ -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;
@@ -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};
@@ -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
@@ -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},
@@ -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;
@@ -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<
+7
View File
@@ -0,0 +1,7 @@
//! 入站 HTTP API 与静态资源。
pub mod admin;
pub mod assets;
pub mod client;
pub mod middleware;
pub mod router;
@@ -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<AppState> {
pub fn router_with_assets(dist_dir: impl AsRef<Path>) -> Router<AppState> {
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<AppState>) -> 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
}
+5
View File
@@ -0,0 +1,5 @@
//! 管理员认证与 Redis 会话。
pub mod service;
pub mod store;
pub mod types;
+192
View File
@@ -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<Mutex<HashMap<String, LoginFailureState>>>,
}
#[derive(Debug, Clone)]
struct LoginFailureState {
failures: u32,
first_failed_at: DateTime<Utc>,
locked_until: Option<DateTime<Utc>>,
}
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<bool, SessionError> {
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<bool, SessionError> {
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<Option<LoginSession>, 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<bool, SessionError> {
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<String, LoginFailureState>,
now: DateTime<Utc>,
) {
failures.retain(|_, state| {
state
.locked_until
.is_some_and(|locked_until| locked_until > now)
|| now - state.first_failed_at <= LOGIN_FAILURE_WINDOW
});
}
+137
View File
@@ -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<T> = Result<T, AuthStoreError>;
/// 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<bool> {
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<Option<StoredAdminUser>> {
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<Utc>,
}
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<bool> {
let mut connection = self.redis.manager();
let value: Option<String> = 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<bool> {
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)))
}
}
+42
View File
@@ -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<Utc>,
}
/// 管理员会话错误。
#[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,
}
@@ -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<dyn Error + Send + Sync>> {
}
async fn build_application(
config: AppConfig,
mut config: AppConfig,
) -> Result<(Router, TaskCoordinator), Box<dyn Error + Send + Sync>> {
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,
))
}
@@ -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<Self, ::config::ConfigError> {
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<Path>) -> Result<Self, ::config::ConfigError> {
load_file(config_dir.as_ref().join("config.yaml"))
}
}
fn load_file(config_file: impl AsRef<Path>) -> Result<AppConfig, ::config::ConfigError> {
::config::Config::builder()
.add_source(::config::File::from(config_file.as_ref()).required(true))
.build()?
.try_deserialize()
}
+233
View File
@@ -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<i64> },
#[error("target PostgreSQL database is not empty")]
TargetNotEmpty,
#[error("invalid timestamp in {table}.{column}: {value}")]
InvalidTimestamp {
table: &'static str,
column: &'static str,
value: String,
},
#[error("invalid JSON in {table}.{column}: {source}")]
InvalidJson {
table: &'static str,
column: &'static str,
#[source]
source: serde_json::Error,
},
}
/// 将只读 SQLite v3 历史库一次性导入空的 PostgreSQL 终态库。
pub async fn import_sqlite(
target: &PgPool,
source_path: impl AsRef<Path>,
) -> Result<ImportSqliteReport, ImportSqliteError> {
let options = SqliteConnectOptions::new()
.filename(source_path.as_ref())
.read_only(true)
.create_if_missing(false);
let source = SqlitePoolOptions::new()
.max_connections(1)
.connect_with(options)
.await?;
validate_source_version(&source).await?;
let mut transaction = target.begin().await?;
reject_nonempty_target(&mut transaction).await?;
let mut report = ImportSqliteReport::default();
core::import_core_tables(&source, &mut transaction, &mut report).await?;
telemetry::import_telemetry_tables(&source, &mut transaction, &mut report).await?;
core::count_discarded_runtime_tables(&source, &mut report).await?;
transaction.commit().await?;
source.close().await;
Ok(report)
}
async fn validate_source_version(source: &SqlitePool) -> Result<(), ImportSqliteError> {
let (actual,): (Option<i64>,) = sqlx::query_as("select max(version) from schema_migrations")
.fetch_one(source)
.await?;
if actual != Some(SQLITE_SCHEMA_VERSION) {
return Err(ImportSqliteError::UnsupportedSourceVersion { actual });
}
Ok(())
}
async fn reject_nonempty_target(
transaction: &mut Transaction<'_, Postgres>,
) -> Result<(), ImportSqliteError> {
let (rows,): (i64,) = sqlx::query_as(
"select
(select count(*) from admin_users)
+ (select count(*) from client_api_keys)
+ (select count(*) from runtime_settings)
+ (select count(*) from accounts)
+ (select count(*) from account_usage)
+ (select count(*) from 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<DateTime<Utc>, ImportSqliteError> {
DateTime::parse_from_rfc3339(&value)
.map(|timestamp| timestamp.with_timezone(&Utc))
.map_err(|_| ImportSqliteError::InvalidTimestamp {
table,
column,
value,
})
}
fn parse_optional_timestamp(
table: &'static str,
column: &'static str,
value: Option<String>,
) -> Result<Option<DateTime<Utc>>, ImportSqliteError> {
value
.map(|value| parse_timestamp(table, column, value))
.transpose()
}
fn parse_cookie_expiration(value: &str) -> Option<DateTime<Utc>> {
DateTime::parse_from_rfc3339(value)
.or_else(|_| DateTime::parse_from_rfc2822(value))
.ok()
.map(|timestamp| timestamp.with_timezone(&Utc))
}
fn parse_json(
table: &'static str,
column: &'static str,
value: &str,
) -> Result<Value, ImportSqliteError> {
serde_json::from_str(value).map_err(|source| ImportSqliteError::InvalidJson {
table,
column,
source,
})
}
fn parse_optional_json(
table: &'static str,
column: &'static str,
value: Option<String>,
) -> Result<Option<Value>, ImportSqliteError> {
value
.map(|value| parse_json(table, column, &value))
.transpose()
}
fn normalized_dimension(value: Option<String>) -> String {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| "__unknown__".to_string())
}
+491
View File
@@ -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::<String, _>("id")?)
.bind(row.try_get::<String, _>("password_hash")?)
.bind(parse_timestamp(
"admin_users",
"created_at",
row.try_get("created_at")?,
)?)
.bind(parse_timestamp(
"admin_users",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("admin_users");
}
Ok(())
}
async fn import_client_api_keys(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, name, prefix, key, label, enabled, created_at, last_used_at
from client_api_keys order by id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
sqlx::query(
"insert into client_api_keys
(id, name, prefix, key, label, enabled, created_at, last_used_at)
values ($1, $2, $3, $4, $5, $6, $7, $8)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("name")?)
.bind(row.try_get::<String, _>("prefix")?)
.bind(row.try_get::<String, _>("key")?)
.bind(row.try_get::<Option<String>, _>("label")?)
.bind(row.try_get::<i64, _>("enabled")? == 1)
.bind(parse_timestamp(
"client_api_keys",
"created_at",
row.try_get("created_at")?,
)?)
.bind(parse_optional_timestamp(
"client_api_keys",
"last_used_at",
row.try_get("last_used_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("client_api_keys");
}
Ok(())
}
async fn import_runtime_settings(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, model_aliases_json, refresh_margin_seconds, refresh_concurrency,
max_concurrent_per_account, request_interval_ms, rotation_strategy,
admin_api_key, updated_at
from runtime_settings order by id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
let aliases: String = row.try_get("model_aliases_json")?;
let admin_key = row
.try_get::<Option<String>, _>("admin_api_key")?
.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::<i64, _>("id")?)
.bind(Json(parse_json(
"runtime_settings",
"model_aliases_json",
&aliases,
)?))
.bind(row.try_get::<i64, _>("refresh_margin_seconds")?)
.bind(row.try_get::<i64, _>("refresh_concurrency")?)
.bind(row.try_get::<i64, _>("max_concurrent_per_account")?)
.bind(row.try_get::<i64, _>("request_interval_ms")?)
.bind(row.try_get::<String, _>("rotation_strategy")?)
.bind(admin_key)
.bind(parse_timestamp(
"runtime_settings",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("runtime_settings");
}
Ok(())
}
async fn import_accounts(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query(
"select id, email, chatgpt_account_id, chatgpt_user_id, label, plan_type,
access_token, refresh_token, access_token_expires_at, next_refresh_at,
status, quota_json, quota_fetched_at, quota_limit_reached,
quota_verify_required, quota_cooldown_until, cloudflare_cooldown_until,
added_at, updated_at
from accounts order by id",
)
.fetch(source);
while let Some(row) = rows.try_next().await? {
let source_status = row.try_get::<String, _>("status")?;
let (status, next_refresh_at) = if source_status == "refreshing" {
report.add_normalized("accounts.refreshing_to_expired");
("expired", None)
} else {
(
source_status.as_str(),
parse_optional_timestamp(
"accounts",
"next_refresh_at",
row.try_get("next_refresh_at")?,
)?,
)
};
let quota = parse_optional_json(
"accounts",
"quota_json",
row.try_get::<Option<String>, _>("quota_json")?,
)?
.map(Json);
sqlx::query(
"insert into accounts
(id, email, chatgpt_account_id, chatgpt_user_id, label, plan_type,
access_token, refresh_token, access_token_expires_at, next_refresh_at,
status, quota_json, quota_fetched_at, quota_limit_reached,
quota_verify_required, quota_cooldown_until, cloudflare_cooldown_until,
added_at, updated_at)
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13,
$14, $15, $16, $17, $18, $19)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<Option<String>, _>("email")?)
.bind(row.try_get::<Option<String>, _>("chatgpt_account_id")?)
.bind(row.try_get::<Option<String>, _>("chatgpt_user_id")?)
.bind(row.try_get::<Option<String>, _>("label")?)
.bind(row.try_get::<Option<String>, _>("plan_type")?)
.bind(row.try_get::<String, _>("access_token")?)
.bind(row.try_get::<Option<String>, _>("refresh_token")?)
.bind(parse_optional_timestamp(
"accounts",
"access_token_expires_at",
row.try_get("access_token_expires_at")?,
)?)
.bind(next_refresh_at)
.bind(status)
.bind(quota)
.bind(parse_optional_timestamp(
"accounts",
"quota_fetched_at",
row.try_get("quota_fetched_at")?,
)?)
.bind(row.try_get::<i64, _>("quota_limit_reached")? == 1)
.bind(row.try_get::<i64, _>("quota_verify_required")? == 1)
.bind(parse_optional_timestamp(
"accounts",
"quota_cooldown_until",
row.try_get("quota_cooldown_until")?,
)?)
.bind(parse_optional_timestamp(
"accounts",
"cloudflare_cooldown_until",
row.try_get("cloudflare_cooldown_until")?,
)?)
.bind(parse_timestamp(
"accounts",
"added_at",
row.try_get("added_at")?,
)?)
.bind(parse_timestamp(
"accounts",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("accounts");
}
Ok(())
}
async fn import_account_usage(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query("select * from account_usage order by account_id").fetch(source);
while let Some(row) = rows.try_next().await? {
sqlx::query(
"insert into account_usage
(account_id, request_count, empty_response_count, input_tokens, output_tokens,
cached_tokens, reasoning_tokens, total_tokens, image_input_tokens,
image_output_tokens, image_request_count, image_request_failed_count,
window_request_count, window_input_tokens, window_output_tokens,
window_cached_tokens, window_image_input_tokens, window_image_output_tokens,
window_image_request_count, window_image_request_failed_count,
window_started_at, window_reset_at, limit_window_seconds, last_used_at)
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12,
$13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24)",
)
.bind(row.try_get::<String, _>("account_id")?)
.bind(row.try_get::<i64, _>("request_count")?)
.bind(row.try_get::<i64, _>("empty_response_count")?)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<i64, _>("cached_tokens")?)
.bind(row.try_get::<i64, _>("reasoning_tokens")?)
.bind(row.try_get::<i64, _>("total_tokens")?)
.bind(row.try_get::<i64, _>("image_input_tokens")?)
.bind(row.try_get::<i64, _>("image_output_tokens")?)
.bind(row.try_get::<i64, _>("image_request_count")?)
.bind(row.try_get::<i64, _>("image_request_failed_count")?)
.bind(row.try_get::<i64, _>("window_request_count")?)
.bind(row.try_get::<i64, _>("window_input_tokens")?)
.bind(row.try_get::<i64, _>("window_output_tokens")?)
.bind(row.try_get::<i64, _>("window_cached_tokens")?)
.bind(row.try_get::<i64, _>("window_image_input_tokens")?)
.bind(row.try_get::<i64, _>("window_image_output_tokens")?)
.bind(row.try_get::<i64, _>("window_image_request_count")?)
.bind(row.try_get::<i64, _>("window_image_request_failed_count")?)
.bind(parse_optional_timestamp(
"account_usage",
"window_started_at",
row.try_get("window_started_at")?,
)?)
.bind(parse_optional_timestamp(
"account_usage",
"window_reset_at",
row.try_get("window_reset_at")?,
)?)
.bind(row.try_get::<Option<i64>, _>("limit_window_seconds")?)
.bind(parse_optional_timestamp(
"account_usage",
"last_used_at",
row.try_get("last_used_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("account_usage");
}
Ok(())
}
async fn import_account_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::<String, _>("account_id")?)
.bind(row.try_get::<String, _>("model")?)
.bind(row.try_get::<i64, _>("request_count")?)
.bind(row.try_get::<i64, _>("error_count")?)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<i64, _>("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::<Option<String>, _>("expires_at")?
.and_then(|value| match parse_cookie_expiration(&value) {
Some(timestamp) => Some(timestamp),
None => {
report.add_cookie_expiration_parse_failure();
None
}
});
sqlx::query(
"insert into account_cookies
(id, account_id, domain, name, value, path, expires_at, updated_at)
values ($1, $2, $3, $4, $5, $6, $7, $8)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("account_id")?)
.bind(row.try_get::<String, _>("domain")?)
.bind(row.try_get::<String, _>("name")?)
.bind(row.try_get::<String, _>("value")?)
.bind(row.try_get::<String, _>("path")?)
.bind(expires_at)
.bind(parse_timestamp(
"account_cookies",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("account_cookies");
}
Ok(())
}
async fn import_fingerprints(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows = sqlx::query("select * from fingerprints order by id").fetch(source);
while let Some(row) = rows.try_next().await? {
let default_headers: String = row.try_get("default_headers_json")?;
let header_order: String = row.try_get("header_order_json")?;
sqlx::query(
"insert into fingerprints
(id, originator, app_version, build_number, platform, arch, chromium_version,
user_agent_template, default_headers_json, header_order_json, source,
created_at, updated_at)
values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("originator")?)
.bind(row.try_get::<String, _>("app_version")?)
.bind(row.try_get::<String, _>("build_number")?)
.bind(row.try_get::<String, _>("platform")?)
.bind(row.try_get::<String, _>("arch")?)
.bind(row.try_get::<String, _>("chromium_version")?)
.bind(row.try_get::<String, _>("user_agent_template")?)
.bind(Json(parse_json(
"fingerprints",
"default_headers_json",
&default_headers,
)?))
.bind(Json(parse_json(
"fingerprints",
"header_order_json",
&header_order,
)?))
.bind(row.try_get::<String, _>("source")?)
.bind(parse_timestamp(
"fingerprints",
"created_at",
row.try_get("created_at")?,
)?)
.bind(parse_timestamp(
"fingerprints",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("fingerprints");
}
Ok(())
}
async fn import_fingerprint_history(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows =
sqlx::query("select * from fingerprint_update_history order by id").fetch(source);
while let Some(row) = rows.try_next().await? {
let manifest = parse_optional_json(
"fingerprint_update_history",
"manifest_json",
row.try_get::<Option<String>, _>("manifest_json")?,
)?
.map(Json);
sqlx::query(
"insert into fingerprint_update_history
(id, current_fingerprint_id, app_version, build_number, chromium_version,
source, manifest_json, created_at)
values ($1, $2, $3, $4, $5, $6, $7, $8)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<String, _>("current_fingerprint_id")?)
.bind(row.try_get::<String, _>("app_version")?)
.bind(row.try_get::<String, _>("build_number")?)
.bind(row.try_get::<Option<String>, _>("chromium_version")?)
.bind(row.try_get::<String, _>("source")?)
.bind(manifest)
.bind(parse_timestamp(
"fingerprint_update_history",
"created_at",
row.try_get("created_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("fingerprint_update_history");
}
Ok(())
}
pub(super) async fn count_discarded_runtime_tables(
source: &SqlitePool,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let (sessions, affinities, leases, snapshots): (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(())
}
@@ -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<String> = row.try_get("account_id")?;
let model: Option<String> = row.try_get("model")?;
let status_code: Option<i64> = row.try_get("status_code")?;
let metadata_text: String = row.try_get("metadata_json")?;
let metadata = parse_json("usage_records", "metadata_json", &metadata_text)?;
if level == "error" {
insert_usage_error(target, &row, metadata).await?;
report.add_imported("ops_error_logs");
continue;
}
let valid_account = account_id
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty());
let valid_model = model
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty());
let valid_status = status_code.is_some_and(|status| (200..=399).contains(&status));
if !valid_account || !valid_model || !valid_status {
report.add_discarded("usage_record_noise", 1);
continue;
}
insert_success_record(
target,
&row,
account_id.expect("validated account id"),
model.expect("validated model"),
status_code.expect("validated status code"),
metadata,
)
.await?;
report.add_imported("usage_records");
}
Ok(())
}
async fn insert_success_record(
target: &mut Transaction<'_, Postgres>,
row: &sqlx::sqlite::SqliteRow,
account_id: String,
model: String,
status_code: i64,
mut metadata: Value,
) -> Result<(), ImportSqliteError> {
let requested_model = metadata_string(&metadata, &["requestedModel", "requested_model"]);
let upstream_model = metadata_string(&metadata, &["upstreamModel", "upstream_model"]);
let service_tier = metadata_string(&metadata, &["serviceTier", "service_tier"]);
let first_token_ms = metadata_nonnegative(&metadata, &["firstTokenMs", "first_token_ms"]);
let input_tokens = usage_number(&metadata, &["inputTokens", "input_tokens"]);
let output_tokens = usage_number(&metadata, &["outputTokens", "output_tokens"]);
let cached_tokens = usage_number(&metadata, &["cachedTokens", "cached_tokens"]);
let reasoning_tokens = usage_number(&metadata, &["reasoningTokens", "reasoning_tokens"]);
strip_promoted_success_fields(&mut metadata);
sqlx::query(
"insert into usage_records
(id, request_id, client_api_key_id, kind, route, provider, account_id, model,
requested_model, upstream_model, service_tier, status_code, transport,
attempt_index, response_id, upstream_request_id, latency_ms, first_token_ms,
input_tokens, output_tokens, cached_tokens, reasoning_tokens, message,
metadata_json, created_at)
values ($1, $2, null, $3, $4, 'openai', $5, $6, $7, $8, $9, $10,
$11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<Option<String>, _>("request_id")?)
.bind(row.try_get::<String, _>("kind")?)
.bind(row.try_get::<Option<String>, _>("route")?)
.bind(account_id)
.bind(model)
.bind(requested_model)
.bind(upstream_model)
.bind(service_tier)
.bind(status_code as i32)
.bind(row.try_get::<Option<String>, _>("transport")?)
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
.bind(row.try_get::<Option<String>, _>("response_id")?)
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
.bind(first_token_ms)
.bind(input_tokens)
.bind(output_tokens)
.bind(cached_tokens)
.bind(reasoning_tokens)
.bind(row.try_get::<String, _>("message")?)
.bind(Json(metadata))
.bind(parse_timestamp(
"usage_records",
"created_at",
row.try_get("created_at")?,
)?)
.execute(&mut **target)
.await?;
Ok(())
}
async fn insert_usage_error(
target: &mut Transaction<'_, Postgres>,
row: &sqlx::sqlite::SqliteRow,
metadata: Value,
) -> Result<(), ImportSqliteError> {
sqlx::query(
"insert into ops_error_logs
(id, request_id, client_api_key_id, kind, provider, account_id, route, model,
status_code, client_status_code, upstream_status_code, transport, attempt_index,
failure_class, response_id, upstream_request_id, latency_ms, message,
metadata_json, created_at)
values ($1, $2, null, $3, 'openai', $4, $5, $6, $7, null, $8, $9,
$10, $11, $12, $13, $14, $15, $16, $17)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<Option<String>, _>("request_id")?)
.bind(row.try_get::<String, _>("kind")?)
.bind(row.try_get::<Option<String>, _>("account_id")?)
.bind(row.try_get::<Option<String>, _>("route")?)
.bind(row.try_get::<Option<String>, _>("model")?)
.bind(optional_status(row.try_get("status_code")?))
.bind(optional_status(row.try_get("upstream_status_code")?))
.bind(row.try_get::<Option<String>, _>("transport")?)
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
.bind(row.try_get::<Option<String>, _>("failure_class")?)
.bind(row.try_get::<Option<String>, _>("response_id")?)
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
.bind(row.try_get::<String, _>("message")?)
.bind(Json(metadata))
.bind(parse_timestamp(
"usage_records",
"created_at",
row.try_get("created_at")?,
)?)
.execute(&mut **target)
.await?;
Ok(())
}
async fn import_ops_error_logs(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows =
sqlx::query("select * from ops_error_logs order by created_at, id").fetch(source);
while let Some(row) = rows.try_next().await? {
let metadata_text: String = row.try_get("metadata_json")?;
sqlx::query(
"insert into ops_error_logs
(id, request_id, client_api_key_id, kind, provider, account_id, route, model,
status_code, client_status_code, upstream_status_code, transport, attempt_index,
failure_class, response_id, upstream_request_id, latency_ms, message,
metadata_json, created_at)
values ($1, $2, null, $3, 'openai', $4, $5, $6, $7, $8, $9, $10,
$11, $12, $13, $14, $15, $16, $17, $18)",
)
.bind(row.try_get::<String, _>("id")?)
.bind(row.try_get::<Option<String>, _>("request_id")?)
.bind(row.try_get::<String, _>("kind")?)
.bind(row.try_get::<Option<String>, _>("account_id")?)
.bind(row.try_get::<Option<String>, _>("route")?)
.bind(row.try_get::<Option<String>, _>("model")?)
.bind(optional_status(row.try_get("status_code")?))
.bind(optional_status(row.try_get("client_status_code")?))
.bind(optional_status(row.try_get("upstream_status_code")?))
.bind(row.try_get::<Option<String>, _>("transport")?)
.bind(row.try_get::<Option<i64>, _>("attempt_index")?)
.bind(row.try_get::<Option<String>, _>("failure_class")?)
.bind(row.try_get::<Option<String>, _>("response_id")?)
.bind(row.try_get::<Option<String>, _>("upstream_request_id")?)
.bind(row.try_get::<Option<i64>, _>("latency_ms")?)
.bind(row.try_get::<String, _>("message")?)
.bind(Json(parse_json(
"ops_error_logs",
"metadata_json",
&metadata_text,
)?))
.bind(parse_timestamp(
"ops_error_logs",
"created_at",
row.try_get("created_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("ops_error_logs");
}
Ok(())
}
async fn import_request_time_buckets(
source: &SqlitePool,
target: &mut Transaction<'_, Postgres>,
report: &mut ImportSqliteReport,
) -> Result<(), ImportSqliteError> {
let mut rows =
sqlx::query("select * from usage_time_buckets order by bucket_start").fetch(source);
while let Some(row) = rows.try_next().await? {
let request_count: i64 = row.try_get("request_count")?;
let error_count: i64 = row.try_get("error_count")?;
let min_latency_ms: i64 = row.try_get("min_latency_ms")?;
sqlx::query(
"insert into request_time_buckets
(bucket_start, provider, account_id, model, service_tier, success_count,
error_count, input_tokens, output_tokens, cached_tokens,
first_token_latency_sum, first_token_latency_count, latency_sum,
latency_count, max_latency_ms, min_latency_ms, updated_at)
values ($1, 'openai', $2, $3, $4, $5, $6, $7, $8, $9, $10, $11,
$12, $13, $14, $15, $16)",
)
.bind(parse_timestamp(
"usage_time_buckets",
"bucket_start",
row.try_get("bucket_start")?,
)?)
.bind(normalized_dimension(row.try_get("account_id")?))
.bind(normalized_dimension(row.try_get("model")?))
.bind(normalized_dimension(row.try_get("service_tier")?))
.bind(request_count.saturating_sub(error_count).max(0))
.bind(error_count)
.bind(row.try_get::<i64, _>("input_tokens")?)
.bind(row.try_get::<i64, _>("output_tokens")?)
.bind(row.try_get::<i64, _>("cached_tokens")?)
.bind(row.try_get::<i64, _>("first_token_latency_sum")?)
.bind(row.try_get::<i64, _>("first_token_latency_count")?)
.bind(row.try_get::<i64, _>("latency_sum")?)
.bind(row.try_get::<i64, _>("latency_count")?)
.bind(row.try_get::<i64, _>("max_latency_ms")?)
.bind((min_latency_ms != 0).then_some(min_latency_ms))
.bind(parse_timestamp(
"usage_time_buckets",
"updated_at",
row.try_get("updated_at")?,
)?)
.execute(&mut **target)
.await?;
report.add_imported("request_time_buckets");
}
Ok(())
}
fn optional_status(status: Option<i64>) -> Option<i32> {
status.and_then(|status| i32::try_from(status).ok())
}
fn metadata_string(metadata: &Value, keys: &[&str]) -> Option<String> {
keys.iter().find_map(|key| {
metadata
.get(*key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
})
}
fn metadata_nonnegative(metadata: &Value, keys: &[&str]) -> Option<i64> {
keys.iter()
.find_map(|key| metadata.get(*key).and_then(Value::as_i64))
.filter(|value| *value >= 0)
}
fn usage_number(metadata: &Value, keys: &[&str]) -> Option<i64> {
metadata
.get("usage")
.and_then(|usage| metadata_nonnegative(usage, keys))
.or_else(|| metadata_nonnegative(metadata, keys))
}
fn strip_promoted_success_fields(metadata: &mut Value) {
let Some(object) = metadata.as_object_mut() else {
return;
};
for key in [
"usage",
"route",
"apiKind",
"requestedModel",
"requested_model",
"upstreamModel",
"upstream_model",
"serviceTier",
"service_tier",
"statusCode",
"transport",
"attemptIndex",
"responseId",
"upstreamRequestId",
"openaiRequestId",
"firstTokenMs",
"first_token_ms",
] {
object.remove(key);
}
}
+9
View File
@@ -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;

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