starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-14 08:06:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

docs: add architecture overview

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-06-19 20:53:26 +08:00
co-authored by Codex
parent f07442a28b
commit 01b538c416
423 changed files with 59073 additions and 48538 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+21 -11
View File
@@ -1,22 +1,32 @@
# Rust build output.
/target/
/data/
/logs/
/local.yaml
/local.yml
/.worktrees/
/worktrees/
/.codex-ws-audit/
# Application runtime state.
/.runtime/
*.sqlite
*.sqlite3
*.db
*.db-shm
*.db-wal
*.log
# Local configuration.
/local.yaml
/local.yml
# Local workspaces and tooling state.
/.worktrees/
/worktrees/
/.codex-ws-audit/
/.codegraph/
*.tmp
*.swp
*~
*.codegraph
*/node_modules
*/dist
.agents
.cargo
/.agents/
/.cargo/
/.reasonix/
# Frontend dependencies and build output.
node_modules/
dist/
Cargo.lock
Generated
+136 -79
View File
@@ -60,9 +60,9 @@ checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3"
[[package]]
name = "alloc-stdlib"
version = "0.2.2"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece"
checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195"
dependencies = [
"alloc-no-stdlib",
]
@@ -262,9 +262,9 @@ dependencies = [
[[package]]
name = "brotli"
version = "8.0.3"
version = "8.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8119e4516436f5708bbc474a9d395bf12f1b5395e93a92a56e647ac3388c8610"
checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3"
dependencies = [
"alloc-no-stdlib",
"alloc-stdlib",
@@ -273,9 +273,9 @@ dependencies = [
[[package]]
name = "brotli-decompressor"
version = "5.0.1"
version = "5.0.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5962523e1b92ce1b5e793d9169b9943eece10d39f62550bc04bb605d75b94924"
checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583"
dependencies = [
"alloc-no-stdlib",
"alloc-stdlib",
@@ -365,27 +365,78 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a"
[[package]]
name = "codex-proxy-rs"
name = "codex-proxy-adapters"
version = "0.1.0"
dependencies = [
"async-trait",
"bytes",
"chrono",
"codex-proxy-core",
"codex-proxy-platform",
"flate2",
"futures",
"indexmap",
"reqwest",
"rustls",
"rustls-native-certs",
"rustls-pki-types",
"secrecy",
"serde",
"serde_json",
"sqlx",
"tempfile",
"thiserror 2.0.18",
"tokio",
"tokio-rustls",
"tokio-tungstenite",
"tracing",
"uuid",
"wiremock",
]
[[package]]
name = "codex-proxy-assets"
version = "0.1.0"
dependencies = [
"axum",
"tempfile",
"tokio",
"tower",
]
[[package]]
name = "codex-proxy-core"
version = "0.1.0"
dependencies = [
"async-trait",
"base64",
"chrono",
"hex",
"indexmap",
"rand 0.10.1",
"serde",
"serde_json",
"sha2 0.11.0",
"thiserror 2.0.18",
"tokio",
"tracing",
"uuid",
]
[[package]]
name = "codex-proxy-platform"
version = "0.1.0"
dependencies = [
"aes-gcm",
"anyhow",
"argon2",
"async-trait",
"axum",
"base64",
"chrono",
"config",
"dirs",
"flate2",
"futures",
"hex",
"hmac",
"indexmap",
"insta",
"pretty_assertions",
"rand 0.10.1",
"reqwest",
"rustls",
"rustls-native-certs",
"rustls-pki-types",
@@ -398,18 +449,77 @@ dependencies = [
"tempfile",
"thiserror 2.0.18",
"tokio",
"tokio-rustls",
"tokio-tungstenite",
"tower",
"tower-http",
"tracing",
"tracing-appender",
"tracing-subscriber",
"uuid",
"wiremock",
"zeroize",
]
[[package]]
name = "codex-proxy-rs"
version = "0.1.0"
dependencies = [
"codex-proxy-core",
"codex-proxy-runtime",
"codex-proxy-server",
]
[[package]]
name = "codex-proxy-runtime"
version = "0.1.0"
dependencies = [
"async-trait",
"base64",
"bytes",
"chrono",
"codex-proxy-adapters",
"codex-proxy-core",
"codex-proxy-platform",
"futures",
"reqwest",
"secrecy",
"serde_json",
"sqlx",
"tempfile",
"thiserror 2.0.18",
"tokio",
"tokio-tungstenite",
"tracing",
"uuid",
"wiremock",
]
[[package]]
name = "codex-proxy-server"
version = "0.1.0"
dependencies = [
"async-trait",
"axum",
"base64",
"chrono",
"codex-proxy-adapters",
"codex-proxy-assets",
"codex-proxy-core",
"codex-proxy-platform",
"codex-proxy-runtime",
"futures",
"secrecy",
"serde",
"serde_json",
"serde_yml",
"sqlx",
"tempfile",
"tokio",
"tokio-tungstenite",
"tower",
"tower-http",
"tracing",
"tracing-subscriber",
"uuid",
"wiremock",
]
[[package]]
name = "compression-codecs"
version = "0.4.38"
@@ -459,17 +569,6 @@ dependencies = [
"yaml-rust2",
]
[[package]]
name = "console"
version = "0.16.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87"
dependencies = [
"encode_unicode",
"libc",
"windows-sys 0.61.2",
]
[[package]]
name = "const-oid"
version = "0.10.2"
@@ -693,12 +792,6 @@ dependencies = [
"powerfmt",
]
[[package]]
name = "diff"
version = "0.1.13"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8"
[[package]]
name = "digest"
version = "0.10.7"
@@ -787,12 +880,6 @@ dependencies = [
"serde",
]
[[package]]
name = "encode_unicode"
version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0"
[[package]]
name = "encoding_rs"
version = "0.8.35"
@@ -1072,9 +1159,9 @@ dependencies = [
[[package]]
name = "h2"
version = "0.4.14"
version = "0.4.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "171fefbc92fe4a4de27e0698d6a5b392d6a0e333506bc49133760b3bcf948733"
checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155"
dependencies = [
"atomic-waker",
"bytes",
@@ -1435,18 +1522,6 @@ dependencies = [
"generic-array",
]
[[package]]
name = "insta"
version = "1.47.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b4a6248eb93a4401ed2f37dfe8ea592d3cf05b7cf4f8efa867b6895af7e094e"
dependencies = [
"console",
"once_cell",
"similar",
"tempfile",
]
[[package]]
name = "ipnet"
version = "2.12.0"
@@ -1850,16 +1925,6 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "pretty_assertions"
version = "1.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3ae130e2f271fbc2ac3a40fb1d07180839cdbbe443c7a27e1e3c13c5cac0116d"
dependencies = [
"diff",
"yansi",
]
[[package]]
name = "prettyplease"
version = "0.2.37"
@@ -2460,12 +2525,6 @@ version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214"
[[package]]
name = "similar"
version = "2.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa"
[[package]]
name = "slab"
version = "0.4.12"
@@ -3771,6 +3830,10 @@ version = "0.6.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4"
[[package]]
name = "xtask"
version = "0.1.0"
[[package]]
name = "yaml-rust2"
version = "0.11.0"
@@ -3782,12 +3845,6 @@ dependencies = [
"hashlink",
]
[[package]]
name = "yansi"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049"
[[package]]
name = "yoke"
version = "0.8.3"
Cargo.toml
+33 -12
View File
@@ -1,25 +1,46 @@
[package]
name = "codex-proxy-rs"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
rust-version.workspace = true
description = "A lean Rust Codex proxy backed only by ChatGPT/Codex accounts."
repository = "https://example.invalid/codex-proxy-rs"
autobins = false
[dev-dependencies]
codex-proxy-core = { path = "crates/core" }
codex-proxy-runtime = { path = "crates/runtime" }
codex-proxy-server = { path = "crates/server" }
[lints]
workspace = true
[workspace]
members = ["crates/*"]
resolver = "2"
[workspace.package]
version = "0.1.0"
edition = "2021"
license = "MIT"
publish = false
description = "A lean Rust Codex proxy backed only by ChatGPT/Codex accounts."
repository = "https://example.invalid/codex-proxy-rs"
rust-version = "1.95"
[dependencies]
[workspace.dependencies]
aes-gcm = "0.10.3"
anyhow = "1.0.102"
argon2 = "0.5.3"
async-trait = "0.1.89"
axum = "0.8.9"
base64 = "0.22.1"
bytes = "1.11.1"
chrono = { version = "0.4.45", features = ["serde", "clock"] }
config = "0.15.23"
dirs = "5.0.1"
futures = "0.3.32"
flate2 = "1.1.9"
futures = "0.3.32"
hex = "0.4.3"
hmac = "0.13.0"
indexmap = { version = "2.0", features = ["serde"] }
@@ -46,19 +67,19 @@ tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] }
uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
zeroize = "1.8.2"
[dev-dependencies]
insta = "1.47.2"
pretty_assertions = "1.4.1"
tempfile = "3.27.0"
wiremock = "0.6.5"
[lints.rust]
[workspace.lints.rust]
future_incompatible = { level = "warn", priority = -1 }
nonstandard_style = { level = "deny", priority = -1 }
unsafe_code = "forbid"
[lints.clippy]
[workspace.lints.clippy]
all = { level = "deny", priority = 10 }
redundant_clone = { level = "deny", priority = 9 }
large_enum_variant = { level = "warn", priority = 8 }
needless_collect = { level = "deny", priority = 7 }
[profile.dev]
debug = true
[profile.release]
debug = true
config.yaml
+4 -4
View File
@@ -32,10 +32,10 @@ quota:
usage_stats:
history_retention_days: null
database:
url: sqlite://data/codex-proxy-rs.sqlite
url: sqlite://.runtime/data/codex-proxy-rs.sqlite
security:
master_key_file: data/master.key
api_key_pepper_file: data/api-key-pepper.key
master_key_file: .runtime/data/master.key
api_key_pepper_file: .runtime/data/api-key-pepper.key
tls:
force_http11: false
ws_pool:
@@ -47,7 +47,7 @@ admin:
default_username: admin
default_password: admin
logging:
directory: logs
directory: .runtime/logs
retention_days: 14
enabled: true
capacity: 2000
crates/adapters/Cargo.toml
+38
View File
@@ -0,0 +1,38 @@
[package]
name = "codex-proxy-adapters"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
rust-version.workspace = true
[dependencies]
async-trait.workspace = true
bytes.workspace = true
codex-proxy-core = { path = "../core" }
codex-proxy-platform = { path = "../platform" }
chrono.workspace = true
flate2.workspace = true
futures.workspace = true
indexmap.workspace = true
reqwest.workspace = true
rustls.workspace = true
rustls-native-certs.workspace = true
rustls-pki-types.workspace = true
secrecy.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
thiserror.workspace = true
tokio.workspace = true
tokio-rustls.workspace = true
tokio-tungstenite.workspace = true
tracing.workspace = true
uuid.workspace = true
[dev-dependencies]
tempfile = "3.27.0"
wiremock = "0.6.5"
[lints]
workspace = true
crates/adapters/src/codex/client.rs
+1497
View File
File diff suppressed because it is too large Load Diff
crates/adapters/src/codex/fingerprint.rs
+477
View File
@@ -0,0 +1,477 @@
//! 指纹更新与持久化适配器。
use std::{path::PathBuf, sync::Arc, time::Duration};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sqlx::{Row, SqlitePool};
use thiserror::Error;
use tokio::sync::Mutex;
use tracing::{info, warn};
use uuid::Uuid;
use codex_proxy_core::gateway::fingerprint::Fingerprint;
/// 指纹历史记录的来源标识。
pub const CODEX_DESKTOP_UPDATE_SOURCE: &str = "codex_desktop_update_source";
const AUTO_UPDATED_FINGERPRINT_ID: &str = "auto_updated";
const AUTO_UPDATE_SOURCE: &str = "auto_update";
const AUTO_UPDATE_PLATFORM: &str = "darwin";
const AUTO_UPDATE_ARCH: &str = "arm64";
const AUTO_UPDATE_USER_AGENT_TEMPLATE: &str = "Codex Desktop/{version} ({platform}; {arch})";
const DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION: &str = "146";
const APPCAST_TIMEOUT: Duration = Duration::from_secs(30);
const POLL_INTERVAL: Duration = Duration::from_secs(3 * 24 * 60 * 60);
type AppcastFields = (Option<String>, Option<String>, Option<String>);
/// 更新清单解析结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FingerprintUpdate {
/// 应用版本。
pub app_version: String,
/// 构建号。
pub build_number: String,
}
/// 指纹更新错误。
#[derive(Debug, Error)]
pub enum FingerprintError {
/// 更新清单 JSON 无效。
#[error("invalid update manifest: {0}")]
InvalidManifest(#[from] serde_json::Error),
/// HTTP 请求失败。
#[error("failed to fetch update manifest: {0}")]
Http(#[from] reqwest::Error),
/// 数据库存储失败。
#[error("failed to persist fingerprint update: {0}")]
Database(#[from] sqlx::Error),
}
/// Appcast 检查错误。
#[derive(Debug, Error)]
pub enum UpdateError {
/// HTTP 请求失败。
#[error("HTTP 请求失败: {0}")]
Http(#[from] reqwest::Error),
/// 获取 appcast 失败。
#[error("获取 appcast 失败，状态码: {0}")]
AppcastFetch(u16),
/// 解析 appcast 失败。
#[error("解析 appcast 失败")]
AppcastParse,
/// JSON 序列化失败。
#[error("JSON 序列化失败: {0}")]
Json(#[from] serde_json::Error),
/// 文件操作失败。
#[error("文件操作失败: {0}")]
Io(#[from] std::io::Error),
/// 数据库存储失败。
#[error("数据库操作失败: {0}")]
Database(#[from] sqlx::Error),
}
#[derive(Deserialize)]
struct Manifest {
version: String,
build_number: String,
}
/// 从更新清单 JSON 中提取版本信息。
pub fn parse_update_manifest(input: &str) -> Result<FingerprintUpdate, FingerprintError> {
// 自动更新只同步桌面端指纹字段，不把远端配置当作运行时业务配置执行。
let manifest: Manifest = serde_json::from_str(input)?;
Ok(FingerprintUpdate {
app_version: manifest.version,
build_number: manifest.build_number,
})
}
/// 指纹历史记录。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoredFingerprint {
/// 应用版本。
pub app_version: String,
/// 构建号。
pub build_number: String,
/// 来源。
pub source: String,
}
/// 指纹自动更新状态。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UpdateState {
/// 最近检查时间。
pub last_check: DateTime<Utc>,
/// 最新版本。
pub latest_version: Option<String>,
/// 最新构建号。
pub latest_build: Option<String>,
/// 下载地址。
pub download_url: Option<String>,
/// 是否有可用更新。
pub update_available: bool,
/// 当前版本。
pub current_version: String,
/// 当前构建号。
pub current_build: String,
}
/// SQLite 指纹仓储。
#[derive(Clone)]
pub struct FingerprintRepository {
pool: SqlitePool,
}
impl FingerprintRepository {
/// 使用给定连接池构造仓储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 插入一条指纹更新历史。
pub async fn insert_update(&self, update: &FingerprintUpdate) -> Result<(), sqlx::Error> {
let mut fingerprint = Fingerprint::default_codex_desktop();
fingerprint.app_version.clone_from(&update.app_version);
fingerprint.build_number.clone_from(&update.build_number);
sqlx::query(
"insert into fingerprints (id, app_version, build_number, platform, arch, chromium_version, user_agent_template, source, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(Uuid::new_v4().to_string())
.bind(fingerprint.app_version)
.bind(fingerprint.build_number)
.bind(fingerprint.platform)
.bind(fingerprint.arch)
.bind(fingerprint.chromium_version)
.bind(fingerprint.user_agent_template)
.bind(CODEX_DESKTOP_UPDATE_SOURCE)
.bind(Utc::now().to_rfc3339())
.execute(&self.pool)
.await?;
Ok(())
}
/// 更新自动应用的最新版本快照。
pub async fn upsert_auto_update(
&self,
app_version: &str,
build_number: &str,
chromium_version: Option<&str>,
) -> Result<(), sqlx::Error> {
sqlx::query(
"insert into fingerprints (id, app_version, build_number, platform, arch, chromium_version, user_agent_template, source, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?) \
on conflict(id) do update set app_version = excluded.app_version, build_number = excluded.build_number, platform = excluded.platform, arch = excluded.arch, chromium_version = excluded.chromium_version, user_agent_template = excluded.user_agent_template, source = excluded.source, created_at = excluded.created_at",
)
.bind(AUTO_UPDATED_FINGERPRINT_ID)
.bind(app_version)
.bind(build_number)
.bind(AUTO_UPDATE_PLATFORM)
.bind(AUTO_UPDATE_ARCH)
.bind(chromium_version.unwrap_or(DEFAULT_AUTO_UPDATE_CHROMIUM_VERSION))
.bind(AUTO_UPDATE_USER_AGENT_TEMPLATE)
.bind(AUTO_UPDATE_SOURCE)
.bind(Utc::now().to_rfc3339())
.execute(&self.pool)
.await?;
Ok(())
}
/// 读取最新一条指纹历史记录。
pub async fn latest(&self) -> Result<Option<StoredFingerprint>, sqlx::Error> {
let row = sqlx::query(
"select app_version, build_number, source from fingerprints order by created_at desc, id desc limit 1",
)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|row| StoredFingerprint {
app_version: row.get("app_version"),
build_number: row.get("build_number"),
source: row.get("source"),
}))
}
/// 读取自动更新应用后的最新指纹。
pub async fn load_latest_auto_updated(&self) -> Result<Option<Fingerprint>, sqlx::Error> {
let row = sqlx::query(
r#"
select app_version, build_number, platform, arch, chromium_version, user_agent_template
from fingerprints
where source = ?
order by created_at desc
limit 1
"#,
)
.bind(AUTO_UPDATE_SOURCE)
.fetch_optional(&self.pool)
.await?;
let Some(row) = row else {
return Ok(None);
};
Ok(Some(Fingerprint {
originator: "Codex Desktop".to_string(),
app_version: row.get("app_version"),
build_number: row.get("build_number"),
platform: row.get("platform"),
arch: row.get("arch"),
chromium_version: row.get("chromium_version"),
user_agent_template: row.get("user_agent_template"),
default_headers: Fingerprint::default_headers(),
header_order: Fingerprint::default_header_order(),
}))
}
}
/// 通过 HTTP 拉取更新清单并写入历史记录。
#[derive(Clone)]
pub struct FingerprintUpdater {
client: reqwest::Client,
repository: FingerprintRepository,
update_url: String,
}
impl FingerprintUpdater {
/// 构造更新器。
pub fn new(
client: reqwest::Client,
repository: FingerprintRepository,
update_url: impl Into<String>,
) -> Self {
Self {
client,
repository,
update_url: update_url.into(),
}
}
/// 执行一次更新清单拉取并持久化历史记录。
pub async fn poll_once(&self) -> Result<FingerprintUpdate, FingerprintError> {
let manifest = self
.client
.get(&self.update_url)
.timeout(APPCAST_TIMEOUT)
.send()
.await?
.error_for_status()?
.text()
.await?;
let update = parse_update_manifest(&manifest)?;
self.repository.insert_update(&update).await?;
Ok(update)
}
}
/// Appcast 指纹更新检查器。
#[derive(Clone)]
pub struct UpdateChecker {
repository: Option<FingerprintRepository>,
client: reqwest::Client,
appcast_url: String,
extracted_fingerprint_path: PathBuf,
state: Arc<Mutex<InternalState>>,
}
struct InternalState {
current_version: String,
current_build: String,
}
impl UpdateChecker {
/// 使用自定义 client 和路径构造检查器。
pub fn with_client(
repository: Option<FingerprintRepository>,
client: reqwest::Client,
appcast_url: impl Into<String>,
extracted_fingerprint_path: PathBuf,
current_version: impl Into<String>,
current_build: impl Into<String>,
) -> Self {
Self {
repository,
client,
appcast_url: appcast_url.into(),
extracted_fingerprint_path,
state: Arc::new(Mutex::new(InternalState {
current_version: current_version.into(),
current_build: current_build.into(),
})),
}
}
/// 检查 appcast 是否存在更新。
pub async fn check_for_update(&self) -> Result<UpdateState, UpdateError> {
let state = self.state.lock().await;
let current_version = state.current_version.clone();
let current_build = state.current_build.clone();
drop(state);
let xml = self.fetch_appcast().await?;
let (version, build, download_url) = parse_appcast(&xml)?;
let update_available = version
.as_ref()
.is_some_and(|value| value != &current_version)
|| build.as_ref().is_some_and(|value| value != &current_build);
Ok(UpdateState {
last_check: Utc::now(),
latest_version: version,
latest_build: build,
download_url,
update_available,
current_version,
current_build,
})
}
/// 检查并在需要时应用更新。
pub async fn check_and_apply_update(&self) -> Result<bool, UpdateError> {
let update_state = self.check_for_update().await?;
if !update_state.update_available {
return Ok(false);
}
let (Some(version), Some(build)) = (
update_state.latest_version.as_deref(),
update_state.latest_build.as_deref(),
) else {
return Ok(false);
};
let state = self.state.lock().await;
let current_version = state.current_version.clone();
let current_build = state.current_build.clone();
drop(state);
info!(
version = %version,
build = %build,
current_version = %current_version,
current_build = %current_build,
"发现新的 fingerprint 版本"
);
self.apply_version_update(version, build).await?;
let mut state = self.state.lock().await;
state.current_version = version.to_string();
state.current_build = build.to_string();
Ok(true)
}
/// 启动后台轮询任务。
pub fn start_background_checker(self) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
let mut ticker = tokio::time::interval(POLL_INTERVAL);
info!(
interval_secs = POLL_INTERVAL.as_secs(),
"fingerprint 后台版本检查器已启动"
);
if let Err(error) = self.check_and_apply_update().await {
warn!(error = %error, "fingerprint 首次版本检查失败");
}
loop {
ticker.tick().await;
if let Err(error) = self.check_and_apply_update().await {
warn!(error = %error, "fingerprint 定期版本检查失败");
}
}
})
}
async fn fetch_appcast(&self) -> Result<String, UpdateError> {
let response = self
.client
.get(&self.appcast_url)
.timeout(APPCAST_TIMEOUT)
.send()
.await?;
if !response.status().is_success() {
return Err(UpdateError::AppcastFetch(response.status().as_u16()));
}
response.text().await.map_err(UpdateError::from)
}
async fn apply_version_update(&self, version: &str, build: &str) -> Result<(), UpdateError> {
let chromium_version =
load_matching_chromium_version(&self.extracted_fingerprint_path, version, build);
if let Some(repository) = &self.repository {
repository
.upsert_auto_update(version, build, chromium_version.as_deref())
.await?;
}
Ok(())
}
}
fn parse_appcast(xml: &str) -> Result<AppcastFields, UpdateError> {
let item_start = xml.find("<item>").ok_or(UpdateError::AppcastParse)?;
let item_end = xml.find("</item>").ok_or(UpdateError::AppcastParse)?;
let item = &xml[item_start..item_end];
let version = extract_sparkle_field(item, "shortVersionString");
let build = extract_sparkle_field(item, "version");
let download_url = item.find("url=\"").and_then(|position| {
let start = position + 5;
item[start..]
.find('"')
.map(|end| item[start..start + end].to_string())
});
Ok((version, build, download_url))
}
fn extract_sparkle_field(item: &str, field: &str) -> Option<String> {
let attr_pattern = format!("sparkle:{}=\"", field);
if let Some(position) = item.find(&attr_pattern) {
let start = position + attr_pattern.len();
if let Some(end) = item[start..].find('"') {
return Some(item[start..start + end].to_string());
}
}
let elem_start = format!("<sparkle:{}>", field);
let elem_end = format!("</sparkle:{}>", field);
if let Some(start_position) = item.find(&elem_start) {
let content_start = start_position + elem_start.len();
if let Some(end_position) = item[content_start..].find(&elem_end) {
return Some(item[content_start..content_start + end_position].to_string());
}
}
None
}
fn load_matching_chromium_version(
extracted_fingerprint_path: &PathBuf,
version: &str,
build: &str,
) -> Option<String> {
let content = std::fs::read_to_string(extracted_fingerprint_path).ok()?;
let parsed: serde_json::Value = serde_json::from_str(&content).ok()?;
let extracted_version = parsed.get("app_version")?.as_str()?;
let extracted_build = parsed.get("build_number")?.as_str()?;
if extracted_version == version && extracted_build == build {
parsed
.get("chromium_version")?
.as_str()
.map(ToString::to_string)
} else {
None
}
}
crates/adapters/src/codex/mod.rs
+15
View File
@@ -0,0 +1,15 @@
//! Codex 上游适配器。
pub mod client;
pub mod fingerprint;
pub mod models;
pub mod websocket;
pub use client::{
endpoint_request_path, endpoint_url, primary_usage_request_path, usage_endpoint_urls,
};
/// 构造默认 Codex HTTP 客户端。
pub fn default_http_client() -> client::HttpClient {
client::HttpClient::new()
}
crates/adapters/src/codex/models.rs
+3
View File
@@ -0,0 +1,3 @@
//! Codex 模型目录适配器。
pub use crate::codex::client::CodexBackendClient;
crates/adapters/src/codex/websocket/connect.rs
+1078
View File
File diff suppressed because it is too large Load Diff
src/codex/gateway/transport/websocket/deflate.rs → crates/adapters/src/codex/websocket/deflate.rs
+64 -28
View File
@@ -1,19 +1,24 @@
//! WebSocket permessage-deflate 处理。
use std::{
collections::VecDeque,
io,
io::{self, Read},
pin::Pin,
task::{Context, Poll},
};
use flate2::{Decompress, FlushDecompress, Status};
use flate2::read::DeflateDecoder;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
const SERVER_FRAME_CHUNK_BYTES: usize = 8192;
const PERMESSAGE_DEFLATE_TAIL: [u8; 4] = [0x00, 0x00, 0xff, 0xff];
const OPCODE_CONTINUATION: u8 = 0x0;
const OPCODE_TEXT: u8 = 0x1;
const OPCODE_BINARY: u8 = 0x2;
pub(super) struct PerMessageDeflateStream<S> {
const SERVER_FRAME_CHUNK_BYTES: usize = 8192;
/// Async stream wrapper that inflates server permessage-deflate data frames
/// before tungstenite sees them.
pub(crate) struct PerMessageDeflateStream<S> {
inner: S,
enabled: bool,
raw_input: Vec<u8>,
@@ -26,6 +31,7 @@ struct CompressedMessage {
payload: Vec<u8>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct FrameParts {
fin: bool,
rsv1: bool,
@@ -37,7 +43,7 @@ struct FrameParts {
}
impl<S> PerMessageDeflateStream<S> {
pub(super) fn new(inner: S, enabled: bool, preloaded: Vec<u8>) -> Self {
pub(crate) fn new(inner: S, enabled: bool, preloaded: Vec<u8>) -> Self {
Self {
inner,
enabled,
@@ -67,8 +73,10 @@ where
return Poll::Ready(Ok(()));
}
if this.rewrite_next_frame()? {
continue;
match this.rewrite_next_frame() {
Ok(true) => continue,
Ok(false) => {}
Err(error) => return Poll::Ready(Err(error)),
}
let mut chunk = [0_u8; SERVER_FRAME_CHUNK_BYTES];
@@ -122,9 +130,9 @@ impl<S> PerMessageDeflateStream<S> {
return Ok(true);
}
if frame.rsv1 && is_data_frame(frame.opcode) {
if frame.rsv1 && is_initial_data_frame(frame.opcode) {
if frame.fin {
let inflated = inflate_message(payload)?;
let inflated = inflate_permessage_deflate_message(payload)?;
self.decoded_input
.extend(encode_server_frame(frame.opcode, false, &inflated));
} else {
@@ -143,7 +151,7 @@ impl<S> PerMessageDeflateStream<S> {
let message = self.compressed_message.take().ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidData, "missing compressed message")
})?;
let inflated = inflate_message(&message.payload)?;
let inflated = inflate_permessage_deflate_message(&message.payload)?;
self.decoded_input.extend(encode_server_frame(
message.opcode,
false,
@@ -170,6 +178,48 @@ fn copy_decoded_input(decoded_input: &mut VecDeque<u8>, buffer: &mut ReadBuf<'_>
true
}
/// 判断扩展响应是否启用了 permessage-deflate。
pub fn permessage_deflate_enabled(extension_header: &str) -> bool {
extension_header
.split(',')
.any(|value| value.trim().starts_with("permessage-deflate"))
}
/// 解压单条 permessage-deflate 消息 payload。
pub fn inflate_permessage_deflate_message(payload: &[u8]) -> io::Result<Vec<u8>> {
let mut compressed = Vec::with_capacity(payload.len() + PERMESSAGE_DEFLATE_TAIL.len());
compressed.extend_from_slice(payload);
compressed.extend_from_slice(&PERMESSAGE_DEFLATE_TAIL);
let mut decoder = DeflateDecoder::new(compressed.as_slice());
let mut inflated = Vec::new();
decoder.read_to_end(&mut inflated)?;
if inflated.is_empty() {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"permessage-deflate frame produced no output",
));
}
Ok(inflated)
}
/// 将启用 RSV1 的压缩服务端 data frame 重写为普通未压缩 frame。
pub fn rewrite_permessage_deflate_server_frame(frame: &[u8]) -> io::Result<Option<Vec<u8>>> {
let Some(parts) = parse_frame_parts(frame)? else {
return Ok(None);
};
if parts.frame_len != frame.len() || parts.masked || !parts.rsv1 || !parts.fin {
return Ok(None);
}
if !is_data_frame(parts.opcode) {
return Ok(None);
}
let payload = &frame[parts.payload_offset..parts.payload_offset + parts.payload_len];
let inflated = inflate_permessage_deflate_message(payload)?;
Ok(Some(encode_server_frame(parts.opcode, false, &inflated)))
}
fn parse_frame_parts(bytes: &[u8]) -> io::Result<Option<FrameParts>> {
if bytes.len() < 2 {
return Ok(None);
@@ -231,24 +281,6 @@ fn parse_frame_parts(bytes: &[u8]) -> io::Result<Option<FrameParts>> {
}))
}
fn inflate_message(payload: &[u8]) -> io::Result<Vec<u8>> {
let mut compressed = Vec::with_capacity(payload.len() + PERMESSAGE_DEFLATE_TAIL.len());
compressed.extend_from_slice(payload);
compressed.extend_from_slice(&PERMESSAGE_DEFLATE_TAIL);
let mut decompressor = Decompress::new(false);
let mut inflated = Vec::with_capacity(payload.len().saturating_mul(4).max(1024));
match decompressor.decompress_vec(&compressed, &mut inflated, FlushDecompress::Sync) {
Ok(Status::Ok | Status::StreamEnd) => Ok(inflated),
Ok(Status::BufError) if !inflated.is_empty() => Ok(inflated),
Ok(Status::BufError) => Err(io::Error::new(
io::ErrorKind::InvalidData,
"permessage-deflate frame produced no output",
)),
Err(error) => Err(io::Error::new(io::ErrorKind::InvalidData, error)),
}
}
fn encode_server_frame(opcode: u8, rsv1: bool, payload: &[u8]) -> Vec<u8> {
let mut frame = Vec::new();
let rsv1_bit = if rsv1 { 0x40 } else { 0 };
@@ -269,5 +301,9 @@ fn encode_server_frame(opcode: u8, rsv1: bool, payload: &[u8]) -> Vec<u8> {
}
fn is_data_frame(opcode: u8) -> bool {
matches!(opcode, OPCODE_TEXT | OPCODE_BINARY | OPCODE_CONTINUATION)
}
fn is_initial_data_frame(opcode: u8) -> bool {
matches!(opcode, OPCODE_TEXT | OPCODE_BINARY)
}
crates/adapters/src/codex/websocket/mod.rs
+6
View File
@@ -0,0 +1,6 @@
//! Codex WebSocket 适配器。
pub mod connect;
pub mod deflate;
pub mod opening;
pub mod pool;
crates/adapters/src/codex/websocket/opening.rs
+56
View File
@@ -0,0 +1,56 @@
//! WebSocket 打开握手。
use std::{
io,
path::{Path, PathBuf},
};
use chrono::Utc;
use codex_proxy_core::protocol::codex::websocket::{OpeningAuditSnapshot, WebSocketAuditArtifact};
use uuid::Uuid;
/// WebSocket audit artifact 输出目录环境变量。
pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR";
/// 构造握手审计快照。
pub fn opening_audit_snapshot(header_order: Vec<String>) -> OpeningAuditSnapshot {
OpeningAuditSnapshot {
header_order,
..OpeningAuditSnapshot::default()
}
}
/// 显式写入 WebSocket audit artifact。
pub async fn write_websocket_audit_artifact_for_dir(
dir: Option<&Path>,
artifact: &WebSocketAuditArtifact,
) -> io::Result<Option<PathBuf>> {
let Some(dir) = dir.filter(|dir| !dir.as_os_str().is_empty()) else {
return Ok(None);
};
tokio::fs::create_dir_all(dir).await?;
let path = dir.join(websocket_audit_file_name());
let body = serde_json::to_vec_pretty(artifact).map_err(io::Error::other)?;
tokio::fs::write(&path, body).await?;
Ok(Some(path))
}
/// 按环境变量配置写入 WebSocket audit artifact。
pub async fn write_websocket_audit_artifact_from_env(
artifact: &WebSocketAuditArtifact,
) -> io::Result<Option<PathBuf>> {
let Some(dir) = std::env::var_os(WS_AUDIT_DIR_ENV)
.filter(|value| !value.is_empty())
.map(PathBuf::from)
else {
return Ok(None);
};
write_websocket_audit_artifact_for_dir(Some(&dir), artifact).await
}
fn websocket_audit_file_name() -> String {
let timestamp = Utc::now().format("%Y%m%dT%H%M%S%.3fZ");
format!("codex-ws-audit-{timestamp}-{}.json", Uuid::new_v4())
}
src/codex/gateway/transport/websocket/pool.rs → crates/adapters/src/codex/websocket/pool.rs
+183 -145
View File
@@ -1,3 +1,5 @@
//! Codex WebSocket 连接池。
use std::{
collections::HashMap,
sync::Arc,
@@ -6,15 +8,18 @@ use std::{
use futures::{SinkExt, StreamExt};
use tokio::{net::TcpStream, sync::Mutex, time::timeout};
use tokio_tungstenite::{
tungstenite::{http::HeaderMap as WsHeaderMap, Message},
MaybeTlsStream, WebSocketStream,
};
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
use super::codec::{rate_limit_headers, set_cookie_headers, turn_state};
use super::deflate::PerMessageDeflateStream;
use super::OpeningAuditSnapshot;
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60);
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5);
const DEFAULT_LIVENESS_TIMEOUT: Duration = Duration::from_millis(62_500);
/// WebSocket 连接池 key。
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct CodexWebSocketPoolKey {
base_url: String,
@@ -23,6 +28,7 @@ pub struct CodexWebSocketPoolKey {
}
impl CodexWebSocketPoolKey {
/// 构造连接池 key。
pub fn new(
base_url: impl Into<String>,
account_id: impl Into<String>,
@@ -35,44 +41,63 @@ impl CodexWebSocketPoolKey {
}
}
fn account_id(&self) -> &str {
pub(crate) fn account_id(&self) -> &str {
&self.account_id
}
}
/// WebSocket 连接池。
#[derive(Clone)]
pub struct CodexWebSocketPool {
inner: Arc<Mutex<WebSocketPoolState>>,
config: CodexWebSocketPoolConfig,
}
impl Default for CodexWebSocketPool {
fn default() -> Self {
Self::with_config(CodexWebSocketPoolConfig::default())
}
}
/// WebSocket 连接池配置。
#[derive(Debug, Clone, Copy)]
pub struct CodexWebSocketPoolConfig {
/// 是否启用连接池。
pub enabled: bool,
/// 单个 socket 的最大生命周期。
pub max_age: Duration,
/// 单个账号允许占用的最大池 slot 数。
pub max_per_account: usize,
/// 后台维护间隔；`None` 表示不启动后台任务。
pub maintenance_interval: Option<Duration>,
/// idle socket 探活 ping 间隔；`None` 表示维护时只做过期清理。
pub ping_interval: Option<Duration>,
/// 发送 ping 后等待上游响应的超时时间。
pub ping_timeout: Duration,
/// idle socket 无活动多久后视为失活。
pub liveness_timeout: Option<Duration>,
}
impl CodexWebSocketPool {
const DEFAULT_MAX_AGE: Duration = Duration::from_secs(55 * 60);
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25);
const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5);
const DEFAULT_LIVENESS_TIMEOUT: Duration = Duration::from_millis(62_500);
pub fn with_default_max_age() -> Self {
Self::with_config(CodexWebSocketPoolConfig::default())
impl Default for CodexWebSocketPoolConfig {
fn default() -> Self {
Self {
enabled: true,
max_age: DEFAULT_MAX_AGE,
max_per_account: DEFAULT_MAX_PER_ACCOUNT,
maintenance_interval: Some(DEFAULT_MAINTENANCE_INTERVAL),
ping_interval: Some(DEFAULT_PING_INTERVAL),
ping_timeout: DEFAULT_PING_TIMEOUT,
liveness_timeout: Some(DEFAULT_LIVENESS_TIMEOUT),
}
}
}
pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self {
impl CodexWebSocketPool {
/// 构造连接池策略和状态。
pub fn new(max_per_account: usize, max_age: Duration) -> Self {
Self::with_config(CodexWebSocketPoolConfig {
max_per_account: max_per_account.max(1),
max_age,
max_per_account,
maintenance_interval: None,
ping_interval: None,
liveness_timeout: None,
@@ -80,6 +105,7 @@ impl CodexWebSocketPool {
})
}
/// 使用完整配置构造连接池。
pub fn with_config(config: CodexWebSocketPoolConfig) -> Self {
let pool = Self {
inner: Arc::new(Mutex::new(WebSocketPoolState::default())),
@@ -89,46 +115,47 @@ impl CodexWebSocketPool {
pool
}
pub async fn evict_account(&self, account_id: &str) {
let mut idle_connections = Vec::new();
{
let mut state = self.inner.lock().await;
let keys = state
.slots
.keys()
.filter(|key| key.account_id() == account_id)
.cloned()
.collect::<Vec<_>>();
for key in keys {
if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) {
idle_connections.push(*connection);
}
}
/// 使用默认配置构造连接池。
pub fn with_default_max_age() -> Self {
Self::with_config(CodexWebSocketPoolConfig::default())
}
/// 使用最大存活时间与账号容量限制构造连接池。
pub fn with_limits(max_age: Duration, max_per_account: usize) -> Self {
Self {
inner: Arc::new(Mutex::new(WebSocketPoolState::default())),
config: CodexWebSocketPoolConfig {
max_age,
max_per_account: max_per_account.max(1),
maintenance_interval: None,
ping_interval: None,
liveness_timeout: None,
..CodexWebSocketPoolConfig::default()
},
}
close_idle_connections(idle_connections).await;
}
pub async fn shutdown(&self) {
let idle_connections = {
let mut state = self.inner.lock().await;
state.shutting_down = true;
state
.slots
.drain()
.filter_map(|(_, slot)| match slot {
WebSocketPoolSlot::Idle(connection) => Some(*connection),
WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None,
})
.collect::<Vec<_>>()
};
close_idle_connections(idle_connections).await;
/// 返回单账号最大连接数。
pub fn max_per_account(&self) -> usize {
self.config.max_per_account
}
pub async fn gc_sweep(&self) {
self.maintain_idle_connections().await;
/// 返回连接最大存活时长。
pub fn max_age(&self) -> Duration {
self.config.max_age
}
pub(super) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
/// 判断当前账号是否还能打开新连接。
pub fn permits_new_connection(&self, current_connections: usize) -> bool {
current_connections < self.config.max_per_account
}
/// 判断连接是否已达到回收年龄。
pub fn should_recycle(&self, age: Duration) -> bool {
age >= self.config.max_age
}
pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
let mut expired_connection = None;
let acquire = {
let mut state = self.inner.lock().await;
@@ -143,7 +170,7 @@ impl CodexWebSocketPool {
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else {
return WebSocketPoolAcquire::Bypass;
};
if connection.created_at.elapsed() <= self.config.max_age {
if connection.created_at.elapsed() < self.config.max_age {
state.slots.insert(key.clone(), WebSocketPoolSlot::Busy);
return WebSocketPoolAcquire::Reused(connection);
}
@@ -163,13 +190,13 @@ impl CodexWebSocketPool {
};
if let Some(connection) = expired_connection {
close_idle_connections(vec![connection]).await;
close_pooled_connection(connection).await;
}
acquire
}
pub(super) async fn put(
pub(crate) async fn put(
&self,
key: CodexWebSocketPoolKey,
connection: PooledWebSocketConnection,
@@ -179,7 +206,7 @@ impl CodexWebSocketPool {
let mut state = self.inner.lock().await;
let expired = connection
.as_ref()
.is_some_and(|connection| connection.created_at.elapsed() > self.config.max_age);
.is_some_and(|connection| connection.created_at.elapsed() >= self.config.max_age);
if expired || state.shutting_down || !self.config.enabled {
state.slots.remove(&key);
} else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Busy)) {
@@ -195,7 +222,7 @@ impl CodexWebSocketPool {
}
}
pub(super) async fn discard(&self, key: &CodexWebSocketPoolKey) {
pub(crate) async fn discard(&self, key: &CodexWebSocketPoolKey) {
let idle_connection = {
let mut state = self.inner.lock().await;
match state.slots.remove(key) {
@@ -204,7 +231,59 @@ impl CodexWebSocketPool {
}
};
if let Some(connection) = idle_connection {
close_idle_connections(vec![connection]).await;
close_pooled_connection(connection).await;
}
}
/// 驱逐指定账号的 idle 连接，并阻止已占用 slot 回收到池中。
pub async fn evict_account(&self, account_id: &str) {
let mut idle_connections = Vec::new();
{
let mut state = self.inner.lock().await;
let keys = state
.slots
.keys()
.filter(|key| key.account_id() == account_id)
.cloned()
.collect::<Vec<_>>();
for key in keys {
if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) {
idle_connections.push(*connection);
}
}
}
close_pooled_connections(idle_connections).await;
}
/// 关闭连接池，关闭所有 idle 连接，并让后续 acquire 直接绕过池。
pub async fn shutdown(&self) {
let idle_connections = {
let mut state = self.inner.lock().await;
state.shutting_down = true;
state
.slots
.drain()
.filter_map(|(_, slot)| match slot {
WebSocketPoolSlot::Idle(connection) => Some(*connection),
WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None,
})
.collect::<Vec<_>>()
};
close_pooled_connections(idle_connections).await;
}
/// 维护 idle 连接：关闭过期/失活连接，并按需 ping 探活。
pub async fn gc_sweep(&self) {
self.maintain_idle_connections().await;
}
/// 维护 idle 连接：关闭过期/失活连接，并按需 ping 探活。
pub async fn maintain_idle_connections(&self) {
let maintenance = self.take_idle_connections_for_maintenance().await;
close_pooled_connections(maintenance.close).await;
for (key, connection) in maintenance.probe {
let connection = probe_idle_connection(connection, self.config.ping_timeout).await;
self.return_maintained_connection(key, connection).await;
}
}
@@ -212,11 +291,11 @@ impl CodexWebSocketPool {
let Some(interval_duration) = self.config.maintenance_interval else {
return;
};
let inner = Arc::downgrade(&self.inner);
let config = self.config;
let Ok(_handle) = tokio::runtime::Handle::try_current() else {
return;
};
let inner = Arc::downgrade(&self.inner);
let config = self.config;
tokio::spawn(async move {
let mut interval = tokio::time::interval(interval_duration);
loop {
@@ -237,15 +316,6 @@ impl CodexWebSocketPool {
self.inner.lock().await.shutting_down
}
pub async fn maintain_idle_connections(&self) {
let maintenance = self.take_idle_connections_for_maintenance().await;
close_idle_connections(maintenance.close).await;
for (key, connection) in maintenance.probe {
let connection = probe_idle_connection(connection, self.config.ping_timeout).await;
self.return_maintained_connection(key, connection).await;
}
}
async fn take_idle_connections_for_maintenance(&self) -> WebSocketPoolMaintenance {
let mut close = Vec::new();
let mut probe = Vec::new();
@@ -255,8 +325,26 @@ impl CodexWebSocketPool {
if state.shutting_down || !self.config.enabled {
return WebSocketPoolMaintenance { close, probe };
}
let expired_keys = expired_idle_keys(&state.slots, self.config.max_age);
for key in expired_keys {
let keys = state
.slots
.iter()
.filter_map(|(key, slot)| match slot {
WebSocketPoolSlot::Idle(connection)
if should_close_idle_connection(
connection,
now,
self.config.max_age,
self.config.liveness_timeout,
) =>
{
Some(key.clone())
}
WebSocketPoolSlot::Idle(_)
| WebSocketPoolSlot::Busy
| WebSocketPoolSlot::Checking => None,
})
.collect::<Vec<_>>();
for key in keys {
if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) {
close.push(*connection);
}
@@ -332,59 +420,28 @@ impl CodexWebSocketPool {
}
}
impl Default for CodexWebSocketPool {
fn default() -> Self {
Self::with_default_max_age()
}
}
impl Default for CodexWebSocketPoolConfig {
fn default() -> Self {
Self {
enabled: true,
max_age: CodexWebSocketPool::DEFAULT_MAX_AGE,
max_per_account: CodexWebSocketPool::DEFAULT_MAX_PER_ACCOUNT,
maintenance_interval: Some(CodexWebSocketPool::DEFAULT_MAINTENANCE_INTERVAL),
ping_interval: Some(CodexWebSocketPool::DEFAULT_PING_INTERVAL),
ping_timeout: CodexWebSocketPool::DEFAULT_PING_TIMEOUT,
liveness_timeout: Some(CodexWebSocketPool::DEFAULT_LIVENESS_TIMEOUT),
}
}
}
#[derive(Default)]
struct WebSocketPoolState {
slots: HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
shutting_down: bool,
}
pub(super) type CodexWsStream = WebSocketStream<PerMessageDeflateStream<MaybeTlsStream<TcpStream>>>;
pub(crate) type CodexWsStream = WebSocketStream<PerMessageDeflateStream<MaybeTlsStream<TcpStream>>>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct CodexWebSocketConnectionMetadata {
pub(super) turn_state: Option<String>,
pub(super) set_cookie_headers: Vec<String>,
pub(super) rate_limit_headers: Vec<(String, String)>,
pub(super) opening_audit: OpeningAuditSnapshot,
#[derive(Clone)]
pub(crate) struct CodexWebSocketConnectionMetadata {
pub(crate) turn_state: Option<String>,
pub(crate) set_cookie_headers: Vec<String>,
pub(crate) rate_limit_headers: Vec<(String, String)>,
pub(crate) handshake_status: u16,
}
impl CodexWebSocketConnectionMetadata {
pub(super) fn from_headers(headers: &WsHeaderMap, opening_audit: OpeningAuditSnapshot) -> Self {
Self {
turn_state: turn_state(headers),
set_cookie_headers: set_cookie_headers(headers),
rate_limit_headers: rate_limit_headers(headers),
opening_audit,
}
}
}
pub(super) struct PooledWebSocketConnection {
pub(super) websocket: CodexWsStream,
pub(super) metadata: CodexWebSocketConnectionMetadata,
pub(super) created_at: Instant,
pub(super) last_activity_at: Instant,
pub(super) last_ping_at: Option<Instant>,
pub(crate) struct PooledWebSocketConnection {
pub(crate) websocket: CodexWsStream,
pub(crate) metadata: CodexWebSocketConnectionMetadata,
pub(crate) created_at: Instant,
pub(crate) last_activity_at: Instant,
pub(crate) last_ping_at: Option<Instant>,
}
enum WebSocketPoolSlot {
@@ -393,64 +450,46 @@ enum WebSocketPoolSlot {
Checking,
}
pub(super) enum WebSocketPoolAcquire {
pub(crate) enum WebSocketPoolAcquire {
Reused(Box<PooledWebSocketConnection>),
FreshReserved,
Bypass,
}
fn account_slot_count(
inner: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
slots: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
account_id: &str,
) -> usize {
inner
slots
.keys()
.filter(|key| key.account_id() == account_id)
.count()
}
async fn close_idle_connections(connections: Vec<PooledWebSocketConnection>) {
async fn close_pooled_connection(mut connection: PooledWebSocketConnection) {
let _ = connection.websocket.send(Message::Close(None)).await;
}
async fn close_pooled_connections(connections: Vec<PooledWebSocketConnection>) {
for connection in connections {
close_pooled_connection(connection).await;
}
}
async fn close_pooled_connection(connection: PooledWebSocketConnection) {
let mut websocket = connection.websocket;
let _ = websocket.close(None).await;
}
struct WebSocketPoolMaintenance {
close: Vec<PooledWebSocketConnection>,
probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>,
}
fn expired_idle_keys(
slots: &HashMap<CodexWebSocketPoolKey, WebSocketPoolSlot>,
max_age: Duration,
) -> Vec<CodexWebSocketPoolKey> {
slots
.iter()
.filter_map(|(key, slot)| match slot {
WebSocketPoolSlot::Idle(connection) if connection.created_at.elapsed() > max_age => {
Some(key.clone())
}
WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => {
None
}
})
.collect()
}
fn should_close_idle_connection(
connection: &PooledWebSocketConnection,
now: Instant,
max_age: Duration,
liveness_timeout: Option<Duration>,
) -> bool {
connection.created_at.elapsed() > max_age
now.duration_since(connection.created_at) >= max_age
|| liveness_timeout.is_some_and(|timeout| {
!timeout.is_zero() && now.duration_since(connection.last_activity_at) > timeout
!timeout.is_zero() && now.duration_since(connection.last_activity_at) >= timeout
})
}
@@ -469,8 +508,7 @@ async fn probe_idle_connection(
mut connection: PooledWebSocketConnection,
ping_timeout: Duration,
) -> Option<PooledWebSocketConnection> {
let now = Instant::now();
connection.last_ping_at = Some(now);
connection.last_ping_at = Some(Instant::now());
if connection
.websocket
.send(Message::Ping(Vec::new().into()))
crates/adapters/src/lib.rs
+7
View File
@@ -0,0 +1,7 @@
#![deny(missing_docs)]
//! 适配器层，承载具体存储、HTTP、OAuth 和 WebSocket 实现。
pub mod codex;
pub mod oauth;
pub mod sqlite;
crates/adapters/src/oauth/mod.rs
+4
View File
@@ -0,0 +1,4 @@
//! OAuth 上游适配器。
/// OpenAI OAuth 的 Reqwest 实现。
pub mod openai;
crates/adapters/src/oauth/openai.rs
+190
View File
@@ -0,0 +1,190 @@
//! OpenAI OAuth Reqwest 适配器。
use async_trait::async_trait;
use reqwest::{Client, StatusCode};
use serde::Deserialize;
use codex_proxy_core::auth::{
oauth::{DeviceCode, OAuthConfig, OAuthError, RefreshFailure, TokenPair},
ports::{OAuthClient, TokenRefresher},
};
/// OpenAI OAuth 上游客户端。
#[derive(Clone)]
pub struct OpenAiOAuthClient {
client: Client,
config: OAuthConfig,
}
impl OpenAiOAuthClient {
/// 使用 Reqwest 客户端和静态配置构造 OAuth 客户端。
pub fn new(client: Client, config: OAuthConfig) -> Self {
Self { client, config }
}
}
/// 使用默认 Reqwest 客户端构造 OpenAI OAuth 客户端。
pub fn default_openai_oauth_client(config: OAuthConfig) -> OpenAiOAuthClient {
OpenAiOAuthClient::new(Client::new(), config)
}
#[derive(Deserialize)]
struct RefreshTokenResponse {
access_token: String,
refresh_token: Option<String>,
}
#[derive(Deserialize)]
struct OAuthErrorResponse {
error: Option<String>,
error_description: Option<String>,
}
#[async_trait]
impl OAuthClient for OpenAiOAuthClient {
async fn exchange_code(
&self,
code: &str,
code_verifier: &str,
redirect_uri: &str,
) -> Result<TokenPair, OAuthError> {
let response = self
.client
.post(&self.config.token_endpoint)
.form(&[
("grant_type", "authorization_code"),
("client_id", self.config.client_id.as_str()),
("code", code),
("redirect_uri", redirect_uri),
("code_verifier", code_verifier),
])
.send()
.await
.map_err(|_| OAuthError::Transport)?;
let status = response.status();
let body = response.text().await.map_err(|_| OAuthError::Transport)?;
if !status.is_success() {
return Err(OAuthError::Rejected(format!(
"Token exchange failed ({status}): {body}"
)));
}
parse_token_pair(&body).map_err(|_| OAuthError::Transport)
}
async fn request_device_code(&self) -> Result<DeviceCode, OAuthError> {
let response = self
.client
.post(&self.config.device_code_endpoint)
.form(&[
("client_id", self.config.client_id.as_str()),
("scope", "openid profile email offline_access"),
])
.send()
.await
.map_err(|_| OAuthError::Transport)?;
let status = response.status();
let body = response.text().await.map_err(|_| OAuthError::Transport)?;
if !status.is_success() {
return Err(OAuthError::Rejected(format!(
"Device code request failed ({status}): {body}"
)));
}
serde_json::from_str::<DeviceCode>(&body).map_err(|_| OAuthError::Transport)
}
async fn poll_device_token(&self, device_code: &str) -> Result<TokenPair, OAuthError> {
let response = self
.client
.post(&self.config.token_endpoint)
.form(&[
("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
("device_code", device_code),
("client_id", self.config.client_id.as_str()),
])
.send()
.await
.map_err(|_| OAuthError::Transport)?;
let status = response.status();
let body = response.text().await.map_err(|_| OAuthError::Transport)?;
if !status.is_success() {
return Err(classify_device_poll_failure(status, &body));
}
parse_token_pair(&body).map_err(|_| OAuthError::Transport)
}
}
#[async_trait]
impl TokenRefresher for OpenAiOAuthClient {
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure> {
let response = self
.client
.post(&self.config.token_endpoint)
.form(&[
("grant_type", "refresh_token"),
("client_id", self.config.client_id.as_str()),
("refresh_token", refresh_token),
])
.send()
.await
.map_err(|_| RefreshFailure::Transport)?;
let status = response.status();
let body = response
.text()
.await
.map_err(|_| RefreshFailure::Transport)?;
if !status.is_success() {
return Err(classify_refresh_failure(status, &body));
}
parse_token_pair(&body).map_err(|_| RefreshFailure::Transport)
}
}
fn parse_token_pair(body: &str) -> Result<TokenPair, ()> {
let tokens = serde_json::from_str::<RefreshTokenResponse>(body).map_err(|_| ())?;
if tokens.access_token.trim().is_empty() {
return Err(());
}
Ok(TokenPair {
access_token: tokens.access_token,
refresh_token: tokens.refresh_token,
})
}
fn classify_device_poll_failure(status: StatusCode, body: &str) -> OAuthError {
let parsed = serde_json::from_str::<OAuthErrorResponse>(body).ok();
match parsed.as_ref().and_then(|data| data.error.as_deref()) {
Some("authorization_pending") => OAuthError::AuthorizationPending,
Some("slow_down") => OAuthError::SlowDown,
Some(code) => OAuthError::Rejected(
parsed
.as_ref()
.and_then(|data| data.error_description.clone())
.unwrap_or_else(|| code.to_string()),
),
None => OAuthError::Rejected(format!("Device token poll failed ({status}): {body}")),
}
}
fn classify_refresh_failure(status: StatusCode, body: &str) -> RefreshFailure {
let lower = body.to_ascii_lowercase();
if lower.contains("quota") {
return RefreshFailure::QuotaExhausted;
}
if lower.contains("account has been deactivated")
|| lower.contains("refresh_token_reused")
|| lower.contains("banned")
{
return RefreshFailure::Banned;
}
if status == StatusCode::BAD_REQUEST
|| status == StatusCode::UNAUTHORIZED
|| lower.contains("invalid_grant")
|| lower.contains("invalid_token")
|| lower.contains("access_denied")
|| lower.contains("refresh_token_expired")
|| lower.contains("token_revoked")
{
return RefreshFailure::InvalidGrant;
}
RefreshFailure::Transport
}
crates/adapters/src/sqlite/account_tokens.rs
+10
View File
@@ -0,0 +1,10 @@
//! SQLite 账号 token 加解密辅助。
use secrecy::SecretString;
use codex_proxy_platform::crypto::{CryptoResult, SecretBox};
/// 加密账号 token。
pub fn encrypt_token(secret_box: &SecretBox, token: &SecretString) -> CryptoResult<String> {
secret_box.encrypt(token)
}
crates/adapters/src/sqlite/account_usage.rs
+21
View File
@@ -0,0 +1,21 @@
//! SQLite 账号用量存储。
use sqlx::SqlitePool;
/// SQLite 用量存储。
#[derive(Clone)]
pub struct SqliteUsageStore {
pool: SqlitePool,
}
impl SqliteUsageStore {
/// 构造存储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 返回底层连接池。
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
}
crates/adapters/src/sqlite/accounts.rs
+1671
View File
File diff suppressed because it is too large Load Diff
src/admin/session/repository.rs → crates/adapters/src/sqlite/admin_sessions.rs
+47 -27
View File
@@ -1,66 +1,81 @@
use chrono::{DateTime, Utc};
use sqlx::{Row, SqlitePool};
use uuid::Uuid;
//! SQLite 管理员会话存储。
use chrono::{DateTime, Utc};
use sqlx::SqlitePool;
/// 已持久化的管理员用户。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AdminUserRecord {
pub struct StoredAdminUser {
/// 管理员用户 ID。
pub id: String,
/// 管理员密码哈希。
pub password_hash: String,
}
/// SQLite 管理员会话存储。
#[derive(Clone)]
pub struct AdminAuthRepository {
pub struct SqliteAdminSessionStore {
pool: SqlitePool,
}
impl AdminAuthRepository {
/// SQLite 管理员会话存储结果。
pub type SqliteAdminSessionStoreResult<T> = Result<T, sqlx::Error>;
impl SqliteAdminSessionStore {
/// 构造存储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
pub async fn ensure_default_admin(&self, password_hash: &str) -> Result<bool, sqlx::Error> {
let mut tx = self.pool.begin().await?;
let count: (i64,) = sqlx::query_as("select count(*) from admin_users")
.fetch_one(&mut *tx)
/// 返回底层连接池。
pub fn pool(&self) -> &SqlitePool {
&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 count.0 > 0 {
tx.commit().await?;
if existing.0 > 0 {
return Ok(false);
}
let admin_id = format!("admin_{}", Uuid::new_v4().simple());
let now = Utc::now().to_rfc3339();
sqlx::query(
"insert into admin_users (id, password_hash, created_at, updated_at) values (?, ?, ?, ?)",
)
.bind(admin_id)
.bind("admin_1")
.bind(password_hash)
.bind(&now)
.bind(&now)
.execute(&mut *tx)
.execute(&self.pool)
.await?;
tx.commit().await?;
Ok(true)
}
pub async fn load_first_admin(&self) -> Result<Option<AdminUserRecord>, sqlx::Error> {
let row = sqlx::query(
"select id, password_hash from admin_users order by created_at asc limit 1",
/// 读取第一个管理员用户。
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(|row| AdminUserRecord {
id: row.get("id"),
password_hash: row.get("password_hash"),
}))
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>,
) -> Result<(), sqlx::Error> {
) -> SqliteAdminSessionStoreResult<()> {
let now = Utc::now().to_rfc3339();
sqlx::query(
"insert into admin_sessions (id, user_id, expires_at, created_at) values (?, ?, ?, ?)",
@@ -74,7 +89,8 @@ impl AdminAuthRepository {
Ok(())
}
pub async fn validate_session(&self, session_id: &str) -> Result<bool, sqlx::Error> {
/// 校验管理员会话是否存在且未过期。
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)
@@ -84,8 +100,12 @@ impl AdminAuthRepository {
Ok(row.is_some())
}
pub async fn cleanup_expired_sessions(&self, now: DateTime<Utc>) -> Result<u64, sqlx::Error> {
let result = sqlx::query("delete from admin_sessions where expires_at < ?")
/// 删除已经过期的管理员会话。
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?;
crates/adapters/src/sqlite/client_keys.rs
+276
View File
@@ -0,0 +1,276 @@
//! SQLite 客户端 API Key 存储适配器。
use async_trait::async_trait;
use chrono::Utc;
use sqlx::{Row, SqlitePool};
use thiserror::Error;
use uuid::Uuid;
use codex_proxy_core::admin::ports::{ClientKeyStore, ClientKeyStoreError, ClientKeyStoreResult};
use codex_proxy_platform::{
identity::{ApiKeyHasher, AuthError},
json::{decode_cursor, encode_cursor, Page},
};
/// SQLite 客户端 API Key 存储错误。
#[derive(Debug, Error)]
pub enum SqliteClientKeyStoreError {
/// 数据库错误。
#[error("sqlite client key database error: {0}")]
Database(#[from] sqlx::Error),
/// API Key 验证错误。
#[error("sqlite client key auth error: {0}")]
Auth(#[from] AuthError),
/// 分页游标无效。
#[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,
/// 是否允许用于 `/v1` 认证。
pub enabled: bool,
/// 创建时间。
pub created_at: String,
/// 最近一次成功使用时间。
pub last_used_at: Option<String>,
}
/// 新建客户端 API Key 后的一次性结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CreatedClientApiKey {
/// API Key 记录 ID。
pub id: String,
/// API Key 名称。
pub name: String,
/// 管理员可见标签。
pub label: Option<String>,
/// 明文 API Key 的短前缀。
pub prefix: String,
/// 是否允许用于 `/v1` 认证。
pub enabled: bool,
/// 创建时间。
pub created_at: String,
/// 最近一次成功使用时间。
pub last_used_at: Option<String>,
/// 仅返回一次的明文 API Key。
pub plaintext: String,
}
/// SQLite 客户端 API Key 存储。
#[derive(Clone)]
pub struct SqliteClientKeyStore {
pool: SqlitePool,
hasher: ApiKeyHasher,
}
impl SqliteClientKeyStore {
/// 构造存储适配器。
pub fn new(pool: SqlitePool, hasher: ApiKeyHasher) -> Self {
Self { pool, hasher }
}
/// 暴露底层连接池，供集成测试和运行时组合层复用。
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
/// 创建新的本地客户端 API Key。
pub async fn create(
&self,
name: &str,
) -> Result<CreatedClientApiKey, SqliteClientKeyStoreError> {
let generated = self.hasher.generate_client_api_key(name);
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_hash, enabled, created_at, last_used_at) values (?, ?, null, ?, ?, 1, ?, null)",
)
.bind(&id)
.bind(name)
.bind(&generated.prefix)
.bind(&generated.key_hash)
.bind(&now)
.execute(&self.pool)
.await?;
Ok(CreatedClientApiKey {
id,
name: name.to_string(),
label: None,
prefix: generated.prefix,
enabled: true,
created_at: now,
last_used_at: None,
plaintext: generated.plaintext,
})
}
/// 按创建时间倒序分页列出客户端 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, 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, 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| 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, enabled, created_at, last_used_at from client_api_keys where id = ?",
)
.bind(id)
.fetch_optional(&self.pool)
.await?;
Ok(row.map(|row| 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)
}
}
#[async_trait]
impl ClientKeyStore for SqliteClientKeyStore {
async fn verify_and_touch(&self, plaintext: &str) -> ClientKeyStoreResult<bool> {
verify_and_touch(&self.pool, &self.hasher, plaintext)
.await
.map_err(map_client_key_store_error)
}
}
async fn verify_and_touch(
pool: &SqlitePool,
hasher: &ApiKeyHasher,
plaintext: &str,
) -> Result<bool, SqliteClientKeyStoreError> {
let prefix = plaintext.chars().take(12).collect::<String>();
let rows =
sqlx::query("select id, key_hash from client_api_keys where prefix = ? and enabled = 1")
.bind(prefix)
.fetch_all(pool)
.await?;
for row in rows {
let key_hash = row.get::<String, _>("key_hash");
if hasher.verify_client_api_key(plaintext, &key_hash)? {
let id = row.get::<String, _>("id");
sqlx::query("update client_api_keys set last_used_at = ? where id = ?")
.bind(Utc::now().to_rfc3339())
.bind(id)
.execute(pool)
.await?;
return Ok(true);
}
}
Ok(false)
}
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"),
enabled: row.get::<i64, _>("enabled") != 0,
created_at: row.get("created_at"),
last_used_at: row.get("last_used_at"),
}
}
fn map_client_key_store_error(error: SqliteClientKeyStoreError) -> ClientKeyStoreError {
ClientKeyStoreError::OperationFailed {
message: error.to_string(),
}
}
src/codex/accounts/cookies/repository.rs → crates/adapters/src/sqlite/cookies.rs
+107 -66
View File
@@ -1,35 +1,63 @@
//! SQLite Cookie 存储。
use chrono::{DateTime, Duration, Utc};
use secrecy::{ExposeSecret, SecretString};
use sqlx::{Row, SqlitePool};
use thiserror::Error;
use uuid::Uuid;
use crate::platform::crypto::{CryptoError, SecretBox};
use codex_proxy_platform::crypto::{CryptoError, SecretBox};
const DEFAULT_COOKIE_DOMAIN: &str = "chatgpt.com";
const CAPTURABLE_COOKIES: &[&str] = &["cf_clearance"];
/// SQLite Cookie 存储错误。
#[derive(Debug, Error)]
pub enum CookieError {
#[error("database error: {0}")]
pub enum SqliteCookieStoreError {
/// 数据库错误。
#[error("sqlite cookie store database error: {0}")]
Database(#[from] sqlx::Error),
#[error("cookie encryption error: {0}")]
/// 加解密错误。
#[error("sqlite cookie store crypto error: {0}")]
Crypto(#[from] CryptoError),
}
pub type CookieResult<T> = Result<T, CookieError>;
/// SQLite Cookie 存储结果。
pub type SqliteCookieStoreResult<T> = Result<T, SqliteCookieStoreError>;
/// SQLite Cookie 存储。
#[derive(Clone)]
pub struct CookieRepository {
pub struct SqliteCookieStore {
pool: SqlitePool,
secret_box: SecretBox,
}
impl CookieRepository {
impl SqliteCookieStore {
/// 构造存储。
pub fn new(pool: SqlitePool, secret_box: SecretBox) -> Self {
Self { pool, secret_box }
}
pub async fn capture_set_cookie(&self, account_id: &str, raw: &str) -> CookieResult<()> {
/// 返回底层连接池。
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
/// 检查账号是否存在。
pub async fn account_exists(&self, account_id: &str) -> SqliteCookieStoreResult<bool> {
let row = sqlx::query("select 1 from accounts where id = ?")
.bind(account_id)
.fetch_optional(&self.pool)
.await?;
Ok(row.is_some())
}
/// 捕获上游 `Set-Cookie` 响应头中允许持久化的 Cookie。
pub async fn capture_set_cookie(
&self,
account_id: &str,
raw: &str,
) -> SqliteCookieStoreResult<()> {
let Some(parsed) = parse_set_cookie(raw) else {
return Ok(());
};
@@ -39,7 +67,12 @@ impl CookieRepository {
self.upsert_cookie(account_id, parsed).await
}
pub async fn set_cookie_header(&self, account_id: &str, raw: &str) -> CookieResult<usize> {
/// 将 Cookie 请求头写入账号 Cookie 存储。
pub async fn set_cookie_header(
&self,
account_id: &str,
raw: &str,
) -> SqliteCookieStoreResult<usize> {
let parsed = parse_cookie_header(raw);
let count = parsed.len();
for cookie in parsed {
@@ -48,42 +81,23 @@ impl CookieRepository {
Ok(count)
}
async fn upsert_cookie(&self, account_id: &str, parsed: ParsedCookie) -> CookieResult<()> {
let now = Utc::now().to_rfc3339();
let value_cipher = self
.secret_box
.encrypt(&SecretString::new(parsed.value.into()))?;
sqlx::query(
"insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?) on conflict(account_id, domain, name, path) do update set value_cipher = excluded.value_cipher, expires_at = excluded.expires_at, updated_at = excluded.updated_at",
)
.bind(Uuid::new_v4().to_string())
.bind(account_id)
.bind(parsed.domain)
.bind(parsed.name)
.bind(value_cipher)
.bind(parsed.path)
.bind(parsed.expires_at)
.bind(now)
.execute(&self.pool)
.await?;
Ok(())
}
/// 为请求域名读取账号 Cookie 请求头。
pub async fn cookie_header(
&self,
account_id: &str,
request_domain: &str,
) -> CookieResult<Option<String>> {
) -> SqliteCookieStoreResult<Option<String>> {
self.cookie_header_for_request(account_id, request_domain, "/")
.await
}
/// 为请求域名和路径读取账号 Cookie 请求头。
pub async fn cookie_header_for_request(
&self,
account_id: &str,
request_domain: &str,
request_path: &str,
) -> CookieResult<Option<String>> {
) -> SqliteCookieStoreResult<Option<String>> {
let rows = sqlx::query(
"select domain, name, value_cipher, path, expires_at from account_cookies where account_id = ?",
)
@@ -101,7 +115,8 @@ impl CookieRepository {
if !path_matches(request_path, &path) {
continue;
}
if cookie_is_expired(row.get::<Option<String>, _>("expires_at").as_deref(), now) {
let expires_at = row.get::<Option<String>, _>("expires_at");
if cookie_is_expired(expires_at.as_deref(), now) {
continue;
}
let name = row.get::<String, _>("name");
@@ -132,7 +147,8 @@ impl CookieRepository {
}
}
pub async fn delete_account_cookies(&self, account_id: &str) -> CookieResult<u64> {
/// 删除账号全部 Cookie。
pub async fn delete_account_cookies(&self, account_id: &str) -> SqliteCookieStoreResult<u64> {
let result = sqlx::query("delete from account_cookies where account_id = ?")
.bind(account_id)
.execute(&self.pool)
@@ -140,27 +156,44 @@ impl CookieRepository {
Ok(result.rows_affected())
}
pub async fn cleanup_expired(&self, now: DateTime<Utc>) -> CookieResult<u64> {
let rows =
sqlx::query("select id, expires_at from account_cookies where expires_at is not null")
.fetch_all(&self.pool)
.await?;
let mut deleted = 0;
for row in rows {
let expires_at = row.get::<String, _>("expires_at");
if !cookie_is_expired(Some(&expires_at), now) {
continue;
}
let result = sqlx::query("delete from account_cookies where id = ?")
.bind(row.get::<String, _>("id"))
.execute(&self.pool)
.await?;
deleted += result.rows_affected();
}
Ok(deleted)
/// 删除已经过期的 Cookie。
pub async fn cleanup_expired(&self, now: DateTime<Utc>) -> SqliteCookieStoreResult<u64> {
let result = sqlx::query(
"delete from account_cookies where expires_at is not null and expires_at <= ?",
)
.bind(now.to_rfc3339())
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
async fn upsert_cookie(
&self,
account_id: &str,
parsed: ParsedCookie,
) -> SqliteCookieStoreResult<()> {
let now = Utc::now().to_rfc3339();
let value_cipher = self
.secret_box
.encrypt(&SecretString::new(parsed.value.into()))?;
sqlx::query(
"insert into account_cookies (id, account_id, domain, name, value_cipher, path, expires_at, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?) on conflict(account_id, domain, name, path) do update set value_cipher = excluded.value_cipher, expires_at = excluded.expires_at, updated_at = excluded.updated_at",
)
.bind(Uuid::new_v4().to_string())
.bind(account_id)
.bind(parsed.domain)
.bind(parsed.name)
.bind(value_cipher)
.bind(parsed.path)
.bind(parsed.expires_at)
.bind(now)
.execute(&self.pool)
.await?;
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ParsedCookie {
domain: String,
name: String,
@@ -169,6 +202,7 @@ struct ParsedCookie {
expires_at: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct CookieHeaderPair {
path_len: usize,
name: String,
@@ -178,7 +212,13 @@ struct CookieHeaderPair {
fn parse_set_cookie(raw: &str) -> Option<ParsedCookie> {
let mut parts = raw.split(';').map(str::trim);
let (name, value) = parts.next()?.split_once('=')?;
let mut domain = "chatgpt.com".to_string();
let name = name.trim();
let value = value.trim();
if name.is_empty() || value.is_empty() {
return None;
}
let mut domain = DEFAULT_COOKIE_DOMAIN.to_string();
let mut path = "/".to_string();
let mut expires_at = None;
for part in parts {
@@ -186,18 +226,19 @@ fn parse_set_cookie(raw: &str) -> Option<ParsedCookie> {
continue;
};
match attribute.trim().to_ascii_lowercase().as_str() {
"domain" => domain = value.trim_start_matches('.').to_string(),
"domain" => domain = value.trim().trim_start_matches('.').to_string(),
"path" => path = normalize_cookie_path(value),
"max-age" => {
if let Ok(seconds) = value.parse::<i64>() {
if let Ok(seconds) = value.trim().parse::<i64>() {
expires_at = Some(max_age_expires_at(seconds));
}
break;
}
"expires" => expires_at = Some(value.to_string()),
"expires" => expires_at = Some(value.trim().to_string()),
_ => {}
}
}
Some(ParsedCookie {
domain,
name: name.to_string(),
@@ -226,7 +267,7 @@ fn parse_cookie_header(raw: &str) -> Vec<ParsedCookie> {
return None;
}
Some(ParsedCookie {
domain: "chatgpt.com".to_string(),
domain: DEFAULT_COOKIE_DOMAIN.to_string(),
name: name.to_string(),
value: value.to_string(),
path: "/".to_string(),
@@ -243,15 +284,6 @@ fn domain_matches(request_domain: &str, cookie_domain: &str) -> bool {
.is_some_and(|prefix| prefix.ends_with('.'))
}
fn normalize_cookie_path(path: &str) -> String {
let path = path.trim();
if path.starts_with('/') {
path.to_string()
} else {
"/".to_string()
}
}
fn path_matches(request_path: &str, cookie_path: &str) -> bool {
let request_path = normalize_request_path(request_path);
let cookie_path = normalize_cookie_path(cookie_path);
@@ -273,6 +305,15 @@ fn normalize_request_path(path: &str) -> String {
}
}
fn normalize_cookie_path(path: &str) -> String {
let path = path.trim();
if path.starts_with('/') {
path.to_string()
} else {
"/".to_string()
}
}
fn cookie_is_expired(expires_at: Option<&str>, now: DateTime<Utc>) -> bool {
expires_at
.and_then(parse_cookie_expires_at)
crates/adapters/src/sqlite/events.rs
+286
View File
@@ -0,0 +1,286 @@
//! SQLite 事件日志存储。
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use sqlx::{QueryBuilder, Row, Sqlite, SqlitePool};
use thiserror::Error;
use codex_proxy_core::events::{
model::{EventLevel, EventLog},
ports::{EventLogStore, EventLogStoreError, EventLogStoreResult},
};
use codex_proxy_platform::json::{decode_cursor, encode_cursor, Page};
/// SQLite 事件日志错误。
#[derive(Debug, Error)]
pub enum SqliteEventLogStoreError {
/// 数据库错误。
#[error("sqlite event log database error: {0}")]
Database(#[from] sqlx::Error),
/// JSON 错误。
#[error("sqlite event log json error: {0}")]
Json(#[from] serde_json::Error),
/// 时间格式错误。
#[error("sqlite event log timestamp error: {0}")]
Timestamp(#[from] chrono::ParseError),
/// 事件等级非法。
#[error("invalid event level: {0}")]
InvalidLevel(String),
/// 分页游标非法。
#[error("invalid event log pagination cursor")]
InvalidCursor,
}
/// SQLite 事件日志结果。
pub type SqliteEventLogStoreResult<T> = Result<T, SqliteEventLogStoreError>;
/// 事件日志查询过滤器。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct EventLogFilter {
/// 事件类别。
pub kind: Option<String>,
/// 事件等级。
pub level: Option<EventLevel>,
/// 请求 ID。
pub request_id: Option<String>,
/// 账号 ID。
pub account_id: Option<String>,
/// 路由。
pub route: Option<String>,
/// 模型。
pub model: Option<String>,
/// HTTP 状态码。
pub status_code: Option<i64>,
/// 搜索关键词。
pub search: Option<String>,
}
/// SQLite 事件日志存储。
#[derive(Clone)]
pub struct SqliteEventLogStore {
pool: SqlitePool,
}
impl SqliteEventLogStore {
/// 构造存储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 写入事件日志。
pub async fn append(&self, event: &EventLog) -> SqliteEventLogStoreResult<()> {
append_event(&self.pool, event).await
}
/// 保留最新的指定数量事件日志。
pub async fn trim_to_capacity(&self, capacity: u32) -> SqliteEventLogStoreResult<u64> {
trim_to_capacity(&self.pool, capacity).await
}
/// 分页查询事件日志。
pub async fn list(
&self,
filter: EventLogFilter,
cursor: Option<String>,
limit: u32,
) -> SqliteEventLogStoreResult<Page<EventLog>> {
let fetch_limit = i64::from(limit) + 1;
let mut builder = QueryBuilder::<Sqlite>::new(
"select id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at from event_logs",
);
push_filter(&mut builder, &filter, cursor.as_deref())?;
builder.push(" order by created_at desc, id desc limit ");
builder.push_bind(fetch_limit);
let rows = builder.build().fetch_all(&self.pool).await?;
let has_next = rows.len() > limit as usize;
let take_count = rows.len().min(limit as usize);
let items = rows
.into_iter()
.take(take_count)
.map(|row| event_from_row(&row))
.collect::<SqliteEventLogStoreResult<Vec<_>>>()?;
let next_cursor = if has_next {
items
.last()
.map(|event| encode_cursor(&event.created_at.to_rfc3339(), &event.id))
} else {
None
};
Ok(Page { items, next_cursor })
}
/// 按 ID 读取事件日志。
pub async fn get(&self, id: &str) -> SqliteEventLogStoreResult<Option<EventLog>> {
let row = sqlx::query(
"select id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at from event_logs where id = ?",
)
.bind(id)
.fetch_optional(&self.pool)
.await?;
row.map(|row| event_from_row(&row)).transpose()
}
/// 统计事件日志数量。
pub async fn count(&self) -> SqliteEventLogStoreResult<u64> {
let count: (i64,) = sqlx::query_as("select count(*) from event_logs")
.fetch_one(&self.pool)
.await?;
Ok(count.0.max(0) as u64)
}
/// 清空事件日志。
pub async fn clear(&self) -> SqliteEventLogStoreResult<u64> {
let result = sqlx::query("delete from event_logs")
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
}
#[async_trait]
impl EventLogStore for SqliteEventLogStore {
async fn append(&self, event: &EventLog) -> EventLogStoreResult<()> {
append_event(&self.pool, event).await.map_err(|error| {
EventLogStoreError::OperationFailed {
message: error.to_string(),
}
})?;
Ok(())
}
}
async fn append_event(pool: &SqlitePool, event: &EventLog) -> SqliteEventLogStoreResult<()> {
sqlx::query(
"insert into event_logs (id, request_id, kind, level, account_id, route, model, status_code, latency_ms, message, metadata_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(&event.id)
.bind(&event.request_id)
.bind(&event.kind)
.bind(level_to_db(event.level))
.bind(&event.account_id)
.bind(&event.route)
.bind(&event.model)
.bind(event.status_code)
.bind(event.latency_ms)
.bind(&event.message)
.bind(event.metadata.to_string())
.bind(event.created_at.to_rfc3339())
.execute(pool)
.await?;
Ok(())
}
async fn trim_to_capacity(pool: &SqlitePool, capacity: u32) -> SqliteEventLogStoreResult<u64> {
let result = sqlx::query(
"delete from event_logs where id in (
select id from event_logs
order by created_at desc, id desc
limit -1 offset ?
)",
)
.bind(i64::from(capacity))
.execute(pool)
.await?;
Ok(result.rows_affected())
}
fn push_filter(
builder: &mut QueryBuilder<Sqlite>,
filter: &EventLogFilter,
cursor: Option<&str>,
) -> SqliteEventLogStoreResult<()> {
let mut separated = builder.separated(" and ");
separated.push(" where 1 = 1");
if let Some(kind) = filter.kind.as_deref() {
separated.push("kind = ");
separated.push_bind_unseparated(kind);
}
if let Some(level) = filter.level {
separated.push("level = ");
separated.push_bind_unseparated(level_to_db(level));
}
if let Some(request_id) = filter.request_id.as_deref() {
separated.push("request_id = ");
separated.push_bind_unseparated(request_id);
}
if let Some(account_id) = filter.account_id.as_deref() {
separated.push("account_id = ");
separated.push_bind_unseparated(account_id);
}
if let Some(route) = filter.route.as_deref() {
separated.push("route = ");
separated.push_bind_unseparated(route);
}
if let Some(model) = filter.model.as_deref() {
separated.push("model = ");
separated.push_bind_unseparated(model);
}
if let Some(status_code) = filter.status_code {
separated.push("status_code = ");
separated.push_bind_unseparated(status_code);
}
if let Some(search) = filter.search.as_deref() {
let pattern = format!("%{search}%");
separated.push("(message like ");
separated.push_bind_unseparated(pattern.clone());
separated.push_unseparated(" or metadata_json like ");
separated.push_bind_unseparated(pattern);
separated.push_unseparated(")");
}
if let Some(cursor) = cursor {
let (created_at, id) =
decode_cursor(cursor).ok_or(SqliteEventLogStoreError::InvalidCursor)?;
separated.push("(created_at < ");
separated.push_bind_unseparated(created_at.clone());
separated.push_unseparated(" or (created_at = ");
separated.push_bind_unseparated(created_at);
separated.push_unseparated(" and id < ");
separated.push_bind_unseparated(id);
separated.push_unseparated("))");
}
Ok(())
}
fn event_from_row(row: &sqlx::sqlite::SqliteRow) -> SqliteEventLogStoreResult<EventLog> {
Ok(EventLog {
id: row.get("id"),
request_id: row.get("request_id"),
kind: row.get("kind"),
level: level_from_db(&row.get::<String, _>("level"))?,
account_id: row.get("account_id"),
route: row.get("route"),
model: row.get("model"),
status_code: row.get("status_code"),
latency_ms: row.get("latency_ms"),
message: row.get("message"),
metadata: serde_json::from_str(&row.get::<String, _>("metadata_json"))?,
created_at: parse_rfc3339(&row.get::<String, _>("created_at"))?,
})
}
fn level_to_db(level: EventLevel) -> &'static str {
match level {
EventLevel::Debug => "debug",
EventLevel::Info => "info",
EventLevel::Warn => "warn",
EventLevel::Error => "error",
}
}
fn level_from_db(value: &str) -> SqliteEventLogStoreResult<EventLevel> {
match value {
"debug" => Ok(EventLevel::Debug),
"info" => Ok(EventLevel::Info),
"warn" => Ok(EventLevel::Warn),
"error" => Ok(EventLevel::Error),
other => Err(SqliteEventLogStoreError::InvalidLevel(other.to_string())),
}
}
fn parse_rfc3339(value: &str) -> SqliteEventLogStoreResult<DateTime<Utc>> {
Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc))
}
crates/adapters/src/sqlite/mod.rs
+22
View File
@@ -0,0 +1,22 @@
//! SQLite 适配器。
/// 账号 token 存储辅助。
pub mod account_tokens;
/// 账号用量存储。
pub mod account_usage;
/// 账号读取适配器。
pub mod accounts;
/// 管理员会话存储。
pub mod admin_sessions;
/// 客户端 API Key 存储。
pub mod client_keys;
/// Cookie 存储。
pub mod cookies;
/// 事件日志存储。
pub mod events;
/// 模型快照存储。
pub mod models;
/// 刷新租约存储。
pub mod refresh_leases;
/// 会话亲和性存储。
pub mod session_affinity;
crates/adapters/src/sqlite/models.rs
+102
View File
@@ -0,0 +1,102 @@
//! 模型快照 SQLite 仓储。
use async_trait::async_trait;
use chrono::Utc;
use sqlx::{Row, SqlitePool};
use thiserror::Error;
use codex_proxy_core::models::{
model::ModelPlanSnapshot,
ports::{ModelSnapshotStore, ModelSnapshotStoreError, ModelSnapshotStoreResult},
};
/// 模型快照仓储实现。
#[derive(Debug, Clone)]
pub struct ModelSnapshotRepository {
pool: SqlitePool,
}
/// 模型快照仓储错误。
#[derive(Debug, Error)]
pub enum ModelSnapshotRepositoryError {
/// 存储层错误。
#[error("model snapshot store error: {0}")]
Store(#[from] ModelSnapshotStoreError),
/// 数据库错误。
#[error("model snapshot database error: {0}")]
Database(#[from] sqlx::Error),
/// JSON 错误。
#[error("model snapshot json error: {0}")]
Json(#[from] serde_json::Error),
}
impl ModelSnapshotRepository {
/// 构造仓储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 直接写入单个计划快照。
pub async fn replace_plan_snapshot(
&self,
snapshot: &ModelPlanSnapshot,
) -> ModelSnapshotStoreResult<()> {
let models_json = serde_json::to_string(&snapshot.models).map_err(|error| {
ModelSnapshotStoreError::OperationFailed {
message: error.to_string(),
}
})?;
let fetched_at = Utc::now().to_rfc3339();
sqlx::query(
"insert into model_plan_snapshots (plan_type, models_json, fetched_at) values (?, ?, ?) \
on conflict(plan_type) do update set models_json = excluded.models_json, fetched_at = excluded.fetched_at",
)
.bind(&snapshot.plan_type)
.bind(models_json)
.bind(fetched_at)
.execute(&self.pool)
.await
.map_err(|error| ModelSnapshotStoreError::OperationFailed {
message: error.to_string(),
})?;
Ok(())
}
/// 列出所有快照。
pub async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult<Vec<ModelPlanSnapshot>> {
let rows = sqlx::query(
"select plan_type, models_json from model_plan_snapshots order by plan_type asc",
)
.fetch_all(&self.pool)
.await
.map_err(|error| ModelSnapshotStoreError::OperationFailed {
message: error.to_string(),
})?;
rows.into_iter()
.map(|row| {
let plan_type: String = row.get("plan_type");
let models_json: String = row.get("models_json");
let models = serde_json::from_str(&models_json).map_err(|error| {
ModelSnapshotStoreError::OperationFailed {
message: error.to_string(),
}
})?;
Ok(ModelPlanSnapshot { plan_type, models })
})
.collect()
}
}
#[async_trait]
impl ModelSnapshotStore for ModelSnapshotRepository {
async fn replace_plan_snapshot(
&self,
snapshot: &ModelPlanSnapshot,
) -> ModelSnapshotStoreResult<()> {
Self::replace_plan_snapshot(self, snapshot).await
}
async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult<Vec<ModelPlanSnapshot>> {
Self::list_plan_snapshots(self).await
}
}
crates/adapters/src/sqlite/refresh_leases.rs
+80
View File
@@ -0,0 +1,80 @@
//! SQLite 刷新租约存储。
use chrono::{DateTime, Utc};
use sqlx::SqlitePool;
use thiserror::Error;
/// SQLite 刷新租约存储结果。
pub type SqliteRefreshLeaseStoreResult<T> = Result<T, SqliteRefreshLeaseStoreError>;
/// SQLite 刷新租约存储错误。
#[derive(Debug, Error)]
pub enum SqliteRefreshLeaseStoreError {
/// SQLite 操作失败。
#[error("sqlite refresh lease query failed: {0}")]
Sqlx(#[from] sqlx::Error),
}
/// SQLite 刷新租约存储。
#[derive(Clone)]
pub struct SqliteRefreshLeaseStore {
pool: SqlitePool,
}
impl SqliteRefreshLeaseStore {
/// 构造存储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 返回底层连接池。
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
/// 尝试获取账号刷新租约。
pub async fn try_acquire(
&self,
account_id: &str,
owner: &str,
expires_at: DateTime<Utc>,
now: DateTime<Utc>,
) -> SqliteRefreshLeaseStoreResult<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,
) -> SqliteRefreshLeaseStoreResult<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)
}
}
crates/adapters/src/sqlite/session_affinity.rs
+170
View File
@@ -0,0 +1,170 @@
//! SQLite 会话亲和性存储。
use chrono::{DateTime, Duration, Utc};
use codex_proxy_core::serving::affinity::{SessionAffinityEntry, StoredSessionAffinity};
use sqlx::SqlitePool;
use thiserror::Error;
/// SQLite 会话亲和性存储。
#[derive(Clone)]
pub struct SqliteSessionAffinityStore {
pool: SqlitePool,
}
/// SQLite 会话亲和性存储错误。
#[derive(Debug, Error)]
pub enum SqliteSessionAffinityStoreError {
/// 数据库错误。
#[error("database error: {0}")]
Database(#[from] sqlx::Error),
/// 时间戳格式无效。
#[error("invalid session affinity timestamp: {0}")]
InvalidTimestamp(#[from] chrono::ParseError),
/// 函数调用 ID JSON 无效。
#[error("invalid session affinity function call ids: {0}")]
InvalidFunctionCallIds(#[from] serde_json::Error),
}
/// SQLite 会话亲和性存储结果。
pub type SqliteSessionAffinityStoreResult<T> = Result<T, SqliteSessionAffinityStoreError>;
impl SqliteSessionAffinityStore {
/// 构造存储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 返回底层连接池。
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
/// 插入或更新响应 ID 的会话亲和性条目。
pub async fn upsert(
&self,
response_id: &str,
entry: &SessionAffinityEntry,
ttl: Duration,
) -> SqliteSessionAffinityStoreResult<()> {
let function_call_ids_json = serde_json::to_string(&entry.function_call_ids)?;
let expires_at = entry
.created_at
.checked_add_signed(ttl)
.unwrap_or(DateTime::<Utc>::MAX_UTC);
sqlx::query(
r"
insert into session_affinities (
response_id,
account_id,
conversation_id,
turn_state,
instructions_hash,
input_tokens,
function_call_ids_json,
variant_hash,
expires_at,
created_at
) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
on conflict(response_id) do update set
account_id = excluded.account_id,
conversation_id = excluded.conversation_id,
turn_state = excluded.turn_state,
instructions_hash = excluded.instructions_hash,
input_tokens = excluded.input_tokens,
function_call_ids_json = excluded.function_call_ids_json,
variant_hash = excluded.variant_hash,
expires_at = excluded.expires_at,
created_at = excluded.created_at",
)
.bind(response_id)
.bind(&entry.account_id)
.bind(&entry.conversation_id)
.bind(&entry.turn_state)
.bind(&entry.instructions_hash)
.bind(entry.input_tokens.map(u64_to_i64_saturating))
.bind(function_call_ids_json)
.bind(&entry.variant_hash)
.bind(expires_at.to_rfc3339())
.bind(entry.created_at.to_rfc3339())
.execute(&self.pool)
.await?;
Ok(())
}
/// 列出当前未过期的会话亲和性条目。
pub async fn list_active(
&self,
now: DateTime<Utc>,
) -> SqliteSessionAffinityStoreResult<Vec<StoredSessionAffinity>> {
let rows = sqlx::query(
r"
select
response_id,
account_id,
conversation_id,
turn_state,
instructions_hash,
input_tokens,
function_call_ids_json,
variant_hash,
expires_at,
created_at
from session_affinities
where expires_at > ?
order by created_at asc, response_id asc",
)
.bind(now.to_rfc3339())
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| stored_session_affinity_from_row(&row))
.collect()
}
/// 删除已过期的会话亲和性条目。
pub async fn delete_expired(
&self,
now: DateTime<Utc>,
) -> SqliteSessionAffinityStoreResult<u64> {
let result = sqlx::query("delete from session_affinities where expires_at <= ?")
.bind(now.to_rfc3339())
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
}
fn stored_session_affinity_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteSessionAffinityStoreResult<StoredSessionAffinity> {
use sqlx::Row as _;
let function_call_ids_json = row.get::<String, _>("function_call_ids_json");
Ok(StoredSessionAffinity {
response_id: row.get("response_id"),
entry: SessionAffinityEntry {
account_id: row.get("account_id"),
conversation_id: row.get("conversation_id"),
turn_state: row.get("turn_state"),
instructions_hash: row.get("instructions_hash"),
input_tokens: optional_nonnegative_i64_to_u64(row.get("input_tokens")),
function_call_ids: serde_json::from_str(&function_call_ids_json)?,
variant_hash: row.get("variant_hash"),
created_at: parse_rfc3339(&row.get::<String, _>("created_at"))?,
},
expires_at: parse_rfc3339(&row.get::<String, _>("expires_at"))?,
})
}
fn parse_rfc3339(value: &str) -> SqliteSessionAffinityStoreResult<DateTime<Utc>> {
Ok(DateTime::parse_from_rfc3339(value)?.with_timezone(&Utc))
}
fn optional_nonnegative_i64_to_u64(value: Option<i64>) -> Option<u64> {
value.and_then(|value| u64::try_from(value).ok())
}
fn u64_to_i64_saturating(value: u64) -> i64 {
value.min(i64::MAX as u64) as i64
}
crates/adapters/tests/account_repository.rs
+322
View File
@@ -0,0 +1,322 @@
use chrono::{Duration, Utc};
use codex_proxy_adapters::sqlite::accounts::{NewAccount, SqliteAccountStore, UsageDelta};
use codex_proxy_core::accounts::{model::AccountStatus, ports::AccountStore};
use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite};
use secrecy::{ExposeSecret, SecretString};
#[tokio::test]
async fn account_repository_should_encrypt_tokens_and_load_decrypted_account() {
let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 4).await;
let repo = SqliteAccountStore::new(pool.clone(), secret_box);
let expires_at = Utc::now() + Duration::hours(1);
repo.insert(NewAccount {
id: "acct_a".to_string(),
email: Some("user@example.com".to_string()),
account_id: Some("chatgpt-account".to_string()),
user_id: Some("chatgpt-user".to_string()),
label: Some("primary".to_string()),
plan_type: Some("plus".to_string()),
access_token: SecretString::new("access-secret".to_string().into()),
refresh_token: Some(SecretString::new("refresh-secret".to_string().into())),
access_token_expires_at: Some(expires_at),
status: AccountStatus::Active,
})
.await
.unwrap();
let stored: (String, String) = sqlx::query_as(
"select access_token_cipher, refresh_token_cipher from accounts where id = ?",
)
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
assert!(stored.0.starts_with("v1:"));
assert!(!stored.0.contains("access-secret"));
assert!(stored.1.starts_with("v1:"));
assert!(!stored.1.contains("refresh-secret"));
let loaded = repo.get("acct_a").await.unwrap().unwrap();
assert_eq!(loaded.access_token.expose_secret(), "access-secret");
assert_eq!(
loaded.refresh_token.unwrap().expose_secret(),
"refresh-secret"
);
assert_eq!(loaded.email.as_deref(), Some("user@example.com"));
assert_eq!(loaded.status, AccountStatus::Active);
}
#[tokio::test]
async fn account_repository_should_cursor_page_accounts_by_added_at_desc() {
let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 5).await;
let repo = SqliteAccountStore::new(pool.clone(), secret_box.clone());
seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await;
seed_repo_account(&pool, &secret_box, "acct_b", "2026-06-11T00:01:00Z").await;
seed_repo_account(&pool, &secret_box, "acct_c", "2026-06-11T00:02:00Z").await;
let first_page = repo.list(None, 2).await.unwrap();
assert_eq!(
first_page
.items
.iter()
.map(|account| account.id.as_str())
.collect::<Vec<_>>(),
["acct_c", "acct_b"]
);
let second_page = repo.list(first_page.next_cursor, 2).await.unwrap();
assert_eq!(
second_page
.items
.iter()
.map(|account| account.id.as_str())
.collect::<Vec<_>>(),
["acct_a"]
);
}
#[tokio::test]
async fn account_repository_should_list_metadata_without_decrypting_tokens() {
let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 35).await;
let repo = SqliteAccountStore::new(pool.clone(), secret_box);
sqlx::query(
"insert into accounts (id, email, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)",
)
.bind("acct_corrupt")
.bind("user@example.com")
.bind("not-a-secret-box-cipher")
.bind("active")
.bind("2026-06-11T00:00:00Z")
.bind("2026-06-11T00:00:00Z")
.execute(&pool)
.await
.unwrap();
let page = repo.list_metadata(None, 10).await.unwrap();
assert_eq!(page.items[0].id, "acct_corrupt");
assert_eq!(page.items[0].email.as_deref(), Some("user@example.com"));
assert_eq!(page.items[0].status, AccountStatus::Active);
}
#[tokio::test]
async fn account_repository_should_update_status_and_label_without_rewriting_tokens() {
let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 6).await;
let repo = SqliteAccountStore::new(pool.clone(), secret_box);
repo.insert(NewAccount {
id: "acct_a".to_string(),
email: None,
account_id: None,
user_id: None,
label: None,
plan_type: None,
access_token: SecretString::new("access-secret".to_string().into()),
refresh_token: None,
access_token_expires_at: None,
status: AccountStatus::Active,
})
.await
.unwrap();
let before: (String,) = sqlx::query_as("select access_token_cipher from accounts where id = ?")
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
assert!(repo
.set_status("acct_a", AccountStatus::Disabled)
.await
.unwrap());
assert!(repo
.set_label("acct_a", Some("work".to_string()))
.await
.unwrap());
let after: (String,) = sqlx::query_as("select access_token_cipher from accounts where id = ?")
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
let loaded = repo.get("acct_a").await.unwrap().unwrap();
assert_eq!(after.0, before.0);
assert_eq!(loaded.status, AccountStatus::Disabled);
assert_eq!(loaded.label.as_deref(), Some("work"));
}
#[tokio::test]
async fn account_repository_should_accumulate_usage_counters() {
let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 7).await;
let repo = SqliteAccountStore::new(pool, secret_box);
seed_new_account(&repo, "acct_a").await;
repo.record_usage(
"acct_a",
UsageDelta {
input_tokens: 10,
output_tokens: 4,
cached_tokens: 3,
empty_response_count: 0,
..UsageDelta::default()
},
)
.await
.unwrap();
repo.record_usage(
"acct_a",
UsageDelta {
input_tokens: 8,
output_tokens: 2,
cached_tokens: 1,
empty_response_count: 1,
..UsageDelta::default()
},
)
.await
.unwrap();
let usage = repo.get_usage("acct_a").await.unwrap().unwrap();
assert_eq!(usage.request_count, 2);
assert_eq!(usage.input_tokens, 18);
assert_eq!(usage.output_tokens, 6);
assert_eq!(usage.cached_tokens, 4);
assert_eq!(usage.empty_response_count, 1);
assert!(usage.last_used_at.is_some());
}
#[tokio::test]
async fn account_repository_should_restore_window_usage_into_runtime_pool_accounts() {
let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 33).await;
let repo = SqliteAccountStore::new(pool.clone(), secret_box.clone());
let window_started_at = Utc::now() - Duration::minutes(2);
let window_reset_at = Utc::now() + Duration::minutes(3);
seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await;
sqlx::query(
"insert into account_usage (account_id, request_count, image_input_tokens, image_output_tokens, image_request_count, image_request_failed_count, window_request_count, window_input_tokens, window_output_tokens, window_cached_tokens, window_image_input_tokens, window_image_output_tokens, window_image_request_count, window_image_request_failed_count, window_started_at, window_reset_at, limit_window_seconds) values (?, 9, 31, 9, 2, 1, 7, 11, 13, 17, 19, 5, 1, 1, ?, ?, 300)",
)
.bind("acct_a")
.bind(window_started_at.to_rfc3339())
.bind(window_reset_at.to_rfc3339())
.execute(&pool)
.await
.unwrap();
let account = repo.list_pool_accounts().await.unwrap().remove(0);
assert_eq!(account.window_request_count, 7);
assert_eq!(account.image_input_tokens, 31);
assert_eq!(account.image_output_tokens, 9);
assert_eq!(account.image_request_count, 2);
assert_eq!(account.image_request_failed_count, 1);
assert_eq!(account.window_input_tokens, 11);
assert_eq!(account.window_output_tokens, 13);
assert_eq!(account.window_cached_tokens, 17);
assert_eq!(account.window_image_input_tokens, 19);
assert_eq!(account.window_image_output_tokens, 5);
assert_eq!(account.window_image_request_count, 1);
assert_eq!(account.window_image_request_failed_count, 1);
assert_eq!(account.window_started_at, Some(window_started_at));
assert_eq!(account.window_reset_at, Some(window_reset_at));
assert_eq!(account.limit_window_seconds, Some(300));
}
#[tokio::test]
async fn account_repository_should_restore_quota_verify_required_into_runtime_pool_accounts() {
let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 34).await;
let repo = SqliteAccountStore::new(pool.clone(), secret_box.clone());
seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await;
sqlx::query("update accounts set quota_verify_required = 1 where id = ?")
.bind("acct_a")
.execute(&pool)
.await
.unwrap();
let account = repo.list_pool_accounts().await.unwrap().remove(0);
assert!(account.quota_verify_required);
}
#[tokio::test]
async fn account_repository_should_update_quota_json_and_fetched_at() {
let (pool, secret_box, _dir) = sqlite_account_store_parts("accounts.sqlite", 36).await;
let repo = SqliteAccountStore::new(pool.clone(), secret_box.clone());
seed_repo_account(&pool, &secret_box, "acct_a", "2026-06-11T00:00:00Z").await;
let updated = repo
.update_quota_json("acct_a", r#"{"rate_limit":{"limit_reached":false}}"#)
.await
.unwrap();
let missing_updated = repo
.update_quota_json("missing", r#"{"rate_limit":null}"#)
.await
.unwrap();
let row: (String, Option<String>) =
sqlx::query_as("select quota_json, quota_fetched_at from accounts where id = ?")
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(
(updated, missing_updated, row.0, row.1.is_some()),
(
true,
false,
r#"{"rate_limit":{"limit_reached":false}}"#.to_string(),
true
)
);
}
async fn sqlite_account_store_parts(
db_name: &str,
key_byte: u8,
) -> (sqlx::SqlitePool, SecretBox, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join(db_name);
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
(pool, SecretBox::new([key_byte; 32]), dir)
}
async fn seed_new_account(repo: &SqliteAccountStore, id: &str) {
repo.insert(NewAccount {
id: id.to_string(),
email: None,
account_id: None,
user_id: None,
label: None,
plan_type: None,
access_token: SecretString::new("access-secret".to_string().into()),
refresh_token: None,
access_token_expires_at: None,
status: AccountStatus::Active,
})
.await
.unwrap();
}
async fn seed_repo_account(
pool: &sqlx::SqlitePool,
secret_box: &SecretBox,
id: &str,
added_at: &str,
) {
let token = SecretString::new(format!("access-{id}").into());
let cipher = secret_box.encrypt(&token).unwrap();
sqlx::query(
"insert into accounts (id, access_token_cipher, status, added_at, updated_at) values (?, ?, ?, ?, ?)",
)
.bind(id)
.bind(cipher)
.bind("active")
.bind(added_at)
.bind(added_at)
.execute(pool)
.await
.unwrap();
}
crates/adapters/tests/admin_sessions.rs
+82
View File
@@ -0,0 +1,82 @@
use chrono::{Duration, Utc};
use codex_proxy_adapters::sqlite::admin_sessions::SqliteAdminSessionStore;
use codex_proxy_platform::{identity::hash_admin_password, storage::connect_sqlite};
#[tokio::test]
async fn admin_session_store_should_create_and_load_default_admin_once() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("admin-auth.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let store = SqliteAdminSessionStore::new(pool);
let password_hash = hash_admin_password("admin-password").expect("password hash");
let created = store
.ensure_default_admin(&password_hash)
.await
.expect("default admin should be created");
let skipped = store
.ensure_default_admin("different-hash")
.await
.expect("second default admin should be skipped");
let admin = store
.load_first_admin()
.await
.expect("default admin should load")
.expect("default admin should exist");
assert!(created);
assert!(!skipped);
assert_eq!(admin.password_hash, password_hash);
}
#[tokio::test]
async fn admin_session_store_should_create_validate_and_cleanup_sessions() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("admin-sessions.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let store = SqliteAdminSessionStore::new(pool);
let password_hash = hash_admin_password("admin-password").expect("password hash");
store
.ensure_default_admin(&password_hash)
.await
.expect("default admin should be created");
let admin = store
.load_first_admin()
.await
.expect("admin should load")
.expect("admin should exist");
let now = Utc::now();
store
.create_session("sess_future", &admin.id, now + Duration::minutes(5))
.await
.expect("future session should be inserted");
store
.create_session("sess_expired", &admin.id, now - Duration::minutes(5))
.await
.expect("expired session should be inserted");
assert!(store
.validate_session("sess_future")
.await
.expect("future session should validate"));
assert!(!store
.validate_session("sess_expired")
.await
.expect("expired session should not validate"));
let deleted = store
.cleanup_expired_sessions(now)
.await
.expect("expired session should be deleted");
assert_eq!(deleted, 1);
assert!(store
.validate_session("sess_future")
.await
.expect("future session should remain"));
}
crates/adapters/tests/client_keys.rs
+89
View File
@@ -0,0 +1,89 @@
use codex_proxy_adapters::sqlite::client_keys::SqliteClientKeyStore;
use codex_proxy_core::admin::ports::ClientKeyStore;
use codex_proxy_platform::{identity::ApiKeyHasher, storage::connect_sqlite};
#[tokio::test]
async fn client_key_store_should_create_list_disable_enable_and_delete_keys() {
let (store, _dir) = client_key_store("client-keys.sqlite", 10).await;
let created = store.create("cursor").await.unwrap();
assert_eq!(created.name, "cursor");
assert!(created.plaintext.starts_with("cpr_"));
assert!(store.verify_and_touch(&created.plaintext).await.unwrap());
let first_page = store.list(None, 10).await.unwrap();
assert_eq!(first_page.items.len(), 1);
assert_eq!(first_page.items[0].name, "cursor");
assert!(first_page.items[0].last_used_at.is_some());
assert!(store.set_enabled(&created.id, false).await.unwrap());
assert!(!store.verify_and_touch(&created.plaintext).await.unwrap());
assert!(store.set_enabled(&created.id, true).await.unwrap());
assert!(store.verify_and_touch(&created.plaintext).await.unwrap());
assert!(store.delete(&created.id).await.unwrap());
assert!(!store.verify_and_touch(&created.plaintext).await.unwrap());
}
#[tokio::test]
async fn client_key_store_should_page_keys_by_created_at_desc() {
let (store, _dir) = client_key_store("client-keys-page.sqlite", 11).await;
let key_a = store.create("a").await.unwrap();
let key_b = store.create("b").await.unwrap();
let key_c = store.create("c").await.unwrap();
sqlx::query("update client_api_keys set created_at = ? where id = ?")
.bind("2026-06-11T00:00:00Z")
.bind(&key_a.id)
.execute(store.pool())
.await
.unwrap();
sqlx::query("update client_api_keys set created_at = ? where id = ?")
.bind("2026-06-11T00:01:00Z")
.bind(&key_b.id)
.execute(store.pool())
.await
.unwrap();
sqlx::query("update client_api_keys set created_at = ? where id = ?")
.bind("2026-06-11T00:02:00Z")
.bind(&key_c.id)
.execute(store.pool())
.await
.unwrap();
let first_page = store.list(None, 2).await.unwrap();
assert_eq!(
first_page
.items
.iter()
.map(|key| key.name.as_str())
.collect::<Vec<_>>(),
["c", "b"]
);
let second_page = store.list(first_page.next_cursor, 2).await.unwrap();
assert_eq!(
second_page
.items
.iter()
.map(|key| key.name.as_str())
.collect::<Vec<_>>(),
["a"]
);
}
async fn client_key_store(
db_name: &str,
key_byte: u8,
) -> (SqliteClientKeyStore, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join(db_name);
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.unwrap();
(
SqliteClientKeyStore::new(pool, ApiKeyHasher::new([key_byte; 32])),
dir,
)
}
crates/adapters/tests/codex.rs
+3641
View File
File diff suppressed because it is too large Load Diff
crates/adapters/tests/cookies.rs
+93
View File
@@ -0,0 +1,93 @@
use chrono::{Duration, Utc};
use codex_proxy_adapters::sqlite::{
accounts::{NewAccount, SqliteAccountStore},
cookies::SqliteCookieStore,
};
use codex_proxy_core::accounts::model::AccountStatus;
use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite};
use secrecy::SecretString;
use sqlx::SqlitePool;
#[tokio::test]
async fn cookie_store_should_capture_cf_clearance_and_replay_by_domain_and_path() {
let (pool, secret_box, _dir) = sqlite_cookie_store_parts("captured-cookies.sqlite").await;
seed_account(&pool, &secret_box, "acct_a").await;
let store = SqliteCookieStore::new(pool, secret_box);
store
.capture_set_cookie("acct_a", "cf_clearance=root; Domain=.chatgpt.com; Path=/")
.await
.unwrap();
store
.capture_set_cookie(
"acct_a",
"cf_clearance=codex; Domain=.chatgpt.com; Path=/codex",
)
.await
.unwrap();
assert_eq!(
store
.cookie_header_for_request("acct_a", "chatgpt.com", "/codex/responses")
.await
.unwrap(),
Some("cf_clearance=codex; cf_clearance=root".to_string())
);
assert_eq!(
store
.cookie_header_for_request("acct_a", "chatgpt.com", "/api/codex/usage")
.await
.unwrap(),
Some("cf_clearance=root".to_string())
);
}
#[tokio::test]
async fn cookie_store_should_ignore_non_capturable_set_cookie_headers() {
let (pool, secret_box, _dir) = sqlite_cookie_store_parts("ignored-cookies.sqlite").await;
seed_account(&pool, &secret_box, "acct_a").await;
let store = SqliteCookieStore::new(pool.clone(), secret_box);
store
.capture_set_cookie("acct_a", "__cf_bm=bot-session; Domain=.chatgpt.com; Path=/")
.await
.unwrap();
let stored_count: (i64,) =
sqlx::query_as("select count(*) from account_cookies where account_id = ?")
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(stored_count.0, 0);
assert_eq!(
store.cookie_header("acct_a", "chatgpt.com").await.unwrap(),
None
);
}
async fn sqlite_cookie_store_parts(name: &str) -> (SqlitePool, SecretBox, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join(name);
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
(pool, SecretBox::new([21u8; 32]), dir)
}
async fn seed_account(pool: &SqlitePool, secret_box: &SecretBox, id: &str) {
SqliteAccountStore::new(pool.clone(), secret_box.clone())
.insert(NewAccount {
id: id.to_string(),
email: Some(format!("{id}@example.com")),
account_id: Some(format!("chatgpt-{id}")),
user_id: Some(format!("user-{id}")),
label: None,
plan_type: Some("plus".to_string()),
access_token: SecretString::new(format!("access-{id}").into()),
refresh_token: None,
access_token_expires_at: Some(Utc::now() + Duration::hours(1)),
status: AccountStatus::Active,
})
.await
.unwrap();
}
crates/adapters/tests/events.rs
+71
View File
@@ -0,0 +1,71 @@
use codex_proxy_adapters::sqlite::events::{EventLogFilter, SqliteEventLogStore};
use codex_proxy_core::events::model::{EventLevel, EventLog};
use codex_proxy_platform::storage::connect_sqlite;
#[tokio::test]
async fn event_log_store_should_filter_and_cursor_page_events() {
let (store, _dir) = event_log_store("event-logs.sqlite").await;
let mut matching = EventLog::new("request", EventLevel::Error, "upstream timeout");
matching.id = "log_matching".to_string();
matching.request_id = Some("req_1".to_string());
matching.route = Some("/v1/responses".to_string());
store.append(&matching).await.unwrap();
store
.append(&EventLog::new(
"request",
EventLevel::Info,
"upstream timeout",
))
.await
.unwrap();
store
.append(&EventLog::new(
"account",
EventLevel::Error,
"upstream timeout",
))
.await
.unwrap();
let page = store
.list(
EventLogFilter {
kind: Some("request".to_string()),
level: Some(EventLevel::Error),
search: Some("timeout".to_string()),
..EventLogFilter::default()
},
None,
1,
)
.await
.unwrap();
assert_eq!(page.items.len(), 1);
assert_eq!(page.items[0].id, "log_matching");
assert_eq!(page.items[0].request_id.as_deref(), Some("req_1"));
assert_eq!(page.items[0].route.as_deref(), Some("/v1/responses"));
}
#[tokio::test]
async fn event_log_store_should_get_count_and_clear_events() {
let (store, _dir) = event_log_store("event-logs-clear.sqlite").await;
let mut event = EventLog::new("request", EventLevel::Warn, "detail");
event.id = "log_detail".to_string();
store.append(&event).await.unwrap();
let loaded = store.get("log_detail").await.unwrap().unwrap();
assert_eq!(loaded.id, "log_detail");
assert_eq!(store.count().await.unwrap(), 1);
assert_eq!(store.clear().await.unwrap(), 1);
assert_eq!(store.count().await.unwrap(), 0);
}
async fn event_log_store(db_name: &str) -> (SqliteEventLogStore, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join(db_name);
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.unwrap();
(SqliteEventLogStore::new(pool), dir)
}
crates/adapters/tests/models.rs
+93
View File
@@ -0,0 +1,93 @@
use chrono::Utc;
use codex_proxy_adapters::sqlite::{accounts::SqliteAccountStore, models::ModelSnapshotRepository};
use codex_proxy_core::accounts::{model::AccountStatus, ports::AccountStore};
use codex_proxy_core::models::model::{BackendModelEntry, ModelConfig, ModelPlanSnapshot};
use codex_proxy_platform::crypto::SecretBox;
use secrecy::SecretString;
#[tokio::test]
async fn model_snapshot_repository_should_replace_and_load_plan_snapshots() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("models.sqlite");
let pool = codex_proxy_platform::storage::connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.unwrap();
let repo = ModelSnapshotRepository::new(pool);
let snapshot = ModelPlanSnapshot::from_backend_entries(
"team",
vec![BackendModelEntry {
id: Some("gpt-team".to_string()),
name: Some("GPT Team".to_string()),
..BackendModelEntry::default()
}],
);
repo.replace_plan_snapshot(&snapshot).await.unwrap();
let loaded = repo.list_plan_snapshots().await.unwrap();
assert_eq!(loaded.len(), 1);
assert_eq!(loaded[0].plan_type, "team");
assert_eq!(loaded[0].models[0].id, "gpt-team");
assert_eq!(loaded[0].models[0].source, "backend");
}
#[test]
fn model_config_should_be_usable_from_core_models() {
let config = ModelConfig {
default_model: "gpt-5.5".to_string(),
default_reasoning_effort: None,
service_tier: None,
aliases: Default::default(),
};
assert_eq!(config.default_model, "gpt-5.5");
}
#[tokio::test]
async fn sqlite_account_store_should_list_pool_accounts() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("accounts.sqlite");
let pool = codex_proxy_platform::storage::connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let secret_box = SecretBox::new([41u8; 32]);
let now = Utc::now().to_rfc3339();
let access_token_cipher = secret_box
.encrypt(&SecretString::new("access-token".to_string().into()))
.expect("encrypt access token");
let refresh_token_cipher = secret_box
.encrypt(&SecretString::new("refresh-token".to_string().into()))
.expect("encrypt refresh token");
sqlx::query(
"insert into accounts (
id, email, account_id, user_id, label, plan_type, access_token_cipher, refresh_token_cipher,
access_token_expires_at, status, added_at, updated_at
) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind("acct_1")
.bind("user@example.com")
.bind("chatgpt-account")
.bind(Option::<String>::None)
.bind("primary")
.bind("plus")
.bind(access_token_cipher)
.bind(Some(refresh_token_cipher))
.bind(Option::<String>::None)
.bind(match AccountStatus::Active {
AccountStatus::Active => "active",
_ => unreachable!(),
})
.bind(&now)
.bind(&now)
.execute(&pool)
.await
.expect("insert account");
let store = SqliteAccountStore::new(pool, secret_box);
let accounts = store.list_pool_accounts().await.expect("list accounts");
assert_eq!(accounts.len(), 1);
assert_eq!(accounts[0].id, "acct_1");
assert_eq!(accounts[0].access_token, "access-token");
assert_eq!(accounts[0].plan_type.as_deref(), Some("plus"));
}
tests/codex_gateway/oauth_refresh.rs → crates/adapters/tests/oauth.rs
+9 -15
View File
@@ -3,12 +3,14 @@ use wiremock::{
Mock, MockServer, ResponseTemplate,
};
use codex_proxy_rs::codex::gateway::oauth::{
OAuthClient, OAuthConfig, OAuthError, OpenAiOAuthRefresher, TokenRefresher,
use codex_proxy_adapters::oauth::openai::OpenAiOAuthClient;
use codex_proxy_core::auth::{
oauth::{OAuthConfig, OAuthError},
ports::{OAuthClient, TokenRefresher},
};
#[tokio::test]
async fn openai_oauth_refresher_should_exchange_refresh_token_with_form_body() {
async fn openai_oauth_client_should_exchange_refresh_token_with_form_body() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/token"))
@@ -18,17 +20,9 @@ async fn openai_oauth_refresher_should_exchange_refresh_token_with_form_body() {
})))
.mount(&server)
.await;
let refresher = OpenAiOAuthRefresher::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
OAuthConfig {
client_id: "codex-client".to_string(),
auth_endpoint: "https://auth.openai.com/oauth/authorize".to_string(),
device_code_endpoint: "https://auth.openai.com/oauth/device/code".to_string(),
token_endpoint: format!("{}/oauth/token", server.uri()),
},
);
let client = test_oauth_client(&server);
let tokens = refresher.refresh("refresh-secret").await.unwrap();
let tokens = client.refresh("refresh-secret").await.unwrap();
assert_eq!(tokens.access_token, "new-access");
assert_eq!(tokens.refresh_token.as_deref(), Some("new-refresh"));
@@ -123,8 +117,8 @@ async fn openai_oauth_client_should_map_device_poll_pending_errors() {
assert!(body.contains("client_id=codex-client"));
}
fn test_oauth_client(server: &MockServer) -> OpenAiOAuthRefresher {
OpenAiOAuthRefresher::new(
fn test_oauth_client(server: &MockServer) -> OpenAiOAuthClient {
OpenAiOAuthClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
OAuthConfig {
client_id: "codex-client".to_string(),
crates/adapters/tests/refresh_leases.rs
+68
View File
@@ -0,0 +1,68 @@
use chrono::{Duration, TimeZone, Utc};
use codex_proxy_adapters::sqlite::{
accounts::{NewAccount, SqliteAccountStore},
refresh_leases::SqliteRefreshLeaseStore,
};
use codex_proxy_core::accounts::model::AccountStatus;
use codex_proxy_platform::{crypto::SecretBox, storage::connect_sqlite};
use secrecy::SecretString;
#[tokio::test]
async fn refresh_lease_store_should_acquire_release_and_respect_expiry() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("refresh-leases.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let secret_box = SecretBox::new([31u8; 32]);
let accounts = SqliteAccountStore::new(pool.clone(), secret_box);
accounts
.insert(NewAccount {
id: "acct-lease".to_string(),
email: None,
account_id: None,
user_id: None,
label: None,
plan_type: None,
access_token: SecretString::new("access".to_string().into()),
refresh_token: Some(SecretString::new("refresh".to_string().into())),
access_token_expires_at: None,
status: AccountStatus::Active,
})
.await
.expect("account should be inserted");
let leases = SqliteRefreshLeaseStore::new(pool);
let now = Utc.with_ymd_and_hms(2026, 6, 19, 8, 0, 0).unwrap();
let first_expires_at = now + Duration::minutes(5);
let second_expires_at = now + Duration::minutes(10);
assert!(leases
.try_acquire("acct-lease", "owner-a", first_expires_at, now)
.await
.expect("first owner should acquire"));
assert!(!leases
.try_acquire("acct-lease", "owner-b", second_expires_at, now)
.await
.expect("second owner should be blocked"));
assert!(leases
.try_acquire(
"acct-lease",
"owner-b",
second_expires_at,
first_expires_at + Duration::seconds(1),
)
.await
.expect("second owner should acquire expired lease"));
assert!(!leases
.release("acct-lease", "owner-a")
.await
.expect("wrong owner should not release"));
assert!(leases
.release("acct-lease", "owner-b")
.await
.expect("current owner should release"));
assert!(leases
.try_acquire("acct-lease", "owner-c", second_expires_at, now)
.await
.expect("lease should be acquirable after release"));
}
crates/adapters/tests/session_affinity.rs
+102
View File
@@ -0,0 +1,102 @@
use chrono::{Duration, TimeZone, Utc};
use codex_proxy_adapters::sqlite::session_affinity::SqliteSessionAffinityStore;
use codex_proxy_core::serving::affinity::SessionAffinityEntry;
use codex_proxy_platform::storage::connect_sqlite;
#[tokio::test]
async fn session_affinity_store_should_upsert_and_list_active_records() {
let dir = tempfile::tempdir().expect("tempdir");
let db = dir.path().join("session-affinity.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let store = SqliteSessionAffinityStore::new(pool);
let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap();
store
.upsert(
"resp_1",
&SessionAffinityEntry {
account_id: "acct_1".to_string(),
conversation_id: "conv_1".to_string(),
turn_state: Some("turn_1".to_string()),
instructions_hash: Some("hash_1".to_string()),
input_tokens: Some(128),
function_call_ids: vec!["call_a".to_string(), "call_b".to_string()],
variant_hash: Some("variant_1".to_string()),
created_at: now,
},
Duration::hours(4),
)
.await
.expect("affinity should be stored");
let records = store
.list_active(now + Duration::minutes(1))
.await
.expect("active affinities should load");
assert_eq!(records.len(), 1);
assert_eq!(records[0].response_id, "resp_1");
assert_eq!(records[0].entry.account_id, "acct_1");
assert_eq!(records[0].entry.function_call_ids, vec!["call_a", "call_b"]);
}
#[tokio::test]
async fn session_affinity_store_should_delete_expired_records() {
let dir = tempfile::tempdir().expect("tempdir");
let db = dir.path().join("session-affinity-cleanup.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let store = SqliteSessionAffinityStore::new(pool);
let now = Utc.with_ymd_and_hms(2026, 6, 18, 12, 0, 0).unwrap();
store
.upsert(
"resp_expired",
&SessionAffinityEntry {
account_id: "acct_1".to_string(),
conversation_id: "conv_1".to_string(),
turn_state: None,
instructions_hash: None,
input_tokens: None,
function_call_ids: Vec::new(),
variant_hash: None,
created_at: now - Duration::hours(5),
},
Duration::hours(1),
)
.await
.expect("expired affinity should be stored");
store
.upsert(
"resp_active",
&SessionAffinityEntry {
account_id: "acct_2".to_string(),
conversation_id: "conv_2".to_string(),
turn_state: None,
instructions_hash: None,
input_tokens: None,
function_call_ids: Vec::new(),
variant_hash: None,
created_at: now,
},
Duration::hours(1),
)
.await
.expect("active affinity should be stored");
let deleted = store
.delete_expired(now)
.await
.expect("expired affinities should be deleted");
let records = store
.list_active(now)
.await
.expect("active affinities should load");
assert_eq!(deleted, 1);
assert_eq!(records.len(), 1);
assert_eq!(records[0].response_id, "resp_active");
}
crates/assets/Cargo.toml
+19
View File
@@ -0,0 +1,19 @@
[package]
name = "codex-proxy-assets"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
rust-version.workspace = true
[dependencies]
axum.workspace = true
tokio.workspace = true
[dev-dependencies]
tempfile = "3.27.0"
tokio.workspace = true
tower.workspace = true
[lints]
workspace = true
crates/assets/src/headers.rs
+60
View File
@@ -0,0 +1,60 @@
//! 静态资源安全头。
use axum::http::{header, HeaderMap, HeaderValue};
const SPA_CACHE_CONTROL: &str = "no-cache";
const ASSET_CACHE_CONTROL: &str = "public, max-age=31536000, immutable";
const CONTENT_SECURITY_POLICY: &str = concat!(
"default-src 'self'; ",
"script-src 'self'; ",
"style-src 'self' 'unsafe-inline'; ",
"img-src 'self' data:; ",
"font-src 'self' data:; ",
"connect-src 'self'; ",
"base-uri 'self'; ",
"frame-ancestors 'none'"
);
/// 返回给定请求路径的缓存策略。
pub fn cache_control_for_path(path: &str) -> &'static str {
if path.starts_with("/assets/") {
ASSET_CACHE_CONTROL
} else {
SPA_CACHE_CONTROL
}
}
/// 给静态资源响应写入缓存和安全头。
pub fn apply_static_headers(headers: &mut HeaderMap, request_path: &str) {
headers.insert(
header::CACHE_CONTROL,
HeaderValue::from_static(cache_control_for_path(request_path)),
);
headers.insert(
header::X_CONTENT_TYPE_OPTIONS,
HeaderValue::from_static("nosniff"),
);
headers.insert("x-frame-options", HeaderValue::from_static("DENY"));
headers.insert("referrer-policy", HeaderValue::from_static("no-referrer"));
headers.insert(
header::CONTENT_SECURITY_POLICY,
HeaderValue::from_static(CONTENT_SECURITY_POLICY),
);
}
/// 根据文件扩展名返回常见前端资源的 content type。
pub fn content_type_for_path(path: &str) -> HeaderValue {
if path.ends_with(".html") || path == "/" {
HeaderValue::from_static("text/html; charset=utf-8")
} else if path.ends_with(".js") {
HeaderValue::from_static("text/javascript; charset=utf-8")
} else if path.ends_with(".css") {
HeaderValue::from_static("text/css; charset=utf-8")
} else if path.ends_with(".json") {
HeaderValue::from_static("application/json")
} else if path.ends_with(".svg") {
HeaderValue::from_static("image/svg+xml")
} else {
HeaderValue::from_static("application/octet-stream")
}
}
crates/assets/src/lib.rs
+4
View File
@@ -0,0 +1,4 @@
//! 前端静态资源层，提供 SPA 回退和静态文件服务辅助。
pub mod headers;
pub mod router;
crates/assets/src/router.rs
+87
View File
@@ -0,0 +1,87 @@
//! 资产路由。
use std::{
path::{Component, Path, PathBuf},
sync::Arc,
};
use axum::{
body::Body,
extract::{Path as AxumPath, State},
http::{header, StatusCode},
response::{IntoResponse, Response},
routing::get,
Router,
};
use crate::headers::{apply_static_headers, content_type_for_path};
#[derive(Clone)]
struct AssetState {
dist_dir: Arc<PathBuf>,
}
/// 构造前端静态资源路由。
pub fn spa_router(dist_dir: impl AsRef<Path>) -> Router {
let state = AssetState {
dist_dir: Arc::new(dist_dir.as_ref().to_path_buf()),
};
Router::new()
.route("/", get(serve_index))
.route("/assets/{*path}", get(serve_asset))
.fallback(serve_spa_fallback)
.with_state(state)
}
async fn serve_index(State(state): State<AssetState>) -> Response {
serve_file(&state.dist_dir.join("index.html"), "/").await
}
async fn serve_spa_fallback(State(state): State<AssetState>) -> Response {
serve_file(&state.dist_dir.join("index.html"), "/").await
}
async fn serve_asset(
State(state): State<AssetState>,
AxumPath(path): AxumPath<String>,
) -> Response {
let Some(relative_path) = safe_relative_path(&path) else {
return StatusCode::NOT_FOUND.into_response();
};
let request_path = format!("/assets/{path}");
serve_file(
&state.dist_dir.join("assets").join(relative_path),
&request_path,
)
.await
}
async fn serve_file(path: &Path, request_path: &str) -> Response {
let Ok(bytes) = tokio::fs::read(path).await else {
return StatusCode::NOT_FOUND.into_response();
};
let mut response = Response::new(Body::from(bytes));
let path_for_content_type = path.to_string_lossy();
let headers = response.headers_mut();
headers.insert(
header::CONTENT_TYPE,
content_type_for_path(&path_for_content_type),
);
apply_static_headers(headers, request_path);
response
}
fn safe_relative_path(path: &str) -> Option<PathBuf> {
let mut sanitized = PathBuf::new();
for component in Path::new(path).components() {
match component {
Component::Normal(segment) => sanitized.push(segment),
_ => return None,
}
}
(!sanitized.as_os_str().is_empty()).then_some(sanitized)
}
crates/assets/tests/assets.rs
+94
View File
@@ -0,0 +1,94 @@
use axum::{
body::{to_bytes, Body},
http::{header, Request, StatusCode},
};
use tower::ServiceExt;
#[test]
fn asset_headers_should_distinguish_spa_and_fingerprinted_assets() {
assert_eq!(
codex_proxy_assets::headers::cache_control_for_path("/"),
"no-cache"
);
assert_eq!(
codex_proxy_assets::headers::cache_control_for_path("/assets/app.abc123.js"),
"public, max-age=31536000, immutable"
);
}
#[tokio::test]
async fn asset_router_should_serve_index_and_spa_fallback() {
let dir = tempfile::tempdir().expect("temp dir");
std::fs::write(
dir.path().join("index.html"),
"<!doctype html><main id=\"app\">Codex Proxy</main>",
)
.expect("index should be written");
std::fs::create_dir(dir.path().join("assets")).expect("assets dir should be created");
let router = codex_proxy_assets::router::spa_router(dir.path());
let index = router
.clone()
.oneshot(request("/"))
.await
.expect("index response");
let fallback = router
.oneshot(request("/admin/accounts"))
.await
.expect("fallback response");
assert_eq!(index.status(), StatusCode::OK);
assert_eq!(fallback.status(), StatusCode::OK);
assert_body_contains(index, "Codex Proxy").await;
assert_body_contains(fallback, "Codex Proxy").await;
}
#[tokio::test]
async fn asset_router_should_serve_assets_with_cache_and_security_headers() {
let dir = tempfile::tempdir().expect("temp dir");
std::fs::write(dir.path().join("index.html"), "<!doctype html>").expect("write index");
std::fs::create_dir(dir.path().join("assets")).expect("assets dir should be created");
std::fs::write(
dir.path().join("assets").join("app.abc123.js"),
"window.__codexProxy = true;",
)
.expect("asset should be written");
let response = codex_proxy_assets::router::spa_router(dir.path())
.oneshot(request("/assets/app.abc123.js"))
.await
.expect("asset response");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers().get(header::CACHE_CONTROL),
Some(&header::HeaderValue::from_static(
"public, max-age=31536000, immutable"
))
);
assert_eq!(
response.headers().get(header::X_CONTENT_TYPE_OPTIONS),
Some(&header::HeaderValue::from_static("nosniff"))
);
assert_body_contains(response, "__codexProxy").await;
}
fn request(uri: &str) -> Request<Body> {
Request::builder()
.uri(uri)
.body(Body::empty())
.expect("request should build")
}
async fn assert_body_contains(response: axum::response::Response, expected: &str) {
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should collect");
let text = std::str::from_utf8(&body).expect("body should be utf8");
assert!(
text.contains(expected),
"body should contain `{expected}`, got `{text}`"
);
}
crates/core/Cargo.toml
+28
View File
@@ -0,0 +1,28 @@
[package]
name = "codex-proxy-core"
version.workspace = true
edition.workspace = true
license.workspace = true
publish.workspace = true
rust-version.workspace = true
[dependencies]
async-trait.workspace = true
base64.workspace = true
chrono.workspace = true
hex.workspace = true
indexmap.workspace = true
rand.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing.workspace = true
uuid.workspace = true
[lints]
workspace = true
[dev-dependencies]
tokio.workspace = true
crates/core/src/accounts/cloudflare.rs
+127
View File
@@ -0,0 +1,127 @@
//! Cloudflare recovery state for accounts.
use std::{collections::HashMap, sync::Arc};
use chrono::{DateTime, Duration, Utc};
use tokio::sync::RwLock;
const PATH_BLOCK_THRESHOLD: u32 = 3;
const PATH_BLOCK_STALE_AFTER: Duration = Duration::hours(1);
#[derive(Debug, Clone, Copy)]
struct PathBlockState {
count: u32,
last_at: DateTime<Utc>,
}
/// Tracks per-account Cloudflare path-block failures.
#[derive(Debug, Clone, Default)]
pub struct CloudflarePathBlockTracker {
counts: Arc<RwLock<HashMap<String, PathBlockState>>>,
}
impl CloudflarePathBlockTracker {
/// Creates an empty path-block tracker.
pub fn new() -> Self {
Self::default()
}
/// Records one path-block failure and returns the current non-stale count.
pub async fn record_path_block(&self, account_id: &str, now: DateTime<Utc>) -> u32 {
let mut counts = self.counts.write().await;
let count = counts
.get(account_id)
.filter(|state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER)
.map(|state| state.count.saturating_add(1))
.unwrap_or(1);
counts.insert(
account_id.to_string(),
PathBlockState {
count,
last_at: now,
},
);
count
}
/// Clears any tracked path-block failures for an account.
pub async fn reset(&self, account_id: &str) {
self.counts.write().await.remove(account_id);
}
/// Returns the current non-stale path-block count for an account.
pub async fn count(&self, account_id: &str, now: DateTime<Utc>) -> u32 {
self.counts
.read()
.await
.get(account_id)
.filter(|state| now.signed_duration_since(state.last_at) <= PATH_BLOCK_STALE_AFTER)
.map(|state| state.count)
.unwrap_or_default()
}
/// Returns whether the account should be disabled at the current count.
pub async fn should_disable(&self, account_id: &str, now: DateTime<Utc>) -> bool {
self.count(account_id, now).await >= PATH_BLOCK_THRESHOLD
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn record_path_block_should_disable_after_three_recent_failures() {
let tracker = CloudflarePathBlockTracker::new();
let now = Utc::now();
assert_eq!(tracker.record_path_block("acct_1", now).await, 1);
assert_eq!(
tracker
.record_path_block("acct_1", now + Duration::seconds(1))
.await,
2
);
assert_eq!(
tracker
.record_path_block("acct_1", now + Duration::seconds(2))
.await,
3
);
assert!(
tracker
.should_disable("acct_1", now + Duration::seconds(2))
.await
);
}
#[tokio::test]
async fn record_path_block_should_ignore_stale_failures() {
let tracker = CloudflarePathBlockTracker::new();
let now = Utc::now();
assert_eq!(tracker.record_path_block("acct_1", now).await, 1);
assert_eq!(
tracker
.record_path_block("acct_1", now + Duration::hours(2))
.await,
1
);
}
#[tokio::test]
async fn reset_should_clear_account_count() {
let tracker = CloudflarePathBlockTracker::new();
let now = Utc::now();
tracker.record_path_block("acct_1", now).await;
tracker
.record_path_block("acct_1", now + Duration::seconds(1))
.await;
tracker.reset("acct_1").await;
assert_eq!(tracker.count("acct_1", now + Duration::seconds(2)).await, 0);
}
}
crates/core/src/accounts/cookies.rs
+21
View File
@@ -0,0 +1,21 @@
//! Cookie 捕获与重放策略。
/// 单个 Cookie 条目。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CookieEntry {
/// Cookie 名称。
pub name: String,
/// Cookie 值。
pub value: String,
}
/// 将 Cookie 条目组装为 HTTP Cookie 头。
pub fn cookie_header(cookies: &[CookieEntry]) -> Option<String> {
let header = cookies
.iter()
.filter(|cookie| !cookie.name.trim().is_empty())
.map(|cookie| format!("{}={}", cookie.name.trim(), cookie.value))
.collect::<Vec<_>>()
.join("; ");
(!header.is_empty()).then_some(header)
}
src/codex/accounts/jwt.rs → crates/core/src/accounts/jwt.rs
+22 -2
View File
@@ -1,14 +1,19 @@
use base64::Engine;
use chrono::{DateTime, Utc};
use chrono::{DateTime, TimeZone, Utc};
use serde_json::Value;
/// JWT 过期状态。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JwtExpiry {
/// token 已过期。
Expired,
/// token 仍然有效。
Valid,
/// token 缺失、格式错误或不包含可解析的 exp。
MissingOrInvalid,
}
/// 按给定时间点判断 JWT 的 `exp` 是否已过期。
pub fn jwt_expiry(token: &str, now: DateTime<Utc>) -> JwtExpiry {
let Some(exp) = jwt_exp(token) else {
return JwtExpiry::MissingOrInvalid;
@@ -20,6 +25,11 @@ pub fn jwt_expiry(token: &str, now: DateTime<Utc>) -> JwtExpiry {
}
}
/// 读取 JWT `exp` 并转换成 UTC 时间。
pub fn jwt_expiration(token: &str) -> Option<DateTime<Utc>> {
Utc.timestamp_opt(jwt_exp(token)?, 0).single()
}
fn jwt_exp(token: &str) -> Option<i64> {
let mut parts = token.split('.');
let _header = parts.next()?;
@@ -42,7 +52,7 @@ mod tests {
use chrono::{TimeZone, Utc};
use serde_json::json;
use super::{jwt_expiry, JwtExpiry};
use super::{jwt_expiration, jwt_expiry, JwtExpiry};
#[test]
fn jwt_expiry_should_detect_expired_valid_and_invalid_tokens() {
@@ -59,6 +69,16 @@ mod tests {
assert_eq!(jwt_expiry("not-a-jwt", now), JwtExpiry::MissingOrInvalid);
}
#[test]
fn jwt_expiration_should_return_exp_as_utc_datetime() {
let expires_at = Utc.with_ymd_and_hms(2026, 6, 14, 8, 0, 0).unwrap();
assert_eq!(
jwt_expiration(&test_jwt(expires_at.timestamp())),
Some(expires_at)
);
}
fn test_jwt(exp: i64) -> String {
let header = json!({"alg": "none", "typ": "JWT"});
let payload = json!({ "exp": exp });
crates/core/src/accounts/lifecycle.rs
+20
View File
@@ -0,0 +1,20 @@
//! 账号生命周期状态转换规则。
use chrono::{DateTime, Utc};
use crate::accounts::model::{Account, AccountStatus};
/// 根据 token 过期时间推导账号运行时状态。
pub fn effective_status(account: &Account, now: DateTime<Utc>) -> AccountStatus {
if account.status != AccountStatus::Active {
return account.status;
}
if account
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now)
{
AccountStatus::Expired
} else {
AccountStatus::Active
}
}
crates/core/src/accounts/mod.rs
+20
View File
@@ -0,0 +1,20 @@
//! 账号领域中的纯逻辑辅助。
/// Cloudflare 阻断恢复状态。
pub mod cloudflare;
/// Cookie 捕获与重放策略。
pub mod cookies;
/// JWT 过期判断。
pub mod jwt;
/// 账号生命周期状态转换。
pub mod lifecycle;
/// 账号领域模型。
pub mod model;
/// 账号池调度策略。
pub mod pool;
/// 账号领域端口。
pub mod ports;
/// 账号管理用例。
pub mod service;
/// 账号用量累积策略。
pub mod usage;
src/codex/accounts/model.rs → crates/core/src/accounts/model.rs
+44
View File
@@ -1,54 +1,98 @@
//! 账号领域模型。
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
/// 账号当前状态。
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum AccountStatus {
/// 账号可正常对外服务。
Active,
/// 访问令牌已经过期，需要重新刷新。
Expired,
/// 当前配额已经耗尽。
QuotaExhausted,
/// 账号正处于刷新流程中。
Refreshing,
/// 账号被显式禁用。
Disabled,
/// 账号被上游封禁。
Banned,
}
/// 账号聚合根。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Account {
/// 账号主键。
pub id: String,
/// 展示用邮箱地址。
pub email: Option<String>,
/// 上游账号标识。
pub account_id: Option<String>,
/// 上游用户标识。
pub user_id: Option<String>,
/// 后台自定义标签。
pub label: Option<String>,
/// 订阅计划类型。
pub plan_type: Option<String>,
/// 当前访问令牌。
pub access_token: String,
/// 刷新令牌。
pub refresh_token: Option<String>,
/// 访问令牌过期时间。
pub access_token_expires_at: Option<DateTime<Utc>>,
/// 当前账号状态。
pub status: AccountStatus,
/// 是否已经触发配额封顶。
pub quota_limit_reached: bool,
/// 是否需要执行额外配额校验。
pub quota_verify_required: bool,
/// 配额冷却结束时间。
pub quota_cooldown_until: Option<DateTime<Utc>>,
/// Cloudflare 冷却结束时间。
pub cloudflare_cooldown_until: Option<DateTime<Utc>>,
/// 历史总请求数。
pub request_count: u64,
/// 历史空响应次数。
pub empty_response_count: u64,
/// 历史图片输入 token 数。
pub image_input_tokens: u64,
/// 历史图片输出 token 数。
pub image_output_tokens: u64,
/// 历史图片请求数。
pub image_request_count: u64,
/// 历史图片请求失败数。
pub image_request_failed_count: u64,
/// 当前窗口请求数。
pub window_request_count: u64,
/// 当前窗口输入 token 数。
pub window_input_tokens: u64,
/// 当前窗口输出 token 数。
pub window_output_tokens: u64,
/// 当前窗口缓存 token 数。
pub window_cached_tokens: u64,
/// 当前窗口图片输入 token 数。
pub window_image_input_tokens: u64,
/// 当前窗口图片输出 token 数。
pub window_image_output_tokens: u64,
/// 当前窗口图片请求数。
pub window_image_request_count: u64,
/// 当前窗口图片请求失败数。
pub window_image_request_failed_count: u64,
/// 当前统计窗口起始时间。
pub window_started_at: Option<DateTime<Utc>>,
/// 当前统计窗口重置时间。
pub window_reset_at: Option<DateTime<Utc>>,
/// 当前限流窗口大小（秒）。
pub limit_window_seconds: Option<u64>,
/// 首次加入时间。
pub added_at: String,
/// 最近一次使用时间。
pub last_used_at: Option<String>,
}
impl Account {
/// 构造测试默认账号。
pub fn test(id: &str, status: AccountStatus) -> Self {
Self {
id: id.to_string(),
src/codex/accounts/pool.rs → crates/core/src/accounts/pool.rs
+87 -17
View File
@@ -1,3 +1,5 @@
//! 账号池调度策略。
use std::{
cmp::Ordering,
collections::{BTreeMap, VecDeque},
@@ -5,25 +7,37 @@ use std::{
use chrono::{DateTime, Duration, Utc};
use crate::codex::accounts::{
use crate::accounts::{
jwt::{jwt_expiry, JwtExpiry},
model::{Account, AccountStatus},
service::AccountService,
};
/// 账号轮转策略。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RotationStrategy {
/// 优先选择当前窗口用量最低的账号。
LeastUsed,
/// 按候选账号顺序循环选择。
RoundRobin,
/// 优先选择最近使用过的账号。
Sticky,
}
/// 账号池配置。
#[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>>,
}
@@ -40,15 +54,21 @@ impl Default for AccountPoolOptions {
}
}
/// 账号获取请求。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AccountAcquireRequest {
/// 本次请求使用的模型。
pub model: String,
/// 本次请求需要排除的账号 ID。
pub exclude_account_ids: Vec<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(),
@@ -58,6 +78,7 @@ impl AccountAcquireRequest {
}
}
/// 设置需要排除的账号 ID。
pub fn with_exclude_account_ids(
mut self,
account_ids: impl IntoIterator<Item = impl Into<String>>,
@@ -66,49 +87,76 @@ impl AccountAcquireRequest {
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
}
}
/// 成功获取到的账号及调度元数据。
#[derive(Debug, Clone)]
pub struct AcquiredAccount {
/// 被选中的账号快照。
pub account: Account,
/// 同一账号上一个在途槽位的创建时间。
pub previous_slot_at: Option<DateTime<Utc>>,
}
/// 账号池容量摘要。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AccountCapacitySummary {
/// 单账号最大并发请求数。
pub max_concurrent_per_account: usize,
/// 所有可用账号合计槽位数。
pub total_slots: usize,
/// 当前已占用槽位数。
pub used_slots: usize,
/// 当前可用槽位数。
pub available_slots: usize,
}
/// 运行时窗口用量增量。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountWindowUsageDelta {
/// 输入 token 增量。
pub input_tokens: u64,
/// 输出 token 增量。
pub output_tokens: u64,
/// 缓存 token 增量。
pub cached_tokens: u64,
/// 图片输入 token 增量。
pub image_input_tokens: u64,
/// 图片输出 token 增量。
pub image_output_tokens: u64,
/// 是否记录一次成功图片请求。
pub image_request_succeeded: bool,
/// 是否记录一次失败图片请求。
pub image_request_failed: bool,
}
/// 账号池状态摘要。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountPoolStatusSummary {
/// 账号总数。
pub total: usize,
/// 可调度账号数。
pub active: usize,
/// token 过期账号数。
pub expired: usize,
/// 配额耗尽账号数。
pub quota_exhausted: usize,
/// 处于限流冷却的账号数。
pub rate_limited: usize,
/// 正在刷新账号数。
pub refreshing: usize,
/// 禁用账号数。
pub disabled: usize,
/// 封禁账号数。
pub banned: usize,
}
/// 纯内存账号池，负责账号调度和运行时状态维护。
#[derive(Debug)]
pub struct AccountPool {
accounts: BTreeMap<String, Account>,
@@ -125,6 +173,7 @@ impl Default for AccountPool {
}
impl AccountPool {
/// 使用指定配置创建账号池。
pub fn with_options(options: AccountPoolOptions) -> Self {
Self {
accounts: BTreeMap::new(),
@@ -135,16 +184,19 @@ impl AccountPool {
}
}
/// 插入或替换账号。
pub fn insert(&mut self, account: Account) {
self.accounts.insert(account.id.clone(), account);
}
/// 移除账号及其在途槽位。
pub fn remove(&mut self, account_id: &str) -> bool {
let removed = self.accounts.remove(account_id).is_some();
self.slots.remove(account_id);
removed
}
/// 清空账号池。
pub fn clear(&mut self) {
self.accounts.clear();
self.slots.clear();
@@ -152,10 +204,12 @@ impl AccountPool {
self.least_used_cursor = 0;
}
/// 替换模型计划 allowlist。
pub fn set_model_plan_allowlist(&mut self, allowlist: BTreeMap<String, Vec<String>>) {
self.options.model_plan_allowlist = allowlist;
}
/// 更新账号标签。
pub fn set_label(&mut self, account_id: &str, label: Option<String>) -> bool {
let Some(account) = self.accounts.get_mut(account_id) else {
return false;
@@ -164,6 +218,7 @@ impl AccountPool {
true
}
/// 更新账号状态。
pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool {
let Some(account) = self.accounts.get_mut(account_id) else {
return false;
@@ -175,6 +230,7 @@ impl AccountPool {
true
}
/// 重置账号累计和窗口用量。
pub fn reset_usage(&mut self, account_id: &str) -> bool {
let Some(account) = self.accounts.get_mut(account_id) else {
return false;
@@ -198,6 +254,7 @@ impl AccountPool {
true
}
/// 标记账号配额限流直到指定时间。
pub fn mark_quota_limited_until(
&mut self,
account_id: &str,
@@ -226,6 +283,7 @@ impl AccountPool {
true
}
/// 设置账号是否需要额外配额校验。
pub fn set_quota_verify_required(&mut self, account_id: &str, required: bool) -> bool {
let Some(account) = self.accounts.get_mut(account_id) else {
return false;
@@ -234,6 +292,7 @@ impl AccountPool {
true
}
/// 同步上游返回的限流窗口信息。
pub fn sync_rate_limit_window(
&mut self,
account_id: &str,
@@ -243,6 +302,7 @@ impl AccountPool {
self.sync_rate_limit_window_at(account_id, new_reset_at, limit_window_seconds, Utc::now())
}
/// 记录账号运行时窗口 token 用量。
pub fn record_window_token_usage(
&mut self,
account_id: &str,
@@ -290,6 +350,7 @@ impl AccountPool {
true
}
/// 标记账号处于 Cloudflare 冷却期。
pub fn set_cloudflare_cooldown_until(
&mut self,
account_id: &str,
@@ -303,11 +364,13 @@ impl AccountPool {
true
}
/// 使用当前时间为指定模型获取账号。
pub fn acquire(&mut self, model: &str) -> Option<Account> {
self.acquire_with(AccountAcquireRequest::new(model, Utc::now()))
.map(|acquired| acquired.account)
}
/// 使用完整调度请求获取账号。
pub fn acquire_with(&mut self, request: AccountAcquireRequest) -> Option<AcquiredAccount> {
self.cleanup_stale_slots(request.now);
self.refresh_account_statuses(request.now);
@@ -337,6 +400,7 @@ impl AccountPool {
})
}
/// 释放指定账号的一个在途槽位。
pub fn release(&mut self, account_id: &str) {
let Some(slots) = self.slots.get_mut(account_id) else {
return;
@@ -347,6 +411,7 @@ impl AccountPool {
}
}
/// 计算账号池容量摘要。
pub fn capacity_summary(&mut self, now: DateTime<Utc>) -> AccountCapacitySummary {
self.cleanup_stale_slots(now);
self.refresh_account_statuses(now);
@@ -354,7 +419,12 @@ impl AccountPool {
.accounts
.values()
.filter(|account| {
account.status == AccountStatus::Active && self.is_quota_available(account, now)
account.status == AccountStatus::Active
&& AccountService::quota_available_at(
account,
now,
self.options.skip_quota_limited,
)
})
.count();
let total_slots = active_accounts * self.options.max_concurrent_per_account;
@@ -377,6 +447,7 @@ impl AccountPool {
}
}
/// 计算账号池状态摘要。
pub fn status_summary(&mut self, now: DateTime<Utc>) -> AccountPoolStatusSummary {
self.cleanup_stale_slots(now);
self.refresh_account_statuses(now);
@@ -386,7 +457,13 @@ impl AccountPool {
};
for account in self.accounts.values() {
match account.status {
AccountStatus::Active if !self.is_quota_available(account, now) => {
AccountStatus::Active
if !AccountService::quota_available_at(
account,
now,
self.options.skip_quota_limited,
) =>
{
summary.rate_limited += 1;
}
AccountStatus::Active => summary.active += 1,
@@ -400,7 +477,7 @@ impl AccountPool {
summary
}
/// 获取所有配额锁定的账户ID列表（用于主动配额刷新）
/// 获取所有处于配额锁定状态的账号 ID。
pub fn list_quota_locked_accounts(&self) -> Vec<String> {
self.accounts
.values()
@@ -458,24 +535,17 @@ impl AccountPool {
fn is_base_available(&self, account: &Account, request: &AccountAcquireRequest) -> bool {
account.status == AccountStatus::Active
&& self.slot_count(&account.id) < self.options.max_concurrent_per_account
&& self.is_quota_available(account, request.now)
&& AccountService::quota_available_at(
account,
request.now,
self.options.skip_quota_limited,
)
&& self.is_model_allowed(account, &request.model)
&& !request
.exclude_account_ids
.iter()
.any(|account_id| account_id == &account.id)
&& account
.cloudflare_cooldown_until
.is_none_or(|cooldown_until| request.now >= cooldown_until)
}
fn is_quota_available(&self, account: &Account, now: DateTime<Utc>) -> bool {
if !self.options.skip_quota_limited || !account.quota_limit_reached {
return true;
}
account
.quota_cooldown_until
.is_some_and(|cooldown_until| now >= cooldown_until)
&& AccountService::cloudflare_available_at(account, request.now)
}
fn slot_count(&self, account_id: &str) -> usize {
crates/core/src/accounts/ports.rs
+100
View File
@@ -0,0 +1,100 @@
//! 账号领域端口。
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use thiserror::Error;
use crate::accounts::{
model::{Account, AccountStatus},
usage::AccountUsageDelta,
};
/// 账号存储错误。
#[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<()>;
/// 读取账号当前配额 JSON。
async fn get_quota_json(&self, _account_id: &str) -> AccountStoreResult<Option<String>> {
Ok(None)
}
/// 更新账号当前配额 JSON。
async fn update_quota_json(
&self,
_account_id: &str,
_quota_json: &str,
) -> AccountStoreResult<bool> {
Ok(false)
}
/// 同步账号当前 rate-limit 统计窗口。
async fn sync_rate_limit_window(
&self,
_account_id: &str,
_reset_at: DateTime<Utc>,
_limit_window_seconds: Option<u64>,
) -> AccountStoreResult<()> {
Ok(())
}
/// 记录账号被用于一次外部请求。
async fn record_request(&self, account_id: &str) -> AccountStoreResult<()> {
self.record_usage_delta(
account_id,
AccountUsageDelta {
requests: 1,
..AccountUsageDelta::default()
},
)
.await
}
}
crates/core/src/accounts/service.rs
+44
View File
@@ -0,0 +1,44 @@
//! 账号管理用例。
use chrono::{DateTime, Utc};
use crate::accounts::model::{Account, AccountStatus};
/// 账号服务。
#[derive(Debug, Clone, Default)]
pub struct AccountService;
impl AccountService {
/// 判断账号是否可参与请求调度。
pub fn can_serve(account: &Account) -> bool {
account.status == AccountStatus::Active && !account.quota_limit_reached
}
/// 判断账号在指定时间是否处于可服务的基础状态。
pub fn can_serve_at(account: &Account, now: DateTime<Utc>, skip_quota_limited: bool) -> bool {
account.status == AccountStatus::Active
&& Self::quota_available_at(account, now, skip_quota_limited)
&& Self::cloudflare_available_at(account, now)
}
/// 判断账号配额锁定是否已经释放。
pub fn quota_available_at(
account: &Account,
now: DateTime<Utc>,
skip_quota_limited: bool,
) -> bool {
if !skip_quota_limited || !account.quota_limit_reached {
return true;
}
account
.quota_cooldown_until
.is_some_and(|cooldown_until| now >= cooldown_until)
}
/// 判断 Cloudflare 冷却是否已经释放。
pub fn cloudflare_available_at(account: &Account, now: DateTime<Utc>) -> bool {
account
.cloudflare_cooldown_until
.is_none_or(|cooldown_until| now >= cooldown_until)
}
}
crates/core/src/accounts/usage.rs
+41
View File
@@ -0,0 +1,41 @@
//! 账号用量累积策略。
/// 账号用量增量。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountUsageDelta {
/// 请求数增量。
pub requests: u64,
/// 输入 token 增量。
pub input_tokens: u64,
/// 输出 token 增量。
pub output_tokens: u64,
/// 缓存 token 增量。
pub cached_tokens: u64,
/// 空响应数增量。
pub empty_responses: u64,
/// 图片工具输入 token 增量。
pub image_input_tokens: u64,
/// 图片工具输出 token 增量。
pub image_output_tokens: u64,
/// 图片工具成功请求数增量。
pub image_requests: u64,
/// 图片工具失败请求数增量。
pub image_request_failures: u64,
}
impl AccountUsageDelta {
/// 合并两个用量增量。
pub fn merged(self, other: Self) -> Self {
Self {
requests: self.requests + other.requests,
input_tokens: self.input_tokens + other.input_tokens,
output_tokens: self.output_tokens + other.output_tokens,
cached_tokens: self.cached_tokens + other.cached_tokens,
empty_responses: self.empty_responses + other.empty_responses,
image_input_tokens: self.image_input_tokens + other.image_input_tokens,
image_output_tokens: self.image_output_tokens + other.image_output_tokens,
image_requests: self.image_requests + other.image_requests,
image_request_failures: self.image_request_failures + other.image_request_failures,
}
}
}
crates/core/src/admin/auth.rs
+30
View File
@@ -0,0 +1,30 @@
//! 管理员认证领域服务。
/// 管理员登录请求。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AdminLoginRequest {
/// 用户名。
pub username: String,
/// 明文密码。
pub password: String,
}
/// 管理员认证服务。
#[derive(Debug, Clone)]
pub struct AdminAuthService {
default_username: String,
}
impl AdminAuthService {
/// 构造管理员认证服务。
pub fn new(default_username: impl Into<String>) -> Self {
Self {
default_username: default_username.into(),
}
}
/// 判断用户名是否匹配配置中的管理员用户名。
pub fn username_matches(&self, username: &str) -> bool {
self.default_username == username
}
}
crates/core/src/admin/client_keys.rs
+26
View File
@@ -0,0 +1,26 @@
//! 客户端 API Key 业务服务。
use std::sync::Arc;
use crate::admin::ports::{ClientKeyStore, ClientKeyStoreResult};
/// 客户端 API Key 服务。
#[derive(Clone)]
pub struct ClientKeyService {
store: Arc<dyn ClientKeyStore>,
}
impl ClientKeyService {
/// 构造服务。
pub fn new(store: Arc<dyn ClientKeyStore>) -> Self {
Self { store }
}
/// 验证客户端 API Key。
pub async fn verify(&self, plaintext: &str) -> ClientKeyStoreResult<bool> {
if !plaintext.starts_with("cpr_") {
return Ok(false);
}
self.store.verify_and_touch(plaintext).await
}
}
crates/core/src/admin/mod.rs
+6
View File
@@ -0,0 +1,6 @@
//! 管理员领域用例和端口。
pub mod auth;
pub mod client_keys;
pub mod ports;
pub mod settings;
crates/core/src/admin/ports.rs
+25
View File
@@ -0,0 +1,25 @@
//! 管理员领域端口。
use async_trait::async_trait;
use thiserror::Error;
/// 客户端 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>;
}
crates/core/src/admin/settings.rs
+339
View File
@@ -0,0 +1,339 @@
//! 管理端设置领域逻辑。
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use thiserror::Error;
const ROTATION_STRATEGIES: [&str; 3] = ["least_used", "round_robin", "sticky"];
/// 管理端可变设置。
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AdminSettings {
/// 默认模型。
pub default_model: String,
/// 默认 reasoning effort。
pub default_reasoning_effort: Option<String>,
/// 默认服务层级。
pub service_tier: Option<String>,
/// 模型别名。
pub model_aliases: BTreeMap<String, String>,
/// 是否启用 token 刷新。
pub refresh_enabled: bool,
/// token 刷新提前秒数。
pub refresh_margin_seconds: u64,
/// token 刷新并发度。
pub refresh_concurrency: u32,
/// 单账号最大并发请求数。
pub max_concurrent_per_account: usize,
/// 单账号请求间隔毫秒。
pub request_interval_ms: u64,
/// 账号轮换策略。
pub rotation_strategy: String,
/// 订阅层级优先级。
pub tier_priority: Vec<String>,
/// 配额刷新间隔分钟。
pub quota_refresh_interval_minutes: u64,
/// 配额预警阈值。
pub quota_warning_thresholds: AdminQuotaWarningThresholds,
/// 是否跳过配额耗尽账号。
pub quota_skip_exhausted: bool,
/// 是否启用日志。
pub logs_enabled: bool,
/// 日志容量。
pub logs_capacity: u32,
/// 是否捕获请求/响应体。
pub logs_capture_body: bool,
/// 用量历史保留天数。
pub usage_history_retention_days: Option<u64>,
}
impl Default for AdminSettings {
fn default() -> Self {
Self {
default_model: "gpt-5.5".to_string(),
default_reasoning_effort: None,
service_tier: None,
model_aliases: BTreeMap::new(),
refresh_enabled: true,
refresh_margin_seconds: 300,
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
rotation_strategy: "least_used".to_string(),
tier_priority: Vec::new(),
quota_refresh_interval_minutes: 5,
quota_warning_thresholds: AdminQuotaWarningThresholds::default(),
quota_skip_exhausted: true,
logs_enabled: false,
logs_capacity: 2_000,
logs_capture_body: false,
usage_history_retention_days: None,
}
}
}
/// 配额预警阈值集合。
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AdminQuotaWarningThresholds {
/// 主阈值列表。
pub primary: Vec<u8>,
/// 次阈值列表。
pub secondary: Vec<u8>,
}
impl Default for AdminQuotaWarningThresholds {
fn default() -> Self {
Self {
primary: vec![80, 90],
secondary: vec![80, 90],
}
}
}
/// 管理端设置补丁。
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AdminSettingsPatch {
/// 默认模型。
pub default_model: Option<String>,
/// 默认 reasoning effort。
pub default_reasoning_effort: Option<String>,
/// 默认服务层级。
pub service_tier: Option<String>,
/// 模型别名。
pub model_aliases: Option<BTreeMap<String, String>>,
/// 是否启用 token 刷新。
pub refresh_enabled: Option<bool>,
/// token 刷新提前秒数。
pub refresh_margin_seconds: Option<u64>,
/// token 刷新并发度。
pub refresh_concurrency: Option<u32>,
/// 单账号最大并发请求数。
pub max_concurrent_per_account: Option<usize>,
/// 单账号请求间隔毫秒。
pub request_interval_ms: Option<u64>,
/// 账号轮换策略。
pub rotation_strategy: Option<String>,
/// 订阅层级优先级。
pub tier_priority: Option<Vec<String>>,
/// 配额刷新间隔分钟。
pub quota_refresh_interval_minutes: Option<u64>,
/// 配额预警阈值。
pub quota_warning_thresholds: Option<AdminQuotaWarningThresholds>,
/// 是否跳过配额耗尽账号。
pub quota_skip_exhausted: Option<bool>,
/// 是否启用日志。
pub logs_enabled: Option<bool>,
/// 日志容量。
pub logs_capacity: Option<u32>,
/// 是否捕获请求/响应体。
pub logs_capture_body: Option<bool>,
/// 用量历史保留天数。
pub usage_history_retention_days: Option<u64>,
}
/// 设置领域服务。
#[derive(Debug, Clone, Default)]
pub struct SettingsService;
impl SettingsService {
/// 合并并验证补丁设置。
pub fn apply_patch(
current: &mut AdminSettings,
patch: AdminSettingsPatch,
) -> Result<(), SettingsServiceError> {
if let Some(default_model) = patch.default_model {
current.default_model = non_empty_string("defaultModel", default_model)?;
}
if let Some(default_reasoning_effort) = patch.default_reasoning_effort {
current.default_reasoning_effort = Some(non_empty_string(
"defaultReasoningEffort",
default_reasoning_effort,
)?);
}
if let Some(service_tier) = patch.service_tier {
current.service_tier = Some(non_empty_string("serviceTier", service_tier)?);
}
if let Some(model_aliases) = patch.model_aliases {
current.model_aliases = validate_aliases(model_aliases)?;
}
if let Some(refresh_enabled) = patch.refresh_enabled {
current.refresh_enabled = refresh_enabled;
}
if let Some(refresh_margin_seconds) = patch.refresh_margin_seconds {
current.refresh_margin_seconds = refresh_margin_seconds;
}
if let Some(refresh_concurrency) = patch.refresh_concurrency {
current.refresh_concurrency = positive_u32("refreshConcurrency", refresh_concurrency)?;
}
if let Some(max_concurrent_per_account) = patch.max_concurrent_per_account {
current.max_concurrent_per_account =
positive_usize("maxConcurrentPerAccount", max_concurrent_per_account)?;
}
if let Some(request_interval_ms) = patch.request_interval_ms {
current.request_interval_ms = request_interval_ms;
}
if let Some(rotation_strategy) = patch.rotation_strategy {
current.rotation_strategy = validate_rotation_strategy(rotation_strategy)?;
}
if let Some(tier_priority) = patch.tier_priority {
current.tier_priority = validate_tier_priority(tier_priority)?;
}
if let Some(quota_refresh_interval_minutes) = patch.quota_refresh_interval_minutes {
current.quota_refresh_interval_minutes = positive_u64(
"quotaRefreshIntervalMinutes",
quota_refresh_interval_minutes,
)?;
}
if let Some(quota_warning_thresholds) = patch.quota_warning_thresholds {
validate_thresholds(&quota_warning_thresholds)?;
current.quota_warning_thresholds = quota_warning_thresholds;
}
if let Some(quota_skip_exhausted) = patch.quota_skip_exhausted {
current.quota_skip_exhausted = quota_skip_exhausted;
}
if let Some(logs_enabled) = patch.logs_enabled {
current.logs_enabled = logs_enabled;
}
if let Some(logs_capacity) = patch.logs_capacity {
current.logs_capacity = positive_u32("logsCapacity", logs_capacity)?;
}
if let Some(logs_capture_body) = patch.logs_capture_body {
current.logs_capture_body = logs_capture_body;
}
if let Some(usage_history_retention_days) = patch.usage_history_retention_days {
current.usage_history_retention_days = Some(positive_u64(
"usageHistoryRetentionDays",
usage_history_retention_days,
)?);
}
Ok(())
}
}
/// 设置领域错误。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum SettingsServiceError {
/// 设置字段非法。
#[error("invalid setting `{field}`: {message}")]
InvalidField {
/// 字段名。
field: String,
/// 错误说明。
message: String,
},
}
impl SettingsServiceError {
/// 返回出错字段名。
pub fn field(&self) -> &str {
match self {
Self::InvalidField { field, .. } => field,
}
}
/// 返回字段错误说明。
pub fn message(&self) -> &str {
match self {
Self::InvalidField { message, .. } => message,
}
}
}
fn validate_aliases(
aliases: BTreeMap<String, String>,
) -> Result<BTreeMap<String, String>, SettingsServiceError> {
let mut validated = BTreeMap::new();
for (alias, target) in aliases {
let alias = non_empty_string("modelAliases", alias)?;
let target = non_empty_string("modelAliases", target)?;
validated.insert(alias, target);
}
Ok(validated)
}
fn validate_tier_priority(tiers: Vec<String>) -> Result<Vec<String>, SettingsServiceError> {
let mut validated = Vec::with_capacity(tiers.len());
for tier in tiers {
validated.push(non_empty_string("tierPriority", tier)?);
}
Ok(validated)
}
fn validate_rotation_strategy(strategy: String) -> Result<String, SettingsServiceError> {
let strategy = non_empty_string("rotationStrategy", strategy)?;
if ROTATION_STRATEGIES.contains(&strategy.as_str()) {
Ok(strategy)
} else {
Err(invalid_field(
"rotationStrategy",
"must be one of least_used, round_robin, sticky",
))
}
}
fn validate_thresholds(
thresholds: &AdminQuotaWarningThresholds,
) -> Result<(), SettingsServiceError> {
if thresholds.primary.is_empty() && thresholds.secondary.is_empty() {
return Err(invalid_field(
"quotaWarningThresholds",
"must include at least one threshold",
));
}
if thresholds
.primary
.iter()
.chain(thresholds.secondary.iter())
.any(|threshold| *threshold > 100)
{
return Err(invalid_field(
"quotaWarningThresholds",
"thresholds must be between 0 and 100",
));
}
Ok(())
}
fn non_empty_string(field: &str, value: String) -> Result<String, SettingsServiceError> {
let value = value.trim().to_string();
if value.is_empty() {
Err(invalid_field(field, "must not be empty"))
} else {
Ok(value)
}
}
fn positive_u32(field: &str, value: u32) -> Result<u32, SettingsServiceError> {
if value == 0 {
Err(invalid_field(field, "must be greater than 0"))
} else {
Ok(value)
}
}
fn positive_u64(field: &str, value: u64) -> Result<u64, SettingsServiceError> {
if value == 0 {
Err(invalid_field(field, "must be greater than 0"))
} else {
Ok(value)
}
}
fn positive_usize(field: &str, value: usize) -> Result<usize, SettingsServiceError> {
if value == 0 {
Err(invalid_field(field, "must be greater than 0"))
} else {
Ok(value)
}
}
fn invalid_field(field: &str, message: impl Into<String>) -> SettingsServiceError {
SettingsServiceError::InvalidField {
field: field.to_string(),
message: message.into(),
}
}
crates/core/src/auth/mod.rs
+8
View File
@@ -0,0 +1,8 @@
//! 认证领域中的纯逻辑与端口。
/// OAuth token 刷新策略。
pub mod oauth;
/// 认证上游端口定义。
pub mod ports;
/// 管理员会话领域逻辑。
pub mod session;
crates/core/src/auth/oauth.rs
+417
View File
@@ -0,0 +1,417 @@
//! OAuth 领域逻辑。
use std::{collections::BTreeMap, sync::Arc};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use chrono::{DateTime, Duration, Utc};
use rand::Rng;
use serde::Deserialize;
use sha2::{Digest, Sha256};
use thiserror::Error;
use tokio::sync::Semaphore;
use crate::accounts::model::{Account, AccountStatus};
use crate::auth::ports::TokenRefresher;
/// 一次刷新返回的 token 对。
#[derive(Debug, Clone)]
pub struct TokenPair {
/// 新的访问令牌。
pub access_token: String,
/// 可选的新刷新令牌。
pub refresh_token: Option<String>,
}
/// OAuth 客户端静态配置。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OAuthConfig {
/// OAuth 客户端 ID。
pub client_id: String,
/// 浏览器授权入口。
pub auth_endpoint: String,
/// 设备码申请入口。
pub device_code_endpoint: String,
/// Token 交换入口。
pub token_endpoint: String,
}
/// 设备码登录响应。
#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
pub struct DeviceCode {
/// 设备码。
pub device_code: String,
/// 用户输入码。
pub user_code: String,
/// 激活入口。
pub verification_uri: String,
/// 带用户码的激活入口。
pub verification_uri_complete: String,
/// 过期秒数。
pub expires_in: u64,
/// 建议轮询间隔秒数。
pub interval: u64,
}
/// OAuth 调用错误。
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum OAuthError {
/// 设备授权尚未完成。
#[error("authorization pending")]
AuthorizationPending,
/// 上游要求放慢轮询速度。
#[error("slow down")]
SlowDown,
/// 上游显式拒绝请求。
#[error("OAuth request rejected: {0}")]
Rejected(String),
/// HTTP 或解析层失败。
#[error("OAuth transport failed")]
Transport,
}
impl OAuthError {
/// 返回可用于设备码轮询响应的标准错误码。
pub fn pending_code(&self) -> Option<&'static str> {
match self {
Self::AuthorizationPending => Some("authorization_pending"),
Self::SlowDown => Some("slow_down"),
Self::Rejected(_) | Self::Transport => None,
}
}
}
/// 一次 PKCE 登录开始后返回给前端的数据。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PkceLogin {
/// 引导用户跳转的授权地址。
pub auth_url: String,
/// 回调校验状态值。
pub state: String,
}
/// 存储在服务端的 PKCE 会话。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PkceSession {
/// PKCE code verifier。
pub code_verifier: String,
/// OAuth 回调地址。
pub redirect_uri: String,
/// 登录完成后跳回的 host。
pub return_host: String,
}
#[derive(Debug, Clone)]
struct PendingPkceSession {
session: PkceSession,
created_at: DateTime<Utc>,
exchanging: bool,
}
/// 管理 PKCE 授权回调状态的内存会话仓。
#[derive(Debug)]
pub struct PkceSessionStore {
pending: BTreeMap<String, PendingPkceSession>,
completed: BTreeMap<String, DateTime<Utc>>,
ttl: Duration,
}
impl Default for PkceSessionStore {
fn default() -> Self {
Self {
pending: BTreeMap::new(),
completed: BTreeMap::new(),
ttl: Duration::minutes(5),
}
}
}
impl PkceSessionStore {
/// 开始一次新的 PKCE 登录流程。
pub fn start_login(&mut self, return_host: &str, config: &OAuthConfig) -> PkceLogin {
self.cleanup(Utc::now());
let (code_verifier, code_challenge) = generate_pkce_pair();
let state = random_state();
let redirect_uri = "http://localhost:1455/auth/callback".to_string();
let auth_url = build_auth_url(config, &redirect_uri, &state, &code_challenge);
self.pending.insert(
state.clone(),
PendingPkceSession {
session: PkceSession {
code_verifier,
redirect_uri,
return_host: return_host.to_string(),
},
created_at: Utc::now(),
exchanging: false,
},
);
PkceLogin { auth_url, state }
}
/// 尝试独占获取一条待完成的 PKCE 会话。
pub fn try_acquire(&mut self, state: &str) -> Option<PkceSession> {
self.cleanup(Utc::now());
let pending = self.pending.get_mut(state)?;
if pending.exchanging {
return None;
}
pending.exchanging = true;
Some(pending.session.clone())
}
/// 释放一次失败的会话占用。
pub fn release(&mut self, state: &str) {
if let Some(pending) = self.pending.get_mut(state) {
pending.exchanging = false;
}
}
/// 将给定 state 标记为已完成。
pub fn complete(&mut self, state: &str) {
self.pending.remove(state);
self.completed.insert(state.to_string(), Utc::now());
}
/// 判断 state 是否已经完成或正在处理中。
pub fn is_completed_or_exchanging(&mut self, state: &str) -> bool {
self.cleanup(Utc::now());
self.completed.contains_key(state)
|| self
.pending
.get(state)
.is_some_and(|pending| pending.exchanging)
}
fn cleanup(&mut self, now: DateTime<Utc>) {
self.pending
.retain(|_, pending| now - pending.created_at <= self.ttl);
self.completed
.retain(|_, completed_at| now - *completed_at <= self.ttl);
}
}
/// 刷新任务的调度策略。
#[derive(Debug, Clone, Copy)]
pub struct RefreshPolicy {
/// 提前多久开始刷新访问令牌。
pub refresh_margin_seconds: u64,
/// 允许并发执行的刷新任务数。
pub refresh_concurrency: u32,
}
/// 触发刷新动作的原因。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RefreshTrigger {
/// 在访问令牌即将过期前触发刷新。
BeforeExpiry,
/// 在上游返回未授权后立即刷新。
Unauthorized,
}
/// 上游刷新失败后的领域结果。
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
pub enum RefreshFailure {
/// 刷新令牌无效或已过期。
#[error("refresh token is invalid or expired")]
InvalidGrant,
/// 账号配额耗尽。
#[error("account quota is exhausted")]
QuotaExhausted,
/// 账号被上游封禁。
#[error("account is banned")]
Banned,
/// 账号被显式禁用。
#[error("account is disabled")]
Disabled,
/// 刷新请求在传输层失败。
#[error("refresh transport failed")]
Transport,
}
/// 调度器自身的执行错误。
#[derive(Debug, Error)]
pub enum RefreshError {
/// 并发限制信号量已关闭。
#[error("refresh task semaphore closed")]
ConcurrencyClosed,
/// 刷新请求在传输层失败。
#[error("refresh transport failed")]
Transport,
}
/// 负责执行单账号刷新策略的调度器。
#[derive(Clone)]
pub struct RefreshScheduler<C> {
policy: RefreshPolicy,
client: Arc<C>,
semaphore: Arc<Semaphore>,
}
impl<C> RefreshScheduler<C>
where
C: TokenRefresher,
{
/// 使用策略和上游刷新端口构造调度器。
pub fn new(policy: RefreshPolicy, client: C) -> Self {
let concurrency = policy.refresh_concurrency.max(1) as usize;
Self {
policy,
client: Arc::new(client),
semaphore: Arc::new(Semaphore::new(concurrency)),
}
}
/// 在给定时间点按触发原因刷新账号。
///
/// 当账号不需要刷新时返回原账号快照；当刷新失败但属于可映射的领域错误时，
/// 返回更新过状态的账号。
///
/// # Errors
///
/// 当并发控制已关闭或刷新传输失败时返回 [`RefreshError`]。
pub async fn refresh_account_at(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> Result<Account, RefreshError> {
if !self.should_refresh(account, trigger, now) {
return Ok(account.clone());
}
let Some(refresh_token) = account.refresh_token.as_deref() else {
let mut expired = account.clone();
expired.status = AccountStatus::Expired;
return Ok(expired);
};
let _permit = self
.semaphore
.acquire()
.await
.map_err(|_| RefreshError::ConcurrencyClosed)?;
match self.client.refresh(refresh_token).await {
Ok(token_pair) => Ok(apply_token_pair(account, token_pair)),
Err(RefreshFailure::Transport) => Err(RefreshError::Transport),
Err(error) => Ok(apply_refresh_failure(account, error)),
}
}
/// 判断账号在给定触发原因下是否需要刷新。
pub fn should_refresh_account_at(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> bool {
self.should_refresh(account, trigger, now)
}
fn should_refresh(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> bool {
if account.status != AccountStatus::Active {
return false;
}
match trigger {
RefreshTrigger::Unauthorized => true,
RefreshTrigger::BeforeExpiry => account
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now + self.refresh_margin()),
}
}
fn refresh_margin(&self) -> Duration {
let seconds = self.policy.refresh_margin_seconds.min(86_400 * 7) as i64;
Duration::seconds(seconds)
}
}
/// 将新的 token 对应用到账号快照上。
pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
let mut refreshed = account.clone();
refreshed.access_token = token_pair.access_token;
// 刷新响应不返回 refresh_token 时，继续保留旧值，避免永久失去刷新能力。
if let Some(refresh_token) = token_pair.refresh_token {
refreshed.refresh_token = Some(refresh_token);
}
refreshed.status = AccountStatus::Active;
refreshed
}
/// 将刷新失败映射为账号状态变更。
pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Account {
let mut updated = account.clone();
updated.status = match failure {
RefreshFailure::InvalidGrant => AccountStatus::Expired,
RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted,
RefreshFailure::Banned => AccountStatus::Banned,
RefreshFailure::Disabled => AccountStatus::Disabled,
RefreshFailure::Transport => AccountStatus::Active,
};
updated
}
fn generate_pkce_pair() -> (String, String) {
let mut bytes = [0u8; 32];
rand::rng().fill_bytes(&mut bytes);
let code_verifier = URL_SAFE_NO_PAD.encode(bytes);
let digest = Sha256::digest(code_verifier.as_bytes());
let code_challenge = URL_SAFE_NO_PAD.encode(digest);
(code_verifier, code_challenge)
}
fn random_state() -> String {
let mut bytes = [0u8; 16];
rand::rng().fill_bytes(&mut bytes);
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
}
fn build_auth_url(
config: &OAuthConfig,
redirect_uri: &str,
state: &str,
code_challenge: &str,
) -> String {
let params = [
("response_type", "code"),
("client_id", config.client_id.as_str()),
("redirect_uri", redirect_uri),
("scope", "openid profile email offline_access"),
("code_challenge", code_challenge),
("code_challenge_method", "S256"),
("id_token_add_organizations", "true"),
("codex_cli_simplified_flow", "true"),
("state", state),
("originator", "codex_cli_rs"),
];
let query = params
.iter()
.map(|(key, value)| format!("{}={}", percent_encode(key), percent_encode(value)))
.collect::<Vec<_>>()
.join("&");
format!("{}?{query}", config.auth_endpoint)
}
fn percent_encode(value: &str) -> String {
let mut encoded = String::new();
for byte in value.bytes() {
match byte {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => {
encoded.push(char::from(byte));
}
b' ' => encoded.push_str("%20"),
_ => encoded.push_str(&format!("%{byte:02X}")),
}
}
encoded
}
crates/core/src/auth/ports.rs
+30
View File
@@ -0,0 +1,30 @@
//! 认证相关的上游端口。
use async_trait::async_trait;
use crate::auth::oauth::{DeviceCode, OAuthError, RefreshFailure, TokenPair};
/// 刷新令牌的上游端口。
#[async_trait]
pub trait TokenRefresher: Send + Sync + 'static {
/// 使用给定刷新令牌换取新的 token 对。
async fn refresh(&self, refresh_token: &str) -> Result<TokenPair, RefreshFailure>;
}
/// OAuth 交互上游端口。
#[async_trait]
pub trait OAuthClient: Send + Sync + 'static {
/// 使用 PKCE 授权码交换 token 对。
async fn exchange_code(
&self,
code: &str,
code_verifier: &str,
redirect_uri: &str,
) -> Result<TokenPair, OAuthError>;
/// 请求设备码登录信息。
async fn request_device_code(&self) -> Result<DeviceCode, OAuthError>;
/// 轮询设备码换取 token 对。
async fn poll_device_token(&self, device_code: &str) -> Result<TokenPair, OAuthError>;
}
crates/core/src/auth/session.rs
+24
View File
@@ -0,0 +1,24 @@
//! 管理员会话领域逻辑。
use chrono::{DateTime, Duration, Utc};
/// 管理员会话。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AdminSession {
/// 会话 ID。
pub id: String,
/// 过期时间。
pub expires_at: DateTime<Utc>,
}
impl AdminSession {
/// 判断会话在给定时间是否已经过期。
pub fn is_expired_at(&self, now: DateTime<Utc>) -> bool {
now >= self.expires_at
}
}
/// 根据 TTL 计算会话过期时间。
pub fn session_expiry(now: DateTime<Utc>, ttl_minutes: i64) -> DateTime<Utc> {
now + Duration::minutes(ttl_minutes.max(0))
}
crates/core/src/error.rs
+23
View File
@@ -0,0 +1,23 @@
//! 核心领域通用错误。
use thiserror::Error;
/// 核心领域错误。
#[derive(Debug, Error)]
pub enum CoreError {
/// 输入不符合领域约束。
#[error("invalid domain input: {message}")]
InvalidInput {
/// 错误说明。
message: String,
},
/// 依赖端口调用失败。
#[error("port operation failed: {message}")]
Port {
/// 错误说明。
message: String,
},
}
/// 核心领域结果。
pub type CoreResult<T> = Result<T, CoreError>;
crates/core/src/events/mod.rs
+8
View File
@@ -0,0 +1,8 @@
//! 事件日志领域。
/// 事件日志模型。
pub mod model;
/// 事件日志端口。
pub mod ports;
/// 事件日志策略服务。
pub mod service;
src/codex/events/event.rs → crates/core/src/events/model.rs
+28 -18
View File
@@ -1,49 +1,59 @@
use chrono::Utc;
//! 事件日志模型。
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use serde_json::Value;
use uuid::Uuid;
/// 事件等级。
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum EventLevel {
/// 调试。
Debug,
/// 信息。
Info,
/// 警告。
Warn,
/// 错误。
Error,
}
impl EventLevel {
pub fn as_str(self) -> &'static str {
match self {
Self::Debug => "debug",
Self::Info => "info",
Self::Warn => "warn",
Self::Error => "error",
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
/// 结构化事件日志。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct EventLog {
/// 事件 ID。
pub id: String,
/// 请求 ID。
pub request_id: Option<String>,
/// 事件类别。
pub kind: String,
/// 事件等级。
pub level: EventLevel,
/// 账号 ID。
pub account_id: Option<String>,
/// HTTP 路由。
pub route: Option<String>,
/// 模型名。
pub model: Option<String>,
/// HTTP 状态码。
pub status_code: Option<i64>,
/// 请求耗时毫秒。
pub latency_ms: Option<i64>,
/// 展示消息。
pub message: String,
/// 结构化元数据。
pub metadata: Value,
pub created_at: String,
/// 创建时间。
pub created_at: DateTime<Utc>,
}
impl EventLog {
/// 构造事件日志。
pub fn new(kind: impl Into<String>, level: EventLevel, message: impl Into<String>) -> Self {
Self {
id: Uuid::new_v4().to_string(),
id: format!("log_{}", Uuid::new_v4().simple()),
request_id: None,
kind: kind.into(),
level,
@@ -53,8 +63,8 @@ impl EventLog {
status_code: None,
latency_ms: None,
message: message.into(),
metadata: json!({}),
created_at: Utc::now().to_rfc3339(),
metadata: Value::Object(Default::default()),
created_at: Utc::now(),
}
}
}
crates/core/src/events/ports.rs
+27
View File
@@ -0,0 +1,27 @@
//! 事件日志端口。
use async_trait::async_trait;
use thiserror::Error;
use crate::events::model::EventLog;
/// 事件日志存储错误。
#[derive(Debug, Error)]
pub enum EventLogStoreError {
/// 底层存储失败。
#[error("event log store operation failed: {message}")]
OperationFailed {
/// 错误说明。
message: String,
},
}
/// 事件日志存储结果。
pub type EventLogStoreResult<T> = Result<T, EventLogStoreError>;
/// 事件日志存储端口。
#[async_trait]
pub trait EventLogStore: Send + Sync + 'static {
/// 写入事件日志。
async fn append(&self, event: &EventLog) -> EventLogStoreResult<()>;
}
crates/core/src/events/service.rs
+21
View File
@@ -0,0 +1,21 @@
//! 事件日志策略服务。
use crate::events::model::{EventLevel, EventLog};
/// 事件日志服务。
#[derive(Debug, Clone)]
pub struct EventLogService {
enabled: bool,
}
impl EventLogService {
/// 构造事件日志服务。
pub fn new(enabled: bool) -> Self {
Self { enabled }
}
/// 判断事件是否应该记录。
pub fn should_record(&self, event: &EventLog) -> bool {
self.enabled || event.level == EventLevel::Error
}
}
src/codex/gateway/conversation_identity.rs → crates/core/src/gateway/conversation.rs
+3 -3
View File
@@ -7,7 +7,7 @@ use serde_json::Value;
use sha2::{Digest, Sha256};
use uuid::Uuid;
use crate::codex::gateway::transport::types::CodexResponsesRequest;
use crate::protocol::codex::responses::CodexResponsesRequest;
const LEADING_SYSTEM_REMINDER_OPEN: &str = "<system-reminder>";
const LEADING_SYSTEM_REMINDER_CLOSE: &str = "</system-reminder>";
@@ -142,10 +142,10 @@ fn normalize_conversation_anchor_text(text: &str) -> String {
rest.to_string()
}
/// 构建账号作用域的身份哈希
/// 构建账号作用域的身份哈希。
///
/// 格式：`{prefix}_{sha256(kind + "\0" + account_scope + "\0" + client_value)[..32]}`
/// 其中 prefix 为 "cp"（conversation）或 "cw"（window）
/// 其中 prefix 为 `cp`（conversation）或 `cw`（window）。
fn build_account_scoped_identity(kind: &str, account_scope: &str, client_value: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(kind.as_bytes());
src/codex/gateway/fingerprint/model.rs → crates/core/src/gateway/fingerprint.rs
+16
View File
@@ -1,20 +1,31 @@
use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
/// 上游请求指纹。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct Fingerprint {
/// 客户端来源名。
pub originator: String,
/// 应用版本。
pub app_version: String,
/// 构建号。
pub build_number: String,
/// 平台名。
pub platform: String,
/// 架构名。
pub arch: String,
/// Chromium 主版本。
pub chromium_version: String,
/// User-Agent 模板。
pub user_agent_template: String,
/// 默认请求头。
pub default_headers: IndexMap<String, String>,
/// 请求头顺序。
pub header_order: Vec<String>,
}
impl Fingerprint {
/// 返回默认的 Codex Desktop 指纹。
pub fn default_codex_desktop() -> Self {
Self {
originator: "Codex Desktop".to_string(),
@@ -29,10 +40,12 @@ impl Fingerprint {
}
}
/// 返回测试用默认指纹。
pub fn default_for_tests() -> Self {
Self::default_codex_desktop()
}
/// 返回默认请求头集合。
pub fn default_headers() -> IndexMap<String, String> {
let mut headers = IndexMap::new();
headers.insert(
@@ -48,6 +61,7 @@ impl Fingerprint {
headers
}
/// 返回默认请求头顺序。
pub fn default_header_order() -> Vec<String> {
vec![
"authorization".to_string(),
@@ -73,6 +87,7 @@ impl Fingerprint {
]
}
/// 根据模板展开最终 User-Agent。
pub fn user_agent(&self) -> String {
self.user_agent_template
.replace("{version}", &self.app_version)
@@ -80,6 +95,7 @@ impl Fingerprint {
.replace("{arch}", &self.arch)
}
/// 生成 `sec-ch-ua` 头值。
pub fn sec_ch_ua(&self) -> String {
format!(
"\"Chromium\";v=\"{}\", \"Not:A-Brand\";v=\"24\"",
crates/core/src/gateway/installation.rs
+50
View File
@@ -0,0 +1,50 @@
//! Installation ID 的纯规则。
//!
//! 该模块只负责 UUID 格式校验和生成。`~/.codex/installation_id` 与数据目录
//! 持久化属于平台文件系统职责，放在 `platform::storage` 中。
use uuid::Uuid;
/// 生成新的 installation ID。
pub fn generate_installation_id() -> String {
Uuid::new_v4().to_string()
}
/// 解析并规范化 installation ID。
///
/// 输入会先去除首尾空白；非 UUID 字符串返回 `None`。
pub fn parse_installation_id(raw: &str) -> Option<String> {
let trimmed = raw.trim();
Uuid::parse_str(trimmed).ok()?;
Some(trimmed.to_string())
}
#[cfg(test)]
mod tests {
use uuid::Uuid;
use super::*;
#[test]
fn parse_installation_id_should_trim_and_accept_uuid() {
let id = Uuid::new_v4().to_string();
let parsed = parse_installation_id(&format!(" {id}\n"));
assert_eq!(parsed, Some(id));
}
#[test]
fn parse_installation_id_should_reject_non_uuid() {
let parsed = parse_installation_id("not-a-uuid");
assert_eq!(parsed, None);
}
#[test]
fn generate_installation_id_should_return_uuid() {
let generated = generate_installation_id();
assert!(Uuid::parse_str(&generated).is_ok());
}
}
crates/core/src/gateway/mod.rs
+10
View File
@@ -0,0 +1,10 @@
//! 上游网关领域中的纯逻辑辅助。
/// 会话与缓存键派生。
pub mod conversation;
/// 上游指纹。
pub mod fingerprint;
/// 安装 ID 生成。
pub mod installation;
/// 上游网关端口。
pub mod ports;
crates/core/src/gateway/ports.rs
+42
View File
@@ -0,0 +1,42 @@
//! 上游 Codex 网关端口。
use async_trait::async_trait;
use thiserror::Error;
use crate::models::model::BackendModelEntry;
/// 拉取上游模型目录时的请求上下文。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CodexModelCatalogRequest<'a> {
/// 当前账号访问令牌。
pub access_token: &'a str,
/// 上游账号 ID。
pub account_id: Option<&'a str>,
/// 请求 ID。
pub request_id: &'a str,
/// Codex installation id。
pub installation_id: Option<&'a str>,
/// 订阅计划类型。
pub plan_type: &'a str,
}
/// 上游模型目录客户端错误。
#[derive(Debug, Error)]
pub enum CodexModelCatalogClientError {
/// 上游请求失败。
#[error("model catalog request failed: {message}")]
RequestFailed {
/// 错误说明。
message: String,
},
}
/// 上游模型目录客户端。
#[async_trait]
pub trait CodexModelCatalogClient: Send + Sync + 'static {
/// 读取当前账号可见的上游模型目录。
async fn fetch_models(
&self,
request: &CodexModelCatalogRequest<'_>,
) -> Result<Vec<BackendModelEntry>, CodexModelCatalogClientError>;
}
crates/core/src/lib.rs
+14
View File
@@ -0,0 +1,14 @@
#![deny(missing_docs)]
//! 核心领域层，承载纯业务模型、协议语义和策略。
pub mod accounts;
pub mod admin;
pub mod auth;
pub mod error;
pub mod events;
pub mod gateway;
pub mod models;
pub mod protocol;
pub mod serving;
pub mod usage;
src/codex/models/catalog.rs → crates/core/src/models/catalog.rs
+25 -112
View File
@@ -1,67 +1,16 @@
//! 模型目录聚合与纯解析规则。
use std::collections::{BTreeMap, BTreeSet};
use serde::{Deserialize, Serialize};
use crate::config::ModelConfig;
use crate::models::model::{
BackendModelEntry, BackendReasoningEffort, CodexModelInfo, ModelConfig, ModelDebugEntry,
ModelPlanSnapshot, ModelStoreDebug, ParsedModelName, ReasoningEffortInfo,
};
const SERVICE_TIER_SUFFIXES: [&str; 2] = ["fast", "flex"];
const REASONING_EFFORT_SUFFIXES: [&str; 6] = ["none", "minimal", "low", "medium", "high", "xhigh"];
const STATIC_REASONING_EFFORTS: [&str; 4] = ["low", "medium", "high", "xhigh"];
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ReasoningEffortInfo {
pub reasoning_effort: String,
pub description: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CodexModelInfo {
pub id: String,
pub display_name: String,
pub description: String,
pub is_default: bool,
pub supported_reasoning_efforts: Vec<ReasoningEffortInfo>,
pub default_reasoning_effort: String,
pub input_modalities: Vec<String>,
pub output_modalities: Vec<String>,
pub supports_personality: bool,
pub upgrade: Option<String>,
pub source: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub context_window: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_context_window: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_output_tokens: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub truncation_policy_limit: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedModelName {
pub model_id: String,
pub reasoning_effort: Option<String>,
pub service_tier: Option<String>,
}
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ModelStoreDebug {
pub total_models: usize,
pub static_models: usize,
pub alias_count: usize,
pub models: Vec<ModelDebugEntry>,
}
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ModelDebugEntry {
pub id: String,
pub source: String,
}
/// 模型目录聚合。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelCatalog {
models: Vec<CodexModelInfo>,
@@ -70,57 +19,8 @@ pub struct ModelCatalog {
model_plan_index: BTreeMap<String, Vec<String>>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendReasoningEffort {
#[serde(default, rename = "reasoning_effort", alias = "reasoningEffort")]
pub reasoning_effort: Option<String>,
#[serde(default)]
pub effort: Option<String>,
#[serde(default)]
pub description: Option<String>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendTruncationPolicy {
pub limit: Option<u64>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendModelEntry {
pub slug: Option<String>,
pub id: Option<String>,
pub name: Option<String>,
pub display_name: Option<String>,
pub description: Option<String>,
pub is_default: Option<bool>,
pub default_reasoning_effort: Option<String>,
pub default_reasoning_level: Option<String>,
#[serde(default)]
pub supported_reasoning_efforts: Vec<BackendReasoningEffort>,
#[serde(default)]
pub supported_reasoning_levels: Vec<BackendReasoningEffort>,
pub input_modalities: Option<Vec<String>>,
pub output_modalities: Option<Vec<String>>,
pub supports_personality: Option<bool>,
pub upgrade: Option<String>,
#[serde(alias = "contextWindow")]
pub context_window: Option<u64>,
#[serde(alias = "maxContextWindow")]
pub max_context_window: Option<u64>,
#[serde(alias = "maxOutputTokens")]
pub max_output_tokens: Option<u64>,
#[serde(alias = "truncationPolicy")]
pub truncation_policy: Option<BackendTruncationPolicy>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ModelPlanSnapshot {
pub plan_type: String,
pub models: Vec<CodexModelInfo>,
}
impl ModelPlanSnapshot {
/// 将上游模型条目标准化为可持久化快照。
pub fn from_backend_entries(
plan_type: impl Into<String>,
entries: Vec<BackendModelEntry>,
@@ -133,6 +33,7 @@ impl ModelPlanSnapshot {
}
impl ModelCatalog {
/// 从静态配置构造模型目录。
pub fn from_config(config: &ModelConfig) -> Self {
let default_model = config.default_model.trim().to_string();
let aliases = normalize_aliases(&config.aliases);
@@ -144,6 +45,7 @@ impl ModelCatalog {
}
}
/// 从静态配置和后端快照构造模型目录。
pub fn from_config_and_snapshots(
config: &ModelConfig,
snapshots: &[ModelPlanSnapshot],
@@ -159,7 +61,7 @@ impl ModelCatalog {
for snapshot in snapshots {
for model in &snapshot.models {
// 后端模型是按账号 plan 返回的；catalog 需要去重展示，调度层仍要保留 model -> plans。
// 后端模型按账号计划返回；目录对外去重，但仍保留 model -> plans 映射。
models_by_id.insert(model.id.clone(), model.clone());
let plans = model_plan_index.entry(model.id.clone()).or_default();
if !plans.contains(&snapshot.plan_type) {
@@ -181,10 +83,12 @@ impl ModelCatalog {
}
}
/// 返回对外可见的模型列表副本。
pub fn models(&self) -> Vec<CodexModelInfo> {
self.models.clone()
}
/// 按模型 ID 查询模型信息。
pub fn model_info(&self, model_id: &str) -> Option<CodexModelInfo> {
self.models
.iter()
@@ -192,6 +96,7 @@ impl ModelCatalog {
.cloned()
}
/// 返回模型目录调试视图。
pub fn debug(&self) -> ModelStoreDebug {
ModelStoreDebug {
total_models: self.models.len(),
@@ -208,10 +113,12 @@ impl ModelCatalog {
}
}
/// 返回 model -> plans allowlist。
pub fn model_plan_allowlist(&self) -> BTreeMap<String, Vec<String>> {
self.model_plan_index.clone()
}
/// 判断输入模型名是否在目录中可识别。
pub fn is_recognized_model_name(&self, input: &str) -> bool {
let trimmed = input.trim();
if trimmed.is_empty() {
@@ -231,6 +138,7 @@ impl ModelCatalog {
|| self.model_info(&stripped.model_name).is_some()
}
/// 解析外部传入的模型名，提取别名、推理强度和服务层级后缀。
pub fn parse_model_name(&self, input: &str) -> ParsedModelName {
let trimmed = input.trim();
if self.aliases.contains_key(trimmed) || self.model_info(trimmed).is_some() {
@@ -249,6 +157,7 @@ impl ModelCatalog {
}
}
/// 将标准化模型名重新拼成展示名。
pub fn build_display_model_name(parsed: &ParsedModelName) -> String {
let mut name = parsed.model_id.clone();
if let Some(reasoning_effort) = &parsed.reasoning_effort {
@@ -361,9 +270,13 @@ fn normalize_backend_model(raw: BackendModelEntry) -> CodexModelInfo {
let id = first_non_empty([raw.slug.as_deref(), raw.id.as_deref(), raw.name.as_deref()])
.unwrap_or("unknown")
.to_string();
let display_name = first_non_empty([raw.display_name.as_deref(), raw.name.as_deref()])
.unwrap_or(&id)
.to_string();
let display_name = first_non_empty([
raw.display_name.as_deref(),
raw.title.as_deref(),
raw.name.as_deref(),
])
.unwrap_or(&id)
.to_string();
let reasoning_entries = if raw.supported_reasoning_efforts.is_empty() {
raw.supported_reasoning_levels
} else {
crates/core/src/models/mod.rs
+10
View File
@@ -0,0 +1,10 @@
//! 模型目录领域。
/// 模型目录聚合与纯解析规则。
pub mod catalog;
/// 模型目录配置与纯数据类型。
pub mod model;
/// 模型快照存储端口。
pub mod ports;
/// 模型目录服务用例。
pub mod service;
crates/core/src/models/model.rs
+183
View File
@@ -0,0 +1,183 @@
//! 模型目录领域中的纯数据类型。
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
/// 模型目录默认值与别名配置。
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct ModelConfig {
/// 默认模型名。
pub default_model: String,
/// 默认推理强度。
pub default_reasoning_effort: Option<String>,
/// 默认服务层级。
pub service_tier: Option<String>,
/// 模型别名映射。
pub aliases: BTreeMap<String, String>,
}
/// 单个 reasoning effort 的标准化展示信息。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ReasoningEffortInfo {
/// 推理强度名。
pub reasoning_effort: String,
/// 展示描述。
pub description: String,
}
/// 对外暴露的单个模型目录条目。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CodexModelInfo {
/// 模型唯一 ID。
pub id: String,
/// 展示名。
pub display_name: String,
/// 描述文本。
pub description: String,
/// 是否默认模型。
pub is_default: bool,
/// 支持的推理强度列表。
pub supported_reasoning_efforts: Vec<ReasoningEffortInfo>,
/// 默认推理强度。
pub default_reasoning_effort: String,
/// 输入模态。
pub input_modalities: Vec<String>,
/// 输出模态。
pub output_modalities: Vec<String>,
/// 是否支持 personality。
pub supports_personality: bool,
/// 升级提示。
pub upgrade: Option<String>,
/// 来源标记。
pub source: String,
/// 当前上下文窗口。
#[serde(skip_serializing_if = "Option::is_none")]
pub context_window: Option<u64>,
/// 最大上下文窗口。
#[serde(skip_serializing_if = "Option::is_none")]
pub max_context_window: Option<u64>,
/// 最大输出 token 数。
#[serde(skip_serializing_if = "Option::is_none")]
pub max_output_tokens: Option<u64>,
/// 截断策略限制。
#[serde(skip_serializing_if = "Option::is_none")]
pub truncation_policy_limit: Option<u64>,
}
/// 解析模型名后得到的标准化结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedModelName {
/// 最终模型 ID。
pub model_id: String,
/// 推理强度后缀。
pub reasoning_effort: Option<String>,
/// 服务层级后缀。
pub service_tier: Option<String>,
}
/// 模型目录调试视图。
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ModelStoreDebug {
/// 模型总数。
pub total_models: usize,
/// 静态模型数。
pub static_models: usize,
/// 别名数量。
pub alias_count: usize,
/// 调试模型条目。
pub models: Vec<ModelDebugEntry>,
}
/// 单个调试模型条目。
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ModelDebugEntry {
/// 模型 ID。
pub id: String,
/// 来源标记。
pub source: String,
}
/// 按计划类型持久化的模型快照。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct ModelPlanSnapshot {
/// 订阅计划类型。
pub plan_type: String,
/// 该计划可见的模型列表。
pub models: Vec<CodexModelInfo>,
}
/// 上游声明的 reasoning effort 条目。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendReasoningEffort {
/// reasoning_effort 字段。
#[serde(default, rename = "reasoning_effort", alias = "reasoningEffort")]
pub reasoning_effort: Option<String>,
/// effort 字段的兼容别名。
#[serde(default)]
pub effort: Option<String>,
/// 展示描述。
#[serde(default)]
pub description: Option<String>,
}
/// 上游截断策略定义。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendTruncationPolicy {
/// 上下文截断限制。
pub limit: Option<u64>,
}
/// 上游 `/codex/models` 返回的原始模型条目。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct BackendModelEntry {
/// slug。
pub slug: Option<String>,
/// id。
pub id: Option<String>,
/// name。
pub name: Option<String>,
/// display_name。
pub display_name: Option<String>,
/// title。
pub title: Option<String>,
/// description。
pub description: Option<String>,
/// 是否默认模型。
pub is_default: Option<bool>,
/// 默认 reasoning effort。
pub default_reasoning_effort: Option<String>,
/// 默认 reasoning level。
pub default_reasoning_level: Option<String>,
/// 支持的 reasoning effort 列表。
#[serde(default)]
pub supported_reasoning_efforts: Vec<BackendReasoningEffort>,
/// 支持的 reasoning level 列表。
#[serde(default)]
pub supported_reasoning_levels: Vec<BackendReasoningEffort>,
/// 输入模态。
pub input_modalities: Option<Vec<String>>,
/// 输出模态。
pub output_modalities: Option<Vec<String>>,
/// 是否支持 personality。
pub supports_personality: Option<bool>,
/// 升级提示。
pub upgrade: Option<String>,
/// contextWindow 别名。
#[serde(alias = "contextWindow")]
pub context_window: Option<u64>,
/// maxContextWindow 别名。
#[serde(alias = "maxContextWindow")]
pub max_context_window: Option<u64>,
/// maxOutputTokens 别名。
#[serde(alias = "maxOutputTokens")]
pub max_output_tokens: Option<u64>,
/// truncationPolicy 别名。
#[serde(alias = "truncationPolicy")]
pub truncation_policy: Option<BackendTruncationPolicy>,
}
crates/core/src/models/ports.rs
+33
View File
@@ -0,0 +1,33 @@
//! 模型快照存储端口。
use async_trait::async_trait;
use thiserror::Error;
use crate::models::model::ModelPlanSnapshot;
/// 模型快照存储错误。
#[derive(Debug, Error)]
pub enum ModelSnapshotStoreError {
/// 底层存储失败。
#[error("model snapshot store operation failed: {message}")]
OperationFailed {
/// 错误说明。
message: String,
},
}
/// 模型快照存储结果类型。
pub type ModelSnapshotStoreResult<T> = Result<T, ModelSnapshotStoreError>;
/// 模型快照存储端口。
#[async_trait]
pub trait ModelSnapshotStore: Send + Sync + 'static {
/// 用单个计划快照替换同名快照。
async fn replace_plan_snapshot(
&self,
snapshot: &ModelPlanSnapshot,
) -> ModelSnapshotStoreResult<()>;
/// 列出所有计划快照。
async fn list_plan_snapshots(&self) -> ModelSnapshotStoreResult<Vec<ModelPlanSnapshot>>;
}
crates/core/src/models/service.rs
+233
View File
@@ -0,0 +1,233 @@
//! 模型目录服务用例。
use std::{collections::BTreeMap, sync::Arc};
use thiserror::Error;
use crate::{
accounts::model::{Account, AccountStatus},
gateway::ports::{CodexModelCatalogClient, CodexModelCatalogRequest},
models::{
catalog::ModelCatalog,
model::{ModelConfig, ModelPlanSnapshot},
ports::{ModelSnapshotStore, ModelSnapshotStoreError},
},
};
/// 模型刷新摘要。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelRefreshResult {
/// 成功刷新并写入的计划数。
pub refreshed_plans: usize,
/// 本次成功写入的模型数。
pub model_count: usize,
/// 刷新失败的计划数。
pub failed_plans: usize,
}
/// 模型服务错误。
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ModelServiceError {
/// 没有注入快照存储。
#[error("model snapshot store is unavailable")]
SnapshotStoreUnavailable,
/// 没有注入上游模型客户端。
#[error("model catalog client is unavailable")]
UpstreamClientUnavailable,
/// 没有可用账号。
#[error("no active accounts available for model refresh")]
NoAccounts,
/// 快照写入失败。
#[error("failed to store model snapshot")]
StoreSnapshot,
/// 刷新后重新读取快照失败。
#[error("failed to load model snapshots")]
LoadSnapshots,
/// 所有计划都刷新失败。
#[error("all model refresh plans failed")]
AllPlansFailed(ModelRefreshResult),
}
/// 模型到可用计划的映射。
pub type ModelPlanAllowlist = BTreeMap<String, Vec<String>>;
/// 可共享更新的模型计划映射缓存。
pub type SharedModelPlanAllowlist = Arc<tokio::sync::Mutex<ModelPlanAllowlist>>;
/// 模型目录服务。
#[derive(Clone)]
pub struct ModelService {
config: ModelConfig,
snapshot_store: Option<Arc<dyn ModelSnapshotStore>>,
upstream_client: Option<Arc<dyn CodexModelCatalogClient>>,
model_plan_allowlist: Option<SharedModelPlanAllowlist>,
}
impl ModelService {
/// 构造模型服务。
pub fn new(
config: ModelConfig,
snapshot_store: Option<Arc<dyn ModelSnapshotStore>>,
upstream_client: Option<Arc<dyn CodexModelCatalogClient>>,
model_plan_allowlist: Option<SharedModelPlanAllowlist>,
) -> Self {
Self {
config,
snapshot_store,
upstream_client,
model_plan_allowlist,
}
}
/// 返回模型目录的静态配置。
pub fn config(&self) -> &ModelConfig {
&self.config
}
/// 构造当前对外暴露的模型目录。
pub async fn catalog(&self) -> ModelCatalog {
let Some(snapshot_store) = self.snapshot_store.as_ref() else {
return ModelCatalog::from_config(&self.config);
};
match snapshot_store.list_plan_snapshots().await {
Ok(snapshots) if !snapshots.is_empty() => {
ModelCatalog::from_config_and_snapshots(&self.config, &snapshots)
}
Ok(_) => ModelCatalog::from_config(&self.config),
Err(error) => {
tracing::warn!(error = %error, "加载模型快照失败，回退到静态目录");
ModelCatalog::from_config(&self.config)
}
}
}
/// 刷新活跃账号对应的后端模型目录。
pub async fn refresh_backend_models(
&self,
accounts: &[Account],
request_id: &str,
) -> Result<ModelRefreshResult, ModelServiceError> {
self.refresh_backend_models_with_installation_id(accounts, request_id, None)
.await
}
/// 使用运行时 installation id 刷新活跃账号对应的后端模型目录。
pub async fn refresh_backend_models_with_installation_id(
&self,
accounts: &[Account],
request_id: &str,
installation_id: Option<&str>,
) -> Result<ModelRefreshResult, ModelServiceError> {
let snapshot_store = self
.snapshot_store
.as_ref()
.ok_or(ModelServiceError::SnapshotStoreUnavailable)?;
let upstream_client = self
.upstream_client
.as_ref()
.ok_or(ModelServiceError::UpstreamClientUnavailable)?;
let plan_accounts = distinct_active_plan_accounts(accounts);
if plan_accounts.is_empty() {
return Err(ModelServiceError::NoAccounts);
}
let mut result = ModelRefreshResult {
refreshed_plans: 0,
model_count: 0,
failed_plans: 0,
};
for (plan_type, account) in plan_accounts {
let request = CodexModelCatalogRequest {
access_token: &account.access_token,
account_id: account.account_id.as_deref(),
request_id,
installation_id,
plan_type: &plan_type,
};
let entries = match upstream_client.fetch_models(&request).await {
Ok(entries) if !entries.is_empty() => entries,
Ok(_) => {
result.failed_plans += 1;
continue;
}
Err(error) => {
tracing::warn!(error = %error, plan_type, "刷新后端模型失败");
result.failed_plans += 1;
continue;
}
};
let snapshot = ModelPlanSnapshot::from_backend_entries(plan_type, entries);
result.model_count += snapshot.models.len();
snapshot_store
.replace_plan_snapshot(&snapshot)
.await
.map_err(map_store_snapshot_error)?;
result.refreshed_plans += 1;
}
if result.refreshed_plans == 0 {
return Err(ModelServiceError::AllPlansFailed(result));
}
let allowlist = self.model_plan_allowlist_from_store().await?;
if let Some(shared_allowlist) = &self.model_plan_allowlist {
*shared_allowlist.lock().await = allowlist;
}
Ok(result)
}
/// 读取当前缓存的 model -> plans allowlist。
pub async fn model_plan_allowlist(&self) -> Result<ModelPlanAllowlist, ModelServiceError> {
self.model_plan_allowlist_from_store().await
}
async fn model_plan_allowlist_from_store(
&self,
) -> Result<ModelPlanAllowlist, ModelServiceError> {
let snapshot_store = self
.snapshot_store
.as_ref()
.ok_or(ModelServiceError::SnapshotStoreUnavailable)?;
let snapshots = snapshot_store
.list_plan_snapshots()
.await
.map_err(map_load_snapshots_error)?;
Ok(
ModelCatalog::from_config_and_snapshots(&self.config, &snapshots)
.model_plan_allowlist(),
)
}
}
fn distinct_active_plan_accounts(accounts: &[Account]) -> Vec<(String, Account)> {
let mut by_plan = BTreeMap::new();
for account in accounts {
if account.status != AccountStatus::Active {
continue;
}
let plan_type = account
.plan_type
.as_deref()
.unwrap_or("unknown")
.to_string();
by_plan.entry(plan_type).or_insert_with(|| account.clone());
}
by_plan.into_iter().collect()
}
fn map_store_snapshot_error(_: ModelSnapshotStoreError) -> ModelServiceError {
ModelServiceError::StoreSnapshot
}
fn map_load_snapshots_error(_: ModelSnapshotStoreError) -> ModelServiceError {
ModelServiceError::LoadSnapshots
}
crates/core/src/protocol/codex/chat.rs
+11
View File
@@ -0,0 +1,11 @@
//! Codex Chat 扩展。
use serde::{Deserialize, Serialize};
/// Codex chat 侧的扩展选项。
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct CodexChatOptions {
/// 是否包含 reasoning 摘要。
pub include_reasoning: bool,
}
src/codex/gateway/transport/rate_limits.rs → crates/core/src/protocol/codex/events.rs
+267 -112
View File
@@ -1,60 +1,234 @@
//! Codex 事件语义、用量提取与限流解析。
use std::{
collections::BTreeMap,
time::{SystemTime, UNIX_EPOCH},
};
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use crate::protocol::codex::sse::{parse_sse_events, SseError};
/// 从 Codex/OpenAI usage 结构中提取出的标准化 token 用量。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TokenUsage {
/// 输入 token 数。
pub input_tokens: u64,
/// 输出 token 数。
pub output_tokens: u64,
/// 命中缓存的输入 token 数。
pub cached_tokens: u64,
/// 图片工具输入 token 数。
pub image_input_tokens: u64,
/// 图片工具输出 token 数。
pub image_output_tokens: u64,
/// 主模型总 token 数，不包含图片工具 token。
pub total_tokens: u64,
}
impl TokenUsage {
/// 构造仅包含主模型 token 的用量快照。
pub fn new(input_tokens: u64, output_tokens: u64, cached_tokens: u64) -> Self {
Self {
input_tokens,
output_tokens,
cached_tokens,
image_input_tokens: 0,
image_output_tokens: 0,
total_tokens: input_tokens + output_tokens,
}
}
/// 将两份用量按字段累加。
pub fn merged(self, other: Self) -> Self {
Self {
input_tokens: self.input_tokens + other.input_tokens,
output_tokens: self.output_tokens + other.output_tokens,
cached_tokens: self.cached_tokens + other.cached_tokens,
image_input_tokens: self.image_input_tokens + other.image_input_tokens,
image_output_tokens: self.image_output_tokens + other.image_output_tokens,
total_tokens: self.total_tokens + other.total_tokens,
}
}
}
/// 从单个 JSON 响应体中提取用量。
///
/// 该函数同时兼容 Codex usage 结构和 OpenAI usage 结构。
pub fn extract_usage(body: &Value) -> Option<TokenUsage> {
let usage = body.get("usage").unwrap_or(body);
if !usage.is_object() {
return None;
}
let input_tokens = number_field(usage, "input_tokens")
.or_else(|| number_field(usage, "prompt_tokens"))
.unwrap_or_default();
let output_tokens = number_field(usage, "output_tokens")
.or_else(|| number_field(usage, "completion_tokens"))
.unwrap_or_default();
let cached_tokens = nested_number_field(usage, &["input_tokens_details", "cached_tokens"])
.or_else(|| nested_number_field(usage, &["prompt_tokens_details", "cached_tokens"]))
.or_else(|| number_field(usage, "cached_tokens"))
.unwrap_or_default();
let image_input_tokens =
nested_number_field(body, &["tool_usage", "image_gen", "input_tokens"]).unwrap_or_default();
let image_output_tokens =
nested_number_field(body, &["tool_usage", "image_gen", "output_tokens"])
.unwrap_or_default();
let total_tokens = number_field(usage, "total_tokens").unwrap_or(input_tokens + output_tokens);
let has_usage = [
"input_tokens",
"output_tokens",
"prompt_tokens",
"completion_tokens",
"cached_tokens",
"total_tokens",
]
.iter()
.any(|field| usage.get(*field).is_some())
|| usage.get("input_tokens_details").is_some()
|| usage.get("prompt_tokens_details").is_some()
|| image_input_tokens > 0
|| image_output_tokens > 0;
has_usage.then_some(TokenUsage {
input_tokens,
output_tokens,
cached_tokens,
image_input_tokens,
image_output_tokens,
total_tokens,
})
}
/// 从完整 SSE 文本中提取最终可见用量。
///
/// 如果存在 `response.completed` 的 usage，则优先返回它；否则回退到最后一条
/// 可见 usage 事件。
///
/// # Errors
///
/// 当输入不是合法 SSE 流时，返回 [`SseError`]。
pub fn extract_sse_usage(body: &str) -> Result<Option<TokenUsage>, SseError> {
let events = parse_sse_events(body)?;
let mut fallback_usage: Option<TokenUsage> = None;
for event in events {
if event.data == "[DONE]" {
continue;
}
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
continue;
};
let event_usage = value
.get("response")
.and_then(extract_usage)
.or_else(|| extract_usage(&value));
if let Some(event_usage) = event_usage {
if event.event.as_deref() == Some("response.completed") {
return Ok(Some(event_usage));
}
fallback_usage = Some(event_usage);
}
}
Ok(fallback_usage)
}
/// 从上游错误响应体中提取 retry-after 秒数。
///
/// 支持结构化的 `resets_in_seconds` / `resets_at`，也支持官方错误消息里的
/// `try again in 11.054s` / `try again in 28ms` 文本。
pub fn retry_after_seconds_from_body(body: &str) -> Option<u64> {
let value = serde_json::from_str::<Value>(body).ok()?;
let error = value
.pointer("/response/error")
.or_else(|| value.get("error"))
.unwrap_or(&value);
if let Some(seconds) = error
.get("resets_in_seconds")
.and_then(Value::as_u64)
.filter(|seconds| *seconds > 0)
{
return Some(seconds);
}
retry_after_seconds_from_resets_at(error)
.or_else(|| retry_after_seconds_from_rate_limit_message(error))
}
/// 标准化的单个限流窗口。
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RateLimitWindow {
/// 已使用百分比。
pub used_percent: f64,
/// 窗口分钟数。
pub window_minutes: Option<u64>,
/// 重置时间戳。
pub reset_at: Option<i64>,
}
impl RateLimitWindow {
/// 以秒返回限流窗口大小。
pub fn limit_window_seconds(self) -> Option<u64> {
self.window_minutes
.and_then(|minutes| minutes.checked_mul(60))
}
/// 将重置时间戳转换为 UTC 时间。
pub fn reset_at_datetime(self) -> Option<DateTime<Utc>> {
let reset_at = self.reset_at?;
DateTime::<Utc>::from_timestamp(reset_at, 0)
}
}
/// 复合限流信息，通常用于 code review 配额。
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RateLimitDetails {
/// 当前请求是否被允许。
pub allowed: Option<bool>,
/// 当前窗口是否已经触顶。
pub limit_reached: Option<bool>,
/// 主限流窗口。
pub primary: Option<RateLimitWindow>,
/// 次级限流窗口。
pub secondary: Option<RateLimitWindow>,
}
/// 从 header 或内部事件中解析出的完整限流状态。
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ParsedRateLimits {
/// 主限流窗口。
pub primary: Option<RateLimitWindow>,
/// 次级限流窗口。
pub secondary: Option<RateLimitWindow>,
/// code review 限流窗口。
pub code_review: Option<RateLimitDetails>,
}
impl ParsedRateLimits {
/// 返回主窗口的 UTC 重置时间。
pub fn primary_reset_at(self) -> Option<DateTime<Utc>> {
self.primary.and_then(RateLimitWindow::reset_at_datetime)
}
/// 返回主窗口的秒级窗口长度。
pub fn primary_limit_window_seconds(self) -> Option<u64> {
self.primary.and_then(RateLimitWindow::limit_window_seconds)
}
/// 返回主窗口是否已经触顶。
pub fn primary_limit_reached(self) -> bool {
self.primary
.is_some_and(|window| window.used_percent >= 100.0)
}
}
/// 从响应头对中解析限流信息。
pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option<ParsedRateLimits> {
let mut normalized = BTreeMap::new();
for (name, value) in headers {
@@ -78,6 +252,7 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option<ParsedRa
})
}
/// 从内部 `codex.rate_limits` 事件中解析限流信息。
pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
if value
.get("type")
@@ -117,12 +292,14 @@ pub fn parse_rate_limits_event(value: &Value) -> Option<ParsedRateLimits> {
})
}
/// 从原始 JSON 文本中解析内部 `codex.rate_limits` 事件。
pub fn parse_rate_limits_event_raw(raw: &str) -> Option<ParsedRateLimits> {
serde_json::from_str::<Value>(raw)
.ok()
.and_then(|value| parse_rate_limits_event(&value))
}
/// 将解析后的限流状态转换为配额响应体。
pub fn rate_limit_quota(
rate_limits: &ParsedRateLimits,
plan_type: Option<&str>,
@@ -142,6 +319,7 @@ pub fn rate_limit_quota(
quota
}
/// 将限流状态转换回 HTTP 头键值对。
pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(String, String)> {
let mut headers = Vec::new();
push_window_headers(&mut headers, "x-codex-primary", rate_limits.primary);
@@ -161,6 +339,7 @@ pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(Strin
headers
}
/// 为退避秒数增加抖动，避免同一时间窗口内集中重试。
pub fn cooldown_with_jitter(seconds: u64, factor_bps: u16) -> Duration {
let factor_bps = u64::from(factor_bps);
let variance = seconds.saturating_mul(factor_bps) / 10_000;
@@ -176,6 +355,94 @@ pub fn cooldown_with_jitter(seconds: u64, factor_bps: u16) -> Duration {
Duration::seconds(lower.saturating_add(offset).min(i64::MAX as u64) as i64)
}
fn number_field(value: &Value, field: &str) -> Option<u64> {
value.get(field)?.as_u64()
}
fn nested_number_field(value: &Value, path: &[&str]) -> Option<u64> {
let mut current = value;
for segment in path {
current = current.get(*segment)?;
}
current.as_u64()
}
fn retry_after_seconds_from_resets_at(error: &Value) -> Option<u64> {
let resets_at = error.get("resets_at").and_then(Value::as_u64)?;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.ok()?
.as_secs();
(resets_at > now).then_some(resets_at - now)
}
fn retry_after_seconds_from_rate_limit_message(error: &Value) -> Option<u64> {
let code = error
.get("code")
.or_else(|| error.get("type"))
.and_then(Value::as_str)?;
if code != "rate_limit_exceeded" {
return None;
}
let message = error.get("message").and_then(Value::as_str)?;
parse_try_again_delay_seconds(message)
}
fn parse_try_again_delay_seconds(message: &str) -> Option<u64> {
let lower = message.to_ascii_lowercase();
let marker = "try again in";
let offset = lower.find(marker)? + marker.len();
let remainder = message.get(offset..)?.trim_start();
let number_end = number_prefix_len(remainder)?;
let value = remainder.get(..number_end)?.parse::<f64>().ok()?;
if !value.is_finite() || value <= 0.0 {
return None;
}
let unit_text = remainder
.get(number_end..)?
.trim_start()
.to_ascii_lowercase();
let unit = unit_token(&unit_text)?;
match unit {
"ms" => positive_seconds_ceil(value / 1000.0),
"s" | "second" | "seconds" => positive_seconds_ceil(value),
_ => None,
}
}
fn number_prefix_len(input: &str) -> Option<usize> {
let mut seen_digit = false;
let mut seen_dot = false;
let mut end = 0;
for (index, ch) in input.char_indices() {
if ch.is_ascii_digit() {
seen_digit = true;
end = index + ch.len_utf8();
} else if ch == '.' && !seen_dot {
seen_dot = true;
end = index + ch.len_utf8();
} else {
break;
}
}
seen_digit.then_some(end)
}
fn unit_token(input: &str) -> Option<&str> {
let end = input
.char_indices()
.find_map(|(index, ch)| (!ch.is_ascii_alphabetic()).then_some(index))
.unwrap_or(input.len());
(end > 0).then_some(&input[..end])
}
fn positive_seconds_ceil(seconds: f64) -> Option<u64> {
if !seconds.is_finite() || seconds <= 0.0 || seconds > u64::MAX as f64 {
return None;
}
Some(seconds.ceil() as u64)
}
fn push_window_headers(
headers: &mut Vec<(String, String)>,
prefix: &str,
@@ -353,115 +620,3 @@ fn is_review_limit_name(value: Option<&str>) -> bool {
"review" | "code_review" | "codex_review" | "codex_code_review"
)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
parse_rate_limit_headers, parse_rate_limits_event, rate_limit_quota, RateLimitWindow,
};
#[test]
fn parse_rate_limit_headers_should_extract_primary_secondary_and_review_windows() {
let headers = vec![
(
"x-codex-primary-used-percent".to_string(),
"100".to_string(),
),
(
"x-codex-primary-window-minutes".to_string(),
"5".to_string(),
),
(
"x-codex-primary-reset-at".to_string(),
"1893456300".to_string(),
),
(
"x-codex-secondary-used-percent".to_string(),
"42.5".to_string(),
),
(
"x-codex-secondary-window-minutes".to_string(),
"10080".to_string(),
),
(
"x-codex-code-review-primary-used-percent".to_string(),
"80".to_string(),
),
(
"x-codex-code-review-primary-reset-at".to_string(),
"1893456600".to_string(),
),
];
let parsed = parse_rate_limit_headers(&headers).unwrap();
assert_eq!(
parsed.primary,
Some(RateLimitWindow {
used_percent: 100.0,
window_minutes: Some(5),
reset_at: Some(1_893_456_300),
})
);
assert_eq!(parsed.secondary.unwrap().window_minutes, Some(10080));
assert_eq!(
parsed.code_review.unwrap().primary.unwrap().reset_at,
Some(1_893_456_600)
);
}
#[test]
fn parse_rate_limits_event_should_extract_internal_websocket_rate_limits() {
let event = json!({
"type": "codex.rate_limits",
"rate_limits": {
"primary": {
"used_percent": 99.5,
"window_minutes": 300,
"reset_at": 1893456300
},
"secondary": {
"used_percent": 10,
"window_minutes": 10080,
"reset_at": 1894056000
}
}
});
let parsed = parse_rate_limits_event(&event).unwrap();
assert_eq!(parsed.primary.unwrap().used_percent, 99.5);
assert_eq!(parsed.secondary.unwrap().reset_at, Some(1_894_056_000));
}
#[test]
fn rate_limit_quota_should_preserve_existing_credits_when_passive_data_lacks_credits() {
let headers = vec![
("x-codex-primary-used-percent".to_string(), "25".to_string()),
(
"x-codex-primary-window-minutes".to_string(),
"5".to_string(),
),
(
"x-codex-primary-reset-at".to_string(),
"1893456300".to_string(),
),
];
let existing = json!({
"credits": {
"has_credits": true,
"unlimited": false,
"balance": 12
}
});
let parsed = parse_rate_limit_headers(&headers).unwrap();
let quota = rate_limit_quota(&parsed, Some("plus"), Some(&existing));
assert_eq!(quota["plan_type"], "plus");
assert_eq!(quota["rate_limit"]["remaining_percent"], 75);
assert_eq!(quota["credits"]["balance"], 12);
}
}
crates/core/src/protocol/codex/mod.rs
+14
View File
@@ -0,0 +1,14 @@
//! Codex 协议类型。
/// Codex Chat 扩展字段。
pub mod chat;
/// Codex 事件语义与用量提取。
pub mod events;
/// Codex Responses / Compact 请求体。
pub mod responses;
/// Codex JSON schema 处理。
pub(crate) mod schema;
/// SSE 事件解析。
pub mod sse;
/// WebSocket 帧类型与纯编解码。
pub mod websocket;
src/codex/gateway/transport/types.rs → crates/core/src/protocol/codex/responses.rs
+40
View File
@@ -1,64 +1,95 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
/// Codex Responses 上游请求体。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodexResponsesRequest {
/// 模型名。
pub model: String,
/// 指令文本。
pub instructions: String,
/// 输入消息与结构化条目。
pub input: Vec<Value>,
/// 是否流式返回。
pub stream: bool,
/// 是否要求上游存储响应。
pub store: bool,
#[serde(skip_serializing_if = "Option::is_none")]
/// reasoning 配置。
pub reasoning: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 工具定义。
pub tools: Option<Vec<Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 工具选择策略。
pub tool_choice: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 是否允许并行工具调用。
pub parallel_tool_calls: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 输出文本格式配置。
pub text: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 图片生成开关。
pub generate: Option<bool>,
#[serde(skip)]
/// tuple schema 原始定义，仅供响应重构时使用。
pub tuple_schema: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// service tier。
pub service_tier: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 用于显式续链的前一个 response ID。
pub previous_response_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 提示缓存键。
pub prompt_cache_key: Option<String>,
#[serde(skip)]
/// 是否由客户端显式提供了 prompt cache key。
pub explicit_prompt_cache_key: bool,
#[serde(skip)]
/// 客户端会话 ID。
pub client_conversation_id: Option<String>,
#[serde(skip)]
/// 变体身份键。
pub variant_identity: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
/// include 列表。
pub include: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 传给上游的 client metadata。
pub client_metadata: Option<Value>,
#[serde(skip)]
/// 是否偏好 WebSocket 传输。
pub use_websocket: bool,
#[serde(skip)]
/// 是否强制 HTTP SSE。
pub force_http_sse: bool,
#[serde(skip)]
/// turn state 透传头。
pub turn_state: Option<String>,
#[serde(skip)]
/// turn metadata 透传头。
pub turn_metadata: Option<String>,
#[serde(skip)]
/// beta features 透传头。
pub beta_features: Option<String>,
#[serde(skip)]
/// 客户端版本头。
pub version: Option<String>,
#[serde(skip)]
/// timing metrics 透传头。
pub include_timing_metrics: Option<String>,
#[serde(skip)]
/// codex window id。
pub codex_window_id: Option<String>,
#[serde(skip)]
/// 父线程 ID。
pub parent_thread_id: Option<String>,
}
impl CodexResponsesRequest {
/// 构造默认的 HTTP SSE 请求。
pub fn new_http_sse(
model: impl Into<String>,
instructions: impl Into<String>,
@@ -97,6 +128,7 @@ impl CodexResponsesRequest {
}
}
/// 判断请求是否声明了图片生成工具。
pub fn expects_image_generation(&self) -> bool {
self.tools.as_deref().is_some_and(|tools| {
tools
@@ -106,17 +138,25 @@ impl CodexResponsesRequest {
}
}
/// Codex compact 端点请求体。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodexCompactRequest {
/// 模型名。
pub model: String,
/// 输入消息与结构化条目。
pub input: Vec<Value>,
/// 指令文本。
pub instructions: String,
#[serde(skip_serializing_if = "Option::is_none")]
/// 工具定义。
pub tools: Option<Vec<Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 是否允许并行工具调用。
pub parallel_tool_calls: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
/// reasoning 配置。
pub reasoning: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
/// 输出文本格式配置。
pub text: Option<Value>,
}
src/codex/gateway/protocol/tuple_schema.rs → crates/core/src/protocol/codex/schema.rs
+4
View File
@@ -1,3 +1,5 @@
//! Codex 协议所需的 JSON schema 纯处理逻辑。
use serde_json::{Map, Value};
pub(crate) struct PreparedSchema {
@@ -18,6 +20,7 @@ pub(crate) fn prepare_schema(schema: Value) -> PreparedSchema {
}
}
#[allow(dead_code)]
pub(crate) fn reconvert_tuple_values(data: Value, schema: &Value) -> Value {
reconvert_tuple_values_with_root(data, schema, schema)
}
@@ -116,6 +119,7 @@ fn has_tuple_schemas(value: &Value) -> bool {
{
return true;
}
object_values_have_tuple_schema(object, "properties")
|| object.get("items").is_some_and(has_tuple_schemas)
|| schema_array_has_tuple_schema(object, "oneOf")
src/codex/gateway/transport/sse.rs → crates/core/src/protocol/codex/sse.rs
+35 -19
View File
@@ -1,21 +1,35 @@
//! SSE 事件解析与编码。
use thiserror::Error;
pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024;
/// 单条 SSE 事件。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SseEvent {
/// 事件名。
pub event: Option<String>,
/// 数据体。
pub data: String,
/// 可选 ID。
pub id: Option<String>,
/// 可选 retry。
pub retry: Option<u64>,
}
/// 单事件缓冲上限。
pub const MAX_SSE_EVENT_BUFFER_BYTES: usize = 64 * 1024 * 1024;
/// SSE 解析错误。
#[derive(Debug, Error, PartialEq, Eq)]
pub enum SseError {
/// retry 字段不是合法整数。
#[error("invalid SSE retry value: {0}")]
InvalidRetry(String),
/// 单个事件缓冲超过上限。
#[error("SSE buffer exceeded {max_bytes} bytes — aborting stream")]
BufferExceeded { max_bytes: usize },
BufferExceeded {
/// 上限字节数。
max_bytes: usize,
},
}
#[derive(Debug, Default)]
@@ -60,6 +74,7 @@ impl EventBuilder {
}
}
/// 解析 SSE 事件流。
pub fn parse_sse_events(input: &str) -> Result<Vec<SseEvent>, SseError> {
let mut events = Vec::new();
let mut builder = EventBuilder::default();
@@ -110,6 +125,23 @@ pub fn parse_sse_events(input: &str) -> Result<Vec<SseEvent>, SseError> {
Ok(events)
}
/// 编码单条 SSE 事件。
pub fn encode_sse_event(event: &str, data: &str) -> String {
let mut frame = String::new();
if !event.is_empty() {
frame.push_str("event: ");
frame.push_str(event);
frame.push('\n');
}
for line in data.split('\n') {
frame.push_str("data: ");
frame.push_str(line);
frame.push('\n');
}
frame.push('\n');
frame
}
fn track_event_buffer_bytes(current_bytes: &mut usize, line: &str) -> Result<(), SseError> {
let line_separator_bytes = usize::from(*current_bytes != 0);
*current_bytes = current_bytes
@@ -166,22 +198,6 @@ fn non_sse_error_message(raw: &str) -> String {
.to_string()
}
pub fn encode_sse_event(event: &str, data: &str) -> String {
let mut frame = String::new();
if !event.is_empty() {
frame.push_str("event: ");
frame.push_str(event);
frame.push('\n');
}
for line in data.split('\n') {
frame.push_str("data: ");
frame.push_str(line);
frame.push('\n');
}
frame.push('\n');
frame
}
fn split_sse_field(line: &str) -> (&str, &str) {
let Some((field, value)) = line.split_once(':') else {
return (line, "");
src/codex/gateway/transport/websocket/codec.rs → crates/core/src/protocol/codex/websocket.rs
+815 -554
View File
File diff suppressed because it is too large Load Diff
crates/core/src/protocol/mod.rs
+4
View File
@@ -0,0 +1,4 @@
//! 协议语义与纯数据类型。
pub mod codex;
pub mod openai;
crates/core/src/protocol/openai/chat.rs
+620
View File
@@ -0,0 +1,620 @@
//! OpenAI Chat Completions 请求到 Codex Responses 请求的纯转换。
use std::collections::{BTreeMap, BTreeSet};
use chrono::Utc;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use thiserror::Error;
use uuid::Uuid;
use crate::protocol::codex::{
responses::CodexResponsesRequest,
schema::{prepare_schema, reconvert_tuple_values},
sse::{parse_sse_events, SseError},
};
/// OpenAI Chat Completions 请求体。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatCompletionRequest {
/// 客户端请求的模型名。
pub model: String,
/// 是否流式返回。
#[serde(default)]
pub stream: bool,
/// 按 OpenAI 语义组织的消息列表。
pub messages: Vec<ChatMessage>,
/// 可选的 reasoning effort。
#[serde(default)]
pub reasoning_effort: Option<String>,
/// 可选的 service tier。
#[serde(default)]
pub service_tier: Option<String>,
/// 原生 tools 字段。
#[serde(default)]
pub tools: Option<Vec<Value>>,
/// tool choice 配置。
#[serde(default)]
pub tool_choice: Option<Value>,
/// parallel tool calls 配置。
#[serde(default)]
pub parallel_tool_calls: Option<bool>,
/// 旧式 functions 字段，会转换为 Codex tools。
#[serde(default)]
pub functions: Option<Vec<Value>>,
/// 响应格式配置。
#[serde(default)]
pub response_format: Option<Value>,
/// 可选的 user 字段，用作稳定会话锚点。
#[serde(default)]
pub user: Option<String>,
}
/// Chat 请求中的单条消息。
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
/// 角色名。
pub role: String,
/// 消息内容，既可能是字符串也可能是结构化数组。
#[serde(default)]
pub content: Option<Value>,
/// function/tool 相关的名称。
#[serde(default)]
pub name: Option<String>,
/// assistant 消息内的 tool_calls。
#[serde(default)]
pub tool_calls: Option<Vec<Value>>,
/// tool 响应对应的调用 ID。
#[serde(default)]
pub tool_call_id: Option<String>,
/// 旧式 function_call 字段。
#[serde(default)]
pub function_call: Option<Value>,
}
/// Chat 转译失败的错误。
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ChatTranslationError {
/// 请求没有任何消息，无法生成有效上游输入。
#[error("messages must not be empty")]
EmptyMessages,
}
/// 将 OpenAI Chat Completions 请求纯转换为 Codex Responses 请求。
pub fn translate_chat_to_codex(
request: ChatCompletionRequest,
) -> Result<CodexResponsesRequest, ChatTranslationError> {
if request.messages.is_empty() {
return Err(ChatTranslationError::EmptyMessages);
}
let client_conversation_id = request
.user
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string);
let instructions = chat_instructions(&request.messages);
let mut input = chat_input(request.messages);
if input.is_empty() {
input.push(json!({"role": "user", "content": ""}));
}
let mut codex_request = CodexResponsesRequest::new_http_sse(request.model, instructions, input);
codex_request.force_http_sse = true;
codex_request.tools = codex_tools(request.tools, request.functions);
codex_request.tool_choice = request.tool_choice;
codex_request.parallel_tool_calls = request.parallel_tool_calls;
let response_format = response_format_text(request.response_format);
codex_request.text = response_format.text;
codex_request.tuple_schema = response_format.tuple_schema;
codex_request.service_tier = request.service_tier;
codex_request.prompt_cache_key = client_conversation_id.clone();
codex_request.client_conversation_id = client_conversation_id;
if let Some(effort) = request.reasoning_effort {
codex_request.reasoning = Some(json!({"effort": effort, "summary": "auto"}));
}
Ok(codex_request)
}
/// 将 Codex SSE 完成响应转换为 OpenAI Chat Completions 响应体。
pub fn chat_completion_from_codex_sse(
body: &str,
model: &str,
include_reasoning: bool,
tuple_schema: Option<&Value>,
) -> Result<Option<Value>, SseError> {
let events = parse_sse_events(body)?;
let mut content = String::new();
let mut reasoning = String::new();
let mut function_call_items = BTreeMap::new();
let mut finished_call_ids = BTreeSet::new();
let mut tool_calls = Vec::new();
let mut completed_response = None;
for event in events {
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
continue;
};
match event.event.as_deref() {
Some("response.output_text.delta") => {
if let Some(delta) = value.get("delta").and_then(Value::as_str) {
content.push_str(delta);
}
}
Some("response.reasoning_summary_text.delta") if include_reasoning => {
if let Some(delta) = value.get("delta").and_then(Value::as_str) {
reasoning.push_str(delta);
}
}
Some("response.output_item.added") => {
if let Some((item_id, call_id, name)) = function_call_item(value.get("item")) {
function_call_items.insert(item_id, FunctionCallInfo { call_id, name });
}
}
Some("response.function_call_arguments.done") => {
if let Some(tool_call) =
tool_call_from_done_event(&value, &function_call_items, &mut finished_call_ids)
{
tool_calls.push(tool_call);
}
}
Some("response.output_item.done") => {
if let Some(tool_call) = tool_call_from_output_item(
value.get("item"),
&function_call_items,
&mut finished_call_ids,
) {
tool_calls.push(tool_call);
}
}
Some("response.completed") => {
completed_response = value.get("response").cloned();
}
_ => {}
}
}
let Some(response) = completed_response else {
return Ok(None);
};
if content.is_empty() {
content = output_text_from_response(&response);
}
if let Some(tuple_schema) = tuple_schema {
content = reconvert_tuple_text(&content, tuple_schema).unwrap_or(content);
}
let mut message = json!({
"role": "assistant",
"content": if content.is_empty() && !tool_calls.is_empty() {
Value::Null
} else {
Value::String(content)
},
});
if include_reasoning && !reasoning.is_empty() {
message["reasoning_content"] = Value::String(reasoning);
}
let has_tool_calls = !tool_calls.is_empty();
if has_tool_calls {
message["tool_calls"] = Value::Array(tool_calls);
}
Ok(Some(json!({
"id": format!("chatcmpl-{}", Uuid::new_v4().simple()),
"object": "chat.completion",
"created": Utc::now().timestamp(),
"model": model,
"choices": [{
"index": 0,
"message": message,
"finish_reason": if has_tool_calls { "tool_calls" } else { "stop" },
}],
"usage": openai_usage(response.get("usage")),
})))
}
fn chat_instructions(messages: &[ChatMessage]) -> String {
let instructions = messages
.iter()
.filter(|message| message.role == "system" || message.role == "developer")
.map(|message| extract_text(&message.content))
.filter(|text| !text.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
if instructions.is_empty() {
"You are a helpful assistant.".to_string()
} else {
instructions
}
}
fn chat_input(messages: Vec<ChatMessage>) -> Vec<Value> {
let mut input = Vec::new();
for message in messages {
match message.role.as_str() {
"system" | "developer" => {}
"assistant" => push_assistant_message(&mut input, message),
"tool" => input.push(json!({
"type": "function_call_output",
"call_id": message.tool_call_id.unwrap_or_else(|| "unknown".to_string()),
"output": extract_text(&message.content),
})),
"function" => input.push(json!({
"type": "function_call_output",
"call_id": format!("fc_{}", message.name.unwrap_or_else(|| "unknown".to_string())),
"output": extract_text(&message.content),
})),
_ => input.push(json!({
"role": "user",
"content": extract_content(&message.content),
})),
}
}
input
}
fn push_assistant_message(input: &mut Vec<Value>, message: ChatMessage) {
let text = extract_text(&message.content);
let has_tool_calls = message
.tool_calls
.as_ref()
.is_some_and(|calls| !calls.is_empty());
if !text.is_empty() || (!has_tool_calls && message.function_call.is_none()) {
input.push(json!({"role": "assistant", "content": text}));
}
if let Some(tool_calls) = message.tool_calls {
for tool_call in tool_calls {
let call_id = tool_call
.get("id")
.and_then(Value::as_str)
.unwrap_or("unknown");
let function = tool_call.get("function").unwrap_or(&Value::Null);
let name = function
.get("name")
.and_then(Value::as_str)
.unwrap_or("unknown");
let arguments = function
.get("arguments")
.and_then(Value::as_str)
.unwrap_or("");
input.push(json!({
"type": "function_call",
"call_id": call_id,
"name": name,
"arguments": arguments,
}));
}
}
if let Some(function_call) = message.function_call {
let name = function_call
.get("name")
.and_then(Value::as_str)
.unwrap_or("unknown");
let arguments = function_call
.get("arguments")
.and_then(Value::as_str)
.unwrap_or("");
input.push(json!({
"type": "function_call",
"call_id": format!("fc_{name}"),
"name": name,
"arguments": arguments,
}));
}
}
fn extract_text(content: &Option<Value>) -> String {
let Some(content) = content else {
return String::new();
};
if let Some(text) = content.as_str() {
return text.to_string();
}
let Some(parts) = content.as_array() else {
return String::new();
};
parts
.iter()
.filter_map(|part| {
(part.get("type").and_then(Value::as_str) == Some("text"))
.then(|| part.get("text").and_then(Value::as_str))
.flatten()
})
.collect::<Vec<_>>()
.join("\n")
}
fn extract_content(content: &Option<Value>) -> Value {
let Some(content) = content else {
return Value::String(String::new());
};
if content.as_str().is_some() {
return content.clone();
}
let Some(parts) = content.as_array() else {
return Value::String(String::new());
};
let has_image = parts
.iter()
.any(|part| part.get("type").and_then(Value::as_str) == Some("image_url"));
if !has_image {
return Value::String(extract_text(&Some(content.clone())));
}
let codex_parts = parts
.iter()
.filter_map(codex_content_part)
.collect::<Vec<_>>();
if codex_parts.is_empty() {
Value::String(String::new())
} else {
Value::Array(codex_parts)
}
}
fn codex_content_part(part: &Value) -> Option<Value> {
match part.get("type").and_then(Value::as_str)? {
"text" => part
.get("text")
.and_then(Value::as_str)
.map(|text| json!({"type": "input_text", "text": text})),
"image_url" => image_url(part).map(|url| json!({"type": "input_image", "image_url": url})),
_ => None,
}
}
fn image_url(part: &Value) -> Option<&str> {
let image = part.get("image_url")?;
image
.as_str()
.or_else(|| image.get("url").and_then(Value::as_str))
}
fn reconvert_tuple_text(text: &str, tuple_schema: &Value) -> Option<String> {
let parsed = serde_json::from_str::<Value>(text).ok()?;
Some(reconvert_tuple_values(parsed, tuple_schema).to_string())
}
fn codex_tools(tools: Option<Vec<Value>>, functions: Option<Vec<Value>>) -> Option<Vec<Value>> {
if let Some(tools) = tools.filter(|tools| !tools.is_empty()) {
return Some(tools);
}
functions
.filter(|functions| !functions.is_empty())
.map(|functions| {
functions
.into_iter()
.map(|function| json!({"type": "function", "function": function}))
.collect()
})
}
struct PreparedResponseFormat {
text: Option<Value>,
tuple_schema: Option<Value>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct FunctionCallInfo {
call_id: String,
name: String,
}
fn response_format_text(response_format: Option<Value>) -> PreparedResponseFormat {
let Some(format) = response_format else {
return PreparedResponseFormat {
text: None,
tuple_schema: None,
};
};
let Some(kind) = format.get("type").and_then(Value::as_str) else {
return PreparedResponseFormat {
text: None,
tuple_schema: None,
};
};
match kind {
"json_object" => PreparedResponseFormat {
text: Some(json!({"format": {"type": "json_object"}})),
tuple_schema: None,
},
"json_schema" => {
let Some(schema) = format.get("json_schema") else {
return PreparedResponseFormat {
text: None,
tuple_schema: None,
};
};
let name = schema
.get("name")
.and_then(Value::as_str)
.unwrap_or("response");
let prepared_schema =
prepare_schema(schema.get("schema").cloned().unwrap_or_else(|| json!({})));
let mut codex_format = json!({
"type": "json_schema",
"name": name,
"schema": prepared_schema.schema,
});
if let Some(strict) = schema.get("strict").and_then(Value::as_bool) {
codex_format["strict"] = Value::Bool(strict);
}
PreparedResponseFormat {
text: Some(json!({"format": codex_format})),
tuple_schema: prepared_schema.original_schema,
}
}
_ => PreparedResponseFormat {
text: None,
tuple_schema: None,
},
}
}
fn output_text_from_response(response: &Value) -> String {
response
.get("output_text")
.and_then(Value::as_str)
.map(ToString::to_string)
.unwrap_or_else(|| {
response
.get("output")
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(output_text_from_item)
.collect::<Vec<_>>()
.join("\n\n")
})
.unwrap_or_default()
})
}
fn output_text_from_item(item: &Value) -> Option<String> {
let content = item.get("content")?.as_array()?;
let text = content
.iter()
.filter_map(|part| {
let part_type = part.get("type")?.as_str()?;
if part_type != "output_text" && part_type != "text" {
return None;
}
part.get("text")?.as_str()
})
.collect::<String>();
(!text.is_empty()).then_some(text)
}
fn function_call_item(item: Option<&Value>) -> Option<(String, String, String)> {
let item = item?;
if item.get("type").and_then(Value::as_str) != Some("function_call") {
return None;
}
let item_id = item.get("id").and_then(Value::as_str)?;
let call_id = item.get("call_id").and_then(Value::as_str)?;
let name = item.get("name").and_then(Value::as_str)?;
Some((item_id.to_string(), call_id.to_string(), name.to_string()))
}
fn tool_call_from_done_event(
event: &Value,
function_call_items: &BTreeMap<String, FunctionCallInfo>,
finished_call_ids: &mut BTreeSet<String>,
) -> Option<Value> {
let event_id = event
.get("call_id")
.or_else(|| event.get("item_id"))
.and_then(Value::as_str)?;
let info = function_call_items.get(event_id);
let call_id = info
.map(|info| info.call_id.as_str())
.unwrap_or(event_id)
.to_string();
if !finished_call_ids.insert(call_id.clone()) {
return None;
}
let name = event
.get("name")
.and_then(Value::as_str)
.filter(|name| !name.is_empty())
.or_else(|| info.map(|info| info.name.as_str()))
.unwrap_or("unknown");
let arguments = event
.get("arguments")
.and_then(Value::as_str)
.unwrap_or_default();
Some(openai_tool_call(call_id, name, arguments))
}
fn tool_call_from_output_item(
item: Option<&Value>,
function_call_items: &BTreeMap<String, FunctionCallInfo>,
finished_call_ids: &mut BTreeSet<String>,
) -> Option<Value> {
let item = item?;
if item.get("type").and_then(Value::as_str) != Some("function_call") {
return None;
}
let item_id = item.get("id").and_then(Value::as_str);
let info = item_id.and_then(|item_id| function_call_items.get(item_id));
let call_id = item
.get("call_id")
.and_then(Value::as_str)
.or_else(|| info.map(|info| info.call_id.as_str()))
.or(item_id)?;
if !finished_call_ids.insert(call_id.to_string()) {
return None;
}
let name = item
.get("name")
.and_then(Value::as_str)
.filter(|name| !name.is_empty())
.or_else(|| info.map(|info| info.name.as_str()))
.unwrap_or("unknown");
let arguments = item
.get("arguments")
.and_then(Value::as_str)
.unwrap_or_default();
Some(openai_tool_call(call_id.to_string(), name, arguments))
}
fn openai_tool_call(call_id: String, name: &str, arguments: &str) -> Value {
json!({
"id": call_id,
"type": "function",
"function": {
"name": name,
"arguments": arguments,
},
})
}
fn openai_usage(usage: Option<&Value>) -> Value {
let prompt_tokens = usage
.and_then(|usage| usage.get("input_tokens"))
.and_then(Value::as_u64)
.unwrap_or_default();
let completion_tokens = usage
.and_then(|usage| usage.get("output_tokens"))
.and_then(Value::as_u64)
.unwrap_or_default();
let cached_tokens = usage
.and_then(|usage| usage.get("input_tokens_details"))
.and_then(|details| details.get("cached_tokens"))
.and_then(Value::as_u64)
.unwrap_or_default();
let reasoning_tokens = usage
.and_then(|usage| usage.get("output_tokens_details"))
.and_then(|details| details.get("reasoning_tokens"))
.and_then(Value::as_u64);
let mut openai_usage = json!({
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens.saturating_add(completion_tokens),
"prompt_tokens_details": {
"cached_tokens": cached_tokens,
},
});
if let Some(reasoning_tokens) = reasoning_tokens {
openai_usage["completion_tokens_details"] = json!({
"reasoning_tokens": reasoning_tokens,
});
}
openai_usage
}
crates/core/src/protocol/openai/errors.rs
+15
View File
@@ -0,0 +1,15 @@
//! OpenAI 错误响应类型。
use serde::{Deserialize, Serialize};
/// OpenAI 错误对象。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct OpenAiError {
/// 错误消息。
pub message: String,
/// 错误类型。
#[serde(rename = "type")]
pub error_type: String,
/// 错误代码。
pub code: Option<String>,
}
crates/core/src/protocol/openai/mod.rs
+6
View File
@@ -0,0 +1,6 @@
//! OpenAI 协议输入类型与纯转换逻辑。
pub mod chat;
pub mod errors;
pub mod models;
pub mod responses;
crates/core/src/protocol/openai/models.rs
+25
View File
@@ -0,0 +1,25 @@
//! OpenAI 模型列表响应类型。
use serde::{Deserialize, Serialize};
/// OpenAI 模型对象。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct OpenAiModel {
/// 模型 ID。
pub id: String,
/// 对象类型。
pub object: String,
/// 创建时间戳。
pub created: i64,
/// 所有者。
pub owned_by: String,
}
/// OpenAI 模型列表。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct OpenAiModelList {
/// 对象类型。
pub object: String,
/// 模型数据。
pub data: Vec<OpenAiModel>,
}
crates/core/src/protocol/openai/responses.rs
+745
View File
@@ -0,0 +1,745 @@
//! OpenAI Responses API 类型。
use std::collections::BTreeSet;
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use crate::protocol::codex::{
responses::{CodexCompactRequest, CodexResponsesRequest},
schema::{prepare_schema, reconvert_tuple_values},
sse::{parse_sse_events, SseError},
};
/// OpenAI 请求声明的响应格式类别。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResponseFormat {
/// 普通文本输出。
Text,
/// JSON object 输出。
JsonObject,
/// JSON schema 输出。
JsonSchema,
}
/// Responses API 请求体。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct OpenAiResponsesRequest {
/// 模型名。
pub model: String,
/// 输入内容。
#[serde(default = "default_responses_input")]
pub input: Value,
/// 指令文本。
#[serde(default)]
pub instructions: Option<String>,
/// reasoning 配置。
#[serde(default)]
pub reasoning: Option<Value>,
/// 工具定义。
#[serde(default)]
pub tools: Option<Value>,
/// service tier。
#[serde(default)]
pub service_tier: Option<String>,
/// 工具选择策略。
#[serde(default)]
pub tool_choice: Option<Value>,
/// 是否允许并行工具调用。
#[serde(default)]
pub parallel_tool_calls: Option<bool>,
/// 输出文本格式配置。
#[serde(default)]
pub text: Option<Value>,
/// 图片生成开关。
#[serde(default)]
pub generate: Option<bool>,
/// 提示缓存键。
#[serde(default)]
pub prompt_cache_key: Option<String>,
/// include 列表。
#[serde(default)]
pub include: Option<Value>,
/// 传给上游的 client metadata。
#[serde(default)]
pub client_metadata: Option<Value>,
/// 前一个响应 ID。
#[serde(default)]
pub previous_response_id: Option<String>,
/// Codex turn state 透传头。
#[serde(default, rename = "turnState", alias = "turn_state")]
pub turn_state: Option<String>,
/// Codex turn metadata 透传头。
#[serde(default, rename = "turnMetadata", alias = "turn_metadata")]
pub turn_metadata: Option<String>,
/// beta features 透传头。
#[serde(default, rename = "betaFeatures", alias = "beta_features")]
pub beta_features: Option<String>,
/// 客户端版本头。
#[serde(default)]
pub version: Option<String>,
/// timing metrics 透传头。
#[serde(
default,
rename = "includeTimingMetrics",
alias = "include_timing_metrics"
)]
pub include_timing_metrics: Option<String>,
/// codex window id。
#[serde(default, rename = "codexWindowId", alias = "codex_window_id")]
pub codex_window_id: Option<String>,
/// 父线程 ID。
#[serde(default, rename = "parentThreadId", alias = "parent_thread_id")]
pub parent_thread_id: Option<String>,
/// 是否偏好 WebSocket 传输。
#[serde(default)]
pub use_websocket: Option<bool>,
/// 是否流式返回。
#[serde(default = "default_responses_stream")]
pub stream: bool,
}
fn default_responses_input() -> Value {
Value::Array(Vec::new())
}
fn default_responses_stream() -> bool {
true
}
/// 将 OpenAI Responses 请求转换为 Codex Responses 请求。
pub fn translate_response_to_codex(request: OpenAiResponsesRequest) -> CodexResponsesRequest {
let prepared_text = prepare_text_format(request.text, true);
let mut codex_request = CodexResponsesRequest::new_http_sse(
request.model,
request.instructions.unwrap_or_default(),
sanitize_responses_input(request.input),
);
codex_request.previous_response_id = request.previous_response_id;
codex_request.turn_state = request.turn_state;
codex_request.turn_metadata = non_empty_string(request.turn_metadata);
codex_request.beta_features = non_empty_string(request.beta_features);
codex_request.include_timing_metrics = non_empty_string(request.include_timing_metrics);
codex_request.version = non_empty_string(request.version);
codex_request.codex_window_id = non_empty_string(request.codex_window_id);
codex_request.parent_thread_id = non_empty_string(request.parent_thread_id);
codex_request.reasoning = responses_reasoning(request.reasoning);
codex_request.tools = non_empty_array(request.tools);
codex_request.tool_choice = request.tool_choice;
codex_request.parallel_tool_calls = request.parallel_tool_calls;
codex_request.text = prepared_text.text;
codex_request.tuple_schema = prepared_text.tuple_schema;
codex_request.generate = request.generate;
codex_request.service_tier = non_empty_string(request.service_tier);
codex_request.prompt_cache_key = non_empty_string(request.prompt_cache_key);
codex_request.explicit_prompt_cache_key = codex_request.prompt_cache_key.is_some();
codex_request.include = string_array(request.include);
codex_request.client_metadata = sanitize_client_metadata(request.client_metadata);
ensure_reasoning_include(&mut codex_request);
match request.use_websocket {
Some(true) => codex_request.use_websocket = true,
Some(false) => codex_request.force_http_sse = true,
None => {}
}
codex_request
}
/// 将 OpenAI Responses 请求转换为 Codex compact 请求。
pub fn translate_response_to_compact(request: OpenAiResponsesRequest) -> CodexCompactRequest {
CodexCompactRequest {
model: request.model,
input: sanitize_responses_input(request.input),
instructions: request.instructions.unwrap_or_default(),
tools: non_empty_array(request.tools),
parallel_tool_calls: request.parallel_tool_calls,
reasoning: compact_reasoning(request.reasoning),
text: prepare_text_format(request.text, false).text,
}
}
struct PreparedTextFormat {
text: Option<Value>,
tuple_schema: Option<Value>,
}
fn prepare_text_format(text: Option<Value>, prepare_tuple_schema: bool) -> PreparedTextFormat {
let Some(Value::Object(text)) = text else {
return PreparedTextFormat {
text: None,
tuple_schema: None,
};
};
let Some(Value::Object(format)) = text.get("format") else {
return PreparedTextFormat {
text: None,
tuple_schema: None,
};
};
let Some(format_type) = format.get("type").and_then(Value::as_str) else {
return PreparedTextFormat {
text: None,
tuple_schema: None,
};
};
let mut sanitized_format = Map::new();
sanitized_format.insert("type".to_string(), Value::String(format_type.to_string()));
if let Some(name) = format.get("name").and_then(Value::as_str) {
sanitized_format.insert("name".to_string(), Value::String(name.to_string()));
}
let mut tuple_schema = None;
if let Some(Value::Object(schema)) = format.get("schema") {
let schema = Value::Object(schema.clone());
let schema = if prepare_tuple_schema {
let prepared = prepare_schema(schema);
tuple_schema = prepared.original_schema;
prepared.schema
} else {
schema
};
sanitized_format.insert("schema".to_string(), schema);
}
if let Some(strict) = format.get("strict").and_then(Value::as_bool) {
sanitized_format.insert("strict".to_string(), Value::Bool(strict));
}
PreparedTextFormat {
text: Some(json!({"format": sanitized_format})),
tuple_schema,
}
}
fn sanitize_responses_input(input: Value) -> Vec<Value> {
match input {
Value::Array(items) => sanitize_codex_input_items(items),
Value::Null => Vec::new(),
value => vec![value],
}
}
fn sanitize_codex_input_items(input: Vec<Value>) -> Vec<Value> {
input
.into_iter()
.filter_map(|item| {
let Value::Object(object) = item else {
return Some(item);
};
match object.get("type").and_then(Value::as_str) {
Some("reasoning") => sanitize_reasoning_item(&object),
Some("compaction") => sanitize_compaction_item(&object),
_ => Some(Value::Object(object)),
}
})
.collect()
}
fn sanitize_reasoning_item(item: &Map<String, Value>) -> Option<Value> {
let id = non_empty_str(item.get("id"))?;
let summary = sanitize_summary(item.get("summary"))?;
let mut sanitized = Map::new();
sanitized.insert("type".to_string(), Value::String("reasoning".to_string()));
sanitized.insert("id".to_string(), Value::String(id.to_string()));
sanitized.insert("summary".to_string(), Value::Array(summary));
if let Some(status) = item
.get("status")
.and_then(Value::as_str)
.filter(|status| matches!(*status, "in_progress" | "completed" | "incomplete"))
{
sanitized.insert("status".to_string(), Value::String(status.to_string()));
}
if let Some(encrypted_content) = non_empty_str(item.get("encrypted_content")) {
sanitized.insert(
"encrypted_content".to_string(),
Value::String(encrypted_content.to_string()),
);
}
if let Some(content) = sanitize_reasoning_content(item.get("content")) {
sanitized.insert("content".to_string(), Value::Array(content));
}
Some(Value::Object(sanitized))
}
fn sanitize_summary(value: Option<&Value>) -> Option<Vec<Value>> {
let Value::Array(parts) = value? else {
return None;
};
Some(
parts
.iter()
.filter_map(|part| {
let Value::Object(part) = part else {
return None;
};
if part.get("type").and_then(Value::as_str) != Some("summary_text") {
return None;
}
let text = part.get("text").and_then(Value::as_str)?;
Some(json!({"type": "summary_text", "text": text}))
})
.collect(),
)
}
fn sanitize_reasoning_content(value: Option<&Value>) -> Option<Vec<Value>> {
let Value::Array(parts) = value? else {
return None;
};
let content = parts
.iter()
.filter_map(|part| {
let Value::Object(part) = part else {
return None;
};
if part.get("type").and_then(Value::as_str) != Some("reasoning_text") {
return None;
}
let text = part.get("text").and_then(Value::as_str)?;
Some(json!({"type": "reasoning_text", "text": text}))
})
.collect::<Vec<_>>();
(!content.is_empty()).then_some(content)
}
fn sanitize_compaction_item(item: &Map<String, Value>) -> Option<Value> {
let encrypted_content = non_empty_str(item.get("encrypted_content"))?;
let mut sanitized = Map::new();
sanitized.insert("type".to_string(), Value::String("compaction".to_string()));
sanitized.insert(
"encrypted_content".to_string(),
Value::String(encrypted_content.to_string()),
);
if let Some(id) = non_empty_str(item.get("id")) {
sanitized.insert("id".to_string(), Value::String(id.to_string()));
}
Some(Value::Object(sanitized))
}
fn responses_reasoning(reasoning: Option<Value>) -> Option<Value> {
let Value::Object(input) = reasoning? else {
return None;
};
let effort = input.get("effort").and_then(Value::as_str);
let summary = input
.get("summary")
.and_then(Value::as_str)
.unwrap_or("auto");
let mut output = Map::new();
output.insert("summary".to_string(), Value::String(summary.to_string()));
if let Some(effort) = effort {
output.insert("effort".to_string(), Value::String(effort.to_string()));
}
Some(Value::Object(output))
}
fn compact_reasoning(reasoning: Option<Value>) -> Option<Value> {
let Value::Object(input) = reasoning? else {
return None;
};
let mut output = Map::new();
if let Some(effort) = input.get("effort").and_then(Value::as_str) {
output.insert("effort".to_string(), Value::String(effort.to_string()));
}
if let Some(summary) = input.get("summary").and_then(Value::as_str) {
output.insert("summary".to_string(), Value::String(summary.to_string()));
}
(!output.is_empty()).then_some(Value::Object(output))
}
fn ensure_reasoning_include(request: &mut CodexResponsesRequest) {
if request.reasoning.is_none() {
return;
}
if request
.include
.as_ref()
.is_some_and(|include| !include.is_empty())
{
return;
}
request.include = Some(vec!["reasoning.encrypted_content".to_string()]);
}
fn sanitize_client_metadata(client_metadata: Option<Value>) -> Option<Value> {
let Value::Object(input) = client_metadata? else {
return None;
};
let metadata = input
.into_iter()
.filter_map(|(key, value)| {
value
.as_str()
.map(|value| (key, Value::String(value.to_string())))
})
.collect::<Map<_, _>>();
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
fn non_empty_array(value: Option<Value>) -> Option<Vec<Value>> {
let Value::Array(values) = value? else {
return None;
};
(!values.is_empty()).then_some(values)
}
fn string_array(value: Option<Value>) -> Option<Vec<String>> {
let Value::Array(values) = value? else {
return None;
};
values
.into_iter()
.map(|value| match value {
Value::String(value) => Some(value),
_ => None,
})
.collect()
}
fn non_empty_string(value: Option<String>) -> Option<String> {
value.and_then(|value| {
let trimmed = value.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
})
}
fn non_empty_str(value: Option<&Value>) -> Option<&str> {
let value = value?.as_str()?;
(!value.trim().is_empty()).then_some(value)
}
/// 从 Codex SSE 收集出的 Responses API 非流式结果。
#[derive(Debug, Clone, PartialEq)]
pub enum CollectedResponse {
/// 上游返回了完成响应。
Completed(Value),
/// 上游返回了失败事件。
Failed(ResponsesSseFailure),
/// 上游没有返回 `response.completed`。
MissingCompleted,
/// 上游完成响应没有可见输出。
Empty,
}
/// Responses SSE 失败事件。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ResponsesSseFailure {
/// SSE 事件名。
pub event: String,
/// 错误消息。
pub message: String,
/// 上游错误代码。
pub upstream_code: Option<String>,
}
impl ResponsesSseFailure {
fn from_event(event: &str, value: &Value) -> Self {
Self {
event: event.to_string(),
message: failure_message(value).unwrap_or_else(|| "Codex upstream SSE failed".into()),
upstream_code: failure_code(value),
}
}
}
/// 从完成 SSE 中提取会话亲和性和 replay 元数据。
#[derive(Debug, Clone, PartialEq)]
pub struct CompletedResponseMetadata {
/// 响应 ID。
pub response_id: String,
/// 完成响应中包含的函数调用 ID。
pub function_call_ids: Vec<String>,
/// 可用于 reasoning replay 的响应条目。
pub replay_items: Vec<Value>,
}
/// 将 Codex Responses SSE 完成响应收集为非流式 Responses API JSON。
///
/// # Errors
///
/// 当输入不是合法 SSE 流时，返回 [`SseError`]。
pub fn response_from_codex_sse(
body: &str,
tuple_schema: Option<&Value>,
) -> Result<CollectedResponse, SseError> {
let events = parse_sse_events(body)?;
let mut output_text = String::new();
let mut output_items = Vec::new();
let mut completed_response = None;
let mut failed_response = None;
for event in events {
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
continue;
};
match event.event.as_deref() {
Some("response.output_text.delta") => {
if let Some(delta) = value.get("delta").and_then(Value::as_str) {
output_text.push_str(delta);
}
}
Some("response.output_item.done") => {
if let Some(item) = value.get("item") {
output_items.push(item.clone());
}
}
Some("response.completed") => {
completed_response = value.get("response").cloned();
}
Some(event_name @ ("error" | "response.failed")) if failed_response.is_none() => {
failed_response = Some(ResponsesSseFailure::from_event(event_name, &value));
}
_ => {}
}
}
if let Some(failure) = failed_response {
return Ok(CollectedResponse::Failed(failure));
}
let Some(mut response) = completed_response else {
return Ok(CollectedResponse::MissingCompleted);
};
if is_empty_response(&response, &output_text, &output_items) {
return Ok(CollectedResponse::Empty);
}
ensure_completed_response_output(&mut response, &output_items, &output_text);
reconvert_completed_response_tuple_values(&mut response, tuple_schema);
sync_output_text_from_output(&mut response);
Ok(CollectedResponse::Completed(response))
}
/// 从 Codex Responses SSE 中提取完成响应元数据。
///
/// # Errors
///
/// 当输入不是合法 SSE 流时，返回 [`SseError`]。
pub fn completed_response_metadata(
body: &str,
) -> Result<Option<CompletedResponseMetadata>, SseError> {
let events = parse_sse_events(body)?;
let mut response_id = None;
let mut function_call_ids = BTreeSet::new();
let mut replay_items = Vec::new();
for event in events {
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
continue;
};
match event.event.as_deref() {
Some("response.output_item.done") => {
if let Some(item) = value.get("item") {
collect_response_replay_items(item, &mut replay_items);
}
if let Some(call_id) = value.pointer("/item/call_id").and_then(Value::as_str) {
function_call_ids.insert(call_id.to_string());
}
}
Some("response.completed") => {
response_id = value
.pointer("/response/id")
.and_then(Value::as_str)
.map(ToString::to_string);
collect_response_function_call_ids(&value, &mut function_call_ids);
if let Some(output) = value.pointer("/response/output").and_then(Value::as_array) {
for item in output {
collect_response_replay_items(item, &mut replay_items);
}
}
}
_ => {}
}
}
Ok(response_id.map(|response_id| CompletedResponseMetadata {
response_id,
function_call_ids: function_call_ids.into_iter().collect(),
replay_items,
}))
}
/// 对单个 Responses SSE 事件的数据执行 tuple schema 回转换。
pub fn reconvert_responses_sse_event_tuple_values(
event_name: Option<&str>,
mut data: Value,
tuple_schema: &Value,
) -> Value {
match responses_event_type(event_name, &data) {
Some("response.output_text.delta") => {
reconvert_output_text_delta_tuple_values(&mut data, tuple_schema);
}
Some("response.output_item.done") => {
if let Some(item) = data.get_mut("item") {
reconvert_output_item_tuple_values(item, tuple_schema);
}
}
Some("response.completed") => {
if let Some(response) = data.get_mut("response") {
reconvert_completed_response_tuple_values(response, Some(tuple_schema));
sync_output_text_from_output(response);
}
}
_ => {}
}
data
}
fn responses_event_type<'a>(event_name: Option<&'a str>, data: &'a Value) -> Option<&'a str> {
event_name.or_else(|| data.get("type").and_then(Value::as_str))
}
fn reconvert_output_text_delta_tuple_values(data: &mut Value, tuple_schema: &Value) {
let Some(delta) = data.get("delta").and_then(Value::as_str) else {
return;
};
let Ok(parsed) = serde_json::from_str::<Value>(delta) else {
return;
};
let reconverted = reconvert_tuple_values(parsed, tuple_schema);
data["delta"] = Value::String(reconverted.to_string());
}
fn is_empty_response(response: &Value, output_text: &str, output_items: &[Value]) -> bool {
if !output_text.trim().is_empty() || !output_items.is_empty() {
return false;
}
response
.pointer("/usage/output_tokens")
.and_then(Value::as_u64)
.unwrap_or_default()
== 0
}
fn ensure_completed_response_output(
response: &mut Value,
output_items: &[Value],
output_text: &str,
) {
let output_is_empty = response
.get("output")
.and_then(Value::as_array)
.is_none_or(Vec::is_empty);
if !output_is_empty {
return;
}
if !output_items.is_empty() {
response["output"] = Value::Array(output_items.to_vec());
return;
}
if output_text.is_empty() {
return;
}
response["output"] = json!([{
"type": "message",
"status": "completed",
"role": "assistant",
"content": [{
"type": "output_text",
"text": output_text,
"annotations": []
}]
}]);
}
fn reconvert_completed_response_tuple_values(response: &mut Value, tuple_schema: Option<&Value>) {
let Some(tuple_schema) = tuple_schema else {
return;
};
let Some(items) = response.get_mut("output").and_then(Value::as_array_mut) else {
return;
};
for item in items {
reconvert_output_item_tuple_values(item, tuple_schema);
}
}
fn reconvert_output_item_tuple_values(item: &mut Value, tuple_schema: &Value) {
let Some(content) = item.get_mut("content").and_then(Value::as_array_mut) else {
return;
};
for part in content {
let part_type = part.get("type").and_then(Value::as_str);
if !matches!(part_type, Some("output_text" | "text")) {
continue;
}
let Some(text) = part.get("text").and_then(Value::as_str) else {
continue;
};
let Ok(parsed) = serde_json::from_str::<Value>(text) else {
continue;
};
let reconverted = reconvert_tuple_values(parsed, tuple_schema);
part["text"] = Value::String(reconverted.to_string());
}
}
fn sync_output_text_from_output(response: &mut Value) {
let Some(items) = response.get("output").and_then(Value::as_array) else {
return;
};
let texts = items
.iter()
.filter_map(output_text_from_item)
.collect::<Vec<_>>();
if texts.is_empty() {
return;
}
response["output_text"] = Value::String(texts.join("\n\n"));
}
fn collect_response_function_call_ids(value: &Value, function_call_ids: &mut BTreeSet<String>) {
let Some(output) = value.pointer("/response/output").and_then(Value::as_array) else {
return;
};
for item in output {
if let Some(call_id) = item.get("call_id").and_then(Value::as_str) {
function_call_ids.insert(call_id.to_string());
}
}
}
fn collect_response_replay_items(item: &Value, replay_items: &mut Vec<Value>) {
if matches!(
item.get("type").and_then(Value::as_str),
Some("reasoning" | "function_call")
) {
replay_items.push(item.clone());
}
}
fn output_text_from_item(item: &Value) -> Option<String> {
let content = item.get("content")?.as_array()?;
let text = content
.iter()
.filter_map(|part| {
let part_type = part.get("type")?.as_str()?;
if part_type != "output_text" && part_type != "text" {
return None;
}
part.get("text")?.as_str()
})
.collect::<String>();
(!text.is_empty()).then_some(text)
}
fn failure_message(value: &Value) -> Option<String> {
value
.pointer("/response/error/message")
.or_else(|| value.pointer("/error/message"))
.or_else(|| value.get("message"))
.and_then(Value::as_str)
.filter(|message| !message.trim().is_empty())
.map(ToString::to_string)
}
fn failure_code(value: &Value) -> Option<String> {
value
.pointer("/response/error/code")
.or_else(|| value.pointer("/error/code"))
.or_else(|| value.get("code"))
.and_then(Value::as_str)
.filter(|code| !code.trim().is_empty())
.map(ToString::to_string)
}
crates/core/src/serving/affinity.rs
+238
View File
@@ -0,0 +1,238 @@
//! 会话亲和性策略。
use std::collections::BTreeMap;
use chrono::{DateTime, Duration, Utc};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::{
gateway::conversation::derive_stable_conversation_key,
protocol::codex::responses::CodexResponsesRequest,
};
/// 会话亲和性键。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionAffinityKey {
/// 客户端会话 ID。
pub conversation_id: String,
}
/// 会话亲和性条目。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionAffinityEntry {
/// 创建该响应的账号 ID。
pub account_id: String,
/// 对话 ID。
pub conversation_id: String,
/// 上游 turn state。
pub turn_state: Option<String>,
/// 指令哈希。
pub instructions_hash: Option<String>,
/// 输入 token 数。
pub input_tokens: Option<u64>,
/// 函数调用 ID 列表。
pub function_call_ids: Vec<String>,
/// 用于区分并发对话分支的变体哈希。
pub variant_hash: Option<String>,
/// 创建时间。
pub created_at: DateTime<Utc>,
}
/// 持久化的会话亲和性条目。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StoredSessionAffinity {
/// 响应 ID。
pub response_id: String,
/// 亲和性条目。
pub entry: SessionAffinityEntry,
/// 持久化过期时间。
pub expires_at: DateTime<Utc>,
}
/// 纯内存会话亲和性映射。
#[derive(Debug, Clone)]
pub struct SessionAffinityMap {
entries: BTreeMap<String, SessionAffinityEntry>,
ttl: Duration,
}
impl SessionAffinityMap {
/// 构造指定 TTL 的会话亲和性映射。
pub fn new(ttl: Duration) -> Self {
Self {
entries: BTreeMap::new(),
ttl,
}
}
/// 记录响应 ID 到亲和性条目的映射。
pub fn record(&mut self, response_id: String, entry: SessionAffinityEntry) {
self.entries.insert(response_id, entry);
}
/// 从持久化记录恢复未过期条目。
pub fn restore(&mut self, records: Vec<StoredSessionAffinity>, now: DateTime<Utc>) -> usize {
let mut restored = 0usize;
for record in records {
if self.stored_record_is_active(&record, now) {
self.entries.insert(record.response_id, record.entry);
restored = restored.saturating_add(1);
}
}
restored
}
/// 根据响应 ID 查找账号 ID。
pub fn lookup_account(&self, response_id: &str, now: DateTime<Utc>) -> Option<String> {
self.active_entry(response_id, now)
.map(|entry| entry.account_id.clone())
}
/// 根据响应 ID 查找对话 ID。
pub fn lookup_conversation_id(&self, response_id: &str, now: DateTime<Utc>) -> Option<String> {
self.active_entry(response_id, now)
.map(|entry| entry.conversation_id.clone())
}
/// 根据响应 ID 查找 turn state。
pub fn lookup_turn_state(&self, response_id: &str, now: DateTime<Utc>) -> Option<String> {
self.active_entry(response_id, now)
.and_then(|entry| entry.turn_state.clone())
}
/// 根据响应 ID 查找指令哈希。
pub fn lookup_instructions_hash(
&self,
response_id: &str,
now: DateTime<Utc>,
) -> Option<String> {
self.active_entry(response_id, now)
.and_then(|entry| entry.instructions_hash.clone())
}
/// 根据响应 ID 查找函数调用 ID 列表。
pub fn lookup_function_call_ids(&self, response_id: &str, now: DateTime<Utc>) -> Vec<String> {
self.active_entry(response_id, now)
.map(|entry| entry.function_call_ids.clone())
.unwrap_or_default()
}
/// 查找指定对话和可选变体下最新的响应 ID。
pub fn lookup_latest_response_by_conversation(
&self,
conversation_id: &str,
max_age: Option<Duration>,
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<String> {
self.entries
.iter()
.filter(|(_, entry)| entry.conversation_id == conversation_id)
.filter(|(_, entry)| self.entry_is_active(entry, now))
.filter(|(_, entry)| {
variant_hash
.is_none_or(|variant_hash| entry.variant_hash.as_deref() == Some(variant_hash))
})
.filter(|(_, entry)| {
max_age.is_none_or(|max_age| now.signed_duration_since(entry.created_at) <= max_age)
})
.max_by_key(|(_, entry)| entry.created_at)
.map(|(response_id, _)| response_id.clone())
}
/// 删除响应 ID 映射。
pub fn forget(&mut self, response_id: &str) -> bool {
self.entries.remove(response_id).is_some()
}
/// 返回当前内存条目数。
pub fn len(&self) -> usize {
self.entries.len()
}
/// 返回是否没有内存条目。
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
fn active_entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<&SessionAffinityEntry> {
self.entries
.get(response_id)
.filter(|entry| self.entry_is_active(entry, now))
}
fn stored_record_is_active(&self, record: &StoredSessionAffinity, now: DateTime<Utc>) -> bool {
record.expires_at > now && self.entry_is_active(&record.entry, now)
}
fn entry_is_active(&self, entry: &SessionAffinityEntry, now: DateTime<Utc>) -> bool {
entry
.created_at
.checked_add_signed(self.ttl)
.is_some_and(|expires_at| expires_at > now)
}
}
/// 准备用于区分并发分支的变体身份。
pub fn prepare_variant_identity(request: &mut CodexResponsesRequest) {
request.variant_identity = build_variant_identity(request);
}
/// 计算请求变体哈希。
pub fn compute_variant_hash(request: &CodexResponsesRequest) -> String {
compute_variant_hash_with_identity(request, request.variant_identity.as_deref())
}
fn compute_variant_hash_with_identity(
request: &CodexResponsesRequest,
identity: Option<&str>,
) -> String {
let tools_json = request
.tools
.as_ref()
.map_or_else(|| "[]".to_string(), |tools| tools_json(tools));
let mut hasher = Sha256::new();
hasher.update(request.instructions.as_bytes());
hasher.update(b"\0");
hasher.update(tools_json.as_bytes());
if let Some(identity) = identity
.map(str::trim)
.filter(|identity| !identity.is_empty())
{
hasher.update(b"\0");
hasher.update(identity.as_bytes());
}
hex::encode(hasher.finalize()).chars().take(12).collect()
}
fn tools_json(tools: &[Value]) -> String {
serde_json::to_string(tools).unwrap_or_else(|_| "[]".to_string())
}
fn build_variant_identity(request: &CodexResponsesRequest) -> Option<String> {
let mut parts = Vec::with_capacity(2);
if let Some(window_id) = non_empty_str(request.codex_window_id.as_deref()) {
parts.push(format!("window:{window_id}"));
}
if request.explicit_prompt_cache_key
|| non_empty_str(request.client_conversation_id.as_deref()).is_some()
{
if let Some(anchor) = derive_stable_conversation_key(request) {
parts.push(format!("anchor:{anchor}"));
}
}
(!parts.is_empty()).then(|| parts.join("\0"))
}
fn non_empty_str(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
/// 计算 instructions 哈希。
pub fn hash_instructions(instructions: Option<&str>) -> String {
let mut hasher = Sha256::new();
hasher.update(instructions.unwrap_or_default().as_bytes());
hex::encode(hasher.finalize())
}
crates/core/src/serving/chat.rs
+13
View File
@@ -0,0 +1,13 @@
//! Chat 补全编排。
use crate::protocol::{
codex::responses::CodexResponsesRequest,
openai::chat::{translate_chat_to_codex, ChatCompletionRequest, ChatTranslationError},
};
/// 将 Chat 请求准备为上游 Responses 请求。
pub fn prepare_chat_request(
request: ChatCompletionRequest,
) -> Result<CodexResponsesRequest, ChatTranslationError> {
translate_chat_to_codex(request)
}
crates/core/src/serving/errors.rs
+17
View File
@@ -0,0 +1,17 @@
//! 服务编排错误。
use thiserror::Error;
/// 服务编排错误。
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ServingError {
/// 没有可用账号。
#[error("no available account")]
NoAvailableAccount,
/// 上游请求失败。
#[error("upstream request failed: {message}")]
Upstream {
/// 错误说明。
message: String,
},
}
crates/core/src/serving/fallback.rs
+16
View File
@@ -0,0 +1,16 @@
//! 上游回退策略。
/// 判断状态码是否表示上游限流。
pub fn status_code_is_rate_limited(status_code: u16) -> bool {
status_code == 429
}
/// 判断状态码是否表示账号配额耗尽。
pub fn status_code_is_quota_exhausted(status_code: u16) -> bool {
status_code == 402
}
/// 判断状态码是否表示可换账号重试的临时上游错误。
pub fn status_code_is_transient_upstream(status_code: u16) -> bool {
matches!(status_code, 500..=599)
}
src/codex/serving/dispatch/implicit_resume.rs → crates/core/src/serving/implicit_resume.rs
+56 -13
View File
@@ -1,20 +1,27 @@
//! Responses 隐式续接策略。
use std::collections::HashSet;
use serde_json::Value;
use crate::codex::gateway::transport::types::CodexResponsesRequest;
#[derive(Debug, Clone)]
pub(crate) struct ImplicitResumeSnapshot {
input: Vec<Value>,
previous_response_id: Option<String>,
turn_state: Option<String>,
use_websocket: bool,
force_http_sse: bool,
/// 隐式续接前的可恢复请求状态。
#[derive(Debug, Clone, PartialEq)]
pub struct ImplicitResumeSnapshot {
/// 完整输入历史。
pub input: Vec<Value>,
/// 原始 previous_response_id。
pub previous_response_id: Option<String>,
/// 原始 turn state。
pub turn_state: Option<String>,
/// 原始 WebSocket 偏好。
pub use_websocket: bool,
/// 原始 HTTP SSE 强制标记。
pub force_http_sse: bool,
}
impl ImplicitResumeSnapshot {
pub(super) fn capture(request: &CodexResponsesRequest) -> Self {
/// 捕获当前请求中隐式续接会改写的字段。
pub fn capture(request: &crate::protocol::codex::responses::CodexResponsesRequest) -> Self {
Self {
input: request.input.clone(),
previous_response_id: request.previous_response_id.clone(),
@@ -24,7 +31,8 @@ impl ImplicitResumeSnapshot {
}
}
pub(crate) fn restore(self, request: &mut CodexResponsesRequest) {
/// 恢复隐式续接前的请求字段。
pub fn restore(self, request: &mut crate::protocol::codex::responses::CodexResponsesRequest) {
request.input = self.input;
request.previous_response_id = self.previous_response_id;
request.turn_state = self.turn_state;
@@ -33,7 +41,8 @@ impl ImplicitResumeSnapshot {
}
}
pub(super) fn continuation_input_start(input: &[Value]) -> usize {
/// 返回续接输入在完整输入中的起始位置。
pub fn continuation_input_start(input: &[Value]) -> usize {
let mut last_model_output_index = None;
for (index, item) in input.iter().enumerate() {
if item.get("role").is_some() {
@@ -49,7 +58,8 @@ pub(super) fn continuation_input_start(input: &[Value]) -> usize {
last_model_output_index.map_or(0, |index| index.saturating_add(1))
}
pub(super) fn implicit_resume_allowed(
/// 判断完整输入历史是否可以隐式续接到已记录响应。
pub fn implicit_resume_allowed(
continuation_input: &[Value],
full_input: &[Value],
stored_function_call_ids: &[String],
@@ -98,3 +108,36 @@ fn inline_function_call_ids(input: &[Value]) -> HashSet<String> {
.map(ToString::to_string)
.collect()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn continuation_input_start_should_follow_last_assistant_message() {
let input = vec![
json!({"role": "user", "content": "first"}),
json!({"role": "assistant", "content": "cached"}),
json!({"role": "user", "content": "continue"}),
];
assert_eq!(continuation_input_start(&input), 2);
}
#[test]
fn implicit_resume_allowed_should_reject_unmatched_function_call_output() {
let input = vec![json!({
"type": "function_call_output",
"call_id": "missing",
"output": "tool output"
})];
assert!(!implicit_resume_allowed(
&input,
&input,
&["expected".to_string()]
));
}
}
crates/core/src/serving/mod.rs
+26
View File
@@ -0,0 +1,26 @@
//! 对外服务编排策略。
/// 亲和性策略。
pub mod affinity;
/// Chat 补全编排。
pub mod chat;
/// 服务错误。
pub mod errors;
/// 回退策略。
pub mod fallback;
/// Responses 隐式续接策略。
pub mod implicit_resume;
/// 配额策略。
pub mod quota;
/// Responses reasoning replay 缓存策略。
pub mod reasoning_replay;
/// 错误恢复策略。
pub mod recovery;
/// Responses 编排。
pub mod responses;
/// 路由策略。
pub mod routing;
/// 流生命周期策略。
pub mod stream;
/// 用量跟踪策略。
pub mod usage;
src/codex/accounts/service/quota.rs → crates/core/src/serving/quota.rs
+8 -148
View File
@@ -1,162 +1,21 @@
use crate::codex::accounts::model::AccountStatus;
use super::{
health::{apply_codex_account_error, fetch_account_usage, public_codex_error},
AccountQuotaError, AccountQuotaResult, AccountQuotaWarning, AccountQuotaWarnings,
AccountService, AccountServiceError, QuotaWarningLevel, QuotaWarningWindow,
};
//! 配额检查策略。
use serde_json::{json, Map, Value};
impl AccountService {
pub async fn account_quota(
&self,
account_id: &str,
request_id: &str,
) -> Result<AccountQuotaResult, AccountQuotaError> {
let repo = self
.repository
.as_ref()
.ok_or(AccountQuotaError::RepositoryUnavailable)?;
let account = match repo.get(account_id).await {
Ok(Some(account)) => account,
Ok(None) => return Err(AccountQuotaError::NotFound),
Err(_) => return Err(AccountQuotaError::Load),
};
if account.status != AccountStatus::Active {
return Err(AccountQuotaError::Inactive(account.status));
}
match fetch_account_usage(self, &account, request_id).await {
Ok(raw) => {
let quota = quota_from_usage(&raw);
if repo
.update_quota_json(&account.id, &quota.to_string())
.await
.is_err()
{
return Err(AccountQuotaError::StoreQuota);
}
Ok(AccountQuotaResult { quota, raw })
}
Err(error) => {
apply_codex_account_error(self, repo, &account, &error).await;
Err(AccountQuotaError::Fetch(public_codex_error(&error)))
}
}
}
pub async fn quota_warnings(&self) -> Result<AccountQuotaWarnings, AccountServiceError> {
let snapshots = self
.repository()?
.list_quota_snapshots()
.await
.map_err(|_| AccountServiceError::QuotaWarnings)?;
let primary_thresholds = sorted_thresholds(&self.config.quota.warning_thresholds.primary);
let secondary_thresholds =
sorted_thresholds(&self.config.quota.warning_thresholds.secondary);
let mut warnings = Vec::new();
let mut updated_at = None;
for snapshot in snapshots {
let Ok(quota) = serde_json::from_str::<Value>(&snapshot.quota_json) else {
continue;
};
let before_len = warnings.len();
if let Some(warning) = warning_from_quota_window(
&snapshot.account_id,
snapshot.email.as_deref(),
&quota,
"rate_limit",
QuotaWarningWindow::Primary,
&primary_thresholds,
) {
warnings.push(warning);
}
if let Some(warning) = warning_from_quota_window(
&snapshot.account_id,
snapshot.email.as_deref(),
&quota,
"secondary_rate_limit",
QuotaWarningWindow::Secondary,
&secondary_thresholds,
) {
warnings.push(warning);
}
if warnings.len() > before_len {
updated_at = max_optional_datetime(updated_at, snapshot.quota_fetched_at);
}
}
Ok(AccountQuotaWarnings {
warnings,
updated_at,
})
}
/// 判断配额百分比是否达到阈值。
pub fn quota_reached(used_percent: f64, threshold: f64) -> bool {
used_percent >= threshold
}
fn warning_from_quota_window(
account_id: &str,
email: Option<&str>,
quota: &Value,
field: &str,
window: QuotaWarningWindow,
thresholds: &[u8],
) -> Option<AccountQuotaWarning> {
let quota_window = quota.get(field).filter(|value| !value.is_null())?;
let used_percent = quota_window
.get("used_percent")
.and_then(Value::as_f64)
.filter(|value| value.is_finite())?;
let level = warning_level(used_percent, thresholds)?;
Some(AccountQuotaWarning {
account_id: account_id.to_string(),
email: email.map(str::to_string),
window,
level,
used_percent,
reset_at: quota_window.get("reset_at").and_then(Value::as_i64),
})
}
fn warning_level(used_percent: f64, thresholds: &[u8]) -> Option<QuotaWarningLevel> {
let matched_index = thresholds
.iter()
.rposition(|threshold| used_percent >= f64::from(*threshold))?;
if matched_index + 1 == thresholds.len() {
Some(QuotaWarningLevel::Critical)
} else {
Some(QuotaWarningLevel::Warning)
}
}
fn sorted_thresholds(thresholds: &[u8]) -> Vec<u8> {
let mut thresholds = thresholds.to_vec();
thresholds.sort_unstable();
thresholds.dedup();
thresholds
}
fn max_optional_datetime(
current: Option<chrono::DateTime<chrono::Utc>>,
candidate: Option<chrono::DateTime<chrono::Utc>>,
) -> Option<chrono::DateTime<chrono::Utc>> {
match (current, candidate) {
(Some(current), Some(candidate)) => Some(current.max(candidate)),
(Some(current), None) => Some(current),
(None, Some(candidate)) => Some(candidate),
(None, None) => None,
}
}
pub(crate) fn quota_from_usage(usage: &Value) -> Value {
/// 从 Codex usage 响应中提取持久化配额快照。
pub fn quota_from_usage(usage: &Value) -> Value {
let additional = usage
.get("additional_rate_limits")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let mut rate_limits_by_limit_id = Map::new();
for item in &additional {
let Some(limit_id) = item
.get("metered_feature")
@@ -185,6 +44,7 @@ pub(crate) fn quota_from_usage(usage: &Value) -> Value {
}),
);
}
let additional_review = additional.iter().find(|item| {
is_review_limit_id(item.get("metered_feature").and_then(Value::as_str))
|| is_review_limit_id(item.get("limit_name").and_then(Value::as_str))
src/codex/serving/dispatch/reasoning_replay.rs → crates/core/src/serving/reasoning_replay.rs
+122 -80
View File
@@ -1,15 +1,16 @@
use std::{collections::HashMap, time::Duration};
//! Responses reasoning replay 缓存策略。
use chrono::{DateTime, Utc};
use std::collections::{BTreeMap, HashSet};
use chrono::{DateTime, Duration, Utc};
use serde_json::{json, Map, Value};
use tokio::sync::RwLock;
const REASONING_REPLAY_TTL: Duration = Duration::from_secs(55 * 60);
const MAX_ENTRIES: usize = 512;
const MAX_ENTRY_BYTES: usize = 256 * 1024;
const MAX_TOTAL_BYTES: usize = 4 * 1024 * 1024;
#[derive(Debug, Clone)]
/// 单条 reasoning replay 缓存记录。
#[derive(Debug, Clone, PartialEq)]
struct ReasoningReplayEntry {
account_id: String,
conversation_id: String,
@@ -20,31 +21,39 @@ struct ReasoningReplayEntry {
sequence: u64,
}
#[derive(Debug, Default)]
struct ReasoningReplayState {
entries: HashMap<String, ReasoningReplayEntry>,
/// 纯内存 reasoning replay 缓存。
#[derive(Debug, Clone)]
pub struct ReasoningReplayCache {
entries: BTreeMap<String, ReasoningReplayEntry>,
total_bytes: usize,
next_sequence: u64,
}
#[derive(Debug, Default)]
pub(super) struct ReasoningReplayCache {
state: RwLock<ReasoningReplayState>,
ttl: Duration,
}
impl ReasoningReplayCache {
pub(super) async fn record(
&self,
response_id: &str,
/// 构造指定 TTL 的缓存。
pub fn new(ttl: Duration) -> Self {
Self {
entries: BTreeMap::new(),
total_bytes: 0,
next_sequence: 0,
ttl,
}
}
/// 记录响应对应的 replay 条目，返回实际保存的条目数。
pub fn record(
&mut self,
response_id: String,
account_id: &str,
conversation_id: &str,
variant_hash: &str,
items: &[Value],
now: DateTime<Utc>,
) -> usize {
let sanitized = sanitize_reasoning_replay_items(items);
let mut state = self.state.write().await;
cleanup_expired(&mut state, Utc::now());
delete_entry(&mut state, response_id);
self.cleanup_expired(now);
self.delete_entry(&response_id);
if sanitized.is_empty() {
return 0;
@@ -54,38 +63,39 @@ impl ReasoningReplayCache {
return 0;
}
let sequence = state.next_sequence;
state.next_sequence = state.next_sequence.saturating_add(1);
state.total_bytes = state.total_bytes.saturating_add(byte_size);
state.entries.insert(
response_id.to_string(),
let sequence = self.next_sequence;
self.next_sequence = self.next_sequence.saturating_add(1);
self.total_bytes = self.total_bytes.saturating_add(byte_size);
self.entries.insert(
response_id,
ReasoningReplayEntry {
account_id: account_id.to_string(),
conversation_id: conversation_id.to_string(),
variant_hash: variant_hash.to_string(),
items: sanitized,
byte_size,
created_at: Utc::now(),
created_at: now,
sequence,
},
);
evict_overflow(&mut state);
state
.entries
.get(response_id)
self.evict_overflow();
self.entries
.values()
.max_by_key(|entry| entry.sequence)
.map_or(0, |entry| entry.items.len())
}
pub(super) async fn lookup(
&self,
/// 查找指定响应与身份对应的 replay 条目。
pub fn lookup(
&mut self,
response_id: &str,
account_id: &str,
conversation_id: &str,
variant_hash: &str,
now: DateTime<Utc>,
) -> Vec<Value> {
let mut state = self.state.write().await;
cleanup_expired(&mut state, Utc::now());
let Some(entry) = state.entries.get(response_id) else {
self.cleanup_expired(now);
let Some(entry) = self.entries.get(response_id) else {
return Vec::new();
};
if entry.account_id != account_id
@@ -97,15 +107,16 @@ impl ReasoningReplayCache {
entry.items.clone()
}
pub(super) async fn evict_by_identity(
&self,
/// 按账号、对话、变体驱逐 replay 条目。
pub fn evict_by_identity(
&mut self,
account_id: &str,
conversation_id: &str,
variant_hash: &str,
now: DateTime<Utc>,
) -> usize {
let mut state = self.state.write().await;
cleanup_expired(&mut state, Utc::now());
let response_ids = state
self.cleanup_expired(now);
let response_ids = self
.entries
.iter()
.filter(|(_, entry)| {
@@ -117,14 +128,50 @@ impl ReasoningReplayCache {
.collect::<Vec<_>>();
let evicted = response_ids.len();
for response_id in response_ids {
delete_entry(&mut state, &response_id);
self.delete_entry(&response_id);
}
evicted
}
fn cleanup_expired(&mut self, now: DateTime<Utc>) {
let ttl = self.ttl;
let response_ids = self
.entries
.iter()
.filter(|(_, entry)| entry_expired(entry.created_at, ttl, now))
.map(|(response_id, _)| response_id.clone())
.collect::<Vec<_>>();
for response_id in response_ids {
self.delete_entry(&response_id);
}
}
fn evict_overflow(&mut self) {
while self.entries.len() > MAX_ENTRIES || self.total_bytes > MAX_TOTAL_BYTES {
let Some(response_id) = self.oldest_response_id() else {
return;
};
self.delete_entry(&response_id);
}
}
fn oldest_response_id(&self) -> Option<String> {
self.entries
.iter()
.min_by_key(|(_, entry)| (entry.created_at, entry.sequence))
.map(|(response_id, _)| response_id.clone())
}
fn delete_entry(&mut self, response_id: &str) {
let Some(entry) = self.entries.remove(response_id) else {
return;
};
self.total_bytes = self.total_bytes.saturating_sub(entry.byte_size);
}
}
fn sanitize_reasoning_replay_items(items: &[Value]) -> Vec<Value> {
let mut seen = std::collections::HashSet::new();
let mut seen = HashSet::new();
let mut sanitized = Vec::new();
for item in items {
let Some(replay_item) =
@@ -245,44 +292,39 @@ fn estimate_items_bytes(items: &[Value]) -> usize {
serde_json::to_vec(items).map_or(usize::MAX, |bytes| bytes.len())
}
fn cleanup_expired(state: &mut ReasoningReplayState, now: DateTime<Utc>) {
let response_ids = state
.entries
.iter()
.filter(|(_, entry)| entry_expired(entry.created_at, now))
.map(|(response_id, _)| response_id.clone())
.collect::<Vec<_>>();
for response_id in response_ids {
delete_entry(state, &response_id);
fn entry_expired(created_at: DateTime<Utc>, ttl: Duration, now: DateTime<Utc>) -> bool {
now.signed_duration_since(created_at) >= ttl
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn reasoning_replay_cache_should_sanitize_reasoning_items() {
let mut cache = ReasoningReplayCache::new(Duration::minutes(55));
let saved = cache.record(
"resp_1".to_string(),
"acct_1",
"conv_1",
"variant_1",
&[json!({
"type": "reasoning",
"id": "rs_1",
"status": "completed",
"summary": [{"type": "summary_text", "text": "kept"}],
"encrypted_content": "enc",
"extra": "drop"
})],
Utc::now(),
);
let replay = cache.lookup("resp_1", "acct_1", "conv_1", "variant_1", Utc::now());
assert_eq!(saved, 1);
assert_eq!(replay[0].get("extra"), None);
assert_eq!(replay[0]["encrypted_content"], "enc");
}
}
fn entry_expired(created_at: DateTime<Utc>, now: DateTime<Utc>) -> bool {
now.signed_duration_since(created_at)
.to_std()
.is_ok_and(|age| age > REASONING_REPLAY_TTL)
}
fn evict_overflow(state: &mut ReasoningReplayState) {
while state.entries.len() > MAX_ENTRIES || state.total_bytes > MAX_TOTAL_BYTES {
let Some(response_id) = oldest_response_id(state) else {
return;
};
delete_entry(state, &response_id);
}
}
fn oldest_response_id(state: &ReasoningReplayState) -> Option<String> {
state
.entries
.iter()
.min_by_key(|(_, entry)| (entry.created_at, entry.sequence))
.map(|(response_id, _)| response_id.clone())
}
fn delete_entry(state: &mut ReasoningReplayState, response_id: &str) {
let Some(entry) = state.entries.remove(response_id) else {
return;
};
state.total_bytes = state.total_bytes.saturating_sub(entry.byte_size);
}

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