Compare commits

...
12 Commits
Author SHA1 Message Date
zyycnandCodex 07829ca385 chore: bump version to 1.0.2
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 19:39:35 +08:00
zyycnandCodex 617bd676f0 refactor: move websocket keepalive into pump
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 19:38:15 +08:00
zyycnandCodex 0c18b44246 style: refine quota usage display
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 16:36:24 +08:00
zyycnandCodex 73658238e5 fix: refine admin usage reporting
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 16:25:01 +08:00
zyycnandCodex be3d648fee style: format responses proxy code
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 05:11:27 +08:00
zyycnandCodex d123704172 chore: bump version to 1.0.1
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 05:05:55 +08:00
zyycnandCodex 7d41151b72 feat: align responses proxy and api keys
Co-authored-by: Codex <noreply@openai.com>
2026-07-05 05:05:05 +08:00
zyycnandCodex 52dd6212ca refactor: remove chat completions path
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 23:04:23 +08:00
zyycnandCodex c0cd7ccbc1 fix: clean up quota dispatch audit issues
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 21:34:05 +08:00
zyycnandCodex 50d9b6eef8 fix: harden account scheduling and client context
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 19:56:32 +08:00
zyycnandCodex 94ef9784ba fix: serve vite public assets
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 16:45:25 +08:00
zyycnandCodex 38ab14fb19 fix: defer update event stream connection
Co-authored-by: Codex <noreply@openai.com>
2026-07-04 16:00:25 +08:00
134 changed files with 5167 additions and 7869 deletions
-1
View File
@@ -106,7 +106,6 @@ docker compose -f deploy/docker-compose.yml up -d
| 端点 | 说明 |
| --- | --- |
| `POST /v1/responses` | OpenAI Responses 兼容接口 |
| `POST /v1/chat/completions` | Chat Completions 兼容接口 |
| `GET /v1/models` | 模型列表 |
| `GET /v1/models/{id}` | 模型详情 |
| `GET /v1/models/catalog` | 管理端可见模型目录 |
+28
View File
@@ -1112,6 +1112,12 @@ dependencies = [
"pin-project-lite",
]
[[package]]
name = "http-range-header"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9171a2ea8a68358193d15dd5d70c1c10a2afc3e7e4c5bc92bc9f025cebd7359c"
[[package]]
name = "httparse"
version = "1.10.1"
@@ -1496,6 +1502,16 @@ version = "0.3.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
[[package]]
name = "mime_guess"
version = "2.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e"
dependencies = [
"mime",
"unicase",
]
[[package]]
name = "miniz_oxide"
version = "0.8.9"
@@ -2176,6 +2192,7 @@ version = "1.0.150"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
dependencies = [
"indexmap",
"itoa",
"memchr",
"serde",
@@ -2839,6 +2856,11 @@ dependencies = [
"http",
"http-body",
"http-body-util",
"http-range-header",
"httpdate",
"mime",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"tokio",
"tokio-util",
@@ -2994,6 +3016,12 @@ version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
[[package]]
name = "unicase"
version = "2.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142"
[[package]]
name = "unicode-bidi"
version = "0.3.18"
+3 -3
View File
@@ -29,7 +29,7 @@ rustls-pki-types = "1.14.1"
secrecy = "0.10.3"
semver = "1.0.27"
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.150"
serde_json = { version = "1.0.150", features = ["preserve_order"] }
sha2 = "0.10.9"
sqlx = { version = "0.9.0", features = ["runtime-tokio", "tls-rustls", "sqlite", "chrono", "uuid"] }
subtle = "2.6.1"
@@ -37,7 +37,8 @@ tar = "0.4.44"
thiserror = "2.0.18"
tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "signal", "fs", "time", "sync", "net", "test-util"] }
tokio-tungstenite = { version = "0.28.0", features = ["proxy", "rustls-tls-native-roots"] }
tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout"] }
tower = { version = "0.5.3", features = ["util"] }
tower-http = { version = "0.6.11", features = ["trace", "cors", "request-id", "timeout", "fs"] }
tracing = "0.1.44"
tracing-appender = "0.2.5"
tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] }
@@ -46,7 +47,6 @@ uuid = { version = "1.23.3", features = ["v4", "v7", "serde"] }
[dev-dependencies]
tempfile = "3.27.0"
tower = { version = "0.5.3", features = ["util"] }
wiremock = "0.6.5"
[patch.crates-io]
+69 -4
View File
@@ -1,11 +1,13 @@
//! 管理端账号 quota 展示视图。
use chrono::{DateTime, Utc};
use std::collections::HashMap;
use chrono::{DateTime, Duration, Utc};
use serde::Serialize;
use serde_json::Value;
use crate::infra::{
format::format_percent,
format::{format_compact_percent, format_tokens},
time::{china_datetime, china_relative_time},
};
@@ -29,8 +31,33 @@ struct AdminAccountQuotaWindowData {
label_display: String,
used_percent: Option<f64>,
used_percent_display: String,
token_usage_display: String,
reset_at_display: String,
window_used_display: String,
#[serde(skip)]
reset_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone)]
pub(crate) struct AdminAccountQuotaUsageWindow {
pub(crate) key: String,
pub(crate) start: DateTime<Utc>,
pub(crate) end: DateTime<Utc>,
}
#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct AdminAccountQuotaWindowTokenUsage {
total_tokens: u64,
}
impl AdminAccountQuotaWindowTokenUsage {
pub(crate) fn add_tokens(&mut self, tokens: u64) {
self.total_tokens = self.total_tokens.saturating_add(tokens);
}
fn display(self) -> String {
format_tokens(self.total_tokens)
}
}
impl Default for AdminAccountQuotaData {
@@ -48,6 +75,40 @@ impl AdminAccountQuotaData {
.iter()
.any(|window| window.used_percent.is_some_and(|percent| percent >= 80.0))
}
pub(crate) fn usage_windows(&self) -> Vec<AdminAccountQuotaUsageWindow> {
self.windows
.iter()
.filter_map(AdminAccountQuotaWindowData::usage_window)
.collect()
}
pub(crate) fn apply_token_usage(
&mut self,
usage_by_window: &HashMap<String, AdminAccountQuotaWindowTokenUsage>,
) {
for window in &mut self.windows {
window.token_usage_display =
match (window.usage_window(), usage_by_window.get(&window.key)) {
(Some(_), Some(usage)) => usage.display(),
(Some(_), None) => AdminAccountQuotaWindowTokenUsage::default().display(),
(None, _) => "-".to_string(),
};
}
}
}
impl AdminAccountQuotaWindowData {
fn usage_window(&self) -> Option<AdminAccountQuotaUsageWindow> {
let end = self.reset_at?;
let seconds = i64::try_from(self.window_seconds?).ok()?;
let start = end.checked_sub_signed(Duration::seconds(seconds))?;
(start <= end).then(|| AdminAccountQuotaUsageWindow {
key: self.key.clone(),
start,
end,
})
}
}
pub(crate) fn quota_data(
@@ -111,11 +172,13 @@ fn push_monthly_quota_window(
window_seconds,
label_display: "月限额".to_string(),
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent),
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
token_usage_display: "-".to_string(),
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
reset_at,
});
true
}
@@ -198,11 +261,13 @@ fn push_quota_window(
window_seconds,
label_display,
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_percent),
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
token_usage_display: "-".to_string(),
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
window_used_display: quota_window_used_display(reset_at, window_seconds),
reset_at,
});
}
+186 -58
View File
@@ -10,7 +10,7 @@ use axum::{
response::{IntoResponse, Response},
Json,
};
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
@@ -25,7 +25,10 @@ use crate::{
},
admin::update_payload::{parse_editable_update, EditableUpdateMessages},
admin::{
accounts::quota_view::{quota_data, AdminAccountQuotaData},
accounts::quota_view::{
quota_data, AdminAccountQuotaData, AdminAccountQuotaUsageWindow,
AdminAccountQuotaWindowTokenUsage,
},
accounts::service::{
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata,
AdminAccountRefreshOutcome, AdminAccountRefreshResult, AdminAccountUpdate,
@@ -38,7 +41,10 @@ use crate::{
format_cost, format_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64,
},
json::{clamp_limit, clamp_page, total_pages, Page},
time::{china_datetime, china_quarter_hour_start, china_relative_time, china_rfc3339},
time::{
china_datetime, china_quarter_hour_start, china_relative_time, china_rfc3339,
parse_rfc3339_utc,
},
},
runtime::state::AppState,
upstream::accounts::model::AccountStatus,
@@ -163,7 +169,7 @@ impl AdminAccountData {
label: a.label,
plan_type: a.plan_type,
has_refresh_token: a.has_refresh_token,
status: account_status_str(a.status).to_string(),
status: a.status.as_str().to_string(),
access_token_expires_at,
access_token_expires_at_display,
added_at: china_rfc3339(&a.added_at),
@@ -297,6 +303,12 @@ struct AccountModelUsageRecord {
last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone)]
struct AccountQuotaUsageWindow {
account_id: String,
window: AdminAccountQuotaUsageWindow,
}
impl AccountListStats {
fn data_for(&self, account: AdminAccountMetadata) -> AdminAccountData {
let account_id = account.id.clone();
@@ -396,7 +408,7 @@ impl From<AdminAccountRefreshResult> for AccountRefreshData {
Self {
id: result.id,
email: result.email,
previous_status: account_status_str(result.previous_status).to_string(),
previous_status: result.previous_status.as_str().to_string(),
result: account_refresh_outcome_str(result.outcome),
error: result.error,
duration_ms: result.duration_ms,
@@ -472,8 +484,8 @@ pub(crate) async fn accounts(
) -> Result<impl IntoResponse, AdminError> {
let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50));
let use_numbered_page = params.page.is_some() || params.page_size.is_some();
let stats = account_list_stats(&state).await;
let summary = account_summary_data(&state, &stats).await;
let quota_by_account = quota_snapshots_by_account(&state).await;
let summary = account_summary_data(&state, &quota_by_account).await;
if use_numbered_page {
return match state
@@ -483,6 +495,7 @@ pub(crate) async fn accounts(
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await;
let page = crate::infra::json::NumberedPage {
items: page
.items
@@ -509,6 +522,7 @@ pub(crate) async fn accounts(
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await;
let page = Page {
items: page
.items
@@ -640,7 +654,8 @@ pub(crate) async fn account_quota(
let quota = data.get("quota").cloned().unwrap_or(Value::Null);
let raw = data.get("raw").cloned().unwrap_or(Value::Null);
let quota_json = quota.to_string();
let quota_data = quota_data(&quota_json, Some(Utc::now()));
let mut quota_data = quota_data(&quota_json, Some(Utc::now()));
apply_account_quota_window_token_usage(&state, &account_id, &mut quota_data).await;
let plan_type = quota
.get("plan_type")
.and_then(Value::as_str)
@@ -658,7 +673,7 @@ pub(crate) async fn account_quota(
Err(AdminAccountError::NotFound) => Err(account_not_found()),
Err(AdminAccountError::Inactive(status)) => Err(AdminError::conflict(format!(
"Account is {}, cannot query quota",
account_status_str(status)
status.as_str()
))),
Err(AdminAccountError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!(
"Failed to fetch quota from Codex API: {msg}"
@@ -841,15 +856,23 @@ pub(crate) async fn update_account(
}
}
async fn account_list_stats(state: &AppState) -> AccountListStats {
let usage_records = list_all_usage_records(state).await;
let quota_snapshots = state
async fn account_list_stats(
state: &AppState,
accounts: &[AdminAccountMetadata],
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
) -> AccountListStats {
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let usage_records = state
.services
.background_tasks
.accounts
.list_quota_snapshots()
.usage
.list_by_account_ids(&account_ids)
.await
.unwrap_or_default();
let mut quota_by_account = quota_by_account.clone();
apply_quota_window_token_usage(state, accounts, &mut quota_by_account).await;
let models_by_account = list_current_window_model_usage(state, &usage_records).await;
AccountListStats {
@@ -857,22 +880,141 @@ async fn account_list_stats(state: &AppState) -> AccountListStats {
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect(),
quota_by_account: quota_snapshots
.into_iter()
.map(|snapshot| {
(
snapshot.account_id,
quota_data(&snapshot.quota_json, snapshot.quota_fetched_at),
)
})
.collect(),
quota_by_account,
models_by_account,
}
}
async fn apply_account_quota_window_token_usage(
state: &AppState,
account_id: &str,
quota: &mut AdminAccountQuotaData,
) {
let windows = quota
.usage_windows()
.into_iter()
.map(|window| AccountQuotaUsageWindow {
account_id: account_id.to_string(),
window,
})
.collect::<Vec<_>>();
let usage_by_account = quota_window_token_usage_by_account(state, &windows).await;
let empty_usage = HashMap::new();
quota.apply_token_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage));
}
async fn apply_quota_window_token_usage(
state: &AppState,
accounts: &[AdminAccountMetadata],
quota_by_account: &mut HashMap<String, AdminAccountQuotaData>,
) {
let windows = accounts
.iter()
.filter_map(|account| {
quota_by_account
.get(&account.id)
.map(|quota| (account.id.as_str(), quota))
})
.flat_map(|(account_id, quota)| {
quota
.usage_windows()
.into_iter()
.map(move |window| AccountQuotaUsageWindow {
account_id: account_id.to_string(),
window,
})
})
.collect::<Vec<_>>();
let usage_by_account = quota_window_token_usage_by_account(state, &windows).await;
for account in accounts {
if let Some(quota) = quota_by_account.get_mut(&account.id) {
let empty_usage = HashMap::new();
quota.apply_token_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage));
}
}
}
async fn quota_window_token_usage_by_account(
state: &AppState,
windows: &[AccountQuotaUsageWindow],
) -> HashMap<String, HashMap<String, AdminAccountQuotaWindowTokenUsage>> {
let Some(min_start) = windows.iter().map(|window| window.window.start).min() else {
return HashMap::new();
};
let Some(max_end) = windows.iter().map(|window| window.window.end).max() else {
return HashMap::new();
};
let account_ids = windows
.iter()
.map(|window| window.account_id.as_str())
.collect::<HashSet<_>>();
let mut windows_by_account = HashMap::<&str, Vec<&AdminAccountQuotaUsageWindow>>::new();
for window in windows {
windows_by_account
.entry(window.account_id.as_str())
.or_default()
.push(&window.window);
}
let mut builder = QueryBuilder::<Sqlite>::new(
"select
account_id,
bucket_start,
coalesce(sum(input_tokens), 0) as input_tokens,
coalesce(sum(output_tokens), 0) as output_tokens
from usage_time_buckets
where account_id in (",
);
let mut separated = builder.separated(", ");
for account_id in &account_ids {
separated.push_bind(*account_id);
}
separated.push_unseparated(")");
builder.push(" and bucket_start >= ");
builder.push_bind(min_start.to_rfc3339());
builder.push(" and bucket_start <= ");
builder.push_bind(max_end.to_rfc3339());
builder.push(" group by account_id, bucket_start");
let rows = builder
.build()
.fetch_all(state.services.background_tasks.accounts.pool())
.await
.unwrap_or_default();
let mut usage_by_account =
HashMap::<String, HashMap<String, AdminAccountQuotaWindowTokenUsage>>::new();
for row in rows {
let account_id: String = row.get("account_id");
let Ok(bucket_start) = parse_rfc3339_utc(row.get::<&str, _>("bucket_start")) else {
continue;
};
let total_tokens = nonnegative_i64_to_u64(row.get::<i64, _>("input_tokens"))
.saturating_add(nonnegative_i64_to_u64(row.get::<i64, _>("output_tokens")));
let Some(account_windows) = windows_by_account.get(account_id.as_str()) else {
continue;
};
for window in account_windows
.iter()
.filter(|window| bucket_start >= window.start && bucket_start <= window.end)
{
usage_by_account
.entry(account_id.clone())
.or_default()
.entry(window.key.clone())
.or_default()
.add_tokens(total_tokens);
}
}
usage_by_account
}
async fn account_summary_data(
state: &AppState,
stats: &AccountListStats,
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
) -> AdminAccountSummaryData {
let accounts = list_all_account_metadata(state).await;
let total = accounts.len() as u64;
@@ -883,8 +1025,7 @@ async fn account_summary_data(
let high_usage = accounts
.iter()
.filter(|account| {
stats
.quota_by_account
quota_by_account
.get(&account.id)
.is_some_and(account_quota_has_high_usage)
})
@@ -902,6 +1043,24 @@ async fn account_summary_data(
}
}
async fn quota_snapshots_by_account(state: &AppState) -> HashMap<String, AdminAccountQuotaData> {
state
.services
.background_tasks
.accounts
.list_quota_snapshots()
.await
.unwrap_or_default()
.into_iter()
.map(|snapshot| {
(
snapshot.account_id,
quota_data(&snapshot.quota_json, snapshot.quota_fetched_at),
)
})
.collect()
}
async fn list_all_account_metadata(state: &AppState) -> Vec<AdminAccountMetadata> {
let mut page = 1;
let mut accounts = Vec::new();
@@ -934,26 +1093,6 @@ fn account_summary_needs_attention(status: AccountStatus) -> bool {
)
}
async fn list_all_usage_records(state: &AppState) -> Vec<AdminUsageRecord> {
let mut cursor = None;
let mut records = Vec::new();
loop {
let Ok(page) = state
.services
.usage
.list(cursor, ACCOUNT_STATS_PAGE_LIMIT)
.await
else {
return Vec::new();
};
records.extend(page.items);
let Some(next_cursor) = page.next_cursor else {
return records;
};
cursor = Some(next_cursor);
}
}
async fn list_current_window_model_usage(
state: &AppState,
usage_records: &[AdminUsageRecord],
@@ -1202,14 +1341,3 @@ fn account_export_ids(value: Option<&str>) -> Vec<String> {
.map(ToString::to_string)
.collect()
}
fn account_status_str(status: crate::upstream::accounts::model::AccountStatus) -> &'static str {
match status {
crate::upstream::accounts::model::AccountStatus::Active => "active",
crate::upstream::accounts::model::AccountStatus::Expired => "expired",
crate::upstream::accounts::model::AccountStatus::QuotaExhausted => "quota_exhausted",
crate::upstream::accounts::model::AccountStatus::Refreshing => "refreshing",
crate::upstream::accounts::model::AccountStatus::Disabled => "disabled",
crate::upstream::accounts::model::AccountStatus::Banned => "banned",
}
}
@@ -113,12 +113,6 @@ pub enum AdminAccountError {
Inactive(AcctStatus),
}
#[derive(Debug, Clone)]
pub struct UpdatedAccountStatus {
pub id: String,
pub status: AcctStatus,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdminAccountRefreshOutcome {
Alive,
@@ -233,6 +227,14 @@ pub(super) struct ManualCreateTokens {
pub(super) access_token: String,
pub(super) refresh_token_for_new: Option<String>,
pub(super) refresh_token_for_existing: Option<String>,
pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims,
}
#[derive(Debug, Clone)]
pub(super) struct RefreshedAccountTokens {
pub(super) access_token: String,
pub(super) refresh_token: Option<String>,
pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims,
}
#[derive(Debug, Clone)]
@@ -4,10 +4,7 @@ use super::{
contracts::{ExportedAccount, ExportedAccounts},
AdminAccountError, AdminAccountService,
};
use crate::{
infra::time::china_rfc3339_str,
upstream::accounts::{model::AccountStatus, store::StoredAccount},
};
use crate::{infra::time::china_rfc3339_str, upstream::accounts::store::StoredAccount};
impl AdminAccountService {
pub async fn export(&self, ids: Vec<String>) -> Result<ExportedAccounts, AdminAccountError> {
@@ -58,20 +55,9 @@ impl From<StoredAccount> for ExportedAccount {
access_token_expires_at: account
.access_token_expires_at
.map(|value| value.to_rfc3339()),
status: export_status_str(account.status).to_string(),
status: account.status.as_str().to_string(),
added_at: china_rfc3339_str(&account.added_at),
updated_at: china_rfc3339_str(&account.updated_at),
}
}
}
fn export_status_str(status: AccountStatus) -> &'static str {
match status {
AccountStatus::Active => "active",
AccountStatus::Expired => "expired",
AccountStatus::QuotaExhausted => "quota_exhausted",
AccountStatus::Refreshing => "refreshing",
AccountStatus::Disabled => "disabled",
AccountStatus::Banned => "banned",
}
}
+11 -38
View File
@@ -281,8 +281,7 @@ impl AdminAccountService {
token: Option<String>,
refresh_token: Option<String>,
) -> Result<Option<ResolvedImportTokens>, AdminAccountError> {
let mut refresh_token =
crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
let refresh_token = crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
let Some(access_token) = crate::upstream::accounts::importing::normalize_nonempty(
token
.as_deref()
@@ -292,26 +291,12 @@ impl AdminAccountService {
return Ok(None);
};
let refreshed = self
.token_refresher
.refresh(&existing_refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(
&refreshed.access_token,
),
))
.ok_or(AdminAccountError::TokenRequired)?;
refresh_token = refreshed.refresh_token;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
.refresh_tokens_from_refresh_token(&existing_refresh_token)
.await?;
return Ok(Some(ResolvedImportTokens {
access_token,
refresh_token,
claims: Some(claims),
access_token: refreshed.access_token,
refresh_token: refreshed.refresh_token,
claims: Some(refreshed.claims),
}));
};
@@ -334,24 +319,12 @@ impl AdminAccountService {
}));
};
let refreshed = self
.token_refresher
.refresh(&existing_refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(&refreshed.access_token),
))
.ok_or(AdminAccountError::TokenRequired)?;
refresh_token = refreshed.refresh_token;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
.refresh_tokens_from_refresh_token(&existing_refresh_token)
.await?;
Ok(Some(ResolvedImportTokens {
access_token,
refresh_token,
claims: Some(claims),
access_token: refreshed.access_token,
refresh_token: refreshed.refresh_token,
claims: Some(refreshed.claims),
}))
}
async fn import_target_account_id(
+31 -86
View File
@@ -20,7 +20,7 @@ use super::{
contracts::{
parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountHealthCheck,
AdminAccountMetadata, AdminAccountRefreshOutcome, AdminAccountRefreshResult,
AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens, UpdatedAccountStatus,
AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens,
},
AdminAccountService,
};
@@ -92,37 +92,32 @@ impl AdminAccountService {
.as_deref()
.map(crate::upstream::accounts::importing::normalize_bearer_token),
) {
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
ManualCreateTokens {
access_token,
refresh_token_for_new: provided_refresh_token.clone(),
refresh_token_for_existing: provided_refresh_token,
claims,
}
} else if let Some(refresh_token) = provided_refresh_token {
let token_pair = self
.token_refresher
.refresh(&refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(
&token_pair.access_token,
),
))
.ok_or(AdminAccountError::TokenRequired)?;
let refreshed = self
.refresh_tokens_from_refresh_token(&refresh_token)
.await?;
ManualCreateTokens {
access_token,
refresh_token_for_new: token_pair.refresh_token.clone(),
refresh_token_for_existing: token_pair.refresh_token,
access_token: refreshed.access_token,
refresh_token_for_new: refreshed.refresh_token.clone(),
refresh_token_for_existing: refreshed.refresh_token,
claims: refreshed.claims,
}
} else {
return Err(AdminAccountError::TokenRequired);
};
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&tokens.access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
let claims = tokens.claims;
let existing = if let Some(account_id) = claims.account_id.as_deref() {
self.store
.find_by_chatgpt_identity(account_id, claims.user_id.as_deref())
@@ -195,45 +190,6 @@ impl AdminAccountService {
.map(stored_to_admin_metadata)
.ok_or(AdminAccountError::NotFound)
}
pub async fn update_label(
&self,
account_id: &str,
label: Option<String>,
) -> Result<bool, AdminAccountError> {
if label.as_ref().is_some_and(|l| l.chars().count() > 64) {
return Err(AdminAccountError::LabelTooLong);
}
let updated = self
.store
.set_label(account_id, label)
.await
.map_err(|_| AdminAccountError::UpdateLabel)?;
if updated {
self.sync_account_pool_best_effort(account_id, "account label update")
.await;
}
Ok(updated)
}
pub async fn update_status(
&self,
account_id: &str,
status: &str,
) -> Result<Option<UpdatedAccountStatus>, AdminAccountError> {
let status = parse_account_status(status)?;
match self.store.set_status(account_id, status).await {
Ok(true) => {
self.sync_account_pool_best_effort(account_id, "account status update")
.await;
self.evict_account_websocket_pool(account_id).await;
Ok(Some(UpdatedAccountStatus {
id: account_id.to_string(),
status,
}))
}
Ok(false) => Ok(None),
Err(_) => Err(AdminAccountError::UpdateStatus),
}
}
pub async fn update_account(
&self,
account_id: &str,
@@ -336,39 +292,28 @@ impl AdminAccountService {
}
match self
.token_refresher
.refresh(refresh_token.expose_secret())
.refresh_tokens_from_refresh_token(refresh_token.expose_secret())
.await
{
Ok(tokens) => {
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(
&tokens.access_token,
),
))
.ok_or(AdminAccountError::TokenRequired)?;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
Ok(refreshed) => {
let updated = self
.store
.update_from_claims(
account_id,
crate::upstream::accounts::store::AccountClaimsUpdate {
email: claims.email,
account_id: claims.account_id.or(account.account_id),
user_id: claims.user_id,
plan_type: claims.plan_type,
access_token: SecretString::new(access_token.into()),
refresh_token: tokens
email: refreshed.claims.email,
account_id: refreshed.claims.account_id.or(account.account_id),
user_id: refreshed.claims.user_id,
plan_type: refreshed.claims.plan_type,
access_token: SecretString::new(refreshed.access_token.into()),
refresh_token: refreshed
.refresh_token
.map(|token| SecretString::new(token.into())),
access_token_expires_at: Some(claims.expires_at),
next_refresh_at: Some(
self.next_refresh_at_for_expires_at(account_id, claims.expires_at),
),
access_token_expires_at: Some(refreshed.claims.expires_at),
next_refresh_at: Some(self.next_refresh_at_for_expires_at(
account_id,
refreshed.claims.expires_at,
)),
status: crate::upstream::accounts::model::AccountStatus::Active,
},
)
@@ -387,7 +332,7 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
})
}
Err(failure) => {
Err(AdminAccountError::RefreshTokenExchange(failure)) => {
if manual_refresh_failure_is_permanent(&failure) {
let status = crate::upstream::accounts::model::AccountStatus::Expired;
let updated = self
@@ -422,6 +367,7 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
})
}
Err(error) => Err(error),
}
}
@@ -494,8 +440,7 @@ fn manual_refresh_failure_is_permanent(
match failure {
crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant
| crate::upstream::accounts::token_refresh::RefreshFailure::Banned => true,
crate::upstream::accounts::token_refresh::RefreshFailure::QuotaExhausted
| crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport
crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport
| crate::upstream::accounts::token_refresh::RefreshFailure::Transport => false,
}
}
+28 -2
View File
@@ -11,7 +11,7 @@ mod testing;
use std::sync::Arc as StdArc;
use chrono::DateTime;
use chrono::{DateTime, Utc};
use crate::{
upstream::accounts::{
@@ -27,7 +27,7 @@ use crate::{
pub use contracts::{
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, AdminAccountRefreshOutcome,
AdminAccountRefreshResult, AdminAccountUpdate, BatchDeleteAccounts, ExportedAccounts,
ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, UpdatedAccountStatus,
ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput,
};
#[derive(Clone)]
@@ -82,6 +82,32 @@ impl AdminAccountService {
)
}
async fn refresh_tokens_from_refresh_token(
&self,
refresh_token: &str,
) -> Result<contracts::RefreshedAccountTokens, AdminAccountError> {
let token_pair = self
.token_refresher
.refresh(refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(&token_pair.access_token),
))
.ok_or(AdminAccountError::TokenRequired)?;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
Ok(contracts::RefreshedAccountTokens {
access_token,
refresh_token: token_pair.refresh_token,
claims,
})
}
pub(crate) async fn sync_account_pool(
&self,
account_id: &str,
+6 -33
View File
@@ -220,11 +220,10 @@ fn stored_account_to_model_refresh_account(stored: StoredAccount) -> Account {
}
fn test_responses_request(model: String) -> CodexResponsesRequest {
CodexResponsesRequest {
let mut request = CodexResponsesRequest::new_http_sse(
model,
instructions: "You are checking whether this Codex account can answer. Reply with ok."
.to_string(),
input: vec![json!({
"You are checking whether this Codex account can answer. Reply with ok.",
vec![json!({
"role": "user",
"content": [
{
@@ -233,35 +232,9 @@ fn test_responses_request(model: String) -> CodexResponsesRequest {
}
]
})],
stream: true,
store: false,
reasoning: None,
tools: None,
tool_choice: None,
parallel_tool_calls: None,
text: None,
generate: None,
tuple_schema: None,
service_tier: None,
previous_response_id: None,
prompt_cache_key: None,
explicit_prompt_cache_key: false,
client_conversation_id: None,
variant_identity: None,
include: None,
client_metadata: None,
client_ip: None,
client_user_agent: None,
use_websocket: false,
force_http_sse: true,
turn_state: None,
turn_metadata: None,
beta_features: None,
version: None,
include_timing_metrics: None,
codex_window_id: None,
parent_thread_id: None,
}
);
request.force_http_sse = true;
request
}
async fn process_upstream_test_stream(
+11 -11
View File
@@ -4,7 +4,7 @@ use axum::{
extract::{FromRequestParts, State},
http::{header::SET_COOKIE, request::Parts, HeaderMap, HeaderValue, StatusCode},
response::{IntoResponse, Response},
Json,
Extension, Json,
};
use serde::{Deserialize, Serialize};
@@ -13,6 +13,7 @@ use crate::{
auth::service::AdminSessionError,
response::{AdminEnvelope, AdminError, AdminResponse},
},
http::middleware::request_id::ClientIp,
infra::time::china_rfc3339,
runtime::state::AppState,
};
@@ -57,10 +58,10 @@ struct AdminSessionStatusData {
/// `POST /api/admin/login`
pub(crate) async fn login(
State(state): State<AppState>,
headers: HeaderMap,
client_ip: Option<Extension<ClientIp>>,
Json(payload): Json<AdminLoginRequest>,
) -> Result<Response, AdminError> {
let source = admin_login_source(&headers);
let source = admin_login_source(optional_client_ip(client_ip));
let login_result = state
.services
.admin_sessions
@@ -191,12 +192,11 @@ fn admin_api_key_header(headers: &HeaderMap) -> Option<String> {
(!value.is_empty()).then(|| value.to_string())
}
fn admin_login_source(headers: &HeaderMap) -> String {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("unknown")
.to_string()
fn optional_client_ip(client_ip: Option<Extension<ClientIp>>) -> Option<String> {
let Extension(client_ip) = client_ip?;
Some(client_ip.as_str().to_string())
}
fn admin_login_source(fallback_client_ip: Option<String>) -> String {
fallback_client_ip.unwrap_or_else(|| "unknown".to_string())
}
@@ -37,6 +37,18 @@ impl AdminUsageService {
})
}
/// 按账号 ID 批量读取账号用量。
pub async fn list_by_account_ids(
&self,
account_ids: &[String],
) -> Result<Vec<AdminUsageRecord>, AdminUsageError> {
self.store
.list_usage_by_account_ids(account_ids)
.await
.map(|records| records.into_iter().map(AdminUsageRecord::from).collect())
.map_err(|_| AdminUsageError::List)
}
/// 汇总账号用量。
pub async fn summary(&self) -> Result<AdminUsageSummary, AdminUsageError> {
self.store
@@ -245,6 +245,27 @@ impl SqliteUsageStore {
Ok(to_page(&rows, limit))
}
/// 按账号 ID 批量读取账号用量。
pub async fn list_usage_by_account_ids(
&self,
account_ids: &[String],
) -> SqliteUsageStoreResult<Vec<UsageListRecord>> {
if account_ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(LIST_USAGE_SELECT_SQL);
builder.push("\nwhere au.account_id in (");
let mut separated = builder.separated(", ");
for account_id in account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await?;
rows.iter().map(usage_list_from_row).collect()
}
/// 汇总账号用量。
pub async fn usage_summary(&self) -> SqliteUsageStoreResult<UsageSummary> {
let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?;
+33 -19
View File
@@ -219,10 +219,10 @@ struct UsageWindow {
errors: u64,
first_token_latency_sum: u64,
first_token_latency_count: u64,
max_first_token_bucket_latency: u64,
min_first_token_bucket_latency: u64,
latency_sum: u64,
latency_count: u64,
max_latency: u64,
min_latency: u64,
}
impl UsageWindow {
@@ -243,6 +243,14 @@ impl UsageWindow {
.checked_div(self.first_token_latency_count)
}
fn max_first_token_bucket_latency(self) -> Option<u64> {
(self.max_first_token_bucket_latency > 0).then_some(self.max_first_token_bucket_latency)
}
fn min_first_token_bucket_latency(self) -> Option<u64> {
(self.min_first_token_bucket_latency > 0).then_some(self.min_first_token_bucket_latency)
}
fn avg_latency(self) -> Option<u64> {
self.latency_sum.checked_div(self.latency_count)
}
@@ -252,6 +260,7 @@ impl UsageWindow {
pub(crate) async fn dashboard_summary(
State(state): State<AppState>,
_auth: AdminAuth,
Query(query): Query<DashboardTrendQuery>,
) -> Result<impl IntoResponse, AdminError> {
let accounts = state
.services
@@ -288,7 +297,7 @@ pub(crate) async fn dashboard_summary(
let now = Utc::now();
let time_buckets = dashboard_time_buckets(&state, now).await;
let pool_summary = account_pool_summary(&accounts);
let trend = dashboard_trend_data(&time_buckets, DashboardTrendKind::Usage);
let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default());
let quota_used_by_account = account_quota_used_percent_by_id(&state, &usage_records).await;
let settings = state.services.settings.current();
let recent_usage_records = recent_events.into_iter().take(10).collect::<Vec<_>>();
@@ -483,7 +492,9 @@ fn dashboard_trend_data_at(
let points = buckets
.iter()
.map(|(bucket_start, bucket)| {
let latency = bucket.avg_latency().unwrap_or(0);
let latency = bucket.avg_first_token_latency().unwrap_or(0);
let max_latency = bucket.max_first_token_bucket_latency().unwrap_or(0);
let min_latency = bucket.min_first_token_bucket_latency().unwrap_or(0);
let tokens = bucket.tokens();
let success_rate = if bucket.requests > 0 {
((bucket.requests - bucket.errors) as f64 / bucket.requests as f64 * 1000.0).round()
@@ -507,10 +518,10 @@ fn dashboard_trend_data_at(
errors_value: bucket.errors,
latency: format_optional_duration_ms(Some(latency)),
latency_value: latency,
max_latency: format_optional_duration_ms(Some(bucket.max_latency)),
max_latency_value: bucket.max_latency,
min_latency: format_optional_duration_ms(Some(bucket.min_latency)),
min_latency_value: bucket.min_latency,
max_latency: format_optional_duration_ms(Some(max_latency)),
max_latency_value: max_latency,
min_latency: format_optional_duration_ms(Some(min_latency)),
min_latency_value: min_latency,
success_rate: format_percent(success_rate),
success_rate_value: success_rate,
}
@@ -772,20 +783,23 @@ fn apply_bucket(window: &mut UsageWindow, record: &AdminUsageTimeBucketRecord) {
window.errors += nonnegative_i64_to_u64(record.error_count);
window.first_token_latency_sum += nonnegative_i64_to_u64(record.first_token_latency_sum);
window.first_token_latency_count += nonnegative_i64_to_u64(record.first_token_latency_count);
window.latency_sum += nonnegative_i64_to_u64(record.latency_sum);
window.latency_count += nonnegative_i64_to_u64(record.latency_count);
let max_latency = nonnegative_i64_to_u64(record.max_latency_ms);
if max_latency > 0 {
window.max_latency = window.max_latency.max(max_latency);
}
let min_latency = nonnegative_i64_to_u64(record.min_latency_ms);
if min_latency > 0 {
window.min_latency = if window.min_latency == 0 {
min_latency
let bucket_first_token_latency = first_token_bucket_latency(record);
if let Some(latency) = bucket_first_token_latency {
window.max_first_token_bucket_latency = window.max_first_token_bucket_latency.max(latency);
window.min_first_token_bucket_latency = if window.min_first_token_bucket_latency == 0 {
latency
} else {
window.min_latency.min(min_latency)
window.min_first_token_bucket_latency.min(latency)
};
}
window.latency_sum += nonnegative_i64_to_u64(record.latency_sum);
window.latency_count += nonnegative_i64_to_u64(record.latency_count);
}
fn first_token_bucket_latency(record: &AdminUsageTimeBucketRecord) -> Option<u64> {
let sum = nonnegative_i64_to_u64(record.first_token_latency_sum);
let count = nonnegative_i64_to_u64(record.first_token_latency_count);
sum.checked_div(count).filter(|latency| *latency > 0)
}
fn account_pool_summary(accounts: &[Account]) -> AccountPoolDiagnostics {
@@ -29,8 +29,8 @@ use crate::{
},
infra::{
format::{
format_compact_number, format_duration_ms, format_duration_ms_f64, format_multiplier,
format_number, format_percent, format_precise_cost, format_rate, format_token_price,
format_compact_number, format_cost, format_duration_ms, format_duration_ms_f64,
format_multiplier, format_number, format_percent, format_rate, format_token_price,
format_tokens,
},
json::{clamp_limit, clamp_page, NumberedPage, Page},
@@ -125,7 +125,6 @@ pub(crate) struct UsageRecordCostDetailsData {
output_cost: f64,
cache_read_cost: f64,
total_cost: f64,
billed_cost: f64,
original_cost: f64,
input_price_per_mtoken: f64,
output_price_per_mtoken: f64,
@@ -137,7 +136,6 @@ pub(crate) struct UsageRecordCostDetailsData {
output_cost_display: String,
cache_read_cost_display: String,
total_cost_display: String,
billed_cost_display: String,
original_cost_display: String,
input_price_display: String,
output_price_display: String,
@@ -332,19 +330,7 @@ pub(crate) async fn usage_records_token_trend(
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
let filter = filter_from_query(query)?;
match state.services.usage_records.token_trend(filter).await {
Ok(points) => Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
points
.into_iter()
.map(UsageRecordTrendPointData::from)
.collect::<Vec<_>>(),
),
)),
Err(error) => Err(log_error(&error)),
}
usage_records_trend_response(state, query).await
}
/// `GET /api/admin/usage/records/insights/latency-trend`
@@ -353,8 +339,15 @@ pub(crate) async fn usage_records_latency_trend(
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
usage_records_trend_response(state, query).await
}
async fn usage_records_trend_response(
state: AppState,
query: UsageRecordsQuery,
) -> Result<AdminResponse<AdminEnvelope<Vec<UsageRecordTrendPointData>>>, AdminError> {
let filter = filter_from_query(query)?;
match state.services.usage_records.latency_trend(filter).await {
match state.services.usage_records.trend(filter).await {
Ok(points) => Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
@@ -607,7 +600,6 @@ fn usage_cost_details(
output_cost: cost.output_cost,
cache_read_cost: cost.cache_read_cost,
total_cost: cost.total_cost,
billed_cost: cost.billed_cost,
original_cost: cost.original_cost,
input_price_per_mtoken: cost.input_price_per_mtoken,
output_price_per_mtoken: cost.output_price_per_mtoken,
@@ -615,12 +607,11 @@ fn usage_cost_details(
service_tier: cost.service_tier,
service_tier_display: cost.service_tier_display,
multiplier: cost.multiplier,
input_cost_display: format_precise_cost(cost.input_cost),
output_cost_display: format_precise_cost(cost.output_cost),
cache_read_cost_display: format_precise_cost(cost.cache_read_cost),
total_cost_display: format_precise_cost(cost.total_cost),
billed_cost_display: format_precise_cost(cost.billed_cost),
original_cost_display: format_precise_cost(cost.original_cost),
input_cost_display: format_cost(cost.input_cost),
output_cost_display: format_cost(cost.output_cost),
cache_read_cost_display: format_cost(cost.cache_read_cost),
total_cost_display: format_cost(cost.total_cost),
original_cost_display: format_cost(cost.original_cost),
input_price_display: format_token_price(cost.input_price_per_mtoken),
output_price_display: format_token_price(cost.output_price_per_mtoken),
cache_read_price_display: format_token_price(cost.cache_read_price_per_mtoken),
@@ -682,9 +673,9 @@ impl UsageRecordBreakdownData {
total_tokens_value: item.total_tokens,
total_tokens_total: format_tokens(total_tokens_total),
total_tokens_total_value: total_tokens_total,
cost: format_precise_cost(item.cost),
actual_cost: format_precise_cost(item.actual_cost),
account_cost: format_precise_cost(item.account_cost),
cost: format_cost(item.cost),
actual_cost: format_cost(item.actual_cost),
account_cost: format_cost(item.account_cost),
average_latency_ms: format_duration_ms_f64(item.average_latency_ms),
}
}
@@ -726,8 +717,8 @@ impl From<UsageRecordTrendPoint> for UsageRecordTrendPointData {
total_tokens_value: point.total_tokens,
cache_hit_rate: format_percent(cache_hit_rate_value),
cache_hit_rate_value,
cost: format_precise_cost(point.cost),
actual_cost: format_precise_cost(point.actual_cost),
cost: format_cost(point.cost),
actual_cost: format_cost(point.actual_cost),
average_latency_ms: format_duration_ms_f64(point.average_latency_ms),
average_latency_ms_value: point.average_latency_ms,
}
@@ -92,7 +92,6 @@ pub(crate) struct AdminUsageRecordCostDetails {
pub output_cost: f64,
pub cache_read_cost: f64,
pub total_cost: f64,
pub billed_cost: f64,
pub original_cost: f64,
pub input_price_per_mtoken: f64,
pub output_price_per_mtoken: f64,
@@ -237,19 +236,8 @@ impl AdminUsageRecordService {
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合 Token 趋势。
pub async fn token_trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
self.store
.trend(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合延迟趋势。
pub async fn latency_trend(
/// 聚合使用记录趋势。
pub async fn trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
@@ -343,7 +331,6 @@ pub(crate) fn usage_record_cost_details(
output_cost: breakdown.output_cost,
cache_read_cost: breakdown.cache_read_cost,
total_cost: breakdown.total_cost,
billed_cost: breakdown.total_cost,
original_cost,
input_price_per_mtoken: breakdown.input_price_per_mtoken,
output_price_per_mtoken: breakdown.output_price_per_mtoken,
+97 -42
View File
@@ -1,6 +1,7 @@
//! 请求上下文中间件。
use std::net::{IpAddr, SocketAddr};
use std::sync::LazyLock;
use axum::{
extract::{connect_info::ConnectInfo, Request, State},
@@ -17,11 +18,25 @@ const MAX_REQUEST_ID_LEN: usize = 128;
const CF_CONNECTING_IP_HEADER: &str = "cf-connecting-ip";
const REAL_IP_HEADER: &str = "x-real-ip";
const X_FORWARDED_FOR_HEADER: &str = "x-forwarded-for";
const UNTRUSTED_CLIENT_IP_HEADERS: &[&str] = &[
CF_CONNECTING_IP_HEADER,
REAL_IP_HEADER,
X_FORWARDED_FOR_HEADER,
];
/// 私有 / 回环 / 唯一本地地址网段,用于在自动模式下从 `X-Forwarded-For` 中跳过内网跳板 IP。
static PRIVATE_NETWORKS: LazyLock<Vec<IpNet>> = LazyLock::new(|| {
[
"10.0.0.0/8",
"172.16.0.0/12",
"192.168.0.0/16",
"127.0.0.0/8",
"::1/128",
"fc00::/7",
]
.iter()
.filter_map(|entry| entry.parse().ok())
.collect()
});
fn is_private_ip(ip: IpAddr) -> bool {
PRIVATE_NETWORKS.iter().any(|network| network.contains(&ip))
}
/// 可信反向代理配置。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -42,6 +57,11 @@ impl TrustedProxyConfig {
fn is_trusted(&self, ip: IpAddr) -> bool {
self.networks.iter().any(|network| network.contains(&ip))
}
/// 是否未配置任何可信代理(决定是否启用自动真实 IP 探测)。
fn is_empty(&self) -> bool {
self.networks.is_empty()
}
}
/// 可信代理配置错误。
@@ -77,6 +97,17 @@ impl RequestId {
}
}
/// 由连接来源和可信代理链解析出的客户端 IP。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientIp(String);
impl ClientIp {
/// 返回 IP 字符串。
pub fn as_str(&self) -> &str {
&self.0
}
}
fn valid_request_id(value: &str) -> bool {
!value.is_empty()
&& value.len() <= MAX_REQUEST_ID_LEN
@@ -117,10 +148,10 @@ async fn run_with_request_id(mut request: Request, next: Next) -> Response {
}
fn attach_request_context(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
attach_real_ip_from_connection(trusted_proxies, request);
attach_client_ip_from_connection(trusted_proxies, request);
}
fn attach_real_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
fn attach_client_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
let peer_ip = request
.extensions()
.get::<ConnectInfo<SocketAddr>>()
@@ -128,68 +159,92 @@ fn attach_real_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request:
let client_ip =
peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers(), trusted_proxies));
for header in UNTRUSTED_CLIENT_IP_HEADERS {
request.headers_mut().remove(*header);
}
let Some(client_ip) = client_ip else {
return;
};
if let Ok(value) = HeaderValue::from_str(&client_ip.to_string()) {
request.headers_mut().insert(REAL_IP_HEADER, value);
}
request.extensions_mut().insert(ClientIp(client_ip));
}
fn client_ip_from_peer(
peer_ip: IpAddr,
headers: &HeaderMap,
trusted_proxies: &TrustedProxyConfig,
) -> IpAddr {
) -> String {
if trusted_proxies.is_empty() {
// 未配置可信代理:自动模式,按 CF-Connecting-IP → X-Real-IP →
// X-Forwarded-For(首个公网 IP)→ peer IP 的优先级解析真实客户端 IP。
return auto_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string());
}
if trusted_proxies.is_trusted(peer_ip) {
trusted_forwarded_client_ip(headers).unwrap_or(peer_ip)
// 已配置可信代理且直连来源可信:只采信转发头。
trusted_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string())
} else {
peer_ip
// 已配置可信代理但直连来源不可信:忽略转发头,使用 peer IP。
peer_ip.to_string()
}
}
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<IpAddr> {
header_ip(headers, CF_CONNECTING_IP_HEADER)
.or_else(|| first_forwarded_for_ip(headers))
.or_else(|| header_ip(headers, REAL_IP_HEADER))
/// 自动模式下解析转发头中的真实客户端 IP。
///
/// 未配置可信代理时使用,对齐常见反代(Cloudflare / Nginx)默认行为:
/// `CF-Connecting-IP` → `X-Real-IP` → `X-Forwarded-For` 首个公网 IP。
fn auto_forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
if let Some(cf_ip) = header_string(headers, CF_CONNECTING_IP_HEADER) {
return Some(cf_ip);
}
if let Some(real_ip) = header_string(headers, REAL_IP_HEADER) {
return Some(real_ip);
}
forwarded_for_public_ip(headers).or_else(|| first_forwarded_for_value(headers))
}
fn header_ip(headers: &HeaderMap, name: &str) -> Option<IpAddr> {
parse_ip(headers.get(name)?.to_str().ok()?)
/// 返回 `X-Forwarded-For` 链中第一个公网 IP(跳过内网跳板)。
fn forwarded_for_public_ip(headers: &HeaderMap) -> Option<String> {
let value = header_string(headers, X_FORWARDED_FOR_HEADER)?;
value
.split(',')
.map(str::trim)
.filter(|entry| !entry.is_empty())
.find(|entry| match entry.parse::<IpAddr>() {
Ok(ip) => !is_private_ip(ip),
Err(_) => false,
})
.map(ToString::to_string)
}
fn first_forwarded_for_ip(headers: &HeaderMap) -> Option<IpAddr> {
let value = headers.get(X_FORWARDED_FOR_HEADER)?.to_str().ok()?;
parse_ip(value.split(',').next()?)
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
header_string(headers, CF_CONNECTING_IP_HEADER)
.or_else(|| first_forwarded_for_value(headers))
.or_else(|| header_string(headers, REAL_IP_HEADER))
}
fn parse_ip(value: &str) -> Option<IpAddr> {
value.trim().parse().ok()
fn header_string(headers: &HeaderMap, name: &str) -> Option<String> {
headers
.get(name)
.and_then(|value| value.to_str().ok())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn first_forwarded_for_value(headers: &HeaderMap) -> Option<String> {
let value = header_string(headers, X_FORWARDED_FOR_HEADER)?;
value
.split(',')
.next()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn parse_trusted_proxy_entry(entry: &str) -> Result<IpNet, TrustedProxyConfigError> {
let entry = entry.trim();
if entry.is_empty() {
return Err(invalid_trusted_proxy_entry(entry));
}
if let Ok(network) = entry.parse::<IpNet>() {
return Ok(network);
}
let ip = entry
entry
.parse::<IpAddr>()
.map_err(|_| invalid_trusted_proxy_entry(entry))?;
IpNet::new(
ip,
match ip {
IpAddr::V4(_) => 32,
IpAddr::V6(_) => 128,
},
)
.map_err(|_| invalid_trusted_proxy_entry(entry))
.map(IpNet::from)
.map_err(|_| invalid_trusted_proxy_entry(entry))
}
fn invalid_trusted_proxy_entry(entry: &str) -> TrustedProxyConfigError {
+11 -10
View File
@@ -63,6 +63,15 @@ pub fn format_percent(value: f64) -> String {
}
}
/// 展示百分比,并去掉无意义的尾随 0。
pub fn format_compact_percent(value: f64) -> String {
if value.is_finite() {
format!("{}%", trim_trailing_zeroes(&format!("{value:.1}")))
} else {
"—".to_string()
}
}
/// 将 0 到 1 的比例展示为百分比。
pub fn format_rate(value: Option<f64>) -> String {
value.map_or_else(|| "—".to_string(), |value| format_percent(value * 100.0))
@@ -86,15 +95,6 @@ pub fn format_cost(value: f64) -> String {
format!("${value:.2}")
}
/// 使用六位小数展示美元成本。
pub fn format_precise_cost(value: f64) -> String {
if value.is_finite() {
format!("${value:.6}")
} else {
"—".to_string()
}
}
/// 展示每百万 token 的美元单价。
pub fn format_token_price(value: f64) -> String {
if value.is_finite() {
@@ -160,7 +160,8 @@ fn format_compact_scaled(value: f64, unit: &str) -> String {
format!("{}{unit}", trim_trailing_zeroes(&rounded))
}
fn trim_trailing_zeroes(value: &str) -> String {
/// 去掉小数字符串中无意义的尾随 0 和小数点。
pub fn trim_trailing_zeroes(value: &str) -> String {
value
.trim_end_matches('0')
.trim_end_matches('.')
-671
View File
@@ -1,671 +0,0 @@
//! OpenAI Chat Completions 调度服务。
use std::sync::Arc;
use std::time::Instant;
use chrono::Utc;
use serde_json::Value;
use thiserror::Error;
use crate::{
admin::monitoring::{
usage_record_model::{ResponseUsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
},
proxy::dispatch::{
auth_recovery::trigger_refresh_after_auth_failure,
cloudflare::{
cloudflare_challenge_error_message, cloudflare_path_block_error_message,
is_cloudflare_challenge_upstream_error, is_cloudflare_path_block_upstream_error,
CloudflareRecovery,
},
errors::{
auth_failure_account_status, backend_transport_name, is_auth_upstream_error,
is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error,
is_rate_limit_upstream_error, rate_limit_cooldown_until, upstream_error_body,
upstream_error_http_status,
},
upstream::{
create_response_with_account, verify_acquired_quota_if_required,
QuotaVerificationContext, QuotaVerificationDecision,
QUOTA_VERIFY_LIMIT_REACHED_MESSAGE,
},
usage_events::{
reasoning_effort_from_request, record_dispatch_error_event, record_response_event,
DispatchErrorUsageRecord,
},
},
proxy::openai::chat::ChatStreamTranslationError,
upstream::accounts::{
model::AccountStatus,
pool::{AccountAcquireRequest, RuntimeAccountPoolService},
token_refresh::RuntimeTokenRefreshService,
},
upstream::{
models::{catalog::ModelCatalog, service::ModelService},
protocol::responses::{apply_response_model_options, CodexResponsesRequest},
token_client::OpenAiTokenClient,
transport::{
backend_transport_for_response_request, is_banned_upstream_error, CodexBackendClient,
CodexClientError,
},
},
};
/// OpenAI Chat Completions 调度服务。
#[derive(Clone)]
pub struct ChatDispatchService {
account_pool: Arc<RuntimeAccountPoolService>,
models: Arc<ModelService>,
codex: Arc<CodexBackendClient>,
usage_records: Arc<AdminUsageRecordService>,
token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
installation_id: Option<String>,
cloudflare: CloudflareRecovery,
}
impl ChatDispatchService {
pub(crate) fn new(
account_pool: Arc<RuntimeAccountPoolService>,
models: Arc<ModelService>,
codex: Arc<CodexBackendClient>,
usage_records: Arc<AdminUsageRecordService>,
token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
installation_id: Option<String>,
cloudflare: CloudflareRecovery,
) -> Self {
Self {
account_pool,
models,
codex,
usage_records,
token_refresh,
installation_id,
cloudflare,
}
}
/// 调度非流式 Chat Completions 请求到 Codex Responses 上游。
pub async fn complete(
&self,
request_id: &str,
mut request: CodexResponsesRequest,
requested_model: &str,
) -> Result<Value, ChatDispatchError> {
let started_at = Instant::now();
let catalog = self.models.catalog().await;
let parsed_model = catalog.parse_model_name(requested_model);
let display_model = ModelCatalog::build_display_model_name(&parsed_model);
apply_response_model_options(&mut request, &parsed_model);
let include_reasoning = request
.reasoning
.as_ref()
.and_then(|reasoning| reasoning.get("effort"))
.and_then(Value::as_str)
.is_some_and(|effort| !effort.trim().is_empty());
let tuple_schema = request.tuple_schema.clone();
let mut excluded_account_ids = Vec::new();
let mut rate_limited_count = 0usize;
let mut last_rate_limit_error = None;
let mut quota_exhausted_count = 0usize;
let mut last_quota_error = None;
let mut expired_count = 0usize;
let mut last_auth_error = None;
let mut disabled_count = 0usize;
let mut last_disabled_auth_error = None;
let mut banned_count = 0usize;
let mut last_banned_auth_error = None;
let mut last_banned_status_code: Option<u16> = None;
let mut cloudflare_challenge_count = 0usize;
let mut last_cloudflare_challenge_error = None;
let mut cloudflare_path_block_count = 0usize;
let mut last_cloudflare_path_block_error = None;
let mut model_unsupported_count = 0usize;
let mut last_model_unsupported_error = None;
let mut model_unsupported_retry_used = false;
let mut quota_verify_attempts = 0usize;
let mut last_failed_account_id = None;
macro_rules! return_dispatch_error {
($error:expr) => {{
let error = $error;
self.record_chat_dispatch_error(
request_id,
requested_model,
started_at,
last_failed_account_id.as_deref(),
Some(backend_transport_name(
backend_transport_for_response_request(&request),
)),
&error,
)
.await;
return Err(error);
}};
($error:expr, account_id: $account_id:expr, transport: $transport:expr) => {{
let error = $error;
self.record_chat_dispatch_error(
request_id,
requested_model,
started_at,
$account_id,
$transport,
&error,
)
.await;
return Err(error);
}};
}
let (account, response) = loop {
let acquire_request = AccountAcquireRequest::new(&request.model, Utc::now())
.with_exclude_account_ids(excluded_account_ids.iter().cloned());
let acquired = match self.account_pool.acquire_with(&acquire_request).await {
Some(acquired) => acquired,
None if quota_exhausted_count > 0 => {
return_dispatch_error!(ChatDispatchError::QuotaExhausted {
count: quota_exhausted_count,
upstream_error: last_quota_error.unwrap_or_default(),
});
}
None if rate_limited_count > 0 => {
return_dispatch_error!(ChatDispatchError::RateLimited {
count: rate_limited_count,
upstream_error: last_rate_limit_error.unwrap_or_default(),
});
}
None if expired_count > 0 => {
return_dispatch_error!(ChatDispatchError::Expired {
count: expired_count,
upstream_error: last_auth_error.unwrap_or_default(),
});
}
None if disabled_count > 0 => {
return_dispatch_error!(ChatDispatchError::Disabled {
count: disabled_count,
upstream_error: last_disabled_auth_error.unwrap_or_default(),
});
}
None if banned_count > 0 => {
return_dispatch_error!(ChatDispatchError::Banned {
count: banned_count,
upstream_error: last_banned_auth_error.unwrap_or_default(),
status_code: last_banned_status_code.unwrap_or(403),
});
}
None if cloudflare_challenge_count > 0 => {
return_dispatch_error!(ChatDispatchError::CloudflareChallenge {
count: cloudflare_challenge_count,
upstream_error: last_cloudflare_challenge_error.unwrap_or_default(),
});
}
None if cloudflare_path_block_count > 0 => {
return_dispatch_error!(ChatDispatchError::CloudflarePathBlocked {
count: cloudflare_path_block_count,
upstream_error: last_cloudflare_path_block_error.unwrap_or_default(),
});
}
None if model_unsupported_count > 0 => {
return_dispatch_error!(ChatDispatchError::ModelUnsupported {
count: model_unsupported_count,
upstream_error: last_model_unsupported_error.unwrap_or_default(),
});
}
None => return_dispatch_error!(ChatDispatchError::NoActiveAccount),
};
let acquired = match verify_acquired_quota_if_required(
QuotaVerificationContext {
account_pool: self.account_pool.as_ref(),
codex: self.codex.as_ref(),
cloudflare: &self.cloudflare,
installation_id: self.installation_id.as_deref(),
request_id,
excluded_account_ids: &mut excluded_account_ids,
verify_attempts: &mut quota_verify_attempts,
},
acquired,
)
.await
{
QuotaVerificationDecision::Ready(acquired) => *acquired,
QuotaVerificationDecision::RetryWithAnotherAccount => {
rate_limited_count += 1;
last_rate_limit_error = Some(QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string());
continue;
}
QuotaVerificationDecision::MaxAttemptsReached => {
return_dispatch_error!(ChatDispatchError::RateLimited {
count: rate_limited_count + 1,
upstream_error: QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string(),
});
}
};
self.account_pool.wait_for_request_interval(&acquired).await;
let account = acquired.account;
let release_account_id = account.id.clone();
let response_result = create_response_with_account(
&self.codex,
self.installation_id.as_deref(),
&self.cloudflare,
&request,
request_id,
&account,
started_at,
)
.await;
self.account_pool.release(&release_account_id).await;
match response_result {
Ok(response) => break (account, response),
Err(error) if is_rate_limit_upstream_error(&error) => {
rate_limited_count += 1;
last_rate_limit_error = Some(upstream_error_body(&error));
last_failed_account_id = Some(release_account_id.clone());
let cooldown_until = rate_limit_cooldown_until(&error, Utc::now());
self.account_pool
.mark_quota_limited_until(&release_account_id, cooldown_until)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_quota_exhausted_upstream_error(&error) => {
quota_exhausted_count += 1;
last_quota_error = Some(upstream_error_body(&error));
last_failed_account_id = Some(release_account_id.clone());
self.account_pool
.set_status(&release_account_id, AccountStatus::QuotaExhausted)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_auth_upstream_error(&error) => {
let upstream_error = upstream_error_body(&error);
let account_status = auth_failure_account_status(&error);
last_failed_account_id = Some(release_account_id.clone());
match account_status {
AccountStatus::Disabled => {
disabled_count += 1;
last_disabled_auth_error = Some(upstream_error);
}
AccountStatus::Banned => {
banned_count += 1;
last_banned_status_code = Some(upstream_error_http_status(&error));
last_banned_auth_error = Some(upstream_error);
}
_ => {
expired_count += 1;
last_auth_error = Some(upstream_error);
}
}
self.account_pool
.set_status(&release_account_id, account_status)
.await;
trigger_refresh_after_auth_failure(
&self.token_refresh,
&release_account_id,
account_status,
);
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_challenge_upstream_error(&error) => {
cloudflare_challenge_count += 1;
last_cloudflare_challenge_error =
Some(cloudflare_challenge_error_message().to_string());
last_failed_account_id = Some(release_account_id.clone());
self.cloudflare
.apply_challenge(self.account_pool.as_ref(), &release_account_id)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_cloudflare_path_block_upstream_error(&error) => {
cloudflare_path_block_count += 1;
last_cloudflare_path_block_error =
Some(cloudflare_path_block_error_message().to_string());
last_failed_account_id = Some(release_account_id.clone());
self.cloudflare
.apply_path_block(self.account_pool.as_ref(), &release_account_id)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_model_unsupported_upstream_error(&error) => {
let upstream_error = upstream_error_body(&error);
if model_unsupported_retry_used {
return_dispatch_error!(
ChatDispatchError::ModelUnsupported {
count: model_unsupported_count + 1,
upstream_error,
},
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(backend_transport_for_response_request(
&request
)))
);
}
model_unsupported_count += 1;
last_model_unsupported_error = Some(upstream_error);
last_failed_account_id = Some(release_account_id.clone());
model_unsupported_retry_used = true;
excluded_account_ids.push(release_account_id);
}
Err(error) if is_banned_upstream_error(&error) => {
banned_count += 1;
last_banned_status_code = Some(upstream_error_http_status(&error));
last_banned_auth_error = Some(upstream_error_body(&error));
last_failed_account_id = Some(release_account_id.clone());
self.account_pool
.set_status(&release_account_id, AccountStatus::Banned)
.await;
excluded_account_ids.push(release_account_id);
}
Err(error) => {
return_dispatch_error!(
ChatDispatchError::Upstream(error),
account_id: Some(&release_account_id),
transport: Some(backend_transport_name(backend_transport_for_response_request(
&request
)))
);
}
}
};
let account_id = account.id.clone();
let body = match crate::proxy::openai::chat::chat_completion_from_codex_sse(
&response.body,
&display_model,
include_reasoning,
tuple_schema.as_ref(),
) {
Ok(Some(body)) => body,
Ok(None) => {
return_dispatch_error!(
ChatDispatchError::EmptyUpstreamResponse,
account_id: Some(&account_id),
transport: Some(backend_transport_name(response.transport))
);
}
Err(error) => {
return_dispatch_error!(
ChatDispatchError::InvalidSse(error),
account_id: Some(&account_id),
transport: Some(backend_transport_name(response.transport))
);
}
};
let response_id = body.get("id").and_then(Value::as_str);
self.cloudflare.reset_account_recovery(&account_id).await;
self.account_pool
.sync_passive_rate_limit_headers(&account, &response.rate_limit_headers)
.await;
if let Some(ref usage) = response.usage {
self.account_pool
.record_token_usage(&account_id, &request.model, usage)
.await;
}
let mut metadata = serde_json::json!({
"responseId": response_id,
"stream": false,
"transport": backend_transport_name(response.transport),
"firstTokenMs": response.first_token_ms,
"usage": response.usage,
});
if let (Some(object), Some(decision)) =
(metadata.as_object_mut(), response.websocket_pool_decision)
{
object.insert("websocketPool".to_string(), decision.metadata_value());
}
record_response_event(ResponseUsageRecord {
usage_records: &self.usage_records,
request_id,
account_id: &account_id,
route: "/v1/chat/completions",
model: &display_model,
requested_model: Some(requested_model),
client_ip: request.client_ip.as_deref(),
client_user_agent: request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(&request),
service_tier: request.service_tier.as_deref(),
started_at,
status_code: 200,
level: UsageRecordLevel::Info,
message: "v1 chat completions completed",
metadata,
rate_limit_headers: &response.rate_limit_headers,
})
.await;
Ok(body)
}
async fn record_chat_dispatch_error(
&self,
request_id: &str,
requested_model: &str,
started_at: Instant,
account_id: Option<&str>,
transport: Option<&str>,
error: &ChatDispatchError,
) {
record_chat_dispatch_error_event(ChatDispatchErrorEventRecord {
usage_records: &self.usage_records,
request_id,
account_id,
route: "/v1/chat/completions",
model: requested_model,
started_at,
transport,
error,
})
.await;
}
}
/// Chat Completions 调度错误。
#[derive(Debug, Error)]
pub enum ChatDispatchError {
#[error("failed to list runtime accounts")]
AccountStore,
#[error("no active account is available")]
NoActiveAccount,
#[error("all accounts exhausted by quota")]
QuotaExhausted {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by rate limit")]
RateLimited {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by expired auth")]
Expired {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by disabled auth")]
Disabled {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by banned auth")]
Banned {
count: usize,
upstream_error: String,
status_code: u16,
},
#[error("all accounts exhausted by Cloudflare challenge")]
CloudflareChallenge {
count: usize,
upstream_error: String,
},
#[error("all accounts exhausted by Cloudflare path-block")]
CloudflarePathBlocked {
count: usize,
upstream_error: String,
},
#[error("all accounts do not support the requested model")]
ModelUnsupported {
count: usize,
upstream_error: String,
},
#[error("upstream request failed: {0}")]
Upstream(#[from] CodexClientError),
#[error("invalid upstream SSE response: {0}")]
InvalidSse(ChatStreamTranslationError),
#[error("upstream response did not include response.completed")]
EmptyUpstreamResponse,
}
impl ChatDispatchError {
pub fn http_status_code(&self) -> u16 {
match self {
Self::NoActiveAccount | Self::AccountStore => 503,
Self::QuotaExhausted { .. } => 429,
Self::RateLimited { .. } => 429,
Self::Expired { .. } | Self::Disabled { .. } => 401,
Self::Banned { status_code, .. } => *status_code,
Self::CloudflareChallenge { .. }
| Self::CloudflarePathBlocked { .. }
| Self::InvalidSse(_)
| Self::EmptyUpstreamResponse => 502,
Self::ModelUnsupported { .. } => 400,
Self::Upstream(error) => upstream_error_http_status(error),
}
}
}
// ====================================================================
// Event recording helpers
// ====================================================================
struct ChatDispatchErrorEventRecord<'a> {
usage_records: &'a AdminUsageRecordService,
request_id: &'a str,
account_id: Option<&'a str>,
route: &'a str,
model: &'a str,
started_at: Instant,
transport: Option<&'a str>,
error: &'a ChatDispatchError,
}
async fn record_chat_dispatch_error_event(record: ChatDispatchErrorEventRecord<'_>) {
let mut metadata = serde_json::json!({
"route": record.route,
"apiKind": "chat",
"stream": false,
"failed": true,
"errorKind": "dispatch",
"error": record.error.to_string(),
});
if let Some(object) = metadata.as_object_mut() {
enrich_chat_error_metadata(object, record.error);
if let Some(transport) = record.transport {
object.insert("transport".to_string(), serde_json::json!(transport));
}
}
record_dispatch_error_event(DispatchErrorUsageRecord {
usage_records: record.usage_records,
request_id: record.request_id,
account_id: record.account_id,
route: record.route,
model: record.model,
started_at: record.started_at,
status_code: i64::from(record.error.http_status_code()),
message: "chat dispatch failed",
metadata,
})
.await;
}
fn enrich_chat_error_metadata(
object: &mut serde_json::Map<String, Value>,
error: &ChatDispatchError,
) {
let (failure_class, exhausted_count, upstream_error, upstream_status) = match error {
ChatDispatchError::AccountStore => ("account_store", None, None, None),
ChatDispatchError::NoActiveAccount => ("no_available_accounts", None, None, None),
ChatDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
"quota_exhausted",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::RateLimited {
count,
upstream_error,
} => (
"rate_limited",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::Expired {
count,
upstream_error,
} => ("expired", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::Disabled {
count,
upstream_error,
} => ("disabled", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::Banned {
count,
upstream_error,
..
} => ("banned", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
"cloudflare_challenge",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
"cloudflare_path_blocked",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
"model_unsupported",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::Upstream(error) => {
let upstream_status = match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
};
(
"upstream",
None,
Some(upstream_error_body(error)),
upstream_status,
)
}
ChatDispatchError::InvalidSse(_) => ("invalid_sse", None, None, None),
ChatDispatchError::EmptyUpstreamResponse => ("empty_upstream_response", None, None, None),
};
object.insert(
"failureClass".to_string(),
Value::String(failure_class.to_string()),
);
if let Some(count) = exhausted_count {
object.insert("exhaustedCount".to_string(), serde_json::json!(count));
}
if let Some(error) = upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = upstream_status {
object.insert("upstreamStatus".to_string(), serde_json::json!(status));
}
}
+124
View File
@@ -1,12 +1,136 @@
//! Shared upstream error classification for dispatch routes.
use chrono::{DateTime, Duration, Utc};
use serde_json::{json, Map, Value};
use crate::{
proxy::dispatch::exhaustion::{ExhaustedAccountKind, ExhaustedAccountRef},
upstream::accounts::model::AccountStatus,
upstream::transport::{is_banned_auth_signal, CodexBackendTransport, CodexClientError},
};
#[derive(Clone, Debug)]
pub(crate) struct DispatchErrorMetadata {
pub failure_class: DispatchFailureClass,
pub exhausted_count: Option<usize>,
pub upstream_error: Option<String>,
pub upstream_status: Option<u16>,
}
#[derive(Clone, Copy, Debug)]
pub(crate) enum DispatchFailureClass {
NoAvailableAccounts,
QuotaExhausted,
RateLimited,
Expired,
Disabled,
Banned,
CloudflareChallenge,
CloudflarePathBlocked,
ModelUnsupported,
Upstream,
InvalidSse,
MissingCompleted,
EmptyUpstreamResponse,
ResponseFailed,
}
impl DispatchFailureClass {
pub(crate) fn as_str(self) -> &'static str {
match self {
Self::NoAvailableAccounts => "no_available_accounts",
Self::QuotaExhausted => "quota_exhausted",
Self::RateLimited => "rate_limited",
Self::Expired => "expired",
Self::Disabled => "disabled",
Self::Banned => "banned",
Self::CloudflareChallenge => "cloudflare_challenge",
Self::CloudflarePathBlocked => "cloudflare_path_blocked",
Self::ModelUnsupported => "model_unsupported",
Self::Upstream => "upstream",
Self::InvalidSse => "invalid_sse",
Self::MissingCompleted => "missing_completed",
Self::EmptyUpstreamResponse => "empty_upstream_response",
Self::ResponseFailed => "response_failed",
}
}
}
impl DispatchErrorMetadata {
pub(crate) fn no_available_accounts() -> Self {
Self {
failure_class: DispatchFailureClass::NoAvailableAccounts,
exhausted_count: None,
upstream_error: None,
upstream_status: None,
}
}
pub(crate) fn exhausted_ref(exhausted: ExhaustedAccountRef<'_>) -> Self {
Self {
failure_class: DispatchFailureClass::from(exhausted.kind),
exhausted_count: Some(exhausted.count),
upstream_error: Some(exhausted.upstream_error.to_string()),
upstream_status: None,
}
}
pub(crate) fn upstream(error: &CodexClientError) -> Self {
Self {
failure_class: DispatchFailureClass::Upstream,
exhausted_count: None,
upstream_error: Some(upstream_error_body(error)),
upstream_status: match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
},
}
}
pub(crate) fn simple(failure_class: DispatchFailureClass) -> Self {
Self {
failure_class,
exhausted_count: None,
upstream_error: None,
upstream_status: None,
}
}
}
impl From<ExhaustedAccountKind> for DispatchFailureClass {
fn from(kind: ExhaustedAccountKind) -> Self {
match kind {
ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted,
ExhaustedAccountKind::RateLimited => Self::RateLimited,
ExhaustedAccountKind::Expired => Self::Expired,
ExhaustedAccountKind::Disabled => Self::Disabled,
ExhaustedAccountKind::Banned => Self::Banned,
ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge,
ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked,
ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported,
}
}
}
pub(crate) fn insert_dispatch_error_metadata(
object: &mut Map<String, Value>,
metadata: DispatchErrorMetadata,
) {
object.insert(
"failureClass".to_string(),
Value::String(metadata.failure_class.as_str().to_string()),
);
if let Some(count) = metadata.exhausted_count {
object.insert("exhaustedCount".to_string(), json!(count));
}
if let Some(error) = metadata.upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = metadata.upstream_status {
object.insert("upstreamStatus".to_string(), json!(status));
}
}
pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
+222
View File
@@ -0,0 +1,222 @@
//! 账号耗尽状态聚合。
use crate::upstream::accounts::model::AccountStatus;
#[derive(Clone, Copy, Debug)]
pub(crate) enum ExhaustedAccountKind {
QuotaExhausted,
RateLimited,
Expired,
Disabled,
Banned,
CloudflareChallenge,
CloudflarePathBlocked,
ModelUnsupported,
}
impl ExhaustedAccountKind {
pub(crate) fn message_reason(self) -> &'static str {
match self {
Self::QuotaExhausted => "quota-exhausted",
Self::RateLimited => "rate-limited",
Self::Expired => "expired",
Self::Disabled => "disabled",
Self::Banned => "banned",
Self::CloudflareChallenge => "cloudflare-challenge",
Self::CloudflarePathBlocked => "cloudflare-path-block",
Self::ModelUnsupported => "model-unsupported",
}
}
}
#[derive(Clone, Debug)]
pub(crate) struct ExhaustedAccount {
pub kind: ExhaustedAccountKind,
pub count: usize,
pub upstream_error: String,
pub status_code: Option<u16>,
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct ExhaustedAccountRef<'a> {
pub kind: ExhaustedAccountKind,
pub count: usize,
pub upstream_error: &'a str,
}
#[derive(Default)]
pub(crate) struct AccountExhaustionTracker {
last_kind: Option<ExhaustedAccountKind>,
last_account_id: Option<String>,
quota_exhausted: ExhaustedCounter,
rate_limited: ExhaustedCounter,
expired: ExhaustedCounter,
disabled: ExhaustedCounter,
banned: ExhaustedCounter,
cloudflare_challenge: ExhaustedCounter,
cloudflare_path_blocked: ExhaustedCounter,
model_unsupported: ExhaustedCounter,
model_unsupported_retry_used: bool,
}
#[derive(Default)]
struct ExhaustedCounter {
count: usize,
upstream_error: Option<String>,
status_code: Option<u16>,
}
impl AccountExhaustionTracker {
pub(crate) fn last_account_id(&self) -> Option<&str> {
self.last_account_id.as_deref()
}
pub(crate) fn record_rate_limited(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::RateLimited,
upstream_error,
None,
);
}
pub(crate) fn record_quota_exhausted(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::QuotaExhausted,
upstream_error,
None,
);
}
pub(crate) fn record_auth_failure(
&mut self,
account_id: Option<&str>,
account_status: AccountStatus,
upstream_error: impl Into<String>,
status_code: Option<u16>,
) {
let kind = match account_status {
AccountStatus::Disabled => ExhaustedAccountKind::Disabled,
AccountStatus::Banned => ExhaustedAccountKind::Banned,
_ => ExhaustedAccountKind::Expired,
};
self.record(account_id, kind, upstream_error, status_code);
}
pub(crate) fn record_cloudflare_challenge(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::CloudflareChallenge,
upstream_error,
None,
);
}
pub(crate) fn record_cloudflare_path_blocked(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::CloudflarePathBlocked,
upstream_error,
None,
);
}
pub(crate) fn record_model_unsupported(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.model_unsupported_retry_used = true;
self.record(
account_id,
ExhaustedAccountKind::ModelUnsupported,
upstream_error,
None,
);
}
pub(crate) fn model_unsupported_retry_exhausted(
&self,
upstream_error: impl Into<String>,
) -> Option<ExhaustedAccount> {
self.model_unsupported_retry_used.then(|| ExhaustedAccount {
kind: ExhaustedAccountKind::ModelUnsupported,
count: self.model_unsupported.count + 1,
upstream_error: upstream_error.into(),
status_code: None,
})
}
pub(crate) fn last_exhausted(&self) -> Option<ExhaustedAccount> {
let kind = self.last_kind?;
let counter = self.counter(kind);
Some(ExhaustedAccount {
kind,
count: counter.count,
upstream_error: counter.upstream_error.clone().unwrap_or_default(),
status_code: counter.status_code,
})
}
fn record(
&mut self,
account_id: Option<&str>,
kind: ExhaustedAccountKind,
upstream_error: impl Into<String>,
status_code: Option<u16>,
) {
self.last_kind = Some(kind);
if let Some(account_id) = account_id {
self.last_account_id = Some(account_id.to_string());
}
let counter = self.counter_mut(kind);
counter.count += 1;
counter.upstream_error = Some(upstream_error.into());
if let Some(status_code) = status_code {
counter.status_code = Some(status_code);
}
}
fn counter(&self, kind: ExhaustedAccountKind) -> &ExhaustedCounter {
match kind {
ExhaustedAccountKind::QuotaExhausted => &self.quota_exhausted,
ExhaustedAccountKind::RateLimited => &self.rate_limited,
ExhaustedAccountKind::Expired => &self.expired,
ExhaustedAccountKind::Disabled => &self.disabled,
ExhaustedAccountKind::Banned => &self.banned,
ExhaustedAccountKind::CloudflareChallenge => &self.cloudflare_challenge,
ExhaustedAccountKind::CloudflarePathBlocked => &self.cloudflare_path_blocked,
ExhaustedAccountKind::ModelUnsupported => &self.model_unsupported,
}
}
fn counter_mut(&mut self, kind: ExhaustedAccountKind) -> &mut ExhaustedCounter {
match kind {
ExhaustedAccountKind::QuotaExhausted => &mut self.quota_exhausted,
ExhaustedAccountKind::RateLimited => &mut self.rate_limited,
ExhaustedAccountKind::Expired => &mut self.expired,
ExhaustedAccountKind::Disabled => &mut self.disabled,
ExhaustedAccountKind::Banned => &mut self.banned,
ExhaustedAccountKind::CloudflareChallenge => &mut self.cloudflare_challenge,
ExhaustedAccountKind::CloudflarePathBlocked => &mut self.cloudflare_path_blocked,
ExhaustedAccountKind::ModelUnsupported => &mut self.model_unsupported,
}
}
}
@@ -20,8 +20,8 @@ impl ImplicitResumeSnapshot {
/// 捕获当前请求中隐式续接会改写的字段。
pub fn capture(request: &CodexResponsesRequest) -> Self {
Self {
input: request.input.clone(),
previous_response_id: request.previous_response_id.clone(),
input: request.input().to_vec(),
previous_response_id: request.previous_response_id().map(ToString::to_string),
turn_state: request.turn_state.clone(),
use_websocket: request.use_websocket,
force_http_sse: request.force_http_sse,
@@ -30,8 +30,8 @@ impl ImplicitResumeSnapshot {
/// 恢复隐式续接前的请求字段。
pub fn restore(self, request: &mut CodexResponsesRequest) {
request.input = self.input;
request.previous_response_id = self.previous_response_id;
request.set_input(self.input);
request.set_previous_response_id(self.previous_response_id);
request.turn_state = self.turn_state;
request.use_websocket = self.use_websocket;
request.force_http_sse = self.force_http_sse;
+1 -1
View File
@@ -1,9 +1,9 @@
//! 请求调度模块(账号选择、回退、恢复、会话亲和性)。
mod auth_recovery;
pub mod chat;
pub mod cloudflare;
pub mod errors;
mod exhaustion;
pub mod implicit_resume;
pub mod reasoning_replay;
pub mod responses;
@@ -207,9 +207,6 @@ fn sanitize_reasoning_item(value: &Value) -> Option<Value> {
}
sanitized.insert("summary".to_string(), Value::Array(summary));
sanitized.insert("encrypted_content".to_string(), json!(encrypted_content));
if let Some(content) = object.get("content").and_then(sanitize_content) {
sanitized.insert("content".to_string(), Value::Array(content));
}
Some(Value::Object(sanitized))
}
@@ -250,22 +247,6 @@ fn sanitize_summary(value: &Value) -> Option<Vec<Value>> {
)
}
fn sanitize_content(value: &Value) -> Option<Vec<Value>> {
let content = value
.as_array()?
.iter()
.filter_map(|part| {
let object = part.as_object()?;
if object.get("type").and_then(Value::as_str) != Some("reasoning_text") {
return None;
}
let text = object.get("text")?.as_str()?;
Some(json!({"type": "reasoning_text", "text": text}))
})
.collect::<Vec<_>>();
(!content.is_empty()).then_some(content)
}
fn replay_item_key(item: &Value) -> Option<String> {
match item.get("type").and_then(Value::as_str)? {
"reasoning" => item
@@ -26,8 +26,7 @@ pub(super) async fn record_response_affinity(
usage: Option<TokenUsage>,
) {
let Some(conversation_id) = request
.prompt_cache_key
.as_deref()
.prompt_cache_key()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
@@ -52,7 +51,7 @@ pub(super) async fn record_response_affinity(
turn_state: turn_state
.filter(|value| !value.trim().is_empty())
.or_else(|| request.turn_state.clone()),
instructions_hash: Some(hash_instructions(Some(&request.instructions))),
instructions_hash: Some(hash_instructions(Some(request.instructions()))),
input_tokens: usage.map(|usage| usage.input_tokens),
function_call_ids: metadata.function_call_ids,
variant_hash: Some(variant_hash.clone()),
@@ -86,11 +85,7 @@ pub(super) async fn evict_reasoning_replay(
account_id: &str,
) {
let variant_hash = compute_variant_hash(request);
let conversation_id = request
.prompt_cache_key
.as_deref()
.unwrap_or("")
.to_string();
let conversation_id = request.prompt_cache_key().unwrap_or("").to_string();
reasoning_replay.lock().await.evict_by_identity(
account_id,
&conversation_id,
+202 -104
View File
@@ -1,8 +1,14 @@
use serde_json::{json, Value};
use serde_json::Value;
use thiserror::Error;
use crate::{
proxy::dispatch::errors::{upstream_error_body, upstream_error_http_status},
proxy::dispatch::{
errors::{
insert_dispatch_error_metadata, upstream_error_http_status, DispatchErrorMetadata,
DispatchFailureClass,
},
exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
},
upstream::{
protocol::{responses::ResponsesSseFailure, sse::SseError},
transport::CodexClientError,
@@ -14,8 +20,6 @@ use super::sse_failure::{sse_failure_error_body, stream_failure_http_status};
/// Responses 调度错误。
#[derive(Debug, Error)]
pub enum ResponseDispatchError {
#[error("failed to list runtime accounts")]
AccountStore,
#[error("no active account is available")]
NoActiveAccount,
#[error("all accounts exhausted by quota")]
@@ -72,9 +76,47 @@ pub enum ResponseDispatchError {
}
impl ResponseDispatchError {
pub(crate) fn from_exhausted_account(exhausted: ExhaustedAccount) -> Self {
match exhausted.kind {
ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::RateLimited => Self::RateLimited {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Expired => Self::Expired {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Disabled => Self::Disabled {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Banned => Self::Banned {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
status_code: exhausted.status_code.unwrap_or(403),
},
ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
}
}
pub fn http_status_code(&self) -> u16 {
match self {
Self::NoActiveAccount | Self::AccountStore => 503,
Self::NoActiveAccount => 503,
Self::QuotaExhausted { .. } => 429,
Self::RateLimited { .. } => 429,
Self::Expired { .. } | Self::Disabled { .. } => 401,
@@ -89,6 +131,160 @@ impl ResponseDispatchError {
Self::Upstream(error) => upstream_error_http_status(error),
}
}
pub(crate) fn metadata(&self) -> DispatchErrorMetadata {
match self {
Self::NoActiveAccount => DispatchErrorMetadata::no_available_accounts(),
Self::QuotaExhausted {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::QuotaExhausted,
*count,
upstream_error,
),
Self::RateLimited {
count,
upstream_error,
} => {
Self::exhausted_metadata(ExhaustedAccountKind::RateLimited, *count, upstream_error)
}
Self::Expired {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Expired, *count, upstream_error),
Self::Disabled {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Disabled, *count, upstream_error),
Self::Banned {
count,
upstream_error,
..
} => Self::exhausted_metadata(ExhaustedAccountKind::Banned, *count, upstream_error),
Self::CloudflareChallenge {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflareChallenge,
*count,
upstream_error,
),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflarePathBlocked,
*count,
upstream_error,
),
Self::ModelUnsupported {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::ModelUnsupported,
*count,
upstream_error,
),
Self::Upstream(error) => DispatchErrorMetadata::upstream(error),
Self::InvalidSse(_) => DispatchErrorMetadata::simple(DispatchFailureClass::InvalidSse),
Self::MissingCompleted => {
DispatchErrorMetadata::simple(DispatchFailureClass::MissingCompleted)
}
Self::EmptyUpstreamResponse => {
DispatchErrorMetadata::simple(DispatchFailureClass::EmptyUpstreamResponse)
}
Self::Failed(failure) => DispatchErrorMetadata {
failure_class: DispatchFailureClass::ResponseFailed,
exhausted_count: None,
upstream_error: Some(sse_failure_error_body(failure)),
upstream_status: None,
},
}
}
pub(crate) fn exhausted_account(&self) -> Option<ExhaustedAccountRef<'_>> {
match self {
Self::QuotaExhausted {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::QuotaExhausted,
count: *count,
upstream_error,
}),
Self::RateLimited {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::RateLimited,
count: *count,
upstream_error,
}),
Self::Expired {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Expired,
count: *count,
upstream_error,
}),
Self::Disabled {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Disabled,
count: *count,
upstream_error,
}),
Self::Banned {
count,
upstream_error,
..
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Banned,
count: *count,
upstream_error,
}),
Self::CloudflareChallenge {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflareChallenge,
count: *count,
upstream_error,
}),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflarePathBlocked,
count: *count,
upstream_error,
}),
Self::ModelUnsupported {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::ModelUnsupported,
count: *count,
upstream_error,
}),
_ => None,
}
}
fn exhausted_metadata(
kind: ExhaustedAccountKind,
count: usize,
upstream_error: &str,
) -> DispatchErrorMetadata {
DispatchErrorMetadata::exhausted_ref(ExhaustedAccountRef {
kind,
count,
upstream_error,
})
}
}
/// Responses live SSE body stream error.
@@ -132,103 +328,5 @@ pub(super) fn enrich_response_dispatch_error_metadata(
let Some(object) = metadata.as_object_mut() else {
return;
};
let (failure_class, exhausted_count, upstream_error, upstream_status) = match error {
ResponseDispatchError::AccountStore => ("account_store", None, None, None),
ResponseDispatchError::NoActiveAccount => ("no_available_accounts", None, None, None),
ResponseDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
"quota_exhausted",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::RateLimited {
count,
upstream_error,
} => (
"rate_limited",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::Expired {
count,
upstream_error,
} => ("expired", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::Disabled {
count,
upstream_error,
} => ("disabled", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::Banned {
count,
upstream_error,
..
} => ("banned", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
"cloudflare_challenge",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
"cloudflare_path_blocked",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
"model_unsupported",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::Upstream(error) => {
let upstream_status = match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
};
(
"upstream",
None,
Some(upstream_error_body(error)),
upstream_status,
)
}
ResponseDispatchError::InvalidSse(_) => ("invalid_sse", None, None, None),
ResponseDispatchError::MissingCompleted => ("missing_completed", None, None, None),
ResponseDispatchError::EmptyUpstreamResponse => {
("empty_upstream_response", None, None, None)
}
ResponseDispatchError::Failed(failure) => (
"response_failed",
None,
Some(sse_failure_error_body(failure)),
None,
),
};
object.insert(
"failureClass".to_string(),
Value::String(failure_class.to_string()),
);
if let Some(count) = exhausted_count {
object.insert("exhaustedCount".to_string(), json!(count));
}
if let Some(error) = upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = upstream_status {
object.insert("upstreamStatus".to_string(), json!(status));
}
insert_dispatch_error_metadata(object, error.metadata());
}
@@ -168,7 +168,7 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
client_ip: record.request.client_ip.as_deref(),
client_user_agent: record.request.client_user_agent.as_deref(),
reasoning_effort: reasoning_effort_from_request(record.request),
service_tier: record.request.service_tier.as_deref(),
service_tier: record.request.service_tier(),
started_at: record.started_at,
status_code: status_code_for_stream_failure(record.failure),
level: UsageRecordLevel::Error,
@@ -246,7 +246,7 @@ pub(super) async fn record_live_response_stream_event(
context.client_ip.as_deref(),
context.request.client_user_agent.as_deref(),
reasoning_effort_from_request(&context.request),
context.request.service_tier.as_deref(),
context.request.service_tier(),
);
event.metadata = metadata;
if let Err(error) = context.usage_records.record(event).await {
File diff suppressed because it is too large Load Diff
+15 -21
View File
@@ -125,14 +125,6 @@ impl SessionAffinityMap {
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)
@@ -169,12 +161,9 @@ 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 tools_json = request.tools().map_or_else(|| "[]".to_string(), tools_json);
let mut hasher = Sha256::new();
hasher.update(request.instructions.as_bytes());
hasher.update(request.instructions().as_bytes());
hasher.update(b"\0");
hasher.update(tools_json.as_bytes());
if let Some(identity) = identity
@@ -263,23 +252,28 @@ pub fn build_conversation_identity(
/// 确保请求拥有上游可复用的 prompt cache key。
pub fn ensure_prompt_cache_key(request: &mut CodexResponsesRequest) {
if let Some(existing) = request
.prompt_cache_key
.as_deref()
.prompt_cache_key()
.map(str::trim)
.filter(|value| !value.is_empty())
{
request.prompt_cache_key = Some(existing.to_string());
let existing = existing.to_string();
request.set_prompt_cache_key(Some(existing));
return;
}
request.prompt_cache_key =
Some(derive_stable_conversation_key(request).unwrap_or_else(|| Uuid::new_v4().to_string()));
request.set_prompt_cache_key(Some(
derive_stable_conversation_key(request).unwrap_or_else(|| Uuid::new_v4().to_string()),
));
}
/// 按原版 `stable-conversation-key.ts` 的规则派生稳定 conversation key。
pub fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option<String> {
let instructions = request.instructions.chars().take(2000).collect::<String>();
let first_user_text = first_user_text(&request.input);
let instructions = request
.instructions()
.chars()
.take(2000)
.collect::<String>();
let first_user_text = first_user_text(request.input());
let normalized_first_user_text = normalize_conversation_anchor_text(&first_user_text);
let first_user_text = if normalized_first_user_text.is_empty() {
first_user_text
@@ -291,7 +285,7 @@ pub fn derive_stable_conversation_key(request: &CodexResponsesRequest) -> Option
}
let mut hasher = Sha256::new();
hasher.update(request.model.as_bytes());
hasher.update(request.model().as_bytes());
hasher.update(b"\0");
hasher.update(instructions.as_bytes());
hasher.update(b"\0");
+2 -2
View File
@@ -152,7 +152,7 @@ pub(crate) async fn create_response_with_account(
.cookie_header_for_request(&account.id, "/codex/responses")
.await;
let identity = build_conversation_identity(
request.prompt_cache_key.as_deref(),
request.prompt_cache_key(),
request.codex_window_id.as_deref(),
&account.id,
);
@@ -192,7 +192,7 @@ async fn create_response_stream_with_account(
.cookie_header_for_request(&account.id, "/codex/responses")
.await;
let identity = build_conversation_identity(
request.prompt_cache_key.as_deref(),
request.prompt_cache_key(),
request.codex_window_id.as_deref(),
&account.id,
);
+6 -12
View File
@@ -88,19 +88,13 @@ pub(super) async fn record_response_event(record: ResponseUsageRecord<'_>) {
}
pub(super) fn event_kind(route: &str) -> &'static str {
if route == "/v1/chat/completions" {
"v1.chat"
} else {
"v1.response"
}
let _ = route;
"v1.response"
}
pub(super) fn api_kind(route: &str) -> &'static str {
if route == "/v1/chat/completions" {
"chat"
} else {
"responses"
}
let _ = route;
"responses"
}
pub(super) fn enrich_event_route_metadata(metadata: &mut Value, route: &str) {
@@ -149,11 +143,11 @@ pub(super) fn enrich_usage_record_identity(
}
pub(super) fn reasoning_effort_from_request(request: &CodexResponsesRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning.as_ref())
reasoning_effort_from_value(request.reasoning())
}
pub(super) fn reasoning_effort_from_compact_request(request: &CodexCompactRequest) -> Option<&str> {
reasoning_effort_from_value(request.reasoning.as_ref())
reasoning_effort_from_value(request.reasoning())
}
fn insert_trimmed_string(
+1 -1
View File
@@ -1,7 +1,7 @@
//! OpenAI 兼容代理入口。
//!
//! 本模块负责:
//! - OpenAI API 协议转换(Chat Completions / Responses / Models)
//! - OpenAI API 协议转换(Responses / Models)
//! - 请求调度(账号选择、回退、恢复)
//! - 会话亲和性与 reasoning replay
File diff suppressed because it is too large Load Diff
+74 -180
View File
@@ -8,7 +8,7 @@ use axum::{
use serde_json::{json, Value};
use crate::proxy::{
dispatch::{chat::ChatDispatchError, responses::errors::ResponseDispatchError},
dispatch::{errors::DispatchFailureClass, responses::errors::ResponseDispatchError},
openai::responses::response_failed_sse_event as encode_response_failed_sse_event,
};
@@ -101,19 +101,9 @@ pub fn openai_error_response(
)
}
pub fn chat_dispatch_error_response(error: &ChatDispatchError) -> Response {
let error = chat_dispatch_openai_error(error);
openai_error_response(error.status, &error.message, error.error_type, error.code)
.into_response()
}
pub fn chat_stream_dispatch_openai_error(error: &ResponseDispatchError) -> OpenAiErrorDetails {
response_dispatch_openai_error(error)
}
pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Response {
match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
ResponseDispatchError::NoActiveAccount => {
responses_no_available_accounts_response().into_response()
}
error => response_dispatch_openai_error_response(&error),
@@ -122,7 +112,7 @@ pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Respon
pub fn responses_compact_dispatch_error_response(error: ResponseDispatchError) -> Response {
match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
ResponseDispatchError::NoActiveAccount => {
responses_no_available_accounts_response().into_response()
}
error => {
@@ -203,11 +193,6 @@ pub fn model_not_found_response() -> (StatusCode, Json<Value>) {
)
}
/// Chat Completions 请求无效。
pub fn invalid_chat_completion_request_response() -> (StatusCode, Json<Value>) {
invalid_openai_request_response("Invalid chat completion request")
}
/// Responses 请求无效。
pub fn invalid_responses_request_response() -> (StatusCode, Json<Value>) {
invalid_openai_request_response("Invalid responses request")
@@ -222,88 +207,6 @@ fn invalid_openai_request_response(message: &str) -> (StatusCode, Json<Value>) {
)
}
fn chat_dispatch_openai_error(error: &ChatDispatchError) -> OpenAiErrorDetails {
let status = chat_dispatch_status(error);
let (message, kind) = match error {
ChatDispatchError::NoActiveAccount | ChatDispatchError::AccountStore => (
NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned(),
UPSTREAM_UNAVAILABLE_ERROR,
),
ChatDispatchError::Upstream(_) => (
UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
UPSTREAM_ERROR,
),
ChatDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "quota-exhausted", upstream_error, false),
INSUFFICIENT_QUOTA_ERROR,
),
ChatDispatchError::RateLimited {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "rate-limited", upstream_error, false),
RATE_LIMIT_ERROR,
),
ChatDispatchError::Expired {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "expired", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::Disabled {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "disabled", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::Banned {
count,
upstream_error,
..
} => (
exhausted_dispatch_message(*count, "banned", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "cloudflare-challenge", upstream_error, false),
UPSTREAM_ERROR,
),
ChatDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "cloudflare-path-block", upstream_error, false),
UPSTREAM_ERROR,
),
ChatDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "model-unsupported", upstream_error, false),
MODEL_NOT_FOUND_ERROR,
),
ChatDispatchError::InvalidSse(_) | ChatDispatchError::EmptyUpstreamResponse => (
INVALID_UPSTREAM_CODEX_RESPONSE_MESSAGE.to_owned(),
INVALID_UPSTREAM_RESPONSE_ERROR,
),
};
OpenAiErrorDetails {
status,
message,
error_type: kind.error_type,
code: kind.code,
}
}
fn response_dispatch_openai_error(error: &ResponseDispatchError) -> OpenAiErrorDetails {
let DispatchHttpError { status, message } = response_dispatch_http_error(
error,
@@ -321,27 +224,10 @@ fn response_dispatch_openai_error(error: &ResponseDispatchError) -> OpenAiErrorD
}
fn response_dispatch_error_kind(error: &ResponseDispatchError) -> OpenAiErrorKind {
match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
UPSTREAM_UNAVAILABLE_ERROR
}
ResponseDispatchError::QuotaExhausted { .. } => INSUFFICIENT_QUOTA_ERROR,
ResponseDispatchError::RateLimited { .. } => RATE_LIMIT_ERROR,
ResponseDispatchError::Expired { .. }
| ResponseDispatchError::Disabled { .. }
| ResponseDispatchError::Banned { .. } => AUTHENTICATION_ERROR,
ResponseDispatchError::ModelUnsupported { .. } => MODEL_NOT_FOUND_ERROR,
ResponseDispatchError::InvalidSse(_)
| ResponseDispatchError::MissingCompleted
| ResponseDispatchError::EmptyUpstreamResponse => INVALID_UPSTREAM_RESPONSE_ERROR,
ResponseDispatchError::Failed(_) => responses_error_kind_for_status(status_from_u16(
error.http_status_code(),
StatusCode::BAD_GATEWAY,
)),
ResponseDispatchError::Upstream(_)
| ResponseDispatchError::CloudflareChallenge { .. }
| ResponseDispatchError::CloudflarePathBlocked { .. } => UPSTREAM_ERROR,
}
dispatch_failure_openai_error_kind(
error.metadata().failure_class,
status_from_u16(error.http_status_code(), StatusCode::BAD_GATEWAY),
)
}
pub fn response_dispatch_http_error(
@@ -372,13 +258,6 @@ fn responses_error_kind_for_status(status: StatusCode) -> OpenAiErrorKind {
}
}
fn chat_dispatch_status(error: &ChatDispatchError) -> StatusCode {
match error {
ChatDispatchError::Upstream(_) => StatusCode::BAD_GATEWAY,
_ => status_from_u16(error.http_status_code(), StatusCode::BAD_GATEWAY),
}
}
fn response_dispatch_status(
error: &ResponseDispatchError,
mode: ResponseDispatchStatusMode,
@@ -395,60 +274,75 @@ fn response_dispatch_message(
error: &ResponseDispatchError,
style: ResponseDispatchMessageStyle,
) -> String {
match error {
ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => {
NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned()
}
ResponseDispatchError::Upstream(_) => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
ResponseDispatchError::QuotaExhausted {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "quota-exhausted", upstream_error, false),
ResponseDispatchError::RateLimited {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "rate-limited", upstream_error, false),
ResponseDispatchError::Expired {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "expired", upstream_error, false),
ResponseDispatchError::Disabled {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "disabled", upstream_error, false),
ResponseDispatchError::Banned {
count,
upstream_error,
..
} => exhausted_dispatch_message(*count, "banned", upstream_error, false),
ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
} => exhausted_dispatch_message(*count, "cloudflare-challenge", upstream_error, false),
ResponseDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => exhausted_dispatch_message(
*count,
"cloudflare-path-block",
upstream_error,
matches!(style, ResponseDispatchMessageStyle::ResponsesStream),
),
ResponseDispatchError::ModelUnsupported {
count,
upstream_error,
} => exhausted_dispatch_message(
*count,
"model-unsupported",
upstream_error,
matches!(style, ResponseDispatchMessageStyle::ResponsesStream),
),
ResponseDispatchError::InvalidSse(_)
| ResponseDispatchError::MissingCompleted
| ResponseDispatchError::EmptyUpstreamResponse => {
let metadata = error.metadata();
if let Some(exhausted) = error.exhausted_account() {
return exhausted_dispatch_message(
exhausted.count,
exhausted.kind.message_reason(),
exhausted.upstream_error,
matches!(
(metadata.failure_class, style),
(
DispatchFailureClass::CloudflarePathBlocked
| DispatchFailureClass::ModelUnsupported,
ResponseDispatchMessageStyle::ResponsesStream,
)
),
);
}
dispatch_failure_message(
metadata.failure_class,
match style {
ResponseDispatchMessageStyle::Standard => DispatchFailureMessageStyle::Standard,
ResponseDispatchMessageStyle::ResponsesStream => {
DispatchFailureMessageStyle::ResponsesStream
}
},
)
}
#[derive(Clone, Copy)]
enum DispatchFailureMessageStyle {
Standard,
ResponsesStream,
}
fn dispatch_failure_message(
failure_class: DispatchFailureClass,
_style: DispatchFailureMessageStyle,
) -> String {
match failure_class {
DispatchFailureClass::NoAvailableAccounts => NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned(),
DispatchFailureClass::Upstream => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
DispatchFailureClass::InvalidSse
| DispatchFailureClass::MissingCompleted
| DispatchFailureClass::EmptyUpstreamResponse => {
INVALID_UPSTREAM_CODEX_RESPONSE_MESSAGE.to_owned()
}
ResponseDispatchError::Failed(_) => UPSTREAM_CODEX_RESPONSE_FAILED_MESSAGE.to_owned(),
DispatchFailureClass::ResponseFailed => UPSTREAM_CODEX_RESPONSE_FAILED_MESSAGE.to_owned(),
_ => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
}
}
fn dispatch_failure_openai_error_kind(
failure_class: DispatchFailureClass,
status: StatusCode,
) -> OpenAiErrorKind {
match failure_class {
DispatchFailureClass::NoAvailableAccounts => UPSTREAM_UNAVAILABLE_ERROR,
DispatchFailureClass::QuotaExhausted => INSUFFICIENT_QUOTA_ERROR,
DispatchFailureClass::RateLimited => RATE_LIMIT_ERROR,
DispatchFailureClass::Expired
| DispatchFailureClass::Disabled
| DispatchFailureClass::Banned => AUTHENTICATION_ERROR,
DispatchFailureClass::ModelUnsupported => MODEL_NOT_FOUND_ERROR,
DispatchFailureClass::InvalidSse
| DispatchFailureClass::MissingCompleted
| DispatchFailureClass::EmptyUpstreamResponse => INVALID_UPSTREAM_RESPONSE_ERROR,
DispatchFailureClass::ResponseFailed => responses_error_kind_for_status(status),
DispatchFailureClass::Upstream
| DispatchFailureClass::CloudflareChallenge
| DispatchFailureClass::CloudflarePathBlocked => UPSTREAM_ERROR,
}
}
-1
View File
@@ -1,6 +1,5 @@
//! OpenAI 兼容 API 模块。
pub mod chat;
pub mod errors;
pub mod models;
pub mod responses;
+194 -519
View File
@@ -1,4 +1,5 @@
//! OpenAI Responses API 类型、协议转换与 HTTP 处理器。
//! OpenAI Responses API 透明代理入口:解析原始 JSON body,提取调度元数据,
//! 做最小 patch 后交给调度层,请求语义原样透传上游。
use axum::{
body::{Body, Bytes},
@@ -7,20 +8,17 @@ use axum::{
response::{IntoResponse, Response},
Extension, Json,
};
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use std::net::IpAddr;
use crate::{
http::middleware::request_id::RequestId,
http::middleware::request_id::{ClientIp, RequestId},
proxy::{
auth::authorize_client_api_key,
dispatch::responses::{errors::ResponseDispatchError, service::ResponseDispatchStream},
},
runtime::state::AppState,
upstream::protocol::{
responses::{ensure_reasoning_include, CodexCompactRequest, CodexResponsesRequest},
schema::prepare_schema,
responses::{CodexCompactRequest, CodexResponsesRequest},
sse::{encode_sse_event, sse_body_has_done},
},
};
@@ -37,419 +35,162 @@ use super::{
const OPENAI_SUBAGENT_HEADER: &str = "x-openai-subagent";
// ====================================================================
// 协议类型
// ====================================================================
/// 代理内部 transport 提示字段:用于本地传输决策,不转发上游。
const TRANSPORT_ONLY_KEYS: [&str; 1] = ["use_websocket"];
/// 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>,
#[serde(default)]
pub reasoning: Option<Value>,
#[serde(default)]
pub tools: Option<Value>,
#[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>,
#[serde(default)]
pub include: Option<Value>,
#[serde(default)]
pub client_metadata: Option<Value>,
#[serde(default)]
pub previous_response_id: Option<String>,
#[serde(default, rename = "turnState", alias = "turn_state")]
pub turn_state: Option<String>,
#[serde(default, rename = "turnMetadata", alias = "turn_metadata")]
pub turn_metadata: Option<String>,
#[serde(default, rename = "betaFeatures", alias = "beta_features")]
pub beta_features: Option<String>,
#[serde(default)]
pub version: Option<String>,
#[serde(
default,
rename = "includeTimingMetrics",
alias = "include_timing_metrics"
)]
pub include_timing_metrics: Option<String>,
#[serde(default, rename = "codexWindowId", alias = "codex_window_id")]
pub codex_window_id: Option<String>,
#[serde(default, rename = "parentThreadId", alias = "parent_thread_id")]
pub parent_thread_id: Option<String>,
#[serde(default)]
pub use_websocket: Option<bool>,
#[serde(default = "default_responses_stream")]
pub stream: bool,
}
/// Codex 上下文字段:从 body 顶层(camelCase)提取到代理控制状态用于生成请求头,
/// 同时保留在 body 中原样转发上游。
const CONTEXT_HEADER_FIELDS: [(&str, ContextField); 7] = [
("turnState", ContextField::TurnState),
("turnMetadata", ContextField::TurnMetadata),
("betaFeatures", ContextField::BetaFeatures),
("version", ContextField::Version),
("includeTimingMetrics", ContextField::IncludeTimingMetrics),
("codexWindowId", ContextField::CodexWindowId),
("parentThreadId", ContextField::ParentThreadId),
];
fn default_responses_input() -> Value {
Value::Array(Vec::new())
}
fn default_responses_stream() -> bool {
true
#[derive(Clone, Copy)]
enum ContextField {
TurnState,
TurnMetadata,
BetaFeatures,
Version,
IncludeTimingMetrics,
CodexWindowId,
ParentThreadId,
}
// ====================================================================
// 协议转换
// 透明代理请求构造
// ====================================================================
/// 将 OpenAI Responses 请求转换为 Codex Responses 请求。
pub fn translate_response_to_codex(request: OpenAiResponsesRequest) -> CodexResponsesRequest {
let prepared_text = prepare_text_format(request.text, true);
let client_metadata = sanitize_client_metadata(request.client_metadata);
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 = first_request_string(
request.turn_metadata,
client_metadata.as_ref(),
"x-codex-turn-metadata",
);
codex_request.beta_features = first_request_string(
request.beta_features,
client_metadata.as_ref(),
"x-codex-beta-features",
);
codex_request.include_timing_metrics = first_request_string(
request.include_timing_metrics,
client_metadata.as_ref(),
"x-responsesapi-include-timing-metrics",
);
codex_request.version = non_empty_string(request.version);
codex_request.codex_window_id = first_request_string(
request.codex_window_id,
client_metadata.as_ref(),
"x-codex-window-id",
);
codex_request.parent_thread_id = first_request_string(
request.parent_thread_id,
client_metadata.as_ref(),
"x-codex-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 = 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,
/// 从客户端原始 Responses JSON body 构造上游请求(透明代理)。
///
/// 只做代理职责范围内的最小处理:
/// - 剥离 transport-only 字段(`use_websocket`),仅用于本地传输决策。
/// - 从 body 与请求头提取 Codex 上下文透传字段到代理控制状态(body 中保留原值)。
/// - review route / 合法请求头时,往 `client_metadata` 注入 forced subagent。
///
/// 其余字段——`input`、`reasoning`、`text`、`tools`、`tool_choice`、`include`、
/// `client_metadata`、`service_tier` 以及一切未知顶层字段——全部原样透传,不重写。
pub fn build_codex_request(
mut body: Map<String, Value>,
headers: &HeaderMap,
forced_subagent: Option<&str>,
) -> CodexResponsesRequest {
let use_websocket = body.get("use_websocket").and_then(Value::as_bool);
for key in TRANSPORT_ONLY_KEYS {
body.remove(key);
}
if let Some(subagent) = forced_subagent
.map(ToString::to_string)
.or_else(|| openai_subagent_from_headers(headers))
{
inject_subagent_metadata(&mut body, &subagent);
}
let explicit_prompt_cache_key = body
.get("prompt_cache_key")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty());
let mut request = CodexResponsesRequest::from_body(body);
request.explicit_prompt_cache_key = explicit_prompt_cache_key;
match use_websocket {
Some(true) => request.use_websocket = true,
Some(false) => request.force_http_sse = true,
None => {}
}
codex_request
apply_context_header_fields(&mut request, headers);
request.client_user_agent = user_agent_from_headers(headers);
request
}
/// 将 OpenAI Responses 请求转换为 Codex compact 请求。
pub fn translate_response_to_compact(request: OpenAiResponsesRequest) -> CodexCompactRequest {
/// 从客户端原始 Responses JSON body 构造 Codex compact 请求。
///
/// compact 端点只返回一次性 JSON(代理按非流式读取),故剥离 `stream` 这一
/// transport 控制字段,避免上游返回代理无法解析的 SSE 形态。业务语义字段
/// (`reasoning`/`text`/`store`/`prompt_cache_key`/`previous_response_id`/
/// `include`/`client_metadata`/未知字段)一律原样透传。
pub fn build_compact_request(
mut body: Map<String, Value>,
headers: &HeaderMap,
) -> CodexCompactRequest {
body.remove(COMPACT_STREAM_KEY);
for key in TRANSPORT_ONLY_KEYS {
body.remove(key);
}
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,
body,
client_ip: None,
client_user_agent: None,
client_user_agent: user_agent_from_headers(headers),
}
}
struct PreparedTextFormat {
text: Option<Value>,
tuple_schema: Option<Value>,
}
/// compact 端点不支持流式,转发前剥离 `stream`(transport 控制字段,非业务语义)。
const COMPACT_STREAM_KEY: &str = "stream";
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()));
fn inject_subagent_metadata(body: &mut Map<String, Value>, subagent: &str) {
let metadata = body
.entry("client_metadata")
.or_insert_with(|| Value::Object(Map::new()));
if !metadata.is_object() {
*metadata = Value::Object(Map::new());
}
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::String(text) => vec![responses_input_text_message(&text)],
value => vec![value],
}
}
fn sanitize_codex_input_items(input: Vec<Value>) -> Vec<Value> {
input
.into_iter()
.filter_map(|item| {
if let Value::String(text) = item {
return Some(responses_input_text_message(&text));
}
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 responses_input_text_message(text: &str) -> Value {
json!({
"type": "message",
"role": "user",
"content": [
{
"type": "input_text",
"text": text
}
]
})
}
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(object) = metadata.as_object_mut() {
object.insert(
OPENAI_SUBAGENT_HEADER.to_string(),
Value::String(subagent.to_string()),
);
}
if let Some(content) = sanitize_reasoning_content(item.get("content")) {
sanitized.insert("content".to_string(), Value::Array(content));
}
/// 从 body 与请求头提取 Codex 上下文透传字段到代理控制状态。
///
/// body 中的原值优先;缺失时回退请求头。提取只填充代理控制状态用于加请求头,
/// **不修改 body**——这些字段在 body 中原样保留转发上游。
fn apply_context_header_fields(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
for (body_key, field) in CONTEXT_HEADER_FIELDS {
let value = body_context_string(request.body(), body_key)
.or_else(|| header_string(headers, field.header_name()));
field.assign(request, value);
}
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")
fn body_context_string(body: &Map<String, Value>, key: &str) -> Option<String> {
body.get(key)
.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()));
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
impl ContextField {
fn header_name(self) -> &'static str {
match self {
Self::TurnState => "x-codex-turn-state",
Self::TurnMetadata => "x-codex-turn-metadata",
Self::BetaFeatures => "x-codex-beta-features",
Self::Version => "version",
Self::IncludeTimingMetrics => "x-responsesapi-include-timing-metrics",
Self::CodexWindowId => "x-codex-window-id",
Self::ParentThreadId => "x-codex-parent-thread-id",
}
}
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()));
fn assign(self, request: &mut CodexResponsesRequest, value: Option<String>) {
match self {
Self::TurnState => request.turn_state = value,
Self::TurnMetadata => request.turn_metadata = value,
Self::BetaFeatures => request.beta_features = value,
Self::Version => request.version = value,
Self::IncludeTimingMetrics => request.include_timing_metrics = value,
Self::CodexWindowId => request.codex_window_id = value,
Self::ParentThreadId => request.parent_thread_id = value,
}
}
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 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 first_request_string(
direct: Option<String>,
client_metadata: Option<&Value>,
metadata_key: &str,
) -> Option<String> {
non_empty_string(direct).or_else(|| metadata_string(client_metadata, metadata_key))
}
fn metadata_string(client_metadata: Option<&Value>, key: &str) -> Option<String> {
client_metadata?
.as_object()?
.get(key)?
.as_str()
.and_then(|value| {
let trimmed = value.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
})
}
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)
}
/// 编码 OpenAI Responses `response.failed` SSE 事件。
@@ -495,22 +236,34 @@ pub fn response_failed_sse_event_with_id(
pub async fn responses(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_responses(state, request_id, headers, body, "/v1/responses", None).await
}
/// `POST /v1/responses/review`
pub async fn review_responses(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
client_ip: Option<Extension<ClientIp>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_responses(
state,
request_id,
client_ip_string(client_ip),
headers,
body,
"/v1/responses",
None,
)
.await
}
/// `POST /v1/responses/review`
pub async fn review_responses(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
client_ip: Option<Extension<ClientIp>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_responses(
state,
request_id,
client_ip_string(client_ip),
headers,
body,
"/v1/responses/review",
@@ -523,15 +276,29 @@ pub async fn review_responses(
pub async fn compact_responses(
State(state): State<AppState>,
Extension(request_id): Extension<RequestId>,
client_ip: Option<Extension<ClientIp>>,
headers: HeaderMap,
body: Bytes,
) -> Response {
handle_compact_responses(state, request_id, headers, body).await
handle_compact_responses(
state,
request_id,
client_ip_string(client_ip),
headers,
body,
)
.await
}
/// 从可选的 `ClientIp` extension 提取 IP 字符串。
fn client_ip_string(client_ip: Option<Extension<ClientIp>>) -> Option<String> {
client_ip.map(|Extension(client_ip)| client_ip.as_str().to_string())
}
async fn handle_responses(
state: AppState,
request_id: RequestId,
client_ip: Option<String>,
headers: HeaderMap,
body: Bytes,
route: &str,
@@ -541,19 +308,10 @@ async fn handle_responses(
return missing_client_api_key_response().into_response();
}
let Ok(mut openai_request) = parse_openai_responses_request(&body) else {
let Some(body) = parse_responses_body(&body) else {
return invalid_responses_request_response().into_response();
};
apply_responses_header_context(&mut openai_request, &headers);
if openai_request.use_websocket.is_none() {
openai_request.use_websocket = Some(true);
}
if let Some(subagent) = forced_subagent {
force_openai_subagent(&mut openai_request, subagent);
} else if let Some(subagent) = openai_subagent_from_headers(&headers) {
force_openai_subagent(&mut openai_request, &subagent);
}
let model = match validated_responses_model(&state, &openai_request.model).await {
let model = match validated_responses_model(&state, request_model(&body)).await {
Ok(model) => model,
Err(ResponsesModelValidationError::InvalidRequest) => {
return invalid_responses_request_response().into_response();
@@ -562,9 +320,13 @@ async fn handle_responses(
return model_not_found_response().into_response();
}
};
let stream = openai_request.stream;
let mut codex_request = translate_response_to_codex(openai_request);
attach_client_context(&mut codex_request, &headers);
let mut codex_request = build_codex_request(body, &headers, forced_subagent);
codex_request.client_ip = client_ip;
// 客户端未显式指定 transport 时默认偏好 WebSocket。
if !codex_request.force_http_sse {
codex_request.use_websocket = true;
}
let stream = codex_request.stream();
if stream {
return match state
@@ -592,6 +354,7 @@ async fn handle_responses(
async fn handle_compact_responses(
state: AppState,
request_id: RequestId,
client_ip: Option<String>,
headers: HeaderMap,
body: Bytes,
) -> Response {
@@ -599,10 +362,10 @@ async fn handle_compact_responses(
return missing_client_api_key_response().into_response();
}
let Ok(openai_request) = parse_openai_responses_request(&body) else {
let Some(body) = parse_responses_body(&body) else {
return invalid_responses_request_response().into_response();
};
let model = match validated_responses_model(&state, &openai_request.model).await {
let model = match validated_responses_model(&state, request_model(&body)).await {
Ok(model) => model,
Err(ResponsesModelValidationError::InvalidRequest) => {
return invalid_responses_request_response().into_response();
@@ -611,9 +374,8 @@ async fn handle_compact_responses(
return model_not_found_response().into_response();
}
};
let mut compact_request = translate_response_to_compact(openai_request);
compact_request.client_ip = client_ip_from_headers(&headers);
compact_request.client_user_agent = user_agent_from_headers(&headers);
let mut compact_request = build_compact_request(body, &headers);
compact_request.client_ip = client_ip;
match state
.services
@@ -626,22 +388,18 @@ async fn handle_compact_responses(
}
}
fn force_openai_subagent(request: &mut OpenAiResponsesRequest, subagent: &str) {
let mut metadata = match request.client_metadata.take() {
Some(Value::Object(metadata)) => metadata,
_ => serde_json::Map::new(),
};
metadata.insert(
OPENAI_SUBAGENT_HEADER.to_string(),
Value::String(subagent.to_string()),
);
request.client_metadata = Some(Value::Object(metadata));
/// 解析请求体:必须是 JSON object,否则视为 invalid request。
fn parse_responses_body(body: &Bytes) -> Option<Map<String, Value>> {
match serde_json::from_slice::<Value>(body) {
Ok(Value::Object(map)) => Some(map),
_ => None,
}
}
fn parse_openai_responses_request(
body: &Bytes,
) -> Result<OpenAiResponsesRequest, serde_json::Error> {
serde_json::from_slice(body)
fn request_model(body: &Map<String, Value>) -> &str {
body.get("model")
.and_then(Value::as_str)
.unwrap_or_default()
}
async fn validated_responses_model(
@@ -665,93 +423,10 @@ enum ResponsesModelValidationError {
ModelNotFound,
}
fn apply_responses_header_context(request: &mut OpenAiResponsesRequest, headers: &HeaderMap) {
fill_string_from_header(&mut request.turn_state, headers, "x-codex-turn-state");
fill_string_from_header(&mut request.turn_metadata, headers, "x-codex-turn-metadata");
fill_string_from_header(&mut request.beta_features, headers, "x-codex-beta-features");
fill_string_from_header(
&mut request.include_timing_metrics,
headers,
"x-responsesapi-include-timing-metrics",
);
fill_string_from_header(&mut request.version, headers, "version");
fill_string_from_header(&mut request.codex_window_id, headers, "x-codex-window-id");
fill_string_from_header(
&mut request.parent_thread_id,
headers,
"x-codex-parent-thread-id",
);
}
fn fill_string_from_header(field: &mut Option<String>, headers: &HeaderMap, name: &str) {
if field
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
{
return;
}
*field = header_string(headers, name);
}
pub(crate) fn attach_client_context(request: &mut CodexResponsesRequest, headers: &HeaderMap) {
request.client_ip = client_ip_from_headers(headers);
request.client_user_agent = user_agent_from_headers(headers);
}
pub(crate) fn user_agent_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, "user-agent")
}
pub(crate) fn client_ip_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, "cf-connecting-ip")
.or_else(|| header_string(headers, "x-real-ip"))
.or_else(|| forwarded_client_ip(headers))
.and_then(|value| normalize_client_ip(&value))
}
fn forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
let value = header_string(headers, "x-forwarded-for")?;
let parts = value
.split(',')
.map(str::trim)
.filter(|part| !part.is_empty())
.collect::<Vec<_>>();
parts
.iter()
.find(|part| !is_private_ip(part))
.or_else(|| parts.first())
.map(|part| (*part).to_string())
}
fn normalize_client_ip(value: &str) -> Option<String> {
let trimmed = value.trim();
if trimmed.is_empty() {
return None;
}
trimmed
.parse::<std::net::SocketAddr>()
.map(|addr| addr.ip().to_string())
.ok()
.or_else(|| Some(trimmed.to_string()))
}
fn is_private_ip(value: &str) -> bool {
let Some(ip) = normalize_client_ip(value).and_then(|value| value.parse::<IpAddr>().ok()) else {
return false;
};
match ip {
IpAddr::V4(ip) => {
ip.is_private()
|| ip.is_loopback()
|| ip.is_link_local()
|| ip.octets()[0] == 100 && (64..=127).contains(&ip.octets()[1])
}
IpAddr::V6(ip) => ip.is_loopback() || ip.is_unique_local() || ip.is_unicast_link_local(),
}
}
fn openai_subagent_from_headers(headers: &HeaderMap) -> Option<String> {
header_string(headers, OPENAI_SUBAGENT_HEADER).and_then(|value| {
if matches!(
-2
View File
@@ -8,7 +8,6 @@ use axum::{
use crate::runtime::state::AppState;
use super::{
chat::chat_completions,
models::{model_catalog, model_detail, model_info, models},
responses::{compact_responses, responses, review_responses},
};
@@ -19,7 +18,6 @@ pub fn router() -> Router<AppState> {
.route("/v1/responses", post(responses))
.route("/v1/responses/review", post(review_responses))
.route("/v1/responses/compact", post(compact_responses))
.route("/v1/chat/completions", post(chat_completions))
.route("/v1/models", get(models))
.route("/v1/models/catalog", get(model_catalog))
.route("/v1/models/{model_id}", get(model_detail))
+1 -1
View File
@@ -113,7 +113,7 @@ async fn build_application(
let task_coordinator = TaskCoordinator::start(&runtime_config, &services);
let trusted_proxies = TrustedProxyConfig::from_entries(&config.server.trusted_proxies)?;
let app_state = AppState { config, services };
let app_state = AppState { services };
Ok((
crate::http::router::router_with_trusted_proxies(trusted_proxies).with_state(app_state),
-12
View File
@@ -25,7 +25,6 @@ use crate::{
settings::{account_pool_options_from_config, RuntimeSettingsService},
},
proxy::dispatch::{
chat::ChatDispatchService,
responses::service::{ResponseDispatchService, ResponseDispatchServiceParts},
session_affinity::{RuntimeSessionAffinityService, SqliteSessionAffinityStore},
},
@@ -98,7 +97,6 @@ pub struct Services {
pub usage: StdArc<AdminUsageService>,
pub account_pool: StdArc<RuntimeAccountPoolService>,
pub token_refresh: StdArc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
pub chat: StdArc<ChatDispatchService>,
pub responses: StdArc<ResponseDispatchService>,
pub session_affinity: StdArc<RuntimeSessionAffinityService>,
pub codex: StdArc<CodexBackendClient>,
@@ -261,15 +259,6 @@ impl Services {
let cloudflare_recovery =
crate::proxy::dispatch::cloudflare::CloudflareRecovery::new(stores.cookies.clone());
let chat = StdArc::new(ChatDispatchService::new(
account_pool.clone(),
models.clone(),
codex.clone(),
usage_records.clone(),
token_refresh.clone(),
installation_id.clone(),
cloudflare_recovery.clone(),
));
let responses = StdArc::new(ResponseDispatchService::new(ResponseDispatchServiceParts {
account_pool: account_pool.clone(),
models: models.clone(),
@@ -294,7 +283,6 @@ impl Services {
usage,
account_pool,
token_refresh,
chat,
responses,
session_affinity,
codex,
-1
View File
@@ -6,7 +6,6 @@ use crate::runtime::services::Services;
/// HTTP handler 共享的运行时状态。
#[derive(Clone)]
pub struct AppState {
pub config: AppConfig,
pub services: Services,
}
+32 -36
View File
@@ -1,12 +1,14 @@
//! 通用过期数据清理任务骨架。
use std::{fmt::Display, future::Future, time::Duration};
use std::{fmt::Display, future::Future};
use chrono::{DateTime, Utc};
use tokio::time::interval;
use tracing::{debug, info, warn};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// 支持按时间清理过期记录的存储。
pub(crate) trait ExpiredCleanupStore: Send + Sync + 'static {
@@ -47,39 +49,12 @@ where
}
pub(crate) fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(
interval_secs = self.interval_secs,
"{}", self.messages.started
);
loop {
tokio::select! {
_ = ticker.tick() => {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "{}", self.messages.deleted);
}
Ok(_) => {
debug!("{}", self.messages.empty);
}
Err(error) => {
warn!(error = %error, "{}", self.messages.failed);
}
}
}
_ = shutdown_rx.recv() => {
info!("{}", self.messages.stopped);
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
let config = PeriodicTaskConfig::new(
self.interval_secs,
self.messages.started,
self.messages.stopped,
);
spawn_periodic_task(self, config)
}
pub(crate) async fn cleanup_once(&self) -> Result<u64, S::Error> {
@@ -90,3 +65,24 @@ where
self.store.delete_expired_at(now).await
}
}
impl<S> PeriodicTaskRunner for CleanupTask<S>
where
S: ExpiredCleanupStore,
{
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "{}", self.messages.deleted);
}
Ok(_) => {
debug!("{}", self.messages.empty);
}
Err(error) => {
warn!(error = %error, "{}", self.messages.failed);
}
}
})
}
}
+44 -15
View File
@@ -1,6 +1,6 @@
//! 后台任务协调器。
use std::{path::PathBuf, sync::Arc};
use std::{path::PathBuf, sync::Arc, time::Duration};
use tokio::task::JoinHandle;
@@ -120,22 +120,21 @@ impl TaskCoordinator {
/// 后台任务关闭句柄。
pub enum SchedulerHandle {
/// 通过 channel 发送关闭信号。
Channel(tokio::sync::mpsc::Sender<()>),
/// 直接持有 tokio 任务句柄。
JoinHandle(JoinHandle<()>),
Channel {
shutdown_tx: tokio::sync::mpsc::Sender<()>,
handle: JoinHandle<()>,
},
/// 关闭上游 WebSocket 连接池。
WebSocketPool(Arc<crate::upstream::transport::CodexWebSocketPool>),
}
impl SchedulerHandle {
/// 使用关闭 channel 构造句柄。
pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self {
Self::Channel(shutdown_tx)
}
/// 使用 `JoinHandle` 构造句柄。
pub(crate) fn from_join_handle(handle: JoinHandle<()>) -> Self {
Self::JoinHandle(handle)
pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>, handle: JoinHandle<()>) -> Self {
Self::Channel {
shutdown_tx,
handle,
}
}
/// 使用 WebSocket 连接池构造句柄。
@@ -146,11 +145,12 @@ impl SchedulerHandle {
/// 关闭任务。
pub async fn shutdown(self) {
match self {
Self::Channel(shutdown_tx) => {
Self::Channel {
shutdown_tx,
handle,
} => {
let _ = shutdown_tx.send(()).await;
}
Self::JoinHandle(handle) => {
handle.abort();
wait_for_task_shutdown(handle).await;
}
Self::WebSocketPool(pool) => {
pool.shutdown().await;
@@ -158,3 +158,32 @@ impl SchedulerHandle {
}
}
}
const TASK_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
async fn wait_for_task_shutdown(handle: JoinHandle<()>) {
let mut handle = handle;
let timeout = tokio::time::sleep(TASK_SHUTDOWN_TIMEOUT);
tokio::pin!(timeout);
tokio::select! {
result = &mut handle => {
log_task_shutdown_result(result);
}
() = &mut timeout => {
tracing::warn!(timeout_secs = TASK_SHUTDOWN_TIMEOUT.as_secs(), "等待后台任务关闭超时");
handle.abort();
log_task_shutdown_result(handle.await);
}
}
}
fn log_task_shutdown_result(result: Result<(), tokio::task::JoinError>) {
match result {
Ok(()) => {}
Err(error) if error.is_cancelled() => {}
Err(error) => {
tracing::warn!(error = %error, "后台任务关闭时异常退出");
}
}
}
@@ -2,9 +2,14 @@
use std::path::PathBuf;
use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker};
use tracing::warn;
use super::coordinator::SchedulerHandle;
use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker, APPCAST_POLL_INTERVAL};
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// Codex Desktop 官方 appcast 地址。
pub const CODEX_DESKTOP_APPCAST_URL: &str =
@@ -28,5 +33,35 @@ pub fn start_fingerprint_update_task(
current_version,
current_build,
);
SchedulerHandle::from_join_handle(update_checker.start_background_checker())
spawn_periodic_task(
FingerprintUpdateTask {
checker: update_checker,
first_tick: true,
},
PeriodicTaskConfig::new(
APPCAST_POLL_INTERVAL.as_secs(),
"fingerprint 后台版本检查器已启动",
"fingerprint 后台版本检查器已关闭",
),
)
}
struct FingerprintUpdateTask {
checker: UpdateChecker,
first_tick: bool,
}
impl PeriodicTaskRunner for FingerprintUpdateTask {
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
let first_tick = std::mem::take(&mut self.first_tick);
if let Err(error) = self.checker.check_and_apply_update().await {
if first_tick {
warn!(error = %error, "fingerprint 首次版本检查失败");
} else {
warn!(error = %error, "fingerprint 定期版本检查失败");
}
}
})
}
}
+1
View File
@@ -5,6 +5,7 @@ pub mod cookie_cleanup;
pub mod coordinator;
pub mod fingerprint_update;
pub mod model_refresh;
mod periodic;
pub mod quota_refresh;
pub mod session_affinity_cleanup;
pub mod session_cleanup;
+85 -81
View File
@@ -3,7 +3,8 @@
use std::{sync::Arc, time::Duration};
use chrono::Utc;
use tokio::time::{interval, sleep};
use tokio::sync::mpsc;
use tokio::time::sleep;
use tracing::{info, warn};
use crate::upstream::accounts::pool::RuntimeAccountPoolService;
@@ -11,7 +12,10 @@ use crate::upstream::models::service::{
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup},
};
/// 模型刷新任务接线器。
pub struct ModelRefreshTask {
@@ -54,85 +58,13 @@ impl ModelRefreshTask {
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
info!("模型刷新任务已启动");
sleep(Duration::from_millis(self.initial_delay_ms)).await;
let mut attempt = 0;
let mut has_fetched_once = false;
while attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return;
}
() = async {} => {}
}
match self.refresh_once().await {
Ok(result) => {
info!(
refreshed_plans = result.refreshed_plans,
model_count = result.model_count,
"首次模型刷新成功"
);
has_fetched_once = true;
break;
}
Err(error) => {
attempt += 1;
if matches!(error, ModelServiceError::NoAccounts) {
warn!("首次模型刷新跳过:没有可用账户,将进入周期重试");
break;
}
warn!(
attempt,
max_retries = self.max_retries,
error = ?error,
"首次模型刷新失败"
);
if attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
}
}
}
}
if !has_fetched_once {
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
}
let mut ticker = interval(Duration::from_secs(self.refresh_interval_secs));
ticker.tick().await;
loop {
tokio::select! {
_ = ticker.tick() => {
if let Err(error) = self.refresh_once().await {
warn!(error = ?error, "刷新模型列表失败");
}
}
_ = shutdown_rx.recv() => {
info!("模型刷新任务正在关闭");
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
let config = PeriodicTaskConfig::new(
self.refresh_interval_secs,
"模型刷新任务已启动",
"模型刷新任务正在关闭",
)
.wait_first_interval();
spawn_periodic_task(self, config)
}
/// 执行一次模型刷新并同步账号池模型计划约束。
@@ -167,3 +99,75 @@ impl ModelRefreshTask {
Ok(result)
}
}
impl PeriodicTaskRunner for ModelRefreshTask {
fn before_loop<'a>(
&'a mut self,
shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(async move {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
}
let mut attempt = 0;
let mut has_fetched_once = false;
while attempt < self.max_retries {
match self.refresh_once().await {
Ok(result) => {
info!(
refreshed_plans = result.refreshed_plans,
model_count = result.model_count,
"首次模型刷新成功"
);
has_fetched_once = true;
break;
}
Err(error) => {
attempt += 1;
if matches!(error, ModelServiceError::NoAccounts) {
warn!("首次模型刷新跳过:没有可用账户,将进入周期重试");
break;
}
warn!(
attempt,
max_retries = self.max_retries,
error = ?error,
"首次模型刷新失败"
);
if attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
}
}
}
}
if !has_fetched_once {
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
}
PeriodicTaskStartup::Continue
})
}
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
if let Err(error) = self.refresh_once().await {
warn!(error = ?error, "刷新模型列表失败");
}
})
}
}
+105
View File
@@ -0,0 +1,105 @@
//! 周期性后台任务运行器。
use std::{future::Future, pin::Pin, time::Duration};
use tokio::sync::mpsc;
use tokio::time::interval;
use tracing::info;
use super::coordinator::SchedulerHandle;
pub(crate) type TaskFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
/// 周期任务进入循环前的决策。
pub(crate) enum PeriodicTaskStartup {
/// 继续进入周期循环。
Continue,
/// 任务已自行结束,不再进入周期循环。
Stop,
}
/// 周期任务调度配置。
pub(crate) struct PeriodicTaskConfig {
interval_secs: u64,
started_message: &'static str,
stopped_message: &'static str,
run_immediately: bool,
}
impl PeriodicTaskConfig {
pub(crate) fn new(
interval_secs: u64,
started_message: &'static str,
stopped_message: &'static str,
) -> Self {
Self {
interval_secs,
started_message,
stopped_message,
run_immediately: true,
}
}
pub(crate) fn wait_first_interval(mut self) -> Self {
self.run_immediately = false;
self
}
}
/// 周期任务行为。
pub(crate) trait PeriodicTaskRunner: Send + 'static {
fn before_loop<'a>(
&'a mut self,
_shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(std::future::ready(PeriodicTaskStartup::Continue))
}
fn tick(&mut self) -> TaskFuture<'_, ()>;
fn shutdown(&mut self) -> TaskFuture<'_, ()> {
Box::pin(std::future::ready(()))
}
}
/// 启动周期任务,并统一处理关闭信号与 JoinHandle。
pub(crate) fn spawn_periodic_task<T>(mut task: T, config: PeriodicTaskConfig) -> SchedulerHandle
where
T: PeriodicTaskRunner,
{
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
let handle = tokio::spawn(async move {
info!(
interval_secs = config.interval_secs,
"{}", config.started_message
);
if matches!(
task.before_loop(&mut shutdown_rx).await,
PeriodicTaskStartup::Stop
) {
return;
}
let mut ticker = interval(Duration::from_secs(config.interval_secs));
if !config.run_immediately {
ticker.tick().await;
}
loop {
tokio::select! {
_ = ticker.tick() => {
task.tick().await;
}
_ = shutdown_rx.recv() => {
task.shutdown().await;
info!("{}", config.stopped_message);
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx, handle)
}
+38 -36
View File
@@ -1,8 +1,8 @@
//! 配额刷新后台任务接线器。
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{collections::HashMap, sync::Arc};
use tokio::time::interval;
use tokio::time::Instant;
use tracing::{debug, info, warn};
use crate::{
@@ -15,12 +15,16 @@ use crate::{
upstream::transport::CodexBackendClient,
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// 主动配额刷新后台任务。
pub struct QuotaRefreshTask {
service: RuntimeQuotaRefreshService,
interval_secs: u64,
last_refreshed: HashMap<String, Instant>,
}
impl QuotaRefreshTask {
@@ -38,6 +42,7 @@ impl QuotaRefreshTask {
min_refresh_interval_secs,
),
interval_secs,
last_refreshed: HashMap::new(),
}
}
@@ -61,42 +66,39 @@ impl QuotaRefreshTask {
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
let config = PeriodicTaskConfig::new(
self.interval_secs,
"quota 刷新任务已启动",
"quota 刷新任务已关闭",
);
spawn_periodic_task(self, config)
}
}
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
let mut last_refreshed = HashMap::new();
info!(interval_secs = self.interval_secs, "quota 刷新任务已启动");
loop {
tokio::select! {
_ = ticker.tick() => {
match self.service.refresh_locked_accounts(&mut last_refreshed).await {
Ok(summary) if summary.refreshed > 0 => {
log_refreshed_summary(summary);
}
Ok(summary) => {
debug!(
candidates = summary.candidates,
skipped_recent = summary.skipped_recent,
failed = summary.failed,
"没有需要刷新的 quota 锁定或待验证账号"
);
}
Err(error) => {
warn!(error = %error, "quota 刷新任务失败");
}
}
}
_ = shutdown_rx.recv() => {
info!("quota 刷新任务已关闭");
break;
}
impl PeriodicTaskRunner for QuotaRefreshTask {
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self
.service
.refresh_locked_accounts(&mut self.last_refreshed)
.await
{
Ok(summary) if summary.refreshed > 0 => {
log_refreshed_summary(summary);
}
Ok(summary) => {
debug!(
candidates = summary.candidates,
skipped_recent = summary.skipped_recent,
failed = summary.failed,
"没有需要刷新的 quota 锁定或待验证账号"
);
}
Err(error) => {
warn!(error = %error, "quota 刷新任务失败");
}
}
});
SchedulerHandle::new(shutdown_tx)
})
}
}
+43 -41
View File
@@ -1,8 +1,5 @@
//! 令牌刷新后台任务接线器。
use std::time::Duration;
use tokio::time::interval;
use tracing::{debug, info, warn};
use crate::upstream::accounts::{
@@ -12,7 +9,10 @@ use crate::upstream::accounts::{
},
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
const DEFAULT_INTERVAL_SECS: u64 = 60;
@@ -53,46 +53,48 @@ where
/// 启动后台任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
let config = PeriodicTaskConfig::new(
self.interval_secs,
"token 刷新任务已启动",
"token 刷新任务已关闭",
);
spawn_periodic_task(self, config)
}
}
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(interval_secs = self.interval_secs, "token 刷新任务已启动");
loop {
tokio::select! {
_ = ticker.tick() => {
match self.service.schedule_account_timers_once().await {
Ok(summary) if summary.changed() > 0 => {
info!(
scheduled = summary.scheduled,
immediate = summary.immediate,
recovery_scheduled = summary.recovery_scheduled,
replaced = summary.replaced,
"token 刷新定时器已调度"
);
}
Ok(summary) => {
debug!(
scanned = summary.scanned,
skipped = summary.skipped,
"没有需要调度刷新的 token"
);
}
Err(error) => {
warn!(error = %error, "token 刷新定时器调度失败");
}
}
}
_ = shutdown_rx.recv() => {
self.service.clear_scheduled_timers().await;
info!("token 刷新任务已关闭");
break;
}
impl<C> PeriodicTaskRunner for TokenRefreshTask<C>
where
C: TokenRefresher,
{
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self.service.schedule_account_timers_once().await {
Ok(summary) if summary.changed() > 0 => {
info!(
scheduled = summary.scheduled,
immediate = summary.immediate,
recovery_scheduled = summary.recovery_scheduled,
replaced = summary.replaced,
"token 刷新定时器已调度"
);
}
Ok(summary) => {
debug!(
scanned = summary.scanned,
skipped = summary.skipped,
"没有需要调度刷新的 token"
);
}
Err(error) => {
warn!(error = %error, "token 刷新定时器调度失败");
}
}
});
})
}
SchedulerHandle::new(shutdown_tx)
fn shutdown(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
self.service.clear_scheduled_timers().await;
})
}
}
-20
View File
@@ -8,26 +8,6 @@ use thiserror::Error;
use tokio::sync::RwLock;
use uuid::Uuid;
/// 单个 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)
}
const DEFAULT_COOKIE_DOMAIN: &str = "chatgpt.com";
const CAPTURABLE_COOKIES: &[&str] = &["cf_clearance"];
+25 -31
View File
@@ -141,61 +141,55 @@ fn explicit_account_import_source(
fn parse_cpr_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::Cpr,
entries: parse_cpr_entries(payload)?,
entries: parse_account_entries(
payload,
cpr_account_import_entry_from_value,
Some(ACCOUNT_IMPORT_CONTAINER_KEYS),
)?,
})
}
fn parse_cpr_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
ensure_account_import_keys(payload, ACCOUNT_IMPORT_CONTAINER_KEYS)?;
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, cpr_account_import_entry_from_value);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, cpr_account_import_entry_from_value);
}
cpr_account_import_entry_from_value(payload).map(option_to_vec)
}
fn parse_sub2api_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::Sub2Api,
entries: parse_sub2api_entries(payload)?,
entries: parse_account_entries(payload, sub2api_account_import_entry_from_value, None)?,
})
}
fn parse_sub2api_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, sub2api_account_import_entry_from_value);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, sub2api_account_import_entry_from_value);
}
sub2api_account_import_entry_from_value(payload).map(option_to_vec)
}
fn parse_cli_proxy_api_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::CliProxyApi,
entries: parse_cli_proxy_api_entries(payload)?,
entries: parse_account_entries(
payload,
cli_proxy_api_account_import_entry_from_value,
None,
)?,
})
}
fn parse_cli_proxy_api_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
type AccountImportEntryParser = fn(&Value) -> Result<Option<AccountImportEntry>, &'static str>;
fn parse_account_entries(
payload: &Value,
parser: AccountImportEntryParser,
container_keys: Option<&[&str]>,
) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
if let Some(container_keys) = container_keys {
ensure_account_import_keys(payload, container_keys)?;
}
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, cli_proxy_api_account_import_entry_from_value);
return parse_entries(accounts, parser);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, cli_proxy_api_account_import_entry_from_value);
return parse_entries(accounts, parser);
}
cli_proxy_api_account_import_entry_from_value(payload).map(option_to_vec)
parser(payload).map(option_to_vec)
}
fn parse_entries(
accounts: &[Value],
parser: fn(&Value) -> Result<Option<AccountImportEntry>, &'static str>,
parser: AccountImportEntryParser,
) -> Result<Vec<AccountImportEntry>, &'static str> {
let mut entries = Vec::new();
for account in accounts {
+15 -8
View File
@@ -20,16 +20,23 @@ pub enum AccountStatus {
Banned,
}
impl AccountStatus {
/// 返回账号状态的持久化/API 字符串。
pub fn as_str(self) -> &'static str {
match self {
Self::Active => "active",
Self::Expired => "expired",
Self::QuotaExhausted => "quota_exhausted",
Self::Refreshing => "refreshing",
Self::Disabled => "disabled",
Self::Banned => "banned",
}
}
}
impl std::fmt::Display for AccountStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Active => write!(f, "active"),
Self::Expired => write!(f, "expired"),
Self::QuotaExhausted => write!(f, "quota_exhausted"),
Self::Refreshing => write!(f, "refreshing"),
Self::Disabled => write!(f, "disabled"),
Self::Banned => write!(f, "banned"),
}
f.write_str(self.as_str())
}
}
+129 -47
View File
@@ -24,18 +24,6 @@ use crate::upstream::protocol::events::{
parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage,
};
/// Token usage statistics.
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct TokenUsage {
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub reasoning_tokens: i64,
pub total_tokens: i64,
pub image_input_tokens: i64,
pub image_output_tokens: i64,
}
/// 账号轮转策略。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RotationStrategy {
@@ -154,13 +142,19 @@ pub struct DistinctPlanAccount {
#[derive(Debug)]
struct AccountAcquireWithStatusRefresh {
acquired: Option<AcquiredAccount>,
expired_account_ids: Vec<String>,
refreshed_accounts: Vec<RefreshedAccountState>,
}
#[derive(Debug)]
struct DistinctPlanAccountsWithStatusRefresh {
accounts: Vec<DistinctPlanAccount>,
expired_account_ids: Vec<String>,
refreshed_accounts: Vec<RefreshedAccountState>,
}
#[derive(Debug)]
struct AccountCapacitySummaryWithStatusRefresh {
summary: AccountCapacitySummary,
refreshed_accounts: Vec<RefreshedAccountState>,
}
/// 账号池容量摘要。
@@ -176,6 +170,71 @@ pub struct AccountCapacitySummary {
pub available_slots: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct RuntimeAccountStateSnapshot {
status: AccountStatus,
quota_limit_reached: bool,
quota_verify_required: bool,
quota_cooldown_until: Option<DateTime<Utc>>,
cloudflare_cooldown_until: Option<DateTime<Utc>>,
window_request_count: u64,
window_input_tokens: u64,
window_output_tokens: u64,
window_cached_tokens: u64,
window_image_input_tokens: u64,
window_image_output_tokens: u64,
window_image_request_count: u64,
window_image_request_failed_count: u64,
window_started_at: Option<DateTime<Utc>>,
window_reset_at: Option<DateTime<Utc>>,
limit_window_seconds: Option<u64>,
}
impl From<&Account> for RuntimeAccountStateSnapshot {
fn from(account: &Account) -> Self {
Self {
status: account.status,
quota_limit_reached: account.quota_limit_reached,
quota_verify_required: account.quota_verify_required,
quota_cooldown_until: account.quota_cooldown_until,
cloudflare_cooldown_until: account.cloudflare_cooldown_until,
window_request_count: account.window_request_count,
window_input_tokens: account.window_input_tokens,
window_output_tokens: account.window_output_tokens,
window_cached_tokens: account.window_cached_tokens,
window_image_input_tokens: account.window_image_input_tokens,
window_image_output_tokens: account.window_image_output_tokens,
window_image_request_count: account.window_image_request_count,
window_image_request_failed_count: account.window_image_request_failed_count,
window_started_at: account.window_started_at,
window_reset_at: account.window_reset_at,
limit_window_seconds: account.limit_window_seconds,
}
}
}
impl RuntimeAccountStateSnapshot {
fn usage_window_changed(&self, other: &Self) -> bool {
self.window_request_count != other.window_request_count
|| self.window_input_tokens != other.window_input_tokens
|| self.window_output_tokens != other.window_output_tokens
|| self.window_cached_tokens != other.window_cached_tokens
|| self.window_image_input_tokens != other.window_image_input_tokens
|| self.window_image_output_tokens != other.window_image_output_tokens
|| self.window_image_request_count != other.window_image_request_count
|| self.window_image_request_failed_count != other.window_image_request_failed_count
|| self.window_started_at != other.window_started_at
|| self.window_reset_at != other.window_reset_at
|| self.limit_window_seconds != other.limit_window_seconds
}
}
#[derive(Debug, Clone)]
struct RefreshedAccountState {
account: Account,
sync_usage_window: bool,
}
/// 运行时窗口用量增量。
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AccountWindowUsageDelta {
@@ -419,7 +478,7 @@ impl AccountPool {
request: &AccountAcquireRequest,
) -> AccountAcquireWithStatusRefresh {
self.cleanup_stale_slots(request.now);
let expired_account_ids = self.refresh_account_statuses(request.now);
let refreshed_accounts = self.refresh_account_statuses(request.now);
let candidates = self.candidates(request);
let selected = if let Some(preferred_account_id) = &request.preferred_account_id {
let preferred = candidates
@@ -458,7 +517,7 @@ impl AccountPool {
let Some((selected, previous_slot_at, selection_source)) = selected else {
return AccountAcquireWithStatusRefresh {
acquired: None,
expired_account_ids,
refreshed_accounts,
};
};
tracing::info!(
@@ -485,13 +544,15 @@ impl AccountPool {
account: selected,
previous_slot_at,
}),
expired_account_ids,
refreshed_accounts,
}
}
fn refresh_account_statuses(&mut self, now: DateTime<Utc>) -> Vec<String> {
fn refresh_account_statuses(&mut self, now: DateTime<Utc>) -> Vec<RefreshedAccountState> {
let mut refreshed_accounts = Vec::new();
let mut expired_account_ids = Vec::new();
for (account_id, account) in &mut self.accounts {
let before = RuntimeAccountStateSnapshot::from(&*account);
if account.status == AccountStatus::Active && access_token_expired(account, now) {
account.status = AccountStatus::Expired;
expired_account_ids.push(account_id.clone());
@@ -499,11 +560,18 @@ impl AccountPool {
refresh_quota_window(account, now);
refresh_cloudflare_cooldown(account, now);
}
let after = RuntimeAccountStateSnapshot::from(&*account);
if before != after {
refreshed_accounts.push(RefreshedAccountState {
account: account.clone(),
sync_usage_window: before.usage_window_changed(&after),
});
}
}
for account_id in &expired_account_ids {
self.slots.remove(account_id);
}
expired_account_ids
refreshed_accounts
}
/// 释放指定账号的一个在途槽位。
@@ -527,8 +595,15 @@ impl AccountPool {
/// 计算账号池容量摘要。
pub fn capacity_summary(&mut self, now: DateTime<Utc>) -> AccountCapacitySummary {
self.capacity_summary_with_status_refresh(now).summary
}
fn capacity_summary_with_status_refresh(
&mut self,
now: DateTime<Utc>,
) -> AccountCapacitySummaryWithStatusRefresh {
self.cleanup_stale_slots(now);
self.refresh_account_statuses(now);
let refreshed_accounts = self.refresh_account_statuses(now);
let is_capacity_account = |account: &Account| {
account.status == AccountStatus::Active
&& AccountService::quota_available_at(account, now, self.options.skip_quota_limited)
@@ -550,11 +625,14 @@ impl AccountPool {
.map(|(_, slots)| slots.len().min(self.options.max_concurrent_per_account))
.sum();
AccountCapacitySummary {
max_concurrent_per_account: self.options.max_concurrent_per_account,
total_slots,
used_slots,
available_slots: total_slots.saturating_sub(used_slots),
AccountCapacitySummaryWithStatusRefresh {
summary: AccountCapacitySummary {
max_concurrent_per_account: self.options.max_concurrent_per_account,
total_slots,
used_slots,
available_slots: total_slots.saturating_sub(used_slots),
},
refreshed_accounts,
}
}
@@ -569,7 +647,7 @@ impl AccountPool {
now: DateTime<Utc>,
) -> DistinctPlanAccountsWithStatusRefresh {
self.cleanup_stale_slots(now);
let expired_account_ids = self.refresh_account_statuses(now);
let refreshed_accounts = self.refresh_account_statuses(now);
let mut by_plan = IndexMap::<String, Vec<Account>>::new();
for account in self.accounts.values() {
@@ -615,7 +693,7 @@ impl AccountPool {
DistinctPlanAccountsWithStatusRefresh {
accounts,
expired_account_ids,
refreshed_accounts,
}
}
@@ -874,7 +952,14 @@ impl RuntimeAccountPoolService {
/// 读取账号池容量摘要。
pub async fn capacity_summary(&self, now: DateTime<Utc>) -> AccountCapacitySummary {
self.pool.lock().await.capacity_summary(now)
let refresh = self
.pool
.lock()
.await
.capacity_summary_with_status_refresh(now);
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.summary
}
/// 使用当前时间读取账号池容量摘要。
@@ -964,7 +1049,7 @@ impl RuntimeAccountPoolService {
/// 按请求上下文获取可用账号。
pub async fn acquire_with(&self, request: &AccountAcquireRequest) -> Option<AcquiredAccount> {
let refresh = self.pool.lock().await.acquire_with_status_refresh(request);
self.persist_expired_statuses(refresh.expired_account_ids)
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.acquired
}
@@ -976,22 +1061,22 @@ impl RuntimeAccountPoolService {
.lock()
.await
.distinct_plan_accounts_with_status_refresh(now);
self.persist_expired_statuses(refresh.expired_account_ids)
self.persist_runtime_account_states(refresh.refreshed_accounts)
.await;
refresh.accounts
}
async fn persist_expired_statuses(&self, account_ids: Vec<String>) {
for account_id in account_ids {
async fn persist_runtime_account_states(&self, refreshed_states: Vec<RefreshedAccountState>) {
for refreshed in refreshed_states {
if let Err(error) = self
.store
.set_status(&account_id, AccountStatus::Expired)
.sync_runtime_account_state(&refreshed.account, refreshed.sync_usage_window)
.await
{
tracing::warn!(
account_id,
account_id = %refreshed.account.id,
error = %error,
"failed to persist runtime-expired account status"
"failed to persist refreshed runtime account state"
);
}
}
@@ -1191,6 +1276,9 @@ impl RuntimeAccountPoolService {
let limit_reached = quota_snapshot_limit_reached(quota);
let reset_at = quota_snapshot_reset_at(quota);
let cooldown_until = limit_reached.then_some(reset_at).flatten();
let limit_window_seconds = reset_at.and_then(|_| {
crate::upstream::accounts::quota::quota_snapshot_limit_window_seconds(quota)
});
let quota_json = quota.to_string();
let persisted = match self
.store
@@ -1207,15 +1295,8 @@ impl RuntimeAccountPoolService {
false
}
};
let in_memory =
self.pool
.lock()
.await
.apply_quota_state(account_id, limit_reached, cooldown_until);
if let Some(reset_at) = reset_at {
let limit_window_seconds =
crate::upstream::accounts::quota::quota_snapshot_limit_window_seconds(quota);
if let Err(error) = self
.store
.sync_rate_limit_window(account_id, reset_at, limit_window_seconds)
@@ -1227,11 +1308,12 @@ impl RuntimeAccountPoolService {
"failed to persist verified quota window"
);
}
self.pool.lock().await.sync_rate_limit_window(
account_id,
reset_at,
limit_window_seconds,
);
}
let mut pool = self.pool.lock().await;
let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until);
if let Some(reset_at) = reset_at {
pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds);
}
persisted || in_memory
+12 -2
View File
@@ -132,6 +132,8 @@ fn window_from_rate_limit(window: Option<&Value>) -> Value {
"reset_at": window.get("reset_at").and_then(positive_i64),
"window_minutes": window_minutes,
"limit_reached": used_percent >= 100.0,
"used": first_present(window, &["used_tokens", "used", "used_credits", "usage", "consumed_tokens", "consumed"]),
"limit": first_present(window, &["limit_tokens", "limit", "limit_credits", "quota", "total_tokens", "total"]),
})
}
@@ -201,8 +203,10 @@ fn monthly_limit_from_snapshots(snapshots: &[Value]) -> Option<Value> {
"reset_at": bucket.get("reset_at").cloned().unwrap_or(Value::Null),
"window_minutes": bucket.get("window_minutes").cloned().unwrap_or(Value::Null),
"limit_reached": bucket.get("limit_reached").cloned().unwrap_or(Value::Bool(false)),
"used_credits": Value::Null,
"limit_credits": Value::Null,
"used": bucket.get("used").cloned().unwrap_or(Value::Null),
"limit": bucket.get("limit").cloned().unwrap_or(Value::Null),
"used_credits": bucket.get("used_credits").cloned().unwrap_or(Value::Null),
"limit_credits": bucket.get("limit_credits").cloned().unwrap_or(Value::Null),
}));
}
}
@@ -214,6 +218,12 @@ fn remaining_percent(used_percent: f64) -> i64 {
(100.0 - used_percent.clamp(0.0, 100.0)).round() as i64
}
fn first_present(value: &Value, keys: &[&str]) -> Value {
keys.iter()
.find_map(|key| value.get(*key).filter(|value| !value.is_null()).cloned())
.unwrap_or(Value::Null)
}
fn rate_limit_blocked(rate_limit: &Value) -> bool {
rate_limit
.get("limit_reached")
+78 -47
View File
@@ -21,7 +21,7 @@ mod rows;
use queries::*;
use rows::{
count_account_metadata, get_pool_account, list_pool_accounts, map_account_store_error,
metadata_from_row, model_usage_from_row, optional_positive_i64_to_u64, optional_update_value,
metadata_from_row, optional_positive_i64_to_u64, optional_update_value,
push_account_metadata_search, quota_plan_type, quota_snapshot_from_row, sqlite_usage_delta,
status_to_db, stored_account_from_row, to_page, u64_to_i64_saturating,
};
@@ -196,27 +196,6 @@ impl AccountMetadataUpdate {
}
}
/// 账号模型维度用量记录。
#[derive(Debug, Clone)]
pub struct AccountModelUsageRecord {
/// 账号 ID。
pub account_id: String,
/// 模型 ID。
pub model: String,
/// 历史请求总数。
pub request_count: i64,
/// 历史错误数。
pub error_count: i64,
/// 累计输入 token。
pub input_tokens: i64,
/// 累计输出 token。
pub output_tokens: i64,
/// 累计缓存 token。
pub cached_tokens: i64,
/// 最近使用时间。
pub last_used_at: Option<DateTime<Utc>>,
}
/// 账号配额快照。
#[derive(Debug, Clone)]
pub struct AccountQuotaSnapshot {
@@ -322,13 +301,6 @@ pub trait AccountStore: Send + Sync + 'static {
/// 读取账号当前配额 JSON。
async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult<Option<String>>;
/// 更新账号当前配额 JSON。
async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> AccountStoreResult<bool>;
/// 应用已经验证过的账号配额快照。
async fn apply_quota_snapshot(
&self,
@@ -338,6 +310,13 @@ pub trait AccountStore: Send + Sync + 'static {
cooldown_until: Option<DateTime<Utc>>,
) -> AccountStoreResult<bool>;
/// 同步运行时自然刷新出来的账号状态。
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool>;
/// 同步账号当前 rate-limit 统计窗口。
async fn sync_rate_limit_window(
&self,
@@ -711,6 +690,66 @@ impl SqliteAccountStore {
Ok(result.rows_affected() > 0)
}
/// 同步运行时自然刷新出来的账号状态。
pub async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> SqliteAccountStoreResult<bool> {
let now = Utc::now().to_rfc3339();
let quota_limit_reached = if account.quota_limit_reached { 1 } else { 0 };
let quota_verify_required = if account.quota_verify_required { 1 } else { 0 };
let quota_cooldown_until = account.quota_cooldown_until.map(|dt| dt.to_rfc3339());
let cloudflare_cooldown_until = account.cloudflare_cooldown_until.map(|dt| dt.to_rfc3339());
let mut tx = self.pool.begin().await?;
let result = sqlx::query(SYNC_RUNTIME_ACCOUNT_STATE_SQL)
.bind(status_to_db(account.status))
.bind(quota_limit_reached)
.bind(&now)
.bind(quota_limit_reached)
.bind(quota_verify_required)
.bind(&now)
.bind(quota_verify_required)
.bind(&quota_cooldown_until)
.bind(&now)
.bind(&quota_cooldown_until)
.bind(&cloudflare_cooldown_until)
.bind(&now)
.bind(&cloudflare_cooldown_until)
.bind(&now)
.bind(&account.id)
.execute(&mut *tx)
.await?;
if result.rows_affected() == 0 {
tx.commit().await?;
return Ok(false);
}
if sync_usage_window {
sqlx::query(SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL)
.bind(&account.id)
.bind(u64_to_i64_saturating(account.window_request_count))
.bind(u64_to_i64_saturating(account.window_input_tokens))
.bind(u64_to_i64_saturating(account.window_output_tokens))
.bind(u64_to_i64_saturating(account.window_cached_tokens))
.bind(u64_to_i64_saturating(account.window_image_input_tokens))
.bind(u64_to_i64_saturating(account.window_image_output_tokens))
.bind(u64_to_i64_saturating(account.window_image_request_count))
.bind(u64_to_i64_saturating(
account.window_image_request_failed_count,
))
.bind(account.window_started_at.map(|dt| dt.to_rfc3339()))
.bind(account.window_reset_at.map(|dt| dt.to_rfc3339()))
.bind(account.limit_window_seconds.map(u64_to_i64_saturating))
.bind(&now)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
Ok(true)
}
/// 标记账号 Cloudflare 冷却期。
pub async fn set_cloudflare_cooldown_until(
&self,
@@ -804,14 +843,6 @@ impl SqliteAccountStore {
Ok(())
}
/// 列出模型维度用量。
pub async fn list_model_usage(&self) -> SqliteAccountStoreResult<Vec<AccountModelUsageRecord>> {
let rows = sqlx::query(LIST_MODEL_USAGE_SQL)
.fetch_all(&self.pool)
.await?;
rows.iter().map(model_usage_from_row).collect()
}
/// 读取配额快照列表。
pub async fn list_quota_snapshots(
&self,
@@ -1054,16 +1085,6 @@ impl AccountStore for SqliteAccountStore {
.map_err(|error| map_account_store_error(&error))
}
async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> AccountStoreResult<bool> {
SqliteAccountStore::update_quota_json(self, account_id, quota_json)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn apply_quota_snapshot(
&self,
account_id: &str,
@@ -1082,6 +1103,16 @@ impl AccountStore for SqliteAccountStore {
.map_err(|error| map_account_store_error(&error))
}
async fn sync_runtime_account_state(
&self,
account: &Account,
sync_usage_window: bool,
) -> AccountStoreResult<bool> {
SqliteAccountStore::sync_runtime_account_state(self, account, sync_usage_window)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn sync_rate_limit_window(
&self,
account_id: &str,
+57 -13
View File
@@ -327,19 +327,6 @@ on conflict(account_id, model) do update set
cached_tokens = cached_tokens + excluded.cached_tokens,
last_used_at = excluded.last_used_at";
pub(super) const LIST_MODEL_USAGE_SQL: &str = r"
select
account_id,
model,
request_count,
error_count,
input_tokens,
output_tokens,
cached_tokens,
last_used_at
from account_model_usage
order by account_id asc, request_count desc, last_used_at desc, model asc";
pub(super) const LIST_QUOTA_SNAPSHOTS_SQL: &str = r"
select
id,
@@ -432,6 +419,63 @@ set
updated_at = ?
where id = ?";
pub(super) const SYNC_RUNTIME_ACCOUNT_STATE_SQL: &str = r"
update accounts
set
status = case
when status in ('disabled', 'banned', 'refreshing') then status
else ?
end,
quota_limit_reached = case
when ? = 0 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_limit_reached
else ?
end,
quota_verify_required = case
when ? = 1 and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_verify_required
else ?
end,
quota_cooldown_until = case
when ? is null and quota_cooldown_until is not null and quota_cooldown_until > ? then quota_cooldown_until
else ?
end,
cloudflare_cooldown_until = case
when ? is null and cloudflare_cooldown_until is not null and cloudflare_cooldown_until > ? then cloudflare_cooldown_until
else ?
end,
updated_at = ?
where id = ?";
pub(super) const SYNC_RUNTIME_ACCOUNT_USAGE_WINDOW_SQL: &str = r"
insert into account_usage (
account_id,
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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
on conflict(account_id) do update set
window_request_count = excluded.window_request_count,
window_input_tokens = excluded.window_input_tokens,
window_output_tokens = excluded.window_output_tokens,
window_cached_tokens = excluded.window_cached_tokens,
window_image_input_tokens = excluded.window_image_input_tokens,
window_image_output_tokens = excluded.window_image_output_tokens,
window_image_request_count = excluded.window_image_request_count,
window_image_request_failed_count = excluded.window_image_request_failed_count,
window_started_at = excluded.window_started_at,
window_reset_at = excluded.window_reset_at,
limit_window_seconds = excluded.limit_window_seconds
where account_usage.window_reset_at is null
or account_usage.window_reset_at <= excluded.window_reset_at
or account_usage.window_reset_at <= ?";
pub(super) const SET_CLOUDFLARE_COOLDOWN_UNTIL_SQL: &str = r"
update accounts
set
+3 -28
View File
@@ -17,9 +17,8 @@ use crate::{
use super::{
queries::{GET_POOL_ACCOUNT_SQL, LIST_POOL_ACCOUNTS_SQL},
AccountModelUsageRecord, AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore,
SqliteAccountStoreError, SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata,
UsageDelta,
AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore, SqliteAccountStoreError,
SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata, UsageDelta,
};
pub(super) async fn list_pool_accounts(
@@ -230,23 +229,6 @@ pub(super) fn quota_plan_type(quota_json: &str) -> Option<String> {
.map(ToString::to_string)
}
pub(super) fn model_usage_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountModelUsageRecord> {
Ok(AccountModelUsageRecord {
account_id: row.get("account_id"),
model: row.get("model"),
request_count: row.get("request_count"),
error_count: row.get("error_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
last_used_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("last_used_at").as_deref(),
)?,
})
}
pub(super) fn quota_snapshot_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountQuotaSnapshot> {
@@ -261,14 +243,7 @@ pub(super) fn quota_snapshot_from_row(
}
pub(super) fn status_to_db(status: AccountStatus) -> &'static str {
match status {
AccountStatus::Active => "active",
AccountStatus::Expired => "expired",
AccountStatus::QuotaExhausted => "quota_exhausted",
AccountStatus::Refreshing => "refreshing",
AccountStatus::Disabled => "disabled",
AccountStatus::Banned => "banned",
}
status.as_str()
}
pub(super) fn optional_update_value(value: &Option<Option<String>>) -> Option<&str> {
@@ -259,9 +259,6 @@ pub enum RefreshFailure {
/// 刷新令牌无效或已过期。
#[error("refresh token is invalid or expired")]
InvalidGrant,
/// 账号配额耗尽。
#[error("account quota is exhausted")]
QuotaExhausted,
/// 账号被上游封禁。
#[error("account is banned")]
Banned,
@@ -350,7 +347,6 @@ where
match self.client.refresh(refresh_token).await {
Ok(token_pair) => Ok(apply_token_pair(account, token_pair)),
Err(RefreshFailure::QuotaExhausted) => Err(RefreshError::Transport),
Err(RefreshFailure::RetryableTransport) => Err(RefreshError::RetryableTransport),
Err(RefreshFailure::Transport) => Err(RefreshError::Transport),
Err(error) => Ok(apply_refresh_failure(account, error)),
@@ -411,7 +407,6 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco
let mut updated = account.clone();
updated.status = match failure {
RefreshFailure::InvalidGrant => AccountStatus::Expired,
RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted,
RefreshFailure::Banned => AccountStatus::Banned,
RefreshFailure::RetryableTransport => AccountStatus::Active,
RefreshFailure::Transport => AccountStatus::Active,
+2 -25
View File
@@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize};
use sqlx::{Row, SqlitePool};
use thiserror::Error;
use tokio::sync::Mutex;
use tracing::{info, warn};
use tracing::info;
use uuid::Uuid;
/// 运行时当前指纹槽位。
@@ -17,7 +17,7 @@ pub const CURRENT_FINGERPRINT_ID: &str = "current";
const AUTO_UPDATE_SOURCE: &str = "auto_update";
const CONFIG_SEED_SOURCE: &str = "config_seed";
const APPCAST_TIMEOUT: Duration = Duration::from_secs(30);
const POLL_INTERVAL: Duration = Duration::from_hours(72);
pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_hours(72);
type AppcastFields = (Option<String>, Option<String>, Option<String>);
@@ -442,29 +442,6 @@ impl UpdateChecker {
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
+8 -89
View File
@@ -9,19 +9,6 @@ use super::{
snapshot::ModelPlanSnapshot,
};
/// 解析模型名后得到的标准化结果。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ParsedModelName {
/// 最终模型 ID。
pub model_id: String,
/// 推理强度后缀。
pub reasoning_effort: Option<String>,
/// 服务层级后缀。
pub service_tier: Option<String>,
}
const SERVICE_TIER_SUFFIXES: [&str; 2] = ["fast", "flex"];
const REASONING_EFFORT_SUFFIXES: [&str; 6] = ["none", "minimal", "low", "medium", "high", "xhigh"];
const BUILTIN_CODEX_MODELS: [(&str, &str, bool); 8] = [
("gpt-5.5", "GPT-5.5", true),
("gpt-5.4", "GPT-5.4", false),
@@ -170,10 +157,9 @@ impl ModelCatalog {
.cloned()
}
/// 按对外模型名查询模型信息,支持别名和已知后缀。
/// 按对外模型名查询模型信息,支持别名映射。
pub fn model_info_for_name(&self, input: &str) -> Option<CodexModelInfo> {
let parsed = self.parse_model_name(input);
self.model_info(&parsed.model_id)
self.model_info(&self.resolve_model_id(input))
}
/// 返回 model -> plans allowlist。
@@ -192,54 +178,14 @@ impl ModelCatalog {
if trimmed.is_empty() {
return false;
}
if self.model_aliases.contains_key(trimmed) || self.model_info(trimmed).is_some() {
return true;
}
let stripped = strip_known_model_suffixes(trimmed);
if stripped.model_name == trimmed
|| (stripped.reasoning_effort.is_none() && stripped.service_tier.is_none())
{
return false;
}
self.model_aliases.contains_key(&stripped.model_name)
|| self.model_info(&stripped.model_name).is_some()
self.model_aliases.contains_key(trimmed) || self.model_info(trimmed).is_some()
}
/// 解析外部传入的模型名,提取别名、推理强度和服务层级后缀。
pub fn parse_model_name(&self, input: &str) -> ParsedModelName {
let trimmed = input.trim();
if self.model_aliases.contains_key(trimmed) || self.model_info(trimmed).is_some() {
return ParsedModelName {
model_id: self.resolve_model_id(trimmed),
reasoning_effort: None,
service_tier: None,
};
}
let stripped = strip_known_model_suffixes(trimmed);
ParsedModelName {
model_id: self.resolve_model_id(&stripped.model_name),
reasoning_effort: stripped.reasoning_effort,
service_tier: stripped.service_tier,
}
}
/// 将标准化模型名重新拼成展示名。
pub fn build_display_model_name(parsed: &ParsedModelName) -> String {
let mut name = parsed.model_id.clone();
if let Some(reasoning_effort) = &parsed.reasoning_effort {
name.push('-');
name.push_str(reasoning_effort);
}
if let Some(service_tier) = &parsed.service_tier {
name.push('-');
name.push_str(service_tier);
}
name
}
fn resolve_model_id(&self, input: &str) -> String {
/// 沿 alias 链将外部模型名解析为真实 model ID。
///
/// 模型名只承载 model 身份;reasoning effort、service tier 等运行参数由请求
/// body 携带并原样透传上游。
pub fn resolve_model_id(&self, input: &str) -> String {
self.resolve_alias_chain(input)
}
@@ -260,33 +206,6 @@ impl ModelCatalog {
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct StrippedModelName {
model_name: String,
reasoning_effort: Option<String>,
service_tier: Option<String>,
}
fn strip_known_model_suffixes(input: &str) -> StrippedModelName {
let mut remaining = input.trim().to_string();
let service_tier = take_suffix(&mut remaining, &SERVICE_TIER_SUFFIXES);
let reasoning_effort = take_suffix(&mut remaining, &REASONING_EFFORT_SUFFIXES);
StrippedModelName {
model_name: remaining,
reasoning_effort,
service_tier,
}
}
fn take_suffix(remaining: &mut String, suffixes: &[&str]) -> Option<String> {
let suffix = suffixes
.iter()
.find(|suffix| remaining.ends_with(&format!("-{suffix}")))?;
let truncate_to = remaining.len() - suffix.len() - 1;
remaining.truncate(truncate_to);
Some((*suffix).to_string())
}
fn normalize_aliases(input: &BTreeMap<String, String>) -> BTreeMap<String, String> {
input
.iter()
+257 -189
View File
@@ -1,112 +1,71 @@
use std::{collections::BTreeSet, time::Instant};
use serde::{Deserialize, Serialize};
use serde::Serialize;
use serde_json::{json, Map, Value};
use crate::{
infra::time::elapsed_millis_i64,
upstream::{
models::catalog::ParsedModelName,
protocol::{
schema::reconvert_tuple_values,
sse::{parse_sse_events, SseError, SseEvent},
},
upstream::protocol::{
schema::reconvert_tuple_values,
sse::{parse_sse_events, SseError, SseEvent},
},
};
/// Codex Responses 上游请求体。
#[derive(Debug, Clone, Serialize, Deserialize)]
///
/// 发往上游的 Responses 请求。`body` 持有客户端原始 JSON object,逐字段(含顺序、
/// 含未知字段)透传上游,是上游请求体的唯一来源;`use_websocket`/`force_http_sse`
/// 仅用于本地传输选择,不写入 body。常用字段通过访问器方法读写。
/// 代理只做最小 patch(`stream`/`store` 由传输层控制,`input` 字符串做官方兼容转换,
/// implicit resume 时替换 `input`,模型后缀路由时写 `model`/`reasoning`/`service_tier`)。
///
/// 其余字段是代理控制状态,不进上游 body(原 `#[serde(skip)]` 字段)。
#[derive(Debug, Clone)]
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)]
/// 上游请求体(唯一真相源)。
body: Map<String, Value>,
/// 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)]
/// 代理侧识别的客户端 IP,仅用于管理端使用记录展示。
pub client_ip: Option<String>,
#[serde(skip)]
/// 客户端 User-Agent,仅用于管理端使用记录展示。
pub client_user_agent: Option<String>,
#[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 Serialize for CodexResponsesRequest {
/// 上游 body 序列化即原始 `body` map(HTTP SSE 直发;WebSocket 在外层前置 `type`)。
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.body.serialize(serializer)
}
}
/// Codex Responses 请求的上游传输决策。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodexTransport {
@@ -124,7 +83,7 @@ pub fn transport_for_request(request: &CodexResponsesRequest) -> CodexTransport
return CodexTransport::HttpSse;
}
if request.previous_response_id.is_some() {
if request.previous_response_id().is_some() {
return CodexTransport::WebSocketRequired;
}
@@ -567,112 +526,22 @@ fn failure_code(value: &Value) -> Option<String> {
.map(ToString::to_string)
}
/// 将模型名后缀和模型配置应用到 Codex Responses 上游请求。
pub fn apply_response_model_options(
request: &mut CodexResponsesRequest,
parsed_model: &ParsedModelName,
) {
request.model = parsed_model.model_id.clone();
apply_reasoning_options(request, parsed_model);
apply_service_tier_options(request, parsed_model);
}
fn apply_reasoning_options(request: &mut CodexResponsesRequest, parsed_model: &ParsedModelName) {
let existing_reasoning = request.reasoning.take();
let existing_object = match existing_reasoning {
Some(Value::Object(object)) => object,
Some(_) | None => Map::new(),
};
let effort = existing_object
.get("effort")
.and_then(Value::as_str)
.map(ToString::to_string)
.or_else(|| non_empty_string(parsed_model.reasoning_effort.as_deref()));
if effort.is_none() && existing_object.is_empty() {
request.reasoning = None;
return;
}
let summary = existing_object
.get("summary")
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.unwrap_or("auto");
let mut reasoning = Map::new();
reasoning.insert("summary".to_string(), Value::String(summary.to_string()));
if let Some(effort) = effort {
reasoning.insert("effort".to_string(), Value::String(effort));
}
request.reasoning = Some(Value::Object(reasoning));
ensure_reasoning_include(request);
}
pub(crate) 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 apply_service_tier_options(request: &mut CodexResponsesRequest, parsed_model: &ParsedModelName) {
request.service_tier = request
.service_tier
.take()
.and_then(|value| non_empty_string(Some(&value)))
.or_else(|| non_empty_string(parsed_model.service_tier.as_deref()))
.map(normalize_service_tier_for_upstream);
}
fn normalize_service_tier_for_upstream(service_tier: String) -> String {
if service_tier == "fast" {
"priority".to_string()
} else {
service_tier
}
}
fn non_empty_string(value: Option<&str>) -> Option<String> {
value.and_then(|value| {
let trimmed = value.trim();
(!trimmed.is_empty()).then(|| trimmed.to_string())
})
}
impl CodexResponsesRequest {
/// 构造默认的 HTTP SSE 请求。
pub fn new_http_sse(
model: impl Into<String>,
instructions: impl Into<String>,
input: Vec<Value>,
) -> Self {
/// 从客户端原始 Responses JSON object 构造上游请求。
///
/// 只做最小规范化:`input` 缺省为 `[]`、`stream` 缺省 `true`、`store` 缺省 `false`。
/// 其余字段(含未知字段)原样保留在 `body` 中透传上游。
pub fn from_body(mut body: Map<String, Value>) -> Self {
body.entry("input")
.or_insert_with(|| Value::Array(Vec::new()));
body.entry("stream").or_insert(Value::Bool(true));
body.entry("store").or_insert(Value::Bool(false));
Self {
model: model.into(),
instructions: instructions.into(),
input,
stream: true,
store: false,
reasoning: None,
tools: None,
tool_choice: None,
parallel_tool_calls: None,
text: None,
generate: None,
body,
tuple_schema: None,
service_tier: None,
previous_response_id: None,
prompt_cache_key: None,
explicit_prompt_cache_key: false,
client_conversation_id: None,
variant_identity: None,
include: None,
client_metadata: None,
client_ip: None,
client_user_agent: None,
use_websocket: false,
@@ -687,9 +556,192 @@ impl CodexResponsesRequest {
}
}
/// 构造默认的 HTTP SSE 请求(测试与内部构造用)。
pub fn new_http_sse(
model: impl Into<String>,
instructions: impl Into<String>,
input: Vec<Value>,
) -> Self {
let mut body = Map::new();
body.insert("model".to_string(), Value::String(model.into()));
body.insert(
"instructions".to_string(),
Value::String(instructions.into()),
);
body.insert("input".to_string(), Value::Array(input));
Self::from_body(body)
}
/// 上游 body 的只读视图。
pub fn body(&self) -> &Map<String, Value> {
&self.body
}
// --- body 字段类型化访问器(上游语义字段,透传不重写)---
/// 模型名。
pub fn model(&self) -> &str {
self.body
.get("model")
.and_then(Value::as_str)
.unwrap_or_default()
}
/// 设置模型名(模型后缀路由归一)。
pub fn set_model(&mut self, model: impl Into<String>) {
self.body
.insert("model".to_string(), Value::String(model.into()));
}
/// 指令文本(缺省空串)。
pub fn instructions(&self) -> &str {
self.body
.get("instructions")
.and_then(Value::as_str)
.unwrap_or_default()
}
/// 输入条目切片(非数组时为空)。
pub fn input(&self) -> &[Value] {
self.body
.get("input")
.and_then(Value::as_array)
.map_or(&[], Vec::as_slice)
}
/// 替换输入条目(implicit resume / 字符串输入兼容)。
pub fn set_input(&mut self, input: Vec<Value>) {
self.body.insert("input".to_string(), Value::Array(input));
}
/// 是否流式返回。
pub fn stream(&self) -> bool {
self.body
.get("stream")
.and_then(Value::as_bool)
.unwrap_or(true)
}
/// 设置流式标志。
pub fn set_stream(&mut self, stream: bool) {
self.body.insert("stream".to_string(), Value::Bool(stream));
}
/// 是否要求上游存储响应。
pub fn store(&self) -> bool {
self.body
.get("store")
.and_then(Value::as_bool)
.unwrap_or(false)
}
/// reasoning 配置(透传,不规整)。
pub fn reasoning(&self) -> Option<&Value> {
self.body.get("reasoning")
}
/// 设置 reasoning 配置(模型后缀路由)。
pub fn set_reasoning(&mut self, reasoning: Option<Value>) {
match reasoning {
Some(value) => {
self.body.insert("reasoning".to_string(), value);
}
None => {
self.body.remove("reasoning");
}
}
}
/// 工具定义数组(非数组或空时 None)。
pub fn tools(&self) -> Option<&[Value]> {
self.body
.get("tools")
.and_then(Value::as_array)
.filter(|tools| !tools.is_empty())
.map(Vec::as_slice)
}
/// include 列表(透传原值)。
pub fn include(&self) -> Option<&Value> {
self.body.get("include")
}
/// service tier(透传原值)。
pub fn service_tier(&self) -> Option<&str> {
self.body.get("service_tier").and_then(Value::as_str)
}
/// 设置 service tier(模型后缀路由 / 归一)。
pub fn set_service_tier(&mut self, service_tier: Option<String>) {
match service_tier {
Some(value) => {
self.body
.insert("service_tier".to_string(), Value::String(value));
}
None => {
self.body.remove("service_tier");
}
}
}
/// 前一个 response ID。
pub fn previous_response_id(&self) -> Option<&str> {
self.body
.get("previous_response_id")
.and_then(Value::as_str)
}
/// 设置 / 清除前一个 response ID。
pub fn set_previous_response_id(&mut self, previous_response_id: Option<String>) {
match previous_response_id {
Some(value) => {
self.body
.insert("previous_response_id".to_string(), Value::String(value));
}
None => {
self.body.remove("previous_response_id");
}
}
}
/// 提示缓存键。
pub fn prompt_cache_key(&self) -> Option<&str> {
self.body.get("prompt_cache_key").and_then(Value::as_str)
}
/// 设置提示缓存键。
pub fn set_prompt_cache_key(&mut self, prompt_cache_key: Option<String>) {
match prompt_cache_key {
Some(value) => {
self.body
.insert("prompt_cache_key".to_string(), Value::String(value));
}
None => {
self.body.remove("prompt_cache_key");
}
}
}
/// client metadata(透传原值)。
pub fn client_metadata(&self) -> Option<&Value> {
self.body.get("client_metadata")
}
/// 设置 / 合并 client metadata。
pub fn set_client_metadata(&mut self, client_metadata: Option<Value>) {
match client_metadata {
Some(value) => {
self.body.insert("client_metadata".to_string(), value);
}
None => {
self.body.remove("client_metadata");
}
}
}
/// 判断请求是否声明了图片生成工具。
pub fn expects_image_generation(&self) -> bool {
self.tools.as_deref().is_some_and(|tools| {
self.tools().is_some_and(|tools| {
tools
.iter()
.any(|tool| tool.get("type").and_then(Value::as_str) == Some("image_generation"))
@@ -698,30 +750,46 @@ impl CodexResponsesRequest {
}
/// Codex compact 端点请求体。
#[derive(Debug, Clone, Serialize, Deserialize)]
///
/// 发往上游的 Responses compact 请求。`body` 持有客户端原始 JSON object
/// (已剥离 compact 上游不接受的字段),逐字段透传上游;
/// `client_ip`/`client_user_agent` 仅供管理端使用记录展示,不进上游 body。
#[derive(Debug, Clone)]
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>,
#[serde(skip)]
/// 上游请求体(唯一真相源)。
pub body: Map<String, Value>,
/// 代理侧识别的客户端 IP,仅用于管理端使用记录展示。
pub client_ip: Option<String>,
#[serde(skip)]
/// 客户端 User-Agent,仅用于管理端使用记录展示。
pub client_user_agent: Option<String>,
}
impl Serialize for CodexCompactRequest {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
self.body.serialize(serializer)
}
}
impl CodexCompactRequest {
/// 模型名。
pub fn model(&self) -> &str {
self.body
.get("model")
.and_then(Value::as_str)
.unwrap_or_default()
}
/// 设置模型名(模型后缀路由归一)。
pub fn set_model(&mut self, model: impl Into<String>) {
self.body
.insert("model".to_string(), Value::String(model.into()));
}
/// reasoning 配置(透传原值)。
pub fn reasoning(&self) -> Option<&Value> {
self.body.get("reasoning")
}
}
-124
View File
@@ -2,24 +2,6 @@
use serde_json::{Map, Value};
pub(crate) struct PreparedSchema {
pub(crate) schema: Value,
pub(crate) original_schema: Option<Value>,
}
pub(crate) fn prepare_schema(schema: Value) -> PreparedSchema {
let original_schema = has_tuple_schemas(&schema).then(|| schema.clone());
let mut schema = schema;
if original_schema.is_some() {
convert_tuple_schemas(&mut schema);
}
inject_additional_properties(&mut schema);
PreparedSchema {
schema,
original_schema,
}
}
pub(crate) fn reconvert_tuple_values(data: Value, schema: &Value) -> Value {
reconvert_tuple_values_with_root(data, schema, schema)
}
@@ -145,109 +127,3 @@ fn schema_array_has_tuple_schema(object: &Map<String, Value>, key: &str) -> bool
};
values.iter().any(has_tuple_schemas)
}
fn convert_tuple_schemas(value: &mut Value) {
let Value::Object(object) = value else {
return;
};
if let Some(Value::Array(prefix_items)) = object.remove("prefixItems") {
let mut properties = Map::new();
let mut required = Vec::new();
for (index, mut item) in prefix_items.into_iter().enumerate() {
convert_tuple_schemas(&mut item);
let key = index.to_string();
properties.insert(key.clone(), item);
required.push(Value::String(key));
}
object.insert("type".to_string(), Value::String("object".to_string()));
object.insert("properties".to_string(), Value::Object(properties));
object.insert("required".to_string(), Value::Array(required));
object.insert("additionalProperties".to_string(), Value::Bool(false));
object.remove("items");
return;
}
convert_object_values(object, "properties");
convert_object_value(object, "items");
convert_schema_array(object, "oneOf");
convert_schema_array(object, "anyOf");
convert_schema_array(object, "allOf");
convert_object_values(object, "$defs");
convert_object_values(object, "definitions");
convert_object_value(object, "if");
convert_object_value(object, "then");
convert_object_value(object, "else");
convert_object_value(object, "not");
}
fn convert_object_value(object: &mut Map<String, Value>, key: &str) {
if let Some(value) = object.get_mut(key) {
convert_tuple_schemas(value);
}
}
fn convert_object_values(object: &mut Map<String, Value>, key: &str) {
let Some(Value::Object(values)) = object.get_mut(key) else {
return;
};
for value in values.values_mut() {
convert_tuple_schemas(value);
}
}
fn convert_schema_array(object: &mut Map<String, Value>, key: &str) {
let Some(Value::Array(values)) = object.get_mut(key) else {
return;
};
for value in values {
convert_tuple_schemas(value);
}
}
fn inject_additional_properties(value: &mut Value) {
let Value::Object(object) = value else {
return;
};
if object.get("type").and_then(Value::as_str) == Some("object")
&& !object.contains_key("additionalProperties")
{
object.insert("additionalProperties".to_string(), Value::Bool(false));
}
inject_object_values(object, "properties");
inject_object_value(object, "items");
inject_schema_array(object, "oneOf");
inject_schema_array(object, "anyOf");
inject_schema_array(object, "allOf");
inject_object_values(object, "$defs");
inject_object_values(object, "definitions");
inject_object_value(object, "if");
inject_object_value(object, "then");
inject_object_value(object, "else");
inject_object_value(object, "not");
}
fn inject_object_value(object: &mut Map<String, Value>, key: &str) {
if let Some(value) = object.get_mut(key) {
inject_additional_properties(value);
}
}
fn inject_object_values(object: &mut Map<String, Value>, key: &str) {
let Some(Value::Object(values)) = object.get_mut(key) else {
return;
};
for value in values.values_mut() {
inject_additional_properties(value);
}
}
fn inject_schema_array(object: &mut Map<String, Value>, key: &str) {
let Some(Value::Array(values)) = object.get_mut(key) else {
return;
};
for value in values {
inject_additional_properties(value);
}
}
+19 -209
View File
@@ -1,4 +1,4 @@
use serde::{ser::SerializeMap, Deserialize, Serialize};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::upstream::protocol::responses::{
@@ -48,17 +48,6 @@ pub struct WebSocketAuditArtifact {
pub opening: Option<OpeningAuditSnapshot>,
/// 首个 `response.create` payload 快照。
pub payload: Option<PayloadAuditSnapshot>,
/// 失败快照。
pub error: Option<WebSocketAuditErrorSnapshot>,
}
/// WebSocket audit error snapshot.
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct WebSocketAuditErrorSnapshot {
/// 错误分类。
pub classification: String,
/// 错误消息。
pub message: String,
}
/// 从 WebSocket 错误帧推导出的上游错误分类。
@@ -153,31 +142,13 @@ fn is_internal_websocket_event(event: &str) -> bool {
event == "codex.rate_limits"
}
/// 判断 WebSocket 事件是否应在转发前剥离。
///
/// 仅剥离传输层内部帧:`response.metadata` 承载 turn_state,由上层提取转存到
/// 会话状态,不下发客户端。业务事件一律原样转发,不做 schema 校验。
/// (`codex.rate_limits` 内部事件由 `is_internal_websocket_event` 单独剥离。)
fn websocket_event_should_skip(raw: &str) -> bool {
websocket_metadata_turn_state(raw).is_some()
|| websocket_event_shape_parse_error(raw)
|| websocket_response_completed_missing_response(raw)
|| websocket_response_created_missing_response(raw)
|| websocket_response_output_text_delta_missing_delta(raw)
|| websocket_delta_event_missing_official_required_fields(raw)
|| websocket_output_item_event_missing_item(raw)
|| websocket_output_item_event_non_object_item(raw)
|| websocket_output_item_event_invalid_item_type_tag(raw)
|| websocket_output_item_event_invalid_metadata(raw)
|| websocket_message_output_item_event_invalid_required_fields(raw)
|| websocket_agent_message_output_item_event_invalid_required_fields(raw)
|| websocket_reasoning_output_item_event_invalid_required_fields(raw)
|| websocket_function_call_output_item_event_invalid_required_fields(raw)
|| websocket_function_call_output_payload_item_event_invalid_required_fields(raw)
|| websocket_custom_tool_call_output_item_event_invalid_required_fields(raw)
|| websocket_custom_tool_call_output_payload_item_event_invalid_required_fields(raw)
|| websocket_tool_search_call_output_item_event_invalid_required_fields(raw)
|| websocket_tool_search_output_item_event_invalid_required_fields(raw)
|| websocket_local_shell_call_output_item_event_invalid_required_fields(raw)
|| websocket_web_search_call_output_item_event_invalid_required_fields(raw)
|| websocket_image_generation_call_output_item_event_invalid_required_fields(raw)
|| websocket_compaction_output_item_event_invalid_required_fields(raw)
|| websocket_reasoning_summary_part_added_missing_summary_index(raw)
}
/// 从 `response.metadata` 帧中提取 `x-codex-turn-state`。
@@ -978,7 +949,6 @@ pub fn websocket_audit_artifact_from_attempt(
fallback_allowed: http_sse_fallback_allowed(request),
opening: Some(opening),
payload: Some(payload),
error: None,
}
}
@@ -991,195 +961,35 @@ fn websocket_transport_mode_name(request: &CodexResponsesRequest) -> &'static st
}
/// 生成 Responses WebSocket `response.create` payload。
///
/// payload = `{"type": "response.create"}` 加上原始上游 body 的全部字段
/// (保持插入顺序,含未知字段),逐字段原样透传。
pub fn websocket_response_create_payload(request: &CodexResponsesRequest) -> Value {
let mut body = Map::new();
insert_value(
&mut body,
"type",
let mut payload = Map::new();
payload.insert(
"type".to_string(),
Value::String("response.create".to_string()),
);
insert_value(&mut body, "model", Value::String(request.model.clone()));
insert_value(
&mut body,
"instructions",
Value::String(request.instructions.clone()),
);
insert_value(&mut body, "input", Value::Array(request.input.clone()));
insert_value(&mut body, "store", Value::Bool(request.store));
insert_value(&mut body, "stream", Value::Bool(request.stream));
if let Some(previous_response_id) = &request.previous_response_id {
insert_value(
&mut body,
"previous_response_id",
Value::String(previous_response_id.clone()),
);
for (key, value) in request.body() {
payload.insert(key.clone(), value.clone());
}
if let Some(reasoning) = &request.reasoning {
insert_value(&mut body, "reasoning", reasoning.clone());
}
if let Some(tools) = non_empty_tools(request) {
insert_value(&mut body, "tools", Value::Array(tools.to_vec()));
}
insert_value(
&mut body,
"tool_choice",
request
.tool_choice
.clone()
.unwrap_or(Value::String("auto".to_string())),
);
insert_value(
&mut body,
"parallel_tool_calls",
Value::Bool(request.parallel_tool_calls.unwrap_or(true)),
);
if let Some(text) = &request.text {
insert_value(&mut body, "text", text.clone());
}
if let Some(service_tier) = &request.service_tier {
insert_value(
&mut body,
"service_tier",
Value::String(service_tier.clone()),
);
}
if let Some(prompt_cache_key) = &request.prompt_cache_key {
insert_value(
&mut body,
"prompt_cache_key",
Value::String(prompt_cache_key.clone()),
);
}
if let Some(include) = non_empty_include(request) {
insert_value(
&mut body,
"include",
include
.iter()
.cloned()
.map(Value::String)
.collect::<Vec<_>>()
.into(),
);
}
if let Some(client_metadata) = &request.client_metadata {
insert_value(&mut body, "client_metadata", client_metadata.clone());
}
Value::Object(body)
Value::Object(payload)
}
/// 生成 Responses WebSocket `response.create` 文本帧内容。
pub fn websocket_response_create_payload_text(
request: &CodexResponsesRequest,
) -> Result<String, serde_json::Error> {
serde_json::to_string(&OrderedResponseCreatePayload(request))
}
struct OrderedResponseCreatePayload<'a>(&'a CodexResponsesRequest);
impl Serialize for OrderedResponseCreatePayload<'_> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
let request = self.0;
let mut map = serializer.serialize_map(Some(websocket_payload_keys(request).len()))?;
map.serialize_entry("type", "response.create")?;
map.serialize_entry("model", &request.model)?;
map.serialize_entry("instructions", &request.instructions)?;
map.serialize_entry("input", &request.input)?;
map.serialize_entry("store", &request.store)?;
map.serialize_entry("stream", &request.stream)?;
if let Some(previous_response_id) = &request.previous_response_id {
map.serialize_entry("previous_response_id", previous_response_id)?;
}
if let Some(reasoning) = &request.reasoning {
map.serialize_entry("reasoning", reasoning)?;
}
if let Some(tools) = non_empty_tools(request) {
map.serialize_entry("tools", tools)?;
}
if let Some(tool_choice) = &request.tool_choice {
map.serialize_entry("tool_choice", tool_choice)?;
} else {
map.serialize_entry("tool_choice", "auto")?;
}
map.serialize_entry(
"parallel_tool_calls",
&request.parallel_tool_calls.unwrap_or(true),
)?;
if let Some(text) = &request.text {
map.serialize_entry("text", text)?;
}
if let Some(service_tier) = &request.service_tier {
map.serialize_entry("service_tier", service_tier)?;
}
if let Some(prompt_cache_key) = &request.prompt_cache_key {
map.serialize_entry("prompt_cache_key", prompt_cache_key)?;
}
if let Some(include) = non_empty_include(request) {
map.serialize_entry("include", include)?;
}
if let Some(client_metadata) = &request.client_metadata {
map.serialize_entry("client_metadata", client_metadata)?;
}
map.end()
}
}
fn insert_value(body: &mut Map<String, Value>, key: &str, value: Value) {
body.insert(key.to_string(), value);
serde_json::to_string(&websocket_response_create_payload(request))
}
fn websocket_payload_keys(request: &CodexResponsesRequest) -> Vec<String> {
let mut keys = vec![
"type".to_string(),
"model".to_string(),
"instructions".to_string(),
"input".to_string(),
"store".to_string(),
"stream".to_string(),
];
if request.previous_response_id.is_some() {
keys.push("previous_response_id".to_string());
}
if request.reasoning.is_some() {
keys.push("reasoning".to_string());
}
if non_empty_tools(request).is_some() {
keys.push("tools".to_string());
}
keys.push("tool_choice".to_string());
keys.push("parallel_tool_calls".to_string());
if request.text.is_some() {
keys.push("text".to_string());
}
if request.service_tier.is_some() {
keys.push("service_tier".to_string());
}
if request.prompt_cache_key.is_some() {
keys.push("prompt_cache_key".to_string());
}
if non_empty_include(request).is_some() {
keys.push("include".to_string());
}
if request.client_metadata.is_some() {
keys.push("client_metadata".to_string());
}
let mut keys = Vec::with_capacity(request.body().len() + 1);
keys.push("type".to_string());
keys.extend(request.body().keys().cloned());
keys
}
fn non_empty_tools(request: &CodexResponsesRequest) -> Option<&[Value]> {
request.tools.as_deref().filter(|tools| !tools.is_empty())
}
fn non_empty_include(request: &CodexResponsesRequest) -> Option<&[String]> {
request
.include
.as_deref()
.filter(|include| !include.is_empty())
}
fn redact_payload_body(body: Value) -> Value {
let Value::Object(body) = body else {
return body;
+15 -30
View File
@@ -24,7 +24,7 @@ use uuid::Uuid;
use crate::upstream::fingerprint::Fingerprint;
use crate::upstream::models::backend_entry::BackendModelEntry;
use crate::upstream::protocol::events::extract_sse_usage;
use crate::upstream::protocol::events::{extract_sse_usage, retry_after_seconds_from_body};
use crate::upstream::protocol::responses::{
http_sse_fallback_allowed, transport_for_request, CodexCompactRequest, CodexResponsesRequest,
CodexTransport,
@@ -678,10 +678,9 @@ impl CodexBackendClient {
) -> Option<CodexWebSocketPoolKey> {
let account_id = pool_account_id.or(context.account_id)?;
let conversation_id = request
.prompt_cache_key
.as_deref()
.prompt_cache_key()
.or(request.client_conversation_id.as_deref())
.or(request.previous_response_id.as_deref())?;
.or(request.previous_response_id())?;
Some(CodexWebSocketPoolKey::new(
&self.base_url,
account_id,
@@ -778,7 +777,7 @@ impl CodexBackendClient {
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let mut headers = self.request_headers(context)?;
if let Some(subagent) = openai_subagent_from_metadata(request.client_metadata.as_ref()) {
if let Some(subagent) = openai_subagent_from_metadata(request.client_metadata()) {
headers.insert(
HeaderName::from_static("x-openai-subagent"),
HeaderValue::from_str(&subagent)?,
@@ -1012,18 +1011,6 @@ pub(super) fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Op
.or_else(|| body.and_then(retry_after_seconds_from_body))
}
pub(super) fn retry_after_seconds_from_body(body: &str) -> Option<u64> {
let value: Value = serde_json::from_str(body).ok()?;
value
.get("retry_after_seconds")
.or_else(|| value.get("retry_after"))
.and_then(|v| {
v.as_u64()
.or_else(|| v.as_str().and_then(|s| s.parse::<u64>().ok()))
})
.filter(|seconds| *seconds > 0)
}
pub(super) fn truncate_for_error(body: &str) -> String {
body.chars().take(200).collect()
}
@@ -1082,9 +1069,9 @@ fn response_upstream_request(
) -> CodexResponsesRequest {
let mut upstream = request.clone();
if let Some(session_id) = context.session_id {
upstream.prompt_cache_key = Some(session_id.to_string());
upstream.set_prompt_cache_key(Some(session_id.to_string()));
}
upstream.client_metadata = response_client_metadata(request.client_metadata.as_ref(), context);
upstream.set_client_metadata(response_client_metadata(request.client_metadata(), context));
upstream
}
@@ -1095,15 +1082,15 @@ fn websocket_upstream_request(request: &CodexResponsesRequest) -> CodexResponses
}
fn stamp_ws_stream_request_start_ms(request: &mut CodexResponsesRequest) {
let mut metadata = match request.client_metadata.take() {
Some(Value::Object(metadata)) => metadata,
let mut metadata = match request.client_metadata() {
Some(Value::Object(metadata)) => metadata.clone(),
_ => Map::new(),
};
metadata.insert(
X_CODEX_WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY.to_string(),
Value::String(now_unix_timestamp_millis().to_string()),
);
request.client_metadata = Some(Value::Object(metadata));
request.set_client_metadata(Some(Value::Object(metadata)));
}
fn now_unix_timestamp_millis() -> u128 {
@@ -1117,14 +1104,12 @@ fn response_client_metadata(
client_metadata: Option<&Value>,
context: CodexRequestContext<'_>,
) -> Option<Value> {
let mut metadata = Map::new();
if let Some(Value::Object(input)) = client_metadata {
for (key, value) in input {
if let Some(value) = value.as_str() {
metadata.insert(key.clone(), Value::String(value.to_string()));
}
}
}
// 以客户端原始 client_metadata 为基础(保留 number/bool/object 等所有值类型),
// 在其上追加代理自身的上下文字段。
let mut metadata = match client_metadata {
Some(Value::Object(input)) => input.clone(),
_ => Map::new(),
};
insert_metadata_string(
&mut metadata,
+1
View File
@@ -8,6 +8,7 @@ pub mod tls;
pub mod usage;
pub mod websocket;
pub mod websocket_pool;
pub(crate) mod websocket_pump;
pub use self::{
client::{
+3 -4
View File
@@ -1,11 +1,10 @@
use reqwest::StatusCode;
use serde_json::Value;
use crate::upstream::protocol::events::retry_after_seconds_from_body;
use super::{
client::{
read_capped_error_body, retry_after_seconds, retry_after_seconds_from_body,
truncate_for_error,
},
client::{read_capped_error_body, retry_after_seconds, truncate_for_error},
endpoints::usage_endpoint_urls,
CodexBackendClient, CodexClientError, CodexClientResult, CodexRequestContext,
};
+114 -97
View File
@@ -10,8 +10,9 @@ use std::{
use bytes::Bytes;
use chrono::Utc;
use futures::{channel::mpsc, SinkExt, Stream, StreamExt};
use futures::Stream;
use thiserror::Error;
use tokio::sync::mpsc;
use tokio::{sync::Mutex, time::timeout};
use tokio_tungstenite::{connect_async_tls_with_config, Connector};
use tungstenite::{
@@ -37,15 +38,18 @@ use crate::upstream::protocol::websocket::{
use super::response_meta;
use super::websocket_pool::{
CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey, CodexWsStream,
CodexWebSocketConnectionMetadata, CodexWebSocketPool, CodexWebSocketPoolKey,
PooledWebSocketConnection, WebSocketPoolAcquire, WebSocketPoolDecision,
};
use super::websocket_pump::{PumpKeepalive, PumpedWebSocket, RawWsStream};
use uuid::Uuid;
const REDACTED_HEADER_VALUE: &str = "<redacted>";
const CODEX_RESPONSES_PATH: &str = "/codex/responses";
const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits";
const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20);
const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const WEBSOCKET_STREAM_BUFFER: usize = 16;
/// WebSocket audit artifact 输出目录环境变量。
pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR";
@@ -110,8 +114,6 @@ pub struct CodexWebSocketExchange {
pub set_cookie_headers: Vec<String>,
/// 上游握手响应里的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// 打开握手响应状态码。
pub handshake_status: u16,
/// 首个有效上游 WebSocket 事件到达代理的耗时。
pub first_token_ms: Option<i64>,
/// WebSocket 连接池决策。
@@ -132,8 +134,6 @@ pub struct CodexWebSocketStreamingExchange {
pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
/// 上游内部 metadata 事件里的动态 turn state。
pub turn_state_update: CodexWebSocketTurnStateUpdate,
/// 打开握手响应状态码。
pub handshake_status: u16,
/// WebSocket 连接池决策。
pub pool_decision: Option<WebSocketPoolDecision>,
}
@@ -319,7 +319,7 @@ fn websocket_host_header(endpoint: &str) -> Option<String> {
async fn connect_websocket(
connection: &CodexWebSocketConnection,
) -> Result<(CodexWsStream, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
) -> Result<(RawWsStream, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
let request = websocket_handshake_request(connection)?;
let connector = super::tls::maybe_build_rustls_client_config_with_custom_ca()
.map_err(|error| tungstenite::Error::Io(std::io::Error::other(error)))?
@@ -331,6 +331,18 @@ async fn connect_websocket(
}
}
/// 建立连接并交给后台 pump 托管;`keepalive` 决定该连接是否做主动保活。
///
/// 非池化(即用即弃)连接用 [`PumpKeepalive::disabled`];池化连接用连接池派生的保活策略,
/// 让空闲期在后台完成 ping/pong 与失活检测,从而复用前可零成本判活。
async fn connect_pumped_websocket(
connection: &CodexWebSocketConnection,
keepalive: PumpKeepalive,
) -> Result<(PumpedWebSocket, WsResponse<Option<Vec<u8>>>), CodexWebSocketExchangeError> {
let (raw, response) = connect_websocket(connection).await?;
Ok((PumpedWebSocket::new(raw, keepalive), response))
}
fn websocket_handshake_request(
connection: &CodexWebSocketConnection,
) -> Result<WsRequest, tungstenite::http::Error> {
@@ -401,7 +413,8 @@ async fn execute_fresh_response_create_request(
request: &CodexWebSocketRequest,
started_at: Instant,
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
let (mut websocket, response) = connect_websocket(request.connection()).await?;
let (websocket, response) =
connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?;
websocket
.send(Message::Text(request.payload_text().to_string().into()))
.await?;
@@ -415,7 +428,8 @@ async fn execute_fresh_response_create_request(
async fn execute_fresh_response_create_request_stream(
request: &CodexWebSocketRequest,
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
let (mut websocket, response) = connect_websocket(request.connection()).await?;
let (websocket, response) =
connect_pumped_websocket(request.connection(), PumpKeepalive::disabled()).await?;
websocket
.send(Message::Text(request.payload_text().to_string().into()))
.await?;
@@ -544,7 +558,8 @@ async fn execute_fresh_pooled_response_create_request(
key: CodexWebSocketPoolKey,
started_at: Instant,
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
let (mut websocket, response) = connect_websocket(request.connection()).await?;
let (websocket, response) =
connect_pumped_websocket(request.connection(), pool.keepalive()).await?;
websocket
.send(Message::Text(request.payload_text().to_string().into()))
.await?;
@@ -553,15 +568,12 @@ async fn execute_fresh_pooled_response_create_request(
let metadata = websocket_connection_metadata(&response);
let (exchange, websocket, metadata) =
collect_websocket_response(websocket, metadata, false, started_at).await?;
let last_activity_at = Instant::now();
pool.put(
key,
PooledWebSocketConnection {
websocket,
metadata,
created_at: now,
last_activity_at,
last_ping_at: None,
},
)
.await;
@@ -573,7 +585,8 @@ async fn execute_fresh_pooled_response_create_request_stream(
pool: CodexWebSocketPool,
key: CodexWebSocketPoolKey,
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
let (mut websocket, response) = connect_websocket(request.connection()).await?;
let (websocket, response) =
connect_pumped_websocket(request.connection(), pool.keepalive()).await?;
websocket
.send(Message::Text(request.payload_text().to_string().into()))
.await?;
@@ -587,7 +600,6 @@ async fn execute_fresh_pooled_response_create_request_stream(
pool,
key,
created_at: now,
last_ping_at: None,
}),
None,
))
@@ -597,7 +609,7 @@ async fn execute_pooled_response_create_request(
request: &CodexWebSocketRequest,
pool: &CodexWebSocketPool,
key: CodexWebSocketPoolKey,
mut connection: PooledWebSocketConnection,
connection: PooledWebSocketConnection,
started_at: Instant,
) -> Result<CodexWebSocketExchange, CodexWebSocketExchangeError> {
if let Err(error) = connection
@@ -614,20 +626,16 @@ async fn execute_pooled_response_create_request(
}
let created_at = connection.created_at;
let last_ping_at = connection.last_ping_at;
match collect_websocket_response(connection.websocket, connection.metadata, true, started_at)
.await
{
Ok((exchange, websocket, metadata)) => {
let last_activity_at = Instant::now();
pool.put(
key,
PooledWebSocketConnection {
websocket,
metadata,
created_at,
last_activity_at,
last_ping_at,
},
)
.await;
@@ -668,28 +676,27 @@ async fn execute_pooled_response_create_request_stream(
pool,
key,
created_at: connection.created_at,
last_ping_at: connection.last_ping_at,
}),
Some(prefetched_text),
))
}
async fn first_reused_stream_text(
websocket: &mut CodexWsStream,
websocket: &mut PumpedWebSocket,
) -> Result<String, CodexWebSocketExchangeError> {
loop {
let message =
match next_websocket_message(websocket, Some(WEBSOCKET_RECEIVE_IDLE_TIMEOUT)).await {
Ok(Some(message)) => message,
Ok(None) => {
return Err(
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
message: "websocket closed".to_string(),
},
);
}
Err(error) => return Err(reused_connection_died_before_first_frame(&error)),
};
let message = match next_websocket_message(websocket, WEBSOCKET_RECEIVE_IDLE_TIMEOUT).await
{
Ok(Some(message)) => message,
Ok(None) => {
return Err(
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
message: "websocket closed".to_string(),
},
);
}
Err(error) => return Err(reused_connection_died_before_first_frame(&error)),
};
match message {
Message::Text(text) => return Ok(text.to_string()),
Message::Close(_) => {
@@ -721,14 +728,14 @@ pub async fn execute_response_create_request(
// ---------------------------------------------------------------------------
async fn collect_websocket_response(
mut websocket: CodexWsStream,
mut websocket: PumpedWebSocket,
mut metadata: CodexWebSocketConnectionMetadata,
reused_connection: bool,
started_at: Instant,
) -> Result<
(
CodexWebSocketExchange,
CodexWsStream,
PumpedWebSocket,
CodexWebSocketConnectionMetadata,
),
CodexWebSocketExchangeError,
@@ -738,18 +745,16 @@ async fn collect_websocket_response(
let mut first_token_ms = None;
loop {
let message = match next_websocket_message(
&mut websocket,
receive_timeout_before_first_response_frame(saw_response_frame),
)
.await
{
Ok(message) => message,
Err(error) if reused_connection && !saw_response_frame => {
return Err(reused_connection_died_before_first_frame(&error));
}
Err(error) => return Err(error),
};
let message =
match next_websocket_message(&mut websocket, receive_idle_timeout(saw_response_frame))
.await
{
Ok(message) => message,
Err(error) if reused_connection && !saw_response_frame => {
return Err(reused_connection_died_before_first_frame(&error));
}
Err(error) => return Err(error),
};
let Some(message) = message else {
break;
};
@@ -808,7 +813,6 @@ async fn collect_websocket_response(
turn_state: metadata.turn_state.clone(),
set_cookie_headers: metadata.set_cookie_headers.clone(),
rate_limit_headers: metadata.rate_limit_headers.clone(),
handshake_status: metadata.handshake_status,
first_token_ms,
pool_decision: None,
};
@@ -842,7 +846,6 @@ fn websocket_connection_metadata(
turn_state: response_meta::turn_state(response.headers()),
set_cookie_headers: response_meta::set_cookie_headers(response.headers()),
rate_limit_headers: response_meta::rate_limit_headers(response.headers()),
handshake_status: response.status().as_u16(),
}
}
@@ -854,11 +857,10 @@ struct WebSocketStreamPoolReturn {
pool: CodexWebSocketPool,
key: CodexWebSocketPoolKey,
created_at: Instant,
last_ping_at: Option<Instant>,
}
fn stream_websocket_response(
websocket: CodexWsStream,
websocket: PumpedWebSocket,
metadata: CodexWebSocketConnectionMetadata,
pool_return: Option<WebSocketStreamPoolReturn>,
prefetched_text: Option<String>,
@@ -868,7 +870,7 @@ fn stream_websocket_response(
let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates);
let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone()));
let turn_state_update_for_task = Arc::clone(&turn_state_update);
let (tx, rx) = mpsc::unbounded();
let (tx, rx) = mpsc::channel(WEBSOCKET_STREAM_BUFFER);
tokio::spawn(async move {
forward_websocket_response_stream(
websocket,
@@ -882,26 +884,29 @@ fn stream_websocket_response(
.await;
});
let body = futures::stream::unfold(rx, |mut rx| async move {
rx.recv().await.map(|item| (item, rx))
});
CodexWebSocketStreamingExchange {
body: Box::pin(rx),
body: Box::pin(body),
turn_state: response_metadata.turn_state,
set_cookie_headers: response_metadata.set_cookie_headers,
rate_limit_headers: response_metadata.rate_limit_headers,
rate_limit_header_updates,
turn_state_update,
handshake_status: response_metadata.handshake_status,
pool_decision: None,
}
}
async fn forward_websocket_response_stream(
mut websocket: CodexWsStream,
mut websocket: PumpedWebSocket,
mut metadata: CodexWebSocketConnectionMetadata,
pool_return: Option<WebSocketStreamPoolReturn>,
prefetched_text: Option<String>,
rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
turn_state_update: CodexWebSocketTurnStateUpdate,
tx: mpsc::UnboundedSender<Result<Bytes, CodexWebSocketExchangeError>>,
tx: mpsc::Sender<Result<Bytes, CodexWebSocketExchangeError>>,
) {
let mut pool_return = pool_return;
let mut prefetched_text = prefetched_text;
@@ -910,16 +915,24 @@ async fn forward_websocket_response_stream(
let raw = if let Some(text) = prefetched_text.take() {
text
} else {
let message = match next_websocket_message(
&mut websocket,
receive_timeout_before_first_response_frame(saw_response_frame),
)
.await
{
let message = tokio::select! {
biased;
// 客户端断开下游 SSE 流:立即丢弃连接并释放池 slot,
// 不再傻等上游 idle 超时(否则同会话后续请求会一直 bypass/busy)。
() = tx.closed() => {
discard_stream_websocket(websocket, pool_return).await;
return;
}
message = next_websocket_message(
&mut websocket,
receive_idle_timeout(saw_response_frame),
) => message,
};
let message = match message {
Ok(message) => message,
Err(error) => {
discard_stream_websocket(websocket, pool_return).await;
let _ = tx.unbounded_send(Err(error));
let _ = tx.send(Err(error)).await;
return;
}
};
@@ -930,14 +943,16 @@ async fn forward_websocket_response_stream(
Message::Text(text) => text.to_string(),
Message::Binary(_) => {
discard_stream_websocket(websocket, pool_return).await;
let _ =
tx.unbounded_send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent));
let _ = tx
.send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent))
.await;
return;
}
Message::Close(_) => {
discard_stream_websocket(websocket, pool_return).await;
let _ =
tx.unbounded_send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal));
let _ = tx
.send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal))
.await;
return;
}
_ => continue,
@@ -947,26 +962,32 @@ async fn forward_websocket_response_stream(
discard_stream_websocket(websocket, pool_return).await;
let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw)
.or_else(|| events::retry_after_seconds_from_body(&raw));
let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::Upstream {
status_code: classified.status_code,
retry_after_seconds,
body: raw,
set_cookie_headers: Vec::new(),
}));
let _ = tx
.send(Err(CodexWebSocketExchangeError::Upstream {
status_code: classified.status_code,
retry_after_seconds,
body: raw,
set_cookie_headers: Vec::new(),
}))
.await;
return;
}
if let Some(reason) = websocket_incomplete_response_reason(&raw) {
discard_stream_websocket(websocket, pool_return).await;
let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::IncompleteResponse {
reason,
}));
let _ = tx
.send(Err(CodexWebSocketExchangeError::IncompleteResponse {
reason,
}))
.await;
return;
}
if let Some(message) = websocket_response_completed_parse_error(&raw) {
discard_stream_websocket(websocket, pool_return).await;
let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::InvalidCompletedResponse {
message,
}));
let _ = tx
.send(Err(CodexWebSocketExchangeError::InvalidCompletedResponse {
message,
}))
.await;
return;
}
if let Some(headers) = websocket_rate_limit_event_headers(&raw) {
@@ -984,7 +1005,7 @@ async fn forward_websocket_response_stream(
continue;
};
let terminal = event.as_deref().is_some_and(is_terminal_websocket_event);
if tx.unbounded_send(Ok(Bytes::from(frame))).is_err() {
if tx.send(Ok(Bytes::from(frame))).await.is_err() {
discard_stream_websocket(websocket, pool_return).await;
return;
}
@@ -996,25 +1017,23 @@ async fn forward_websocket_response_stream(
}
discard_stream_websocket(websocket, pool_return).await;
let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal));
let _ = tx
.send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal))
.await;
}
fn receive_timeout_before_first_response_frame(saw_response_frame: bool) -> Option<Duration> {
fn receive_idle_timeout(saw_response_frame: bool) -> Duration {
if saw_response_frame {
None
WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT
} else {
Some(WEBSOCKET_RECEIVE_IDLE_TIMEOUT)
WEBSOCKET_RECEIVE_IDLE_TIMEOUT
}
}
async fn next_websocket_message(
websocket: &mut CodexWsStream,
receive_timeout: Option<Duration>,
websocket: &mut PumpedWebSocket,
receive_timeout: Duration,
) -> Result<Option<Message>, CodexWebSocketExchangeError> {
let Some(receive_timeout) = receive_timeout else {
return websocket.next().await.transpose().map_err(Into::into);
};
match timeout(receive_timeout, websocket.next()).await {
Ok(message) => message.transpose().map_err(Into::into),
Err(_) => Err(CodexWebSocketExchangeError::ReceiveIdleTimeout {
@@ -1024,12 +1043,12 @@ async fn next_websocket_message(
}
async fn finish_stream_websocket(
mut websocket: CodexWsStream,
mut websocket: PumpedWebSocket,
metadata: CodexWebSocketConnectionMetadata,
pool_return: Option<WebSocketStreamPoolReturn>,
) {
let Some(pool_return) = pool_return else {
let _ = websocket.close(None).await;
websocket.close().await;
return;
};
pool_return
@@ -1040,21 +1059,19 @@ async fn finish_stream_websocket(
websocket,
metadata,
created_at: pool_return.created_at,
last_activity_at: Instant::now(),
last_ping_at: pool_return.last_ping_at,
},
)
.await;
}
async fn discard_stream_websocket(
mut websocket: CodexWsStream,
mut websocket: PumpedWebSocket,
pool_return: Option<WebSocketStreamPoolReturn>,
) {
if let Some(pool_return) = pool_return {
pool_return.pool.discard(&pool_return.key).await;
}
let _ = websocket.close(None).await;
websocket.close().await;
}
// ---------------------------------------------------------------------------
+46 -156
View File
@@ -6,18 +6,17 @@ use std::{
time::{Duration, Instant},
};
use futures::{SinkExt, StreamExt};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use tokio::{net::TcpStream, sync::Mutex, time::timeout};
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
use tokio::sync::Mutex;
use super::websocket_pump::{PumpKeepalive, PumpedWebSocket};
const DEFAULT_MAX_PER_ACCOUNT: usize = 8;
const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55);
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)]
@@ -99,7 +98,19 @@ impl Default for CodexWebSocketPoolConfig {
maintenance_interval: Some(DEFAULT_MAINTENANCE_INTERVAL),
ping_interval: Some(DEFAULT_PING_INTERVAL),
ping_timeout: DEFAULT_PING_TIMEOUT,
liveness_timeout: Some(DEFAULT_LIVENESS_TIMEOUT),
// idle 连接不设失活截断:靠 ping/pong 保活,只在 max_age(55 分钟)
// 或 ping 失败时关闭,最大化跨轮复用(对齐 Codex CLI 的长连接策略)。
liveness_timeout: None,
}
}
}
impl CodexWebSocketPoolConfig {
/// pump 后台任务的保活策略:从连接池配置派生出 ping 间隔与 liveness 超时。
pub(crate) fn keepalive(&self) -> PumpKeepalive {
PumpKeepalive {
ping_interval: self.ping_interval,
liveness_timeout: self.liveness_timeout,
}
}
}
@@ -127,6 +138,11 @@ impl CodexWebSocketPool {
pool
}
/// pump 后台任务的保活策略(供建连时传入)。
pub(crate) fn keepalive(&self) -> PumpKeepalive {
self.config.keepalive()
}
pub(crate) async fn acquire(&self, key: &CodexWebSocketPoolKey) -> WebSocketPoolAcquire {
let mut expired_connection = None;
let acquire = {
@@ -135,14 +151,18 @@ impl CodexWebSocketPool {
return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Disabled);
}
match state.slots.get(key) {
Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) => {
Some(WebSocketPoolSlot::Busy) => {
return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Busy);
}
Some(WebSocketPoolSlot::Idle(_)) => {
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(key) else {
return WebSocketPoolAcquire::Bypass(WebSocketPoolBypassReason::Busy);
};
if connection.created_at.elapsed() < self.config.max_age {
// 零成本探活:后台 pump 已实时感知连接死亡(RST/Close/EOF/失活),
// 复用前只需读取 is_closed 标志,避免复用到静默死连接后卡到超时。
if connection.created_at.elapsed() < self.config.max_age
&& !connection.websocket.is_closed()
{
state.slots.insert(key.clone(), WebSocketPoolSlot::Busy);
return WebSocketPoolAcquire::Reused(connection);
}
@@ -199,7 +219,7 @@ impl CodexWebSocketPool {
let mut state = self.inner.lock().await;
match state.slots.remove(key) {
Some(WebSocketPoolSlot::Idle(connection)) => Some(*connection),
Some(WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking) | None => None,
Some(WebSocketPoolSlot::Busy) | None => None,
}
};
if let Some(connection) = idle_connection {
@@ -237,21 +257,21 @@ impl CodexWebSocketPool {
.drain()
.filter_map(|(_, slot)| match slot {
WebSocketPoolSlot::Idle(connection) => Some(*connection),
WebSocketPoolSlot::Busy | WebSocketPoolSlot::Checking => None,
WebSocketPoolSlot::Busy => None,
})
.collect::<Vec<_>>()
};
close_pooled_connections(idle_connections).await;
}
/// 维护 idle 连接:关闭过期/失活连接,并按需 ping 探活。
/// 维护 idle 连接:清扫已被后台 pump 判定死亡或已超龄的连接。
///
/// 保活(ping/pong)与失活检测已下沉到每条连接的 pump 任务内部,此处不再做
/// 同步 ping 探活,只负责把「后台已标记 closed」或「超过 max_age」的 idle
/// 连接从池中摘除并关闭,避免它们占用 slot。
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;
}
let close = self.take_dead_idle_connections().await;
close_pooled_connections(close).await;
}
fn spawn_maintenance_task(&self) {
@@ -284,107 +304,31 @@ impl CodexWebSocketPool {
self.inner.lock().await.shutting_down
}
async fn take_idle_connections_for_maintenance(&self) -> WebSocketPoolMaintenance {
async fn take_dead_idle_connections(&self) -> Vec<PooledWebSocketConnection> {
let mut close = Vec::new();
let mut probe = Vec::new();
let now = Instant::now();
let Some(ping_interval) = self.config.ping_interval else {
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
return WebSocketPoolMaintenance { close, probe };
}
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);
}
}
return WebSocketPoolMaintenance { close, probe };
};
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
return WebSocketPoolMaintenance { close, probe };
return close;
}
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,
) || should_probe_idle_connection(connection, now, ping_interval) =>
if should_close_idle_connection(connection, now, self.config.max_age) =>
{
Some(key.clone())
}
WebSocketPoolSlot::Idle(_)
| WebSocketPoolSlot::Busy
| WebSocketPoolSlot::Checking => None,
WebSocketPoolSlot::Idle(_) | WebSocketPoolSlot::Busy => None,
})
.collect::<Vec<_>>();
for key in keys {
let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) else {
continue;
};
if should_close_idle_connection(
&connection,
now,
self.config.max_age,
self.config.liveness_timeout,
) {
if let Some(WebSocketPoolSlot::Idle(connection)) = state.slots.remove(&key) {
close.push(*connection);
} else {
state.slots.insert(key.clone(), WebSocketPoolSlot::Checking);
probe.push((key, *connection));
}
}
WebSocketPoolMaintenance { close, probe }
}
async fn return_maintained_connection(
&self,
key: CodexWebSocketPoolKey,
connection: Option<PooledWebSocketConnection>,
) {
let mut connection = connection;
{
let mut state = self.inner.lock().await;
if state.shutting_down || !self.config.enabled {
state.slots.remove(&key);
} else if matches!(state.slots.get(&key), Some(WebSocketPoolSlot::Checking)) {
if let Some(connection) = connection.take() {
state
.slots
.insert(key, WebSocketPoolSlot::Idle(Box::new(connection)));
} else {
state.slots.remove(&key);
}
}
}
if let Some(connection) = connection {
close_pooled_connection(connection).await;
}
close
}
}
@@ -394,28 +338,22 @@ struct WebSocketPoolState {
shutting_down: bool,
}
pub(crate) type CodexWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
#[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,
}
pub(crate) struct PooledWebSocketConnection {
pub(crate) websocket: CodexWsStream,
pub(crate) websocket: PumpedWebSocket,
pub(crate) metadata: CodexWebSocketConnectionMetadata,
pub(crate) created_at: Instant,
pub(crate) last_activity_at: Instant,
pub(crate) last_ping_at: Option<Instant>,
}
enum WebSocketPoolSlot {
Idle(Box<PooledWebSocketConnection>),
Busy,
Checking,
}
pub(crate) enum WebSocketPoolAcquire {
@@ -538,7 +476,7 @@ fn account_slot_count(
}
async fn close_pooled_connection(mut connection: PooledWebSocketConnection) {
let _ = connection.websocket.send(Message::Close(None)).await;
connection.websocket.close().await;
}
async fn close_pooled_connections(connections: Vec<PooledWebSocketConnection>) {
@@ -547,59 +485,11 @@ async fn close_pooled_connections(connections: Vec<PooledWebSocketConnection>) {
}
}
struct WebSocketPoolMaintenance {
close: Vec<PooledWebSocketConnection>,
probe: Vec<(CodexWebSocketPoolKey, PooledWebSocketConnection)>,
}
/// idle 连接是否应从池中摘除:被后台 pump 标记死亡,或已超过 `max_age`。
fn should_close_idle_connection(
connection: &PooledWebSocketConnection,
now: Instant,
max_age: Duration,
liveness_timeout: Option<Duration>,
) -> bool {
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
})
}
fn should_probe_idle_connection(
connection: &PooledWebSocketConnection,
now: Instant,
ping_interval: Duration,
) -> bool {
!ping_interval.is_zero()
&& connection
.last_ping_at
.is_none_or(|last_ping_at| now.duration_since(last_ping_at) >= ping_interval)
}
async fn probe_idle_connection(
mut connection: PooledWebSocketConnection,
ping_timeout: Duration,
) -> Option<PooledWebSocketConnection> {
connection.last_ping_at = Some(Instant::now());
if connection
.websocket
.send(Message::Ping(Vec::new().into()))
.await
.is_err()
{
close_pooled_connection(connection).await;
return None;
}
if ping_timeout.is_zero() {
return Some(connection);
}
match timeout(ping_timeout, connection.websocket.next()).await {
Ok(Some(Ok(Message::Close(_)) | Err(_)) | None) | Err(_) => {
close_pooled_connection(connection).await;
None
}
Ok(Some(Ok(_message))) => {
connection.last_activity_at = Instant::now();
Some(connection)
}
}
connection.websocket.is_closed() || now.duration_since(connection.created_at) >= max_age
}
@@ -0,0 +1,236 @@
//! 常驻后台的 WebSocket 泵(对齐官方 Codex CLI 的连接管理)。
//!
//! 每条上游 WebSocket 连接都由一个后台 pump 任务独占:
//! - 持续读取 socket:一旦观察到 `Close` / EOF / 传输错误,立即把连接标记为 `closed`。
//! - 自动回应上游 `Ping`(保活),并按 `ping_interval` 主动 `Ping`(穿透 NAT/中间盒空闲计时器)。
//! - 可选 `liveness_timeout`:长时间无任何入站活动时判定连接失活并退出。
//!
//! 因此空闲连接的“死没死”在后台被实时感知;复用方只需零成本读取 [`PumpedWebSocket::is_closed`],
//! 无需在请求路径上做同步探活 ping(消除“复用到静默死连接才卡住超时”的长尾)。
//!
//! 收发都通过 channel 与 pump 任务交互:
//! - 发送:`send` 走 command channel,等待 pump 回执。
//! - 接收:`next` 从 message channel 取出 pump 转发的入站帧(`Ping`/`Pong` 已被 pump 吞掉)。
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
use std::time::{Duration, Instant};
use futures::{SinkExt, StreamExt};
use tokio::{
net::TcpStream,
sync::{mpsc, oneshot},
task::JoinHandle,
time::MissedTickBehavior,
};
use tokio_tungstenite::{tungstenite::Message, MaybeTlsStream, WebSocketStream};
/// 底层 tungstenite WebSocket 流。
pub(crate) type RawWsStream = WebSocketStream<MaybeTlsStream<TcpStream>>;
const PUMP_COMMAND_BUFFER: usize = 32;
/// pump 保活策略。
#[derive(Debug, Clone, Copy)]
pub(crate) struct PumpKeepalive {
/// 主动 `Ping` 间隔;`None` 表示 pump 不主动 ping(仅被动读取 + 回应上游 ping)。
pub(crate) ping_interval: Option<Duration>,
/// 无入站活动多久后判定失活;`None` 表示只靠显式 close/传输错误发现死亡。
pub(crate) liveness_timeout: Option<Duration>,
}
impl PumpKeepalive {
/// 不做主动保活(用于即用即弃的非池化连接)。
pub(crate) fn disabled() -> Self {
Self {
ping_interval: None,
liveness_timeout: None,
}
}
}
enum PumpCommand {
Send {
message: Message,
ack: oneshot::Sender<Result<(), tungstenite::Error>>,
},
}
/// 由后台 pump 任务托管的 WebSocket 连接句柄。
pub(crate) struct PumpedWebSocket {
tx_command: mpsc::Sender<PumpCommand>,
rx_message: mpsc::UnboundedReceiver<Result<Message, tungstenite::Error>>,
closed: Arc<AtomicBool>,
pump: Option<JoinHandle<()>>,
}
impl std::fmt::Debug for PumpedWebSocket {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PumpedWebSocket")
.field("closed", &self.closed.load(Ordering::Acquire))
.finish()
}
}
impl PumpedWebSocket {
/// 用底层流启动一个 pump 任务并返回句柄。
pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive) -> Self {
let (tx_command, rx_command) = mpsc::channel::<PumpCommand>(PUMP_COMMAND_BUFFER);
let (tx_message, rx_message) = mpsc::unbounded_channel();
let closed = Arc::new(AtomicBool::new(false));
let closed_for_task = Arc::clone(&closed);
let pump = tokio::spawn(async move {
pump_loop(inner, rx_command, tx_message, keepalive).await;
closed_for_task.store(true, Ordering::Release);
});
Self {
tx_command,
rx_message,
closed,
pump: Some(pump),
}
}
/// 通过 pump 发送一帧,返回底层 `send` 的结果。
pub(crate) async fn send(&self, message: Message) -> Result<(), tungstenite::Error> {
let (ack, rx_ack) = oneshot::channel();
if self
.tx_command
.send(PumpCommand::Send { message, ack })
.await
.is_err()
{
return Err(tungstenite::Error::ConnectionClosed);
}
rx_ack
.await
.unwrap_or(Err(tungstenite::Error::ConnectionClosed))
}
/// 取出下一帧入站消息(`Ping`/`Pong` 已被 pump 处理,不会到达这里)。
///
/// 返回 `None` 表示连接已结束(pump 已退出且缓冲已排空)。
pub(crate) async fn next(&mut self) -> Option<Result<Message, tungstenite::Error>> {
self.rx_message.recv().await
}
/// 连接是否已被后台 pump 判定关闭/失活。零成本,用于复用前探活。
pub(crate) fn is_closed(&self) -> bool {
self.closed.load(Ordering::Acquire) || self.tx_command.is_closed()
}
/// 主动关闭连接(best-effort 发送 Close 帧)。
pub(crate) async fn close(&mut self) {
let _ = self.send(Message::Close(None)).await;
}
}
impl Drop for PumpedWebSocket {
fn drop(&mut self) {
if let Some(pump) = self.pump.take() {
pump.abort();
}
}
}
async fn pump_loop(
mut inner: RawWsStream,
mut rx_command: mpsc::Receiver<PumpCommand>,
tx_message: mpsc::UnboundedSender<Result<Message, tungstenite::Error>>,
keepalive: PumpKeepalive,
) {
let mut last_activity = Instant::now();
let ping_interval = keepalive.ping_interval.filter(|d| !d.is_zero());
let liveness_timeout = keepalive.liveness_timeout.filter(|d| !d.is_zero());
// ticker 由 ping 间隔或 liveness 检查间隔中较小者驱动:即使不主动 ping,
// 也要周期性醒来检查 liveness。两者都未配置时不启动 ticker。
let tick_interval = match (ping_interval, liveness_timeout) {
(Some(ping), Some(live)) => Some(ping.min(live)),
(Some(ping), None) => Some(ping),
(None, Some(live)) => Some(live),
(None, None) => None,
};
let mut ticker = tick_interval.map(|d| {
// 首个 tick 推迟一整个间隔:tokio::time::interval 默认会让首个 tick 立即就绪,
// 否则连接一建立就会立刻发一帧 Ping,与首个请求 Text 抢跑、打乱帧序。
let mut ticker = tokio::time::interval_at(tokio::time::Instant::now() + d, d);
ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);
ticker
});
loop {
tokio::select! {
command = rx_command.recv() => {
match command {
// 所有句柄已丢弃:连接不再被使用,退出 pump。
None => break,
Some(PumpCommand::Send { message, ack }) => {
let result = inner.send(message).await;
let should_break = result.is_err();
let _ = ack.send(result);
if should_break {
break;
}
}
}
}
message = inner.next() => {
match message {
None => break,
Some(Ok(Message::Ping(payload))) => {
last_activity = Instant::now();
if inner.send(Message::Pong(payload)).await.is_err() {
break;
}
}
Some(Ok(Message::Pong(_))) => {
last_activity = Instant::now();
}
Some(Ok(message)) => {
last_activity = Instant::now();
let is_close = matches!(message, Message::Close(_));
if tx_message.send(Ok(message)).is_err() {
break;
}
if is_close {
break;
}
}
Some(Err(err)) => {
let _ = tx_message.send(Err(err));
break;
}
}
}
_ = tick(&mut ticker) => {
// 先检查 liveness:长时间无入站活动则判定失活退出。
if let Some(timeout) = liveness_timeout {
if last_activity.elapsed() >= timeout {
break;
}
}
// 再按需主动 ping 保活(穿透 NAT/中间盒空闲计时器)。
if ping_interval.is_some()
&& inner.send(Message::Ping(Vec::new().into())).await.is_err()
{
break;
}
}
}
}
// best-effort:尝试通知对端关闭,随后 drop 底层流释放 socket。
let _ = inner.send(Message::Close(None)).await;
}
/// 等待 ping ticker;`None` 时永远挂起,让 `select!` 分支实际禁用。
async fn tick(ticker: &mut Option<tokio::time::Interval>) {
match ticker {
Some(ticker) => {
ticker.tick().await;
}
None => std::future::pending::<()>().await,
}
}
+28 -66
View File
@@ -1,41 +1,32 @@
use std::{
path::{Component, Path, PathBuf},
sync::Arc,
convert::Infallible,
path::{Path, PathBuf},
};
use axum::{
body::Body,
extract::{Path as AxumPath, State},
http::{header, HeaderValue, StatusCode, Uri},
http::{Request, StatusCode},
response::{IntoResponse, Response},
routing::get,
Json, Router,
};
use serde_json::json;
#[derive(Clone)]
struct AssetState {
dist_dir: Arc<PathBuf>,
}
use tower::service_fn;
use tower_http::services::ServeDir;
pub fn spa_router(dist_dir: impl AsRef<Path>) -> Router {
let state = AssetState {
dist_dir: Arc::new(dist_dir.as_ref().to_path_buf()),
};
let dist_dir = dist_dir.as_ref().to_path_buf();
let index_file = dist_dir.join("index.html");
let fallback = service_fn(move |request: Request<Body>| {
let index_file = index_file.clone();
async move { Ok::<_, Infallible>(spa_fallback_response(request, index_file).await) }
});
Router::new()
.route("/", get(serve_index))
.route("/assets/{*path}", get(serve_asset))
.fallback(serve_spa_fallback)
.with_state(state)
Router::new().fallback_service(ServeDir::new(dist_dir).fallback(fallback))
}
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>, uri: Uri) -> Response {
if is_api_path(uri.path()) {
async fn spa_fallback_response(request: Request<Body>, index_file: PathBuf) -> Response {
let path = request.uri().path();
if is_api_path(path) {
return (
StatusCode::NOT_FOUND,
Json(json!({
@@ -47,61 +38,32 @@ async fn serve_spa_fallback(State(state): State<AssetState>, uri: Uri) -> Respon
.into_response();
}
serve_file(&state.dist_dir.join("index.html")).await
if is_static_file_path(path) {
return StatusCode::NOT_FOUND.into_response();
}
serve_index_file(&index_file).await
}
fn is_api_path(path: &str) -> bool {
path == "/api" || path.starts_with("/api/") || path == "/v1" || path.starts_with("/v1/")
}
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();
};
serve_file(&state.dist_dir.join("assets").join(relative_path)).await
}
async fn serve_file(path: &Path) -> Response {
async fn serve_index_file(path: &Path) -> 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 headers = response.headers_mut();
headers.insert(
header::CONTENT_TYPE,
content_type_for_path(&path.to_string_lossy()),
response.headers_mut().insert(
axum::http::header::CONTENT_TYPE,
axum::http::HeaderValue::from_static("text/html; charset=utf-8"),
);
response
}
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")
}
}
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)
fn is_static_file_path(path: &str) -> bool {
Path::new(path.trim_start_matches('/'))
.extension()
.is_some()
}
@@ -32,7 +32,6 @@ async fn admin_accounts_import_should_store_cpr_account_tokens() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -642,7 +641,6 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -859,7 +857,6 @@ async fn admin_accounts_import_should_require_admin_session_cookie() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -25,7 +25,6 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -768,7 +767,6 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state.clone());
@@ -812,7 +810,6 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
+19 -3
View File
@@ -29,7 +29,6 @@ async fn admin_accounts_list_should_not_expose_account_tokens() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -90,8 +89,8 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
window_input_tokens: 120_000,
window_output_tokens: 4_000,
window_cached_tokens: 40_000,
window_started_at: "2026-06-20T00:00:00Z",
window_reset_at: "2026-06-27T00:00:00Z",
window_started_at: "2026-06-22T12:51:00Z",
window_reset_at: "2026-06-29T12:51:00Z",
limit_window_seconds: 604_800,
last_used_at: "2026-06-23T08:50:13Z",
},
@@ -218,6 +217,20 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
sqlx::query(
"insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind("2026-06-22T10:00:00+00:00")
.bind("acct_stats")
.bind("gpt-five-hour-window")
.bind(1)
.bind(2_000)
.bind(500)
.bind(100)
.bind("2026-06-22T10:00:00Z")
.execute(&pool)
.await
.unwrap();
sqlx::query(
"insert into usage_time_buckets (bucket_start, account_id, model, request_count, input_tokens, output_tokens, cached_tokens, updated_at) values (?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind("2026-06-19T23:45:00+00:00")
.bind("acct_stats")
.bind("gpt-outside-window")
@@ -254,10 +267,13 @@ async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
assert_eq!(item["quota"]["windows"].as_array().unwrap().len(), 3);
assert_eq!(item["quota"]["windows"][0]["labelDisplay"], "月限额");
assert_eq!(item["quota"]["windows"][0]["group"], "monthly");
assert_eq!(item["quota"]["windows"][0]["tokenUsageDisplay"], "1.12M");
assert_eq!(item["quota"]["windows"][1]["labelDisplay"], "5小时限额");
assert_eq!(item["quota"]["windows"][1]["group"], "shortTerm");
assert_eq!(item["quota"]["windows"][1]["tokenUsageDisplay"], "2.5K");
assert_eq!(item["quota"]["windows"][2]["labelDisplay"], "周限额");
assert_eq!(item["quota"]["windows"][2]["group"], "shortTerm");
assert_eq!(item["quota"]["windows"][2]["tokenUsageDisplay"], "124K");
assert!(item["quota"]["windows"][2]["windowUsedDisplay"]
.as_str()
.is_some_and(|value| value.contains(" / 7.0d")));
-1
View File
@@ -172,7 +172,6 @@ async fn admin_accounts_test_app_with_overrides(
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state.clone());
+40
View File
@@ -18,6 +18,11 @@ async fn admin_account_quota_should_send_usage_cookie() {
"used_percent": 12,
"reset_at": 1_800_000_000,
"limit_window_seconds": 18_000
},
"secondary_window": {
"used_percent": 34,
"reset_at": 1_800_604_800,
"limit_window_seconds": 604_800
}
}
})))
@@ -46,6 +51,29 @@ async fn admin_account_quota_should_send_usage_cookie() {
},
)
.await;
sqlx::query(
"insert into usage_time_buckets (
bucket_start,
account_id,
model,
request_count,
input_tokens,
output_tokens,
cached_tokens,
updated_at
) values (?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind("2027-01-15T04:00:00+00:00")
.bind("acct_quota_cookie")
.bind("gpt-5")
.bind(1)
.bind(1_500)
.bind(500)
.bind(300)
.bind("2027-01-15T04:00:00Z")
.execute(&pool)
.await
.unwrap();
SqliteCookieStore::new(pool.clone())
.capture_set_cookie(
"acct_quota_cookie",
@@ -87,4 +115,16 @@ async fn admin_account_quota_should_send_usage_cookie() {
body["data"]["quotaData"]["windows"][0]["group"],
"shortTerm"
);
assert_eq!(
body["data"]["quotaData"]["windows"][0]["tokenUsageDisplay"],
"2K"
);
assert_eq!(
body["data"]["quotaData"]["windows"][1]["labelDisplay"],
"周限额"
);
assert_eq!(
body["data"]["quotaData"]["windows"][1]["tokenUsageDisplay"],
"0"
);
}
-1
View File
@@ -205,7 +205,6 @@ async fn admin_login_test_app(
fingerprint,
));
let state = codex_proxy_rs::runtime::state::AppState {
config,
services: (*services).clone(),
};
(
-1
View File
@@ -43,7 +43,6 @@ async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, tempfile::Te
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
(
+85 -6
View File
@@ -14,7 +14,7 @@ use codex_proxy_rs::{
},
infra::{
database::connect_sqlite,
time::{china_day_start, china_hour},
time::{china_day_start, china_hour, china_hour_start},
},
proxy::dispatch::session_affinity::SqliteSessionAffinityStore,
runtime::{
@@ -155,7 +155,7 @@ async fn dashboard_trend_should_bucket_usage_by_china_hour() {
.await
.unwrap();
let body = dashboard_trend(app).await;
let body = dashboard_trend(app, "usage").await;
let point = body["data"]["points"]
.as_array()
.unwrap()
@@ -167,6 +167,78 @@ async fn dashboard_trend_should_bucket_usage_by_china_hour() {
assert_eq!(point["tokensValue"], 10);
}
#[tokio::test]
async fn dashboard_latency_trend_should_use_first_token_latency() {
let (app, store, _pool, _dir) = dashboard_test_app(
"dashboard-trend-first-token-latency.sqlite",
crate::support::fingerprint::test_fingerprint(),
)
.await;
let now = Utc::now();
let current_hour_start = china_hour_start(now);
let sample_hour_start = if now - current_hour_start > Duration::minutes(31) {
current_hour_start
} else {
current_hour_start - Duration::hours(1)
};
let mut slow_completion =
usage_record_with_tokens(sample_hour_start + Duration::minutes(1), 10);
slow_completion.latency_ms = Some(10_000);
slow_completion.metadata["firstTokenMs"] = json!(100_i64);
let mut fast_completion =
usage_record_with_tokens(sample_hour_start + Duration::minutes(16), 10);
fast_completion.latency_ms = Some(1_000);
fast_completion.metadata["firstTokenMs"] = json!(500_i64);
store.append(&slow_completion).await.unwrap();
store.append(&fast_completion).await.unwrap();
let body = dashboard_trend(app, "latency").await;
let point = body["data"]["points"]
.as_array()
.unwrap()
.iter()
.find(|point| point["requestsValue"] == 2)
.expect("latency trend should include the inserted usage bucket");
let summary = body["data"]["summary"].as_array().unwrap();
assert_eq!(body["data"]["kind"], "latency");
assert_eq!(point["latencyValue"], 300);
assert_eq!(point["latency"], "300 ms");
assert_eq!(point["maxLatencyValue"], 500);
assert_eq!(point["maxLatency"], "500 ms");
assert_eq!(point["minLatencyValue"], 100);
assert_eq!(point["minLatency"], "100 ms");
assert_eq!(summary[0]["value"], "300 ms");
assert_eq!(summary[1]["value"], "500 ms");
assert_eq!(summary[2]["value"], "100 ms");
}
#[tokio::test]
async fn dashboard_summary_should_return_requested_trend_kind() {
let (app, store, _pool, _dir) = dashboard_test_app(
"dashboard-summary-requested-trend-kind.sqlite",
crate::support::fingerprint::test_fingerprint(),
)
.await;
let now = Utc::now();
let mut record = usage_record_with_tokens(now - Duration::seconds(1), 10);
record.latency_ms = Some(10_000);
record.metadata["firstTokenMs"] = json!(250_i64);
store.append(&record).await.unwrap();
let body = dashboard_summary_with_kind(app, "latency").await;
assert_eq!(body["data"]["trend"]["kind"], "latency");
assert!(
body["data"]["trend"]["summary"]
.as_array()
.unwrap()
.iter()
.any(|item| item["value"] == "250 ms"),
"summary should include first-token latency trend values: {body:?}"
);
}
#[tokio::test]
async fn dashboard_summary_should_calculate_priority_cost_from_event_service_tier() {
let cost = dashboard_summary_total_cost_for_usage_record(
@@ -397,7 +469,6 @@ async fn dashboard_test_app(
};
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
(
@@ -409,11 +480,19 @@ async fn dashboard_test_app(
}
async fn dashboard_summary(app: axum::Router) -> Value {
dashboard_summary_uri(app, "/api/admin/dashboard/summary").await
}
async fn dashboard_summary_with_kind(app: axum::Router, kind: &str) -> Value {
dashboard_summary_uri(app, &format!("/api/admin/dashboard/summary?kind={kind}")).await
}
async fn dashboard_summary_uri(app: axum::Router, uri: &str) -> Value {
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/dashboard/summary")
.uri(uri)
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_dashboard_summary")
.body(Body::empty())
@@ -425,12 +504,12 @@ async fn dashboard_summary(app: axum::Router) -> Value {
response_json(response).await
}
async fn dashboard_trend(app: axum::Router) -> Value {
async fn dashboard_trend(app: axum::Router, kind: &str) -> Value {
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/dashboard/trend?kind=usage")
.uri(format!("/api/admin/dashboard/trend?kind={kind}"))
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_dashboard_trend")
.body(Body::empty())
@@ -183,10 +183,10 @@ async fn admin_usage_records_should_return_table_display_fields() {
assert_eq!(row["tokenDetails"]["reasoningTokens"], 12);
assert_eq!(row["tokenDetails"]["totalTokens"], 1200);
assert_eq!(row["tokenDetails"]["inputTokensDisplay"], "1,000");
assert_eq!(row["costDetails"]["inputCostDisplay"], "$0.004500");
assert_eq!(row["costDetails"]["outputCostDisplay"], "$0.006000");
assert_eq!(row["costDetails"]["cacheReadCostDisplay"], "$0.000050");
assert_eq!(row["costDetails"]["totalCostDisplay"], "$0.010550");
assert_eq!(row["costDetails"]["inputCostDisplay"], "$0.00");
assert_eq!(row["costDetails"]["outputCostDisplay"], "$0.01");
assert_eq!(row["costDetails"]["cacheReadCostDisplay"], "$0.00");
assert_eq!(row["costDetails"]["totalCostDisplay"], "$0.01");
assert_eq!(
row["costDetails"]["inputPriceDisplay"],
"$5.0000 / 1M Token"
@@ -307,7 +307,7 @@ async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimens
let mut second = UsageRecord::new("request", UsageRecordLevel::Error, "insight failure");
second.id = "usage_insight_failure".to_string();
second.route = Some("/v1/chat/completions".to_string());
second.route = Some("/v1/responses/review".to_string());
second.model = Some("gpt-5.5-mini".to_string());
second.status_code = Some(500);
second.latency_ms = Some(120);
@@ -315,8 +315,8 @@ async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimens
"stream": false,
"requestedModel": "gpt-5.5-mini",
"upstreamModel": "gpt-5.5-mini",
"upstreamEndpoint": "/backend/chat",
"endpointPath": "/v1/chat/completions",
"upstreamEndpoint": "/backend/responses/review",
"endpointPath": "/v1/responses/review",
"usage": {
"inputTokens": 5,
"outputTokens": 1
@@ -343,7 +343,7 @@ async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimens
assert_eq!(requested_body["data"][0]["totalTokensValue"], 42);
assert_eq!(requested_body["data"][0]["totalTokensTotal"], "48");
assert_eq!(requested_body["data"][0]["totalTokensTotalValue"], 48);
assert_eq!(requested_body["data"][0]["actualCost"], "$0.000483");
assert_eq!(requested_body["data"][0]["actualCost"], "$0.00");
let upstream_models = app
.clone()
@@ -402,7 +402,7 @@ async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimens
.map(|item| item["name"].as_str().unwrap())
.collect::<Vec<_>>();
assert!(endpoint_paths.contains(&"/v1/responses -> /backend/responses"));
assert!(endpoint_paths.contains(&"/v1/chat/completions -> /backend/chat"));
assert!(endpoint_paths.contains(&"/v1/responses/review -> /backend/responses/review"));
let token_trend = app
.clone()
@@ -422,8 +422,8 @@ async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimens
assert_eq!(token_body["data"][0]["outputTokens"], "13");
assert_eq!(token_body["data"][0]["cacheCreationTokens"], "0");
assert_eq!(token_body["data"][0]["cachedTokens"], "6");
assert_eq!(token_body["data"][0]["actualCost"], "$0.000538");
assert_eq!(token_body["data"][0]["cost"], "$0.000538");
assert_eq!(token_body["data"][0]["actualCost"], "$0.00");
assert_eq!(token_body["data"][0]["cost"], "$0.00");
let latency_trend = app
.oneshot(
@@ -532,14 +532,14 @@ async fn admin_usage_records_insight_card_endpoints_should_return_source_specifi
let mut second = UsageRecord::new("request", UsageRecordLevel::Info, "card mini");
second.id = "usage_card_mini".to_string();
second.route = Some("/v1/chat/completions".to_string());
second.route = Some("/v1/responses/compact".to_string());
second.model = Some("gpt-5.5-mini".to_string());
second.latency_ms = Some(120);
second.metadata = json!({
"stream": false,
"requestedModel": "gpt-5.5-mini",
"upstreamModel": "gpt-5.5-mini",
"upstreamEndpoint": "/backend/chat",
"upstreamEndpoint": "/backend/responses/compact",
"usage": {
"inputTokens": 5,
"outputTokens": 1
@@ -602,7 +602,7 @@ async fn admin_usage_records_insight_card_endpoints_should_return_source_specifi
.map(|item| item["name"].as_str().unwrap())
.collect::<Vec<_>>();
assert!(endpoint_paths.contains(&"/v1/responses -> /backend/responses"));
assert!(endpoint_paths.contains(&"/v1/chat/completions -> /backend/chat"));
assert!(endpoint_paths.contains(&"/v1/responses/compact -> /backend/responses/compact"));
let token_trend = app
.clone()
@@ -885,7 +885,6 @@ async fn admin_usage_records_test_app_with_config(
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
(
+3 -7
View File
@@ -283,7 +283,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
let services = Services::new(&config, stores, fingerprint);
let services = std::sync::Arc::new(services);
let state = AppState {
config: config.clone(),
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -417,7 +416,6 @@ async fn admin_settings_update_should_apply_runtime_services() {
.await
.unwrap();
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -453,12 +451,12 @@ async fn admin_settings_update_should_apply_runtime_services() {
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let parsed = services
let resolved = services
.models
.catalog()
.await
.parse_model_name("runtime-alias");
assert_eq!(parsed.model_id, "gpt-5.5");
.resolve_model_id("runtime-alias");
assert_eq!(resolved, "gpt-5.5");
let refresh_policy = services.refresh_policy.snapshot();
assert_eq!(refresh_policy.refresh_margin_seconds, 120);
@@ -526,7 +524,6 @@ async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -569,7 +566,6 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
(
-1
View File
@@ -1220,7 +1220,6 @@ async fn admin_system_test_app(db_name: &str) -> (axum::Router, tempfile::TempDi
let stores = stores(pool);
let services = Arc::new(Services::new(&config, stores, test_fingerprint()));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
-5
View File
@@ -1,5 +0,0 @@
event: response.output_text.delta
data: {"delta":"hello"}
event: response.completed
data: {"response":{"id":"resp_1","output_text":"hello","usage":{"input_tokens":9,"output_tokens":3}}}
@@ -21,8 +21,6 @@
}
],
"store": false,
"stream": true,
"previous_response_id": "resp_implicit_resume_first",
"tool_choice": "auto",
"parallel_tool_calls": true
"stream": false,
"previous_response_id": "resp_implicit_resume_first"
}
@@ -14,6 +14,12 @@
"text": "cached reasoning"
}
],
"content": [
{
"type": "reasoning_text",
"text": "raw reasoning should not be replayed"
}
],
"encrypted_content": "enc_reasoning_replay"
}
],
+130 -36
View File
@@ -11,10 +11,12 @@ use axum::{
http::{Request, StatusCode},
middleware::{from_fn, from_fn_with_state},
routing::get,
Router,
Extension, Router,
};
use codex_proxy_rs::http::middleware::{
request_id::{attach_request_id, attach_request_id_with_proxy_config, TrustedProxyConfig},
request_id::{
attach_request_id, attach_request_id_with_proxy_config, ClientIp, TrustedProxyConfig,
},
trace::http_trace_layer,
};
use tokio::time::{sleep, Duration};
@@ -219,16 +221,12 @@ async fn request_context_should_generate_new_request_id_for_invalid_header() {
}
#[tokio::test]
async fn request_context_should_attach_real_ip_from_connection() {
async fn request_context_should_attach_client_ip_from_connection() {
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string()
get(|Extension(client_ip): Extension<ClientIp>| async move {
client_ip.as_str().to_string()
}),
)
.layer(from_fn(attach_request_id));
@@ -247,17 +245,30 @@ async fn request_context_should_attach_real_ip_from_connection() {
}
#[tokio::test]
async fn request_context_should_replace_untrusted_forwarded_ip_headers() {
async fn request_context_should_auto_detect_forwarded_ip_without_trusted_proxy_config() {
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing")
.to_string()
}),
get(
|Extension(client_ip): Extension<ClientIp>, headers: HeaderMap| async move {
let cf_ip = headers
.get("cf-connecting-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
let forwarded_for = headers
.get("x-forwarded-for")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
let real_ip = headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
format!(
"{}|cf={cf_ip}|xff={forwarded_for}|real={real_ip}",
client_ip.as_str()
)
},
),
)
.layer(from_fn(attach_request_id));
let mut request = Request::builder()
@@ -274,7 +285,37 @@ async fn request_context_should_replace_untrusted_forwarded_ip_headers() {
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "127.0.0.1");
// 未配置可信代理时进入自动模式:CF-Connecting-IP 优先于 X-Real-IP / X-Forwarded-For。
assert_eq!(
std::str::from_utf8(&bytes).unwrap(),
"203.0.113.43|cf=203.0.113.43|xff=203.0.113.42, 10.0.0.2|real=203.0.113.44"
);
}
#[tokio::test]
async fn request_context_should_skip_private_forwarded_hops_in_auto_mode() {
let app = Router::new()
.route(
"/ip",
get(|Extension(client_ip): Extension<ClientIp>| async move {
client_ip.as_str().to_string()
}),
)
.layer(from_fn(attach_request_id));
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "10.0.0.2, 172.16.0.3, 203.0.113.42")
.body(Body::empty())
.expect("request should build");
request
.extensions_mut()
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
// 自动模式下应跳过 X-Forwarded-For 链中的内网跳板,取首个公网 IP。
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "203.0.113.42");
}
#[tokio::test]
@@ -284,14 +325,12 @@ async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() {
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
let real_ip = headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing");
let forwarded_present = headers.contains_key("x-forwarded-for");
format!("{real_ip}|forwarded={forwarded_present}")
}),
get(
|Extension(client_ip): Extension<ClientIp>, headers: HeaderMap| async move {
let forwarded_present = headers.contains_key("x-forwarded-for");
format!("{}|forwarded={forwarded_present}", client_ip.as_str())
},
),
)
.layer(from_fn_with_state(
trusted_proxies,
@@ -311,23 +350,19 @@ async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() {
assert_eq!(
std::str::from_utf8(&bytes).unwrap(),
"203.0.113.42|forwarded=false"
"203.0.113.42|forwarded=true"
);
}
#[tokio::test]
async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() {
async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded_for_missing() {
let trusted_proxies =
TrustedProxyConfig::from_entries(&["127.0.0.1".to_string()]).expect("trusted proxy");
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing")
.to_string()
get(|Extension(client_ip): Extension<ClientIp>| async move {
client_ip.as_str().to_string()
}),
)
.layer(from_fn_with_state(
@@ -336,8 +371,7 @@ async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() {
));
let mut request = Request::builder()
.uri("/ip")
.header("cf-connecting-ip", "198.51.100.23")
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
.header("x-real-ip", "198.51.100.23")
.body(Body::empty())
.expect("request should build");
request
@@ -350,6 +384,66 @@ async fn request_context_should_prefer_cloudflare_ip_from_trusted_proxy() {
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "198.51.100.23");
}
#[tokio::test]
async fn request_context_should_ignore_forwarded_ip_from_untrusted_peer() {
let trusted_proxies =
TrustedProxyConfig::from_entries(&["10.0.0.0/8".to_string()]).expect("trusted proxy");
let app = Router::new()
.route(
"/ip",
get(|Extension(client_ip): Extension<ClientIp>| async move {
client_ip.as_str().to_string()
}),
)
.layer(from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
));
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "203.0.113.42")
.body(Body::empty())
.expect("request should build");
request
.extensions_mut()
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "127.0.0.1");
}
#[tokio::test]
async fn request_context_should_not_inject_real_ip_header() {
let app = Router::new()
.route(
"/ip",
get(|headers: HeaderMap| async move {
headers
.get("x-real-ip")
.and_then(|value| value.to_str().ok())
.unwrap_or("missing")
.to_string()
}),
)
.layer(from_fn(attach_request_id));
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
.header("cf-connecting-ip", "203.0.113.43")
.body(Body::empty())
.expect("request should build");
request
.extensions_mut()
.insert(ConnectInfo(SocketAddr::from(([127, 0, 0, 1], 61234))));
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "missing");
}
#[test]
fn trusted_proxy_config_should_reject_invalid_entries() {
let error = TrustedProxyConfig::from_entries(&["not-an-ip".to_string()]).unwrap_err();
+31 -1
View File
@@ -22,6 +22,11 @@ async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes
"<!doctype html><main id=\"app\">Codex Proxy Admin</main>",
)
.expect("index should be written");
std::fs::write(
dist.join("favicon.svg"),
"<svg><title>Codex Proxy</title></svg>",
)
.expect("favicon should be written");
std::fs::write(dist.join("assets").join("app.js"), "window.__asset = true;")
.expect("asset should be written");
let db = dir.path().join("assets-routes.sqlite");
@@ -29,6 +34,21 @@ async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes
let (app, _pool) = router_with_dist_and_database(&dist, &database_url).await;
let index = app.clone().oneshot(request("/")).await.expect("index");
let route_fallback = app
.clone()
.oneshot(request("/dashboard"))
.await
.expect("route fallback");
let favicon = app
.clone()
.oneshot(request("/favicon.svg"))
.await
.expect("favicon");
let missing_static = app
.clone()
.oneshot(request("/missing.svg"))
.await
.expect("missing static response");
let asset = app
.clone()
.oneshot(request("/assets/app.js"))
@@ -45,12 +65,23 @@ async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes
.expect("health response");
assert_eq!(index.status(), StatusCode::OK);
assert_eq!(route_fallback.status(), StatusCode::OK);
assert_eq!(favicon.status(), StatusCode::OK);
assert_eq!(missing_static.status(), StatusCode::NOT_FOUND);
assert_eq!(asset.status(), StatusCode::OK);
assert_eq!(api.status(), StatusCode::UNAUTHORIZED);
assert_eq!(health.status(), StatusCode::NO_CONTENT);
assert_eq!(
favicon.headers().get(header::CONTENT_TYPE),
Some(&header::HeaderValue::from_static("image/svg+xml"))
);
assert_no_static_policy_headers(&index);
assert_no_static_policy_headers(&route_fallback);
assert_no_static_policy_headers(&favicon);
assert_no_static_policy_headers(&asset);
assert_body_contains(index, "Codex Proxy Admin").await;
assert_body_contains(route_fallback, "Codex Proxy Admin").await;
assert_body_contains(favicon, "Codex Proxy").await;
assert_body_contains(asset, "__asset").await;
assert_eq!(response_json(api).await["code"], Value::from(40101));
}
@@ -172,7 +203,6 @@ async fn router_with_dist_and_database(dist: &Path, database_url: &str) -> (Rout
fingerprint,
));
let state = codex_proxy_rs::runtime::state::AppState {
config,
services: (*services).clone(),
};
(
File diff suppressed because it is too large Load Diff
@@ -132,7 +132,7 @@ async fn responses_compact_should_post_json_to_codex_compact_upstream() {
.header("x-request-id", "req_compact_success_log")
.body(Body::from(
json!({
"model": "gpt-5.5-fast",
"model": "gpt-5.5",
"instructions": "compress the session",
"input": [
{"role": "user", "content": "hello"},
@@ -160,7 +160,7 @@ async fn responses_compact_should_post_json_to_codex_compact_upstream() {
"stream": true,
"use_websocket": false,
"store": true,
"prompt_cache_key": "must_not_forward"
"prompt_cache_key": "session-seed"
})
.to_string(),
))
@@ -193,24 +193,27 @@ async fn responses_compact_should_post_json_to_codex_compact_upstream() {
assert_eq!(upstream_body["model"], "gpt-5.5");
assert_eq!(upstream_body["instructions"], "compress the session");
assert_eq!(upstream_body["parallel_tool_calls"], false);
// compact 原样透传 reasoning,含未知键。
assert_eq!(
upstream_body["reasoning"],
json!({"effort": "high", "summary": "auto"})
json!({"effort": "high", "summary": "auto", "extra": "drop"})
);
assert_eq!(
upstream_body["tools"],
json!([{"type": "function", "name": "lookup"}])
);
assert_eq!(upstream_body["text"]["format"]["type"], "json_schema");
// compact 只剥离 stream(端点不支持流式);store/prompt_cache_key 等业务字段原样透传。
assert!(upstream_body.get("stream").is_none());
assert!(upstream_body.get("store").is_none());
assert!(upstream_body.get("prompt_cache_key").is_none());
assert_eq!(upstream_body["input"].as_array().unwrap().len(), 3);
assert!(upstream_body["input"][1].get("ignored").is_none());
assert_eq!(upstream_body["store"], true);
assert_eq!(upstream_body["prompt_cache_key"], "session-seed");
assert_eq!(upstream_body["input"].as_array().unwrap().len(), 4);
assert_eq!(upstream_body["input"][1]["ignored"], "drop");
assert_eq!(
upstream_body["input"][2]["encrypted_content"],
"enc_compact"
);
assert_eq!(upstream_body["input"][3]["id"], "drop_missing_encrypted");
let event = latest_response_usage_record(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.request_id.as_deref(), Some("req_compact_success_log"));
@@ -111,7 +111,6 @@ where
.await
}
const CHAT_SUCCESS_SSE: &str = concat!(include_str!("../../../fixtures/chat/success.sse"), "\n");
const RESPONSES_SUCCESS_SSE: &str = concat!(
include_str!("../../../fixtures/responses/http_sse/success.sse"),
"\n"
@@ -234,7 +233,6 @@ const WEBSOCKET_UNANSWERED_FUNCTION_CALL: &str =
const REASONING_REPLAY_REQUEST_GOLDEN: &str =
include_str!("../../../fixtures/responses/golden/reasoning_replay_request.json");
mod chat_routes;
mod compact_routes;
mod responses_http;
mod responses_recovery;
@@ -283,7 +281,7 @@ async fn test_app_state_with_pool_installation_id_and_usage_record_options(
.initialize_hot_path_state()
.await
.expect("hot path state should initialize");
AppState { config, services }
AppState { services }
}
async fn test_app_with_account(base_url: String) -> (axum::Router, String, tempfile::TempDir) {
@@ -435,35 +433,6 @@ async fn test_app_with_account_pool_config(
(router::router().with_state(state), api_key, pool, dir)
}
async fn test_app_with_restored_pool_then_disabled_account(
base_url: String,
) -> (axum::Router, String, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("openai-chat-restored-pool.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let api_key = insert_client_api_key(&pool).await;
insert_account(&pool).await;
let state = test_app_state_with_pool_and_installation_id(
&test_config(url, base_url),
pool.clone(),
TEST_INSTALLATION_ID.to_string(),
)
.await;
state
.services
.account_pool
.restore_from_repository()
.await
.expect("account pool should restore");
sqlx::query("update accounts set status = 'disabled' where id = ?")
.bind("acct_chat")
.execute(&pool)
.await
.unwrap();
(router::router().with_state(state), api_key, dir)
}
async fn test_app_with_two_accounts_and_affinity(
base_url: String,
) -> (axum::Router, String, tempfile::TempDir) {
@@ -251,7 +251,8 @@ async fn responses_should_use_websocket_upstream_by_default_while_serving_sse()
assert!(body.contains("\"id\":\"resp_ws_default\""));
assert_eq!(captured.payload["type"], "response.create");
assert_eq!(captured.payload["model"], "gpt-5.5");
assert!(captured.payload.get("generate").is_none());
// 透明代理:`generate` 原样透传上游,不再被剥离。
assert_eq!(captured.payload["generate"], false);
assert!(captured.payload.get("previous_response_id").is_none());
assert!(captured.payload["prompt_cache_key"]
.as_str()
@@ -328,14 +329,18 @@ async fn responses_should_ignore_camel_case_use_websocket_field() {
.unwrap();
let captured = upstream.await.unwrap();
// 透明代理:只有 snake_case `use_websocket` 是代理 transport 提示并被剥离;
// camelCase `useWebSocket` 是未知字段,不影响传输决策(请求仍走 WebSocket),
// 且原样透传上游。
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(captured.payload["type"], "response.create");
assert!(captured.payload.get("useWebSocket").is_none());
assert_eq!(captured.payload["useWebSocket"], false);
assert!(captured.payload.get("use_websocket").is_none());
}
#[tokio::test]
async fn responses_should_convert_tuple_schema_before_upstream() {
async fn responses_should_forward_text_schema_verbatim_to_upstream() {
// 透明代理:`text.format.schema` 原样透传上游,不再把 tuple prefixItems 转成 object。
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
@@ -368,24 +373,23 @@ async fn responses_should_convert_tuple_schema_before_upstream() {
let requests = server.received_requests().await.unwrap();
let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap();
let schema = &upstream_body["text"]["format"]["schema"];
assert!(schema["properties"]["point"].get("prefixItems").is_none());
assert!(schema["properties"]["point"].get("items").is_none());
assert_eq!(
schema["properties"]["point"],
json!({
"type": "object",
"properties": {
"0": {"type": "number"},
"1": {"type": "number"}
},
"required": ["0", "1"],
"additionalProperties": false
"type": "array",
"prefixItems": [
{"type": "number"},
{"type": "number"}
],
"items": false
})
);
}
#[tokio::test]
async fn responses_should_reconvert_tuple_schema_output_for_client() {
async fn responses_should_pass_through_tuple_schema_output_verbatim() {
// 透明代理:不再重写请求 schema,因此也不再对响应做 tuple 回转换——
// 上游返回的 object 形态 `{"point":{"0":1,"1":2}}` 原样透传给客户端。
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
@@ -416,12 +420,16 @@ async fn responses_should_reconvert_tuple_schema_output_for_client() {
assert_eq!(status, StatusCode::OK);
assert_eq!(body["id"], "resp_tuple");
assert_eq!(body["output_text"], "{\"point\":[1,2]}");
assert_eq!(body["output"][0]["content"][0]["text"], "{\"point\":[1,2]}");
assert_eq!(body["output_text"], "{\"point\":{\"0\":1,\"1\":2}}");
assert_eq!(
body["output"][0]["content"][0]["text"],
"{\"point\":{\"0\":1,\"1\":2}}"
);
}
#[tokio::test]
async fn responses_stream_should_reconvert_tuple_schema_output_for_client() {
async fn responses_stream_should_pass_through_tuple_schema_output_verbatim() {
// 透明代理:流式响应同样原样透传上游 object 形态,不做 tuple 回转换。
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
@@ -451,9 +459,9 @@ async fn responses_stream_should_reconvert_tuple_schema_output_for_client() {
let body = response_text(response).await;
assert_eq!(status, StatusCode::OK);
assert!(body.contains(r#""delta":"{\"point\":[1,2]}""#));
assert!(body.contains(r#""output_text":"{\"point\":[1,2]}""#));
assert!(!body.contains(r#""point\":{\"0":1,"1":2}"#));
// 透明代理 HTTP SSE 为字节透传:上游 object 形态的 delta 与 output item 原样出现。
assert!(body.contains(r#""delta":"{\"point\":{\"0\":1,\"1\":2}}""#));
assert!(body.contains(r#""text":"{\"point\":{\"0\":1,\"1\":2}}""#));
assert!(body.ends_with("data: [DONE]\n\n"));
}
@@ -486,10 +494,12 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
.header("x-openai-subagent", "review")
.body(Body::from(
json!({
"model": "gpt-5.5-high-fast",
"model": "gpt-5.5",
"stream": false,
"use_websocket": false,
"input": [],
"reasoning": {"effort": "high"},
"service_tier": "priority",
"prompt_cache_key": "pcache",
"client_metadata": {
"safe": "yes",
@@ -521,8 +531,9 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
.find(|request| request.url.path() == "/codex/responses")
.expect("responses upstream request should be sent");
let upstream_body: Value = serde_json::from_slice(&upstream.body).unwrap();
let identity =
build_conversation_identity(Some("pcache"), Some("window-metadata"), "acct_chat");
// 透明代理:window-id 不再从 client_metadata 回退读取,客户端未通过 body
// `codexWindowId` 或 `x-codex-window-id` 请求头提供时,window-id 由会话派生(`:0` 后缀)。
let identity = build_conversation_identity(Some("pcache"), None, "acct_chat");
let conversation_id = identity
.conversation_id
.as_deref()
@@ -538,16 +549,11 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
.and_then(|value| value.to_str().ok())
};
// 透明代理:model/reasoning/service_tier 原样透传,不做后缀路由、不注入默认值。
assert_eq!(upstream_body["model"], "gpt-5.5");
assert_eq!(upstream_body["service_tier"], "priority");
assert_eq!(
upstream_body["reasoning"],
json!({"summary": "auto", "effort": "high"})
);
assert_eq!(
upstream_body["include"],
json!(["reasoning.encrypted_content"])
);
assert_eq!(upstream_body["reasoning"], json!({"effort": "high"}));
assert!(upstream_body.get("include").is_none());
assert_eq!(upstream_body["prompt_cache_key"], conversation_id);
assert_eq!(upstream_body["client_metadata"]["safe"], "yes");
assert_eq!(
@@ -583,37 +589,32 @@ async fn responses_should_forward_parity_fields_context_headers_and_account_scop
assert_eq!(upstream_header("session_id"), None);
assert_eq!(upstream_header("x-codex-window-id"), Some(window_id));
assert_eq!(upstream_header("x-codex-turn-state"), Some("turn-body"));
assert_eq!(
upstream_header("x-codex-turn-metadata"),
Some("meta-metadata")
);
// 透明代理:turnMetadata/parentThreadId 仅出现在 client_metadata 中(旧的
// client_metadata→context fallback 已移除),不再落到请求头。
assert_eq!(upstream_header("x-codex-turn-metadata"), None);
assert_eq!(upstream_header("x-codex-beta-features"), Some("beta-body"));
assert_eq!(
upstream_header("x-responsesapi-include-timing-metrics"),
Some("true")
);
assert_eq!(upstream_header("version"), Some("2026-06-12"));
assert_eq!(
upstream_header("x-codex-parent-thread-id"),
Some("parent-metadata")
);
assert_eq!(upstream_header("x-codex-parent-thread-id"), None);
assert_eq!(upstream_header("x-openai-subagent"), Some("review"));
for local_field in [
"turnState",
"turnMetadata",
"betaFeatures",
"includeTimingMetrics",
"version",
"codexWindowId",
"parentThreadId",
"use_websocket",
// 透明代理:上下文透传字段提取到代理控制状态用于加请求头,body 中同时原样保留。
for context_field in [
("turnState", "turn-body"),
("betaFeatures", "beta-body"),
("includeTimingMetrics", "true"),
("version", "2026-06-12"),
] {
assert!(upstream_body.get(local_field).is_none());
assert_eq!(upstream_body[context_field.0], context_field.1);
}
// transport-only 字段不进上游 body。
assert!(upstream_body.get("use_websocket").is_none());
}
#[tokio::test]
async fn responses_should_preserve_non_empty_include_when_reasoning_defaults_apply() {
async fn responses_should_pass_through_reasoning_and_include_verbatim() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
@@ -636,12 +637,12 @@ async fn responses_should_preserve_non_empty_include_when_reasoning_defaults_app
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "gpt-5.5-high",
"model": "gpt-5.5",
"stream": false,
"use_websocket": false,
"input": [],
"include": ["file_search_call.results"],
"use_websocket": false
"reasoning": {"effort": "high"},
"include": ["file_search_call.results"]
})
.to_string(),
))
@@ -653,10 +654,8 @@ async fn responses_should_preserve_non_empty_include_when_reasoning_defaults_app
assert_eq!(response.status(), StatusCode::OK);
let requests = server.received_requests().await.unwrap();
let upstream_body: Value = serde_json::from_slice(&requests[0].body).unwrap();
assert_eq!(
upstream_body["reasoning"],
json!({"summary": "auto", "effort": "high"})
);
// 透明代理:reasoning 与 include 原样透传上游,不改写、不注入。
assert_eq!(upstream_body["reasoning"], json!({"effort": "high"}));
assert_eq!(
upstream_body["include"],
json!(["file_search_call.results"])
@@ -664,7 +663,7 @@ async fn responses_should_preserve_non_empty_include_when_reasoning_defaults_app
}
#[tokio::test]
async fn responses_should_sanitize_reasoning_and_compaction_input_before_upstream() {
async fn responses_should_preserve_input_items_before_upstream() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
@@ -737,16 +736,30 @@ async fn responses_should_sanitize_reasoning_and_compaction_input_before_upstrea
{
"type": "reasoning",
"id": "rs_1",
"summary": [{"type": "summary_text", "text": "valid summary"}],
"status": "completed",
"summary": [
{"type": "summary_text", "text": "valid summary"},
{"type": "ignored", "text": "drop"}
],
"encrypted_content": "enc_reasoning",
"content": [{"type": "reasoning_text", "text": "valid reasoning"}]
"content": [
{"type": "reasoning_text", "text": "valid reasoning"},
{"type": "ignored", "text": "drop"}
],
"extra": "drop"
},
{
"type": "reasoning",
"id": "",
"summary": [{"type": "summary_text", "text": "drop"}]
},
{
"type": "compaction",
"id": "cmp_1",
"encrypted_content": "enc_compaction",
"id": "cmp_1"
"extra": "drop"
},
{"type": "compaction", "id": "cmp_drop"},
{"type": "message", "role": "user", "content": "keep me", "extra": 42}
])
);
@@ -133,6 +133,92 @@ async fn responses_websocket_stream_should_synthesize_response_failed_when_close
assert!(body.contains("stream_disconnected"));
}
#[tokio::test]
async fn responses_websocket_stream_first_error_429_should_retry_fallback_account() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base_url = format!("http://{}", listener.local_addr().unwrap());
let upstream = tokio::spawn(async move {
let (first_stream, _) = listener.accept().await.unwrap();
let mut first_websocket =
accept_websocket_with_authorization(first_stream, "Bearer access-primary").await;
let first_payload = send_websocket_response_and_capture_payload(
&mut first_websocket,
json!({
"type": "error",
"error": {
"type": "usage_limit_reached",
"message": "The usage limit has been reached",
"retry_after_seconds": 77
},
"status_code": 429
})
.to_string(),
)
.await;
first_websocket.close(None).await.unwrap();
let second_payload = accept_successful_websocket_response_with_authorization(
&listener,
"Bearer access-secondary",
"resp_after_ws_stream_rate_limit",
)
.await;
(first_payload, second_payload)
});
let (app, api_key, pool, _dir) = test_app_with_two_accounts(base_url).await;
let response = app
.oneshot(responses_json_request(
&api_key,
&json!({
"model": "gpt-5.5",
"input": [],
"stream": true,
"use_websocket": true
}),
))
.await
.unwrap();
let status = response.status();
let content_type = response
.headers()
.get("content-type")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string();
let body = response_text(response).await;
let (first_payload, second_payload) = upstream.await.unwrap();
let primary_quota_state: (i64, Option<String>) = sqlx::query_as(
"select quota_limit_reached, quota_cooldown_until from accounts where id = ?",
)
.bind("acct_primary")
.fetch_one(&pool)
.await
.unwrap();
let primary_usage: (i64,) =
sqlx::query_as("select request_count from account_usage where account_id = ?")
.bind("acct_primary")
.fetch_one(&pool)
.await
.unwrap();
let secondary_usage: (i64,) =
sqlx::query_as("select request_count from account_usage where account_id = ?")
.bind("acct_secondary")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(status, StatusCode::OK);
assert!(content_type.starts_with("text/event-stream"));
assert!(body.contains("resp_after_ws_stream_rate_limit"));
assert_eq!(first_payload["type"], "response.create");
assert_eq!(second_payload["type"], "response.create");
assert_eq!(primary_quota_state.0, 1);
assert!(primary_quota_state.1.is_some());
assert_eq!(primary_usage.0, 1);
assert_eq!(secondary_usage.0, 1);
}
#[tokio::test]
async fn responses_websocket_should_reuse_connection_for_recorded_conversation() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -677,6 +763,7 @@ async fn responses_websocket_should_implicitly_resume_full_history_with_reasonin
second_payload["input"][0]["encrypted_content"],
"enc_reasoning_replay"
);
assert!(second_payload["input"][0].get("content").is_none());
assert_eq!(second_payload["input"][1]["content"], "continue");
assert_eq!(second_payload["input"].as_array().unwrap().len(), 2);
@@ -86,58 +86,6 @@ async fn responses_should_use_imported_account_record_usage_cookie_and_usage_rec
assert_eq!(metadata["usage"]["inputTokens"], 7);
}
#[tokio::test]
async fn chat_completions_non_stream_should_record_first_token_latency() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/codex/responses"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(RESPONSES_COMPLETED_USAGE_SSE),
)
.expect(1)
.mount(&server)
.await;
let (app, api_key, pool, _dir) = test_app_with_account_pool_and_logging(server.uri()).await;
let response = app
.clone()
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.header("x-request-id", "req_chat_non_stream_first_token")
.body(Body::from(
json!({
"model": "gpt-5.5",
"messages": [{"role": "user", "content": "Say hello"}]
})
.to_string(),
))
.unwrap(),
)
.await
.unwrap();
let status = response.status();
let body = response_json(response).await;
let event = latest_usage_record(&pool, "v1.chat").await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(status, StatusCode::OK);
assert_eq!(body["usage"]["prompt_tokens"], 7);
assert_eq!(metadata["stream"], false);
assert_eq!(metadata["transport"], "http_sse");
assert!(
metadata["firstTokenMs"]
.as_i64()
.is_some_and(|value| value > 0),
"chat non-stream usage metadata should include first token latency: {metadata:?}",
);
}
#[tokio::test]
async fn responses_usage_record_should_store_resolved_upstream_model_after_alias_mapping() {
let server = MockServer::start().await;
@@ -170,7 +170,6 @@ async fn build_test_app(
.await
.expect("hot path state should initialize");
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -111,24 +111,6 @@ async fn responses_route_should_default_omitted_stream_to_sse() {
assert!(body.ends_with("data: [DONE]\n\n"));
}
#[tokio::test]
async fn chat_completions_route_should_reject_empty_messages_with_bad_request() {
let (app, api_key, _dir) = test_app_with_client_api_key().await;
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header("authorization", format!("Bearer {api_key}"))
.header("content-type", "application/json")
.body(Body::from(r#"{"model":"gpt-5.5","messages":[]}"#))
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}
async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("openai-responses-routes.sqlite");
@@ -154,7 +136,6 @@ async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::Temp
.await
.expect("hot path state should initialize");
let state = AppState {
config,
services: (*services).clone(),
};
(
+8 -8
View File
@@ -79,14 +79,6 @@ impl AccountStore for FakeAccountStore {
Ok(None)
}
async fn update_quota_json(
&self,
_account_id: &str,
_quota_json: &str,
) -> AccountStoreResult<bool> {
Ok(false)
}
async fn apply_quota_snapshot(
&self,
_account_id: &str,
@@ -97,6 +89,14 @@ impl AccountStore for FakeAccountStore {
Ok(false)
}
async fn sync_runtime_account_state(
&self,
_account: &Account,
_sync_usage_window: bool,
) -> AccountStoreResult<bool> {
Ok(false)
}
async fn sync_rate_limit_window(
&self,
_account_id: &str,
+8 -8
View File
@@ -192,14 +192,6 @@ impl AccountStore for PlanAccountStore {
Ok(None)
}
async fn update_quota_json(
&self,
_account_id: &str,
_quota_json: &str,
) -> AccountStoreResult<bool> {
Ok(false)
}
async fn apply_quota_snapshot(
&self,
_account_id: &str,
@@ -210,6 +202,14 @@ impl AccountStore for PlanAccountStore {
Ok(false)
}
async fn sync_runtime_account_state(
&self,
_account: &Account,
_sync_usage_window: bool,
) -> AccountStoreResult<bool> {
Ok(false)
}
async fn sync_rate_limit_window(
&self,
_account_id: &str,
@@ -210,6 +210,150 @@ async fn runtime_account_pool_should_persist_expired_status_when_jwt_expiry_is_d
assert_eq!(status.0, "expired");
}
#[tokio::test]
async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_expires() {
let now = fixed_time();
let (pool, _dir) =
crate::support::sqlite::init_test_db("runtime-pool-quota-window-reset.sqlite").await;
let store = codex_proxy_rs::upstream::accounts::store::SqliteAccountStore::new(pool.clone());
store
.insert(codex_proxy_rs::upstream::accounts::store::NewAccount {
id: "acct_quota_reset".to_string(),
email: None,
account_id: Some("chatgpt-quota-reset".to_string()),
user_id: None,
label: None,
plan_type: Some("plus".to_string()),
access_token: secrecy::SecretString::new("access-token".to_string().into()),
refresh_token: None,
access_token_expires_at: Some(now + Duration::hours(1)),
status: AccountStatus::QuotaExhausted,
added_at: None,
})
.await
.unwrap();
sqlx::query(
r#"
update accounts
set
quota_limit_reached = 1,
quota_verify_required = 0,
quota_cooldown_until = ?,
cloudflare_cooldown_until = ?
where id = ?
"#,
)
.bind((now - Duration::seconds(1)).to_rfc3339())
.bind((now - Duration::seconds(1)).to_rfc3339())
.bind("acct_quota_reset")
.execute(&pool)
.await
.unwrap();
sqlx::query(
r#"
insert into account_usage (
account_id,
request_count,
window_request_count,
window_input_tokens,
window_output_tokens,
window_cached_tokens,
window_started_at,
window_reset_at,
limit_window_seconds
) values (?, 9, 7, 11, 13, 17, ?, ?, 60)
"#,
)
.bind("acct_quota_reset")
.bind((now - Duration::seconds(120)).to_rfc3339())
.bind((now - Duration::seconds(1)).to_rfc3339())
.execute(&pool)
.await
.unwrap();
let runtime_pool = codex_proxy_rs::upstream::accounts::pool::RuntimeAccountPoolService::new(
std::sync::Arc::new(store),
AccountPoolOptions::default(),
0,
);
runtime_pool.restore_from_repository().await.unwrap();
let acquired = runtime_pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.await
.unwrap();
type StoredQuotaResetState = (
String,
i64,
i64,
Option<String>,
Option<String>,
i64,
i64,
i64,
i64,
Option<String>,
Option<i64>,
);
let (
stored_status,
stored_quota_limit_reached,
stored_quota_verify_required,
stored_quota_cooldown_until,
stored_cloudflare_cooldown_until,
stored_window_request_count,
stored_window_input_tokens,
stored_window_output_tokens,
stored_window_cached_tokens,
stored_window_reset_at,
stored_limit_window_seconds,
): StoredQuotaResetState = sqlx::query_as(
r#"
select
a.status,
a.quota_limit_reached,
a.quota_verify_required,
a.quota_cooldown_until,
a.cloudflare_cooldown_until,
au.window_request_count,
au.window_input_tokens,
au.window_output_tokens,
au.window_cached_tokens,
au.window_reset_at,
au.limit_window_seconds
from accounts a
left join account_usage au on au.account_id = a.id
where a.id = ?
"#,
)
.bind("acct_quota_reset")
.fetch_one(&pool)
.await
.unwrap();
let persisted_reset_at = stored_window_reset_at
.as_deref()
.and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok())
.map(|value| value.with_timezone(&Utc));
assert_eq!(acquired.account.id, "acct_quota_reset");
assert!(acquired.account.quota_verify_required);
assert_eq!(stored_status, "active");
assert_eq!(stored_quota_limit_reached, 0);
assert_eq!(stored_quota_verify_required, 1);
assert!(stored_quota_cooldown_until.is_none());
assert!(stored_cloudflare_cooldown_until.is_none());
assert_eq!(
(
stored_window_request_count,
stored_window_input_tokens,
stored_window_output_tokens,
stored_window_cached_tokens,
),
(0, 0, 0, 0)
);
assert!(persisted_reset_at.is_some_and(|reset_at| reset_at > now));
assert_eq!(stored_limit_window_seconds, Some(60));
}
#[test]
fn account_pool_should_prefer_configured_tier_priority() {
let mut pool = AccountPool::with_options(AccountPoolOptions {
@@ -518,6 +518,99 @@ async fn account_repository_should_map_quota_snapshot_to_account_status() {
assert!(restored_status.2.is_none());
}
#[tokio::test]
async fn account_repository_runtime_state_sync_should_not_clear_newer_future_quota_cooldown() {
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 42).await;
let repo = SqliteAccountStore::new(pool.clone());
seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await;
let future_cooldown = Utc::now() + Duration::hours(1);
sqlx::query(
"update accounts set quota_limit_reached = 1, quota_verify_required = 0, quota_cooldown_until = ? where id = ?",
)
.bind(future_cooldown.to_rfc3339())
.bind("acct_a")
.execute(&pool)
.await
.unwrap();
let mut stale_account = AccountStore::list_pool_accounts(&repo)
.await
.unwrap()
.into_iter()
.find(|account| account.id == "acct_a")
.unwrap();
stale_account.quota_limit_reached = false;
stale_account.quota_verify_required = true;
stale_account.quota_cooldown_until = None;
repo.sync_runtime_account_state(&stale_account, false)
.await
.unwrap();
let stored: (i64, i64, Option<String>) = sqlx::query_as(
"select quota_limit_reached, quota_verify_required, quota_cooldown_until from accounts where id = ?",
)
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(stored.0, 1);
assert_eq!(stored.1, 0);
assert_eq!(stored.2, Some(future_cooldown.to_rfc3339()));
}
#[tokio::test]
async fn account_repository_runtime_state_sync_should_not_regress_newer_usage_window() {
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 43).await;
let repo = SqliteAccountStore::new(pool.clone());
seed_repo_account(&pool, "acct_a", "2026-06-11T00:00:00Z").await;
let now = Utc::now();
let future_reset = now + Duration::hours(1);
let stale_reset = now + Duration::minutes(1);
sqlx::query(
r#"
insert into account_usage (
account_id,
window_request_count,
window_input_tokens,
window_reset_at,
limit_window_seconds
) values (?, 5, 17, ?, 3600)
"#,
)
.bind("acct_a")
.bind(future_reset.to_rfc3339())
.execute(&pool)
.await
.unwrap();
let mut stale_account = AccountStore::list_pool_accounts(&repo)
.await
.unwrap()
.into_iter()
.find(|account| account.id == "acct_a")
.unwrap();
stale_account.window_request_count = 0;
stale_account.window_input_tokens = 0;
stale_account.window_reset_at = Some(stale_reset);
stale_account.window_started_at = Some(now);
stale_account.limit_window_seconds = Some(60);
repo.sync_runtime_account_state(&stale_account, true)
.await
.unwrap();
let stored: (i64, i64, Option<String>, Option<i64>) = sqlx::query_as(
"select window_request_count, window_input_tokens, window_reset_at, limit_window_seconds from account_usage where account_id = ?",
)
.bind("acct_a")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(stored.0, 5);
assert_eq!(stored.1, 17);
assert_eq!(stored.2, Some(future_reset.to_rfc3339()));
assert_eq!(stored.3, Some(3600));
}
#[tokio::test]
async fn account_repository_should_not_override_disabled_status_from_quota_snapshot() {
let (pool, _dir) = sqlite_account_store_parts("accounts.sqlite", 38).await;
+8 -31
View File
@@ -12,7 +12,7 @@ use codex_proxy_rs::upstream::accounts::model::{Account, AccountStatus};
use codex_proxy_rs::upstream::models::backend_entry::{
BackendModelEntry, BackendReasoningEffort, BackendTruncationPolicy,
};
use codex_proxy_rs::upstream::models::catalog::{ModelCatalog, ParsedModelName};
use codex_proxy_rs::upstream::models::catalog::ModelCatalog;
use codex_proxy_rs::upstream::models::config::ModelConfig;
use codex_proxy_rs::upstream::models::service::{
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
@@ -26,13 +26,12 @@ use codex_proxy_rs::upstream::transport::{
};
#[test]
fn model_catalog_should_parse_alias_reasoning_and_service_tier_suffixes() {
fn model_catalog_should_resolve_alias_chain_to_model_id() {
let mut model_aliases = BTreeMap::new();
model_aliases.insert("codex-fast".to_string(), "gpt-5.5".to_string());
let catalog = ModelCatalog::from_config(&ModelConfig { model_aliases });
let parsed = catalog.parse_model_name("codex-fast-high-flex");
// 透明代理:只解析 alias,不再拆分 `-high`/`-flex` 等后缀。
assert!(catalog
.public_model_ids()
.contains(&"codex-fast".to_string()));
@@ -40,40 +39,18 @@ fn model_catalog_should_parse_alias_reasoning_and_service_tier_suffixes() {
catalog.model_info_for_name("codex-fast").unwrap().id,
"gpt-5.5"
);
assert_eq!(
parsed,
ParsedModelName {
model_id: "gpt-5.5".to_string(),
reasoning_effort: Some("high".to_string()),
service_tier: Some("flex".to_string())
}
);
assert_eq!(catalog.resolve_model_id("codex-fast"), "gpt-5.5");
}
#[test]
fn model_catalog_should_keep_unknown_model_names_without_fallback() {
fn model_catalog_should_keep_unknown_model_names_verbatim() {
let catalog = ModelCatalog::from_config(&ModelConfig {
model_aliases: BTreeMap::new(),
});
let parsed = catalog.parse_model_name("unknown-medium");
assert_eq!(parsed.model_id, "unknown");
assert_eq!(parsed.reasoning_effort, Some("medium".to_string()));
}
#[test]
fn model_catalog_should_parse_none_and_minimal_reasoning_suffixes() {
let catalog = ModelCatalog::from_config(&ModelConfig {
model_aliases: BTreeMap::new(),
});
let none = catalog.parse_model_name("gpt-5.5-none");
let minimal = catalog.parse_model_name("gpt-5.5-minimal-fast");
assert_eq!(none.reasoning_effort, Some("none".to_string()));
assert_eq!(minimal.reasoning_effort, Some("minimal".to_string()));
assert_eq!(minimal.service_tier, Some("fast".to_string()));
// 无 alias 命中时原样返回,不再剥离后缀。
assert_eq!(catalog.resolve_model_id("unknown-medium"), "unknown-medium");
assert_eq!(catalog.resolve_model_id("gpt-5.5-high"), "gpt-5.5-high");
}
#[test]
@@ -10,25 +10,19 @@ fn codex_websocket_payload_audit_snapshot_should_redact_user_content() {
"content": "private prompt",
})],
);
request.previous_response_id = Some("resp_secret".to_string());
request.reasoning = Some(json!({"effort": "medium"}));
request.tools = Some(vec![json!({
"type": "function",
"name": "private_tool",
"description": "private schema",
})]);
request.text = Some(json!({"format": {"type": "text"}}));
request.service_tier = Some("flex".to_string());
request.prompt_cache_key = Some("cache-secret".to_string());
request.generate = Some(false);
request.include = Some(vec!["reasoning.encrypted_content".to_string()]);
request.client_metadata = Some(json!({
request.set_previous_response_id(Some("resp_secret".to_string()));
request.set_reasoning(Some(json!({"effort": "medium"})));
request.set_service_tier(Some("flex".to_string()));
request.set_prompt_cache_key(Some("cache-secret".to_string()));
request.set_client_metadata(Some(json!({
"thread_id": "thread-secret",
"safe": "value",
}));
})));
let snapshot = websocket_payload_audit_snapshot(&request);
// 透明代理:payload key 顺序即 `type` 前置到原始 body 的插入顺序,
// 不再重排为固定 canonical 顺序,也不注入 tool_choice/parallel_tool_calls 等默认字段。
assert_eq!(
snapshot.top_level_keys,
vec![
@@ -36,17 +30,12 @@ fn codex_websocket_payload_audit_snapshot_should_redact_user_content() {
"model",
"instructions",
"input",
"store",
"stream",
"store",
"previous_response_id",
"reasoning",
"tools",
"tool_choice",
"parallel_tool_calls",
"text",
"service_tier",
"prompt_cache_key",
"include",
"client_metadata",
]
);
@@ -58,7 +47,6 @@ fn codex_websocket_payload_audit_snapshot_should_redact_user_content() {
assert_eq!(snapshot.body["previous_response_id"], "<redacted>");
assert_eq!(snapshot.body["prompt_cache_key"], "<redacted>");
assert_eq!(snapshot.body["client_metadata"], "<redacted>");
assert_eq!(snapshot.body["tools"], "<redacted>");
}
#[test]
@@ -71,12 +59,11 @@ fn codex_websocket_response_create_payload_text_should_preserve_canonical_field_
"content": "private capture prompt",
})],
);
request.prompt_cache_key = Some("session-1".to_string());
request.generate = Some(false);
request.client_metadata = Some(json!({
request.set_prompt_cache_key(Some("session-1".to_string()));
request.set_client_metadata(Some(json!({
"thread_id": "capture-thread-secret",
"safe": "capture",
}));
})));
let payload =
codex_proxy_rs::upstream::protocol::websocket::websocket_response_create_payload_text(
@@ -84,6 +71,8 @@ fn codex_websocket_response_create_payload_text_should_preserve_canonical_field_
)
.expect("payload should serialize");
// 透明代理:payload 即 `type` 前置到原始 body 的插入顺序,字段原样透传,
// 不注入 tool_choice/parallel_tool_calls 默认值。
assert_substrings_appear_in_order(
&payload,
&[
@@ -91,10 +80,8 @@ fn codex_websocket_response_create_payload_text_should_preserve_canonical_field_
"\"model\":\"gpt-5.5\"",
"\"instructions\":\"private capture instructions\"",
"\"input\":",
"\"store\":false",
"\"stream\":true",
"\"tool_choice\":\"auto\"",
"\"parallel_tool_calls\":true",
"\"store\":false",
"\"prompt_cache_key\":\"session-1\"",
"\"client_metadata\":",
],
@@ -128,7 +115,7 @@ fn codex_websocket_audit_artifact_should_record_transport_opening_and_payload()
})],
);
request.force_http_sse = false;
request.previous_response_id = Some("resp_secret".to_string());
request.set_previous_response_id(Some("resp_secret".to_string()));
let opening = OpeningAuditSnapshot {
request_line: "GET /backend-api/codex/responses HTTP/1.1".to_string(),
header_order: vec!["Host".to_string(), "authorization".to_string()],
@@ -154,7 +141,6 @@ fn codex_websocket_audit_artifact_should_record_transport_opening_and_payload()
artifact.payload.expect("payload").body["input"],
"<redacted>"
);
assert_eq!(artifact.error, None);
}
#[test]
@@ -907,7 +893,9 @@ fn codex_websocket_reasoning_summary_part_should_require_summary_index() {
}
#[test]
fn codex_websocket_event_to_sse_frame_should_skip_invalid_json_and_shape_mismatches() {
fn codex_websocket_event_to_sse_frame_should_forward_typed_events_without_schema_filtering() {
// 透明代理:WS 输出侧不再按官方 schema 校验丢弃事件。只有无法确定 type 的帧
// (非 JSON / 缺 type)和内部事件(codex.rate_limits)会被丢弃,其余原样转发。
let invalid_delta = json!({
"type": "response.output_text.delta",
"delta": 42
@@ -919,7 +907,17 @@ fn codex_websocket_event_to_sse_frame_should_skip_invalid_json_and_shape_mismatc
})
.to_string();
// 无 type / 非 JSON:无法编码,丢弃。
assert!(websocket_event_to_sse_frame("not-json-from-upstream").is_none());
assert!(websocket_event_to_sse_frame(&invalid_delta).is_none());
assert!(websocket_event_to_sse_frame(&non_object_item).is_none());
// 内部事件:transport 剥离,丢弃。
assert!(websocket_event_to_sse_frame(r#"{"type":"codex.rate_limits"}"#).is_none());
// 有 type 的畸形事件:原样透传,不再做 schema 过滤。
let invalid_delta_frame =
websocket_event_to_sse_frame(&invalid_delta).expect("typed event should be forwarded");
assert!(invalid_delta_frame.contains("event: response.output_text.delta"));
assert!(invalid_delta_frame.contains(&invalid_delta));
let non_object_item_frame =
websocket_event_to_sse_frame(&non_object_item).expect("typed event should be forwarded");
assert!(non_object_item_frame.contains("event: response.output_item.done"));
assert!(non_object_item_frame.contains(&non_object_item));
}

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