mirror of
https://github.com/zyycn/codex-proxy-rs.git
synced 2026-11-06 08:06:19 +00:00
refactor: complete postgresql redis migration
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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/ 架构和维护文档
|
||||
|
||||
Generated
+99
@@ -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
@@ -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"
|
||||
|
||||
Executable
+59
@@ -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"
|
||||
@@ -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)]
|
||||
+2
-2
@@ -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")
|
||||
+9
-9
@@ -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);
|
||||
}
|
||||
+53
-63
@@ -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)))
|
||||
}
|
||||
}
|
||||
+82
-91
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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::*;
|
||||
+24
-22
@@ -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> {
|
||||
+17
-17
@@ -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;
|
||||
+28
-40
@@ -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("a);
|
||||
|
||||
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,
|
||||
+46
-60
@@ -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,
|
||||
+27
-29
@@ -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(),
|
||||
)),
|
||||
}
|
||||
@@ -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>("a_json).ok(),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id = %account_id,
|
||||
error = %error,
|
||||
"failed to read existing quota json before passive rate-limit sync"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref());
|
||||
self.apply_quota_snapshot(account_id, "a).await;
|
||||
}
|
||||
|
||||
/// 读取运行时账号快照。
|
||||
pub async fn account_snapshot(&self, account_id: &str) -> Option<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, "a_json, limit_reached, cooldown_until)
|
||||
.await
|
||||
{
|
||||
Ok(persisted) => persisted,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"failed to persist verified quota snapshot"
|
||||
);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(reset_at) = reset_at {
|
||||
if let Err(error) = self
|
||||
.store
|
||||
.sync_rate_limit_window(account_id, reset_at, limit_window_seconds)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"failed to persist verified quota window"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let in_memory = {
|
||||
let mut pool = self.pool.lock().await;
|
||||
let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until);
|
||||
if let Some(reset_at) = reset_at {
|
||||
pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds);
|
||||
}
|
||||
in_memory
|
||||
};
|
||||
|
||||
persisted || in_memory
|
||||
}
|
||||
|
||||
/// 从运行时账号池移除账号。
|
||||
pub async fn remove_account(&self, account_id: &str) -> bool {
|
||||
self.pool.lock().await.remove(account_id)
|
||||
}
|
||||
|
||||
/// 从仓储同步单个账号到运行时账号池。
|
||||
pub async fn sync_account_from_repository(
|
||||
&self,
|
||||
account_id: &str,
|
||||
) -> Result<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,
|
||||
};
|
||||
@@ -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>("a_json).ok(),
|
||||
Ok(None) => None,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id = %account_id,
|
||||
error = %error,
|
||||
"failed to read existing quota json before passive rate-limit sync"
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
let quota = rate_limit_quota(&rate_limits, plan_type, existing_quota.as_ref());
|
||||
self.apply_quota_snapshot(account_id, "a).await;
|
||||
}
|
||||
|
||||
/// 读取运行时账号快照。
|
||||
pub async fn account_snapshot(&self, account_id: &str) -> Option<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, "a_json, limit_reached, cooldown_until)
|
||||
.await
|
||||
{
|
||||
Ok(persisted) => persisted,
|
||||
Err(error) => {
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"failed to persist verified quota snapshot"
|
||||
);
|
||||
false
|
||||
}
|
||||
};
|
||||
|
||||
if let Some(reset_at) = reset_at {
|
||||
if let Err(error) = self
|
||||
.store
|
||||
.sync_rate_limit_window(account_id, reset_at, limit_window_seconds)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(
|
||||
account_id,
|
||||
error = %error,
|
||||
"failed to persist verified quota window"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let in_memory = {
|
||||
let mut pool = self.pool.lock().await;
|
||||
let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until);
|
||||
if let Some(reset_at) = reset_at {
|
||||
pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds);
|
||||
}
|
||||
in_memory
|
||||
};
|
||||
|
||||
persisted || in_memory
|
||||
}
|
||||
|
||||
/// 从运行时账号池移除账号。
|
||||
pub async fn remove_account(&self, account_id: &str) -> bool {
|
||||
self.pool.lock().await.remove(account_id)
|
||||
}
|
||||
|
||||
/// 从仓储同步单个账号到运行时账号池。
|
||||
pub async fn sync_account_from_store(
|
||||
&self,
|
||||
account_id: &str,
|
||||
) -> Result<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;
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
//! 账号配额解析与刷新。
|
||||
|
||||
mod runtime;
|
||||
mod service;
|
||||
|
||||
pub use runtime::*;
|
||||
pub use service::*;
|
||||
+13
-13
@@ -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,
|
||||
};
|
||||
@@ -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}"))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
//! Token 刷新策略、租约与运行时服务。
|
||||
|
||||
mod lease;
|
||||
mod runtime;
|
||||
mod service;
|
||||
|
||||
pub use lease::*;
|
||||
pub use runtime::*;
|
||||
pub use service::*;
|
||||
+19
-130
@@ -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,
|
||||
}
|
||||
}
|
||||
+4
-147
@@ -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};
|
||||
+2
-2
@@ -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};
|
||||
+1
-1
@@ -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};
|
||||
|
||||
+1
-1
@@ -3,7 +3,7 @@
|
||||
//! 最简单的策略:忽略账号负载与用量,纯粹按候选集合的顺序用轮转游标依次选中,
|
||||
//! 保证请求在候选账号间均匀轮转。
|
||||
|
||||
use crate::upstream::accounts::model::Account;
|
||||
use crate::accounts::account::Account;
|
||||
|
||||
use super::SelectionInput;
|
||||
|
||||
+2
-2
@@ -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};
|
||||
|
||||
+1
-1
@@ -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};
|
||||
|
||||
+7
-12
@@ -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)]
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
+161
-162
@@ -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>("a_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)
|
||||
})
|
||||
@@ -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("a_json)
|
||||
.bind(quota_fetched_at)
|
||||
.bind(quota_fetched_at)
|
||||
.bind(update.quota_verify_required)
|
||||
.bind(now)
|
||||
.bind(&update.account.id)
|
||||
.execute(&self.pool)
|
||||
.await?
|
||||
} else {
|
||||
sqlx::query(UPDATE_IMPORTED_ACCOUNT_PRESERVING_REFRESH_SQL)
|
||||
.bind(&update.account.email)
|
||||
.bind(&update.account.account_id)
|
||||
.bind(&update.account.user_id)
|
||||
.bind(&update.account.label)
|
||||
.bind(&update.account.plan_type)
|
||||
.bind(update.account.access_token.expose_secret())
|
||||
.bind(update.account.access_token_expires_at)
|
||||
.bind(status_to_db(update.account.status))
|
||||
.bind("a_json)
|
||||
.bind(quota_fetched_at)
|
||||
.bind(quota_fetched_at)
|
||||
.bind(update.quota_verify_required)
|
||||
.bind(now)
|
||||
.bind(&update.account.id)
|
||||
.execute(&self.pool)
|
||||
.await?
|
||||
};
|
||||
Ok(result.rows_affected() > 0)
|
||||
}
|
||||
|
||||
/// 设置下一次刷新时间。
|
||||
pub async fn set_next_refresh_at(
|
||||
&self,
|
||||
account_id: &str,
|
||||
next_refresh_at: Option<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)
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -1,4 +0,0 @@
|
||||
//! 管理端认证与本地会话模块。
|
||||
|
||||
pub mod service;
|
||||
pub mod session;
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -1,4 +0,0 @@
|
||||
//! 管理端 v1 接口访问 Key 模块。
|
||||
|
||||
pub mod routes;
|
||||
pub mod service;
|
||||
@@ -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(),
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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
@@ -1,3 +0,0 @@
|
||||
//! 管理端设置 HTTP 模块。
|
||||
|
||||
pub mod routes;
|
||||
@@ -1,5 +0,0 @@
|
||||
//! 系统版本与在线更新处理器。
|
||||
|
||||
pub mod routes;
|
||||
mod state;
|
||||
mod updater;
|
||||
@@ -1,5 +0,0 @@
|
||||
//! 管理端系统路由入口。
|
||||
|
||||
pub(crate) use super::updater::{
|
||||
perform_update, restart, rollback, update_detail, update_event_stream, update_status, version,
|
||||
};
|
||||
@@ -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, "a_by_account).await;
|
||||
|
||||
if use_numbered_page {
|
||||
return match state
|
||||
.services
|
||||
.admin_accounts
|
||||
.list_page(clamp_page(params.page.unwrap_or(1)), limit, params.search)
|
||||
.await
|
||||
{
|
||||
Ok(page) => {
|
||||
let stats = account_list_stats(&state, &page.items, "a_by_account).await?;
|
||||
let page = crate::infra::json::NumberedPage {
|
||||
items: page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|item| stats.data_for(item))
|
||||
.collect(),
|
||||
total: page.total,
|
||||
page: page.page,
|
||||
page_size: page.page_size,
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminAccountPageEnvelope::numbered(page, summary),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(account_error(&error)),
|
||||
};
|
||||
}
|
||||
|
||||
match state
|
||||
.services
|
||||
.admin_accounts
|
||||
.list(params.cursor, limit)
|
||||
.await
|
||||
{
|
||||
Ok(page) => {
|
||||
let stats = account_list_stats(&state, &page.items, "a_by_account).await?;
|
||||
let page = Page {
|
||||
items: page
|
||||
.items
|
||||
.into_iter()
|
||||
.map(|item| stats.data_for(item))
|
||||
.collect(),
|
||||
next_cursor: page.next_cursor,
|
||||
};
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminAccountPageEnvelope::cursor(page, limit, summary),
|
||||
))
|
||||
}
|
||||
Err(error) => Err(account_error(&error)),
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /api/admin/accounts`
|
||||
pub(crate) async fn create_account(
|
||||
State(state): State<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("a_json, Some(Utc::now()));
|
||||
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
|
||||
let account =
|
||||
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
|
||||
Ok(AdminResponse::new(
|
||||
StatusCode::OK,
|
||||
AdminEnvelope::ok(AdminAccountQuotaResponseData::from_account(
|
||||
quota, raw, quota_data, account,
|
||||
)),
|
||||
))
|
||||
}
|
||||
Err(AccountManageError::NotFound) => Err(account_not_found()),
|
||||
Err(AccountManageError::Inactive(status)) => Err(AdminError::conflict(format!(
|
||||
"Account is {}, cannot query quota",
|
||||
status.as_str()
|
||||
))),
|
||||
Err(AccountManageError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!(
|
||||
"Failed to fetch quota from Codex API: {msg}"
|
||||
))),
|
||||
Err(e) => Err(account_error(&e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn account_data_for_quota_refresh(
|
||||
state: &AppState,
|
||||
account_id: &str,
|
||||
quota_data: AccountQuotaData,
|
||||
) -> Result<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), "a_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, "a_by_account);
|
||||
let quota_window_usage_stats =
|
||||
quota_window_local_usage_by_account(state, &all_quota_windows).await;
|
||||
apply_quota_window_local_usage(accounts, &mut quota_by_account, "a_window_usage_stats);
|
||||
let selected_quota_windows = selected_quota_windows_by_account(accounts, "a_by_account);
|
||||
let usage_records = usage_records
|
||||
.into_iter()
|
||||
.map(|usage| {
|
||||
apply_selected_quota_window_usage(
|
||||
usage,
|
||||
&selected_quota_windows,
|
||||
"a_window_usage_stats,
|
||||
)
|
||||
})
|
||||
.collect::<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()
|
||||
}
|
||||
@@ -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,
|
||||
"a_used_by_account,
|
||||
now,
|
||||
),
|
||||
service_statuses: service_status_data(&state),
|
||||
usage_records: dashboard_usage_records,
|
||||
pool_summary,
|
||||
capacity_info: AccountCapacityDiagnostics::from(capacity),
|
||||
rotation_strategy: Some(settings.rotation_strategy.clone()),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
fn today_usage_record_filter(now: DateTime<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::*;
|
||||
+48
-441
@@ -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,
|
||||
"a_used_by_account,
|
||||
now,
|
||||
),
|
||||
service_statuses: service_status_data(&state),
|
||||
usage_records: dashboard_usage_records,
|
||||
pool_summary,
|
||||
capacity_info: AccountCapacityDiagnostics::from(capacity),
|
||||
rotation_strategy: Some(settings.auth.rotation_strategy.clone()),
|
||||
}),
|
||||
))
|
||||
}
|
||||
|
||||
fn today_usage_record_filter(now: DateTime<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(¤t_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()),
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)]
|
||||
+57
-110
@@ -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<
|
||||
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
//! 管理员认证与 Redis 会话。
|
||||
|
||||
pub mod service;
|
||||
pub mod store;
|
||||
pub mod types;
|
||||
@@ -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
|
||||
});
|
||||
}
|
||||
@@ -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)))
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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
Reference in New Issue
Block a user