Compare commits

...
16 Commits
Author SHA1 Message Date
ZyyandCodex 3c9ec50bf2 chore: bump version to 1.0.10
Co-authored-by: Codex <noreply@openai.com>
2026-07-09 13:08:08 +08:00
ZyyandCodex c1b81809a6 refactor: derive account refresh status from leases
Co-authored-by: Codex <noreply@openai.com>
2026-07-09 13:06:54 +08:00
Zyy 939f392438 docs: refresh architecture and database docs 2026-07-08 21:35:26 +08:00
ZyyandCodex ff0de4742d chore: bump version to 1.0.9
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 18:06:08 +08:00
ZyyandCodex 44f45101b8 test: assert account quota local usage
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 17:44:01 +08:00
ZyyandCodex a240668a47 fix: polish sidebar focus state
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 16:45:52 +08:00
ZyyandCodex 30f354965c fix: align account quota usage tracking
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 16:37:32 +08:00
ZyyandCodex 3b462f7ecf refactor: remove trusted proxy config
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 15:22:23 +08:00
ZyyandCodex 5682d4fbe9 test: simplify ops error log snapshot helper
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 14:55:05 +08:00
ZyyandCodex 146776c742 fix: separate ops errors from usage records
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 14:04:49 +08:00
ZyyandCodex e8c14d4324 fix: complete update log streaming
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 13:58:28 +08:00
zyycnandCodex 5483e37b00 refactor: extract upstream scheduler
Co-authored-by: Codex <noreply@openai.com>
2026-07-08 13:56:27 +08:00
ZyyandCodex daad7c746c fix: strengthen response account affinity
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 18:31:50 +08:00
ZyyandCodex 6e24037c4d chore: bump version to 1.0.8
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 17:06:08 +08:00
ZyyandCodex 95daf7093e feat: add smart account scheduling
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 17:00:56 +08:00
ZyyandCodex 2e121c0670 fix: hot reload fingerprint updates
Co-authored-by: Codex <noreply@openai.com>
2026-07-07 15:01:28 +08:00
116 changed files with 5583 additions and 1488 deletions
+56 -19
View File
@@ -7,7 +7,7 @@ use serde::Serialize;
use serde_json::Value;
use crate::infra::{
format::{format_compact_percent, format_tokens},
format::{format_compact_percent, format_plain_number, format_tokens, nonnegative_i64_to_u64},
time::{china_datetime, china_relative_time},
};
@@ -31,13 +31,29 @@ struct AdminAccountQuotaWindowData {
label_display: String,
used_percent: Option<f64>,
used_percent_display: String,
token_usage_display: String,
#[serde(skip_serializing_if = "Option::is_none")]
local_usage: Option<AdminAccountQuotaWindowLocalUsageData>,
reset_at_display: String,
window_used_display: String,
#[serde(skip)]
reset_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
struct AdminAccountQuotaWindowLocalUsageData {
request_count: i64,
request_count_display: String,
input_tokens: i64,
input_tokens_display: String,
output_tokens: i64,
output_tokens_display: String,
cached_tokens: i64,
cached_tokens_display: String,
total_tokens: i64,
total_tokens_display: String,
}
#[derive(Debug, Clone)]
pub(crate) struct AdminAccountQuotaUsageWindow {
pub(crate) key: String,
@@ -56,17 +72,37 @@ impl AdminAccountQuotaUsageWindow {
}
#[derive(Debug, Clone, Copy, Default)]
pub(crate) struct AdminAccountQuotaWindowTokenUsage {
total_tokens: u64,
pub(crate) struct AdminAccountQuotaWindowLocalUsage {
pub(crate) request_count: i64,
pub(crate) input_tokens: i64,
pub(crate) output_tokens: i64,
pub(crate) cached_tokens: i64,
}
impl AdminAccountQuotaWindowTokenUsage {
pub(crate) fn add_tokens(&mut self, tokens: u64) {
self.total_tokens = self.total_tokens.saturating_add(tokens);
impl AdminAccountQuotaWindowLocalUsage {
pub(crate) fn add(&mut self, other: Self) {
self.request_count = self.request_count.saturating_add(other.request_count);
self.input_tokens = self.input_tokens.saturating_add(other.input_tokens);
self.output_tokens = self.output_tokens.saturating_add(other.output_tokens);
self.cached_tokens = self.cached_tokens.saturating_add(other.cached_tokens);
}
}
fn display(self) -> String {
format_tokens(self.total_tokens)
impl From<AdminAccountQuotaWindowLocalUsage> for AdminAccountQuotaWindowLocalUsageData {
fn from(usage: AdminAccountQuotaWindowLocalUsage) -> Self {
let total_tokens = usage.input_tokens.saturating_add(usage.output_tokens);
Self {
request_count: usage.request_count,
request_count_display: format_plain_number(nonnegative_i64_to_u64(usage.request_count)),
input_tokens: usage.input_tokens,
input_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.input_tokens)),
output_tokens: usage.output_tokens,
output_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.output_tokens)),
cached_tokens: usage.cached_tokens,
cached_tokens_display: format_tokens(nonnegative_i64_to_u64(usage.cached_tokens)),
total_tokens,
total_tokens_display: format_tokens(nonnegative_i64_to_u64(total_tokens)),
}
}
}
@@ -93,17 +129,18 @@ impl AdminAccountQuotaData {
.collect()
}
pub(crate) fn apply_token_usage(
pub(crate) fn apply_local_usage(
&mut self,
usage_by_window: &HashMap<String, AdminAccountQuotaWindowTokenUsage>,
usage_by_window: &HashMap<String, AdminAccountQuotaWindowLocalUsage>,
) {
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(),
};
window.local_usage = window.usage_window().map(|_| {
usage_by_window
.get(&window.key)
.copied()
.unwrap_or_default()
.into()
});
}
}
}
@@ -185,7 +222,7 @@ fn push_monthly_quota_window(
label_display: "月限额".to_string(),
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
token_usage_display: "-".to_string(),
local_usage: None,
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
@@ -274,7 +311,7 @@ fn push_quota_window(
label_display,
used_percent,
used_percent_display: used_percent.map_or_else(|| "-".to_string(), format_compact_percent),
token_usage_display: "-".to_string(),
local_usage: None,
reset_at_display: reset_at
.as_ref()
.map_or_else(|| "-".to_string(), china_datetime),
+106 -132
View File
@@ -27,7 +27,7 @@ use crate::{
admin::{
accounts::quota_view::{
quota_data, AdminAccountQuotaData, AdminAccountQuotaUsageWindow,
AdminAccountQuotaWindowTokenUsage,
AdminAccountQuotaWindowLocalUsage,
},
accounts::service::{
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata,
@@ -135,6 +135,8 @@ struct AdminAccountData {
plan_type: Option<String>,
has_refresh_token: bool,
status: String,
display_status: String,
token_refreshing: bool,
access_token_expires_at: Option<String>,
access_token_expires_at_display: Option<String>,
added_at: String,
@@ -147,7 +149,7 @@ struct AdminAccountData {
impl From<AdminAccountMetadata> for AdminAccountData {
fn from(a: AdminAccountMetadata) -> Self {
Self::from_parts(a, None, None, Vec::new())
Self::from_parts(a, None, None, Vec::new(), false)
}
}
@@ -157,10 +159,12 @@ impl AdminAccountData {
usage: Option<&AdminUsageRecord>,
quota: Option<AdminAccountQuotaData>,
models: Vec<AdminAccountModelUsageData>,
token_refreshing: bool,
) -> Self {
let access_token_expires_at = a.access_token_expires_at.as_ref().map(china_rfc3339);
let access_token_expires_at_display =
a.access_token_expires_at.as_ref().map(china_datetime);
let display_status = account_display_status(a.status, token_refreshing).to_string();
Self {
id: a.id,
email: a.email,
@@ -170,6 +174,8 @@ impl AdminAccountData {
plan_type: a.plan_type,
has_refresh_token: a.has_refresh_token,
status: a.status.as_str().to_string(),
display_status,
token_refreshing,
access_token_expires_at,
access_token_expires_at_display,
added_at: china_rfc3339(&a.added_at),
@@ -289,6 +295,7 @@ struct AccountListStats {
usage_by_account: HashMap<String, AdminUsageRecord>,
quota_by_account: HashMap<String, AdminAccountQuotaData>,
models_by_account: HashMap<String, Vec<AdminAccountModelUsageData>>,
refreshing_account_ids: HashSet<String>,
}
struct AccountModelUsageRecord {
@@ -303,14 +310,6 @@ struct AccountModelUsageRecord {
last_used_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, Copy, Default)]
struct AccountUsageWindowStats {
request_count: i64,
input_tokens: i64,
output_tokens: i64,
cached_tokens: i64,
}
#[derive(Debug, Clone)]
struct AccountQuotaUsageWindow {
account_id: String,
@@ -328,6 +327,7 @@ impl AccountListStats {
.get(&account_id)
.cloned()
.unwrap_or_default(),
self.refreshing_account_ids.contains(&account_id),
)
}
}
@@ -531,7 +531,7 @@ pub(crate) async fn accounts(
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await;
let stats = account_list_stats(&state, &page.items, &quota_by_account).await?;
let page = crate::infra::json::NumberedPage {
items: page
.items
@@ -558,7 +558,7 @@ pub(crate) async fn accounts(
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await;
let stats = account_list_stats(&state, &page.items, &quota_by_account).await?;
let page = Page {
items: page
.items
@@ -691,7 +691,7 @@ pub(crate) async fn account_quota(
let raw = data.get("raw").cloned().unwrap_or(Value::Null);
let quota_json = quota.to_string();
let mut quota_data = quota_data(&quota_json, Some(Utc::now()));
apply_account_quota_window_token_usage(&state, &account_id, &mut quota_data).await;
apply_account_quota_window_local_usage(&state, &account_id, &mut quota_data).await;
let account =
account_data_for_quota_refresh(&state, &account_id, quota_data.clone()).await?;
Ok(AdminResponse::new(
@@ -728,7 +728,8 @@ async fn account_data_for_quota_refresh(
return Err(account_not_found());
};
let quota_by_account = HashMap::from([(account_id.to_string(), quota_data)]);
let stats = account_list_stats(state, std::slice::from_ref(&account), &quota_by_account).await;
let stats =
account_list_stats(state, std::slice::from_ref(&account), &quota_by_account).await?;
Ok(stats.data_for(account))
}
@@ -910,7 +911,7 @@ async fn account_list_stats(
state: &AppState,
accounts: &[AdminAccountMetadata],
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
) -> AccountListStats {
) -> Result<AccountListStats, AdminError> {
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
@@ -922,10 +923,11 @@ async fn account_list_stats(
.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 selected_quota_windows = selected_quota_windows_by_account(accounts, &quota_by_account);
let all_quota_windows = quota_usage_windows_by_account(accounts, &quota_by_account);
let quota_window_usage_stats =
account_usage_stats_by_quota_window(state, &selected_quota_windows).await;
quota_window_local_usage_by_account(state, &all_quota_windows).await;
apply_quota_window_local_usage(accounts, &mut quota_by_account, &quota_window_usage_stats);
let selected_quota_windows = selected_quota_windows_by_account(accounts, &quota_by_account);
let usage_records = usage_records
.into_iter()
.map(|usage| {
@@ -937,14 +939,29 @@ async fn account_list_stats(
})
.collect::<Vec<_>>();
let models_by_account = list_current_window_model_usage(state, &usage_records).await;
let refreshing_account_ids = state
.services
.token_refresh
.refreshing_account_ids(&account_ids, Utc::now())
.await
.map_err(|error| AdminError::internal(error.to_string()))?;
AccountListStats {
Ok(AccountListStats {
usage_by_account: usage_records
.into_iter()
.map(|usage| (usage.account_id.clone(), usage))
.collect(),
quota_by_account,
models_by_account,
refreshing_account_ids,
})
}
fn account_display_status(status: AccountStatus, token_refreshing: bool) -> &'static str {
if token_refreshing && matches!(status, AccountStatus::Active | AccountStatus::Expired) {
"refreshing"
} else {
status.as_str()
}
}
@@ -972,105 +989,11 @@ fn selected_quota_window(quota: &AdminAccountQuotaData) -> Option<AdminAccountQu
})
}
async fn account_usage_stats_by_quota_window(
state: &AppState,
windows_by_account: &HashMap<String, AdminAccountQuotaUsageWindow>,
) -> HashMap<String, AccountUsageWindowStats> {
if windows_by_account.is_empty() {
return HashMap::new();
}
let mut builder = QueryBuilder::<Sqlite>::new(
"select
account_id,
coalesce(sum(request_count), 0) as request_count,
coalesce(sum(input_tokens), 0) as input_tokens,
coalesce(sum(output_tokens), 0) as output_tokens,
coalesce(sum(cached_tokens), 0) as cached_tokens
from usage_time_buckets
where ",
);
for (index, (account_id, window)) in windows_by_account.iter().enumerate() {
if index > 0 {
builder.push(" or ");
}
let bucket_start = china_quarter_hour_start(window.start);
builder.push("(account_id = ");
builder.push_bind(account_id);
builder.push(" and bucket_start >= ");
builder.push_bind(bucket_start.to_rfc3339());
builder.push(" and bucket_start <= ");
builder.push_bind(window.end.to_rfc3339());
builder.push(")");
}
builder.push(" group by account_id");
let rows = builder
.build()
.fetch_all(state.services.background_tasks.accounts.pool())
.await
.unwrap_or_default();
rows.into_iter()
.map(|row| {
(
row.get::<String, _>("account_id"),
AccountUsageWindowStats {
request_count: row.get("request_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
},
)
})
.collect()
}
fn apply_selected_quota_window_usage(
mut usage: AdminUsageRecord,
selected_quota_windows: &HashMap<String, AdminAccountQuotaUsageWindow>,
quota_window_usage_stats: &HashMap<String, AccountUsageWindowStats>,
) -> AdminUsageRecord {
let Some(window) = selected_quota_windows.get(&usage.account_id) else {
return usage;
};
let stats = quota_window_usage_stats
.get(&usage.account_id)
.copied()
.unwrap_or_default();
usage.window_request_count = stats.request_count;
usage.window_input_tokens = stats.input_tokens;
usage.window_output_tokens = stats.output_tokens;
usage.window_cached_tokens = stats.cached_tokens;
usage.window_started_at = Some(window.start);
usage.window_reset_at = Some(window.end);
usage.limit_window_seconds = Some(window.window_seconds);
usage
}
async fn 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,
fn quota_usage_windows_by_account(
accounts: &[AdminAccountMetadata],
quota_by_account: &mut HashMap<String, AdminAccountQuotaData>,
) {
let windows = accounts
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
) -> Vec<AccountQuotaUsageWindow> {
accounts
.iter()
.filter_map(|account| {
quota_by_account
@@ -1086,21 +1009,44 @@ async fn apply_quota_window_token_usage(
window,
})
})
.collect::<Vec<_>>();
let usage_by_account = quota_window_token_usage_by_account(state, &windows).await;
.collect()
}
async fn apply_account_quota_window_local_usage(
state: &AppState,
account_id: &str,
quota: &mut AdminAccountQuotaData,
) {
let windows = quota
.usage_windows()
.into_iter()
.map(|window| AccountQuotaUsageWindow {
account_id: account_id.to_string(),
window,
})
.collect::<Vec<_>>();
let usage_by_account = quota_window_local_usage_by_account(state, &windows).await;
let empty_usage = HashMap::new();
quota.apply_local_usage(usage_by_account.get(account_id).unwrap_or(&empty_usage));
}
fn apply_quota_window_local_usage(
accounts: &[AdminAccountMetadata],
quota_by_account: &mut HashMap<String, AdminAccountQuotaData>,
usage_by_account: &HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>,
) {
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));
quota.apply_local_usage(usage_by_account.get(&account.id).unwrap_or(&empty_usage));
}
}
}
async fn quota_window_token_usage_by_account(
async fn quota_window_local_usage_by_account(
state: &AppState,
windows: &[AccountQuotaUsageWindow],
) -> HashMap<String, HashMap<String, AdminAccountQuotaWindowTokenUsage>> {
) -> HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>> {
let Some(min_start) = windows.iter().map(|window| window.window.start).min() else {
return HashMap::new();
};
@@ -1123,8 +1069,10 @@ async fn quota_window_token_usage_by_account(
"select
account_id,
bucket_start,
coalesce(sum(request_count), 0) as request_count,
coalesce(sum(input_tokens), 0) as input_tokens,
coalesce(sum(output_tokens), 0) as output_tokens
coalesce(sum(output_tokens), 0) as output_tokens,
coalesce(sum(cached_tokens), 0) as cached_tokens
from usage_time_buckets
where account_id in (",
);
@@ -1134,7 +1082,7 @@ async fn quota_window_token_usage_by_account(
}
separated.push_unseparated(")");
builder.push(" and bucket_start >= ");
builder.push_bind(min_start.to_rfc3339());
builder.push_bind(china_quarter_hour_start(min_start).to_rfc3339());
builder.push(" and bucket_start <= ");
builder.push_bind(max_end.to_rfc3339());
builder.push(" group by account_id, bucket_start");
@@ -1145,35 +1093,61 @@ async fn quota_window_token_usage_by_account(
.await
.unwrap_or_default();
let mut usage_by_account =
HashMap::<String, HashMap<String, AdminAccountQuotaWindowTokenUsage>>::new();
HashMap::<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>::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 bucket_usage = AdminAccountQuotaWindowLocalUsage {
request_count: row.get("request_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
};
let Some(account_windows) = windows_by_account.get(account_id.as_str()) else {
continue;
};
for window in account_windows
.iter()
.filter(|window| bucket_start >= window.start && bucket_start <= window.end)
{
for window in account_windows.iter().filter(|window| {
bucket_start >= china_quarter_hour_start(window.start) && bucket_start <= window.end
}) {
usage_by_account
.entry(account_id.clone())
.or_default()
.entry(window.key.clone())
.or_default()
.add_tokens(total_tokens);
.add(bucket_usage);
}
}
usage_by_account
}
fn apply_selected_quota_window_usage(
mut usage: AdminUsageRecord,
selected_quota_windows: &HashMap<String, AdminAccountQuotaUsageWindow>,
quota_window_usage_stats: &HashMap<String, HashMap<String, AdminAccountQuotaWindowLocalUsage>>,
) -> AdminUsageRecord {
let Some(window) = selected_quota_windows.get(&usage.account_id) else {
return usage;
};
let stats = quota_window_usage_stats
.get(&usage.account_id)
.and_then(|usage_by_window| usage_by_window.get(&window.key))
.copied()
.unwrap_or_default();
usage.window_request_count = stats.request_count;
usage.window_input_tokens = stats.input_tokens;
usage.window_output_tokens = stats.output_tokens;
usage.window_cached_tokens = stats.cached_tokens;
usage.window_started_at = Some(window.start);
usage.window_reset_at = Some(window.end);
usage.limit_window_seconds = Some(window.window_seconds);
usage
}
async fn account_summary_data(
state: &AppState,
quota_by_account: &HashMap<String, AdminAccountQuotaData>,
@@ -75,6 +75,8 @@ pub enum AdminAccountError {
Delete,
#[error("failed to update claims")]
UpdateClaims,
#[error("failed to coordinate token refresh lease")]
RefreshLease,
#[error("failed to sync account pool")]
SyncAccountPool,
#[error("failed to store quota")]
@@ -3,10 +3,11 @@ use std::{
hash::{Hash, Hasher},
};
use chrono::{DateTime, Utc};
use chrono::{DateTime, Duration, Utc};
use futures::{stream, StreamExt};
use secrecy::{ExposeSecret, SecretString};
use tokio::time::sleep;
use uuid::Uuid;
use crate::{
infra::{
@@ -25,6 +26,8 @@ use super::{
AdminAccountService,
};
const ADMIN_REFRESH_LEASE_TTL_SECONDS: i64 = 5 * 60;
impl AdminAccountService {
pub async fn list(
&self,
@@ -281,20 +284,44 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
};
let Some(refresh_token) = account.refresh_token.as_ref() else {
let Some(refresh_token) = account
.refresh_token
.as_ref()
.map(|token| token.expose_secret().to_string())
else {
return Ok(skipped("no refresh token"));
};
if account.status == crate::upstream::accounts::model::AccountStatus::Disabled {
return Ok(skipped("manually disabled"));
}
if account.status == crate::upstream::accounts::model::AccountStatus::Refreshing {
let lease_now = Utc::now();
let lease_owner = format!(
"{}:{}",
self.refresh_lease_owner_prefix,
Uuid::new_v4().simple()
);
let lease_acquired = self
.refresh_leases
.try_acquire(
account_id,
&lease_owner,
lease_now + Duration::seconds(ADMIN_REFRESH_LEASE_TTL_SECONDS),
lease_now,
)
.await
.map_err(|_| AdminAccountError::RefreshLease)?;
if !lease_acquired {
return Ok(skipped("refresh already in progress"));
}
match self
.refresh_tokens_from_refresh_token(refresh_token.expose_secret())
.await
{
let refresh_result = self.refresh_tokens_from_refresh_token(&refresh_token).await;
let release_result = self.refresh_leases.release(account_id, &lease_owner).await;
if release_result.is_err() {
return Err(AdminAccountError::RefreshLease);
}
match refresh_result {
Ok(refreshed) => {
let updated = self
.store
+10 -1
View File
@@ -12,13 +12,14 @@ mod testing;
use std::sync::Arc as StdArc;
use chrono::{DateTime, Utc};
use uuid::Uuid;
use crate::{
upstream::accounts::{
cookies::SqliteCookieStore,
pool::RuntimeAccountPoolService,
store::SqliteAccountStore,
token_refresh::{RuntimeRefreshPolicy, TokenRefresher},
token_refresh::{RefreshLeaseStore, RuntimeRefreshPolicy, TokenRefresher},
},
upstream::models::service::ModelService,
upstream::transport::CodexBackendClient,
@@ -38,6 +39,8 @@ pub struct AdminAccountService {
pub(crate) models: StdArc<ModelService>,
pub(crate) account_pool: StdArc<RuntimeAccountPoolService>,
pub(crate) token_refresher: StdArc<dyn TokenRefresher>,
pub(crate) refresh_leases: RefreshLeaseStore,
pub(crate) refresh_lease_owner_prefix: String,
pub(crate) oauth: oauth::AccountOAuthService,
pub(crate) refresh_policy: RuntimeRefreshPolicy,
pub(crate) installation_id: Option<String>,
@@ -50,6 +53,7 @@ pub struct AdminAccountServiceParts {
pub models: StdArc<ModelService>,
pub account_pool: StdArc<RuntimeAccountPoolService>,
pub token_refresher: StdArc<dyn TokenRefresher>,
pub refresh_leases: RefreshLeaseStore,
pub oauth: oauth::AccountOAuthService,
pub refresh_policy: RuntimeRefreshPolicy,
pub installation_id: Option<String>,
@@ -64,6 +68,11 @@ impl AdminAccountService {
models: parts.models,
account_pool: parts.account_pool,
token_refresher: parts.token_refresher,
refresh_leases: parts.refresh_leases,
refresh_lease_owner_prefix: format!(
"admin-account-refresh:{}",
Uuid::new_v4().simple()
),
oauth: parts.oauth,
refresh_policy: parts.refresh_policy,
installation_id: parts.installation_id,
+36 -35
View File
@@ -1,6 +1,6 @@
//! 管理端 Dashboard 聚合视图。
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use axum::{
extract::{Query, State},
@@ -15,7 +15,7 @@ use crate::{
admin::{
auth::session::AdminAuth,
monitoring::{
account_usage_service::{AdminUsageSummary, AdminUsageTimeBucketRecord},
account_usage_service::AdminUsageTimeBucketRecord,
diagnostics::{
fingerprint_diagnostics, AccountCapacityDiagnostics, AccountPoolDiagnostics,
},
@@ -27,7 +27,7 @@ use crate::{
UsageRecordCostDetailsData, UsageRecordTokenDetailsData,
},
usage_record_service::AdminUsageRecordFilter,
usage_record_store::UsageRecordAccountUsage,
usage_record_store::{UsageRecordAccountUsage, UsageRecordSummary},
},
response::{AdminEnvelope, AdminError, AdminResponse},
},
@@ -292,10 +292,10 @@ pub(crate) async fn dashboard_summary(
let today_filter = today_usage_record_filter(now);
let usage_summary = state
.services
.usage
.summary()
.usage_records
.summary(AdminUsageRecordFilter::default())
.await
.unwrap_or_else(|_| empty_usage_summary());
.unwrap_or_else(|_| UsageRecordSummary::default());
let account_usage_records = state
.services
.usage_records
@@ -311,7 +311,17 @@ pub(crate) async fn dashboard_summary(
.unwrap_or_default();
let time_buckets = dashboard_time_buckets(&state, now).await;
let pool_summary = account_pool_summary(&accounts);
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let refreshing_account_ids = state
.services
.token_refresh
.refreshing_account_ids(&account_ids, now)
.await
.map_err(|error| AdminError::internal(error.to_string()))?;
let pool_summary = account_pool_summary(&accounts, &refreshing_account_ids);
let trend = dashboard_trend_data(&time_buckets, query.kind.unwrap_or_default());
let quota_used_by_account =
account_quota_used_percent_by_id(&state, &account_usage_records).await;
@@ -375,7 +385,7 @@ pub(crate) async fn dashboard_trend(
fn dashboard_cards(
accounts: &[Account],
buckets: &[AdminUsageTimeBucketRecord],
usage_summary: &AdminUsageSummary,
usage_summary: &UsageRecordSummary,
) -> DashboardCardsData {
dashboard_cards_at(accounts, buckets, usage_summary, Utc::now())
}
@@ -383,7 +393,7 @@ fn dashboard_cards(
fn dashboard_cards_at(
accounts: &[Account],
buckets: &[AdminUsageTimeBucketRecord],
usage_summary: &AdminUsageSummary,
usage_summary: &UsageRecordSummary,
now: DateTime<Utc>,
) -> DashboardCardsData {
let today_start = china_day_start(now);
@@ -391,9 +401,9 @@ fn dashboard_cards_at(
let today = usage_window(buckets, today_start, now);
let yesterday = usage_window(buckets, yesterday_start, today_start);
let today_cost = cost_window(buckets, today_start, now).unwrap_or(0.0);
let total_requests = nonnegative_i64_to_u64(usage_summary.request_count);
let total_tokens = nonnegative_i64_to_u64(usage_summary.total_tokens);
let total_cached_tokens = nonnegative_i64_to_u64(usage_summary.cached_tokens);
let total_requests = usage_summary.total_requests;
let total_tokens = usage_summary.total_tokens;
let total_cached_tokens = usage_summary.cached_tokens;
let total_hit_rate = summary_cache_hit_rate(usage_summary);
let total_accounts = accounts.len() as u64;
@@ -449,27 +459,10 @@ fn dashboard_cards_at(
}
}
fn empty_usage_summary() -> AdminUsageSummary {
AdminUsageSummary {
account_count: 0,
request_count: 0,
empty_response_count: 0,
input_tokens: 0,
output_tokens: 0,
cached_tokens: 0,
reasoning_tokens: 0,
total_tokens: 0,
image_input_tokens: 0,
image_output_tokens: 0,
image_request_count: 0,
image_request_failed_count: 0,
}
}
fn summary_cache_hit_rate(summary: &AdminUsageSummary) -> Option<f64> {
let input_tokens = nonnegative_i64_to_u64(summary.input_tokens);
fn summary_cache_hit_rate(summary: &UsageRecordSummary) -> Option<f64> {
let input_tokens = summary.input_tokens;
if input_tokens > 0 {
Some(nonnegative_i64_to_u64(summary.cached_tokens) as f64 / input_tokens as f64)
Some(summary.cached_tokens as f64 / input_tokens as f64)
} else {
None
}
@@ -822,17 +815,24 @@ fn first_token_bucket_latency(record: &AdminUsageTimeBucketRecord) -> Option<u64
sum.checked_div(count).filter(|latency| *latency > 0)
}
fn account_pool_summary(accounts: &[Account]) -> AccountPoolDiagnostics {
fn account_pool_summary(
accounts: &[Account],
refreshing_account_ids: &HashSet<String>,
) -> AccountPoolDiagnostics {
let mut summary = AccountPoolDiagnostics {
total: accounts.len(),
..AccountPoolDiagnostics::default()
};
for account in accounts {
match account.status {
AccountStatus::Active | AccountStatus::Expired
if refreshing_account_ids.contains(&account.id) =>
{
summary.refreshing += 1;
}
AccountStatus::Active => summary.active += 1,
AccountStatus::Expired => summary.expired += 1,
AccountStatus::QuotaExhausted => summary.quota_exhausted += 1,
AccountStatus::Refreshing => summary.refreshing += 1,
AccountStatus::Disabled => summary.disabled += 1,
AccountStatus::Banned => summary.banned += 1,
}
@@ -1033,7 +1033,8 @@ fn dashboard_usage_record_data(
}
fn service_status_data(state: &AppState) -> Vec<DashboardServiceStatusData> {
let fingerprint = fingerprint_diagnostics(&state.services.fingerprint);
let current_fingerprint = state.services.fingerprint.snapshot();
let fingerprint = fingerprint_diagnostics(&current_fingerprint);
vec![
DashboardServiceStatusData {
label: "客户端版本".to_string(),
+3
View File
@@ -5,6 +5,9 @@ pub mod account_usage_store;
pub(crate) mod billing;
pub mod dashboard;
pub mod diagnostics;
pub mod ops_error_model;
pub mod ops_error_service;
pub mod ops_error_store;
pub mod usage_record_model;
pub mod usage_record_routes;
pub mod usage_record_service;
@@ -0,0 +1,75 @@
//! 运维错误事件模型。
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use uuid::Uuid;
/// 结构化运维错误事件。
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct OpsErrorLog {
/// 事件 ID。
pub id: String,
/// 请求 ID。
pub request_id: Option<String>,
/// 事件类别。
pub kind: String,
/// 账号 ID。预账号调度失败时为空。
pub account_id: Option<String>,
/// HTTP 路由。
pub route: Option<String>,
/// 模型名。
pub model: Option<String>,
/// 事件状态码。
pub status_code: Option<i64>,
/// 客户端看到的状态码。
pub client_status_code: Option<i64>,
/// 上游 HTTP 状态码。
pub upstream_status_code: Option<i64>,
/// 上游传输方式,例如 http_sse 或 websocket。
pub transport: Option<String>,
/// 同一请求内的上游尝试序号。
pub attempt_index: Option<i64>,
/// 失败分类。
pub failure_class: Option<String>,
/// 上游响应 ID。
pub response_id: Option<String>,
/// 上游请求 ID。
pub upstream_request_id: Option<String>,
/// 请求耗时毫秒。
pub latency_ms: Option<i64>,
/// 展示消息。
pub message: String,
/// 结构化元数据。
pub metadata: Value,
/// 创建时间。
#[serde(serialize_with = "crate::infra::time::serialize_china_rfc3339")]
pub created_at: DateTime<Utc>,
}
impl OpsErrorLog {
/// 构造运维错误事件。
pub fn new(kind: impl Into<String>, message: impl Into<String>) -> Self {
Self {
id: format!("ops_error_{}", Uuid::new_v4().simple()),
request_id: None,
kind: kind.into(),
account_id: None,
route: None,
model: None,
status_code: None,
client_status_code: None,
upstream_status_code: None,
transport: None,
attempt_index: None,
failure_class: None,
response_id: None,
upstream_request_id: None,
latency_ms: None,
message: message.into(),
metadata: Value::Object(Map::default()),
created_at: Utc::now(),
}
}
}
@@ -0,0 +1,68 @@
//! 管理端运维错误事件服务。
use chrono::Utc;
use thiserror::Error;
use crate::admin::monitoring::{
ops_error_model::OpsErrorLog, ops_error_store::SqliteOpsErrorLogStore,
};
/// 管理端运维错误事件错误。
#[derive(Debug, Error)]
pub enum AdminOpsErrorLogError {
/// 写入失败。
#[error("failed to append ops error log")]
Append,
/// 保留期清理失败。
#[error("failed to trim expired ops error logs")]
Retention,
}
/// 管理端运维错误事件服务。
#[derive(Clone)]
pub struct AdminOpsErrorLogService {
store: SqliteOpsErrorLogStore,
capture_body: bool,
}
impl AdminOpsErrorLogService {
/// 构造管理端运维错误事件服务。
pub fn new(store: SqliteOpsErrorLogStore, capture_body: bool) -> Self {
Self {
store,
capture_body,
}
}
/// 记录运维错误事件。
pub async fn record(&self, mut event: OpsErrorLog) -> Result<(), AdminOpsErrorLogError> {
apply_capture_body_policy(&mut event, self.capture_body);
self.store
.append(&event)
.await
.map_err(|_| AdminOpsErrorLogError::Append)?;
self.store
.trim_to_retention(Utc::now())
.await
.map_err(|_| AdminOpsErrorLogError::Retention)?;
Ok(())
}
}
fn apply_capture_body_policy(event: &mut OpsErrorLog, capture_body: bool) {
if capture_body {
return;
}
let Some(metadata) = event.metadata.as_object_mut() else {
return;
};
for key in [
"body",
"rawBody",
"requestBody",
"responseBody",
"upstreamBody",
] {
metadata.remove(key);
}
}
@@ -0,0 +1,261 @@
//! 运维错误事件存储实现(SQLite)。
use chrono::{DateTime, Duration, Utc};
use serde_json::Value;
use sqlx::SqlitePool;
use thiserror::Error;
use crate::{
admin::monitoring::{
ops_error_model::OpsErrorLog,
usage_record_model::{metadata_i64, metadata_service_tier, metadata_string},
usage_record_store::USAGE_RECORD_RETENTION_DAYS,
},
infra::time::china_quarter_hour_start,
};
/// SQLite 运维错误事件错误。
#[derive(Debug, Error)]
pub enum SqliteOpsErrorLogStoreError {
/// 数据库错误。
#[error("sqlite ops error log database error: {0}")]
Database(#[from] sqlx::Error),
}
/// SQLite 运维错误事件结果。
pub type SqliteOpsErrorLogStoreResult<T> = Result<T, SqliteOpsErrorLogStoreError>;
/// SQLite 运维错误事件存储。
#[derive(Clone)]
pub struct SqliteOpsErrorLogStore {
pool: SqlitePool,
}
impl SqliteOpsErrorLogStore {
/// 构造存储。
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
/// 写入运维错误事件。
pub async fn append(&self, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> {
append_event(&self.pool, event).await
}
/// 按保留期清理运维错误事件。
pub async fn trim_to_retention(&self, now: DateTime<Utc>) -> SqliteOpsErrorLogStoreResult<u64> {
trim_to_retention(&self.pool, now).await
}
/// 清空运维错误事件。
pub async fn clear(&self) -> SqliteOpsErrorLogStoreResult<u64> {
let result = sqlx::query("delete from ops_error_logs")
.execute(&self.pool)
.await?;
Ok(result.rows_affected())
}
/// 读取错误事件数量。
pub async fn count(&self) -> SqliteOpsErrorLogStoreResult<u64> {
let (total,): (i64,) = sqlx::query_as("select count(*) from ops_error_logs")
.fetch_one(&self.pool)
.await?;
Ok(total.max(0) as u64)
}
}
#[derive(Debug, Clone)]
struct OpsErrorSummaryFields {
client_status_code: Option<i64>,
upstream_status_code: Option<i64>,
transport: Option<String>,
attempt_index: Option<i64>,
failure_class: Option<String>,
response_id: Option<String>,
upstream_request_id: Option<String>,
}
fn ops_error_summary_fields(event: &OpsErrorLog) -> OpsErrorSummaryFields {
OpsErrorSummaryFields {
client_status_code: event.client_status_code.or_else(|| {
metadata_i64(
&event.metadata,
&["clientStatusCode", "client_status_code", "clientStatus"],
)
}),
upstream_status_code: event.upstream_status_code.or_else(|| {
metadata_i64(
&event.metadata,
&[
"upstreamStatus",
"upstreamStatusCode",
"upstream_status_code",
],
)
}),
transport: event
.transport
.clone()
.or_else(|| metadata_string(&event.metadata, &["transport"])),
attempt_index: event
.attempt_index
.or_else(|| metadata_i64(&event.metadata, &["attemptIndex", "attempt_index"])),
failure_class: event
.failure_class
.clone()
.or_else(|| metadata_string(&event.metadata, &["failureClass", "failure_class"])),
response_id: event
.response_id
.clone()
.or_else(|| metadata_string(&event.metadata, &["responseId", "response_id"])),
upstream_request_id: event.upstream_request_id.clone().or_else(|| {
metadata_string(
&event.metadata,
&[
"upstreamRequestId",
"upstream_request_id",
"openaiRequestId",
],
)
}),
}
}
async fn append_event(pool: &SqlitePool, event: &OpsErrorLog) -> SqliteOpsErrorLogStoreResult<()> {
let summary = ops_error_summary_fields(event);
sqlx::query(
"insert into ops_error_logs (id, request_id, kind, account_id, route, model, status_code, client_status_code, upstream_status_code, transport, attempt_index, failure_class, response_id, upstream_request_id, latency_ms, message, metadata_json, created_at) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
)
.bind(&event.id)
.bind(&event.request_id)
.bind(&event.kind)
.bind(&event.account_id)
.bind(&event.route)
.bind(&event.model)
.bind(event.status_code)
.bind(summary.client_status_code)
.bind(summary.upstream_status_code)
.bind(summary.transport)
.bind(summary.attempt_index)
.bind(summary.failure_class)
.bind(summary.response_id)
.bind(summary.upstream_request_id)
.bind(event.latency_ms)
.bind(&event.message)
.bind(event.metadata.to_string())
.bind(event.created_at.to_rfc3339())
.execute(pool)
.await?;
append_error_time_bucket(pool, event).await?;
Ok(())
}
async fn append_error_time_bucket(
pool: &SqlitePool,
event: &OpsErrorLog,
) -> SqliteOpsErrorLogStoreResult<()> {
let bucket_start = china_quarter_hour_start(event.created_at);
let input_tokens = metadata_usage_i64(&event.metadata, "inputTokens");
let output_tokens = metadata_usage_i64(&event.metadata, "outputTokens");
let cached_tokens = metadata_usage_i64(&event.metadata, "cachedTokens");
let first_token_latency = metadata_first_token_latency(&event.metadata).unwrap_or(0);
let first_token_latency_count = i64::from(first_token_latency > 0);
let latency = event.latency_ms.filter(|value| *value > 0).unwrap_or(0);
let latency_count = i64::from(latency > 0);
let now = Utc::now().to_rfc3339();
sqlx::query(
r"
insert into usage_time_buckets (
bucket_start,
account_id,
model,
service_tier,
request_count,
error_count,
input_tokens,
output_tokens,
cached_tokens,
first_token_latency_sum,
first_token_latency_count,
latency_sum,
latency_count,
max_latency_ms,
min_latency_ms,
updated_at
) values (?, ?, ?, ?, 1, 1, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
on conflict(bucket_start, account_id, model, service_tier) do update set
request_count = request_count + excluded.request_count,
error_count = error_count + excluded.error_count,
input_tokens = input_tokens + excluded.input_tokens,
output_tokens = output_tokens + excluded.output_tokens,
cached_tokens = cached_tokens + excluded.cached_tokens,
first_token_latency_sum = first_token_latency_sum + excluded.first_token_latency_sum,
first_token_latency_count = first_token_latency_count + excluded.first_token_latency_count,
latency_sum = latency_sum + excluded.latency_sum,
latency_count = latency_count + excluded.latency_count,
max_latency_ms = max(max_latency_ms, excluded.max_latency_ms),
min_latency_ms = case
when min_latency_ms = 0 then excluded.min_latency_ms
when excluded.min_latency_ms = 0 then min_latency_ms
else min(min_latency_ms, excluded.min_latency_ms)
end,
updated_at = excluded.updated_at",
)
.bind(bucket_start.to_rfc3339())
.bind(event.account_id.as_deref().unwrap_or_default())
.bind(event.model.as_deref().unwrap_or_default())
.bind(metadata_service_tier(&event.metadata).unwrap_or_default())
.bind(input_tokens)
.bind(output_tokens)
.bind(cached_tokens)
.bind(first_token_latency)
.bind(first_token_latency_count)
.bind(latency)
.bind(latency_count)
.bind(latency)
.bind(latency)
.bind(now)
.execute(pool)
.await?;
Ok(())
}
fn metadata_usage_i64(metadata: &Value, field: &str) -> i64 {
metadata
.get("usage")
.and_then(|usage| usage.get(field))
.or_else(|| metadata.get(field))
.and_then(Value::as_i64)
.filter(|value| *value > 0)
.unwrap_or(0)
}
fn metadata_first_token_latency(metadata: &Value) -> Option<i64> {
[
"firstTokenMs",
"first_token_ms",
"firstTokenLatencyMs",
"first_token_latency_ms",
]
.into_iter()
.find_map(|field| {
metadata
.get(field)
.or_else(|| metadata.get("usage").and_then(|usage| usage.get(field)))
.and_then(Value::as_i64)
.filter(|value| *value > 0)
})
}
async fn trim_to_retention(
pool: &SqlitePool,
now: DateTime<Utc>,
) -> SqliteOpsErrorLogStoreResult<u64> {
let cutoff = now - Duration::days(USAGE_RECORD_RETENTION_DAYS);
let result = sqlx::query("delete from ops_error_logs where created_at < ?")
.bind(cutoff.to_rfc3339())
.execute(pool)
.await?;
Ok(result.rows_affected())
}
@@ -265,7 +265,19 @@ impl AdminUsageRecordService {
/// 记录使用记录。
pub async fn record(&self, event: UsageRecord) -> Result<(), AdminUsageRecordError> {
let settings = *self.settings.read().await;
if !settings.enabled && event.level != UsageRecordLevel::Error {
if !settings.enabled {
return Ok(());
}
if !is_usage_log_event(&event) {
tracing::warn!(
usage_record_id = %event.id,
request_id = event.request_id.as_deref().unwrap_or(""),
account_id = event.account_id.as_deref().unwrap_or(""),
model = event.model.as_deref().unwrap_or(""),
level = ?event.level,
status_code = ?event.status_code,
"skipped non-usage event in usage record service"
);
return Ok(());
}
self.append_with_settings(event, settings).await
@@ -289,6 +301,21 @@ impl AdminUsageRecordService {
}
}
fn is_usage_log_event(event: &UsageRecord) -> bool {
event.level != UsageRecordLevel::Error
&& event.status_code.is_none_or(|status| status < 400)
&& event
.account_id
.as_deref()
.map(str::trim)
.is_some_and(|account_id| !account_id.is_empty())
&& event
.model
.as_deref()
.map(str::trim)
.is_some_and(|model| !model.is_empty())
}
fn apply_capture_body_policy(event: &mut UsageRecord, capture_body: bool) {
if capture_body {
return;
@@ -100,6 +100,11 @@ impl SqliteUsageRecordStore {
Self { pool }
}
/// 暴露底层连接池,供运行时组合层复用。
pub fn pool(&self) -> &SqlitePool {
&self.pool
}
/// 写入使用记录。
pub async fn append(&self, event: &UsageRecord) -> SqliteUsageRecordStoreResult<()> {
append_event(&self.pool, event).await
+40 -9
View File
@@ -110,6 +110,8 @@ struct SystemUpdateEvent {
level: UpdateLogLevel,
step: Option<String>,
message: String,
#[serde(skip_serializing_if = "is_false")]
terminal: bool,
at: String,
}
@@ -157,18 +159,23 @@ pub(crate) async fn update_event_stream(
let receiver = update_event_sender().subscribe();
let shutdown = shutdown_subscription();
let stream = futures::stream::unfold(
(receiver, shutdown),
|(mut receiver, mut shutdown)| async move {
(receiver, shutdown, false),
|(mut receiver, mut shutdown, close_after_send)| async move {
if close_after_send {
return None;
}
loop {
tokio::select! {
_ = shutdown.recv() => return None,
received = receiver.recv() => {
match received {
Ok(message) => {
let close_after_send = message.terminal;
let id = message.id.clone();
let data = serde_json::to_string(&message).unwrap_or_else(|_| "{}".to_string());
let event = Event::default().event("update").id(id).data(data);
return Some((Ok(event), (receiver, shutdown)));
return Some((Ok(event), (receiver, shutdown, close_after_send)));
}
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(broadcast::error::RecvError::Closed) => return None,
@@ -192,7 +199,7 @@ pub(crate) async fn perform_update(
.map_err(|_| AdminError::conflict("System update already running"))?;
let config = SystemUpdateConfig::from_env();
if let Some(reason) = config.update_support_error() {
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Error,
None,
Some("preflight"),
@@ -217,7 +224,7 @@ pub(crate) async fn perform_update(
)
.await
.map_err(|error| {
emit_update_event(UpdateLogLevel::Error, None, Some("release"), error.clone());
emit_terminal_update_event(UpdateLogLevel::Error, None, Some("release"), error.clone());
AdminError::bad_gateway(error)
})?;
let info = update_info_from_release(&config, release.clone());
@@ -226,7 +233,7 @@ pub(crate) async fn perform_update(
"远端最新版本已变更为 v{},请重新检查并确认",
info.latest_version
);
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Warning,
None,
Some("release"),
@@ -235,7 +242,7 @@ pub(crate) async fn perform_update(
return Err(AdminError::conflict(message));
}
if !info.has_update {
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Warning,
None,
Some("release"),
@@ -261,14 +268,14 @@ pub(crate) async fn perform_update(
);
let result = perform_release_update(&config, &release, &target_version, &operation_id).await;
if let Err(error) = &result {
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Error,
Some(&operation_id),
Some("failed"),
error.to_string(),
);
} else {
emit_update_event(
emit_terminal_update_event(
UpdateLogLevel::Success,
Some(&operation_id),
Some("done"),
@@ -732,6 +739,25 @@ fn emit_update_event(
operation_id: Option<&str>,
step: Option<&str>,
message: impl Into<String>,
) {
emit_update_event_with_terminal(level, operation_id, step, message, false);
}
fn emit_terminal_update_event(
level: UpdateLogLevel,
operation_id: Option<&str>,
step: Option<&str>,
message: impl Into<String>,
) {
emit_update_event_with_terminal(level, operation_id, step, message, true);
}
fn emit_update_event_with_terminal(
level: UpdateLogLevel,
operation_id: Option<&str>,
step: Option<&str>,
message: impl Into<String>,
terminal: bool,
) {
let now = Utc::now();
let event = SystemUpdateEvent {
@@ -744,11 +770,16 @@ fn emit_update_event(
level,
step: step.map(ToString::to_string),
message: message.into(),
terminal,
at: now.to_rfc3339(),
};
let _ = update_event_sender().send(event);
}
fn is_false(value: &bool) -> bool {
!*value
}
fn format_bytes(bytes: u64) -> String {
const KIB: u64 = 1024;
const MIB: u64 = KIB * 1024;
+2 -5
View File
@@ -57,9 +57,6 @@ pub struct ServerConfig {
pub host: String,
/// 监听端口。
pub port: u16,
/// 可信反向代理 IP 或 CIDR。非空时只有这些 peer 的转发 IP 头会被采信。
#[serde(default)]
pub trusted_proxies: Vec<String>,
}
/// 上游 API 地址配置。
@@ -100,7 +97,7 @@ impl Default for AuthConfig {
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
rotation_strategy: "least_used".to_string(),
rotation_strategy: "smart".to_string(),
tier_priority: Vec::new(),
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
@@ -171,7 +168,7 @@ pub struct WebSocketPoolConfig {
}
fn default_ws_first_token_timeout_ms() -> u64 {
15_000
20_000
}
impl Default for WebSocketPoolConfig {
+6 -4
View File
@@ -18,7 +18,7 @@ use crate::upstream::accounts::{
token_refresh::{RefreshPolicy, RuntimeRefreshPolicy},
};
const ROTATION_STRATEGIES: [&str; 3] = ["least_used", "round_robin", "sticky"];
const ROTATION_STRATEGIES: [&str; 4] = ["smart", "quota_reset_priority", "round_robin", "sticky"];
const RUNTIME_SETTINGS_ID: i64 = 1;
/// 管理端可变设置。
@@ -46,7 +46,7 @@ impl Default for AdminSettings {
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
rotation_strategy: "least_used".to_string(),
rotation_strategy: "smart".to_string(),
}
}
}
@@ -150,7 +150,7 @@ fn validate_rotation_strategy(strategy: &str) -> Result<String, SettingsServiceE
} else {
Err(invalid_field(
"rotationStrategy",
"must be one of least_used, round_robin, sticky",
"must be one of smart, quota_reset_priority, round_robin, sticky",
))
}
}
@@ -551,9 +551,11 @@ pub fn account_pool_options_from_config(config: &AppConfig) -> AccountPoolOption
fn rotation_strategy_from_config(strategy: &str) -> RotationStrategy {
match strategy {
"smart" => RotationStrategy::Smart,
"quota_reset_priority" => RotationStrategy::QuotaResetPriority,
"round_robin" => RotationStrategy::RoundRobin,
"sticky" => RotationStrategy::Sticky,
_ => RotationStrategy::LeastUsed,
_ => RotationStrategy::Smart,
}
}
+11 -95
View File
@@ -4,7 +4,7 @@ use std::net::{IpAddr, SocketAddr};
use std::sync::LazyLock;
use axum::{
extract::{connect_info::ConnectInfo, Request, State},
extract::{connect_info::ConnectInfo, Request},
http::{HeaderMap, HeaderValue},
middleware::Next,
response::Response,
@@ -19,7 +19,7 @@ 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";
/// 私有 / 回环 / 唯一本地地址网段,用于在自动模式下从 `X-Forwarded-For` 中跳过内网跳板 IP。
/// 私有 / 回环 / 唯一本地地址网段,用于从 `X-Forwarded-For` 中跳过内网跳板 IP。
static PRIVATE_NETWORKS: LazyLock<Vec<IpNet>> = LazyLock::new(|| {
[
"10.0.0.0/8",
@@ -38,40 +38,6 @@ fn is_private_ip(ip: IpAddr) -> bool {
PRIVATE_NETWORKS.iter().any(|network| network.contains(&ip))
}
/// 可信反向代理配置。
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct TrustedProxyConfig {
networks: Vec<IpNet>,
}
impl TrustedProxyConfig {
/// 从 IP 或 CIDR 字符串构造可信代理配置。
pub fn from_entries(entries: &[String]) -> Result<Self, TrustedProxyConfigError> {
let networks = entries
.iter()
.map(|entry| parse_trusted_proxy_entry(entry))
.collect::<Result<Vec<_>, _>>()?;
Ok(Self { networks })
}
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()
}
}
/// 可信代理配置错误。
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum TrustedProxyConfigError {
/// 配置项不是 IP 地址或 CIDR。
#[error("invalid trusted proxy entry `{entry}`; expected an IP address or CIDR network")]
InvalidEntry { entry: String },
}
/// 请求 ID。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RequestId(String);
@@ -97,7 +63,7 @@ impl RequestId {
}
}
/// 由连接来源和可信代理链解析出的客户端 IP。
/// 由连接来源和转发头解析出的客户端 IP。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ClientIp(String);
@@ -118,17 +84,7 @@ fn valid_request_id(value: &str) -> bool {
/// 为请求附加请求 ID,并在响应头中回写。
pub async fn attach_request_id(mut request: Request, next: Next) -> Response {
attach_request_context(&TrustedProxyConfig::default(), &mut request);
run_with_request_id(request, next).await
}
/// 为请求附加请求上下文,并按可信代理配置解析真实客户端 IP。
pub async fn attach_request_id_with_proxy_config(
State(trusted_proxies): State<TrustedProxyConfig>,
mut request: Request,
next: Next,
) -> Response {
attach_request_context(&trusted_proxies, &mut request);
attach_request_context(&mut request);
run_with_request_id(request, next).await
}
@@ -147,17 +103,16 @@ async fn run_with_request_id(mut request: Request, next: Next) -> Response {
response
}
fn attach_request_context(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
attach_client_ip_from_connection(trusted_proxies, request);
fn attach_request_context(request: &mut Request) {
attach_client_ip_from_connection(request);
}
fn attach_client_ip_from_connection(trusted_proxies: &TrustedProxyConfig, request: &mut Request) {
fn attach_client_ip_from_connection(request: &mut Request) {
let peer_ip = request
.extensions()
.get::<ConnectInfo<SocketAddr>>()
.map(|ConnectInfo(addr)| addr.ip());
let client_ip =
peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers(), trusted_proxies));
let client_ip = peer_ip.map(|peer_ip| client_ip_from_peer(peer_ip, request.headers()));
let Some(client_ip) = client_ip else {
return;
@@ -165,28 +120,12 @@ fn attach_client_ip_from_connection(trusted_proxies: &TrustedProxyConfig, reques
request.extensions_mut().insert(ClientIp(client_ip));
}
fn client_ip_from_peer(
peer_ip: IpAddr,
headers: &HeaderMap,
trusted_proxies: &TrustedProxyConfig,
) -> 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_else(|| peer_ip.to_string())
} else {
// 已配置可信代理但直连来源不可信:忽略转发头,使用 peer IP。
peer_ip.to_string()
}
fn client_ip_from_peer(peer_ip: IpAddr, headers: &HeaderMap) -> String {
auto_forwarded_client_ip(headers).unwrap_or_else(|| peer_ip.to_string())
}
/// 自动模式下解析转发头中的真实客户端 IP。
/// 解析转发头中的真实客户端 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) {
@@ -212,13 +151,6 @@ fn forwarded_for_public_ip(headers: &HeaderMap) -> Option<String> {
.map(ToString::to_string)
}
fn trusted_forwarded_client_ip(headers: &HeaderMap) -> Option<String> {
header_string(headers, CF_CONNECTING_IP_HEADER)
.or_else(|| header_string(headers, REAL_IP_HEADER))
.or_else(|| forwarded_for_public_ip(headers))
.or_else(|| first_forwarded_for_value(headers))
}
fn header_string(headers: &HeaderMap, name: &str) -> Option<String> {
headers
.get(name)
@@ -237,19 +169,3 @@ fn first_forwarded_for_value(headers: &HeaderMap) -> Option<String> {
.filter(|value| !value.is_empty())
.map(ToString::to_string)
}
fn parse_trusted_proxy_entry(entry: &str) -> Result<IpNet, TrustedProxyConfigError> {
if let Ok(network) = entry.parse::<IpNet>() {
return Ok(network);
}
entry
.parse::<IpAddr>()
.map(IpNet::from)
.map_err(|_| invalid_trusted_proxy_entry(entry))
}
fn invalid_trusted_proxy_entry(entry: &str) -> TrustedProxyConfigError {
TrustedProxyConfigError::InvalidEntry {
entry: entry.to_string(),
}
}
+3 -22
View File
@@ -5,10 +5,7 @@ use std::path::Path;
use axum::{extract::State, http::StatusCode, middleware, routing::get, Router};
use crate::admin;
use crate::http::middleware::{
request_id::{attach_request_id_with_proxy_config, TrustedProxyConfig},
trace::http_trace_layer,
};
use crate::http::middleware::{request_id::attach_request_id, trace::http_trace_layer};
use crate::proxy;
use crate::runtime::state::AppState;
use crate::web;
@@ -18,34 +15,18 @@ pub const DEFAULT_ASSET_DIST_DIR: &str = "web/dist";
/// 构造整个 HTTP 服务路由。
pub fn router() -> Router<AppState> {
router_with_trusted_proxies(TrustedProxyConfig::default())
}
/// 使用指定可信代理配置构造整个 HTTP 服务路由。
pub fn router_with_trusted_proxies(trusted_proxies: TrustedProxyConfig) -> Router<AppState> {
router_with_assets_and_trusted_proxies(DEFAULT_ASSET_DIST_DIR, trusted_proxies)
router_with_assets(DEFAULT_ASSET_DIST_DIR)
}
/// 使用指定前端构建产物目录构造整个 HTTP 服务路由。
pub fn router_with_assets(dist_dir: impl AsRef<Path>) -> Router<AppState> {
router_with_assets_and_trusted_proxies(dist_dir, TrustedProxyConfig::default())
}
/// 使用指定前端构建产物目录和可信代理配置构造整个 HTTP 服务路由。
pub fn router_with_assets_and_trusted_proxies(
dist_dir: impl AsRef<Path>,
trusted_proxies: TrustedProxyConfig,
) -> Router<AppState> {
Router::new()
.route("/healthz", get(healthz))
.merge(proxy::router::router())
.merge(admin::router::router())
.fallback_service(web::assets::spa_router(dist_dir))
.layer(http_trace_layer())
.layer(middleware::from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
))
.layer(middleware::from_fn(attach_request_id))
}
async fn healthz(State(state): State<AppState>) -> StatusCode {
+18 -6
View File
@@ -4,12 +4,24 @@ use std::{fs, path::Path, str::FromStr, time::Duration};
use sqlx::{sqlite::SqliteConnectOptions, SqlitePool};
const CURRENT_SCHEMA_VERSION: i64 = 1;
const MIGRATIONS: &[Migration] = &[Migration {
version: 1,
name: "initial",
sql: include_str!("migrations/0001_initial.sql"),
}];
const CURRENT_SCHEMA_VERSION: i64 = 3;
const MIGRATIONS: &[Migration] = &[
Migration {
version: 1,
name: "initial",
sql: include_str!("migrations/0001_initial.sql"),
},
Migration {
version: 2,
name: "rename_rotation_strategies",
sql: include_str!("migrations/0002_rename_rotation_strategies.sql"),
},
Migration {
version: 3,
name: "ops_error_logs",
sql: include_str!("migrations/0003_ops_error_logs.sql"),
},
];
struct Migration {
version: i64,
@@ -49,7 +49,7 @@ create table runtime_settings (
refresh_concurrency integer not null check (refresh_concurrency > 0),
max_concurrent_per_account integer not null check (max_concurrent_per_account > 0),
request_interval_ms integer not null check (request_interval_ms >= 0),
rotation_strategy text not null check (rotation_strategy in ('least_used', 'round_robin', 'sticky')),
rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')),
admin_api_key text,
updated_at text not null
);
@@ -68,7 +68,7 @@ create table accounts (
refresh_token text,
access_token_expires_at text,
next_refresh_at text,
status text not null check (status in ('active', 'expired', 'quota_exhausted', 'refreshing', 'disabled', 'banned')),
status text not null check (status in ('active', 'expired', 'quota_exhausted', 'disabled', 'banned')),
quota_json text,
quota_fetched_at text,
quota_limit_reached integer not null default 0 check (quota_limit_reached in (0, 1)),
@@ -0,0 +1,40 @@
create table runtime_settings_new (
id integer primary key check (id = 1),
model_aliases_json text not null default '{}',
refresh_margin_seconds integer not null check (refresh_margin_seconds > 0),
refresh_concurrency integer not null check (refresh_concurrency > 0),
max_concurrent_per_account integer not null check (max_concurrent_per_account > 0),
request_interval_ms integer not null check (request_interval_ms >= 0),
rotation_strategy text not null check (rotation_strategy in ('smart', 'quota_reset_priority', 'round_robin', 'sticky')),
admin_api_key text,
updated_at text not null
);
insert into runtime_settings_new (
id,
model_aliases_json,
refresh_margin_seconds,
refresh_concurrency,
max_concurrent_per_account,
request_interval_ms,
rotation_strategy,
admin_api_key,
updated_at
)
select
id,
model_aliases_json,
refresh_margin_seconds,
refresh_concurrency,
max_concurrent_per_account,
request_interval_ms,
case rotation_strategy
when 'least_used' then 'quota_reset_priority'
else rotation_strategy
end,
admin_api_key,
updated_at
from runtime_settings;
drop table runtime_settings;
alter table runtime_settings_new rename to runtime_settings;
@@ -0,0 +1,31 @@
create table ops_error_logs (
id text primary key,
request_id text,
kind text not null,
account_id text,
route text,
model text,
status_code integer check (status_code is null or (status_code >= 100 and status_code <= 599)),
client_status_code integer check (client_status_code is null or (client_status_code >= 100 and client_status_code <= 599)),
upstream_status_code integer check (upstream_status_code is null or (upstream_status_code >= 100 and upstream_status_code <= 599)),
transport text,
attempt_index integer check (attempt_index is null or attempt_index >= 0),
failure_class text,
response_id text,
upstream_request_id text,
latency_ms integer check (latency_ms is null or latency_ms >= 0),
message text not null,
metadata_json text not null,
created_at text not null
);
create index idx_ops_error_logs_created_id on ops_error_logs(created_at desc, id desc);
create index idx_ops_error_logs_request_id on ops_error_logs(request_id) where request_id is not null;
create index idx_ops_error_logs_account on ops_error_logs(account_id, created_at desc) where account_id is not null;
create index idx_ops_error_logs_route_created on ops_error_logs(route, created_at desc) where route is not null;
create index idx_ops_error_logs_model_created on ops_error_logs(model, created_at desc) where model is not null;
create index idx_ops_error_logs_status_created on ops_error_logs(status_code, created_at desc) where status_code is not null;
create index idx_ops_error_logs_transport_created on ops_error_logs(transport, created_at desc) where transport is not null;
create index idx_ops_error_logs_failure_class on ops_error_logs(failure_class, created_at desc) where failure_class is not null;
create index idx_ops_error_logs_response_id on ops_error_logs(response_id) where response_id is not null;
create index idx_ops_error_logs_upstream_request_id on ops_error_logs(upstream_request_id) where upstream_request_id is not null;
@@ -4,8 +4,9 @@ use serde_json::{json, Map, Value};
use crate::{
admin::monitoring::{
usage_record_model::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
ops_error_model::OpsErrorLog,
ops_error_service::AdminOpsErrorLogService,
usage_record_model::{UsageRecord, UsageRecordLevel},
},
infra::time::elapsed_millis_i64,
proxy::dispatch::{
@@ -17,7 +18,7 @@ use crate::{
usage_events::{
enrich_event_route_metadata, enrich_usage_record_identity,
event_kind as response_event_kind, reasoning_effort_from_request,
record_dispatch_error_event, record_response_event, DispatchErrorUsageRecord,
record_dispatch_error_event, DispatchErrorLogRecord,
},
},
upstream::{
@@ -39,7 +40,7 @@ use super::{
};
pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
pub(super) usage_records: &'a AdminUsageRecordService,
pub(super) ops_errors: &'a AdminOpsErrorLogService,
pub(super) request_id: &'a str,
pub(super) account_id: &'a str,
pub(super) account_email: Option<&'a str>,
@@ -55,7 +56,7 @@ pub(super) struct ResponseUpstreamErrorEventRecord<'a> {
}
pub(super) struct ResponseStreamFailureEventRecord<'a> {
pub(super) usage_records: &'a AdminUsageRecordService,
pub(super) ops_errors: &'a AdminOpsErrorLogService,
pub(super) request_id: &'a str,
pub(super) account_id: &'a str,
pub(super) route: &'a str,
@@ -73,7 +74,7 @@ pub(super) struct ResponseStreamFailureEventRecord<'a> {
}
pub(super) struct ResponseDispatchErrorEventRecord<'a> {
pub(super) usage_records: &'a AdminUsageRecordService,
pub(super) ops_errors: &'a AdminOpsErrorLogService,
pub(super) request_id: &'a str,
pub(super) account_id: Option<&'a str>,
pub(super) route: &'a str,
@@ -110,8 +111,8 @@ pub(super) async fn record_response_dispatch_error_event(
dispatch_error_upstream_status_code(record.error),
);
}
record_dispatch_error_event(DispatchErrorUsageRecord {
usage_records: record.usage_records,
record_dispatch_error_event(DispatchErrorLogRecord {
ops_errors: record.ops_errors,
request_id: record.request_id,
account_id: record.account_id,
route: record.route,
@@ -135,11 +136,7 @@ pub(super) async fn record_response_upstream_error_event(
Some(backend_transport_name(record.transport)),
);
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
"v1.response",
UsageRecordLevel::Error,
"v1 responses upstream request failed",
);
let mut event = OpsErrorLog::new("v1.response", "v1 responses upstream request failed");
event.request_id = Some(record.request_id.to_string());
event.account_id = Some(record.account_id.to_string());
event.route = Some(record.route.to_string());
@@ -167,8 +164,17 @@ pub(super) async fn record_response_upstream_error_event(
);
}
}
enrich_usage_record_identity(
&mut metadata,
Some(record.request.model()),
record.model,
record.request.client_ip.as_deref(),
record.request.client_user_agent.as_deref(),
reasoning_effort_from_request(record.request),
record.request.service_tier(),
);
event.metadata = metadata;
if let Err(error) = record.usage_records.record(event).await {
if let Err(error) = record.ops_errors.record(event).await {
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record upstream error event");
}
}
@@ -204,25 +210,38 @@ pub(super) async fn record_prefetched_response_stream_failure_event(
Some(record.attempt),
);
}
record_response_event(ResponseUsageRecord {
usage_records: record.usage_records,
request_id: record.request_id,
account_id: record.account_id,
route: record.route,
model: record.model,
requested_model: Some(record.requested_model),
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(),
started_at: record.started_at,
status_code: event_status_code,
level: UsageRecordLevel::Error,
message: "v1 responses stream failed",
metadata,
rate_limit_headers: record.rate_limit_headers,
})
.await;
enrich_event_route_metadata(&mut metadata, record.route);
enrich_usage_record_identity(
&mut metadata,
Some(record.requested_model),
record.model,
record.request.client_ip.as_deref(),
record.request.client_user_agent.as_deref(),
reasoning_effort_from_request(record.request),
record.request.service_tier(),
);
let mut event = OpsErrorLog::new(
response_event_kind(record.route),
"v1 responses stream failed",
);
event.request_id = Some(record.request_id.to_string());
event.account_id = Some(record.account_id.to_string());
event.route = Some(record.route.to_string());
event.model = Some(record.model.to_string());
event.status_code = Some(event_status_code);
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
if !record.rate_limit_headers.is_empty() {
if let Some(object) = metadata.as_object_mut() {
object.insert(
"rateLimitHeaders".to_string(),
serde_json::json!(record.rate_limit_headers),
);
}
}
event.metadata = metadata;
if let Err(error) = record.ops_errors.record(event).await {
tracing::warn!(account_id = %record.account_id, error = %error, "failed to record prefetched stream error event");
}
}
fn ensure_stream_metadata_flag(metadata: &mut Value) {
@@ -469,13 +488,6 @@ pub(super) async fn record_live_response_stream_event(
body,
);
log_live_response_stream_finalized(context, status_code, level, message, &metadata, body);
let mut event = UsageRecord::new(response_event_kind(&context.route), level, message);
event.request_id = Some(context.request_id.clone());
event.account_id = Some(context.account_id.clone());
event.route = Some(context.route.clone());
event.model = Some(context.display_model.clone());
event.status_code = Some(status_code);
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
enrich_usage_record_identity(
&mut metadata,
Some(&context.requested_model),
@@ -485,12 +497,38 @@ pub(super) async fn record_live_response_stream_event(
reasoning_effort_from_request(&context.request),
context.request.service_tier(),
);
event.metadata = metadata;
if let Err(error) = context.usage_records.record(event).await {
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream event");
if is_success_usage_event(level, status_code) {
let mut event = UsageRecord::new(response_event_kind(&context.route), level, message);
event.request_id = Some(context.request_id.clone());
event.account_id = Some(context.account_id.clone());
event.route = Some(context.route.clone());
event.model = Some(context.display_model.clone());
event.status_code = Some(status_code);
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
event.metadata = metadata;
if let Err(error) = context.usage_records.record(event).await {
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream event");
}
} else {
let mut event = OpsErrorLog::new(response_event_kind(&context.route), message);
event.request_id = Some(context.request_id.clone());
event.account_id = Some(context.account_id.clone());
event.route = Some(context.route.clone());
event.model = Some(context.display_model.clone());
event.status_code = Some(status_code);
event.latency_ms = Some(elapsed_millis_i64(context.started_at));
event.metadata = metadata;
if let Err(error) = context.ops_errors.record(event).await {
tracing::warn!(account_id = %context.account_id, error = %error, "failed to record live response stream error event");
}
}
}
fn is_success_usage_event(level: UsageRecordLevel, status_code: i64) -> bool {
level != UsageRecordLevel::Error && status_code < 400
}
fn log_live_response_stream_finalized(
context: &LiveResponseStreamContext,
status_code: i64,
+123 -26
View File
@@ -17,6 +17,7 @@ use tokio::sync::Mutex;
use crate::{
admin::monitoring::{
ops_error_service::AdminOpsErrorLogService,
usage_record_model::{ResponseUsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
},
@@ -110,6 +111,7 @@ pub struct ResponseDispatchService {
session_affinity: Arc<RuntimeSessionAffinityService>,
reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
usage_records: Arc<AdminUsageRecordService>,
ops_errors: Arc<AdminOpsErrorLogService>,
token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
installation_id: Option<String>,
cloudflare: CloudflareRecovery,
@@ -121,6 +123,7 @@ pub(crate) struct ResponseDispatchServiceParts {
pub codex: Arc<CodexBackendClient>,
pub session_affinity: Arc<RuntimeSessionAffinityService>,
pub usage_records: Arc<AdminUsageRecordService>,
pub ops_errors: Arc<AdminOpsErrorLogService>,
pub token_refresh: Arc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
pub installation_id: Option<String>,
pub cloudflare: CloudflareRecovery,
@@ -130,6 +133,34 @@ pub(crate) struct ResponseDispatchServiceParts {
const DEFAULT_REASONING_REPLAY_TTL_SECS: i64 = 55 * 60;
const IMPLICIT_RESUME_MAX_AGE_SECS: i64 = DEFAULT_REASONING_REPLAY_TTL_SECS;
#[derive(Debug, Clone, Default)]
struct AccountAffinityDecision {
preferred_account_id: Option<String>,
request_history_owner: RequestHistoryOwner,
}
#[derive(Debug, Clone, Default)]
enum RequestHistoryOwner {
#[default]
Absent,
Known(String),
Unknown,
}
impl AccountAffinityDecision {
fn preferred_account_id(&self) -> Option<&str> {
self.preferred_account_id.as_deref()
}
fn should_strip_request_history_for(&self, account_id: &str) -> bool {
match &self.request_history_owner {
RequestHistoryOwner::Absent => false,
RequestHistoryOwner::Known(history_account_id) => history_account_id != account_id,
RequestHistoryOwner::Unknown => true,
}
}
}
/// Responses live SSE 响应体流。
pub type ResponseBodyStream =
Pin<Box<dyn Stream<Item = Result<Bytes, ResponseDispatchStreamError>> + Send + 'static>>;
@@ -150,6 +181,7 @@ impl ResponseDispatchService {
DEFAULT_REASONING_REPLAY_TTL_SECS,
)))),
usage_records: parts.usage_records,
ops_errors: parts.ops_errors,
token_refresh: parts.token_refresh,
installation_id: parts.installation_id,
cloudflare: parts.cloudflare,
@@ -260,7 +292,31 @@ impl ResponseDispatchService {
Some(snapshot)
}
async fn preferred_account_id_for_request(
async fn account_affinity_for_request(
&self,
request: &CodexResponsesRequest,
now: DateTime<Utc>,
) -> AccountAffinityDecision {
let has_previous_response_id = request.previous_response_id().is_some();
if let Some(account_id) = self.previous_response_account_id(request, now).await {
return AccountAffinityDecision {
preferred_account_id: Some(account_id.clone()),
request_history_owner: RequestHistoryOwner::Known(account_id),
};
}
let conversation_account_id = self.conversation_account_id(request, now).await;
AccountAffinityDecision {
preferred_account_id: conversation_account_id,
request_history_owner: if has_previous_response_id {
RequestHistoryOwner::Unknown
} else {
RequestHistoryOwner::Absent
},
}
}
async fn previous_response_account_id(
&self,
request: &CodexResponsesRequest,
now: DateTime<Utc>,
@@ -271,6 +327,21 @@ impl ResponseDispatchService {
.await
}
async fn conversation_account_id(
&self,
request: &CodexResponsesRequest,
now: DateTime<Utc>,
) -> Option<String> {
let conversation_id = request
.prompt_cache_key()
.map(str::trim)
.filter(|value| !value.is_empty())?;
let variant_hash = compute_variant_hash(request);
self.session_affinity
.lookup_latest_account_by_conversation(conversation_id, None, Some(&variant_hash), now)
.await
}
async fn recover_request_history(
&self,
request: &mut CodexResponsesRequest,
@@ -279,13 +350,22 @@ impl ResponseDispatchService {
if let Some(previous_response_id) = request.previous_response_id() {
self.session_affinity.forget(previous_response_id).await;
}
if let Some(snapshot) = implicit_resume.take() {
snapshot.restore(request);
request.set_previous_response_id(None);
request.turn_state = None;
} else {
strip_request_history(request);
restore_request_without_history(request, implicit_resume);
}
fn strip_history_if_account_changed(
request: &mut CodexResponsesRequest,
implicit_resume: &mut Option<ImplicitResumeSnapshot>,
account_affinity: &AccountAffinityDecision,
acquired_account_id: &str,
) {
if request.previous_response_id().is_none()
|| !account_affinity.should_strip_request_history_for(acquired_account_id)
{
return;
}
restore_request_without_history(request, implicit_resume);
}
async fn apply_cascading_ban_defense(
@@ -326,13 +406,7 @@ impl ResponseDispatchService {
let response_id_to_forget = explicit_previous_response_id
.or(request.previous_response_id())
.map(str::to_string);
if let Some(snapshot) = implicit_resume.take() {
snapshot.restore(request);
request.set_previous_response_id(None);
request.turn_state = None;
} else {
strip_request_history(request);
}
restore_request_without_history(request, implicit_resume);
if let Some(response_id) = response_id_to_forget {
self.session_affinity.forget(&response_id).await;
}
@@ -380,9 +454,9 @@ impl ResponseDispatchService {
let now = Utc::now();
let explicit_previous_response_id = request.previous_response_id().map(ToString::to_string);
let mut implicit_resume = self.prepare_response_session(&mut request).await;
let preferred_account_id = self.preferred_account_id_for_request(&request, now).await;
let account_affinity = self.account_affinity_for_request(&request, now).await;
let mut acquire_request = AccountAcquireRequest::new(request.model(), now);
if let Some(preferred_account_id) = preferred_account_id.as_deref() {
if let Some(preferred_account_id) = account_affinity.preferred_account_id() {
acquire_request = acquire_request.with_preferred_account_id(preferred_account_id);
}
let mut excluded_account_ids = Vec::new();
@@ -493,11 +567,17 @@ impl ResponseDispatchService {
self.apply_cascading_ban_defense(
&mut request,
&mut implicit_resume,
preferred_account_id.as_deref(),
account_affinity.preferred_account_id(),
&acquired.account.id,
explicit_previous_response_id.as_deref(),
)
.await;
Self::strip_history_if_account_changed(
&mut request,
&mut implicit_resume,
&account_affinity,
&acquired.account.id,
);
self.account_pool.wait_for_request_interval(&acquired).await;
let account = acquired.account;
let release_account_id = account.id.clone();
@@ -758,7 +838,7 @@ impl ResponseDispatchService {
}
Err(error) => {
record_response_upstream_error_event(ResponseUpstreamErrorEventRecord {
usage_records: &self.usage_records,
ops_errors: &self.ops_errors,
request_id,
account_id: &release_account_id,
account_email: account.email.as_deref(),
@@ -904,7 +984,7 @@ impl ResponseDispatchService {
error: &ResponseDispatchError,
) {
record_response_dispatch_error_event(ResponseDispatchErrorEventRecord {
usage_records: &self.usage_records,
ops_errors: &self.ops_errors,
request_id,
account_id: details.account_id,
route,
@@ -935,9 +1015,9 @@ impl ResponseDispatchService {
let now = Utc::now();
let explicit_previous_response_id = request.previous_response_id().map(ToString::to_string);
let mut implicit_resume = self.prepare_response_session(&mut request).await;
let preferred_account_id = self.preferred_account_id_for_request(&request, now).await;
let account_affinity = self.account_affinity_for_request(&request, now).await;
let mut acquire_request = AccountAcquireRequest::new(request.model(), now);
if let Some(preferred_account_id) = preferred_account_id.as_deref() {
if let Some(preferred_account_id) = account_affinity.preferred_account_id() {
acquire_request = acquire_request.with_preferred_account_id(preferred_account_id);
}
let mut excluded_account_ids = Vec::new();
@@ -1053,11 +1133,17 @@ impl ResponseDispatchService {
self.apply_cascading_ban_defense(
&mut request,
&mut implicit_resume,
preferred_account_id.as_deref(),
account_affinity.preferred_account_id(),
&acquired.account.id,
explicit_previous_response_id.as_deref(),
)
.await;
Self::strip_history_if_account_changed(
&mut request,
&mut implicit_resume,
&account_affinity,
&acquired.account.id,
);
self.account_pool.wait_for_request_interval(&acquired).await;
let account = acquired.account;
@@ -1236,7 +1322,7 @@ impl ResponseDispatchService {
if let ResponseDispatchError::Upstream(upstream_error) = &error {
record_response_upstream_error_event(
ResponseUpstreamErrorEventRecord {
usage_records: &self.usage_records,
ops_errors: &self.ops_errors,
request_id,
account_id: &release_account_id,
account_email: account.email.as_deref(),
@@ -1341,7 +1427,7 @@ impl ResponseDispatchService {
self.account_pool.release(&release_account_id).await;
record_prefetched_response_stream_failure_event(
ResponseStreamFailureEventRecord {
usage_records: &self.usage_records,
ops_errors: &self.ops_errors,
request_id,
account_id: &release_account_id,
route,
@@ -1367,6 +1453,7 @@ impl ResponseDispatchService {
session_affinity: Arc::clone(&self.session_affinity),
reasoning_replay: Arc::clone(&self.reasoning_replay),
usage_records: Arc::clone(&self.usage_records),
ops_errors: Arc::clone(&self.ops_errors),
cloudflare: self.cloudflare.clone(),
account_id: account.id,
account_plan_type: account.plan_type,
@@ -1503,7 +1590,7 @@ impl ResponseDispatchService {
Err(error) => {
self.account_pool.release(&release_account_id).await;
record_response_upstream_error_event(ResponseUpstreamErrorEventRecord {
usage_records: &self.usage_records,
ops_errors: &self.ops_errors,
request_id,
account_id: &release_account_id,
account_email: account.email.as_deref(),
@@ -1795,7 +1882,7 @@ impl ResponseDispatchService {
error: &ResponseDispatchError,
) {
record_response_dispatch_error_event(ResponseDispatchErrorEventRecord {
usage_records: &self.usage_records,
ops_errors: &self.ops_errors,
request_id,
account_id,
route: "/v1/responses/compact",
@@ -1814,3 +1901,13 @@ fn strip_request_history(request: &mut CodexResponsesRequest) {
request.set_previous_response_id(None);
request.turn_state = None;
}
fn restore_request_without_history(
request: &mut CodexResponsesRequest,
implicit_resume: &mut Option<ImplicitResumeSnapshot>,
) {
if let Some(snapshot) = implicit_resume.take() {
snapshot.restore(request);
}
strip_request_history(request);
}
@@ -5,7 +5,8 @@ use tokio::sync::Mutex;
use crate::{
admin::monitoring::{
usage_record_model::UsageRecordLevel, usage_record_service::AdminUsageRecordService,
ops_error_service::AdminOpsErrorLogService, usage_record_model::UsageRecordLevel,
usage_record_service::AdminUsageRecordService,
},
infra::time::elapsed_millis_i64,
proxy::{
@@ -49,6 +50,7 @@ pub(super) struct LiveResponseStreamContext {
pub(super) session_affinity: Arc<RuntimeSessionAffinityService>,
pub(super) reasoning_replay: Arc<Mutex<ReasoningReplayCache>>,
pub(super) usage_records: Arc<AdminUsageRecordService>,
pub(super) ops_errors: Arc<AdminOpsErrorLogService>,
pub(super) cloudflare: CloudflareRecovery,
pub(super) account_id: String,
pub(super) account_plan_type: Option<String>,
+47 -11
View File
@@ -106,6 +106,38 @@ impl SessionAffinityMap {
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<String> {
self.latest_entry_by_conversation(conversation_id, max_age, variant_hash, now)
.map(|(response_id, _)| response_id.clone())
}
pub fn lookup_latest_account_by_conversation(
&self,
conversation_id: &str,
max_age: Option<Duration>,
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<String> {
self.latest_entry_by_conversation(conversation_id, max_age, variant_hash, now)
.map(|(_, entry)| entry.account_id.clone())
}
pub fn forget(&mut self, response_id: &str) -> bool {
self.entries.remove(response_id).is_some()
}
fn active_entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<&SessionAffinityEntry> {
self.entries
.get(response_id)
.filter(|entry| self.entry_is_active(entry, now))
}
fn latest_entry_by_conversation(
&self,
conversation_id: &str,
max_age: Option<Duration>,
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<(&String, &SessionAffinityEntry)> {
self.entries
.iter()
.filter(|(_, entry)| entry.conversation_id == conversation_id)
@@ -118,17 +150,6 @@ impl SessionAffinityMap {
max_age.is_none_or(|max_age| now.signed_duration_since(entry.created_at) <= max_age)
})
.max_by_key(|(_, entry)| entry.created_at)
.map(|(response_id, _)| response_id.clone())
}
pub fn forget(&mut self, response_id: &str) -> bool {
self.entries.remove(response_id).is_some()
}
fn active_entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<&SessionAffinityEntry> {
self.entries
.get(response_id)
.filter(|entry| self.entry_is_active(entry, now))
}
fn stored_record_is_active(&self, record: &StoredSessionAffinity, now: DateTime<Utc>) -> bool {
@@ -630,6 +651,21 @@ impl RuntimeSessionAffinityService {
.lookup_latest_response_by_conversation(conversation_id, max_age, variant_hash, now)
}
pub async fn lookup_latest_account_by_conversation(
&self,
conversation_id: &str,
max_age: Option<Duration>,
variant_hash: Option<&str>,
now: DateTime<Utc>,
) -> Option<String> {
self.map.read().await.lookup_latest_account_by_conversation(
conversation_id,
max_age,
variant_hash,
now,
)
}
pub async fn forget(&self, response_id: &str) -> bool {
let in_memory = self.map.write().await.forget(response_id);
match self.store.forget(response_id).await {
+4 -1
View File
@@ -103,7 +103,10 @@ pub(crate) async fn verify_acquired_quota_if_required(
.apply_quota_snapshot(&account_id, &quota)
.await;
if quota_snapshot_limit_reached(&quota) {
context.account_pool.release(&account_id).await;
context
.account_pool
.release_without_request_usage(&account_id)
.await;
context.excluded_account_ids.push(account_id.clone());
*context.verify_attempts += 1;
let max_attempts_reached = *context.verify_attempts >= MAX_QUOTA_VERIFY_ATTEMPTS;
+8 -11
View File
@@ -6,15 +6,16 @@ use serde_json::Value;
use crate::{
admin::monitoring::{
usage_record_model::{ResponseUsageRecord, UsageRecord, UsageRecordLevel},
usage_record_service::AdminUsageRecordService,
ops_error_model::OpsErrorLog,
ops_error_service::AdminOpsErrorLogService,
usage_record_model::{ResponseUsageRecord, UsageRecord},
},
infra::time::elapsed_millis_i64,
upstream::protocol::responses::{CodexCompactRequest, CodexResponsesRequest},
};
pub(super) struct DispatchErrorUsageRecord<'a> {
pub usage_records: &'a AdminUsageRecordService,
pub(super) struct DispatchErrorLogRecord<'a> {
pub ops_errors: &'a AdminOpsErrorLogService,
pub request_id: &'a str,
pub account_id: Option<&'a str>,
pub route: &'a str,
@@ -25,15 +26,11 @@ pub(super) struct DispatchErrorUsageRecord<'a> {
pub metadata: Value,
}
pub(super) async fn record_dispatch_error_event(record: DispatchErrorUsageRecord<'_>) {
pub(super) async fn record_dispatch_error_event(record: DispatchErrorLogRecord<'_>) {
let mut metadata = record.metadata;
enrich_event_route_metadata(&mut metadata, record.route);
let mut event = UsageRecord::new(
event_kind(record.route),
UsageRecordLevel::Error,
record.message,
);
let mut event = OpsErrorLog::new(event_kind(record.route), record.message);
event.request_id = Some(record.request_id.to_string());
event.account_id = record.account_id.map(ToString::to_string);
event.route = Some(record.route.to_string());
@@ -42,7 +39,7 @@ pub(super) async fn record_dispatch_error_event(record: DispatchErrorUsageRecord
event.latency_ms = Some(elapsed_millis_i64(record.started_at));
event.metadata = metadata;
if let Err(error) = record.usage_records.record(event).await {
if let Err(error) = record.ops_errors.record(event).await {
tracing::warn!(
account_id = record.account_id.unwrap_or(""),
error = %error,
+3 -4
View File
@@ -5,7 +5,6 @@ use chrono::Utc;
use crate::config::schema::AppConfig;
use crate::config::settings::RuntimeSettingsService;
use crate::http::middleware::request_id::TrustedProxyConfig;
use crate::infra::database::connect_sqlite;
use crate::infra::logging::{init_tracing, LogGuard, RotationConfig};
use crate::infra::paths::{ensure_data_dir, load_or_create_installation_id};
@@ -13,7 +12,7 @@ use crate::runtime::services::{BackgroundTaskStores, Services};
use crate::runtime::shutdown::{restart_executable_path, shutdown_signal};
use crate::runtime::state::{AppState, RuntimeConfig};
use crate::runtime::tasks::coordinator::TaskCoordinator;
use crate::upstream::fingerprint::{Fingerprint, FingerprintRepository};
use crate::upstream::fingerprint::{Fingerprint, FingerprintRepository, RuntimeFingerprint};
const RESTART_DELAY_ENV: &str = "CPR_RESTART_DELAY_MS";
@@ -81,6 +80,7 @@ async fn build_application(
let runtime_fingerprint = fingerprint_repository
.ensure_current_seed(&default_fingerprint)
.await?;
let runtime_fingerprint = RuntimeFingerprint::new(runtime_fingerprint);
let runtime_config = RuntimeConfig::from(config.clone());
let services = Services::try_with_installation_id(
&config,
@@ -112,11 +112,10 @@ 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 { services };
Ok((
crate::http::router::router_with_trusted_proxies(trusted_proxies).with_state(app_state),
crate::http::router::router().with_state(app_state),
task_coordinator,
))
}
+20 -6
View File
@@ -6,6 +6,8 @@ use crate::{
admin::monitoring::{
account_usage_service::AdminUsageService,
account_usage_store::SqliteUsageStore,
ops_error_service::AdminOpsErrorLogService,
ops_error_store::SqliteOpsErrorLogStore,
usage_record_service::AdminUsageRecordService,
usage_record_store::{SqliteUsageRecordStore, DEFAULT_USAGE_RECORD_CAPTURE_BODY},
},
@@ -35,7 +37,7 @@ use crate::{
},
},
upstream::{
fingerprint::{Fingerprint, FingerprintRepository},
fingerprint::{FingerprintRepository, RuntimeFingerprint},
models::{
config::ModelConfig,
service::{ModelService, ModelServiceError},
@@ -91,6 +93,7 @@ pub struct Services {
pub refresh_policy: RuntimeRefreshPolicy,
pub admin_accounts: StdArc<AdminAccountService>,
pub usage_records: StdArc<AdminUsageRecordService>,
pub ops_errors: StdArc<AdminOpsErrorLogService>,
pub usage: StdArc<AdminUsageService>,
pub account_pool: StdArc<RuntimeAccountPoolService>,
pub token_refresh: StdArc<RuntimeTokenRefreshService<OpenAiTokenClient>>,
@@ -98,7 +101,7 @@ pub struct Services {
pub session_affinity: StdArc<RuntimeSessionAffinityService>,
pub codex: StdArc<CodexBackendClient>,
pub websocket_pool: Option<StdArc<CodexWebSocketPool>>,
pub fingerprint: Fingerprint,
pub fingerprint: RuntimeFingerprint,
pub installation_id: Option<String>,
pub background_tasks: BackgroundTaskStores,
}
@@ -120,7 +123,11 @@ impl UsageRecordOptions {
}
impl Services {
pub fn new(config: &AppConfig, stores: BackgroundTaskStores, fingerprint: Fingerprint) -> Self {
pub fn new(
config: &AppConfig,
stores: BackgroundTaskStores,
fingerprint: RuntimeFingerprint,
) -> Self {
Self::try_new(config, stores, fingerprint)
.expect("failed to build runtime services with configured TLS transport")
}
@@ -128,7 +135,7 @@ impl Services {
pub fn try_new(
config: &AppConfig,
stores: BackgroundTaskStores,
fingerprint: Fingerprint,
fingerprint: RuntimeFingerprint,
) -> Result<Self, CustomCaError> {
Self::try_with_installation_id(config, stores, fingerprint, None)
}
@@ -136,7 +143,7 @@ impl Services {
pub(crate) fn try_with_installation_id(
config: &AppConfig,
stores: BackgroundTaskStores,
fingerprint: Fingerprint,
fingerprint: RuntimeFingerprint,
installation_id: Option<String>,
) -> Result<Self, CustomCaError> {
Self::try_with_installation_id_and_usage_record_options(
@@ -151,7 +158,7 @@ impl Services {
pub fn try_with_installation_id_and_usage_record_options(
config: &AppConfig,
stores: BackgroundTaskStores,
fingerprint: Fingerprint,
fingerprint: RuntimeFingerprint,
installation_id: Option<String>,
usage_record_options: UsageRecordOptions,
) -> Result<Self, CustomCaError> {
@@ -189,6 +196,10 @@ impl Services {
usage_record_options.enabled,
usage_record_options.capture_body,
));
let ops_errors = StdArc::new(AdminOpsErrorLogService::new(
SqliteOpsErrorLogStore::new(stores.usage_records.pool().clone()),
usage_record_options.capture_body,
));
let account_pool = StdArc::new(RuntimeAccountPoolService::new(
account_store_trait.clone(),
account_pool_options_from_config(config),
@@ -243,6 +254,7 @@ impl Services {
models: models.clone(),
account_pool: account_pool.clone(),
token_refresher: StdArc::new(token_client),
refresh_leases: stores.refresh_leases.clone(),
oauth: crate::admin::accounts::service::oauth::AccountOAuthService::new(
reqwest::Client::new(),
config.auth.oauth_client_id.clone(),
@@ -265,6 +277,7 @@ impl Services {
codex: codex.clone(),
session_affinity: session_affinity.clone(),
usage_records: usage_records.clone(),
ops_errors: ops_errors.clone(),
token_refresh: token_refresh.clone(),
installation_id: installation_id.clone(),
cloudflare: cloudflare_recovery,
@@ -280,6 +293,7 @@ impl Services {
refresh_policy,
admin_accounts,
usage_records,
ops_errors,
usage,
account_pool,
token_refresh,
+8 -7
View File
@@ -86,16 +86,17 @@ impl TaskCoordinator {
.with_account_pool(services.account_pool.clone())
.start(),
);
coordinator.push(
"fingerprint_update",
coordinator.push("fingerprint_update", {
let fingerprint = services.fingerprint.snapshot();
start_fingerprint_update_task(
Some(stores.fingerprints.clone()),
stores.fingerprints.clone(),
services.fingerprint.clone(),
CODEX_DESKTOP_APPCAST_URL.to_string(),
PathBuf::from(DEFAULT_EXTRACTED_FINGERPRINT_PATH),
services.fingerprint.app_version.clone(),
services.fingerprint.build_number.clone(),
),
);
fingerprint.app_version,
fingerprint.build_number,
)
});
if let Some(pool) = &services.websocket_pool {
coordinator.push(
"websocket_pool",
@@ -2,9 +2,11 @@
use std::path::PathBuf;
use tracing::warn;
use tracing::{info, warn};
use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker, APPCAST_POLL_INTERVAL};
use crate::upstream::fingerprint::{
FingerprintRepository, RuntimeFingerprint, UpdateChecker, APPCAST_POLL_INTERVAL,
};
use super::{
coordinator::SchedulerHandle,
@@ -19,7 +21,8 @@ pub const DEFAULT_EXTRACTED_FINGERPRINT_PATH: &str = "data/extracted-fingerprint
/// 启动指纹自动更新后台任务。
pub fn start_fingerprint_update_task(
repository: Option<FingerprintRepository>,
repository: FingerprintRepository,
runtime_fingerprint: RuntimeFingerprint,
appcast_url: String,
extracted_fingerprint_path: PathBuf,
current_version: String,
@@ -36,6 +39,7 @@ pub fn start_fingerprint_update_task(
spawn_periodic_task(
FingerprintUpdateTask {
checker: update_checker,
runtime_fingerprint,
first_tick: true,
},
PeriodicTaskConfig::new(
@@ -48,6 +52,7 @@ pub fn start_fingerprint_update_task(
struct FingerprintUpdateTask {
checker: UpdateChecker,
runtime_fingerprint: RuntimeFingerprint,
first_tick: bool,
}
@@ -55,11 +60,24 @@ 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 定期版本检查失败");
match self.checker.check_and_apply_update().await {
Ok(Some(fingerprint)) => {
let app_version = fingerprint.app_version.clone();
let build_number = fingerprint.build_number.clone();
self.runtime_fingerprint.replace(fingerprint);
info!(
app_version = %app_version,
build_number = %build_number,
"fingerprint 运行时版本已更新"
);
}
Ok(None) => {}
Err(error) => {
if first_tick {
warn!(error = %error, "fingerprint 首次版本检查失败");
} else {
warn!(error = %error, "fingerprint 定期版本检查失败");
}
}
}
})
+3 -1
View File
@@ -88,7 +88,9 @@ impl ModelRefreshTask {
)
.await;
for selected in &selected_accounts {
self.account_pool.release(&selected.account.id).await;
self.account_pool
.release_without_request_usage(&selected.account.id)
.await;
}
let result = result?;
@@ -613,7 +613,6 @@ pub fn parse_account_status(status: &str) -> Result<AccountStatus, &'static str>
"disabled" => Ok(AccountStatus::Disabled),
"expired" => Ok(AccountStatus::Expired),
"quota_exhausted" => Ok(AccountStatus::QuotaExhausted),
"refreshing" => Ok(AccountStatus::Refreshing),
"banned" => Ok(AccountStatus::Banned),
_ => Err("invalid account status"),
}
-3
View File
@@ -12,8 +12,6 @@ pub enum AccountStatus {
Expired,
/// 当前配额已经耗尽。
QuotaExhausted,
/// 账号正处于刷新流程中。
Refreshing,
/// 账号被显式禁用。
Disabled,
/// 账号被上游封禁。
@@ -27,7 +25,6 @@ impl AccountStatus {
Self::Active => "active",
Self::Expired => "expired",
Self::QuotaExhausted => "quota_exhausted",
Self::Refreshing => "refreshing",
Self::Disabled => "disabled",
Self::Banned => "banned",
}
+129 -196
View File
@@ -1,7 +1,6 @@
//! 账号池调度策略。
use std::{
cmp::Ordering,
collections::{BTreeMap, BTreeSet, VecDeque},
sync::Arc as StdArc,
time::Duration as StdDuration,
@@ -11,7 +10,6 @@ use chrono::{DateTime, Duration, Utc};
use indexmap::IndexMap;
use thiserror::Error;
use crate::infra::time::parse_rfc3339_utc;
use crate::upstream::accounts::quota::{quota_snapshot_limit_reached, quota_snapshot_reset_at};
use crate::upstream::accounts::store::AccountStore;
use crate::upstream::accounts::token_refresh::{jwt_expiry, JwtExpiry};
@@ -23,27 +21,12 @@ use crate::upstream::accounts::{
use crate::upstream::protocol::events::{
parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage,
};
use crate::upstream::scheduler::{
candidates::{self, CandidateFilter, CandidateRequest},
AccountScheduler, ScoreWeights,
};
/// 账号轮转策略。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RotationStrategy {
/// 优先选择当前窗口用量最低的账号。
LeastUsed,
/// 按候选账号顺序循环选择。
RoundRobin,
/// 优先选择最近使用过的账号。
Sticky,
}
impl RotationStrategy {
fn as_str(self) -> &'static str {
match self {
Self::LeastUsed => "least_used",
Self::RoundRobin => "round_robin",
Self::Sticky => "sticky",
}
}
}
pub use crate::upstream::scheduler::RotationStrategy;
/// 账号池配置。
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -69,7 +52,7 @@ impl Default for AccountPoolOptions {
Self {
max_concurrent_per_account: 3,
stale_slot_ttl: Duration::minutes(5),
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::Smart,
skip_quota_limited: true,
tier_priority: Vec::new(),
model_plan_allowlist: BTreeMap::new(),
@@ -260,13 +243,17 @@ pub struct AccountWindowUsageDelta {
pub image_request_failed: bool,
}
/// 纯内存账号池,负责账号调度和运行时状态维护。
/// 纯内存账号池:负责账号存储、在途槽位与运行时状态维护。
///
/// 账号**选择**已剥离到 [`crate::upstream::scheduler`]:候选过滤走
/// [`candidates::filter`],策略选择走 [`AccountScheduler::select`]。账号池只提供选择
/// 所需的只读视图(候选切片 + `slot_count`),并持有调度器实例。
#[derive(Debug)]
pub struct AccountPool {
accounts: IndexMap<String, Account>,
slots: BTreeMap<String, VecDeque<AccountSlot>>,
options: AccountPoolOptions,
rotation_cursor: usize,
scheduler: AccountScheduler,
}
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -295,7 +282,7 @@ impl AccountPool {
accounts: IndexMap::new(),
slots: BTreeMap::new(),
options,
rotation_cursor: 0,
scheduler: AccountScheduler::new(ScoreWeights::default()),
}
}
@@ -309,10 +296,11 @@ impl AccountPool {
self.accounts.get(account_id).cloned()
}
/// 移除账号及其在途槽位。
/// 移除账号及其在途槽位与运行时反馈。
pub fn remove(&mut self, account_id: &str) -> bool {
let removed = self.accounts.shift_remove(account_id).is_some();
self.slots.remove(account_id);
self.scheduler.forget_feedback(account_id);
removed
}
@@ -320,17 +308,24 @@ impl AccountPool {
pub fn clear(&mut self) {
self.accounts.clear();
self.slots.clear();
self.rotation_cursor = 0;
self.scheduler.reset_cursor();
self.scheduler.clear_feedback();
}
/// 替换账号池运行参数。
pub fn set_options(&mut self, options: AccountPoolOptions) {
if self.options.rotation_strategy != options.rotation_strategy {
self.rotation_cursor = 0;
self.scheduler.reset_cursor();
}
self.options = options;
}
/// 回灌一次请求结果到调度器的运行时反馈(供 Smart 打分)。
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
self.scheduler
.report_feedback(account_id, success, first_token_ms);
}
/// 更新账号状态。
pub fn set_status(&mut self, account_id: &str, status: AccountStatus) -> bool {
let Some(account) = self.accounts.get_mut(account_id) else {
@@ -485,7 +480,9 @@ impl AccountPool {
) -> AccountAcquireWithStatusRefresh {
self.cleanup_stale_slots(request.now);
let refreshed_accounts = self.refresh_account_statuses(request.now);
let candidates = self.candidates(request);
let candidates = self.filter_candidates(request);
// 亲和优先:required 硬锁定该账号;preferred 若在候选内则直接选中;
// 二者都不命中时,回退到调度器按当前策略选择。
let selected = if let Some(required_account_id) = &request.required_account_id {
candidates.first().cloned().map(|account| {
(
@@ -495,38 +492,24 @@ impl AccountPool {
)
})
} else if let Some(preferred_account_id) = &request.preferred_account_id {
let preferred = candidates
candidates
.iter()
.find(|account| account.id == *preferred_account_id)
.cloned();
preferred.map(|account| {
(
account,
self.previous_slot_at(preferred_account_id),
"preferred",
)
})
.cloned()
.map(|account| {
(
account,
self.previous_slot_at(preferred_account_id),
"preferred",
)
})
} else {
None
}
.or_else(|| {
Some(match self.options.rotation_strategy {
RotationStrategy::LeastUsed => {
let account = self.select_least_used(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous, "least_used")
}
RotationStrategy::RoundRobin => {
let account = self.select_round_robin(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous, "round_robin")
}
RotationStrategy::Sticky => {
let account = Self::select_sticky(&candidates)?;
let previous = self.previous_slot_at(&account.id);
(account, previous, "sticky")
}
})
let account = self.select_candidate(&candidates, request.now)?;
let previous = self.previous_slot_at(&account.id);
Some((account, previous, self.options.rotation_strategy.as_str()))
});
let Some((selected, previous_slot_at, selection_source)) = selected else {
return AccountAcquireWithStatusRefresh {
@@ -540,6 +523,7 @@ impl AccountPool {
selection_source,
account_id = %selected.id,
candidate_count = candidates.len(),
slot_count = self.slot_count(&selected.id),
request_count = selected.request_count,
window_request_count = selected.window_request_count,
window_reset_at = ?selected.window_reset_at,
@@ -547,7 +531,6 @@ impl AccountPool {
quota_limit_reached = selected.quota_limit_reached,
quota_cooldown_until = ?selected.quota_cooldown_until,
previous_slot_at = ?previous_slot_at,
rotation_cursor = self.rotation_cursor,
"account selected for upstream request"
);
@@ -591,6 +574,22 @@ impl AccountPool {
/// 释放指定账号的一个在途槽位。
pub fn release(&mut self, account_id: &str) -> Option<ReleasedAccountUsage> {
self.release_slot(account_id, true)
}
/// 释放指定账号的一个在途槽位,不累计请求用量。
pub fn release_without_request_usage(
&mut self,
account_id: &str,
) -> Option<ReleasedAccountUsage> {
self.release_slot(account_id, false)
}
fn release_slot(
&mut self,
account_id: &str,
record_request_usage: bool,
) -> Option<ReleasedAccountUsage> {
self.accounts.get(account_id)?;
let mut remove_slots = false;
@@ -602,7 +601,9 @@ impl AccountPool {
if remove_slots {
self.slots.remove(account_id);
}
self.mark_request_usage(account_id, Utc::now());
if record_request_usage {
self.mark_request_usage(account_id, Utc::now());
}
Some(ReleasedAccountUsage {
model: slot.and_then(|slot| slot.model),
})
@@ -665,41 +666,48 @@ impl AccountPool {
let refreshed_accounts = self.refresh_account_statuses(now);
let mut by_plan = IndexMap::<String, Vec<Account>>::new();
for account in self.accounts.values() {
if !self.is_model_refresh_available(account, now) {
continue;
let max_concurrent = self.options.max_concurrent_per_account;
let skip_quota_limited = self.options.skip_quota_limited;
{
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
for account in self.accounts.values() {
if !candidates::is_model_refresh_available(
account,
max_concurrent,
skip_quota_limited,
&slot_count,
now,
) {
continue;
}
let Some(plan_type) = account.plan_type.as_ref() else {
continue;
};
by_plan
.entry(plan_type.clone())
.or_default()
.push(account.clone());
}
let Some(plan_type) = account.plan_type.as_ref() else {
continue;
};
by_plan
.entry(plan_type.clone())
.or_default()
.push(account.clone());
}
let strategy = self.options.rotation_strategy;
let mut accounts = Vec::new();
for (plan_type, group) in by_plan {
let selected = match self.options.rotation_strategy {
RotationStrategy::LeastUsed => self.select_least_used(&group),
RotationStrategy::RoundRobin => self.select_round_robin(&group),
RotationStrategy::Sticky => Self::select_sticky(&group),
};
let Some(account) = selected else {
let Some(account) = self.select_candidate(&group, now) else {
continue;
};
tracing::info!(
plan_type,
rotation_strategy = self.options.rotation_strategy.as_str(),
rotation_strategy = strategy.as_str(),
account_id = %account.id,
candidate_count = group.len(),
slot_count = self.slot_count(&account.id),
request_count = account.request_count,
window_request_count = account.window_request_count,
window_reset_at = ?account.window_reset_at,
last_used_at = account.last_used_at.as_deref().unwrap_or_default(),
quota_limit_reached = account.quota_limit_reached,
quota_cooldown_until = ?account.quota_cooldown_until,
rotation_cursor = self.rotation_cursor,
"account selected for model refresh"
);
self.push_slot(&account.id, now, None);
@@ -712,76 +720,6 @@ impl AccountPool {
}
}
fn select_least_used(&mut self, candidates: &[Account]) -> Option<Account> {
let mut sorted = candidates.iter().collect::<Vec<_>>();
sorted.sort_by(|a, b| compare_least_used(a, b));
let best = *sorted.first()?;
let tied_count = sorted
.iter()
.take_while(|account| compare_least_used(best, account) == Ordering::Equal)
.count();
let index = self.rotation_cursor % tied_count;
self.rotation_cursor = self.rotation_cursor.wrapping_add(1);
Some((*sorted[index]).clone())
}
fn select_sticky(candidates: &[Account]) -> Option<Account> {
let mut selected = candidates.first()?;
for candidate in &candidates[1..] {
if compare_last_used(
candidate.last_used_at.as_deref(),
selected.last_used_at.as_deref(),
) == Ordering::Greater
{
selected = candidate;
}
}
Some((*selected).clone())
}
fn select_round_robin(&mut self, candidates: &[Account]) -> Option<Account> {
if candidates.is_empty() {
return None;
}
self.rotation_cursor %= candidates.len();
let index = self.rotation_cursor;
self.rotation_cursor = self.rotation_cursor.wrapping_add(1);
Some(candidates[index].clone())
}
fn candidates(&self, request: &AccountAcquireRequest) -> Vec<Account> {
let mut candidates = self
.accounts
.values()
.filter(|account| self.is_base_available(account, request))
.cloned()
.collect::<Vec<_>>();
if let Some(required_account_id) = &request.required_account_id {
candidates.retain(|account| account.id == *required_account_id);
return candidates;
}
if let Some(best_tier) = self.best_available_tier(&candidates) {
candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str()));
}
candidates
}
fn is_base_available(&self, account: &Account, request: &AccountAcquireRequest) -> bool {
account.status == AccountStatus::Active
&& self.slot_count(&account.id) < self.options.max_concurrent_per_account
&& AccountService::quota_available_at(
account,
request.now,
self.options.skip_quota_limited,
)
&& self.is_model_allowed(account, &request.model)
&& !request
.exclude_account_ids
.iter()
.any(|account_id| account_id == &account.id)
&& AccountService::cloudflare_available_at(account, request.now)
}
fn slot_count(&self, account_id: &str) -> usize {
self.slots.get(account_id).map_or(0, VecDeque::len)
}
@@ -803,13 +741,6 @@ impl AccountPool {
});
}
fn is_model_refresh_available(&self, account: &Account, now: DateTime<Utc>) -> bool {
account.status == AccountStatus::Active
&& self.slot_count(&account.id) < self.options.max_concurrent_per_account
&& AccountService::quota_available_at(account, now, self.options.skip_quota_limited)
&& AccountService::cloudflare_available_at(account, now)
}
fn mark_request_usage(&mut self, account_id: &str, now: DateTime<Utc>) -> Option<Account> {
let account = self.accounts.get_mut(account_id)?;
account.last_used_at = Some(now.to_rfc3339());
@@ -859,24 +790,33 @@ impl AccountPool {
});
}
fn is_model_allowed(&self, account: &Account, model: &str) -> bool {
let Some(allowed_plans) = self.options.model_plan_allowlist.get(model) else {
return true;
};
let Some(plan_type) = account.plan_type.as_deref() else {
return false;
};
allowed_plans.iter().any(|plan| plan == plan_type)
|| !self.options.fetched_model_plan_types.contains(plan_type)
/// 过滤出本次请求的可用候选集合(委托 [`candidates::filter`])。
fn filter_candidates(&self, request: &AccountAcquireRequest) -> Vec<Account> {
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
candidates::filter(
self.accounts.values(),
&CandidateFilter {
max_concurrent_per_account: self.options.max_concurrent_per_account,
skip_quota_limited: self.options.skip_quota_limited,
tier_priority: &self.options.tier_priority,
model_plan_allowlist: &self.options.model_plan_allowlist,
fetched_model_plan_types: &self.options.fetched_model_plan_types,
slot_count: &slot_count,
},
&CandidateRequest {
model: &request.model,
exclude_account_ids: &request.exclude_account_ids,
required_account_id: request.required_account_id.as_deref(),
now: request.now,
},
)
}
fn best_available_tier(&self, candidates: &[Account]) -> Option<String> {
self.options.tier_priority.iter().find_map(|tier| {
candidates
.iter()
.find(|account| account.plan_type.as_deref() == Some(tier.as_str()))
.and_then(|account| account.plan_type.clone())
})
/// 按当前策略从候选集合中选择一个账号(委托 [`AccountScheduler::select`])。
fn select_candidate(&self, candidates: &[Account], now: DateTime<Utc>) -> Option<Account> {
let slot_count = |account_id: &str| self.slots.get(account_id).map_or(0, VecDeque::len);
self.scheduler
.select(self.options.rotation_strategy, candidates, &slot_count, now)
}
}
@@ -978,10 +918,29 @@ impl RuntimeAccountPoolService {
/// 释放账号在途槽位。
pub async fn release(&self, account_id: &str) {
let released = self.pool.lock().await.release(account_id);
self.release_slot(account_id, true).await;
}
/// 释放账号在途槽位,不累计请求用量。
pub async fn release_without_request_usage(&self, account_id: &str) {
self.release_slot(account_id, false).await;
}
async fn release_slot(&self, account_id: &str, record_request_usage: bool) {
let released = if record_request_usage {
self.pool.lock().await.release(account_id)
} else {
self.pool
.lock()
.await
.release_without_request_usage(account_id)
};
let Some(released) = released else {
return;
};
if !record_request_usage {
return;
}
if let Err(error) = self.store.record_request(account_id).await {
tracing::warn!(
account_id,
@@ -1459,29 +1418,3 @@ fn refresh_cloudflare_cooldown(account: &mut Account, now: DateTime<Utc>) {
account.cloudflare_cooldown_until = None;
}
}
fn compare_least_used(a: &Account, b: &Account) -> Ordering {
a.quota_limit_reached
.cmp(&b.quota_limit_reached)
.then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at))
.then_with(|| a.request_count.cmp(&b.request_count))
.then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref()))
}
fn compare_window_reset(a: Option<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> Ordering {
match (a, b) {
(Some(a), Some(b)) => a.cmp(&b),
(Some(_) | None, None) | (None, Some(_)) => Ordering::Equal,
}
}
fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering {
last_used_millis(a).cmp(&last_used_millis(b))
}
fn last_used_millis(value: Option<&str>) -> i64 {
value
.and_then(|value| parse_rfc3339_utc(value).ok())
.map(|datetime| datetime.timestamp_millis())
.unwrap_or(0)
}
@@ -354,7 +354,7 @@ set
quota_fetched_at = ?,
plan_type = coalesce(?, plan_type),
status = case
when status in ('disabled', 'banned', 'expired', 'refreshing') then status
when status in ('disabled', 'banned', 'expired') then status
when ? = 1 then 'quota_exhausted'
when status = 'quota_exhausted' then 'active'
else status
@@ -423,7 +423,7 @@ pub(super) const SYNC_RUNTIME_ACCOUNT_STATE_SQL: &str = r"
update accounts
set
status = case
when status in ('disabled', 'banned', 'refreshing') then status
when status in ('disabled', 'banned') then status
else ?
end,
quota_limit_reached = case
@@ -255,7 +255,6 @@ pub(super) fn status_from_db(value: &str) -> SqliteAccountStoreResult<AccountSta
"active" => Ok(AccountStatus::Active),
"expired" => Ok(AccountStatus::Expired),
"quota_exhausted" => Ok(AccountStatus::QuotaExhausted),
"refreshing" => Ok(AccountStatus::Refreshing),
"disabled" => Ok(AccountStatus::Disabled),
"banned" => Ok(AccountStatus::Banned),
other => Err(SqliteAccountStoreError::InvalidStatus(other.to_string())),
@@ -1,7 +1,7 @@
//! 刷新租约存储、JWT 解码与账号 claims 验证。
use std::{
collections::hash_map::DefaultHasher,
collections::{hash_map::DefaultHasher, HashSet},
hash::{Hash, Hasher},
sync::{Arc, RwLock},
};
@@ -369,22 +369,53 @@ where
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> bool {
if account.status != AccountStatus::Active {
return false;
}
match trigger {
RefreshTrigger::Unauthorized => true,
RefreshTrigger::BeforeExpiry => account
.access_token_expires_at
.is_some_and(|expires_at| expires_at <= now + self.refresh_margin()),
RefreshTrigger::Unauthorized => can_refresh_after_auth_failure(account.status),
RefreshTrigger::BeforeExpiry => {
can_refresh_before_expiry(account.status)
&& self.should_refresh_before_expiry(account, now)
}
}
}
fn refresh_margin(&self) -> Duration {
fn should_refresh_before_expiry(&self, account: &Account, now: DateTime<Utc>) -> bool {
if let Some(next_refresh_at) = account.next_refresh_at {
return next_refresh_at <= now;
}
account
.access_token_expires_at
.is_some_and(|expires_at| self.refresh_at(account, expires_at) <= now)
}
fn refresh_at(&self, account: &Account, expires_at: DateTime<Utc>) -> DateTime<Utc> {
let policy = self.policy.snapshot();
let seconds = policy.refresh_margin_seconds.min(86_400 * 7) as i64;
Duration::seconds(seconds)
jittered_refresh_at(
account.id.as_str(),
expires_at,
policy.refresh_margin_seconds,
)
}
}
fn can_refresh_after_auth_failure(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Active | AccountStatus::Expired | AccountStatus::QuotaExhausted
)
}
fn can_refresh_before_expiry(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Active | AccountStatus::QuotaExhausted
)
}
fn status_after_successful_token_refresh(status: AccountStatus) -> AccountStatus {
match status {
AccountStatus::Expired => AccountStatus::Active,
status => status,
}
}
@@ -398,7 +429,7 @@ pub fn apply_token_pair(account: &Account, token_pair: TokenPair) -> Account {
refreshed.refresh_token = Some(refresh_token);
}
refreshed.status = AccountStatus::Active;
refreshed.status = status_after_successful_token_refresh(account.status);
refreshed
}
@@ -418,7 +449,7 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco
// 刷新租约存储(SQLite)
// ---------------------------------------------------------------------------
use sqlx::SqlitePool;
use sqlx::{QueryBuilder, Sqlite, SqlitePool};
use thiserror::Error;
/// SQLite 刷新租约存储结果。
@@ -485,6 +516,35 @@ where account_refresh_leases.expires_at <= ?
.await?;
Ok(result.rows_affected() > 0)
}
/// 返回给定账号集合中仍有效的刷新租约。
pub async fn active_account_ids(
&self,
account_ids: &[String],
now: DateTime<Utc>,
) -> RefreshLeaseStoreResult<HashSet<String>> {
if account_ids.is_empty() {
return Ok(HashSet::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(
"select account_id from account_refresh_leases where datetime(expires_at) > datetime(",
);
builder.push_bind(now.to_rfc3339());
builder.push(") and account_id in (");
let mut separated = builder.separated(", ");
for account_id in account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder
.build_query_as::<(String,)>()
.fetch_all(&self.pool)
.await?;
Ok(rows.into_iter().map(|(account_id,)| account_id).collect())
}
}
mod runtime;
@@ -127,6 +127,41 @@ where
self
}
/// 返回给定账号集合中当前正在执行 token 刷新的账号 ID。
pub async fn refreshing_account_ids(
&self,
account_ids: &[String],
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<HashSet<String>> {
if account_ids.is_empty() {
return Ok(HashSet::new());
}
let requested_ids = account_ids
.iter()
.map(String::as_str)
.collect::<HashSet<_>>();
let mut refreshing = self
.in_flight
.lock()
.await
.iter()
.filter(|account_id| requested_ids.contains(account_id.as_str()))
.cloned()
.collect::<HashSet<_>>();
if let Some(refresh_leases) = self.refresh_leases.as_ref() {
refreshing.extend(
refresh_leases
.active_account_ids(account_ids, now)
.await
.map_err(TokenRefreshServiceError::Lease)?,
);
}
Ok(refreshing)
}
/// 执行一次账号级刷新定时器调度。
pub async fn schedule_account_timers_once(
&self,
@@ -165,7 +200,6 @@ where
}
let timer = match account.status {
AccountStatus::Refreshing => Some((Duration::ZERO, RefreshTrigger::Unauthorized)),
AccountStatus::Expired => {
let delay = match account.next_refresh_at {
Some(next_refresh_at) if next_refresh_at > now => {
@@ -543,7 +577,7 @@ where
account: &Account,
now: DateTime<Utc>,
) -> Option<RefreshTrigger> {
if account.status == AccountStatus::Refreshing {
if account.status == AccountStatus::Expired {
return Some(RefreshTrigger::Unauthorized);
}
@@ -562,7 +596,7 @@ where
let mut permanent_failures = 0;
for attempt_index in 0..max_attempts {
let Some(attempt) = self.prepare_refresh_attempt(account, now).await? else {
let Some(attempt) = self.prepare_refresh_attempt(account, trigger, now).await? else {
return Ok(TokenRefreshOutcome::Skipped);
};
@@ -577,32 +611,34 @@ where
.refresh_account_at(&attempt, trigger, now)
.await
{
Ok(mut updated) if account_requires_token_write(&attempt, &updated) => {
updated.access_token_expires_at =
jwt_expiration(&updated.access_token).or(updated.access_token_expires_at);
updated.next_refresh_at = self.computed_refresh_at(&updated);
persist_token_update(&self.store, &updated).await?;
return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated)));
}
Ok(updated) if attempt.status != updated.status => {
if is_permanent_refresh_status(updated.status) {
permanent_failures += 1;
if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD
&& self.sleep_before_retry(&attempt.id, attempt_index).await
{
continue;
}
Ok(updated) if is_permanent_refresh_failure_status(updated.status) => {
permanent_failures += 1;
if permanent_failures < PERMANENT_FAILURE_CONFIRMATION_THRESHOLD
&& self.sleep_before_retry(&attempt.id, attempt_index).await
{
continue;
}
persist_status(&self.store, &attempt.id, updated.status).await?;
persist_next_refresh_at(&self.store, &attempt.id, None).await?;
return Ok(TokenRefreshOutcome::StatusUpdated);
}
Ok(_) => return Ok(TokenRefreshOutcome::Skipped),
Ok(mut updated) => {
updated.access_token_expires_at =
jwt_expiration(&updated.access_token).or(updated.access_token_expires_at);
updated.next_refresh_at = self.computed_refresh_at(&updated);
persist_token_update(&self.store, &updated).await?;
return Ok(TokenRefreshOutcome::Refreshed(Box::new(updated)));
}
Err(RefreshError::RetryableTransport | RefreshError::Transport)
if self.sleep_before_retry(&attempt.id, attempt_index).await => {}
Err(RefreshError::RetryableTransport | RefreshError::Transport) => {
persist_status(&self.store, &attempt.id, AccountStatus::Active).await?;
persist_status(
&self.store,
&attempt.id,
status_after_temporary_refresh_failure(attempt.status),
)
.await?;
let recovery_delay = stable_jittered_duration(
&attempt.id,
Duration::from_secs(RECOVERY_DELAY_SECONDS as u64),
@@ -624,6 +660,7 @@ where
async fn prepare_refresh_attempt(
&self,
account: &Account,
trigger: RefreshTrigger,
now: DateTime<Utc>,
) -> TokenRefreshServiceResult<Option<Account>> {
let Some(stored) = self.store.get(&account.id).await? else {
@@ -631,7 +668,7 @@ where
account.id.clone(),
));
};
let mut attempt = stored_account_to_refresh_account(stored);
let attempt = stored_account_to_refresh_account(stored);
if matches!(
attempt.status,
AccountStatus::Disabled | AccountStatus::Banned
@@ -641,11 +678,13 @@ where
return Ok(None);
}
if attempt.status != AccountStatus::Refreshing {
persist_status(&self.store, &attempt.id, AccountStatus::Refreshing).await?;
if !self
.scheduler
.should_refresh_account_at(&attempt, trigger, now)
{
return Ok(None);
}
attempt.status = AccountStatus::Active;
Ok(Some(attempt))
}
@@ -839,20 +878,16 @@ async fn persist_next_refresh_at(
}
}
fn account_requires_token_write(previous: &Account, updated: &Account) -> bool {
previous.access_token != updated.access_token
|| previous.refresh_token != updated.refresh_token
|| (updated.status == AccountStatus::Active && previous.status != updated.status)
fn is_permanent_refresh_failure_status(status: AccountStatus) -> bool {
matches!(status, AccountStatus::Expired | AccountStatus::Banned)
}
fn is_permanent_refresh_status(status: AccountStatus) -> bool {
matches!(
status,
AccountStatus::Expired
| AccountStatus::QuotaExhausted
| AccountStatus::Disabled
| AccountStatus::Banned
)
fn status_after_temporary_refresh_failure(status: AccountStatus) -> AccountStatus {
match status {
AccountStatus::Expired | AccountStatus::QuotaExhausted => status,
AccountStatus::Active => AccountStatus::Active,
AccountStatus::Disabled | AccountStatus::Banned => status,
}
}
fn refresh_lease_owner() -> String {
+61 -19
View File
@@ -1,6 +1,10 @@
//! 指纹类型、更新与持久化。
use std::{path::PathBuf, sync::Arc, time::Duration};
use std::{
path::PathBuf,
sync::{Arc, RwLock},
time::Duration,
};
use chrono::{DateTime, Utc};
use indexmap::IndexMap;
@@ -89,6 +93,38 @@ impl Fingerprint {
}
}
/// 运行时共享指纹快照。
#[derive(Debug, Clone)]
pub struct RuntimeFingerprint {
current: Arc<RwLock<Arc<Fingerprint>>>,
}
impl RuntimeFingerprint {
pub fn new(fingerprint: Fingerprint) -> Self {
Self {
current: Arc::new(RwLock::new(Arc::new(fingerprint))),
}
}
pub fn current(&self) -> Arc<Fingerprint> {
self.current
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
pub fn snapshot(&self) -> Fingerprint {
self.current().as_ref().clone()
}
pub fn replace(&self, fingerprint: Fingerprint) {
*self
.current
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Arc::new(fingerprint);
}
}
// ---------------------------------------------------------------------------
// 指纹历史记录与仓储
// ---------------------------------------------------------------------------
@@ -345,7 +381,7 @@ fn json_error(error: serde_json::Error) -> sqlx::Error {
/// Appcast 指纹更新检查器。
#[derive(Clone)]
pub struct UpdateChecker {
repository: Option<FingerprintRepository>,
repository: FingerprintRepository,
client: reqwest::Client,
appcast_url: String,
extracted_fingerprint_path: PathBuf,
@@ -360,7 +396,7 @@ struct InternalState {
impl UpdateChecker {
/// 使用自定义 client 和路径构造检查器。
pub fn with_client(
repository: Option<FingerprintRepository>,
repository: FingerprintRepository,
client: reqwest::Client,
appcast_url: impl Into<String>,
extracted_fingerprint_path: PathBuf,
@@ -406,18 +442,18 @@ impl UpdateChecker {
}
/// 检查并在需要时应用更新。
pub async fn check_and_apply_update(&self) -> Result<bool, UpdateError> {
pub async fn check_and_apply_update(&self) -> Result<Option<Fingerprint>, UpdateError> {
let update_state = self.check_for_update().await?;
if !update_state.update_available {
return Ok(false);
return Ok(None);
}
let (Some(version), Some(build)) = (
update_state.latest_version.as_deref(),
update_state.latest_build.as_deref(),
) else {
return Ok(false);
return Ok(None);
};
let state = self.state.lock().await;
@@ -433,13 +469,13 @@ impl UpdateChecker {
"发现新的 fingerprint 版本"
);
self.apply_version_update(version, build).await?;
let updated_fingerprint = self.apply_version_update(version, build).await?;
let mut state = self.state.lock().await;
state.current_version = version.to_string();
state.current_build = build.to_string();
Ok(true)
Ok(Some(updated_fingerprint))
}
async fn fetch_appcast(&self) -> Result<String, UpdateError> {
@@ -457,20 +493,26 @@ impl UpdateChecker {
response.text().await.map_err(UpdateError::from)
}
async fn apply_version_update(&self, version: &str, build: &str) -> Result<(), UpdateError> {
async fn apply_version_update(
&self,
version: &str,
build: &str,
) -> Result<Fingerprint, UpdateError> {
let chromium_version =
load_matching_chromium_version(&self.extracted_fingerprint_path, version, build);
if let Some(repository) = &self.repository {
repository
.update_current_version(version, build, chromium_version.as_deref())
.await?;
repository
.insert_update_history(version, build, chromium_version.as_deref(), None)
.await?;
}
Ok(())
self.repository
.update_current_version(version, build, chromium_version.as_deref())
.await?;
self.repository
.insert_update_history(version, build, chromium_version.as_deref(), None)
.await?;
let updated = self
.repository
.load_current()
.await?
.ok_or(sqlx::Error::RowNotFound)?;
Ok(updated)
}
}
+1
View File
@@ -4,5 +4,6 @@ pub mod accounts;
pub mod fingerprint;
pub mod models;
pub mod protocol;
pub mod scheduler;
pub mod token_client;
pub mod transport;
@@ -0,0 +1,137 @@
//! 候选账号过滤。
//!
//! 所有调度策略共享的前置:从账号池中筛出「本次请求可用」的候选集合。这是选择流程
//! 的第一步,与具体策略无关——策略只在过滤后的候选之间做选择。
//!
//! 过滤规则(全部满足才入选):
//!
//! 1. 账号状态为 [`AccountStatus::Active`];
//! 2. 在途槽位数未达单账号并发上限;
//! 3. 配额可用([`AccountService::quota_available_at`],受 `skip_quota_limited` 控制);
//! 4. 模型在账号订阅计划的允许列表内;
//! 5. 不在本次请求的排除列表内;
//! 6. 不处于 Cloudflare 冷却期。
//!
//! 之后再按订阅层级优先级(`tier_priority`)收窄到最高可用层级;若请求指定了
//! `required_account_id`,则只保留该账号(跳过层级收窄)。
use chrono::{DateTime, Utc};
use crate::upstream::accounts::model::{Account, AccountStatus};
use crate::upstream::accounts::service::AccountService;
/// 候选过滤所需的账号池策略参数(只读视图)。
pub struct CandidateFilter<'a> {
/// 单账号最大并发请求数。
pub max_concurrent_per_account: usize,
/// 是否跳过已标记配额耗尽的账号。
pub skip_quota_limited: bool,
/// 订阅层级优先级,越靠前优先级越高。
pub tier_priority: &'a [String],
/// 模型到允许订阅计划的映射。
pub model_plan_allowlist: &'a std::collections::BTreeMap<String, Vec<String>>,
/// 已成功拉取过模型列表的订阅计划。
pub fetched_model_plan_types: &'a std::collections::BTreeSet<String>,
/// 按账号 ID 读取当前在途槽位数。
pub slot_count: &'a dyn Fn(&str) -> usize,
}
/// 本次账号获取请求的过滤条件。
pub struct CandidateRequest<'a> {
/// 本次请求使用的模型。
pub model: &'a str,
/// 本次请求需要排除的账号 ID。
pub exclude_account_ids: &'a [String],
/// 本次请求必须使用的账号 ID(存在时跳过层级收窄,只保留该账号)。
pub required_account_id: Option<&'a str>,
/// 本次调度使用的当前时间。
pub now: DateTime<Utc>,
}
/// 从账号池筛出本次请求的可用候选集合。
pub fn filter<'a>(
accounts: impl Iterator<Item = &'a Account>,
filter: &CandidateFilter<'_>,
request: &CandidateRequest<'_>,
) -> Vec<Account> {
let mut candidates = accounts
.filter(|account| is_base_available(account, filter, request))
.cloned()
.collect::<Vec<_>>();
if let Some(required_account_id) = request.required_account_id {
candidates.retain(|account| account.id == required_account_id);
return candidates;
}
if let Some(best_tier) = best_available_tier(&candidates, filter.tier_priority) {
candidates.retain(|account| account.plan_type.as_deref() == Some(best_tier.as_str()));
}
candidates
}
/// 判断账号是否满足全部基础可用条件。
fn is_base_available(
account: &Account,
filter: &CandidateFilter<'_>,
request: &CandidateRequest<'_>,
) -> bool {
account.status == AccountStatus::Active
&& (filter.slot_count)(&account.id) < filter.max_concurrent_per_account
&& AccountService::quota_available_at(account, request.now, filter.skip_quota_limited)
&& is_model_allowed(
account,
request.model,
filter.model_plan_allowlist,
filter.fetched_model_plan_types,
)
&& !request
.exclude_account_ids
.iter()
.any(|account_id| account_id == &account.id)
&& AccountService::cloudflare_available_at(account, request.now)
}
/// 判断模型是否允许该账号的订阅计划使用。
fn is_model_allowed(
account: &Account,
model: &str,
model_plan_allowlist: &std::collections::BTreeMap<String, Vec<String>>,
fetched_model_plan_types: &std::collections::BTreeSet<String>,
) -> bool {
let Some(allowed_plans) = model_plan_allowlist.get(model) else {
return true;
};
let Some(plan_type) = account.plan_type.as_deref() else {
return false;
};
allowed_plans.iter().any(|plan| plan == plan_type)
|| !fetched_model_plan_types.contains(plan_type)
}
/// 在候选集合中找出订阅层级优先级最高的可用层级。
fn best_available_tier(candidates: &[Account], tier_priority: &[String]) -> Option<String> {
for tier in tier_priority {
let matched = candidates
.iter()
.any(|account| account.plan_type.as_deref() == Some(tier.as_str()));
if matched {
return Some(tier.clone());
}
}
None
}
/// 判断账号是否可用于模型列表刷新(不涉及排除列表 / 模型允许列表 / 层级收窄)。
pub fn is_model_refresh_available(
account: &Account,
max_concurrent_per_account: usize,
skip_quota_limited: bool,
slot_count: &dyn Fn(&str) -> usize,
now: DateTime<Utc>,
) -> bool {
account.status == AccountStatus::Active
&& slot_count(&account.id) < max_concurrent_per_account
&& AccountService::quota_available_at(account, now, skip_quota_limited)
&& AccountService::cloudflare_available_at(account, now)
}
+171
View File
@@ -0,0 +1,171 @@
//! 账号运行时反馈:错误率与 TTFT 的 EWMA 统计。
//!
//! 每个账号维护两个指数加权移动平均(EWMA):错误率与首 token 延迟(TTFT,
//! 毫秒)。请求结束后经 [`FeedbackStats::report`] 回灌,直接进 [`super::score`]
//! 打分并可驱动 sticky escape。统计以 `AtomicU64`(存 `f64` 位模式)+ CAS 无锁更新,
//! 对齐 sub2api 的设计,避免选择热路径上的锁与 DB 往返。
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::RwLock;
/// EWMA 平滑系数(新样本权重)。与 sub2api 一致取 0.2。
const EWMA_ALPHA: f64 = 0.2;
/// 单账号的运行时 EWMA 统计。
///
/// `*_bits` 存 `f64` 的位模式;`NaN`(全 1 的哨兵)表示"尚无样本"。
#[derive(Debug)]
struct AccountFeedback {
/// 错误率 EWMA,范围 `[0, 1]`。
error_rate_bits: AtomicU64,
/// TTFT EWMA,单位毫秒。
ttft_ms_bits: AtomicU64,
}
/// 表示"尚无样本"的哨兵位模式。
const EMPTY_BITS: u64 = f64::NAN.to_bits();
impl Default for AccountFeedback {
fn default() -> Self {
Self {
error_rate_bits: AtomicU64::new(EMPTY_BITS),
ttft_ms_bits: AtomicU64::new(EMPTY_BITS),
}
}
}
impl AccountFeedback {
/// 读取当前错误率 EWMA;无样本时返回 `None`。
fn error_rate(&self) -> Option<f64> {
load_sample(&self.error_rate_bits)
}
/// 读取当前 TTFT EWMA(毫秒);无样本时返回 `None`。
fn ttft_ms(&self) -> Option<f64> {
load_sample(&self.ttft_ms_bits)
}
/// 回灌一次请求结果。
fn report(&self, success: bool, first_token_ms: Option<u64>) {
let error_sample = if success { 0.0 } else { 1.0 };
update_ewma(&self.error_rate_bits, error_sample);
if let Some(ms) = first_token_ms {
update_ewma(&self.ttft_ms_bits, ms as f64);
}
}
}
/// 读取 EWMA 样本;哨兵(NaN)返回 `None`。
fn load_sample(bits: &AtomicU64) -> Option<f64> {
let value = f64::from_bits(bits.load(Ordering::Relaxed));
if value.is_nan() {
None
} else {
Some(value)
}
}
/// 无锁 CAS 更新 EWMA:首个样本直接落值,其后按 alpha 混合。
fn update_ewma(bits: &AtomicU64, sample: f64) {
let mut current = bits.load(Ordering::Relaxed);
loop {
let previous = f64::from_bits(current);
let next = if previous.is_nan() {
sample
} else {
EWMA_ALPHA * sample + (1.0 - EWMA_ALPHA) * previous
};
match bits.compare_exchange_weak(
current,
next.to_bits(),
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => return,
Err(observed) => current = observed,
}
}
}
/// 账号 EWMA 反馈的读数快照,供打分使用。
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct FeedbackSample {
/// 错误率 EWMA,范围 `[0, 1]`;无样本时为 `None`。
pub error_rate: Option<f64>,
/// TTFT EWMA(毫秒);无样本时为 `None`。
pub ttft_ms: Option<f64>,
}
/// 账号运行时反馈存储:按账号 ID 聚合 EWMA 统计。
///
/// 选择热路径只读(`sample`),请求结束后写(`report`)。内部按账号分片为独立的
/// 原子对象,`RwLock` 仅保护 map 结构本身(插入新账号时短暂写锁)。
#[derive(Debug, Default)]
pub struct FeedbackStats {
accounts: RwLock<HashMap<String, AccountFeedback>>,
}
impl FeedbackStats {
/// 创建空的反馈存储。
pub fn new() -> Self {
Self::default()
}
/// 读取指定账号的 EWMA 快照;账号无记录时返回全 `None`。
pub fn sample(&self, account_id: &str) -> FeedbackSample {
let accounts = self
.accounts
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
match accounts.get(account_id) {
Some(feedback) => FeedbackSample {
error_rate: feedback.error_rate(),
ttft_ms: feedback.ttft_ms(),
},
None => FeedbackSample::default(),
}
}
/// 回灌一次请求结果:成功/失败 + 首 token 延迟(毫秒)。
pub fn report(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
// 快路径:账号已存在时只需读锁 + 原子更新。
{
let accounts = self
.accounts
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if let Some(feedback) = accounts.get(account_id) {
feedback.report(success, first_token_ms);
return;
}
}
// 慢路径:首次见到该账号,取写锁插入后回灌。
let mut accounts = self
.accounts
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
accounts
.entry(account_id.to_string())
.or_default()
.report(success, first_token_ms);
}
/// 移除指定账号的反馈记录(账号被删除时调用)。
pub fn remove(&self, account_id: &str) {
let mut accounts = self
.accounts
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
accounts.remove(account_id);
}
/// 清空全部反馈记录(账号池清空时调用)。
pub fn clear(&self) {
let mut accounts = self
.accounts
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
accounts.clear();
}
}
+112
View File
@@ -0,0 +1,112 @@
//! 智能账号调度模块。
//!
//! 账号选择的唯一归属地。把「选择」从账号池存储中剥离,集中于此:
//!
//! - [`candidates`] —— 候选过滤(可用性 / 层级 / 模型允许列表),所有策略共享的前置。
//! - [`strategy`] —— 全部选择策略,每个策略一个文件([`strategy::smart`] /
//! [`strategy::quota_reset`] / [`strategy::round_robin`] / [`strategy::sticky`])。
//! - [`feedback`] —— 跨策略的运行时 EWMA 反馈(错误率 / TTFT),喂给 Smart 打分。
//!
//! [`AccountScheduler`] 是对外门面:持有轮转游标、打分权重与 EWMA 反馈存储,接收
//! **已过滤的候选切片**并返回选中账号。它不 lock 账号池、不刷状态、不写库——存储与
//! 状态维护仍归 `AccountPool`,二者经候选切片与 `slot_count` 闭包解耦。
pub mod candidates;
pub mod feedback;
pub mod strategy;
use std::sync::Mutex;
use chrono::{DateTime, Utc};
use crate::upstream::accounts::model::Account;
pub use candidates::{CandidateFilter, CandidateRequest};
pub use feedback::{FeedbackSample, FeedbackStats};
pub use strategy::smart::rank_candidates;
pub use strategy::{RotationStrategy, ScoreBreakdown, ScoreWeights};
/// 账号调度器:账号选择的对外门面。
///
/// 持有轮转游标(并列破并列用)、打分权重(Smart 用)与 EWMA 反馈存储。选择热路径
/// 只读反馈与权重,请求结束后经 [`AccountScheduler::report_feedback`] 回灌。游标用
/// [`Mutex`] 保护(选择需可变,但 `AccountScheduler` 以共享引用参与并发请求)。
#[derive(Debug)]
pub struct AccountScheduler {
/// 轮转游标:在并列最优账号之间做公平轮转。
cursor: Mutex<usize>,
/// Smart 策略的打分权重。
weights: ScoreWeights,
/// 账号运行时 EWMA 反馈(错误率 / TTFT)。
feedback: FeedbackStats,
}
impl Default for AccountScheduler {
fn default() -> Self {
Self::new(ScoreWeights::default())
}
}
impl AccountScheduler {
/// 使用指定打分权重创建调度器。
pub fn new(weights: ScoreWeights) -> Self {
Self {
cursor: Mutex::new(0),
weights,
feedback: FeedbackStats::new(),
}
}
/// 替换打分权重(策略切换或热更时调用)。
pub fn set_weights(&mut self, weights: ScoreWeights) {
self.weights = weights;
}
/// 重置轮转游标(策略切换时调用,保证轮转从头开始)。
pub fn reset_cursor(&self) {
*self.lock_cursor() = 0;
}
/// 从已过滤的候选集合中按指定策略选择一个账号。
///
/// `candidates` 必须是 [`candidates::filter`] 的输出;`slot_count` 按账号 ID 读取
/// 当前在途槽位数(Smart 打分与候选无关的运行时状态由调用方提供)。
pub fn select(
&self,
strategy: RotationStrategy,
candidates: &[Account],
slot_count: &dyn Fn(&str) -> usize,
now: DateTime<Utc>,
) -> Option<Account> {
let input = strategy::SelectionInput {
candidates,
slot_count,
feedback: &self.feedback,
weights: &self.weights,
now,
};
let mut cursor = self.lock_cursor();
strategy::select(strategy, &input, &mut cursor)
}
/// 回灌一次请求结果到运行时 EWMA 反馈(错误率 / TTFT),供 Smart 打分。
pub fn report_feedback(&self, account_id: &str, success: bool, first_token_ms: Option<u64>) {
self.feedback.report(account_id, success, first_token_ms);
}
/// 移除指定账号的运行时反馈(账号被删除时调用)。
pub fn forget_feedback(&self, account_id: &str) {
self.feedback.remove(account_id);
}
/// 清空全部运行时反馈(账号池清空时调用)。
pub fn clear_feedback(&self) {
self.feedback.clear();
}
fn lock_cursor(&self) -> std::sync::MutexGuard<'_, usize> {
self.cursor
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
@@ -0,0 +1,142 @@
//! 账号调度策略层。
//!
//! 所有账号选择策略集中于此,每个策略一个文件,互不依赖:
//!
//! - [`smart`] —— 加权线性打分 + EWMA 反馈(“智能调度”)。
//! - [`quota_reset`] —— 优先消耗即将重置的额度窗口。
//! - [`round_robin`] —— 按候选顺序循环。
//! - [`sticky`] —— 优先最近使用过的账号。
//!
//! 本模块只承担三件事:定义 [`RotationStrategy`] 枚举、把选择请求分派到对应策略、
//! 提供各策略共享的 tie-break 比较原语([`select_by`] / [`compare_window_reset`] /
//! [`compare_last_used`])。策略本身不触碰账号池存储:在途槽位数经 `slot_count`
//! 闭包按需读取,EWMA 反馈与打分权重经 [`SelectionInput`] 传入。
use std::cmp::Ordering;
use chrono::{DateTime, Utc};
use crate::infra::time::parse_rfc3339_utc;
use crate::upstream::accounts::model::Account;
use crate::upstream::scheduler::feedback::FeedbackStats;
pub mod quota_reset;
pub mod round_robin;
pub mod smart;
pub mod sticky;
pub use smart::{ScoreBreakdown, ScoreWeights};
/// 账号轮转策略。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RotationStrategy {
/// 加权线性打分调度:多因子(负载/窗口/错误率/TTFT/重置)归一化后加权求和,
/// 配合 EWMA 反馈。
Smart,
/// 优先消耗即将重置的额度窗口。
QuotaResetPriority,
/// 按候选账号顺序循环选择。
RoundRobin,
/// 优先选择最近使用过的账号。
Sticky,
}
impl RotationStrategy {
/// 返回策略的持久化/日志字符串。
pub fn as_str(self) -> &'static str {
match self {
Self::Smart => "smart",
Self::QuotaResetPriority => "quota_reset_priority",
Self::RoundRobin => "round_robin",
Self::Sticky => "sticky",
}
}
}
/// 策略选择的运行时输入。
///
/// `slot_count` 按账号 ID 读取当前在途槽位数;`feedback` / `weights` 仅 [`smart`]
/// 使用,其余策略忽略;`now` 用于窗口重置相关的打分。
pub struct SelectionInput<'a> {
/// 已过滤的候选账号集合。
pub candidates: &'a [Account],
/// 按账号 ID 读取当前在途槽位数。
pub slot_count: &'a dyn Fn(&str) -> usize,
/// 账号 EWMA 反馈存储(Smart 打分用)。
pub feedback: &'a FeedbackStats,
/// 打分权重(Smart 打分用)。
pub weights: &'a ScoreWeights,
/// 当前时刻(窗口重置因子用)。
pub now: DateTime<Utc>,
}
/// 按指定策略从候选集合中选择一个账号。
///
/// `cursor` 是账号池持有的轮转游标,用于在并列最优账号之间做公平轮转;不同策略对它
/// 的读写方式不同(详见各策略模块),本函数只负责透传可变引用。
pub fn select(
strategy: RotationStrategy,
input: &SelectionInput<'_>,
cursor: &mut usize,
) -> Option<Account> {
match strategy {
RotationStrategy::Smart => smart::select(input, cursor),
RotationStrategy::QuotaResetPriority => quota_reset::select(input, cursor),
RotationStrategy::RoundRobin => round_robin::select(input, cursor),
RotationStrategy::Sticky => sticky::select(input, cursor),
}
}
// ====================================================================
// 共享 tie-break 原语
// ====================================================================
/// 按比较函数升序排序后,在并列最优(`Ordering::Equal`)的账号之间用轮转游标破并列。
///
/// 这是 [`smart`] 与 [`quota_reset`] 共享的选择骨架:先按各自的比较规则排序,再在得分
/// 相同的账号之间轮转,保证并列账号被公平地依次选中。
pub(crate) fn select_by(
candidates: &[Account],
cursor: &mut usize,
compare: impl Fn(&Account, &Account) -> Ordering,
) -> Option<Account> {
let mut sorted = candidates.iter().collect::<Vec<_>>();
sorted.sort_by(|a, b| compare(a, b));
let best = *sorted.first()?;
let tied_count = sorted
.iter()
.take_while(|account| compare(best, account) == Ordering::Equal)
.count();
let index = *cursor % tied_count;
*cursor = cursor.wrapping_add(1);
Some((*sorted[index]).clone())
}
/// 计算账号当前窗口的总 token 数(输入 + 输出 + 图片输入 + 图片输出)。
pub(crate) fn account_window_token_count(account: &Account) -> u64 {
account
.window_input_tokens
.saturating_add(account.window_output_tokens)
.saturating_add(account.window_image_input_tokens)
.saturating_add(account.window_image_output_tokens)
}
/// 比较两个窗口重置时间:越早重置越优先;缺失重置时间视为并列(不惩罚)。
pub(crate) fn compare_window_reset(a: Option<DateTime<Utc>>, b: Option<DateTime<Utc>>) -> Ordering {
match (a, b) {
(Some(a), Some(b)) => a.cmp(&b),
(Some(_) | None, None) | (None, Some(_)) => Ordering::Equal,
}
}
/// 比较两个账号的最近使用时间:越久未用越优先(升序,最小者排前)。
pub(crate) fn compare_last_used(a: Option<&str>, b: Option<&str>) -> Ordering {
last_used_millis(a).cmp(&last_used_millis(b))
}
fn last_used_millis(value: Option<&str>) -> i64 {
value
.and_then(|value| parse_rfc3339_utc(value).ok())
.map(|datetime| datetime.timestamp_millis())
.unwrap_or(0)
}
@@ -0,0 +1,29 @@
//! QuotaResetPriority 策略:优先消耗即将重置的额度窗口。
//!
//! 排序规则(升序,越靠前越优先):
//! 1. 未封顶配额优先于已封顶(`quota_limit_reached` false < true);
//! 2. 窗口重置时间越早越优先;
//! 3. 历史请求数越少越优先;
//! 4. 最近使用时间越久越优先(LRU)。
//!
//! 并列最优的账号之间由 [`super::select_by`] 用轮转游标破并列。
use std::cmp::Ordering;
use crate::upstream::accounts::model::Account;
use super::{compare_last_used, compare_window_reset, select_by, SelectionInput};
/// QuotaResetPriority 策略选择。
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
select_by(input.candidates, cursor, compare_quota_reset_priority)
}
/// 配额重置优先级比较:升序,最优(最应先用)排在最前。
fn compare_quota_reset_priority(a: &Account, b: &Account) -> Ordering {
a.quota_limit_reached
.cmp(&b.quota_limit_reached)
.then_with(|| compare_window_reset(a.window_reset_at, b.window_reset_at))
.then_with(|| a.request_count.cmp(&b.request_count))
.then_with(|| compare_last_used(a.last_used_at.as_deref(), b.last_used_at.as_deref()))
}
@@ -0,0 +1,20 @@
//! RoundRobin 策略:按候选顺序循环选择。
//!
//! 最简单的策略:忽略账号负载与用量,纯粹按候选集合的顺序用轮转游标依次选中,
//! 保证请求在候选账号间均匀轮转。
use crate::upstream::accounts::model::Account;
use super::SelectionInput;
/// 按候选顺序做纯轮转选择。
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
let candidates = input.candidates;
if candidates.is_empty() {
return None;
}
*cursor %= candidates.len();
let index = *cursor;
*cursor = cursor.wrapping_add(1);
Some(candidates[index].clone())
}
@@ -0,0 +1,291 @@
//! Smart 策略:加权线性打分 + EWMA 反馈。
//!
//! 参考 sub2api 的调度设计:把多个归一化因子(负载、窗口 token 用量、窗口缓存、窗口
//! 请求数、错误率、TTFT、窗口重置)按可配置权重线性求和,得分越高越优先。相较字典序
//! 元组比较,加权打分是更平滑的负载均衡,不会让单一维度无条件压制其余维度,同时通过
//! EWMA 反馈感知账号健康度与延迟。
//!
//! 各归一化因子都落在 `[0, 1]` 且方向统一(越大越好),因此权重之间可直接比较。配额
//! 已封顶(`quota_limit_reached`)的账号会被施加一个强惩罚,确保只要还有未封顶的候选就
//! 不会选中它。并列最高分的账号之间由 [`super::select_by`] 用轮转游标破并列。
use crate::upstream::accounts::model::Account;
use crate::upstream::scheduler::feedback::{FeedbackSample, FeedbackStats};
use super::{account_window_token_count, SelectionInput};
/// 配额封顶账号的惩罚分。远大于各因子加权和的理论上限(各因子 ∈ [0,1],权重和通常为
/// 个位数),保证未封顶账号总是优先于封顶账号。
const QUOTA_LIMITED_PENALTY: f64 = 1_000.0;
/// 打分权重。每个字段对应一个归一化因子的权重,权重为 0 即关闭该维度。
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ScoreWeights {
/// 在途槽位压力(越空闲越优先)。
pub load: f64,
/// 当前窗口实际 token 用量(用得越少越优先)。
pub window: f64,
/// 当前窗口缓存 token 用量(缓存越少越优先,权重小,仅在实际用量相当时区分)。
pub cached: f64,
/// 当前窗口请求数(请求越少越优先)。
pub window_requests: f64,
/// 错误率 EWMA(越健康越优先)。
pub error_rate: f64,
/// 首 token 延迟 EWMA(越快越优先)。
pub ttft: f64,
/// 窗口重置临近度(越快重置越优先,use-it-or-lose-it,默认关闭)。
pub reset: f64,
}
impl Default for ScoreWeights {
fn default() -> Self {
// 默认权重对齐 sub2api 的取向:负载与健康度为主,窗口用量次之,TTFT 再次,
// 窗口重置默认关闭(避免过早榨干即将重置的账号)。cached 权重远小于 window,
// 确保实际 token 用量主导、缓存仅作细粒度区分。
Self {
load: 1.0,
window: 0.8,
cached: 0.05,
window_requests: 0.6,
error_rate: 1.0,
ttft: 0.5,
reset: 0.0,
}
}
}
/// 单个候选账号的打分明细,供日志与可观测使用。
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct ScoreBreakdown {
/// 各因子加权后的总分(已含配额封顶惩罚)。
pub total: f64,
/// 负载因子(归一化后,未乘权重)。
pub load: f64,
/// 窗口实际 token 用量因子。
pub window: f64,
/// 窗口缓存 token 用量因子。
pub cached: f64,
/// 窗口请求数因子。
pub window_requests: f64,
/// 错误率因子。
pub error_rate: f64,
/// TTFT 因子。
pub ttft: f64,
/// 窗口重置因子。
pub reset: f64,
/// 是否被施加配额封顶惩罚。
pub quota_limited: bool,
}
/// Smart 策略选择:按加权得分降序排序,在并列最高分之间用轮转游标破并列。
pub fn select(input: &SelectionInput<'_>, cursor: &mut usize) -> Option<Account> {
let ranked = rank_candidates(
input.candidates,
input.weights,
input.now.timestamp(),
input.slot_count,
input.feedback,
);
let best_total = ranked.first()?.1.total;
// 并列最高分的账号数量(浮点相等用位相等判定,得分由相同输入算出,稳定)。
let tied_count = ranked
.iter()
.take_while(|(_, breakdown)| breakdown.total.to_bits() == best_total.to_bits())
.count();
let index = *cursor % tied_count;
*cursor = cursor.wrapping_add(1);
let account_index = ranked[index].0;
Some(input.candidates[account_index].clone())
}
/// 对候选集合按加权得分排序(降序),返回 `(账号索引, 打分明细)` 列表。
///
/// 只做打分与排序,最终的并列破并列交给 [`select`]。公开供测试直接断言排序结果。
pub fn rank_candidates(
candidates: &[Account],
weights: &ScoreWeights,
now_secs: i64,
slot_count: &dyn Fn(&str) -> usize,
feedback: &FeedbackStats,
) -> Vec<(usize, ScoreBreakdown)> {
let samples = candidates
.iter()
.map(|account| feedback.sample(&account.id))
.collect::<Vec<_>>();
let ctx = NormalizationContext::build(candidates, &samples, now_secs, slot_count);
let mut ranked = candidates
.iter()
.zip(&samples)
.enumerate()
.map(|(index, (account, sample))| {
(
index,
score_account(account, sample, weights, &ctx, now_secs, slot_count),
)
})
.collect::<Vec<_>>();
// 降序:得分高在前。NaN 不应出现(因子均有界),保守用 total_cmp。
ranked.sort_by(|a, b| b.1.total.total_cmp(&a.1.total));
ranked
}
/// 候选集合的归一化上下文:预先算好各维度的 min/max,供逐账号归一化。
struct NormalizationContext {
max_slot_count: f64,
max_window_tokens: f64,
max_window_cached: f64,
max_window_requests: f64,
min_ttft_ms: f64,
max_ttft_ms: f64,
min_reset_secs: f64,
max_reset_secs: f64,
}
impl NormalizationContext {
fn build(
candidates: &[Account],
samples: &[FeedbackSample],
now_secs: i64,
slot_count: &dyn Fn(&str) -> usize,
) -> Self {
let mut max_slot_count = 0.0_f64;
let mut max_window_tokens = 0.0_f64;
let mut max_window_cached = 0.0_f64;
let mut max_window_requests = 0.0_f64;
let mut min_ttft_ms = f64::INFINITY;
let mut max_ttft_ms = 0.0_f64;
let mut min_reset_secs = f64::INFINITY;
let mut max_reset_secs = 0.0_f64;
for (account, sample) in candidates.iter().zip(samples) {
max_slot_count = max_slot_count.max(slot_count(&account.id) as f64);
max_window_tokens = max_window_tokens.max(account_window_token_count(account) as f64);
max_window_cached = max_window_cached.max(account.window_cached_tokens as f64);
max_window_requests = max_window_requests.max(account.window_request_count as f64);
if let Some(ttft) = sample.ttft_ms {
if ttft > 0.0 {
min_ttft_ms = min_ttft_ms.min(ttft);
max_ttft_ms = max_ttft_ms.max(ttft);
}
}
if let Some(reset_at) = account.window_reset_at {
let secs = (reset_at.timestamp() - now_secs).max(0) as f64;
min_reset_secs = min_reset_secs.min(secs);
max_reset_secs = max_reset_secs.max(secs);
}
}
Self {
max_slot_count,
max_window_tokens,
max_window_cached,
max_window_requests,
min_ttft_ms: if min_ttft_ms.is_finite() {
min_ttft_ms
} else {
0.0
},
max_ttft_ms,
min_reset_secs: if min_reset_secs.is_finite() {
min_reset_secs
} else {
0.0
},
max_reset_secs,
}
}
}
/// 把 `value` 从 `[min, max]` 线性映射到 `[0, 1]`;区间退化时返回中性值 `1.0`
/// (表示该维度对所有候选无区分度,不应影响排序)。
fn normalize_ascending_good(value: f64, min: f64, max: f64) -> f64 {
if max <= min {
return 1.0;
}
((value - min) / (max - min)).clamp(0.0, 1.0)
}
/// 计算指定账号在候选集合中的加权得分明细。
fn score_account(
account: &Account,
sample: &FeedbackSample,
weights: &ScoreWeights,
ctx: &NormalizationContext,
now_secs: i64,
slot_count: &dyn Fn(&str) -> usize,
) -> ScoreBreakdown {
// 负载:槽位越少越好。归一化到 [0,1] 后取反(1 - x),使“越空闲得分越高”。
let load =
1.0 - normalize_ascending_good(slot_count(&account.id) as f64, 0.0, ctx.max_slot_count);
// 窗口实际 token 用量:token 越少越好。
let window = 1.0
- normalize_ascending_good(
account_window_token_count(account) as f64,
0.0,
ctx.max_window_tokens,
);
// 窗口缓存 token 用量:缓存越少越好(权重小,仅在实际用量相当时细分)。
let cached = 1.0
- normalize_ascending_good(
account.window_cached_tokens as f64,
0.0,
ctx.max_window_cached,
);
// 窗口请求数:请求越少越好。
let window_requests = 1.0
- normalize_ascending_good(
account.window_request_count as f64,
0.0,
ctx.max_window_requests,
);
// 错误率 EWMA:本身即 [0,1],越低越健康。无样本时视为 0(满分健康)。
let error_rate = 1.0 - sample.error_rate.unwrap_or(0.0).clamp(0.0, 1.0);
// TTFT:延迟越低越好。无样本的账号视为中性满分,避免冷启动被压。
let ttft = match sample.ttft_ms {
Some(ms) if ms > 0.0 => {
1.0 - normalize_ascending_good(ms, ctx.min_ttft_ms, ctx.max_ttft_ms)
}
_ => 1.0,
};
// 窗口重置:越快重置越优先(use-it-or-lose-it)。无 reset 时间视为中性满分。
let reset = match account.window_reset_at {
Some(reset_at) => {
let secs = (reset_at.timestamp() - now_secs).max(0) as f64;
1.0 - normalize_ascending_good(secs, ctx.min_reset_secs, ctx.max_reset_secs)
}
None => 1.0,
};
let mut total = weights.load * load
+ weights.window * window
+ weights.cached * cached
+ weights.window_requests * window_requests
+ weights.error_rate * error_rate
+ weights.ttft * ttft
+ weights.reset * reset;
// 配额封顶硬惩罚:确保只要还有未封顶候选就不会选中封顶账号。
if account.quota_limit_reached {
total -= QUOTA_LIMITED_PENALTY;
}
ScoreBreakdown {
total,
load,
window,
cached,
window_requests,
error_rate,
ttft,
reset,
quota_limited: account.quota_limit_reached,
}
}
@@ -0,0 +1,28 @@
//! Sticky 策略:优先选择最近使用过的账号。
//!
//! 选中 `last_used_at` 最大(最近使用)的账号,倾向于持续复用同一账号(利于上游的
//! prompt 缓存命中)。该策略不使用轮转游标——最近使用时间本身即确定性顺序。
use std::cmp::Ordering;
use crate::upstream::accounts::model::Account;
use super::{compare_last_used, SelectionInput};
/// 选中最近使用过的账号(`last_used_at` 最大者)。
///
/// `_cursor` 未使用:Sticky 由 `last_used_at` 决定顺序,不参与轮转破并列。
pub fn select(input: &SelectionInput<'_>, _cursor: &mut usize) -> Option<Account> {
let candidates = input.candidates;
let mut selected = candidates.first()?;
for candidate in &candidates[1..] {
if compare_last_used(
candidate.last_used_at.as_deref(),
selected.last_used_at.as_deref(),
) == Ordering::Greater
{
selected = candidate;
}
}
Some(selected.clone())
}
+20 -14
View File
@@ -22,7 +22,7 @@ use thiserror::Error;
use tokio_tungstenite::tungstenite::handshake::client::generate_key;
use uuid::Uuid;
use crate::upstream::fingerprint::Fingerprint;
use crate::upstream::fingerprint::{Fingerprint, RuntimeFingerprint};
use crate::upstream::models::backend_entry::BackendModelEntry;
use crate::upstream::protocol::events::{extract_sse_usage, retry_after_seconds_from_body};
use crate::upstream::protocol::responses::{
@@ -333,14 +333,18 @@ pub struct CodexCompactResponse {
pub struct CodexBackendClient {
pub(super) client: Client,
pub(super) base_url: String,
fingerprint: Fingerprint,
fingerprint: RuntimeFingerprint,
websocket_pool: Option<Arc<CodexWebSocketPool>>,
websocket_first_token_timeout: Option<Duration>,
}
impl CodexBackendClient {
/// 构造客户端。
pub fn new(client: Client, base_url: impl Into<String>, fingerprint: Fingerprint) -> Self {
pub fn new(
client: Client,
base_url: impl Into<String>,
fingerprint: RuntimeFingerprint,
) -> Self {
Self {
client,
base_url: base_url.into().trim_end_matches('/').to_string(),
@@ -798,17 +802,18 @@ impl CodexBackendClient {
&self,
context: CodexRequestContext<'_>,
) -> CodexClientResult<Vec<BackendModelEntry>> {
let fingerprint = self.fingerprint.current();
let endpoints = [
format!(
"{}/codex/models?client_version={}",
self.base_url, self.fingerprint.app_version
self.base_url, fingerprint.app_version
),
format!("{}/models", self.base_url),
format!("{}/sentinel/chat-requirements", self.base_url),
];
for endpoint in endpoints {
let headers = self.auxiliary_request_headers(context)?;
let headers = self.auxiliary_request_headers(&fingerprint, context)?;
let response = self.client.get(endpoint).headers(headers).send().await?;
if !response.status().is_success() {
continue;
@@ -845,9 +850,10 @@ impl CodexBackendClient {
}
fn request_headers(&self, context: CodexRequestContext<'_>) -> CodexClientResult<HeaderMap> {
let fingerprint = self.fingerprint.current();
let request_id = context.session_id.unwrap_or(context.request_id);
let ordered_headers = build_ordered_codex_base_headers(
&self.fingerprint,
&fingerprint,
context.access_token,
context.account_id,
);
@@ -897,13 +903,11 @@ impl CodexBackendClient {
fn auxiliary_request_headers(
&self,
fingerprint: &Fingerprint,
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let ordered_headers = build_ordered_codex_base_headers(
&self.fingerprint,
context.access_token,
context.account_id,
);
let ordered_headers =
build_ordered_codex_base_headers(fingerprint, context.access_token, context.account_id);
let mut headers = HeaderMap::new();
insert_ordered_headers(&mut headers, &ordered_headers)?;
@@ -924,10 +928,11 @@ impl CodexBackendClient {
&self,
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let fingerprint = self.fingerprint.current();
let mut headers = HeaderMap::new();
headers.insert(
USER_AGENT,
HeaderValue::from_str(&self.fingerprint.user_agent())?,
HeaderValue::from_str(&fingerprint.user_agent())?,
);
headers.insert(
AUTHORIZATION,
@@ -935,7 +940,7 @@ impl CodexBackendClient {
);
headers.insert(
HeaderName::from_static("originator"),
HeaderValue::from_str(&self.fingerprint.originator)?,
HeaderValue::from_str(&fingerprint.originator)?,
);
insert_optional_header(&mut headers, "chatgpt-account-id", context.account_id)?;
insert_optional_header(&mut headers, "cookie", context.cookie_header)?;
@@ -947,8 +952,9 @@ impl CodexBackendClient {
&self,
context: CodexRequestContext<'_>,
) -> CodexClientResult<HeaderMap> {
let fingerprint = self.fingerprint.current();
let ordered_headers = build_ordered_codex_base_headers(
&self.fingerprint,
&fingerprint,
context.access_token,
context.account_id,
);
+12 -5
View File
@@ -755,7 +755,7 @@ async fn execute_fresh_pooled_response_create_request(
key,
PooledWebSocketConnection {
websocket,
metadata,
metadata: reusable_websocket_metadata(metadata),
created_at: now,
},
)
@@ -827,7 +827,7 @@ async fn execute_pooled_response_create_request(
let created_at = connection.created_at;
match collect_websocket_response(
connection.websocket,
connection.metadata,
reusable_websocket_metadata(connection.metadata),
true,
started_at,
first_token_started_at,
@@ -840,7 +840,7 @@ async fn execute_pooled_response_create_request(
key,
PooledWebSocketConnection {
websocket,
metadata,
metadata: reusable_websocket_metadata(metadata),
created_at,
},
)
@@ -861,7 +861,7 @@ async fn execute_pooled_response_create_request_stream(
connection: PooledWebSocketConnection,
) -> Result<CodexWebSocketStreamingExchange, CodexWebSocketExchangeError> {
let mut websocket = connection.websocket;
let mut metadata = connection.metadata;
let mut metadata = reusable_websocket_metadata(connection.metadata);
let created_at = connection.created_at;
let first_token_started_at = Instant::now();
if let Err(error) = websocket
@@ -1144,6 +1144,13 @@ fn websocket_connection_metadata(
}
}
fn reusable_websocket_metadata(
mut metadata: CodexWebSocketConnectionMetadata,
) -> CodexWebSocketConnectionMetadata {
metadata.rate_limit_headers.clear();
metadata
}
// ---------------------------------------------------------------------------
// Streaming
// ---------------------------------------------------------------------------
@@ -1406,7 +1413,7 @@ async fn finish_stream_websocket(
pool_return.key,
PooledWebSocketConnection {
websocket,
metadata,
metadata: reusable_websocket_metadata(metadata),
created_at: pool_return.created_at,
},
)
@@ -17,7 +17,7 @@ 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);
pub(crate) const DEFAULT_FIRST_TOKEN_TIMEOUT: Duration = Duration::from_secs(15);
pub(crate) const DEFAULT_FIRST_TOKEN_TIMEOUT: Duration = Duration::from_secs(20);
/// WebSocket 连接池 key。
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+15 -3
View File
@@ -30,7 +30,11 @@ async fn admin_accounts_import_should_store_cpr_account_tokens() {
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
@@ -639,7 +643,11 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
@@ -855,7 +863,11 @@ async fn admin_accounts_import_should_require_admin_session_cookie() {
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
+15 -3
View File
@@ -23,7 +23,11 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() {
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
@@ -748,7 +752,11 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() {
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
@@ -791,7 +799,11 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
+62 -4
View File
@@ -27,7 +27,11 @@ async fn admin_accounts_list_should_not_expose_account_tokens() {
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
@@ -69,6 +73,51 @@ async fn admin_accounts_list_should_not_expose_account_tokens() {
assert!(!response_body.contains("refresh-secret-token"));
}
#[tokio::test]
async fn admin_accounts_list_should_derive_refreshing_display_status_from_active_lease() {
let (app, _state, pool, _dir) =
admin_accounts_test_app("admin-accounts-refreshing-display.sqlite", 72).await;
sqlx::query("insert into accounts (id, email, access_token, refresh_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?, ?)")
.bind("acct_refreshing")
.bind("refreshing@example.com")
.bind("access-token")
.bind("refresh-token")
.bind("active")
.bind("2026-07-09T00:00:00Z")
.bind("2026-07-09T00:00:00Z")
.execute(&pool)
.await
.unwrap();
sqlx::query("insert into account_refresh_leases (account_id, owner, expires_at, updated_at) values (?, ?, ?, ?)")
.bind("acct_refreshing")
.bind("test-owner")
.bind("2099-01-01T00:00:00Z")
.bind("2026-07-09T00:00:00Z")
.execute(&pool)
.await
.unwrap();
let response = app
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/accounts?page=1&pageSize=10")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_accounts_refreshing_display")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response_json(response).await;
let item = &body["data"]["items"][0];
assert_eq!(item["status"], "active");
assert_eq!(item["displayStatus"], "refreshing");
assert_eq!(item["tokenRefreshing"], true);
}
#[tokio::test]
async fn admin_accounts_list_should_include_usage_quota_and_model_stats() {
let (app, _state, pool, _dir) =
@@ -267,13 +316,22 @@ 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"][0]["localUsage"]["totalTokensDisplay"],
"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"][1]["localUsage"]["totalTokensDisplay"],
"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_eq!(
item["quota"]["windows"][2]["localUsage"]["totalTokensDisplay"],
"124K"
);
assert!(item["quota"]["windows"][2]["windowUsedDisplay"]
.as_str()
.is_some_and(|value| value.contains(" / 7.0d")));
+9 -2
View File
@@ -26,7 +26,10 @@ use sqlx::SqlitePool;
use tower::util::ServiceExt;
use crate::support::jwt::unsigned_jwt;
use crate::support::{admin::seed_admin_session, config::test_config, http::response_json};
use crate::support::{
admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint,
http::response_json,
};
mod exporting;
mod importing;
@@ -170,7 +173,11 @@ async fn admin_accounts_test_app_with_overrides(
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
+2 -2
View File
@@ -175,7 +175,7 @@ async fn admin_account_quota_should_send_usage_cookie() {
"shortTerm"
);
assert_eq!(
body["data"]["quotaData"]["windows"][0]["tokenUsageDisplay"],
body["data"]["quotaData"]["windows"][0]["localUsage"]["totalTokensDisplay"],
"2K"
);
assert_eq!(
@@ -183,7 +183,7 @@ async fn admin_account_quota_should_send_usage_cookie() {
"周限额"
);
assert_eq!(
body["data"]["quotaData"]["windows"][1]["tokenUsageDisplay"],
body["data"]["quotaData"]["windows"][1]["localUsage"]["totalTokensDisplay"],
"2K"
);
}
+2 -2
View File
@@ -7,7 +7,7 @@ use codex_proxy_rs::infra::{database::connect_sqlite, identity::hash_admin_passw
use sqlx::SqlitePool;
use tower::util::ServiceExt;
use crate::support::{config::test_config, http::response_json};
use crate::support::{config::test_config, fingerprint::runtime_fingerprint, http::response_json};
#[tokio::test]
async fn admin_login_should_issue_http_only_session_cookie() {
@@ -202,7 +202,7 @@ async fn admin_login_test_app(
let services = std::sync::Arc::new(codex_proxy_rs::runtime::services::Services::new(
&config,
stores,
fingerprint,
runtime_fingerprint(fingerprint),
));
let state = codex_proxy_rs::runtime::state::AppState {
services: (*services).clone(),
+9 -2
View File
@@ -17,7 +17,10 @@ use codex_proxy_rs::{
use serde_json::json;
use tower::util::ServiceExt;
use crate::support::{admin::seed_admin_session, config::test_config, http::response_json};
use crate::support::{
admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint,
http::response_json,
};
mod authorization;
mod lifecycle;
@@ -41,7 +44,11 @@ async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, tempfile::Te
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let services = std::sync::Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
+14 -41
View File
@@ -32,7 +32,9 @@ use serde_json::{json, Value};
use sqlx::SqlitePool;
use tower::util::ServiceExt;
use crate::support::{admin::seed_admin_session, http::response_json};
use crate::support::{
admin::seed_admin_session, fingerprint::runtime_fingerprint, http::response_json,
};
#[tokio::test]
async fn dashboard_summary_should_render_fingerprint_updated_at_as_beijing_time() {
@@ -68,7 +70,7 @@ async fn dashboard_summary_should_render_dash_when_fingerprint_updated_at_is_mis
#[tokio::test]
async fn dashboard_summary_should_calculate_today_traffic_and_average_first_token_latency_from_time_buckets(
) {
let (app, store, pool, _dir) = dashboard_test_app(
let (app, store, _pool, _dir) = dashboard_test_app(
"dashboard-traffic-latency.sqlite",
crate::support::fingerprint::test_fingerprint(),
)
@@ -94,23 +96,22 @@ async fn dashboard_summary_should_calculate_today_traffic_and_average_first_toke
.append(&usage_record_with_tokens(yesterday_record_at, 30))
.await
.unwrap();
seed_usage_summary_with_cached_tokens(&pool, 120, 30, 24).await;
let body = dashboard_summary(app).await;
assert_eq!(body["data"]["cards"]["traffic"]["todayRequests"], "2");
assert_eq!(body["data"]["cards"]["traffic"]["todayRequestsValue"], 2);
assert_eq!(body["data"]["cards"]["traffic"]["totalRequests"], "42");
assert_eq!(body["data"]["cards"]["traffic"]["totalRequests"], "3");
assert_eq!(
body["data"]["cards"]["traffic"]["yesterdayRequestsValue"],
1
);
assert_eq!(body["data"]["cards"]["tokens"]["todayTokens"], "30");
assert_eq!(body["data"]["cards"]["tokens"]["todayTokensValue"], 30);
assert_eq!(body["data"]["cards"]["tokens"]["totalTokens"], "150");
assert_eq!(body["data"]["cards"]["tokens"]["totalTokens"], "60");
assert_eq!(body["data"]["cards"]["tokens"]["yesterdayTokensValue"], 30);
assert_eq!(body["data"]["cards"]["cache"]["totalHitRate"], "20.0%");
assert_eq!(body["data"]["cards"]["cache"]["totalCachedTokens"], "24");
assert_eq!(body["data"]["cards"]["cache"]["totalHitRate"], "0.0%");
assert_eq!(body["data"]["cards"]["cache"]["totalCachedTokens"], "0");
assert_eq!(
body["data"]["cards"]["cache"]["averageFirstTokenLatencyMs"],
"300 ms"
@@ -654,7 +655,11 @@ async fn dashboard_test_app(
client_keys: SqliteClientKeyStore::new(pool.clone()),
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let services = std::sync::Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
@@ -713,39 +718,6 @@ async fn seed_usage_summary(pool: &SqlitePool, input_tokens: u64, output_tokens:
.await;
}
async fn seed_usage_summary_with_cached_tokens(
pool: &SqlitePool,
input_tokens: u64,
output_tokens: u64,
cached_tokens: u64,
) {
sqlx::query(
"insert into accounts (id, email, access_token, status, added_at, updated_at) values (?, ?, ?, ?, ?, ?)",
)
.bind("acct_usage_cached")
.bind("acct-usage-cached@example.com")
.bind("cipher")
.bind("active")
.bind("2026-06-18T00:00:00Z")
.bind("2026-06-18T00:00:00Z")
.execute(pool)
.await
.unwrap();
sqlx::query(
"insert into account_usage (account_id, request_count, input_tokens, output_tokens, cached_tokens, total_tokens, last_used_at) values (?, ?, ?, ?, ?, ?, ?)",
)
.bind("acct_usage_cached")
.bind(42_i64)
.bind(input_tokens as i64)
.bind(output_tokens as i64)
.bind(cached_tokens as i64)
.bind((input_tokens + output_tokens) as i64)
.bind("2026-06-18T00:00:00Z")
.execute(pool)
.await
.unwrap();
}
async fn seed_usage_summary_with_last_used(
pool: &SqlitePool,
input_tokens: u64,
@@ -851,6 +823,7 @@ fn usage_record_with_account_tokens(
fn usage_record(created_at: chrono::DateTime<Utc>, model: &str, metadata: Value) -> UsageRecord {
let mut record = UsageRecord::new("v1.response", UsageRecordLevel::Info, "completed");
record.model = Some(model.to_string());
record.status_code = Some(200);
record.created_at = created_at;
record.metadata = metadata;
record
@@ -23,7 +23,10 @@ use serde_json::json;
use sqlx::SqlitePool;
use tower::util::ServiceExt;
use crate::support::{admin::seed_admin_session, config::test_config, http::response_json};
use crate::support::{
admin::seed_admin_session, config::test_config, fingerprint::runtime_fingerprint,
http::response_json,
};
#[tokio::test]
async fn admin_usage_records_should_require_admin_session_cookie() {
@@ -51,10 +54,16 @@ async fn admin_usage_records_should_cursor_page_events_and_include_request_id()
let now = Utc::now();
let mut older = UsageRecord::new("request", UsageRecordLevel::Info, "older");
older.id = "usage_older".to_string();
older.account_id = Some("acct_usage".to_string());
older.model = Some("gpt-5.5".to_string());
older.status_code = Some(200);
older.created_at = now;
store.append(&older).await.unwrap();
let mut newer = UsageRecord::new("request", UsageRecordLevel::Info, "newer");
newer.id = "usage_newer".to_string();
newer.account_id = Some("acct_usage".to_string());
newer.model = Some("gpt-5.5".to_string());
newer.status_code = Some(200);
newer.created_at = now + Duration::seconds(1);
store.append(&newer).await.unwrap();
@@ -85,8 +94,11 @@ async fn admin_usage_records_should_return_numbered_page_metadata() {
("usage_old", "older timeout", 0),
("usage_new", "newer timeout", 1),
] {
let mut event = UsageRecord::new("request", UsageRecordLevel::Error, message);
let mut event = UsageRecord::new("request", UsageRecordLevel::Info, message);
event.id = id.to_string();
event.account_id = Some("acct_usage".to_string());
event.model = Some("gpt-5.5".to_string());
event.status_code = Some(200);
event.created_at = now + Duration::seconds(offset);
store.append(&event).await.unwrap();
}
@@ -95,7 +107,7 @@ async fn admin_usage_records_should_return_numbered_page_metadata() {
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/usage/records?page=1&pageSize=1&level=error&search=timeout")
.uri("/api/admin/usage/records?page=1&pageSize=1&search=timeout")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_usage_records_numbered")
.body(Body::empty())
@@ -200,12 +212,13 @@ async fn admin_usage_records_should_return_table_display_fields() {
async fn admin_usage_records_should_use_record_account_email_snapshot_when_account_is_missing() {
let (app, store, _dir) =
admin_usage_records_test_app("admin-usage-records-email-snapshot.sqlite").await;
let mut event = UsageRecord::new("v1.response", UsageRecordLevel::Error, "upstream failed");
let mut event = UsageRecord::new("v1.response", UsageRecordLevel::Info, "completed");
event.id = "usage_email_snapshot".to_string();
event.account_id = Some("acct_deleted".to_string());
event.model = Some("gpt-5.5".to_string());
event.status_code = Some(200);
event.metadata = json!({
"accountEmail": "deleted@example.com",
"failed": true
"accountEmail": "deleted@example.com"
});
store.append(&event).await.unwrap();
@@ -234,7 +247,9 @@ async fn admin_usage_records_summary_should_aggregate_filtered_usage() {
admin_usage_records_test_app("admin-usage-records-summary.sqlite").await;
let mut success = UsageRecord::new("request", UsageRecordLevel::Info, "summary success");
success.id = "usage_summary_success".to_string();
success.account_id = Some("acct_summary".to_string());
success.model = Some("gpt-5.5".to_string());
success.status_code = Some(200);
success.latency_ms = Some(1200);
success.metadata = json!({
"usage": {
@@ -247,6 +262,7 @@ async fn admin_usage_records_summary_should_aggregate_filtered_usage() {
let mut failure = UsageRecord::new("request", UsageRecordLevel::Error, "summary failure");
failure.id = "usage_summary_failure".to_string();
failure.account_id = Some("acct_summary".to_string());
failure.model = Some("gpt-5.5-mini".to_string());
failure.status_code = Some(429);
failure.latency_ms = Some(600);
@@ -288,8 +304,10 @@ async fn admin_usage_record_insight_cards_should_aggregate_filtered_usage_dimens
admin_usage_records_test_app_with_pool("admin-usage-records-insights.sqlite").await;
let mut first = UsageRecord::new("request", UsageRecordLevel::Info, "insight success");
first.id = "usage_insight_success".to_string();
first.account_id = Some("acct_insight".to_string());
first.route = Some("/v1/responses".to_string());
first.model = Some("gpt-5.5".to_string());
first.status_code = Some(200);
first.latency_ms = Some(240);
first.metadata = json!({
"stream": true,
@@ -307,6 +325,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.account_id = Some("acct_insight".to_string());
second.route = Some("/v1/responses/review".to_string());
second.model = Some("gpt-5.5-mini".to_string());
second.status_code = Some(500);
@@ -448,6 +467,9 @@ async fn admin_usage_record_trends_should_bucket_by_china_calendar_day() {
admin_usage_records_test_app_with_pool("admin-usage-records-china-day-trend.sqlite").await;
let mut june_29 = UsageRecord::new("request", UsageRecordLevel::Info, "china-day trend");
june_29.id = "usage_china_day_29".to_string();
june_29.account_id = Some("acct_trend".to_string());
june_29.model = Some("gpt-5.5".to_string());
june_29.status_code = Some(200);
june_29.created_at = chrono::DateTime::parse_from_rfc3339("2026-06-29T15:30:00Z")
.unwrap()
.with_timezone(&Utc);
@@ -462,6 +484,9 @@ async fn admin_usage_record_trends_should_bucket_by_china_calendar_day() {
let mut june_30 = UsageRecord::new("request", UsageRecordLevel::Info, "china-day trend");
june_30.id = "usage_china_day_30".to_string();
june_30.account_id = Some("acct_trend".to_string());
june_30.model = Some("gpt-5.5".to_string());
june_30.status_code = Some(200);
june_30.created_at = chrono::DateTime::parse_from_rfc3339("2026-06-29T16:30:00Z")
.unwrap()
.with_timezone(&Utc);
@@ -514,8 +539,10 @@ async fn admin_usage_records_insight_card_endpoints_should_return_source_specifi
admin_usage_records_test_app_with_pool("admin-usage-records-insight-cards.sqlite").await;
let mut first = UsageRecord::new("request", UsageRecordLevel::Info, "card success");
first.id = "usage_card_success".to_string();
first.account_id = Some("acct_card".to_string());
first.route = Some("/v1/responses".to_string());
first.model = Some("gpt-5.5".to_string());
first.status_code = Some(200);
first.latency_ms = Some(240);
first.metadata = json!({
"stream": true,
@@ -532,8 +559,10 @@ 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.account_id = Some("acct_card".to_string());
second.route = Some("/v1/responses/compact".to_string());
second.model = Some("gpt-5.5-mini".to_string());
second.status_code = Some(200);
second.latency_ms = Some(120);
second.metadata = json!({
"stream": false,
@@ -659,11 +688,17 @@ async fn admin_usage_records_should_filter_summary_and_table_by_time_range() {
let now = Utc::now();
let mut older = UsageRecord::new("request", UsageRecordLevel::Info, "outside range");
older.id = "usage_outside_range".to_string();
older.account_id = Some("acct_range".to_string());
older.model = Some("gpt-5.5".to_string());
older.status_code = Some(200);
older.created_at = now - Duration::days(3);
store.append(&older).await.unwrap();
let mut current = UsageRecord::new("request", UsageRecordLevel::Info, "inside range");
current.id = "usage_inside_range".to_string();
current.account_id = Some("acct_range".to_string());
current.model = Some("gpt-5.5".to_string());
current.status_code = Some(200);
current.created_at = now;
store.append(&current).await.unwrap();
@@ -707,18 +742,18 @@ async fn admin_usage_records_should_filter_summary_and_table_by_time_range() {
#[tokio::test]
async fn admin_usage_records_should_filter_and_cursor_page_events() {
let (app, store, _dir) = admin_usage_records_test_app("admin-usage-records.sqlite").await;
let mut matching = UsageRecord::new("request", UsageRecordLevel::Error, "upstream timeout");
let mut matching = UsageRecord::new("request", UsageRecordLevel::Info, "upstream timeout");
matching.id = "usage_matching".to_string();
matching.account_id = Some("acct_filter".to_string());
matching.route = Some("/v1/responses".to_string());
matching.model = Some("gpt-5.5".to_string());
matching.status_code = Some(200);
store.append(&matching).await.unwrap();
store
.append(&UsageRecord::new(
"request",
UsageRecordLevel::Info,
"upstream timeout",
))
.await
.unwrap();
let mut non_matching_message = UsageRecord::new("request", UsageRecordLevel::Info, "completed");
non_matching_message.account_id = Some("acct_filter".to_string());
non_matching_message.model = Some("gpt-5.5".to_string());
non_matching_message.status_code = Some(200);
store.append(&non_matching_message).await.unwrap();
store
.append(&UsageRecord::new(
"account",
@@ -732,7 +767,7 @@ async fn admin_usage_records_should_filter_and_cursor_page_events() {
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/usage/records?kind=request&level=error&search=timeout&limit=1")
.uri("/api/admin/usage/records?kind=request&level=info&search=timeout&limit=1")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_usage_records_filter")
.body(Body::empty())
@@ -883,7 +918,11 @@ async fn admin_usage_records_test_app_with_config(
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let services = std::sync::Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
+1 -1
View File
@@ -7,7 +7,7 @@ fn settings_service_should_apply_retained_settings_patch() {
let mut current = AdminSettings {
max_concurrent_per_account: 4,
request_interval_ms: 50,
rotation_strategy: "least_used".to_string(),
rotation_strategy: "smart".to_string(),
..AdminSettings::default()
};
+19 -6
View File
@@ -28,7 +28,8 @@ use sqlx::SqlitePool;
use tower::util::ServiceExt;
use crate::support::{
admin::seed_admin_session, config::test_config as base_test_config, http::response_json,
admin::seed_admin_session, config::test_config as base_test_config,
fingerprint::runtime_fingerprint, http::response_json,
};
#[tokio::test]
@@ -68,7 +69,7 @@ async fn admin_settings_should_return_runtime_fields() {
assert_eq!(body["data"]["modelAliases"]["codex-fast"], "gpt-5.5");
assert_eq!(body["data"]["refreshMarginSeconds"], 240);
assert_eq!(body["data"]["refreshConcurrency"], 2);
assert_eq!(body["data"]["rotationStrategy"], "least_used");
assert_eq!(body["data"]["rotationStrategy"], "smart");
}
#[tokio::test]
@@ -277,7 +278,7 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Services::new(&config, stores, fingerprint);
let services = Services::new(&config, stores, runtime_fingerprint(fingerprint));
let services = std::sync::Arc::new(services);
let state = AppState {
services: (*services).clone(),
@@ -382,7 +383,11 @@ async fn admin_settings_update_should_apply_runtime_services() {
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let services = std::sync::Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
services
.account_pool
.restore_from_repository()
@@ -480,7 +485,11 @@ async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() {
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let services = std::sync::Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
@@ -522,7 +531,11 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let services = std::sync::Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
let state = AppState {
services: (*services).clone(),
};
+72 -2
View File
@@ -35,7 +35,9 @@ use wiremock::{
};
use crate::support::{
admin::seed_admin_session, config::test_config, fingerprint::test_fingerprint,
admin::seed_admin_session,
config::test_config,
fingerprint::{runtime_fingerprint, test_fingerprint},
http::response_json,
};
@@ -858,6 +860,70 @@ async fn update_events_should_close_on_shutdown() {
assert!(body.is_empty());
}
#[tokio::test]
async fn update_events_should_close_after_terminal_update_log() {
let _guard = SYSTEM_ENV_LOCK.lock().await;
let repository = "zyycn/codex-proxy-rs-events-terminal-route";
let github = MockServer::start().await;
let deploy = tempfile::tempdir().unwrap();
let exe_path = deploy.path().join("codex-proxy-rs");
let web_dist = deploy.path().join("web/dist");
std::fs::create_dir_all(&web_dist).unwrap();
std::fs::write(&exe_path, "old-binary").unwrap();
std::fs::write(web_dist.join("index.html"), "old-web").unwrap();
mount_latest_release_with_archive(&github, repository, "0.4.0", "new-binary", "new-web").await;
set_system_update_env(repository, &github.uri());
set_system_update_paths(deploy.path(), &exe_path, &web_dist);
let (app, _dir) = admin_system_test_app("system-events-terminal.sqlite").await;
let sse_response = app
.clone()
.oneshot(
Request::builder()
.method("GET")
.uri("/api/admin/system/update-events")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_system_update_events_terminal")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
let update_app = app.clone();
let update_task = tokio::spawn(async move {
update_app
.oneshot(
Request::builder()
.method("POST")
.uri("/api/admin/system/update")
.header("cookie", "cpr_admin_session=session_1")
.header("x-request-id", "req_system_update_events_terminal_update")
.header(CONTENT_TYPE, "application/json")
.body(confirmed_update_body("0.4.0"))
.unwrap(),
)
.await
.unwrap()
});
let body = tokio::time::timeout(
Duration::from_secs(2),
to_bytes(sse_response.into_body(), 128 * 1024),
)
.await
.expect("update event stream should end after terminal update log")
.unwrap();
let update_response = update_task.await.unwrap();
let text = String::from_utf8(body.to_vec()).unwrap();
assert!(
update_response.status() == StatusCode::OK
&& text.contains("\"step\":\"done\"")
&& text.contains("\"terminal\":true")
&& text.contains("更新文件已替换,等待服务重启生效")
);
}
async fn get_update_detail(app: &axum::Router, refresh: bool, request_id: &str) -> Value {
let response = app
.clone()
@@ -1218,7 +1284,11 @@ async fn admin_system_test_app(db_name: &str) -> (axum::Router, tempfile::TempDi
seed_admin_session(&pool, "session_1").await;
let config = test_config(url);
let stores = stores(pool);
let services = Arc::new(Services::new(&config, stores, test_fingerprint()));
let services = Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(test_fingerprint()),
));
let state = AppState {
services: (*services).clone(),
};
+5 -23
View File
@@ -8,7 +8,6 @@ const DEFAULT_CONFIG_YAML: &str = r#"
server:
host: 0.0.0.0
port: 8080
trusted_proxies: []
api:
base_url: https://chatgpt.com/backend-api
database:
@@ -80,11 +79,10 @@ logging:
fn default_config_keeps_only_codex_backend() {
let cfg: AppConfig = parse_yaml_config(DEFAULT_CONFIG_YAML).unwrap();
assert_eq!(cfg.server.host, "0.0.0.0");
assert!(cfg.server.trusted_proxies.is_empty());
assert_eq!(cfg.api.base_url, "https://chatgpt.com/backend-api");
assert!(cfg.model_aliases.is_empty());
assert_eq!(cfg.auth.refresh_margin_seconds, 300);
assert_eq!(cfg.auth.rotation_strategy, "least_used");
assert_eq!(cfg.auth.rotation_strategy, "smart");
assert_eq!(cfg.auth.oauth_client_id, "app_EMoamEEZ73f0CkXaXp7hrann");
assert_eq!(
cfg.auth.oauth_token_endpoint,
@@ -259,7 +257,6 @@ ws_pool:
assert_eq!(cfg.server.host, "127.0.0.1");
assert_eq!(cfg.server.port, 8080);
assert!(cfg.server.trusted_proxies.is_empty());
assert_eq!(
cfg.database.url,
"sqlite://.runtime/data/codex-proxy-rs.sqlite"
@@ -280,30 +277,15 @@ ws_pool:
assert_eq!(cfg.logging.retention_days, 14);
}
#[test]
fn config_should_parse_trusted_proxy_entries() {
let cfg: AppConfig = parse_yaml_config(&DEFAULT_CONFIG_YAML.replace(
"trusted_proxies: []",
"trusted_proxies:\n - 127.0.0.1\n - 10.0.0.0/8",
))
.unwrap();
assert_eq!(
cfg.server.trusted_proxies,
vec!["127.0.0.1".to_string(), "10.0.0.0/8".to_string()]
);
}
#[test]
fn server_config_should_reject_unknown_fields() {
let err = parse_yaml_config::<AppConfig>(&DEFAULT_CONFIG_YAML.replace(
"trusted_proxies: []",
"trusted_proxies: []\n trusted_proxy: 127.0.0.1",
))
let err = parse_yaml_config::<AppConfig>(
&DEFAULT_CONFIG_YAML.replace(" port: 8080", " port: 8080\n trusted_proxies: []"),
)
.unwrap_err();
assert!(
err.to_string().contains("trusted_proxy"),
err.to_string().contains("trusted_proxies"),
"expected unknown server field to be rejected, got {err}"
);
}
+7 -95
View File
@@ -9,14 +9,12 @@ use axum::{
extract::connect_info::ConnectInfo,
http::HeaderMap,
http::{Request, StatusCode},
middleware::{from_fn, from_fn_with_state},
middleware::from_fn,
routing::get,
Extension, Router,
};
use codex_proxy_rs::http::middleware::{
request_id::{
attach_request_id, attach_request_id_with_proxy_config, ClientIp, TrustedProxyConfig,
},
request_id::{attach_request_id, ClientIp},
trace::http_trace_layer,
};
use tokio::time::{sleep, Duration};
@@ -285,7 +283,7 @@ async fn request_context_should_auto_detect_forwarded_ip_without_trusted_proxy_c
let response = app.oneshot(request).await.expect("response");
let bytes = to_bytes(response.into_body(), 1024).await.unwrap();
// 未配置可信代理时进入自动模式:CF-Connecting-IP 优先于 X-Real-IP / X-Forwarded-For。
// 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"
@@ -319,45 +317,7 @@ async fn request_context_should_skip_private_forwarded_hops_in_auto_mode() {
}
#[tokio::test]
async fn request_context_should_accept_forwarded_ip_from_trusted_proxy() {
let trusted_proxies =
TrustedProxyConfig::from_entries(&["127.0.0.0/8".to_string()]).expect("trusted proxy");
let app = Router::new()
.route(
"/ip",
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,
attach_request_id_with_proxy_config,
));
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "203.0.113.42, 10.0.0.2")
.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(),
"203.0.113.42|forwarded=true"
);
}
#[tokio::test]
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");
async fn request_context_should_accept_real_ip_when_forwarded_for_missing() {
let app = Router::new()
.route(
"/ip",
@@ -365,10 +325,7 @@ async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded
client_ip.as_str().to_string()
}),
)
.layer(from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
));
.layer(from_fn(attach_request_id));
let mut request = Request::builder()
.uri("/ip")
.header("x-real-ip", "198.51.100.23")
@@ -385,9 +342,7 @@ async fn request_context_should_accept_real_ip_from_trusted_proxy_when_forwarded
}
#[tokio::test]
async fn request_context_should_prefer_real_ip_over_forwarded_for_from_trusted_proxy() {
let trusted_proxies =
TrustedProxyConfig::from_entries(&["127.0.0.1".to_string()]).expect("trusted proxy");
async fn request_context_should_prefer_real_ip_over_forwarded_for() {
let app = Router::new()
.route(
"/ip",
@@ -395,10 +350,7 @@ async fn request_context_should_prefer_real_ip_over_forwarded_for_from_trusted_p
client_ip.as_str().to_string()
}),
)
.layer(from_fn_with_state(
trusted_proxies,
attach_request_id_with_proxy_config,
));
.layer(from_fn(attach_request_id));
let mut request = Request::builder()
.uri("/ip")
.header("x-forwarded-for", "198.51.100.99, 203.0.113.42")
@@ -415,36 +367,6 @@ async fn request_context_should_prefer_real_ip_over_forwarded_for_from_trusted_p
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "203.0.113.42");
}
#[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()
@@ -474,13 +396,3 @@ async fn request_context_should_not_inject_real_ip_header() {
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();
assert!(
error.to_string().contains("not-an-ip"),
"unexpected error: {error}"
);
}
+4 -2
View File
@@ -10,7 +10,9 @@ use serde_json::Value;
use sqlx::SqlitePool;
use tower::ServiceExt;
use crate::support::{config::test_config as base_test_config, http::response_json};
use crate::support::{
config::test_config as base_test_config, fingerprint::runtime_fingerprint, http::response_json,
};
#[tokio::test]
async fn server_router_should_serve_frontend_assets_without_shadowing_api_routes() {
@@ -200,7 +202,7 @@ async fn router_with_dist_and_database(dist: &Path, database_url: &str) -> (Rout
let services = std::sync::Arc::new(codex_proxy_rs::runtime::services::Services::new(
&config,
stores,
fingerprint,
runtime_fingerprint(fingerprint),
));
let state = codex_proxy_rs::runtime::state::AppState {
services: (*services).clone(),
+85 -5
View File
@@ -10,7 +10,7 @@ async fn sqlite_schema_should_create_current_tables() {
let pool = connect_sqlite(&url).await.unwrap();
let rows: Vec<(String,)> = sqlx::query_as(
"select name from sqlite_master where type = 'table' and name in ('schema_migrations', 'admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'model_plan_snapshots', 'session_affinities') order by name",
"select name from sqlite_master where type = 'table' and name in ('schema_migrations', 'admin_users', 'admin_sessions', 'client_api_keys', 'runtime_settings', 'accounts', 'account_refresh_leases', 'account_usage', 'account_model_usage', 'usage_time_buckets', 'account_cookies', 'fingerprints', 'fingerprint_update_history', 'usage_records', 'ops_error_logs', 'model_plan_snapshots', 'session_affinities') order by name",
)
.fetch_all(&pool)
.await
@@ -29,6 +29,7 @@ async fn sqlite_schema_should_create_current_tables() {
"fingerprint_update_history",
"fingerprints",
"model_plan_snapshots",
"ops_error_logs",
"runtime_settings",
"schema_migrations",
"session_affinities",
@@ -51,7 +52,32 @@ async fn sqlite_schema_should_record_applied_migrations() {
.await
.unwrap();
assert_eq!(rows, [(1, "initial".to_string())]);
assert_eq!(
rows,
[
(1, "initial".to_string()),
(2, "rename_rotation_strategies".to_string()),
(3, "ops_error_logs".to_string()),
]
);
}
#[tokio::test]
async fn sqlite_schema_should_reject_refreshing_account_status() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("account-status.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let error = sqlx::query(
"insert into accounts (id, access_token, status, added_at, updated_at)
values ('acct_refreshing', 'access', 'refreshing', '2026-07-09T00:00:00Z', '2026-07-09T00:00:00Z')",
)
.execute(&pool)
.await
.unwrap_err();
assert!(error.to_string().contains("CHECK constraint failed"));
}
#[test]
@@ -155,7 +181,7 @@ async fn sqlite_schema_should_persist_diagnostic_filter_indexes() {
let pool = connect_sqlite(&url).await.unwrap();
let rows: Vec<(String,)> = sqlx::query_as(
"select name from sqlite_master where type = 'index' and name in ('idx_account_cookies_expires', 'idx_usage_records_level_created', 'idx_usage_records_model_created', 'idx_usage_records_route_created', 'idx_usage_records_status_created', 'idx_usage_records_upstream_status_created', 'idx_session_affinities_active_order', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name",
"select name from sqlite_master where type = 'index' and name in ('idx_account_cookies_expires', 'idx_usage_records_level_created', 'idx_usage_records_model_created', 'idx_usage_records_route_created', 'idx_usage_records_status_created', 'idx_usage_records_upstream_status_created', 'idx_ops_error_logs_created_id', 'idx_ops_error_logs_status_created', 'idx_ops_error_logs_failure_class', 'idx_session_affinities_active_order', 'idx_usage_time_buckets_bucket', 'idx_usage_time_buckets_model_bucket') order by name",
)
.fetch_all(&pool)
.await
@@ -165,6 +191,9 @@ async fn sqlite_schema_should_persist_diagnostic_filter_indexes() {
rows.into_iter().map(|row| row.0).collect::<Vec<_>>(),
[
"idx_account_cookies_expires",
"idx_ops_error_logs_created_id",
"idx_ops_error_logs_failure_class",
"idx_ops_error_logs_status_created",
"idx_session_affinities_active_order",
"idx_usage_records_level_created",
"idx_usage_records_model_created",
@@ -469,6 +498,41 @@ async fn sqlite_schema_should_persist_structured_event_diagnostic_columns() {
);
}
#[tokio::test]
async fn sqlite_schema_should_persist_ops_error_log_columns() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("ops-error-columns.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let rows: Vec<(String,)> = sqlx::query_as(
"select name from pragma_table_info('ops_error_logs') where name in ('request_id', 'account_id', 'route', 'model', 'status_code', 'client_status_code', 'upstream_status_code', 'transport', 'attempt_index', 'failure_class', 'response_id', 'upstream_request_id', 'latency_ms', 'metadata_json') order by name",
)
.fetch_all(&pool)
.await
.unwrap();
assert_eq!(
rows.into_iter().map(|row| row.0).collect::<Vec<_>>(),
[
"account_id",
"attempt_index",
"client_status_code",
"failure_class",
"latency_ms",
"metadata_json",
"model",
"request_id",
"response_id",
"route",
"status_code",
"transport",
"upstream_request_id",
"upstream_status_code",
]
);
}
#[tokio::test]
async fn sqlite_schema_should_reject_invalid_account_status() {
let dir = tempfile::tempdir().unwrap();
@@ -501,6 +565,22 @@ async fn sqlite_schema_should_reject_invalid_usage_record_level() {
assert!(result.is_err());
}
#[tokio::test]
async fn sqlite_schema_should_reject_invalid_ops_error_status_code() {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("ops-error-status-check.sqlite");
let url = format!("sqlite://{}", db.display());
let pool = connect_sqlite(&url).await.unwrap();
let result = sqlx::query(
"insert into ops_error_logs (id, kind, status_code, message, metadata_json, created_at) values ('ops_bad', 'request', 99, 'bad', '{}', '2026-06-14T00:00:00Z')",
)
.execute(&pool)
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn sqlite_schema_should_reject_non_boolean_flags() {
let dir = tempfile::tempdir().unwrap();
@@ -536,12 +616,12 @@ async fn sqlite_schema_should_reject_invalid_runtime_settings() {
let pool = connect_sqlite(&url).await.unwrap();
let invalid_id_result = sqlx::query(
"insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (2, 300, 2, 3, 50, 'least_used', '2026-06-14T00:00:00Z')",
"insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (2, 300, 2, 3, 50, 'smart', '2026-06-14T00:00:00Z')",
)
.execute(&pool)
.await;
let invalid_refresh_result = sqlx::query(
"insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (1, 0, 2, 3, 50, 'least_used', '2026-06-14T00:00:00Z')",
"insert into runtime_settings (id, refresh_margin_seconds, refresh_concurrency, max_concurrent_per_account, request_interval_ms, rotation_strategy, updated_at) values (1, 0, 2, 3, 50, 'smart', '2026-06-14T00:00:00Z')",
)
.execute(&pool)
.await;
@@ -274,7 +274,7 @@ async fn responses_compact_should_return_rate_limit_error_when_fallback_is_exhau
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
assert!(message.contains("All accounts exhausted (1 rate-limited)"));
assert!(message.contains("compact quota reached"));
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(
event.request_id.as_deref(),
@@ -339,7 +339,7 @@ async fn responses_compact_should_preserve_upstream_client_error_status() {
assert_eq!(status, StatusCode::BAD_REQUEST);
assert_eq!(body["error"]["type"], "invalid_request_error");
assert_eq!(body["error"]["code"], "codex_api_error");
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(
event.request_id.as_deref(),
@@ -361,7 +361,7 @@ async fn responses_compact_should_preserve_upstream_client_error_status() {
#[tokio::test]
async fn responses_compact_should_return_responses_error_when_no_accounts_are_available() {
let server = MockServer::start().await;
let (app, api_key, _dir) = test_app_without_accounts(server.uri()).await;
let (app, api_key, _pool, _dir) = test_app_without_accounts(server.uri()).await;
let response = app
.oneshot(
@@ -22,7 +22,8 @@ use codex_proxy_rs::{
http::router,
infra::database::connect_sqlite,
proxy::dispatch::session_affinity::{
build_conversation_identity, SessionAffinityEntry, SqliteSessionAffinityStore,
build_conversation_identity, compute_variant_hash, prepare_variant_identity,
SessionAffinityEntry, SqliteSessionAffinityStore,
},
runtime::services::{BackgroundTaskStores, Services, UsageRecordOptions},
runtime::state::AppState,
@@ -31,6 +32,7 @@ use codex_proxy_rs::{
store::SqliteAccountStore, token_refresh::RefreshLeaseStore,
},
upstream::fingerprint::FingerprintRepository,
upstream::protocol::responses::CodexResponsesRequest,
};
use futures::{SinkExt, StreamExt};
use serde_json::{json, Value};
@@ -271,7 +273,7 @@ async fn test_app_state_with_pool_installation_id_and_usage_record_options(
let services = Services::try_with_installation_id_and_usage_record_options(
&config,
stores,
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
Some(installation_id),
usage_record_options,
)
@@ -318,7 +320,16 @@ async fn test_app_with_account_and_pool(
test_app_with_account_pool_config(base_url, |_| {}).await
}
async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, tempfile::TempDir) {
async fn test_app_with_account_pool_and_affinity(
base_url: String,
response_id: &str,
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
test_app_with_account_pool_config_and_affinity(base_url, |_| {}, Some(response_id)).await
}
async fn test_app_without_accounts(
base_url: String,
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("openai-responses-no-accounts.sqlite");
let url = format!("sqlite://{}", db.display());
@@ -327,7 +338,7 @@ async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, t
insert_model_snapshot(&pool).await;
let state = test_app_state_with_pool_and_installation_id(
&test_config(url, base_url),
pool,
pool.clone(),
TEST_INSTALLATION_ID.to_string(),
)
.await;
@@ -337,7 +348,7 @@ async fn test_app_without_accounts(base_url: String) -> (axum::Router, String, t
.restore_from_repository()
.await
.expect("empty account pool should restore");
(router::router().with_state(state), api_key, dir)
(router::router().with_state(state), api_key, pool, dir)
}
async fn test_app_with_account_pool_and_logging(
@@ -407,6 +418,14 @@ async fn test_app_with_account_pool_and_logging_capture_body(
async fn test_app_with_account_pool_config(
base_url: String,
configure: impl FnOnce(&mut AppConfig),
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
test_app_with_account_pool_config_and_affinity(base_url, configure, None).await
}
async fn test_app_with_account_pool_config_and_affinity(
base_url: String,
configure: impl FnOnce(&mut AppConfig),
affinity_response_id: Option<&str>,
) -> (axum::Router, String, SqlitePool, tempfile::TempDir) {
let dir = tempfile::tempdir().unwrap();
let db = dir.path().join("openai-record-affinity.sqlite");
@@ -414,6 +433,9 @@ async fn test_app_with_account_pool_config(
let pool = connect_sqlite(&url).await.unwrap();
let api_key = insert_client_api_key(&pool).await;
insert_account(&pool).await;
if let Some(response_id) = affinity_response_id {
insert_session_affinity(&pool, response_id, "acct_chat").await;
}
let mut config = test_config(url, base_url);
configure(&mut config);
let state = test_app_state_with_pool_and_installation_id(
@@ -428,6 +450,14 @@ async fn test_app_with_account_pool_config(
.restore_from_repository()
.await
.expect("account pool should restore");
if affinity_response_id.is_some() {
state
.services
.session_affinity
.restore_from_repository(Utc::now())
.await
.expect("session affinity should restore");
}
(router::router().with_state(state), api_key, pool, dir)
}
@@ -713,6 +743,26 @@ async fn insert_model_snapshot(pool: &SqlitePool) {
.unwrap();
}
async fn insert_session_affinity(pool: &SqlitePool, response_id: &str, account_id: &str) {
SqliteSessionAffinityStore::new(pool.clone())
.upsert(
response_id,
&SessionAffinityEntry {
account_id: account_id.to_string(),
conversation_id: format!("conv_{response_id}"),
turn_state: Some("turn-stale".to_string()),
instructions_hash: None,
input_tokens: None,
function_call_ids: Vec::new(),
variant_hash: None,
created_at: Utc::now(),
},
Duration::hours(4),
)
.await
.unwrap();
}
async fn insert_named_account(
pool: &SqlitePool,
id: &str,
@@ -1393,10 +1443,52 @@ struct ResponseUsageRecordSnapshot {
metadata_json: String,
}
#[derive(sqlx::FromRow)]
struct OpsErrorLogSnapshot {
request_id: Option<String>,
account_id: Option<String>,
route: Option<String>,
model: Option<String>,
status_code: Option<i64>,
response_id: Option<String>,
metadata_json: String,
}
impl From<OpsErrorLogSnapshot> for ResponseUsageRecordSnapshot {
fn from(snapshot: OpsErrorLogSnapshot) -> Self {
Self {
level: "error".to_string(),
request_id: snapshot.request_id,
account_id: snapshot.account_id,
route: snapshot.route,
model: snapshot.model,
status_code: snapshot.status_code,
response_id: snapshot.response_id,
metadata_json: snapshot.metadata_json,
}
}
}
async fn latest_response_usage_record(pool: &SqlitePool) -> ResponseUsageRecordSnapshot {
latest_usage_record(pool, "v1.response").await
}
async fn latest_response_ops_error_log(pool: &SqlitePool) -> ResponseUsageRecordSnapshot {
let row = sqlx::query_as::<_, OpsErrorLogSnapshot>(
"select request_id, account_id, route, model, status_code, response_id, metadata_json
from ops_error_logs
where kind = ?
order by created_at desc, id desc
limit 1",
)
.bind("v1.response")
.fetch_optional(pool)
.await
.unwrap();
row.unwrap_or_else(|| panic!("expected a v1.response ops error log"))
.into()
}
async fn latest_usage_record(pool: &SqlitePool, kind: &str) -> ResponseUsageRecordSnapshot {
let page = SqliteUsageRecordStore::new(pool.clone())
.list(
@@ -1249,7 +1249,7 @@ async fn responses_should_strip_banned_affinity_history_when_switching_to_active
}
#[tokio::test]
async fn responses_should_keep_history_when_affinity_account_is_quota_exhausted() {
async fn responses_should_strip_quota_history_when_switching_to_active_account() {
let (base_url, upstream) =
spawn_single_websocket_completed_upstream("resp_after_quota_affinity").await;
let (app, api_key, pool, _dir) =
@@ -1287,13 +1287,7 @@ async fn responses_should_keep_history_when_affinity_account_is_quota_exhausted(
captured_header(&captured.headers, "authorization"),
Some("Bearer access-primary")
);
assert_eq!(
captured_header(&captured.headers, "x-codex-turn-state"),
Some("turn_affinity_risk")
);
assert_eq!(
captured.payload["previous_response_id"],
"resp_affinity_risk"
);
assert!(captured_header(&captured.headers, "x-codex-turn-state").is_none());
assert!(captured.payload.get("previous_response_id").is_none());
assert_eq!(affinity_count.0, 1);
}
@@ -93,7 +93,7 @@ async fn responses_should_classify_sse_cyber_policy_failure_as_bad_request() {
#[tokio::test]
async fn responses_should_return_no_available_accounts_error_when_no_accounts_are_available() {
let server = MockServer::start().await;
let (app, api_key, _dir) = test_app_without_accounts(server.uri()).await;
let (app, api_key, pool, _dir) = test_app_without_accounts(server.uri()).await;
let response = app
.oneshot(
@@ -121,6 +121,35 @@ async fn responses_should_return_no_available_accounts_error_when_no_accounts_ar
assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(body["error"]["type"], "server_error");
assert_eq!(body["error"]["code"], "no_available_accounts");
let (usage_count,): (i64,) = sqlx::query_as("select count(*) from usage_records")
.fetch_one(&pool)
.await
.unwrap();
let (ops_count, account_id, model, failure_class): (
i64,
Option<String>,
Option<String>,
Option<String>,
) = sqlx::query_as(
"select count(*), max(account_id), max(model), max(failure_class) from ops_error_logs",
)
.fetch_one(&pool)
.await
.unwrap();
let (bucket_requests, bucket_errors): (i64, i64) = sqlx::query_as(
"select coalesce(sum(request_count), 0), coalesce(sum(error_count), 0) from usage_time_buckets",
)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(usage_count, 0);
assert_eq!(ops_count, 1);
assert_eq!(account_id, None);
assert_eq!(model.as_deref(), Some("gpt-5.5"));
assert_eq!(failure_class.as_deref(), Some("no_available_accounts"));
assert_eq!((bucket_requests, bucket_errors), (1, 1));
}
#[tokio::test]
@@ -1021,7 +1050,8 @@ async fn responses_should_strip_history_after_previous_response_not_found() {
"Previous response with id resp_stale was not found",
))
.await;
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
let (app, api_key, _pool, _dir) =
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
let response = app
.oneshot(
Request::builder()
@@ -1073,7 +1103,8 @@ async fn responses_should_strip_history_after_unanswered_function_call() {
"No tool output found for function call call_1",
))
.await;
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
let (app, api_key, _pool, _dir) =
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
let response = app
.oneshot(
Request::builder()
@@ -1117,7 +1148,8 @@ async fn responses_should_strip_history_after_websocket_invalid_encrypted_reason
"Invalid encrypted content in reasoning replay",
))
.await;
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
let (app, api_key, _pool, _dir) =
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
let response = app
.oneshot(
Request::builder()
@@ -1161,7 +1193,8 @@ async fn responses_should_strip_history_after_sse_invalid_encrypted_reasoning_re
"Invalid encrypted content in reasoning replay",
))
.await;
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
let (app, api_key, _pool, _dir) =
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
let response = app
.oneshot(
Request::builder()
@@ -1205,7 +1238,8 @@ async fn responses_stream_should_strip_history_after_previous_response_not_found
"Previous response with id resp_stale was not found",
))
.await;
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
let (app, api_key, _pool, _dir) =
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
let response = app
.oneshot(
Request::builder()
@@ -1258,7 +1292,8 @@ async fn responses_stream_should_strip_history_after_websocket_invalid_encrypted
"Invalid encrypted content in reasoning replay",
))
.await;
let (app, api_key, _pool, _dir) = test_app_with_account_and_pool(base_url).await;
let (app, api_key, _pool, _dir) =
test_app_with_account_pool_and_affinity(base_url, "resp_stale").await;
let response = app
.oneshot(
Request::builder()
@@ -2122,7 +2157,7 @@ async fn responses_should_return_rate_limit_error_when_429_fallback_is_exhausted
assert_eq!(body["error"]["code"], "rate_limit_exceeded");
assert_eq!(quota_state.0, 1);
assert!(quota_state.1.is_some());
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.level, "error");
assert_eq!(
@@ -2206,7 +2241,7 @@ async fn responses_stream_should_return_rate_limit_error_when_429_fallback_is_ex
assert_eq!(quota_state.0, 1);
assert!(quota_state.1.is_some());
assert!(!response_request_id.is_empty());
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.level, "error");
assert_eq!(
@@ -1539,24 +1539,76 @@ async fn responses_websocket_should_route_previous_response_id_to_recorded_accou
}
#[tokio::test]
async fn responses_websocket_non_stream_previous_response_not_found_should_strip_history_and_retry_same_account(
async fn responses_websocket_should_prefer_conversation_account_without_previous_response_id() {
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 {
accept_successful_websocket_response_with_authorization(
&listener,
"Bearer access-secondary",
"resp_conversation_affinity",
)
.await
});
let (app, state, api_key, _pool, _dir) = test_app_with_two_accounts_and_state(base_url).await;
let request_body = json!({
"model": "gpt-5.5",
"prompt_cache_key": "conv_conversation_affinity",
"input": [],
"stream": false,
"use_websocket": true
});
let mut affinity_body = request_body.as_object().unwrap().clone();
affinity_body.remove("use_websocket");
let mut affinity_request = CodexResponsesRequest::from_body(affinity_body);
affinity_request.explicit_prompt_cache_key = true;
prepare_variant_identity(&mut affinity_request);
let variant_hash = compute_variant_hash(&affinity_request);
state
.services
.session_affinity
.record(
"resp_conversation_latest".to_string(),
SessionAffinityEntry {
account_id: "acct_secondary".to_string(),
conversation_id: "conv_conversation_affinity".to_string(),
turn_state: None,
instructions_hash: None,
input_tokens: None,
function_call_ids: Vec::new(),
variant_hash: Some(variant_hash),
created_at: Utc::now(),
},
)
.await
.unwrap();
let response = app
.oneshot(responses_json_request(&api_key, &request_body))
.await
.unwrap();
let body = response_text(response).await;
let payload = upstream.await.unwrap();
assert!(body.contains("\"id\":\"resp_conversation_affinity\""));
assert!(payload["prompt_cache_key"]
.as_str()
.is_some_and(|value| value.starts_with("cp_")));
assert!(payload.get("previous_response_id").is_none());
}
#[tokio::test]
async fn responses_websocket_non_stream_unknown_previous_response_should_strip_history_before_upstream(
) {
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_payload = accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secret",
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
)
.await;
let second_payload = accept_websocket_response_with_authorization_and_message(
accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secret",
websocket_completed_response("resp_after_history_strip", 3, 1),
)
.await;
(first_payload, second_payload)
.await
});
let (app, api_key, _dir) = test_app_with_account(base_url).await;
@@ -1573,32 +1625,23 @@ async fn responses_websocket_non_stream_previous_response_not_found_should_strip
.await
.unwrap();
let body = response_json(response).await;
let (first_payload, second_payload) = upstream.await.unwrap();
let payload = upstream.await.unwrap();
assert_eq!(body["id"], "resp_after_history_strip");
assert_eq!(first_payload["previous_response_id"], "resp_missing");
assert!(second_payload.get("previous_response_id").is_none());
assert!(payload.get("previous_response_id").is_none());
}
#[tokio::test]
async fn responses_websocket_history_recovery_should_stay_on_fallback_account_when_preferred_is_full(
) {
async fn responses_websocket_should_strip_known_history_when_preferred_account_is_full() {
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_payload = accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secondary",
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
)
.await;
let second_payload = accept_websocket_response_with_authorization_and_message(
accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secondary",
websocket_completed_response("resp_after_fallback_history_strip", 3, 1),
)
.await;
(first_payload, second_payload)
.await
});
let (app, state, api_key, _pool, _dir) =
test_app_with_two_accounts_and_state_config(base_url, |config| {
@@ -1620,35 +1663,23 @@ async fn responses_websocket_history_recovery_should_stay_on_fallback_account_wh
.await
.unwrap();
let body = response_json(response).await;
let (first_payload, second_payload) = upstream.await.unwrap();
let payload = upstream.await.unwrap();
assert_eq!(body["id"], "resp_after_fallback_history_strip");
assert_eq!(
first_payload["previous_response_id"],
"resp_owned_by_primary"
);
assert!(second_payload.get("previous_response_id").is_none());
assert!(payload.get("previous_response_id").is_none());
}
#[tokio::test]
async fn responses_websocket_stream_history_recovery_should_stay_on_fallback_account_when_preferred_is_full(
) {
async fn responses_websocket_stream_should_strip_known_history_when_preferred_account_is_full() {
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_payload = accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secondary",
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
)
.await;
let second_payload = accept_websocket_response_with_authorization_and_message(
accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secondary",
websocket_completed_response("resp_stream_after_fallback_history_strip", 3, 1),
)
.await;
(first_payload, second_payload)
.await
});
let (app, state, api_key, _pool, _dir) =
test_app_with_two_accounts_and_state_config(base_url, |config| {
@@ -1670,35 +1701,24 @@ async fn responses_websocket_stream_history_recovery_should_stay_on_fallback_acc
.await
.unwrap();
let body = response_text(response).await;
let (first_payload, second_payload) = upstream.await.unwrap();
let payload = upstream.await.unwrap();
assert!(body.contains("resp_stream_after_fallback_history_strip"));
assert_eq!(
first_payload["previous_response_id"],
"resp_owned_by_primary"
);
assert!(second_payload.get("previous_response_id").is_none());
assert!(payload.get("previous_response_id").is_none());
}
#[tokio::test]
async fn responses_websocket_stream_previous_response_not_found_should_strip_history_and_retry_same_account(
) {
async fn responses_websocket_stream_unknown_previous_response_should_strip_history_before_upstream()
{
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_payload = accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secret",
WEBSOCKET_PREVIOUS_RESPONSE_NOT_FOUND.trim().to_string(),
)
.await;
let second_payload = accept_websocket_response_with_authorization_and_message(
accept_websocket_response_with_authorization_and_message(
&listener,
"Bearer access-secret",
websocket_completed_response("resp_stream_after_history_strip", 3, 1),
)
.await;
(first_payload, second_payload)
.await
});
let (app, api_key, _dir) = test_app_with_account(base_url).await;
@@ -1715,11 +1735,10 @@ async fn responses_websocket_stream_previous_response_not_found_should_strip_his
.await
.unwrap();
let body = response_text(response).await;
let (first_payload, second_payload) = upstream.await.unwrap();
let payload = upstream.await.unwrap();
assert!(body.contains("\"id\":\"resp_stream_after_history_strip\""));
assert_eq!(first_payload["previous_response_id"], "resp_missing");
assert!(second_payload.get("previous_response_id").is_none());
assert!(payload.get("previous_response_id").is_none());
}
#[tokio::test]
@@ -1742,7 +1761,8 @@ async fn responses_websocket_non_stream_unanswered_function_call_should_strip_hi
.await;
(first_payload, second_payload)
});
let (app, api_key, _dir) = test_app_with_account(base_url).await;
let (app, api_key, _pool, _dir) =
test_app_with_account_pool_and_affinity(base_url, "resp_with_call").await;
let response = app
.oneshot(responses_json_request(
@@ -1808,7 +1828,9 @@ async fn responses_websocket_previous_response_id_should_retry_fallback_account_
.unwrap();
assert!(body.contains("\"id\":\"resp_history_fallback\""));
assert_eq!(fallback_payload["previous_response_id"], "resp_prev");
// 亲和缓存无 `resp_prev` 记录(Unknown owner),换账号前历史即被剥离,
// 由 `input` 让备用账号重建上下文,不把上一账号作用域的 id 透传给它。
assert_eq!(fallback_payload.get("previous_response_id"), None);
assert_eq!(secondary_usage.0, 1);
}
@@ -1858,7 +1880,9 @@ async fn responses_websocket_non_stream_previous_response_id_should_retry_fallba
.unwrap();
assert_eq!(body["id"], "resp_history_fallback_non_stream");
assert_eq!(fallback_payload["previous_response_id"], "resp_prev");
// 亲和缓存无 `resp_prev` 记录(Unknown owner),换账号前历史即被剥离,
// 由 `input` 让备用账号重建上下文,不把上一账号作用域的 id 透传给它。
assert_eq!(fallback_payload.get("previous_response_id"), None);
assert_eq!(secondary_usage.0, 1);
}
@@ -2304,10 +2328,9 @@ async fn responses_stream_with_previous_response_id_should_forward_websocket_chu
captured_header(&captured.headers, "authorization"),
Some("Bearer access-secret")
);
assert_eq!(
captured.payload["previous_response_id"],
"resp_runtime_pool_previous"
);
// 亲和缓存无该 previous_response_id 的记录(Unknown),故上游 payload 里应被剥离,
// 靠已展开的 input 让选中账号自行重建续链——与 sub2api/codex2api 的默认行为一致。
assert_eq!(captured.payload["previous_response_id"], Value::Null);
}
#[tokio::test]
@@ -2335,10 +2358,8 @@ async fn responses_stream_with_previous_response_id_should_record_websocket_audi
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert!(body.contains("resp_live_websocket_stream"));
assert_eq!(
captured.payload["previous_response_id"],
"resp_runtime_pool_previous"
);
// 亲和缓存无记录(Unknown)→ 上游 payload 里 previous_response_id 被剥离。
assert_eq!(captured.payload["previous_response_id"], Value::Null);
assert_eq!(event.level, "info");
assert_eq!(metadata["stream"], true);
assert_eq!(metadata["transport"], "websocket");
@@ -678,7 +678,7 @@ async fn responses_stream_should_record_usage_record_after_late_disconnect() {
let response = response_task.await.unwrap();
close_upstream.send(()).unwrap();
let body = response_text(response).await;
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert!(body.contains("stream_disconnected"));
@@ -744,7 +744,7 @@ async fn responses_stream_should_record_failure_detail_after_transport_error() {
let response = response_task.await.unwrap();
close_upstream.send(()).unwrap();
let _body = response_text(response).await;
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(metadata["failureSource"], "proxy");
@@ -788,7 +788,7 @@ async fn responses_stream_should_not_record_first_token_when_metadata_only_prefi
let response = response_task.await.unwrap();
close_upstream.send(()).unwrap();
let body = response_text(response).await;
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
let first_token_bucket: (i64, i64) = sqlx::query_as(
"select coalesce(sum(first_token_latency_sum), 0), coalesce(sum(first_token_latency_count), 0) from usage_time_buckets",
@@ -855,7 +855,7 @@ async fn responses_should_record_request_count_when_5xx_retries_are_exhausted()
assert_eq!(status, StatusCode::BAD_GATEWAY);
assert_eq!(authorizations.len(), 3, "requests: {authorizations:?}");
assert_eq!(failed_usage.map(|row| row.0).unwrap_or_default(), 1);
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(event.account_id.as_deref(), Some("acct_chat"));
assert_eq!(metadata["accountEmail"], "user@example.com");
@@ -906,7 +906,7 @@ async fn responses_should_record_attempt_trace_after_retrying_another_account()
)
.await
.unwrap();
let event = latest_response_usage_record(&pool).await;
let event = latest_response_ops_error_log(&pool).await;
let metadata: Value = serde_json::from_str(&event.metadata_json).unwrap();
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
@@ -261,6 +261,14 @@ fn session_affinity_map_should_find_latest_response_by_conversation_and_variant(
map.lookup_latest_response_by_conversation("conv", None, Some("variant_a"), now),
Some("resp_new".to_string())
);
assert_eq!(
map.lookup_latest_account_by_conversation("conv", None, Some("variant_a"), now),
Some("acct_b".to_string())
);
assert_eq!(
map.lookup_latest_account_by_conversation("conv", None, Some("variant_b"), now),
Some("acct_c".to_string())
);
assert_eq!(
map.lookup_turn_state("resp_new", now),
Some("turn_new".to_string())
+7 -2
View File
@@ -20,7 +20,8 @@ use codex_proxy_rs::{
use tower::util::ServiceExt;
use crate::support::{
client_keys::insert_client_api_key, config::test_config, http::response_json,
client_keys::insert_client_api_key, config::test_config, fingerprint::runtime_fingerprint,
http::response_json,
};
#[tokio::test]
@@ -164,7 +165,11 @@ async fn build_test_app(
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let services = std::sync::Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
services
.initialize_hot_path_state()
.await
@@ -18,7 +18,8 @@ use sqlx::SqlitePool;
use tower::util::ServiceExt;
use crate::support::{
client_keys::insert_client_api_key, config::test_config, http::response_json,
client_keys::insert_client_api_key, config::test_config, fingerprint::runtime_fingerprint,
http::response_json,
};
#[tokio::test]
@@ -154,7 +155,11 @@ async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::Temp
usage_records: SqliteUsageRecordStore::new(pool.clone()),
};
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let services = std::sync::Arc::new(Services::new(
&config,
stores,
runtime_fingerprint(fingerprint),
));
services
.initialize_hot_path_state()
.await
+3 -3
View File
@@ -49,7 +49,7 @@ async fn services_try_new_should_use_configured_tls_transport_builder() {
let Err(error) = Services::try_new(
&config,
stores,
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
) else {
panic!("invalid custom CA should fail service transport construction");
};
@@ -90,7 +90,7 @@ async fn services_try_new_should_expose_websocket_pool_when_enabled() {
let services = Services::try_new(
&config,
stores(pool),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.expect("services should build");
@@ -111,7 +111,7 @@ async fn services_try_new_should_not_expose_websocket_pool_when_disabled() {
let services = Services::try_new(
&config,
stores(pool),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.expect("services should build");
+18 -33
View File
@@ -11,9 +11,11 @@ async fn fingerprint_update_task_should_start_background_checker() {
repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint())
.await
.expect("seed current fingerprint");
let runtime_fingerprint = runtime_fingerprint_from_repo(&repo).await;
let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task(
Some(repo),
repo,
runtime_fingerprint,
"https://example.invalid/appcast.xml".to_string(),
dir.path().join("extracted-fingerprint.json"),
"26.800.1".to_string(),
@@ -49,6 +51,7 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito
repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint())
.await
.expect("seed current fingerprint");
let runtime_fingerprint = runtime_fingerprint_from_repo(&repo).await;
let extracted_path = dir.path().join("extracted-fingerprint.json");
std::fs::write(
&extracted_path,
@@ -57,7 +60,8 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito
.expect("extracted fingerprint should be written");
let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task(
Some(repo.clone()),
repo.clone(),
runtime_fingerprint.clone(),
format!("{}/appcast.xml", server.uri()),
extracted_path,
"26.800.1".to_string(),
@@ -75,6 +79,15 @@ async fn fingerprint_update_task_should_apply_initial_appcast_update_to_reposito
),
("26.900.1", "7001", "147")
);
let runtime_snapshot = runtime_fingerprint.snapshot();
assert_eq!(
(
runtime_snapshot.app_version.as_str(),
runtime_snapshot.build_number.as_str(),
runtime_snapshot.chromium_version.as_str()
),
("26.900.1", "7001", "147")
);
}
#[tokio::test]
@@ -103,8 +116,10 @@ async fn fingerprint_update_task_should_not_persist_when_appcast_matches_current
repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint())
.await
.expect("seed current fingerprint");
let runtime_fingerprint = runtime_fingerprint_from_repo(&repo).await;
let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task(
Some(repo.clone()),
repo.clone(),
runtime_fingerprint,
format!("{}/appcast.xml", server.uri()),
dir.path().join("missing-extracted-fingerprint.json"),
"26.900.1".to_string(),
@@ -120,33 +135,3 @@ async fn fingerprint_update_task_should_not_persist_when_appcast_matches_current
assert_eq!(history_count, 0);
}
#[tokio::test]
async fn fingerprint_update_task_should_check_appcast_without_repository() {
let server = MockServer::start().await;
mount_appcast(
&server,
r#"
<rss>
<channel>
<item>
<enclosure url="https://example.invalid/download" sparkle:shortVersionString="26.900.1" sparkle:version="7001" />
</item>
</channel>
</rss>
"#,
)
.await;
let dir = tempfile::tempdir().expect("temp dir");
let handle = codex_proxy_rs::runtime::tasks::fingerprint_update::start_fingerprint_update_task(
None,
format!("{}/appcast.xml", server.uri()),
dir.path().join("missing-extracted-fingerprint.json"),
"26.800.1".to_string(),
"6001".to_string(),
);
wait_for_appcast_requests(&server, 1).await;
handle.shutdown().await;
}
+10 -1
View File
@@ -10,7 +10,7 @@ use codex_proxy_rs::upstream::accounts::cookies::SqliteCookieStore;
use codex_proxy_rs::upstream::accounts::model::{Account, AccountStatus};
use codex_proxy_rs::upstream::accounts::model::{AccountModelUsageDelta, AccountUsageDelta};
use codex_proxy_rs::upstream::accounts::store::{AccountStore, AccountStoreResult};
use codex_proxy_rs::upstream::fingerprint::FingerprintRepository;
use codex_proxy_rs::upstream::fingerprint::{FingerprintRepository, RuntimeFingerprint};
use codex_proxy_rs::upstream::models::config::ModelConfig;
use codex_proxy_rs::upstream::models::service::ModelService;
use wiremock::{
@@ -137,6 +137,15 @@ async fn wait_for_current_fingerprint_version(
}
}
async fn runtime_fingerprint_from_repo(repo: &FingerprintRepository) -> RuntimeFingerprint {
RuntimeFingerprint::new(
repo.load_current()
.await
.expect("current fingerprint should load")
.expect("current fingerprint should exist"),
)
}
async fn wait_for_appcast_requests(server: &MockServer, expected_count: usize) {
let deadline = tokio::time::Instant::now() + StdDuration::from_secs(2);
loop {
+27 -3
View File
@@ -1,6 +1,9 @@
use super::*;
use std::collections::BTreeMap;
use std::{
collections::BTreeMap,
sync::atomic::{AtomicUsize, Ordering},
};
use codex_proxy_rs::upstream::{
accounts::pool::{AccountAcquireRequest, AccountPoolOptions, RuntimeAccountPoolService},
@@ -50,6 +53,8 @@ async fn model_refresh_task_should_sync_model_plan_allowlist_to_account_pool() {
plan_account("acct-plus", "plus"),
plan_account("acct-team", "team"),
],
request_usage_records: Arc::new(AtomicUsize::new(0)),
model_request_usage_records: Arc::new(AtomicUsize::new(0)),
});
let account_pool = Arc::new(RuntimeAccountPoolService::new(
account_store.clone(),
@@ -74,6 +79,16 @@ async fn model_refresh_task_should_sync_model_plan_allowlist_to_account_pool() {
.await
.expect("team model should acquire an account");
assert_eq!(acquired.account.id, "acct-team");
assert_eq!(
account_store.request_usage_records.load(Ordering::SeqCst),
0
);
assert_eq!(
account_store
.model_request_usage_records
.load(Ordering::SeqCst),
0
);
}
fn empty_model_config() -> ModelConfig {
@@ -139,6 +154,8 @@ impl CodexModelCatalogClient for FakeModelCatalogClient {
struct PlanAccountStore {
accounts: Vec<Account>,
request_usage_records: Arc<AtomicUsize>,
model_request_usage_records: Arc<AtomicUsize>,
}
#[async_trait]
@@ -174,8 +191,11 @@ impl AccountStore for PlanAccountStore {
async fn record_usage_delta(
&self,
_account_id: &str,
_usage: AccountUsageDelta,
usage: AccountUsageDelta,
) -> AccountStoreResult<()> {
if usage.requests > 0 {
self.request_usage_records.fetch_add(1, Ordering::SeqCst);
}
Ok(())
}
@@ -183,8 +203,12 @@ impl AccountStore for PlanAccountStore {
&self,
_account_id: &str,
_model: &str,
_usage: AccountModelUsageDelta,
usage: AccountModelUsageDelta,
) -> AccountStoreResult<()> {
if usage.requests > 0 {
self.model_request_usage_records
.fetch_add(1, Ordering::SeqCst);
}
Ok(())
}
+1 -2
View File
@@ -8,7 +8,6 @@ pub(crate) fn test_config(database_url: String) -> AppConfig {
server: ServerConfig {
host: "127.0.0.1".to_string(),
port: 0,
trusted_proxies: Vec::new(),
},
api: ApiConfig {
base_url: "https://chatgpt.com/backend-api".to_string(),
@@ -20,7 +19,7 @@ pub(crate) fn test_config(database_url: String) -> AppConfig {
refresh_concurrency: 2,
max_concurrent_per_account: 3,
request_interval_ms: 50,
rotation_strategy: "least_used".to_string(),
rotation_strategy: "smart".to_string(),
tier_priority: Vec::new(),
oauth_client_id: "app_EMoamEEZ73f0CkXaXp7hrann".to_string(),
oauth_token_endpoint: "https://auth.openai.com/oauth/token".to_string(),
+12 -1
View File
@@ -1,4 +1,7 @@
use codex_proxy_rs::{config::schema::FingerprintConfig, upstream::fingerprint::Fingerprint};
use codex_proxy_rs::{
config::schema::FingerprintConfig,
upstream::fingerprint::{Fingerprint, RuntimeFingerprint},
};
pub(crate) fn test_fingerprint() -> Fingerprint {
Fingerprint::from_config(&FingerprintConfig::default())
@@ -9,3 +12,11 @@ pub(crate) fn test_fingerprint_with_updated_at(updated_at: Option<&str>) -> Fing
fingerprint.updated_at = updated_at.map(ToString::to_string);
fingerprint
}
pub(crate) fn runtime_test_fingerprint() -> RuntimeFingerprint {
RuntimeFingerprint::new(test_fingerprint())
}
pub(crate) fn runtime_fingerprint(fingerprint: Fingerprint) -> RuntimeFingerprint {
RuntimeFingerprint::new(fingerprint)
}
@@ -153,10 +153,10 @@ fn account_pool_should_skip_accounts_in_cloudflare_cooldown() {
}
#[test]
fn least_used_should_deprioritize_quota_limited_accounts_when_skip_is_disabled() {
fn quota_reset_priority_should_deprioritize_quota_limited_accounts_when_skip_is_disabled() {
let mut pool = AccountPool::with_options(AccountPoolOptions {
skip_quota_limited: false,
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
let mut limited = crate::support::accounts::test_account("limited", AccountStatus::Active);
@@ -172,10 +172,10 @@ fn least_used_should_deprioritize_quota_limited_accounts_when_skip_is_disabled()
}
#[test]
fn least_used_should_prefer_earlier_rate_limit_window_reset() {
fn quota_reset_priority_should_prefer_earlier_rate_limit_window_reset() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
let mut soon = crate::support::accounts::test_account("soon", AccountStatus::Active);
@@ -543,10 +543,244 @@ fn account_pool_should_cleanup_stale_slots_before_acquire() {
}
#[test]
fn account_pool_should_rotate_tied_least_used_accounts() {
fn smart_should_spread_concurrent_requests_by_slot_pressure() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
max_concurrent_per_account: 3,
rotation_strategy: RotationStrategy::Smart,
..AccountPoolOptions::default()
});
pool.insert(crate::support::accounts::test_account(
"acct_a",
AccountStatus::Active,
));
pool.insert(crate::support::accounts::test_account(
"acct_b",
AccountStatus::Active,
));
pool.insert(crate::support::accounts::test_account(
"acct_c",
AccountStatus::Active,
));
let first = pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
let second = pool
.acquire_with(&AccountAcquireRequest::new(
"gpt-5.5",
now + Duration::milliseconds(1),
))
.unwrap();
let third = pool
.acquire_with(&AccountAcquireRequest::new(
"gpt-5.5",
now + Duration::milliseconds(2),
))
.unwrap();
assert_eq!(
[first.account.id, second.account.id, third.account.id]
.into_iter()
.collect::<BTreeSet<_>>()
.len(),
3
);
}
#[test]
fn smart_should_balance_released_requests_by_runtime_usage() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::Smart,
..AccountPoolOptions::default()
});
pool.insert(crate::support::accounts::test_account(
"acct_a",
AccountStatus::Active,
));
pool.insert(crate::support::accounts::test_account(
"acct_b",
AccountStatus::Active,
));
let first = pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
pool.release(&first.account.id);
let second = pool
.acquire_with(&AccountAcquireRequest::new(
"gpt-5.5",
now + Duration::seconds(1),
))
.unwrap();
assert_ne!(first.account.id, second.account.id);
}
#[test]
fn smart_should_not_let_reset_window_override_lower_request_count() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::Smart,
..AccountPoolOptions::default()
});
let mut reset_soon =
crate::support::accounts::test_account("reset_soon", AccountStatus::Active);
reset_soon.request_count = 10;
reset_soon.window_request_count = 10;
reset_soon.window_reset_at = Some(now + Duration::seconds(30));
let mut quieter = crate::support::accounts::test_account("quieter", AccountStatus::Active);
quieter.request_count = 2;
quieter.window_request_count = 2;
quieter.window_reset_at = Some(now + Duration::seconds(300));
pool.insert(reset_soon);
pool.insert(quieter);
let acquired = pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
assert_eq!(acquired.account.id, "quieter");
}
#[test]
fn smart_should_not_drain_earliest_reset_account_when_other_accounts_are_idle() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::Smart,
..AccountPoolOptions::default()
});
let mut reset_first =
crate::support::accounts::test_account("reset_first", AccountStatus::Active);
reset_first.window_reset_at = Some(now + Duration::seconds(30));
let mut reset_later =
crate::support::accounts::test_account("reset_later", AccountStatus::Active);
reset_later.window_reset_at = Some(now + Duration::minutes(10));
let mut reset_latest =
crate::support::accounts::test_account("reset_latest", AccountStatus::Active);
reset_latest.window_reset_at = Some(now + Duration::minutes(20));
pool.insert(reset_first);
pool.insert(reset_later);
pool.insert(reset_latest);
let mut selected_ids = Vec::new();
for offset in 0..4 {
let acquired = pool
.acquire_with(&AccountAcquireRequest::new(
"gpt-5.5",
now + Duration::seconds(offset),
))
.unwrap();
selected_ids.push(acquired.account.id.clone());
pool.release(&acquired.account.id);
}
assert_eq!(
selected_ids
.iter()
.take(3)
.cloned()
.collect::<BTreeSet<_>>()
.len(),
3
);
assert!(selected_ids.windows(2).all(|pair| pair[0] != pair[1]));
}
#[test]
fn smart_should_use_window_token_load_when_request_counts_tie() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::Smart,
..AccountPoolOptions::default()
});
let mut heavy = crate::support::accounts::test_account("heavy", AccountStatus::Active);
heavy.request_count = 5;
heavy.window_request_count = 3;
heavy.window_input_tokens = 1_000_000;
heavy.window_output_tokens = 500_000;
let mut light = crate::support::accounts::test_account("light", AccountStatus::Active);
light.request_count = 5;
light.window_request_count = 3;
light.window_input_tokens = 8_000;
light.window_output_tokens = 2_000;
pool.insert(heavy);
pool.insert(light);
let acquired = pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
assert_eq!(acquired.account.id, "light");
}
#[test]
fn smart_should_not_double_count_cached_tokens_when_request_counts_tie() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::Smart,
..AccountPoolOptions::default()
});
let mut cache_heavy =
crate::support::accounts::test_account("cache_heavy", AccountStatus::Active);
cache_heavy.request_count = 5;
cache_heavy.window_request_count = 3;
cache_heavy.window_input_tokens = 10_000;
cache_heavy.window_output_tokens = 1_000;
cache_heavy.window_cached_tokens = 100_000;
let mut heavier_actual =
crate::support::accounts::test_account("heavier_actual", AccountStatus::Active);
heavier_actual.request_count = 5;
heavier_actual.window_request_count = 3;
heavier_actual.window_input_tokens = 20_000;
heavier_actual.window_output_tokens = 1_000;
pool.insert(cache_heavy);
pool.insert(heavier_actual);
let acquired = pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
assert_eq!(acquired.account.id, "cache_heavy");
}
#[test]
fn smart_should_use_cached_tokens_as_tie_break_when_actual_token_load_ties() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::Smart,
..AccountPoolOptions::default()
});
let mut cache_heavy =
crate::support::accounts::test_account("cache_heavy", AccountStatus::Active);
cache_heavy.request_count = 5;
cache_heavy.window_request_count = 3;
cache_heavy.window_input_tokens = 10_000;
cache_heavy.window_output_tokens = 1_000;
cache_heavy.window_cached_tokens = 100_000;
let mut cache_light =
crate::support::accounts::test_account("cache_light", AccountStatus::Active);
cache_light.request_count = 5;
cache_light.window_request_count = 3;
cache_light.window_input_tokens = 10_000;
cache_light.window_output_tokens = 1_000;
cache_light.window_cached_tokens = 1_000;
pool.insert(cache_heavy);
pool.insert(cache_light);
let acquired = pool
.acquire_with(&AccountAcquireRequest::new("gpt-5.5", now))
.unwrap();
assert_eq!(acquired.account.id, "cache_light");
}
#[test]
fn account_pool_should_rotate_tied_quota_reset_priority_accounts() {
let mut pool = AccountPool::with_options(AccountPoolOptions {
max_concurrent_per_account: 1,
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
pool.insert(crate::support::accounts::test_account(
@@ -566,11 +800,11 @@ fn account_pool_should_rotate_tied_least_used_accounts() {
}
#[test]
fn least_used_should_prefer_lru_before_tie_rotation() {
fn quota_reset_priority_should_prefer_lru_before_tie_rotation() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
max_concurrent_per_account: 1,
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
pool.insert(crate::support::accounts::test_account(
@@ -603,10 +837,10 @@ fn least_used_should_prefer_lru_before_tie_rotation() {
}
#[test]
fn least_used_should_prefer_lower_runtime_request_count() {
fn quota_reset_priority_should_prefer_lower_runtime_request_count() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
let mut busy = crate::support::accounts::test_account("busy", AccountStatus::Active);
@@ -624,10 +858,10 @@ fn least_used_should_prefer_lower_runtime_request_count() {
}
#[test]
fn least_used_should_not_penalize_missing_window_reset() {
fn quota_reset_priority_should_not_penalize_missing_window_reset() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
let mut fresh = crate::support::accounts::test_account("fresh", AccountStatus::Active);
@@ -647,10 +881,10 @@ fn least_used_should_not_penalize_missing_window_reset() {
}
#[test]
fn least_used_should_fall_through_to_request_count_without_window_resets() {
fn quota_reset_priority_should_fall_through_to_request_count_without_window_resets() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
let mut busier = crate::support::accounts::test_account("busier", AccountStatus::Active);
@@ -668,11 +902,11 @@ fn least_used_should_fall_through_to_request_count_without_window_resets() {
}
#[test]
fn least_used_should_sort_quota_limited_accounts_by_reset_when_not_skipping() {
fn quota_reset_priority_should_sort_quota_limited_accounts_by_reset_when_not_skipping() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
skip_quota_limited: false,
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
let mut later = crate::support::accounts::test_account("later", AccountStatus::Active);
@@ -692,10 +926,10 @@ fn least_used_should_sort_quota_limited_accounts_by_reset_when_not_skipping() {
}
#[test]
fn least_used_should_not_mutate_candidate_order() {
fn quota_reset_priority_should_not_mutate_candidate_order() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
let mut first = crate::support::accounts::test_account("first", AccountStatus::Active);
@@ -68,10 +68,10 @@ fn capacity_summary_should_not_exclude_cloudflare_cooldown_accounts() {
}
#[test]
fn least_used_should_compare_request_count_when_only_one_account_has_window_reset() {
fn quota_reset_priority_should_compare_request_count_when_only_one_account_has_window_reset() {
let now = fixed_time();
let mut pool = AccountPool::with_options(AccountPoolOptions {
rotation_strategy: RotationStrategy::LeastUsed,
rotation_strategy: RotationStrategy::QuotaResetPriority,
..AccountPoolOptions::default()
});
let mut reset_known =
@@ -312,7 +312,7 @@ async fn quota_refresh_service_should_send_usage_cookie_when_cookie_store_is_con
let codex = CodexBackendClient::new(
reqwest::Client::new(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let service =
RuntimeQuotaRefreshService::new(store, Arc::new(codex)).with_cookie_store(cookies);
@@ -393,7 +393,7 @@ async fn quota_refresh_service_should_fetch_usage_for_quota_locked_accounts_and_
let codex = CodexBackendClient::new(
reqwest::Client::new(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let service = RuntimeQuotaRefreshService::new(store.clone(), Arc::new(codex));
@@ -468,7 +468,7 @@ async fn quota_refresh_service_should_refresh_quota_exhausted_accounts() {
let codex = CodexBackendClient::new(
reqwest::Client::new(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let service = RuntimeQuotaRefreshService::new(store, Arc::new(codex));
@@ -505,7 +505,7 @@ async fn quota_refresh_service_should_skip_recent_locked_account_before_cooldown
let codex = CodexBackendClient::new(
reqwest::Client::new(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let service =
RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800);
@@ -546,7 +546,7 @@ async fn quota_refresh_service_should_skip_recent_locked_account_inside_cooldown
let codex = CodexBackendClient::new(
reqwest::Client::new(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let service =
RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800);
@@ -604,7 +604,7 @@ async fn quota_refresh_service_should_bypass_recent_skip_after_cooldown_grace()
let codex = CodexBackendClient::new(
reqwest::Client::new(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let service =
RuntimeQuotaRefreshService::with_min_refresh_interval_secs(store, Arc::new(codex), 1_800);
@@ -658,7 +658,7 @@ async fn quota_refresh_service_should_stagger_multiple_locked_account_requests()
let codex = CodexBackendClient::new(
reqwest::Client::new(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let service = RuntimeQuotaRefreshService::new(store, Arc::new(codex))
.with_request_spacing(StdDuration::from_millis(200));
@@ -68,3 +68,66 @@ async fn refresh_lease_store_should_acquire_release_and_respect_expiry() {
.await
.expect("lease should be acquirable after release"));
}
#[tokio::test]
async fn refresh_lease_store_should_list_active_account_ids() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("refresh-lease-active-ids.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let accounts = SqliteAccountStore::new(pool.clone());
for id in ["acct-active-lease", "acct-expired-lease"] {
accounts
.insert(NewAccount {
id: id.to_string(),
email: None,
account_id: None,
user_id: None,
label: None,
plan_type: None,
access_token: SecretString::new("access".to_string().into()),
refresh_token: Some(SecretString::new("refresh".to_string().into())),
access_token_expires_at: None,
status: AccountStatus::Active,
added_at: None,
})
.await
.expect("account should be inserted");
}
let leases = RefreshLeaseStore::new(pool);
let now = Utc.with_ymd_and_hms(2026, 7, 9, 8, 0, 0).unwrap();
leases
.try_acquire(
"acct-active-lease",
"owner-a",
now + Duration::minutes(5),
now,
)
.await
.expect("active lease should be acquired");
leases
.try_acquire(
"acct-expired-lease",
"owner-b",
now - Duration::minutes(5),
now - Duration::minutes(10),
)
.await
.expect("expired lease should be acquired");
let account_ids = vec![
"acct-active-lease".to_string(),
"acct-expired-lease".to_string(),
"acct-missing".to_string(),
];
let active = leases
.active_account_ids(&account_ids, now)
.await
.expect("active leases should load");
assert_eq!(
active,
std::collections::HashSet::from(["acct-active-lease".to_string()])
);
}
@@ -1,7 +1,7 @@
use super::*;
#[tokio::test]
async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_call() {
async fn token_refresh_task_should_keep_business_status_before_refresher_call() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-status.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
@@ -13,7 +13,7 @@ async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_
let new_access_token = test_jwt(new_expires_at.timestamp());
store
.insert(NewAccount {
id: "acct-refreshing-marker".to_string(),
id: "acct-status-marker".to_string(),
email: Some("marker@example.com".to_string()),
account_id: Some("chatgpt-marker".to_string()),
user_id: Some("user-marker".to_string()),
@@ -32,7 +32,7 @@ async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_
let observed_statuses = Arc::new(Mutex::new(Vec::new()));
let refresher = StoreObservingTokenRefresher {
store: store.clone(),
account_id: "acct-refreshing-marker".to_string(),
account_id: "acct-status-marker".to_string(),
observed_statuses: observed_statuses.clone(),
response: Ok(TokenPair {
access_token: new_access_token.clone(),
@@ -53,14 +53,14 @@ async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_
.await
.expect("refresh should succeed");
let stored = store
.get("acct-refreshing-marker")
.get("acct-status-marker")
.await
.expect("account should load")
.expect("account should exist");
let observed_statuses = observed_statuses.lock().await.clone();
assert_eq!(summary.refreshed, 1);
assert_eq!(observed_statuses, [AccountStatus::Refreshing]);
assert_eq!(observed_statuses, [AccountStatus::Active]);
assert_eq!(stored.status, AccountStatus::Active);
assert_eq!(
stored.access_token.expose_secret(),
@@ -69,7 +69,7 @@ async fn token_refresh_task_should_mark_due_account_refreshing_before_refresher_
}
#[tokio::test]
async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
async fn token_refresh_task_should_recover_expired_account_after_restart() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-recovery.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
@@ -81,7 +81,7 @@ async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
let new_access_token = test_jwt((now + Duration::hours(2)).timestamp());
store
.insert(NewAccount {
id: "acct-refreshing-recovery".to_string(),
id: "acct-expired-recovery".to_string(),
email: Some("recovery@example.com".to_string()),
account_id: Some("chatgpt-recovery".to_string()),
user_id: Some("user-recovery".to_string()),
@@ -91,7 +91,7 @@ async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
refresh_token: Some(SecretString::new("refresh-recovery".to_string().into())),
added_at: None,
access_token_expires_at: Some(now + Duration::hours(1)),
status: AccountStatus::Refreshing,
status: AccountStatus::Expired,
})
.await
.expect("account should be inserted");
@@ -113,9 +113,9 @@ async fn token_refresh_task_should_recover_refreshing_account_after_restart() {
let summary = task
.refresh_due_accounts_once_at(now)
.await
.expect("refreshing account should recover");
.expect("expired account should recover");
let stored = store
.get("acct-refreshing-recovery")
.get("acct-expired-recovery")
.await
.expect("account should load")
.expect("account should exist");
@@ -183,7 +183,7 @@ async fn token_refresh_task_should_retry_transport_failure_before_recovery() {
let observed_statuses = observed_statuses.lock().await.clone();
assert_eq!(summary.failed, 1);
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 5]);
assert_eq!(observed_statuses, [AccountStatus::Active; 5]);
assert_eq!(stored.status, AccountStatus::Active);
assert!(stored.next_refresh_at.is_some_and(|next| next > now));
assert_eq!(
@@ -275,7 +275,7 @@ async fn token_refresh_task_should_delay_recovery_after_retry_exhaustion() {
assert!(recovery_at <= now + Duration::minutes(12));
assert_eq!(delayed.skipped, 1);
assert_eq!(refreshed.refreshed, 1);
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 6]);
assert_eq!(observed_statuses, [AccountStatus::Active; 6]);
assert_eq!(stored.status, AccountStatus::Active);
assert_eq!(
stored.access_token.expose_secret(),
@@ -412,7 +412,7 @@ async fn token_refresh_task_should_confirm_invalid_grant_before_expiring_account
let observed_statuses = observed_statuses.lock().await.clone();
assert_eq!(summary.status_updated, 1);
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 2]);
assert_eq!(observed_statuses, [AccountStatus::Active; 2]);
assert_eq!(stored.status, AccountStatus::Expired);
assert_eq!(
stored.access_token.expose_secret(),
@@ -494,6 +494,92 @@ async fn token_refresh_task_should_fire_per_account_timer_at_refresh_time() {
);
}
#[tokio::test]
async fn token_refresh_task_should_fire_quota_exhausted_timer_without_clearing_status() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir
.path()
.join("token-refresh-quota-exhausted-timer.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let store = SqliteAccountStore::new(pool);
let now = Utc::now();
let old_expires_at = now + Duration::milliseconds(50);
let new_access_token = test_jwt((now + Duration::hours(1)).timestamp());
store
.insert(NewAccount {
id: "acct-refresh-quota-exhausted-timer".to_string(),
email: Some("quota-exhausted-timer@example.com".to_string()),
account_id: Some("chatgpt-quota-exhausted-timer".to_string()),
user_id: Some("user-quota-exhausted-timer".to_string()),
label: None,
plan_type: Some("plus".to_string()),
access_token: SecretString::new(test_jwt(old_expires_at.timestamp()).into()),
refresh_token: Some(SecretString::new(
"refresh-quota-exhausted-timer".to_string().into(),
)),
added_at: None,
access_token_expires_at: Some(old_expires_at),
status: AccountStatus::QuotaExhausted,
})
.await
.expect("account should be inserted");
let refresher = NotifyingTokenRefresher {
calls: Arc::new(AtomicUsize::new(0)),
started: Arc::new(Notify::new()),
response: Ok(TokenPair {
access_token: new_access_token.clone(),
refresh_token: None,
}),
};
let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new(
store.clone(),
RefreshPolicy {
refresh_margin_seconds: 0,
refresh_concurrency: 1,
},
refresher.clone(),
)
.with_retry_delays(vec![StdDuration::ZERO; 4]);
let started = refresher.started.notified();
tokio::pin!(started);
let summary = task
.schedule_account_timers_once_at(now)
.await
.expect("quota-exhausted timer should be scheduled");
timeout(StdDuration::from_secs(2), &mut started)
.await
.expect("quota-exhausted timer should trigger refresh");
let mut stored = store
.get("acct-refresh-quota-exhausted-timer")
.await
.expect("account should load")
.expect("account should exist");
for _ in 0..50 {
if stored.access_token.expose_secret() == new_access_token.as_str() {
break;
}
sleep(StdDuration::from_millis(20)).await;
stored = store
.get("acct-refresh-quota-exhausted-timer")
.await
.expect("account should load")
.expect("account should exist");
}
assert_eq!(summary.scheduled, 1);
assert_eq!(refresher.calls.load(Ordering::SeqCst), 1);
assert_eq!(stored.status, AccountStatus::QuotaExhausted);
assert_eq!(
stored.access_token.expose_secret(),
new_access_token.as_str()
);
assert!(stored.next_refresh_at.is_some());
}
#[tokio::test]
async fn token_refresh_task_should_reschedule_next_timer_after_scheduled_refresh() {
let dir = tempfile::tempdir().expect("temp dir");
@@ -75,6 +75,78 @@ async fn token_refresh_task_should_persist_refreshed_access_token_and_keep_refre
);
}
#[tokio::test]
async fn token_refresh_task_should_refresh_quota_exhausted_account_without_clearing_status() {
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("token-refresh-quota-exhausted.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let store = SqliteAccountStore::new(pool);
let now = Utc.with_ymd_and_hms(2026, 6, 19, 11, 0, 0).unwrap();
let new_expires_at = now + Duration::hours(1);
let new_access_token = test_jwt(new_expires_at.timestamp());
store
.insert(NewAccount {
id: "acct-refresh-quota-exhausted".to_string(),
email: Some("quota-exhausted@example.com".to_string()),
account_id: Some("chatgpt-quota-exhausted".to_string()),
user_id: Some("user-quota-exhausted".to_string()),
label: None,
plan_type: Some("plus".to_string()),
access_token: SecretString::new(
test_jwt((now + Duration::seconds(30)).timestamp()).into(),
),
refresh_token: Some(SecretString::new(
"refresh-quota-exhausted".to_string().into(),
)),
added_at: None,
access_token_expires_at: Some(now + Duration::seconds(30)),
status: AccountStatus::QuotaExhausted,
})
.await
.expect("account should be inserted");
let observed_statuses = Arc::new(Mutex::new(Vec::new()));
let refresher = StoreObservingTokenRefresher {
store: store.clone(),
account_id: "acct-refresh-quota-exhausted".to_string(),
observed_statuses: observed_statuses.clone(),
response: Ok(TokenPair {
access_token: new_access_token.clone(),
refresh_token: None,
}),
};
let task = codex_proxy_rs::upstream::accounts::token_refresh::RuntimeTokenRefreshService::new(
store.clone(),
RefreshPolicy {
refresh_margin_seconds: 300,
refresh_concurrency: 1,
},
refresher,
);
let summary = task
.refresh_due_accounts_once_at(now)
.await
.expect("quota-exhausted account should refresh token");
let stored = store
.get("acct-refresh-quota-exhausted")
.await
.expect("account should load")
.expect("account should exist");
let observed_statuses = observed_statuses.lock().await.clone();
assert_eq!(summary.refreshed, 1);
assert_eq!(observed_statuses, [AccountStatus::QuotaExhausted]);
assert_eq!(stored.status, AccountStatus::QuotaExhausted);
assert_eq!(
stored.access_token.expose_secret(),
new_access_token.as_str()
);
assert_eq!(stored.access_token_expires_at, Some(new_expires_at));
assert!(stored.next_refresh_at.is_some());
}
#[tokio::test]
async fn token_refresh_task_should_persist_success_after_transient_transport_retry() {
let dir = tempfile::tempdir().expect("temp dir");
@@ -139,7 +211,7 @@ async fn token_refresh_task_should_persist_success_after_transient_transport_ret
assert_eq!(summary.refreshed, 1);
assert_eq!(summary.failed, 0);
assert_eq!(observed_statuses, [AccountStatus::Refreshing; 2]);
assert_eq!(observed_statuses, [AccountStatus::Active; 2]);
assert_eq!(stored.status, AccountStatus::Active);
assert_eq!(
stored.access_token.expose_secret(),
+19 -9
View File
@@ -57,14 +57,21 @@ async fn update_checker_should_report_available_update_from_appcast() {
.mount(&server)
.await;
let dir = tempfile::tempdir().expect("temp dir");
let db = dir.path().join("fingerprints-check.sqlite");
let pool = connect_sqlite(&format!("sqlite://{}", db.display()))
.await
.expect("sqlite pool");
let repo = FingerprintRepository::new(pool);
repo.ensure_current_seed(&crate::support::fingerprint::test_fingerprint())
.await
.expect("seed current fingerprint");
let checker = UpdateChecker::with_client(
None,
repo,
reqwest::Client::new(),
format!("{}/appcast.xml", server.uri()),
tempfile::tempdir()
.expect("temp dir")
.path()
.join("extracted-fingerprint.json"),
dir.path().join("extracted-fingerprint.json"),
"26.800.1",
"6001",
);
@@ -113,7 +120,7 @@ async fn update_checker_should_apply_available_update_to_repository() {
.expect("write extracted fingerprint");
let checker = UpdateChecker::with_client(
Some(repo.clone()),
repo.clone(),
reqwest::Client::new(),
format!("{}/appcast.xml", server.uri()),
extracted_path,
@@ -121,17 +128,20 @@ async fn update_checker_should_apply_available_update_to_repository() {
"6001",
);
let applied = checker
let updated = checker
.check_and_apply_update()
.await
.expect("apply update");
.expect("apply update")
.expect("available update should return current fingerprint");
let stored = repo
.load_current()
.await
.expect("load current")
.expect("stored fingerprint");
assert!(applied);
assert_eq!(updated.app_version, "26.900.1");
assert_eq!(updated.build_number, "7001");
assert_eq!(updated.chromium_version, "147");
assert_eq!(stored.app_version, "26.900.1");
assert_eq!(stored.build_number, "7001");
assert_eq!(stored.chromium_version, "147");
+1
View File
@@ -2,5 +2,6 @@ pub(crate) mod accounts;
pub(crate) mod fingerprint;
pub(crate) mod models;
pub(crate) mod protocol;
pub(crate) mod scheduler;
pub(crate) mod token_client;
pub(crate) mod transport;
@@ -0,0 +1,44 @@
use codex_proxy_rs::upstream::scheduler::{FeedbackSample, FeedbackStats};
#[test]
fn sample_of_unknown_account_is_empty() {
let stats = FeedbackStats::new();
assert_eq!(stats.sample("missing"), FeedbackSample::default());
}
#[test]
fn first_sample_lands_directly() {
let stats = FeedbackStats::new();
stats.report("acct", false, Some(400));
let sample = stats.sample("acct");
assert_eq!(sample.error_rate, Some(1.0));
assert_eq!(sample.ttft_ms, Some(400.0));
}
#[test]
fn subsequent_samples_blend_by_alpha() {
let stats = FeedbackStats::new();
stats.report("acct", false, Some(1000)); // error_rate=1.0, ttft=1000
stats.report("acct", true, Some(0)); // error_rate=0.8, ttft=800
let sample = stats.sample("acct");
assert!((sample.error_rate.unwrap() - 0.8).abs() < 1e-9);
assert!((sample.ttft_ms.unwrap() - 800.0).abs() < 1e-9);
}
#[test]
fn success_only_reports_decay_error_rate() {
let stats = FeedbackStats::new();
stats.report("acct", false, None); // error_rate=1.0, ttft untouched
stats.report("acct", true, None); // error_rate=0.8
let sample = stats.sample("acct");
assert!((sample.error_rate.unwrap() - 0.8).abs() < 1e-9);
assert_eq!(sample.ttft_ms, None);
}
#[test]
fn remove_clears_account() {
let stats = FeedbackStats::new();
stats.report("acct", false, Some(100));
stats.remove("acct");
assert_eq!(stats.sample("acct"), FeedbackSample::default());
}
+2
View File
@@ -0,0 +1,2 @@
pub(crate) mod feedback;
pub(crate) mod score;
+144
View File
@@ -0,0 +1,144 @@
use chrono::Utc;
use codex_proxy_rs::upstream::accounts::model::AccountStatus;
use codex_proxy_rs::upstream::scheduler::{rank_candidates, FeedbackStats, ScoreWeights};
use crate::support::accounts::test_account;
#[test]
fn ranks_idle_account_above_loaded_one() {
let mut loaded = test_account("loaded", AccountStatus::Active);
loaded.window_input_tokens = 1_000_000;
let idle = test_account("idle", AccountStatus::Active);
let candidates = vec![loaded, idle];
let feedback = FeedbackStats::new();
let slot_count = |_: &str| 0;
let ranked = rank_candidates(
&candidates,
&ScoreWeights::default(),
Utc::now().timestamp(),
&slot_count,
&feedback,
);
// idle 账号(索引 1)窗口用量为 0,应排在负载重的账号之前。
assert_eq!(ranked.first().unwrap().0, 1);
}
#[test]
fn high_error_rate_account_is_penalized() {
let healthy = test_account("healthy", AccountStatus::Active);
let flaky = test_account("flaky", AccountStatus::Active);
let candidates = vec![healthy, flaky];
let feedback = FeedbackStats::new();
// flaky 连续失败,错误率 EWMA 拉高。
for _ in 0..5 {
feedback.report("flaky", false, None);
}
let slot_count = |_: &str| 0;
let ranked = rank_candidates(
&candidates,
&ScoreWeights::default(),
Utc::now().timestamp(),
&slot_count,
&feedback,
);
// healthy(索引 0)应排在 flaky 之前。
assert_eq!(ranked.first().unwrap().0, 0);
}
#[test]
fn quota_limited_account_ranks_last_despite_being_idle() {
let mut limited = test_account("limited", AccountStatus::Active);
limited.quota_limit_reached = true;
let mut busy = test_account("busy", AccountStatus::Active);
busy.window_input_tokens = 1_000_000;
busy.window_request_count = 100;
let candidates = vec![limited, busy];
let feedback = FeedbackStats::new();
let slot_count = |_: &str| 0;
let ranked = rank_candidates(
&candidates,
&ScoreWeights::default(),
Utc::now().timestamp(),
&slot_count,
&feedback,
);
// 即使 limited 完全空闲,配额封顶惩罚也应让繁忙但未封顶的 busy(索引 1)排前。
assert_eq!(ranked.first().unwrap().0, 1);
}
#[test]
fn actual_token_load_dominates_cached() {
// cache_heavy 实际 token 更少(但缓存高),heavier_actual 实际 token 更多。
// 实际用量权重远大于缓存,cache_heavy(索引 0)应排前。
let mut cache_heavy = test_account("cache_heavy", AccountStatus::Active);
cache_heavy.window_input_tokens = 10_000;
cache_heavy.window_cached_tokens = 100_000;
let mut heavier_actual = test_account("heavier_actual", AccountStatus::Active);
heavier_actual.window_input_tokens = 20_000;
let candidates = vec![cache_heavy, heavier_actual];
let feedback = FeedbackStats::new();
let slot_count = |_: &str| 0;
let ranked = rank_candidates(
&candidates,
&ScoreWeights::default(),
Utc::now().timestamp(),
&slot_count,
&feedback,
);
assert_eq!(ranked.first().unwrap().0, 0);
}
#[test]
fn cached_breaks_ties_when_actual_load_matches() {
// 实际 token 相同,缓存少者(cache_light,索引 1)应排前。
let mut cache_heavy = test_account("cache_heavy", AccountStatus::Active);
cache_heavy.window_input_tokens = 10_000;
cache_heavy.window_cached_tokens = 100_000;
let mut cache_light = test_account("cache_light", AccountStatus::Active);
cache_light.window_input_tokens = 10_000;
cache_light.window_cached_tokens = 1_000;
let candidates = vec![cache_heavy, cache_light];
let feedback = FeedbackStats::new();
let slot_count = |_: &str| 0;
let ranked = rank_candidates(
&candidates,
&ScoreWeights::default(),
Utc::now().timestamp(),
&slot_count,
&feedback,
);
assert_eq!(ranked.first().unwrap().0, 1);
}
#[test]
fn zero_weights_produce_flat_scores() {
let candidates = vec![
test_account("a", AccountStatus::Active),
test_account("b", AccountStatus::Active),
];
let feedback = FeedbackStats::new();
let weights = ScoreWeights {
load: 0.0,
window: 0.0,
cached: 0.0,
window_requests: 0.0,
error_rate: 0.0,
ttft: 0.0,
reset: 0.0,
};
let slot_count = |_: &str| 0;
let ranked = rank_candidates(
&candidates,
&weights,
Utc::now().timestamp(),
&slot_count,
&feedback,
);
assert!(ranked.iter().all(|(_, b)| b.total == 0.0));
}
+1 -1
View File
@@ -59,7 +59,7 @@ async fn codex_backend_client_should_apply_configured_websocket_pool() {
let backend = CodexBackendClient::new(
reqwest::Client::new(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(Arc::new(pool));
let mut request = CodexResponsesRequest::new_http_sse("gpt-5.5", "be brief", Vec::new());
+92 -6
View File
@@ -1,4 +1,5 @@
use super::*;
use codex_proxy_rs::upstream::fingerprint::RuntimeFingerprint;
use codex_proxy_rs::upstream::transport::build_ordered_codex_headers;
use codex_proxy_rs::upstream::transport::websocket::CodexWebSocketConnection;
use serde_json::Value;
@@ -279,7 +280,7 @@ async fn codex_backend_client_websocket_should_forward_security_chain_headers_an
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(pool);
@@ -420,7 +421,7 @@ async fn codex_backend_client_should_send_desktop_headers_and_capture_response_m
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let mut request =
codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse(
@@ -471,6 +472,91 @@ async fn codex_backend_client_should_send_desktop_headers_and_capture_response_m
);
}
#[tokio::test]
async fn codex_backend_client_should_use_latest_runtime_fingerprint_for_new_requests() {
let server = wiremock::MockServer::start().await;
let sse_body = include_str!("../../fixtures/responses/http_sse/completed_usage_basic.sse");
wiremock::Mock::given(wiremock::matchers::method("POST"))
.and(wiremock::matchers::path("/codex/responses"))
.respond_with(
wiremock::ResponseTemplate::new(200)
.insert_header("content-type", "text/event-stream")
.set_body_string(sse_body),
)
.mount(&server)
.await;
let mut initial = crate::support::fingerprint::test_fingerprint();
initial.app_version = "26.800.1".to_string();
let runtime_fingerprint = RuntimeFingerprint::new(initial);
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
runtime_fingerprint.clone(),
);
let mut request =
codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse(
"gpt-5.5",
"",
Vec::new(),
);
request.force_http_sse = true;
let request_context = |request_id| CodexRequestContext {
access_token: "access-token",
account_id: Some("chatgpt-account"),
request_id,
turn_state: None,
turn_metadata: None,
beta_features: None,
include_timing_metrics: None,
version: None,
codex_window_id: None,
parent_thread_id: None,
cookie_header: None,
installation_id: None,
session_id: None,
};
client
.create_response(&request, request_context("req_fingerprint_1"))
.await
.expect("first response should succeed");
let mut updated = runtime_fingerprint.snapshot();
updated.app_version = "26.900.1".to_string();
updated.build_number = "7001".to_string();
runtime_fingerprint.replace(updated);
client
.create_response(&request, request_context("req_fingerprint_2"))
.await
.expect("second response should succeed");
let user_agents = server
.received_requests()
.await
.expect("received requests should load")
.iter()
.map(|request| {
request
.headers
.get("user-agent")
.expect("user-agent should be sent")
.to_str()
.expect("user-agent should be valid")
.to_string()
})
.collect::<Vec<_>>();
assert_eq!(
user_agents,
vec![
"Codex Desktop/26.800.1 (darwin; arm64)".to_string(),
"Codex Desktop/26.900.1 (darwin; arm64)".to_string(),
]
);
}
#[tokio::test]
async fn codex_backend_client_usage_should_use_wham_usage_headers() {
let server = wiremock::MockServer::start().await;
@@ -486,7 +572,7 @@ async fn codex_backend_client_usage_should_use_wham_usage_headers() {
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let usage = client
@@ -571,7 +657,7 @@ async fn codex_backend_client_models_should_use_original_auxiliary_headers() {
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let models = client
@@ -638,7 +724,7 @@ async fn codex_backend_client_should_send_http_sse_headers_in_fingerprint_order(
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
client
@@ -713,7 +799,7 @@ async fn codex_backend_client_should_send_compact_headers_in_fingerprint_order()
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
client
@@ -107,7 +107,7 @@ async fn codex_backend_client_should_cap_non_success_error_body_at_one_mib() {
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let mut request =
codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse(
@@ -161,7 +161,7 @@ async fn codex_backend_client_should_parse_retry_after_from_rate_limit_error_bod
let client = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
server.uri(),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let mut request =
codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse(
+98 -7
View File
@@ -4,6 +4,29 @@ use codex_proxy_rs::upstream::{
transport::websocket::write_websocket_audit_artifact_for_dir,
};
fn rate_limit_event(used_percent: u64) -> String {
json!({
"type": "codex.rate_limits",
"rate_limits": {
"primary": {
"used_percent": used_percent,
"window_minutes": 43200,
"reset_at": 1893456000 + used_percent as i64,
}
}
})
.to_string()
}
fn primary_used_percent_values(headers: &[(String, String)]) -> Vec<&str> {
headers
.iter()
.filter_map(|(name, value)| {
(name == "x-codex-primary-used-percent").then_some(value.as_str())
})
.collect()
}
#[tokio::test]
async fn websocket_audit_artifact_should_require_explicit_directory() {
let dir = tempfile::tempdir().expect("temp dir");
@@ -107,7 +130,7 @@ async fn codex_backend_client_should_decode_permessage_deflate_context_takeover_
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(pool);
let mut request =
@@ -663,6 +686,74 @@ async fn websocket_execute_response_create_request_should_capture_internal_metad
.any(|(name, value)| name == "x-codex-primary-reset-at" && value == "1893456300"));
}
#[tokio::test]
async fn codex_backend_client_should_not_reuse_websocket_rate_limit_headers() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.unwrap();
let mut websocket = accept_codex_test_websocket(stream).await;
let _first_message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(rate_limit_event(10).into()))
.await
.unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_rate_limit_first", 3, 1).into(),
))
.await
.unwrap();
let _second_message = websocket.next().await.unwrap().unwrap();
websocket
.send(Message::Text(rate_limit_event(20).into()))
.await
.unwrap();
websocket
.send(Message::Text(
completed_websocket_response("resp_rate_limit_second", 5, 2).into(),
))
.await
.unwrap();
websocket.close(None).await.unwrap();
});
let pool = Arc::new(CodexWebSocketPool::new(8, Duration::from_mins(1)));
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(pool);
let request = pooled_websocket_request("conversation-rate-limit-reuse");
let first = backend
.create_response(
&request,
request_context("req_rate_limit_first", Some("chatgpt-account")),
)
.await
.expect("first pooled websocket response should succeed");
let second = backend
.create_response(
&request,
request_context("req_rate_limit_second", Some("chatgpt-account")),
)
.await
.expect("second pooled websocket response should reuse connection");
server.await.unwrap();
assert_eq!(
primary_used_percent_values(&first.rate_limit_headers),
vec!["10"]
);
assert_eq!(
primary_used_percent_values(&second.rate_limit_headers),
vec!["20"]
);
}
#[tokio::test]
async fn websocket_execute_response_create_request_should_surface_incomplete_response() {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
@@ -1000,7 +1091,7 @@ async fn codex_backend_client_stream_should_reject_binary_websocket_event() {
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let Err(error) = backend
@@ -1073,7 +1164,7 @@ async fn codex_backend_client_stream_should_retry_without_pool_when_first_conten
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_first_token_timeout(Some(Duration::from_millis(30)));
let request = pooled_websocket_request("conversation-first-token-no-pool-timeout");
@@ -1133,7 +1224,7 @@ async fn codex_backend_client_stream_should_error_when_websocket_closes_before_t
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let mut response = backend
@@ -1203,7 +1294,7 @@ async fn codex_backend_client_stream_should_wait_for_terminal_after_active_webso
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let mut response = backend
@@ -1267,7 +1358,7 @@ async fn codex_backend_client_stream_should_timeout_when_active_websocket_stalls
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
);
let mut response = backend
@@ -1363,7 +1454,7 @@ async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_
let backend = CodexBackendClient::new(
reqwest::Client::builder().no_proxy().build().unwrap(),
format!("http://{addr}"),
crate::support::fingerprint::test_fingerprint(),
crate::support::fingerprint::runtime_test_fingerprint(),
)
.with_websocket_pool(pool);

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