diff --git a/backend/src/admin/accounts/routes.rs b/backend/src/admin/accounts/routes.rs index 61b85bdf9..4744a56b9 100644 --- a/backend/src/admin/accounts/routes.rs +++ b/backend/src/admin/accounts/routes.rs @@ -163,7 +163,7 @@ impl AdminAccountData { label: a.label, plan_type: a.plan_type, has_refresh_token: a.has_refresh_token, - status: account_status_str(a.status).to_string(), + status: a.status.as_str().to_string(), access_token_expires_at, access_token_expires_at_display, added_at: china_rfc3339(&a.added_at), @@ -396,7 +396,7 @@ impl From for AccountRefreshData { Self { id: result.id, email: result.email, - previous_status: account_status_str(result.previous_status).to_string(), + previous_status: result.previous_status.as_str().to_string(), result: account_refresh_outcome_str(result.outcome), error: result.error, duration_ms: result.duration_ms, @@ -472,8 +472,8 @@ pub(crate) async fn accounts( ) -> Result { let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50)); let use_numbered_page = params.page.is_some() || params.page_size.is_some(); - let stats = account_list_stats(&state).await; - let summary = account_summary_data(&state, &stats).await; + let quota_by_account = quota_snapshots_by_account(&state).await; + let summary = account_summary_data(&state, "a_by_account).await; if use_numbered_page { return match state @@ -483,6 +483,7 @@ pub(crate) async fn accounts( .await { Ok(page) => { + let stats = account_list_stats(&state, &page.items, "a_by_account).await; let page = crate::infra::json::NumberedPage { items: page .items @@ -509,6 +510,7 @@ pub(crate) async fn accounts( .await { Ok(page) => { + let stats = account_list_stats(&state, &page.items, "a_by_account).await; let page = Page { items: page .items @@ -658,7 +660,7 @@ pub(crate) async fn account_quota( Err(AdminAccountError::NotFound) => Err(account_not_found()), Err(AdminAccountError::Inactive(status)) => Err(AdminError::conflict(format!( "Account is {}, cannot query quota", - account_status_str(status) + status.as_str() ))), Err(AdminAccountError::FetchQuota(msg)) => Err(AdminError::bad_gateway(format!( "Failed to fetch quota from Codex API: {msg}" @@ -841,13 +843,19 @@ pub(crate) async fn update_account( } } -async fn account_list_stats(state: &AppState) -> AccountListStats { - let usage_records = list_all_usage_records(state).await; - let quota_snapshots = state +async fn account_list_stats( + state: &AppState, + accounts: &[AdminAccountMetadata], + quota_by_account: &HashMap, +) -> AccountListStats { + let account_ids = accounts + .iter() + .map(|account| account.id.clone()) + .collect::>(); + let usage_records = state .services - .background_tasks - .accounts - .list_quota_snapshots() + .usage + .list_by_account_ids(&account_ids) .await .unwrap_or_default(); let models_by_account = list_current_window_model_usage(state, &usage_records).await; @@ -857,22 +865,14 @@ async fn account_list_stats(state: &AppState) -> AccountListStats { .into_iter() .map(|usage| (usage.account_id.clone(), usage)) .collect(), - quota_by_account: quota_snapshots - .into_iter() - .map(|snapshot| { - ( - snapshot.account_id, - quota_data(&snapshot.quota_json, snapshot.quota_fetched_at), - ) - }) - .collect(), + quota_by_account: quota_by_account.clone(), models_by_account, } } async fn account_summary_data( state: &AppState, - stats: &AccountListStats, + quota_by_account: &HashMap, ) -> AdminAccountSummaryData { let accounts = list_all_account_metadata(state).await; let total = accounts.len() as u64; @@ -883,8 +883,7 @@ async fn account_summary_data( let high_usage = accounts .iter() .filter(|account| { - stats - .quota_by_account + quota_by_account .get(&account.id) .is_some_and(account_quota_has_high_usage) }) @@ -902,6 +901,24 @@ async fn account_summary_data( } } +async fn quota_snapshots_by_account(state: &AppState) -> HashMap { + state + .services + .background_tasks + .accounts + .list_quota_snapshots() + .await + .unwrap_or_default() + .into_iter() + .map(|snapshot| { + ( + snapshot.account_id, + quota_data(&snapshot.quota_json, snapshot.quota_fetched_at), + ) + }) + .collect() +} + async fn list_all_account_metadata(state: &AppState) -> Vec { let mut page = 1; let mut accounts = Vec::new(); @@ -934,26 +951,6 @@ fn account_summary_needs_attention(status: AccountStatus) -> bool { ) } -async fn list_all_usage_records(state: &AppState) -> Vec { - let mut cursor = None; - let mut records = Vec::new(); - loop { - let Ok(page) = state - .services - .usage - .list(cursor, ACCOUNT_STATS_PAGE_LIMIT) - .await - else { - return Vec::new(); - }; - records.extend(page.items); - let Some(next_cursor) = page.next_cursor else { - return records; - }; - cursor = Some(next_cursor); - } -} - async fn list_current_window_model_usage( state: &AppState, usage_records: &[AdminUsageRecord], @@ -1202,14 +1199,3 @@ fn account_export_ids(value: Option<&str>) -> Vec { .map(ToString::to_string) .collect() } - -fn account_status_str(status: crate::upstream::accounts::model::AccountStatus) -> &'static str { - match status { - crate::upstream::accounts::model::AccountStatus::Active => "active", - crate::upstream::accounts::model::AccountStatus::Expired => "expired", - crate::upstream::accounts::model::AccountStatus::QuotaExhausted => "quota_exhausted", - crate::upstream::accounts::model::AccountStatus::Refreshing => "refreshing", - crate::upstream::accounts::model::AccountStatus::Disabled => "disabled", - crate::upstream::accounts::model::AccountStatus::Banned => "banned", - } -} diff --git a/backend/src/admin/accounts/service/contracts.rs b/backend/src/admin/accounts/service/contracts.rs index ba5c1c933..a6bdca161 100644 --- a/backend/src/admin/accounts/service/contracts.rs +++ b/backend/src/admin/accounts/service/contracts.rs @@ -113,12 +113,6 @@ pub enum AdminAccountError { Inactive(AcctStatus), } -#[derive(Debug, Clone)] -pub struct UpdatedAccountStatus { - pub id: String, - pub status: AcctStatus, -} - #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AdminAccountRefreshOutcome { Alive, @@ -233,6 +227,14 @@ pub(super) struct ManualCreateTokens { pub(super) access_token: String, pub(super) refresh_token_for_new: Option, pub(super) refresh_token_for_existing: Option, + pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims, +} + +#[derive(Debug, Clone)] +pub(super) struct RefreshedAccountTokens { + pub(super) access_token: String, + pub(super) refresh_token: Option, + pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims, } #[derive(Debug, Clone)] diff --git a/backend/src/admin/accounts/service/exporting.rs b/backend/src/admin/accounts/service/exporting.rs index 7b54c8266..298d43591 100644 --- a/backend/src/admin/accounts/service/exporting.rs +++ b/backend/src/admin/accounts/service/exporting.rs @@ -4,10 +4,7 @@ use super::{ contracts::{ExportedAccount, ExportedAccounts}, AdminAccountError, AdminAccountService, }; -use crate::{ - infra::time::china_rfc3339_str, - upstream::accounts::{model::AccountStatus, store::StoredAccount}, -}; +use crate::{infra::time::china_rfc3339_str, upstream::accounts::store::StoredAccount}; impl AdminAccountService { pub async fn export(&self, ids: Vec) -> Result { @@ -58,20 +55,9 @@ impl From for ExportedAccount { access_token_expires_at: account .access_token_expires_at .map(|value| value.to_rfc3339()), - status: export_status_str(account.status).to_string(), + status: account.status.as_str().to_string(), added_at: china_rfc3339_str(&account.added_at), updated_at: china_rfc3339_str(&account.updated_at), } } } - -fn export_status_str(status: AccountStatus) -> &'static str { - match status { - AccountStatus::Active => "active", - AccountStatus::Expired => "expired", - AccountStatus::QuotaExhausted => "quota_exhausted", - AccountStatus::Refreshing => "refreshing", - AccountStatus::Disabled => "disabled", - AccountStatus::Banned => "banned", - } -} diff --git a/backend/src/admin/accounts/service/importing.rs b/backend/src/admin/accounts/service/importing.rs index 78ba7d47e..d36e0615d 100644 --- a/backend/src/admin/accounts/service/importing.rs +++ b/backend/src/admin/accounts/service/importing.rs @@ -281,8 +281,7 @@ impl AdminAccountService { token: Option, refresh_token: Option, ) -> Result, AdminAccountError> { - let mut refresh_token = - crate::upstream::accounts::importing::normalize_nonempty(refresh_token); + let refresh_token = crate::upstream::accounts::importing::normalize_nonempty(refresh_token); let Some(access_token) = crate::upstream::accounts::importing::normalize_nonempty( token .as_deref() @@ -292,26 +291,12 @@ impl AdminAccountService { return Ok(None); }; let refreshed = self - .token_refresher - .refresh(&existing_refresh_token) - .await - .map_err(AdminAccountError::RefreshTokenExchange)?; - let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( - crate::upstream::accounts::importing::normalize_bearer_token( - &refreshed.access_token, - ), - )) - .ok_or(AdminAccountError::TokenRequired)?; - refresh_token = refreshed.refresh_token; - let claims = crate::upstream::accounts::token_refresh::manual_account_claims( - &access_token, - Utc::now(), - ) - .map_err(AdminAccountError::InvalidToken)?; + .refresh_tokens_from_refresh_token(&existing_refresh_token) + .await?; return Ok(Some(ResolvedImportTokens { - access_token, - refresh_token, - claims: Some(claims), + access_token: refreshed.access_token, + refresh_token: refreshed.refresh_token, + claims: Some(refreshed.claims), })); }; @@ -334,24 +319,12 @@ impl AdminAccountService { })); }; let refreshed = self - .token_refresher - .refresh(&existing_refresh_token) - .await - .map_err(AdminAccountError::RefreshTokenExchange)?; - let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( - crate::upstream::accounts::importing::normalize_bearer_token(&refreshed.access_token), - )) - .ok_or(AdminAccountError::TokenRequired)?; - refresh_token = refreshed.refresh_token; - let claims = crate::upstream::accounts::token_refresh::manual_account_claims( - &access_token, - Utc::now(), - ) - .map_err(AdminAccountError::InvalidToken)?; + .refresh_tokens_from_refresh_token(&existing_refresh_token) + .await?; Ok(Some(ResolvedImportTokens { - access_token, - refresh_token, - claims: Some(claims), + access_token: refreshed.access_token, + refresh_token: refreshed.refresh_token, + claims: Some(refreshed.claims), })) } async fn import_target_account_id( diff --git a/backend/src/admin/accounts/service/lifecycle.rs b/backend/src/admin/accounts/service/lifecycle.rs index e0485b6cb..9ce7cf672 100644 --- a/backend/src/admin/accounts/service/lifecycle.rs +++ b/backend/src/admin/accounts/service/lifecycle.rs @@ -20,7 +20,7 @@ use super::{ contracts::{ parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, AdminAccountRefreshOutcome, AdminAccountRefreshResult, - AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens, UpdatedAccountStatus, + AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens, }, AdminAccountService, }; @@ -92,37 +92,32 @@ impl AdminAccountService { .as_deref() .map(crate::upstream::accounts::importing::normalize_bearer_token), ) { + let claims = crate::upstream::accounts::token_refresh::manual_account_claims( + &access_token, + Utc::now(), + ) + .map_err(AdminAccountError::InvalidToken)?; ManualCreateTokens { access_token, refresh_token_for_new: provided_refresh_token.clone(), refresh_token_for_existing: provided_refresh_token, + claims, } } else if let Some(refresh_token) = provided_refresh_token { - let token_pair = self - .token_refresher - .refresh(&refresh_token) - .await - .map_err(AdminAccountError::RefreshTokenExchange)?; - let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( - crate::upstream::accounts::importing::normalize_bearer_token( - &token_pair.access_token, - ), - )) - .ok_or(AdminAccountError::TokenRequired)?; + let refreshed = self + .refresh_tokens_from_refresh_token(&refresh_token) + .await?; ManualCreateTokens { - access_token, - refresh_token_for_new: token_pair.refresh_token.clone(), - refresh_token_for_existing: token_pair.refresh_token, + access_token: refreshed.access_token, + refresh_token_for_new: refreshed.refresh_token.clone(), + refresh_token_for_existing: refreshed.refresh_token, + claims: refreshed.claims, } } else { return Err(AdminAccountError::TokenRequired); }; - let claims = crate::upstream::accounts::token_refresh::manual_account_claims( - &tokens.access_token, - Utc::now(), - ) - .map_err(AdminAccountError::InvalidToken)?; + let claims = tokens.claims; let existing = if let Some(account_id) = claims.account_id.as_deref() { self.store .find_by_chatgpt_identity(account_id, claims.user_id.as_deref()) @@ -195,45 +190,6 @@ impl AdminAccountService { .map(stored_to_admin_metadata) .ok_or(AdminAccountError::NotFound) } - pub async fn update_label( - &self, - account_id: &str, - label: Option, - ) -> Result { - if label.as_ref().is_some_and(|l| l.chars().count() > 64) { - return Err(AdminAccountError::LabelTooLong); - } - let updated = self - .store - .set_label(account_id, label) - .await - .map_err(|_| AdminAccountError::UpdateLabel)?; - if updated { - self.sync_account_pool_best_effort(account_id, "account label update") - .await; - } - Ok(updated) - } - pub async fn update_status( - &self, - account_id: &str, - status: &str, - ) -> Result, AdminAccountError> { - let status = parse_account_status(status)?; - match self.store.set_status(account_id, status).await { - Ok(true) => { - self.sync_account_pool_best_effort(account_id, "account status update") - .await; - self.evict_account_websocket_pool(account_id).await; - Ok(Some(UpdatedAccountStatus { - id: account_id.to_string(), - status, - })) - } - Ok(false) => Ok(None), - Err(_) => Err(AdminAccountError::UpdateStatus), - } - } pub async fn update_account( &self, account_id: &str, @@ -336,39 +292,28 @@ impl AdminAccountService { } match self - .token_refresher - .refresh(refresh_token.expose_secret()) + .refresh_tokens_from_refresh_token(refresh_token.expose_secret()) .await { - Ok(tokens) => { - let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( - crate::upstream::accounts::importing::normalize_bearer_token( - &tokens.access_token, - ), - )) - .ok_or(AdminAccountError::TokenRequired)?; - let claims = crate::upstream::accounts::token_refresh::manual_account_claims( - &access_token, - Utc::now(), - ) - .map_err(AdminAccountError::InvalidToken)?; + Ok(refreshed) => { let updated = self .store .update_from_claims( account_id, crate::upstream::accounts::store::AccountClaimsUpdate { - email: claims.email, - account_id: claims.account_id.or(account.account_id), - user_id: claims.user_id, - plan_type: claims.plan_type, - access_token: SecretString::new(access_token.into()), - refresh_token: tokens + email: refreshed.claims.email, + account_id: refreshed.claims.account_id.or(account.account_id), + user_id: refreshed.claims.user_id, + plan_type: refreshed.claims.plan_type, + access_token: SecretString::new(refreshed.access_token.into()), + refresh_token: refreshed .refresh_token .map(|token| SecretString::new(token.into())), - access_token_expires_at: Some(claims.expires_at), - next_refresh_at: Some( - self.next_refresh_at_for_expires_at(account_id, claims.expires_at), - ), + access_token_expires_at: Some(refreshed.claims.expires_at), + next_refresh_at: Some(self.next_refresh_at_for_expires_at( + account_id, + refreshed.claims.expires_at, + )), status: crate::upstream::accounts::model::AccountStatus::Active, }, ) @@ -387,7 +332,7 @@ impl AdminAccountService { duration_ms: elapsed_millis_i64(started_at), }) } - Err(failure) => { + Err(AdminAccountError::RefreshTokenExchange(failure)) => { if manual_refresh_failure_is_permanent(&failure) { let status = crate::upstream::accounts::model::AccountStatus::Expired; let updated = self @@ -422,6 +367,7 @@ impl AdminAccountService { duration_ms: elapsed_millis_i64(started_at), }) } + Err(error) => Err(error), } } @@ -494,8 +440,7 @@ fn manual_refresh_failure_is_permanent( match failure { crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant | crate::upstream::accounts::token_refresh::RefreshFailure::Banned => true, - crate::upstream::accounts::token_refresh::RefreshFailure::QuotaExhausted - | crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport + crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport | crate::upstream::accounts::token_refresh::RefreshFailure::Transport => false, } } diff --git a/backend/src/admin/accounts/service/mod.rs b/backend/src/admin/accounts/service/mod.rs index 8cf3541b9..1d536dd35 100644 --- a/backend/src/admin/accounts/service/mod.rs +++ b/backend/src/admin/accounts/service/mod.rs @@ -11,7 +11,7 @@ mod testing; use std::sync::Arc as StdArc; -use chrono::DateTime; +use chrono::{DateTime, Utc}; use crate::{ upstream::accounts::{ @@ -27,7 +27,7 @@ use crate::{ pub use contracts::{ AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, AdminAccountRefreshOutcome, AdminAccountRefreshResult, AdminAccountUpdate, BatchDeleteAccounts, ExportedAccounts, - ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, UpdatedAccountStatus, + ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, }; #[derive(Clone)] @@ -82,6 +82,32 @@ impl AdminAccountService { ) } + async fn refresh_tokens_from_refresh_token( + &self, + refresh_token: &str, + ) -> Result { + let token_pair = self + .token_refresher + .refresh(refresh_token) + .await + .map_err(AdminAccountError::RefreshTokenExchange)?; + let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some( + crate::upstream::accounts::importing::normalize_bearer_token(&token_pair.access_token), + )) + .ok_or(AdminAccountError::TokenRequired)?; + let claims = crate::upstream::accounts::token_refresh::manual_account_claims( + &access_token, + Utc::now(), + ) + .map_err(AdminAccountError::InvalidToken)?; + + Ok(contracts::RefreshedAccountTokens { + access_token, + refresh_token: token_pair.refresh_token, + claims, + }) + } + pub(crate) async fn sync_account_pool( &self, account_id: &str, diff --git a/backend/src/admin/monitoring/account_usage_service.rs b/backend/src/admin/monitoring/account_usage_service.rs index f76b7dc1a..860f823b4 100644 --- a/backend/src/admin/monitoring/account_usage_service.rs +++ b/backend/src/admin/monitoring/account_usage_service.rs @@ -37,6 +37,18 @@ impl AdminUsageService { }) } + /// 按账号 ID 批量读取账号用量。 + pub async fn list_by_account_ids( + &self, + account_ids: &[String], + ) -> Result, AdminUsageError> { + self.store + .list_usage_by_account_ids(account_ids) + .await + .map(|records| records.into_iter().map(AdminUsageRecord::from).collect()) + .map_err(|_| AdminUsageError::List) + } + /// 汇总账号用量。 pub async fn summary(&self) -> Result { self.store diff --git a/backend/src/admin/monitoring/account_usage_store.rs b/backend/src/admin/monitoring/account_usage_store.rs index 3893dca97..68978460d 100644 --- a/backend/src/admin/monitoring/account_usage_store.rs +++ b/backend/src/admin/monitoring/account_usage_store.rs @@ -245,6 +245,27 @@ impl SqliteUsageStore { Ok(to_page(&rows, limit)) } + /// 按账号 ID 批量读取账号用量。 + pub async fn list_usage_by_account_ids( + &self, + account_ids: &[String], + ) -> SqliteUsageStoreResult> { + if account_ids.is_empty() { + return Ok(Vec::new()); + } + + let mut builder = QueryBuilder::::new(LIST_USAGE_SELECT_SQL); + builder.push("\nwhere au.account_id in ("); + let mut separated = builder.separated(", "); + for account_id in account_ids { + separated.push_bind(account_id); + } + separated.push_unseparated(")"); + + let rows = builder.build().fetch_all(&self.pool).await?; + rows.iter().map(usage_list_from_row).collect() + } + /// 汇总账号用量。 pub async fn usage_summary(&self) -> SqliteUsageStoreResult { let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?; diff --git a/backend/src/admin/monitoring/usage_record_routes.rs b/backend/src/admin/monitoring/usage_record_routes.rs index 263f3f20e..a0259b072 100644 --- a/backend/src/admin/monitoring/usage_record_routes.rs +++ b/backend/src/admin/monitoring/usage_record_routes.rs @@ -125,7 +125,6 @@ pub(crate) struct UsageRecordCostDetailsData { output_cost: f64, cache_read_cost: f64, total_cost: f64, - billed_cost: f64, original_cost: f64, input_price_per_mtoken: f64, output_price_per_mtoken: f64, @@ -137,7 +136,6 @@ pub(crate) struct UsageRecordCostDetailsData { output_cost_display: String, cache_read_cost_display: String, total_cost_display: String, - billed_cost_display: String, original_cost_display: String, input_price_display: String, output_price_display: String, @@ -332,19 +330,7 @@ pub(crate) async fn usage_records_token_trend( _auth: AdminAuth, Query(query): Query, ) -> Result { - let filter = filter_from_query(query)?; - match state.services.usage_records.token_trend(filter).await { - Ok(points) => Ok(AdminResponse::new( - StatusCode::OK, - AdminEnvelope::ok( - points - .into_iter() - .map(UsageRecordTrendPointData::from) - .collect::>(), - ), - )), - Err(error) => Err(log_error(&error)), - } + usage_records_trend_response(state, query).await } /// `GET /api/admin/usage/records/insights/latency-trend` @@ -353,8 +339,15 @@ pub(crate) async fn usage_records_latency_trend( _auth: AdminAuth, Query(query): Query, ) -> Result { + usage_records_trend_response(state, query).await +} + +async fn usage_records_trend_response( + state: AppState, + query: UsageRecordsQuery, +) -> Result>>, AdminError> { let filter = filter_from_query(query)?; - match state.services.usage_records.latency_trend(filter).await { + match state.services.usage_records.trend(filter).await { Ok(points) => Ok(AdminResponse::new( StatusCode::OK, AdminEnvelope::ok( @@ -607,7 +600,6 @@ fn usage_cost_details( output_cost: cost.output_cost, cache_read_cost: cost.cache_read_cost, total_cost: cost.total_cost, - billed_cost: cost.billed_cost, original_cost: cost.original_cost, input_price_per_mtoken: cost.input_price_per_mtoken, output_price_per_mtoken: cost.output_price_per_mtoken, @@ -619,7 +611,6 @@ fn usage_cost_details( output_cost_display: format_precise_cost(cost.output_cost), cache_read_cost_display: format_precise_cost(cost.cache_read_cost), total_cost_display: format_precise_cost(cost.total_cost), - billed_cost_display: format_precise_cost(cost.billed_cost), original_cost_display: format_precise_cost(cost.original_cost), input_price_display: format_token_price(cost.input_price_per_mtoken), output_price_display: format_token_price(cost.output_price_per_mtoken), diff --git a/backend/src/admin/monitoring/usage_record_service.rs b/backend/src/admin/monitoring/usage_record_service.rs index 16aed423b..4cdb3c180 100644 --- a/backend/src/admin/monitoring/usage_record_service.rs +++ b/backend/src/admin/monitoring/usage_record_service.rs @@ -92,7 +92,6 @@ pub(crate) struct AdminUsageRecordCostDetails { pub output_cost: f64, pub cache_read_cost: f64, pub total_cost: f64, - pub billed_cost: f64, pub original_cost: f64, pub input_price_per_mtoken: f64, pub output_price_per_mtoken: f64, @@ -237,19 +236,8 @@ impl AdminUsageRecordService { .map_err(|_| AdminUsageRecordError::List) } - /// 聚合 Token 趋势。 - pub async fn token_trend( - &self, - filter: AdminUsageRecordFilter, - ) -> Result, AdminUsageRecordError> { - self.store - .trend(filter.into()) - .await - .map_err(|_| AdminUsageRecordError::List) - } - - /// 聚合延迟趋势。 - pub async fn latency_trend( + /// 聚合使用记录趋势。 + pub async fn trend( &self, filter: AdminUsageRecordFilter, ) -> Result, AdminUsageRecordError> { @@ -343,7 +331,6 @@ pub(crate) fn usage_record_cost_details( output_cost: breakdown.output_cost, cache_read_cost: breakdown.cache_read_cost, total_cost: breakdown.total_cost, - billed_cost: breakdown.total_cost, original_cost, input_price_per_mtoken: breakdown.input_price_per_mtoken, output_price_per_mtoken: breakdown.output_price_per_mtoken, diff --git a/backend/src/proxy/dispatch/chat.rs b/backend/src/proxy/dispatch/chat.rs index 9ab73cf0f..3e0c72fab 100644 --- a/backend/src/proxy/dispatch/chat.rs +++ b/backend/src/proxy/dispatch/chat.rs @@ -20,10 +20,14 @@ use crate::{ CloudflareRecovery, }, errors::{ - auth_failure_account_status, backend_transport_name, is_auth_upstream_error, - is_model_unsupported_upstream_error, is_quota_exhausted_upstream_error, - is_rate_limit_upstream_error, rate_limit_cooldown_until, upstream_error_body, - upstream_error_http_status, + auth_failure_account_status, backend_transport_name, insert_dispatch_error_metadata, + is_auth_upstream_error, is_model_unsupported_upstream_error, + is_quota_exhausted_upstream_error, is_rate_limit_upstream_error, + rate_limit_cooldown_until, upstream_error_body, upstream_error_http_status, + DispatchErrorMetadata, DispatchFailureClass, + }, + exhaustion::{ + AccountExhaustionTracker, ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef, }, upstream::{ create_response_with_account, verify_acquired_quota_if_required, @@ -105,26 +109,8 @@ impl ChatDispatchService { .is_some_and(|effort| !effort.trim().is_empty()); let tuple_schema = request.tuple_schema.clone(); let mut excluded_account_ids = Vec::new(); - let mut rate_limited_count = 0usize; - let mut last_rate_limit_error = None; - let mut quota_exhausted_count = 0usize; - let mut last_quota_error = None; - let mut expired_count = 0usize; - let mut last_auth_error = None; - let mut disabled_count = 0usize; - let mut last_disabled_auth_error = None; - let mut banned_count = 0usize; - let mut last_banned_auth_error = None; - let mut last_banned_status_code: Option = None; - let mut cloudflare_challenge_count = 0usize; - let mut last_cloudflare_challenge_error = None; - let mut cloudflare_path_block_count = 0usize; - let mut last_cloudflare_path_block_error = None; - let mut model_unsupported_count = 0usize; - let mut last_model_unsupported_error = None; - let mut model_unsupported_retry_used = false; + let mut exhausted_accounts = AccountExhaustionTracker::default(); let mut quota_verify_attempts = 0usize; - let mut last_failed_account_id = None; macro_rules! return_dispatch_error { ($error:expr) => {{ @@ -133,7 +119,7 @@ impl ChatDispatchService { request_id, requested_model, started_at, - last_failed_account_id.as_deref(), + exhausted_accounts.last_account_id(), Some(backend_transport_name( backend_transport_for_response_request(&request), )), @@ -162,56 +148,13 @@ impl ChatDispatchService { .with_exclude_account_ids(excluded_account_ids.iter().cloned()); let acquired = match self.account_pool.acquire_with(&acquire_request).await { Some(acquired) => acquired, - None if quota_exhausted_count > 0 => { - return_dispatch_error!(ChatDispatchError::QuotaExhausted { - count: quota_exhausted_count, - upstream_error: last_quota_error.unwrap_or_default(), - }); + None => { + let error = exhausted_accounts + .last_exhausted() + .map(ChatDispatchError::from_exhausted_account) + .unwrap_or(ChatDispatchError::NoActiveAccount); + return_dispatch_error!(error); } - None if rate_limited_count > 0 => { - return_dispatch_error!(ChatDispatchError::RateLimited { - count: rate_limited_count, - upstream_error: last_rate_limit_error.unwrap_or_default(), - }); - } - None if expired_count > 0 => { - return_dispatch_error!(ChatDispatchError::Expired { - count: expired_count, - upstream_error: last_auth_error.unwrap_or_default(), - }); - } - None if disabled_count > 0 => { - return_dispatch_error!(ChatDispatchError::Disabled { - count: disabled_count, - upstream_error: last_disabled_auth_error.unwrap_or_default(), - }); - } - None if banned_count > 0 => { - return_dispatch_error!(ChatDispatchError::Banned { - count: banned_count, - upstream_error: last_banned_auth_error.unwrap_or_default(), - status_code: last_banned_status_code.unwrap_or(403), - }); - } - None if cloudflare_challenge_count > 0 => { - return_dispatch_error!(ChatDispatchError::CloudflareChallenge { - count: cloudflare_challenge_count, - upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), - }); - } - None if cloudflare_path_block_count > 0 => { - return_dispatch_error!(ChatDispatchError::CloudflarePathBlocked { - count: cloudflare_path_block_count, - upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), - }); - } - None if model_unsupported_count > 0 => { - return_dispatch_error!(ChatDispatchError::ModelUnsupported { - count: model_unsupported_count, - upstream_error: last_model_unsupported_error.unwrap_or_default(), - }); - } - None => return_dispatch_error!(ChatDispatchError::NoActiveAccount), }; let acquired = match verify_acquired_quota_if_required( QuotaVerificationContext { @@ -229,15 +172,18 @@ impl ChatDispatchService { { QuotaVerificationDecision::Ready(acquired) => *acquired, QuotaVerificationDecision::RetryWithAnotherAccount => { - rate_limited_count += 1; - last_rate_limit_error = Some(QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string()); + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); continue; } QuotaVerificationDecision::MaxAttemptsReached => { - return_dispatch_error!(ChatDispatchError::RateLimited { - count: rate_limited_count + 1, - upstream_error: QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string(), - }); + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); + let error = exhausted_accounts + .last_exhausted() + .map(ChatDispatchError::from_exhausted_account) + .unwrap_or(ChatDispatchError::NoActiveAccount); + return_dispatch_error!(error); } }; self.account_pool.wait_for_request_interval(&acquired).await; @@ -258,9 +204,10 @@ impl ChatDispatchService { match response_result { Ok(response) => break (account, response), Err(error) if is_rate_limit_upstream_error(&error) => { - rate_limited_count += 1; - last_rate_limit_error = Some(upstream_error_body(&error)); - last_failed_account_id = Some(release_account_id.clone()); + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); self.account_pool .mark_quota_limited_until(&release_account_id, cooldown_until) @@ -268,9 +215,10 @@ impl ChatDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_quota_exhausted_upstream_error(&error) => { - quota_exhausted_count += 1; - last_quota_error = Some(upstream_error_body(&error)); - last_failed_account_id = Some(release_account_id.clone()); + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); self.account_pool .set_status(&release_account_id, AccountStatus::QuotaExhausted) .await; @@ -279,22 +227,12 @@ impl ChatDispatchService { Err(error) if is_auth_upstream_error(&error) => { let upstream_error = upstream_error_body(&error); let account_status = auth_failure_account_status(&error); - last_failed_account_id = Some(release_account_id.clone()); - match account_status { - AccountStatus::Disabled => { - disabled_count += 1; - last_disabled_auth_error = Some(upstream_error); - } - AccountStatus::Banned => { - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error); - } - _ => { - expired_count += 1; - last_auth_error = Some(upstream_error); - } - } + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, account_status) .await; @@ -306,20 +244,20 @@ impl ChatDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { - cloudflare_challenge_count += 1; - last_cloudflare_challenge_error = - Some(cloudflare_challenge_error_message().to_string()); - last_failed_account_id = Some(release_account_id.clone()); + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); self.cloudflare .apply_challenge(self.account_pool.as_ref(), &release_account_id) .await; excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_path_block_upstream_error(&error) => { - cloudflare_path_block_count += 1; - last_cloudflare_path_block_error = - Some(cloudflare_path_block_error_message().to_string()); - last_failed_account_id = Some(release_account_id.clone()); + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); self.cloudflare .apply_path_block(self.account_pool.as_ref(), &release_account_id) .await; @@ -327,29 +265,28 @@ impl ChatDispatchService { } Err(error) if is_model_unsupported_upstream_error(&error) => { let upstream_error = upstream_error_body(&error); - if model_unsupported_retry_used { + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { return_dispatch_error!( - ChatDispatchError::ModelUnsupported { - count: model_unsupported_count + 1, - upstream_error, - }, + ChatDispatchError::from_exhausted_account(exhausted), account_id: Some(&release_account_id), transport: Some(backend_transport_name(backend_transport_for_response_request( &request ))) ); } - model_unsupported_count += 1; - last_model_unsupported_error = Some(upstream_error); - last_failed_account_id = Some(release_account_id.clone()); - model_unsupported_retry_used = true; + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); excluded_account_ids.push(release_account_id); } Err(error) if is_banned_upstream_error(&error) => { - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error_body(&error)); - last_failed_account_id = Some(release_account_id.clone()); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, AccountStatus::Banned) .await; @@ -459,8 +396,6 @@ impl ChatDispatchService { /// Chat Completions 调度错误。 #[derive(Debug, Error)] pub enum ChatDispatchError { - #[error("failed to list runtime accounts")] - AccountStore, #[error("no active account is available")] NoActiveAccount, #[error("all accounts exhausted by quota")] @@ -513,9 +448,47 @@ pub enum ChatDispatchError { } impl ChatDispatchError { + pub(crate) fn from_exhausted_account(exhausted: ExhaustedAccount) -> Self { + match exhausted.kind { + ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::RateLimited => Self::RateLimited { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::Expired => Self::Expired { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::Disabled => Self::Disabled { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::Banned => Self::Banned { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + status_code: exhausted.status_code.unwrap_or(403), + }, + ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + } + } + pub fn http_status_code(&self) -> u16 { match self { - Self::NoActiveAccount | Self::AccountStore => 503, + Self::NoActiveAccount => 503, Self::QuotaExhausted { .. } => 429, Self::RateLimited { .. } => 429, Self::Expired { .. } | Self::Disabled { .. } => 401, @@ -528,6 +501,151 @@ impl ChatDispatchError { Self::Upstream(error) => upstream_error_http_status(error), } } + + pub(crate) fn metadata(&self) -> DispatchErrorMetadata { + match self { + Self::NoActiveAccount => DispatchErrorMetadata::no_available_accounts(), + Self::QuotaExhausted { + count, + upstream_error, + } => Self::exhausted_metadata( + ExhaustedAccountKind::QuotaExhausted, + *count, + upstream_error, + ), + Self::RateLimited { + count, + upstream_error, + } => { + Self::exhausted_metadata(ExhaustedAccountKind::RateLimited, *count, upstream_error) + } + Self::Expired { + count, + upstream_error, + } => Self::exhausted_metadata(ExhaustedAccountKind::Expired, *count, upstream_error), + Self::Disabled { + count, + upstream_error, + } => Self::exhausted_metadata(ExhaustedAccountKind::Disabled, *count, upstream_error), + Self::Banned { + count, + upstream_error, + .. + } => Self::exhausted_metadata(ExhaustedAccountKind::Banned, *count, upstream_error), + Self::CloudflareChallenge { + count, + upstream_error, + } => Self::exhausted_metadata( + ExhaustedAccountKind::CloudflareChallenge, + *count, + upstream_error, + ), + Self::CloudflarePathBlocked { + count, + upstream_error, + } => Self::exhausted_metadata( + ExhaustedAccountKind::CloudflarePathBlocked, + *count, + upstream_error, + ), + Self::ModelUnsupported { + count, + upstream_error, + } => Self::exhausted_metadata( + ExhaustedAccountKind::ModelUnsupported, + *count, + upstream_error, + ), + Self::Upstream(error) => DispatchErrorMetadata::upstream(error), + Self::InvalidSse(_) => DispatchErrorMetadata::simple(DispatchFailureClass::InvalidSse), + Self::EmptyUpstreamResponse => { + DispatchErrorMetadata::simple(DispatchFailureClass::EmptyUpstreamResponse) + } + } + } + + pub(crate) fn exhausted_account(&self) -> Option> { + match self { + Self::QuotaExhausted { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::QuotaExhausted, + count: *count, + upstream_error, + }), + Self::RateLimited { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::RateLimited, + count: *count, + upstream_error, + }), + Self::Expired { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::Expired, + count: *count, + upstream_error, + }), + Self::Disabled { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::Disabled, + count: *count, + upstream_error, + }), + Self::Banned { + count, + upstream_error, + .. + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::Banned, + count: *count, + upstream_error, + }), + Self::CloudflareChallenge { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::CloudflareChallenge, + count: *count, + upstream_error, + }), + Self::CloudflarePathBlocked { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::CloudflarePathBlocked, + count: *count, + upstream_error, + }), + Self::ModelUnsupported { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::ModelUnsupported, + count: *count, + upstream_error, + }), + _ => None, + } + } + + fn exhausted_metadata( + kind: ExhaustedAccountKind, + count: usize, + upstream_error: &str, + ) -> DispatchErrorMetadata { + DispatchErrorMetadata::exhausted_ref(ExhaustedAccountRef { + kind, + count, + upstream_error, + }) + } } // ==================================================================== @@ -555,7 +673,7 @@ async fn record_chat_dispatch_error_event(record: ChatDispatchErrorEventRecord<' "error": record.error.to_string(), }); if let Some(object) = metadata.as_object_mut() { - enrich_chat_error_metadata(object, record.error); + insert_dispatch_error_metadata(object, record.error.metadata()); if let Some(transport) = record.transport { object.insert("transport".to_string(), serde_json::json!(transport)); } @@ -573,99 +691,3 @@ async fn record_chat_dispatch_error_event(record: ChatDispatchErrorEventRecord<' }) .await; } - -fn enrich_chat_error_metadata( - object: &mut serde_json::Map, - error: &ChatDispatchError, -) { - let (failure_class, exhausted_count, upstream_error, upstream_status) = match error { - ChatDispatchError::AccountStore => ("account_store", None, None, None), - ChatDispatchError::NoActiveAccount => ("no_available_accounts", None, None, None), - ChatDispatchError::QuotaExhausted { - count, - upstream_error, - } => ( - "quota_exhausted", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ChatDispatchError::RateLimited { - count, - upstream_error, - } => ( - "rate_limited", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ChatDispatchError::Expired { - count, - upstream_error, - } => ("expired", Some(*count), Some(upstream_error.clone()), None), - ChatDispatchError::Disabled { - count, - upstream_error, - } => ("disabled", Some(*count), Some(upstream_error.clone()), None), - ChatDispatchError::Banned { - count, - upstream_error, - .. - } => ("banned", Some(*count), Some(upstream_error.clone()), None), - ChatDispatchError::CloudflareChallenge { - count, - upstream_error, - } => ( - "cloudflare_challenge", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ChatDispatchError::CloudflarePathBlocked { - count, - upstream_error, - } => ( - "cloudflare_path_blocked", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ChatDispatchError::ModelUnsupported { - count, - upstream_error, - } => ( - "model_unsupported", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ChatDispatchError::Upstream(error) => { - let upstream_status = match error { - CodexClientError::Upstream { status, .. } => Some(status.as_u16()), - _ => None, - }; - ( - "upstream", - None, - Some(upstream_error_body(error)), - upstream_status, - ) - } - ChatDispatchError::InvalidSse(_) => ("invalid_sse", None, None, None), - ChatDispatchError::EmptyUpstreamResponse => ("empty_upstream_response", None, None, None), - }; - - object.insert( - "failureClass".to_string(), - Value::String(failure_class.to_string()), - ); - if let Some(count) = exhausted_count { - object.insert("exhaustedCount".to_string(), serde_json::json!(count)); - } - if let Some(error) = upstream_error { - object.insert("upstreamError".to_string(), Value::String(error)); - } - if let Some(status) = upstream_status { - object.insert("upstreamStatus".to_string(), serde_json::json!(status)); - } -} diff --git a/backend/src/proxy/dispatch/errors.rs b/backend/src/proxy/dispatch/errors.rs index 5a3cc151c..16c192bba 100644 --- a/backend/src/proxy/dispatch/errors.rs +++ b/backend/src/proxy/dispatch/errors.rs @@ -1,12 +1,136 @@ //! Shared upstream error classification for dispatch routes. use chrono::{DateTime, Duration, Utc}; +use serde_json::{json, Map, Value}; use crate::{ + proxy::dispatch::exhaustion::{ExhaustedAccountKind, ExhaustedAccountRef}, upstream::accounts::model::AccountStatus, upstream::transport::{is_banned_auth_signal, CodexBackendTransport, CodexClientError}, }; +#[derive(Clone, Debug)] +pub(crate) struct DispatchErrorMetadata { + pub failure_class: DispatchFailureClass, + pub exhausted_count: Option, + pub upstream_error: Option, + pub upstream_status: Option, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) enum DispatchFailureClass { + NoAvailableAccounts, + QuotaExhausted, + RateLimited, + Expired, + Disabled, + Banned, + CloudflareChallenge, + CloudflarePathBlocked, + ModelUnsupported, + Upstream, + InvalidSse, + MissingCompleted, + EmptyUpstreamResponse, + ResponseFailed, +} + +impl DispatchFailureClass { + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::NoAvailableAccounts => "no_available_accounts", + Self::QuotaExhausted => "quota_exhausted", + Self::RateLimited => "rate_limited", + Self::Expired => "expired", + Self::Disabled => "disabled", + Self::Banned => "banned", + Self::CloudflareChallenge => "cloudflare_challenge", + Self::CloudflarePathBlocked => "cloudflare_path_blocked", + Self::ModelUnsupported => "model_unsupported", + Self::Upstream => "upstream", + Self::InvalidSse => "invalid_sse", + Self::MissingCompleted => "missing_completed", + Self::EmptyUpstreamResponse => "empty_upstream_response", + Self::ResponseFailed => "response_failed", + } + } +} + +impl DispatchErrorMetadata { + pub(crate) fn no_available_accounts() -> Self { + Self { + failure_class: DispatchFailureClass::NoAvailableAccounts, + exhausted_count: None, + upstream_error: None, + upstream_status: None, + } + } + + pub(crate) fn exhausted_ref(exhausted: ExhaustedAccountRef<'_>) -> Self { + Self { + failure_class: DispatchFailureClass::from(exhausted.kind), + exhausted_count: Some(exhausted.count), + upstream_error: Some(exhausted.upstream_error.to_string()), + upstream_status: None, + } + } + + pub(crate) fn upstream(error: &CodexClientError) -> Self { + Self { + failure_class: DispatchFailureClass::Upstream, + exhausted_count: None, + upstream_error: Some(upstream_error_body(error)), + upstream_status: match error { + CodexClientError::Upstream { status, .. } => Some(status.as_u16()), + _ => None, + }, + } + } + + pub(crate) fn simple(failure_class: DispatchFailureClass) -> Self { + Self { + failure_class, + exhausted_count: None, + upstream_error: None, + upstream_status: None, + } + } +} + +impl From for DispatchFailureClass { + fn from(kind: ExhaustedAccountKind) -> Self { + match kind { + ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted, + ExhaustedAccountKind::RateLimited => Self::RateLimited, + ExhaustedAccountKind::Expired => Self::Expired, + ExhaustedAccountKind::Disabled => Self::Disabled, + ExhaustedAccountKind::Banned => Self::Banned, + ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge, + ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked, + ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported, + } + } +} + +pub(crate) fn insert_dispatch_error_metadata( + object: &mut Map, + metadata: DispatchErrorMetadata, +) { + object.insert( + "failureClass".to_string(), + Value::String(metadata.failure_class.as_str().to_string()), + ); + if let Some(count) = metadata.exhausted_count { + object.insert("exhaustedCount".to_string(), json!(count)); + } + if let Some(error) = metadata.upstream_error { + object.insert("upstreamError".to_string(), Value::String(error)); + } + if let Some(status) = metadata.upstream_status { + object.insert("upstreamStatus".to_string(), json!(status)); + } +} + pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool { matches!( error, diff --git a/backend/src/proxy/dispatch/exhaustion.rs b/backend/src/proxy/dispatch/exhaustion.rs new file mode 100644 index 000000000..d89a1b524 --- /dev/null +++ b/backend/src/proxy/dispatch/exhaustion.rs @@ -0,0 +1,222 @@ +//! 账号耗尽状态聚合。 + +use crate::upstream::accounts::model::AccountStatus; + +#[derive(Clone, Copy, Debug)] +pub(crate) enum ExhaustedAccountKind { + QuotaExhausted, + RateLimited, + Expired, + Disabled, + Banned, + CloudflareChallenge, + CloudflarePathBlocked, + ModelUnsupported, +} + +impl ExhaustedAccountKind { + pub(crate) fn message_reason(self) -> &'static str { + match self { + Self::QuotaExhausted => "quota-exhausted", + Self::RateLimited => "rate-limited", + Self::Expired => "expired", + Self::Disabled => "disabled", + Self::Banned => "banned", + Self::CloudflareChallenge => "cloudflare-challenge", + Self::CloudflarePathBlocked => "cloudflare-path-block", + Self::ModelUnsupported => "model-unsupported", + } + } +} + +#[derive(Clone, Debug)] +pub(crate) struct ExhaustedAccount { + pub kind: ExhaustedAccountKind, + pub count: usize, + pub upstream_error: String, + pub status_code: Option, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct ExhaustedAccountRef<'a> { + pub kind: ExhaustedAccountKind, + pub count: usize, + pub upstream_error: &'a str, +} + +#[derive(Default)] +pub(crate) struct AccountExhaustionTracker { + last_kind: Option, + last_account_id: Option, + quota_exhausted: ExhaustedCounter, + rate_limited: ExhaustedCounter, + expired: ExhaustedCounter, + disabled: ExhaustedCounter, + banned: ExhaustedCounter, + cloudflare_challenge: ExhaustedCounter, + cloudflare_path_blocked: ExhaustedCounter, + model_unsupported: ExhaustedCounter, + model_unsupported_retry_used: bool, +} + +#[derive(Default)] +struct ExhaustedCounter { + count: usize, + upstream_error: Option, + status_code: Option, +} + +impl AccountExhaustionTracker { + pub(crate) fn last_account_id(&self) -> Option<&str> { + self.last_account_id.as_deref() + } + + pub(crate) fn record_rate_limited( + &mut self, + account_id: Option<&str>, + upstream_error: impl Into, + ) { + self.record( + account_id, + ExhaustedAccountKind::RateLimited, + upstream_error, + None, + ); + } + + pub(crate) fn record_quota_exhausted( + &mut self, + account_id: Option<&str>, + upstream_error: impl Into, + ) { + self.record( + account_id, + ExhaustedAccountKind::QuotaExhausted, + upstream_error, + None, + ); + } + + pub(crate) fn record_auth_failure( + &mut self, + account_id: Option<&str>, + account_status: AccountStatus, + upstream_error: impl Into, + status_code: Option, + ) { + let kind = match account_status { + AccountStatus::Disabled => ExhaustedAccountKind::Disabled, + AccountStatus::Banned => ExhaustedAccountKind::Banned, + _ => ExhaustedAccountKind::Expired, + }; + self.record(account_id, kind, upstream_error, status_code); + } + + pub(crate) fn record_cloudflare_challenge( + &mut self, + account_id: Option<&str>, + upstream_error: impl Into, + ) { + self.record( + account_id, + ExhaustedAccountKind::CloudflareChallenge, + upstream_error, + None, + ); + } + + pub(crate) fn record_cloudflare_path_blocked( + &mut self, + account_id: Option<&str>, + upstream_error: impl Into, + ) { + self.record( + account_id, + ExhaustedAccountKind::CloudflarePathBlocked, + upstream_error, + None, + ); + } + + pub(crate) fn record_model_unsupported( + &mut self, + account_id: Option<&str>, + upstream_error: impl Into, + ) { + self.model_unsupported_retry_used = true; + self.record( + account_id, + ExhaustedAccountKind::ModelUnsupported, + upstream_error, + None, + ); + } + + pub(crate) fn model_unsupported_retry_exhausted( + &self, + upstream_error: impl Into, + ) -> Option { + self.model_unsupported_retry_used.then(|| ExhaustedAccount { + kind: ExhaustedAccountKind::ModelUnsupported, + count: self.model_unsupported.count + 1, + upstream_error: upstream_error.into(), + status_code: None, + }) + } + + pub(crate) fn last_exhausted(&self) -> Option { + let kind = self.last_kind?; + let counter = self.counter(kind); + Some(ExhaustedAccount { + kind, + count: counter.count, + upstream_error: counter.upstream_error.clone().unwrap_or_default(), + status_code: counter.status_code, + }) + } + + fn record( + &mut self, + account_id: Option<&str>, + kind: ExhaustedAccountKind, + upstream_error: impl Into, + status_code: Option, + ) { + self.last_kind = Some(kind); + if let Some(account_id) = account_id { + self.last_account_id = Some(account_id.to_string()); + } + let counter = self.counter_mut(kind); + counter.count += 1; + counter.upstream_error = Some(upstream_error.into()); + if let Some(status_code) = status_code { + counter.status_code = Some(status_code); + } + } + + fn counter(&self, kind: ExhaustedAccountKind) -> &ExhaustedCounter { + match kind { + ExhaustedAccountKind::QuotaExhausted => &self.quota_exhausted, + ExhaustedAccountKind::RateLimited => &self.rate_limited, + ExhaustedAccountKind::Expired => &self.expired, + ExhaustedAccountKind::Disabled => &self.disabled, + ExhaustedAccountKind::Banned => &self.banned, + ExhaustedAccountKind::CloudflareChallenge => &self.cloudflare_challenge, + ExhaustedAccountKind::CloudflarePathBlocked => &self.cloudflare_path_blocked, + ExhaustedAccountKind::ModelUnsupported => &self.model_unsupported, + } + } + + fn counter_mut(&mut self, kind: ExhaustedAccountKind) -> &mut ExhaustedCounter { + match kind { + ExhaustedAccountKind::QuotaExhausted => &mut self.quota_exhausted, + ExhaustedAccountKind::RateLimited => &mut self.rate_limited, + ExhaustedAccountKind::Expired => &mut self.expired, + ExhaustedAccountKind::Disabled => &mut self.disabled, + ExhaustedAccountKind::Banned => &mut self.banned, + ExhaustedAccountKind::CloudflareChallenge => &mut self.cloudflare_challenge, + ExhaustedAccountKind::CloudflarePathBlocked => &mut self.cloudflare_path_blocked, + ExhaustedAccountKind::ModelUnsupported => &mut self.model_unsupported, + } + } +} diff --git a/backend/src/proxy/dispatch/mod.rs b/backend/src/proxy/dispatch/mod.rs index fb3046557..0a6c634ae 100644 --- a/backend/src/proxy/dispatch/mod.rs +++ b/backend/src/proxy/dispatch/mod.rs @@ -4,6 +4,7 @@ mod auth_recovery; pub mod chat; pub mod cloudflare; pub mod errors; +mod exhaustion; pub mod implicit_resume; pub mod reasoning_replay; pub mod responses; diff --git a/backend/src/proxy/dispatch/responses/errors.rs b/backend/src/proxy/dispatch/responses/errors.rs index b8389a305..abfbd623b 100644 --- a/backend/src/proxy/dispatch/responses/errors.rs +++ b/backend/src/proxy/dispatch/responses/errors.rs @@ -1,8 +1,14 @@ -use serde_json::{json, Value}; +use serde_json::Value; use thiserror::Error; use crate::{ - proxy::dispatch::errors::{upstream_error_body, upstream_error_http_status}, + proxy::dispatch::{ + errors::{ + insert_dispatch_error_metadata, upstream_error_http_status, DispatchErrorMetadata, + DispatchFailureClass, + }, + exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef}, + }, upstream::{ protocol::{responses::ResponsesSseFailure, sse::SseError}, transport::CodexClientError, @@ -14,8 +20,6 @@ use super::sse_failure::{sse_failure_error_body, stream_failure_http_status}; /// Responses 调度错误。 #[derive(Debug, Error)] pub enum ResponseDispatchError { - #[error("failed to list runtime accounts")] - AccountStore, #[error("no active account is available")] NoActiveAccount, #[error("all accounts exhausted by quota")] @@ -72,9 +76,47 @@ pub enum ResponseDispatchError { } impl ResponseDispatchError { + pub(crate) fn from_exhausted_account(exhausted: ExhaustedAccount) -> Self { + match exhausted.kind { + ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::RateLimited => Self::RateLimited { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::Expired => Self::Expired { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::Disabled => Self::Disabled { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::Banned => Self::Banned { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + status_code: exhausted.status_code.unwrap_or(403), + }, + ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported { + count: exhausted.count, + upstream_error: exhausted.upstream_error, + }, + } + } + pub fn http_status_code(&self) -> u16 { match self { - Self::NoActiveAccount | Self::AccountStore => 503, + Self::NoActiveAccount => 503, Self::QuotaExhausted { .. } => 429, Self::RateLimited { .. } => 429, Self::Expired { .. } | Self::Disabled { .. } => 401, @@ -89,6 +131,160 @@ impl ResponseDispatchError { Self::Upstream(error) => upstream_error_http_status(error), } } + + pub(crate) fn metadata(&self) -> DispatchErrorMetadata { + match self { + Self::NoActiveAccount => DispatchErrorMetadata::no_available_accounts(), + Self::QuotaExhausted { + count, + upstream_error, + } => Self::exhausted_metadata( + ExhaustedAccountKind::QuotaExhausted, + *count, + upstream_error, + ), + Self::RateLimited { + count, + upstream_error, + } => { + Self::exhausted_metadata(ExhaustedAccountKind::RateLimited, *count, upstream_error) + } + Self::Expired { + count, + upstream_error, + } => Self::exhausted_metadata(ExhaustedAccountKind::Expired, *count, upstream_error), + Self::Disabled { + count, + upstream_error, + } => Self::exhausted_metadata(ExhaustedAccountKind::Disabled, *count, upstream_error), + Self::Banned { + count, + upstream_error, + .. + } => Self::exhausted_metadata(ExhaustedAccountKind::Banned, *count, upstream_error), + Self::CloudflareChallenge { + count, + upstream_error, + } => Self::exhausted_metadata( + ExhaustedAccountKind::CloudflareChallenge, + *count, + upstream_error, + ), + Self::CloudflarePathBlocked { + count, + upstream_error, + } => Self::exhausted_metadata( + ExhaustedAccountKind::CloudflarePathBlocked, + *count, + upstream_error, + ), + Self::ModelUnsupported { + count, + upstream_error, + } => Self::exhausted_metadata( + ExhaustedAccountKind::ModelUnsupported, + *count, + upstream_error, + ), + Self::Upstream(error) => DispatchErrorMetadata::upstream(error), + Self::InvalidSse(_) => DispatchErrorMetadata::simple(DispatchFailureClass::InvalidSse), + Self::MissingCompleted => { + DispatchErrorMetadata::simple(DispatchFailureClass::MissingCompleted) + } + Self::EmptyUpstreamResponse => { + DispatchErrorMetadata::simple(DispatchFailureClass::EmptyUpstreamResponse) + } + Self::Failed(failure) => DispatchErrorMetadata { + failure_class: DispatchFailureClass::ResponseFailed, + exhausted_count: None, + upstream_error: Some(sse_failure_error_body(failure)), + upstream_status: None, + }, + } + } + + pub(crate) fn exhausted_account(&self) -> Option> { + match self { + Self::QuotaExhausted { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::QuotaExhausted, + count: *count, + upstream_error, + }), + Self::RateLimited { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::RateLimited, + count: *count, + upstream_error, + }), + Self::Expired { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::Expired, + count: *count, + upstream_error, + }), + Self::Disabled { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::Disabled, + count: *count, + upstream_error, + }), + Self::Banned { + count, + upstream_error, + .. + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::Banned, + count: *count, + upstream_error, + }), + Self::CloudflareChallenge { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::CloudflareChallenge, + count: *count, + upstream_error, + }), + Self::CloudflarePathBlocked { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::CloudflarePathBlocked, + count: *count, + upstream_error, + }), + Self::ModelUnsupported { + count, + upstream_error, + } => Some(ExhaustedAccountRef { + kind: ExhaustedAccountKind::ModelUnsupported, + count: *count, + upstream_error, + }), + _ => None, + } + } + + fn exhausted_metadata( + kind: ExhaustedAccountKind, + count: usize, + upstream_error: &str, + ) -> DispatchErrorMetadata { + DispatchErrorMetadata::exhausted_ref(ExhaustedAccountRef { + kind, + count, + upstream_error, + }) + } } /// Responses live SSE body stream error. @@ -132,103 +328,5 @@ pub(super) fn enrich_response_dispatch_error_metadata( let Some(object) = metadata.as_object_mut() else { return; }; - let (failure_class, exhausted_count, upstream_error, upstream_status) = match error { - ResponseDispatchError::AccountStore => ("account_store", None, None, None), - ResponseDispatchError::NoActiveAccount => ("no_available_accounts", None, None, None), - ResponseDispatchError::QuotaExhausted { - count, - upstream_error, - } => ( - "quota_exhausted", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ResponseDispatchError::RateLimited { - count, - upstream_error, - } => ( - "rate_limited", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ResponseDispatchError::Expired { - count, - upstream_error, - } => ("expired", Some(*count), Some(upstream_error.clone()), None), - ResponseDispatchError::Disabled { - count, - upstream_error, - } => ("disabled", Some(*count), Some(upstream_error.clone()), None), - ResponseDispatchError::Banned { - count, - upstream_error, - .. - } => ("banned", Some(*count), Some(upstream_error.clone()), None), - ResponseDispatchError::CloudflareChallenge { - count, - upstream_error, - } => ( - "cloudflare_challenge", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ResponseDispatchError::CloudflarePathBlocked { - count, - upstream_error, - } => ( - "cloudflare_path_blocked", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ResponseDispatchError::ModelUnsupported { - count, - upstream_error, - } => ( - "model_unsupported", - Some(*count), - Some(upstream_error.clone()), - None, - ), - ResponseDispatchError::Upstream(error) => { - let upstream_status = match error { - CodexClientError::Upstream { status, .. } => Some(status.as_u16()), - _ => None, - }; - ( - "upstream", - None, - Some(upstream_error_body(error)), - upstream_status, - ) - } - ResponseDispatchError::InvalidSse(_) => ("invalid_sse", None, None, None), - ResponseDispatchError::MissingCompleted => ("missing_completed", None, None, None), - ResponseDispatchError::EmptyUpstreamResponse => { - ("empty_upstream_response", None, None, None) - } - ResponseDispatchError::Failed(failure) => ( - "response_failed", - None, - Some(sse_failure_error_body(failure)), - None, - ), - }; - - object.insert( - "failureClass".to_string(), - Value::String(failure_class.to_string()), - ); - if let Some(count) = exhausted_count { - object.insert("exhaustedCount".to_string(), json!(count)); - } - if let Some(error) = upstream_error { - object.insert("upstreamError".to_string(), Value::String(error)); - } - if let Some(status) = upstream_status { - object.insert("upstreamStatus".to_string(), json!(status)); - } + insert_dispatch_error_metadata(object, error.metadata()); } diff --git a/backend/src/proxy/dispatch/responses/service.rs b/backend/src/proxy/dispatch/responses/service.rs index 6fb4d2164..be259f447 100644 --- a/backend/src/proxy/dispatch/responses/service.rs +++ b/backend/src/proxy/dispatch/responses/service.rs @@ -34,6 +34,7 @@ use crate::{ rate_limit_cooldown_until, upstream_error_body, upstream_error_http_status, upstream_error_set_cookie_headers, }, + exhaustion::AccountExhaustionTracker, reasoning_replay::ReasoningReplayCache, session_affinity::{ compute_variant_hash, ensure_prompt_cache_key, hash_instructions, @@ -97,18 +98,6 @@ use super::{ stream_lifecycle::LiveResponseStreamContext, }; -#[derive(Clone, Copy)] -enum ExhaustedAccountClass { - QuotaExhausted, - RateLimited, - Expired, - Disabled, - Banned, - CloudflareChallenge, - CloudflarePathBlocked, - ModelUnsupported, -} - /// OpenAI Responses 调度服务。 #[derive(Clone)] pub struct ResponseDispatchService { @@ -394,29 +383,10 @@ impl ResponseDispatchService { acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); } let mut excluded_account_ids = Vec::new(); - let mut rate_limited_count = 0usize; - let mut last_rate_limit_error = None; - let mut quota_exhausted_count = 0usize; - let mut last_quota_error = None; - let mut expired_count = 0usize; - let mut last_auth_error = None; - let mut disabled_count = 0usize; - let mut last_disabled_auth_error = None; - let mut banned_count = 0usize; - let mut last_banned_auth_error = None; - let mut last_banned_status_code: Option = None; - let mut cloudflare_challenge_count = 0usize; - let mut last_cloudflare_challenge_error = None; - let mut cloudflare_path_block_count = 0usize; - let mut last_cloudflare_path_block_error = None; - let mut model_unsupported_count = 0usize; - let mut last_model_unsupported_error = None; - let mut model_unsupported_retry_used = false; + let mut exhausted_accounts = AccountExhaustionTracker::default(); let mut history_recovery_used = false; - let mut last_exhausted_account_class = None; let mut empty_response_retries = 0u8; let mut quota_verify_attempts = 0usize; - let mut last_attempted_account_id = None; const MAX_EMPTY_RESPONSE_RETRIES: u8 = 2; let (account, response, collected_response): ( Account, @@ -432,59 +402,17 @@ impl ResponseDispatchService { .acquire_with(&attempt_acquire_request) .await else { - let error = match last_exhausted_account_class { - Some(ExhaustedAccountClass::QuotaExhausted) => { - ResponseDispatchError::QuotaExhausted { - count: quota_exhausted_count, - upstream_error: last_quota_error.unwrap_or_default(), - } - } - Some(ExhaustedAccountClass::RateLimited) => { - ResponseDispatchError::RateLimited { - count: rate_limited_count, - upstream_error: last_rate_limit_error.unwrap_or_default(), - } - } - Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired { - count: expired_count, - upstream_error: last_auth_error.unwrap_or_default(), - }, - Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled { - count: disabled_count, - upstream_error: last_disabled_auth_error.unwrap_or_default(), - }, - Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned { - count: banned_count, - upstream_error: last_banned_auth_error.unwrap_or_default(), - status_code: last_banned_status_code.unwrap_or(403), - }, - Some(ExhaustedAccountClass::CloudflareChallenge) => { - ResponseDispatchError::CloudflareChallenge { - count: cloudflare_challenge_count, - upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), - } - } - Some(ExhaustedAccountClass::CloudflarePathBlocked) => { - ResponseDispatchError::CloudflarePathBlocked { - count: cloudflare_path_block_count, - upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), - } - } - Some(ExhaustedAccountClass::ModelUnsupported) => { - ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count, - upstream_error: last_model_unsupported_error.unwrap_or_default(), - } - } - None => ResponseDispatchError::NoActiveAccount, - }; + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); self.record_response_dispatch_error( request_id, route, requested_model, started_at, ResponseDispatchErrorDetails { - account_id: last_attempted_account_id.as_deref(), + account_id: exhausted_accounts.last_account_id(), stream: false, compact: false, transport: Some(backend_transport_name( @@ -515,16 +443,19 @@ impl ResponseDispatchService { { QuotaVerificationDecision::Ready(acquired) => *acquired, QuotaVerificationDecision::RetryWithAnotherAccount => { - rate_limited_count += 1; - last_rate_limit_error = Some(QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string()); - last_exhausted_account_class = Some(ExhaustedAccountClass::RateLimited); + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); continue; } QuotaVerificationDecision::MaxAttemptsReached => { - let error = ResponseDispatchError::RateLimited { - count: rate_limited_count + 1, - upstream_error: QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string(), - }; + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); self.record_response_dispatch_error( request_id, route, @@ -556,7 +487,6 @@ impl ResponseDispatchService { self.account_pool.wait_for_request_interval(&acquired).await; let account = acquired.account; let release_account_id = account.id.clone(); - last_attempted_account_id = Some(release_account_id.clone()); let response_result = create_response_with_account_retrying_5xx( &self.codex, self.installation_id.as_deref(), @@ -636,11 +566,11 @@ impl ResponseDispatchService { } if is_model_unsupported_sse_failure(failure) { let upstream_error = sse_failure_error_body(failure); - if model_unsupported_retry_used { - let error = ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count + 1, - upstream_error, - }; + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = + ResponseDispatchError::from_exhausted_account(exhausted); self.record_response_dispatch_error( request_id, route, @@ -657,19 +587,18 @@ impl ResponseDispatchService { .await; return Err(error); } - model_unsupported_count += 1; - last_model_unsupported_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::ModelUnsupported); - model_unsupported_retry_used = true; + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); excluded_account_ids.push(release_account_id); continue; } if is_quota_exhausted_sse_failure(failure) { - quota_exhausted_count += 1; - last_quota_error = Some(failure.message.clone()); - last_exhausted_account_class = - Some(ExhaustedAccountClass::QuotaExhausted); + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + failure.message.clone(), + ); self.account_pool .set_status(&release_account_id, AccountStatus::QuotaExhausted) .await; @@ -679,28 +608,12 @@ impl ResponseDispatchService { if is_auth_sse_failure(failure) { let upstream_error = sse_failure_error_body(failure); let account_status = auth_sse_failure_account_status(failure); - match account_status { - AccountStatus::Disabled => { - disabled_count += 1; - last_disabled_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Disabled); - } - AccountStatus::Banned => { - banned_count += 1; - last_banned_status_code = - Some(stream_failure_http_status(failure)); - last_banned_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Banned); - } - _ => { - expired_count += 1; - last_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Expired); - } - } + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(stream_failure_http_status(failure)), + ); self.account_pool .set_status(&release_account_id, account_status) .await; @@ -716,9 +629,10 @@ impl ResponseDispatchService { break (account, response, collected_response); } Err(error) if is_rate_limit_upstream_error(&error) => { - rate_limited_count += 1; - last_rate_limit_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = Some(ExhaustedAccountClass::RateLimited); + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); self.account_pool .mark_quota_limited_until(&release_account_id, cooldown_until) @@ -726,9 +640,10 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_quota_exhausted_upstream_error(&error) => { - quota_exhausted_count += 1; - last_quota_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = Some(ExhaustedAccountClass::QuotaExhausted); + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); self.account_pool .set_status(&release_account_id, AccountStatus::QuotaExhausted) .await; @@ -748,24 +663,12 @@ impl ResponseDispatchService { Err(error) if is_auth_upstream_error(&error) => { let upstream_error = upstream_error_body(&error); let account_status = auth_failure_account_status(&error); - match account_status { - AccountStatus::Disabled => { - disabled_count += 1; - last_disabled_auth_error = Some(upstream_error); - last_exhausted_account_class = Some(ExhaustedAccountClass::Disabled); - } - AccountStatus::Banned => { - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error); - last_exhausted_account_class = Some(ExhaustedAccountClass::Banned); - } - _ => { - expired_count += 1; - last_auth_error = Some(upstream_error); - last_exhausted_account_class = Some(ExhaustedAccountClass::Expired); - } - } + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, account_status) .await; @@ -777,21 +680,20 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { - cloudflare_challenge_count += 1; - last_cloudflare_challenge_error = - Some(cloudflare_challenge_error_message().to_string()); - last_exhausted_account_class = Some(ExhaustedAccountClass::CloudflareChallenge); + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); self.cloudflare .apply_challenge(self.account_pool.as_ref(), &release_account_id) .await; excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_path_block_upstream_error(&error) => { - cloudflare_path_block_count += 1; - last_cloudflare_path_block_error = - Some(cloudflare_path_block_error_message().to_string()); - last_exhausted_account_class = - Some(ExhaustedAccountClass::CloudflarePathBlocked); + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); self.cloudflare .apply_path_block(self.account_pool.as_ref(), &release_account_id) .await; @@ -799,11 +701,10 @@ impl ResponseDispatchService { } Err(error) if is_model_unsupported_upstream_error(&error) => { let upstream_error = upstream_error_body(&error); - if model_unsupported_retry_used { - let error = ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count + 1, - upstream_error, - }; + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = ResponseDispatchError::from_exhausted_account(exhausted); self.record_response_dispatch_error( request_id, route, @@ -822,17 +723,17 @@ impl ResponseDispatchService { .await; return Err(error); } - model_unsupported_count += 1; - last_model_unsupported_error = Some(upstream_error); - last_exhausted_account_class = Some(ExhaustedAccountClass::ModelUnsupported); - model_unsupported_retry_used = true; + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); excluded_account_ids.push(release_account_id); } Err(error) if is_banned_upstream_error(&error) => { - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = Some(ExhaustedAccountClass::Banned); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, AccountStatus::Banned) .await; @@ -1013,28 +914,9 @@ impl ResponseDispatchService { acquire_request = acquire_request.with_preferred_account_id(preferred_account_id); } let mut excluded_account_ids = Vec::new(); - let mut rate_limited_count = 0usize; - let mut last_rate_limit_error = None; - let mut quota_exhausted_count = 0usize; - let mut last_quota_error = None; - let mut expired_count = 0usize; - let mut last_auth_error = None; - let mut disabled_count = 0usize; - let mut last_disabled_auth_error = None; - let mut banned_count = 0usize; - let mut last_banned_auth_error = None; - let mut last_banned_status_code: Option = None; - let mut cloudflare_challenge_count = 0usize; - let mut last_cloudflare_challenge_error = None; - let mut cloudflare_path_block_count = 0usize; - let mut last_cloudflare_path_block_error = None; - let mut model_unsupported_count = 0usize; - let mut last_model_unsupported_error = None; - let mut model_unsupported_retry_used = false; + let mut exhausted_accounts = AccountExhaustionTracker::default(); let mut history_recovery_used = false; - let mut last_exhausted_account_class = None; let mut quota_verify_attempts = 0usize; - let mut last_attempted_account_id = None::; macro_rules! return_stream_dispatch_error { ($error:expr) => {{ let error = $error; @@ -1044,7 +926,7 @@ impl ResponseDispatchService { requested_model, started_at, ResponseDispatchErrorDetails { - account_id: last_attempted_account_id.as_deref(), + account_id: exhausted_accounts.last_account_id(), stream: true, compact: false, transport: Some(backend_transport_name( @@ -1085,52 +967,10 @@ impl ResponseDispatchService { .acquire_with(&attempt_acquire_request) .await else { - let error = match last_exhausted_account_class { - Some(ExhaustedAccountClass::QuotaExhausted) => { - ResponseDispatchError::QuotaExhausted { - count: quota_exhausted_count, - upstream_error: last_quota_error.unwrap_or_default(), - } - } - Some(ExhaustedAccountClass::RateLimited) => { - ResponseDispatchError::RateLimited { - count: rate_limited_count, - upstream_error: last_rate_limit_error.unwrap_or_default(), - } - } - Some(ExhaustedAccountClass::Expired) => ResponseDispatchError::Expired { - count: expired_count, - upstream_error: last_auth_error.unwrap_or_default(), - }, - Some(ExhaustedAccountClass::Disabled) => ResponseDispatchError::Disabled { - count: disabled_count, - upstream_error: last_disabled_auth_error.unwrap_or_default(), - }, - Some(ExhaustedAccountClass::Banned) => ResponseDispatchError::Banned { - count: banned_count, - upstream_error: last_banned_auth_error.unwrap_or_default(), - status_code: last_banned_status_code.unwrap_or(403), - }, - Some(ExhaustedAccountClass::CloudflareChallenge) => { - ResponseDispatchError::CloudflareChallenge { - count: cloudflare_challenge_count, - upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), - } - } - Some(ExhaustedAccountClass::CloudflarePathBlocked) => { - ResponseDispatchError::CloudflarePathBlocked { - count: cloudflare_path_block_count, - upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), - } - } - Some(ExhaustedAccountClass::ModelUnsupported) => { - ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count, - upstream_error: last_model_unsupported_error.unwrap_or_default(), - } - } - None => ResponseDispatchError::NoActiveAccount, - }; + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); return_stream_dispatch_error!(error); }; let acquired_account_id = acquired.account.id.clone(); @@ -1150,17 +990,21 @@ impl ResponseDispatchService { { QuotaVerificationDecision::Ready(acquired) => *acquired, QuotaVerificationDecision::RetryWithAnotherAccount => { - rate_limited_count += 1; - last_rate_limit_error = Some(QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string()); - last_exhausted_account_class = Some(ExhaustedAccountClass::RateLimited); + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); continue; } QuotaVerificationDecision::MaxAttemptsReached => { + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); return_stream_dispatch_error!( - ResponseDispatchError::RateLimited { - count: rate_limited_count + 1, - upstream_error: QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string(), - }, + error, account_id: Some(&acquired_account_id), transport: Some(backend_transport_name(backend_transport_for_response_request( &request @@ -1181,7 +1025,6 @@ impl ResponseDispatchService { self.account_pool.wait_for_request_interval(&acquired).await; let account = acquired.account; let release_account_id = account.id.clone(); - last_attempted_account_id = Some(release_account_id.clone()); let response_result = create_response_stream_with_account_retrying_5xx( &self.codex, self.installation_id.as_deref(), @@ -1235,9 +1078,10 @@ impl ResponseDispatchService { if is_rate_limit_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - rate_limited_count += 1; - last_rate_limit_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = Some(ExhaustedAccountClass::RateLimited); + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); self.account_pool .mark_quota_limited_until(&release_account_id, cooldown_until) @@ -1249,10 +1093,10 @@ impl ResponseDispatchService { if is_quota_exhausted_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - quota_exhausted_count += 1; - last_quota_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = - Some(ExhaustedAccountClass::QuotaExhausted); + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); self.account_pool .set_status(&release_account_id, AccountStatus::QuotaExhausted) .await; @@ -1265,28 +1109,12 @@ impl ResponseDispatchService { self.account_pool.release(&release_account_id).await; let upstream_error = upstream_error_body(&error); let account_status = auth_failure_account_status(&error); - match account_status { - AccountStatus::Disabled => { - disabled_count += 1; - last_disabled_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Disabled); - } - AccountStatus::Banned => { - banned_count += 1; - last_banned_status_code = - Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Banned); - } - _ => { - expired_count += 1; - last_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Expired); - } - } + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, account_status) .await; @@ -1302,11 +1130,10 @@ impl ResponseDispatchService { if is_cloudflare_challenge_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - cloudflare_challenge_count += 1; - last_cloudflare_challenge_error = - Some(cloudflare_challenge_error_message().to_string()); - last_exhausted_account_class = - Some(ExhaustedAccountClass::CloudflareChallenge); + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); self.cloudflare .apply_challenge(self.account_pool.as_ref(), &release_account_id) .await; @@ -1317,11 +1144,10 @@ impl ResponseDispatchService { if is_cloudflare_path_block_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - cloudflare_path_block_count += 1; - last_cloudflare_path_block_error = - Some(cloudflare_path_block_error_message().to_string()); - last_exhausted_account_class = - Some(ExhaustedAccountClass::CloudflarePathBlocked); + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); self.cloudflare .apply_path_block(self.account_pool.as_ref(), &release_account_id) .await; @@ -1333,21 +1159,19 @@ impl ResponseDispatchService { { self.account_pool.release(&release_account_id).await; let upstream_error = upstream_error_body(&error); - if model_unsupported_retry_used { + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { return_stream_dispatch_error!( - ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count + 1, - upstream_error, - }, + ResponseDispatchError::from_exhausted_account(exhausted), account_id: Some(&release_account_id), transport: Some(backend_transport_name(transport)) ); } - model_unsupported_count += 1; - last_model_unsupported_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::ModelUnsupported); - model_unsupported_retry_used = true; + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); excluded_account_ids.push(release_account_id); continue; } @@ -1355,10 +1179,12 @@ impl ResponseDispatchService { if is_banned_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = Some(ExhaustedAccountClass::Banned); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, AccountStatus::Banned) .await; @@ -1417,29 +1243,29 @@ impl ResponseDispatchService { } if is_model_unsupported_sse_failure(&failure) { let upstream_error = sse_failure_error_body(&failure); - if model_unsupported_retry_used { + if let Some(exhausted) = exhausted_accounts + .model_unsupported_retry_exhausted(upstream_error.clone()) + { self.account_pool.release(&release_account_id).await; return_stream_dispatch_error!( - ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count + 1, - upstream_error, - }, + ResponseDispatchError::from_exhausted_account(exhausted), account_id: Some(&release_account_id), transport: Some(backend_transport_name(transport)) ); } - model_unsupported_count += 1; - last_model_unsupported_error = Some(upstream_error); - model_unsupported_retry_used = true; + exhausted_accounts.record_model_unsupported( + Some(&release_account_id), + upstream_error, + ); excluded_account_ids.push(release_account_id); self.account_pool.release(&account.id).await; continue; } if is_quota_exhausted_sse_failure(&failure) { - quota_exhausted_count += 1; - last_quota_error = Some(failure.message.clone()); - last_exhausted_account_class = - Some(ExhaustedAccountClass::QuotaExhausted); + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + failure.message.clone(), + ); self.account_pool .set_status(&release_account_id, AccountStatus::QuotaExhausted) .await; @@ -1450,28 +1276,12 @@ impl ResponseDispatchService { if is_auth_sse_failure(&failure) { let upstream_error = sse_failure_error_body(&failure); let account_status = auth_sse_failure_account_status(&failure); - match account_status { - AccountStatus::Disabled => { - disabled_count += 1; - last_disabled_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Disabled); - } - AccountStatus::Banned => { - banned_count += 1; - last_banned_status_code = - Some(stream_failure_http_status(&failure)); - last_banned_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Banned); - } - _ => { - expired_count += 1; - last_auth_error = Some(upstream_error); - last_exhausted_account_class = - Some(ExhaustedAccountClass::Expired); - } - } + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(stream_failure_http_status(&failure)), + ); self.account_pool .set_status(&release_account_id, account_status) .await; @@ -1533,9 +1343,10 @@ impl ResponseDispatchService { } Err(error) if is_rate_limit_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - rate_limited_count += 1; - last_rate_limit_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = Some(ExhaustedAccountClass::RateLimited); + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); self.account_pool .mark_quota_limited_until(&release_account_id, cooldown_until) @@ -1544,9 +1355,10 @@ impl ResponseDispatchService { } Err(error) if is_quota_exhausted_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - quota_exhausted_count += 1; - last_quota_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = Some(ExhaustedAccountClass::QuotaExhausted); + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); self.account_pool .set_status(&release_account_id, AccountStatus::QuotaExhausted) .await; @@ -1568,24 +1380,12 @@ impl ResponseDispatchService { self.account_pool.release(&release_account_id).await; let upstream_error = upstream_error_body(&error); let account_status = auth_failure_account_status(&error); - match account_status { - AccountStatus::Disabled => { - disabled_count += 1; - last_disabled_auth_error = Some(upstream_error); - last_exhausted_account_class = Some(ExhaustedAccountClass::Disabled); - } - AccountStatus::Banned => { - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error); - last_exhausted_account_class = Some(ExhaustedAccountClass::Banned); - } - _ => { - expired_count += 1; - last_auth_error = Some(upstream_error); - last_exhausted_account_class = Some(ExhaustedAccountClass::Expired); - } - } + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, account_status) .await; @@ -1598,10 +1398,10 @@ impl ResponseDispatchService { } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - cloudflare_challenge_count += 1; - last_cloudflare_challenge_error = - Some(cloudflare_challenge_error_message().to_string()); - last_exhausted_account_class = Some(ExhaustedAccountClass::CloudflareChallenge); + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); self.cloudflare .apply_challenge(self.account_pool.as_ref(), &release_account_id) .await; @@ -1609,11 +1409,10 @@ impl ResponseDispatchService { } Err(error) if is_cloudflare_path_block_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - cloudflare_path_block_count += 1; - last_cloudflare_path_block_error = - Some(cloudflare_path_block_error_message().to_string()); - last_exhausted_account_class = - Some(ExhaustedAccountClass::CloudflarePathBlocked); + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); self.cloudflare .apply_path_block(self.account_pool.as_ref(), &release_account_id) .await; @@ -1622,30 +1421,29 @@ impl ResponseDispatchService { Err(error) if is_model_unsupported_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; let upstream_error = upstream_error_body(&error); - if model_unsupported_retry_used { + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { return_stream_dispatch_error!( - ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count + 1, - upstream_error, - }, + ResponseDispatchError::from_exhausted_account(exhausted), account_id: Some(&release_account_id), transport: Some(backend_transport_name(backend_transport_for_response_request( &request ))) ); } - model_unsupported_count += 1; - last_model_unsupported_error = Some(upstream_error); - last_exhausted_account_class = Some(ExhaustedAccountClass::ModelUnsupported); - model_unsupported_retry_used = true; + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); excluded_account_ids.push(release_account_id); } Err(error) if is_banned_upstream_error(&error) => { self.account_pool.release(&release_account_id).await; - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error_body(&error)); - last_exhausted_account_class = Some(ExhaustedAccountClass::Banned); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, AccountStatus::Banned) .await; @@ -1685,167 +1483,31 @@ impl ResponseDispatchService { let display_model = ModelCatalog::build_display_model_name(&parsed_model); request.model = parsed_model.model_id; let mut excluded_account_ids = Vec::new(); - let mut rate_limited_count = 0usize; - let mut last_rate_limit_error = None; - let mut quota_exhausted_count = 0usize; - let mut last_quota_error = None; - let mut expired_count = 0usize; - let mut last_auth_error = None; - let mut disabled_count = 0usize; - let mut last_disabled_auth_error = None; - let mut banned_count = 0usize; - let mut last_banned_auth_error = None; - let mut last_banned_status_code: Option = None; - let mut cloudflare_challenge_count = 0usize; - let mut last_cloudflare_challenge_error = None; - let mut cloudflare_path_block_count = 0usize; - let mut last_cloudflare_path_block_error = None; - let mut model_unsupported_count = 0usize; - let mut last_model_unsupported_error = None; - let mut model_unsupported_retry_used = false; + let mut exhausted_accounts = AccountExhaustionTracker::default(); let mut quota_verify_attempts = 0usize; - let mut last_attempted_account_id = None::; loop { let acquire_request = AccountAcquireRequest::new(&request.model, Utc::now()) .with_exclude_account_ids(excluded_account_ids.iter().cloned()); let acquired = match self.account_pool.acquire_with(&acquire_request).await { Some(acquired) => acquired, - None if quota_exhausted_count > 0 => { - let error = ResponseDispatchError::QuotaExhausted { - count: quota_exhausted_count, - upstream_error: last_quota_error.unwrap_or_default(), - }; - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - &error, - ) - .await; - return Err(error); - } - None if rate_limited_count > 0 => { - let error = ResponseDispatchError::RateLimited { - count: rate_limited_count, - upstream_error: last_rate_limit_error.unwrap_or_default(), - }; - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - &error, - ) - .await; - return Err(error); - } - None if expired_count > 0 => { - let error = ResponseDispatchError::Expired { - count: expired_count, - upstream_error: last_auth_error.unwrap_or_default(), - }; - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - &error, - ) - .await; - return Err(error); - } - None if disabled_count > 0 => { - let error = ResponseDispatchError::Disabled { - count: disabled_count, - upstream_error: last_disabled_auth_error.unwrap_or_default(), - }; - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - &error, - ) - .await; - return Err(error); - } - None if banned_count > 0 => { - let error = ResponseDispatchError::Banned { - count: banned_count, - upstream_error: last_banned_auth_error.unwrap_or_default(), - status_code: last_banned_status_code.unwrap_or(403), - }; - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - &error, - ) - .await; - return Err(error); - } - None if cloudflare_challenge_count > 0 => { - let error = ResponseDispatchError::CloudflareChallenge { - count: cloudflare_challenge_count, - upstream_error: last_cloudflare_challenge_error.unwrap_or_default(), - }; - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - &error, - ) - .await; - return Err(error); - } - None if cloudflare_path_block_count > 0 => { - let error = ResponseDispatchError::CloudflarePathBlocked { - count: cloudflare_path_block_count, - upstream_error: last_cloudflare_path_block_error.unwrap_or_default(), - }; - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - &error, - ) - .await; - return Err(error); - } - None if model_unsupported_count > 0 => { - let error = ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count, - upstream_error: last_model_unsupported_error.unwrap_or_default(), - }; - self.record_compact_dispatch_error( - request_id, - requested_model, - started_at, - last_attempted_account_id.as_deref(), - &error, - ) - .await; - return Err(error); - } None => { - let error = ResponseDispatchError::NoActiveAccount; + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); self.record_compact_dispatch_error( request_id, requested_model, started_at, - last_attempted_account_id.as_deref(), + exhausted_accounts.last_account_id(), &error, ) .await; return Err(error); } }; - last_attempted_account_id = Some(acquired.account.id.clone()); + let acquired_account_id = acquired.account.id.clone(); let acquired = match verify_acquired_quota_if_required( QuotaVerificationContext { account_pool: self.account_pool.as_ref(), @@ -1862,20 +1524,24 @@ impl ResponseDispatchService { { QuotaVerificationDecision::Ready(acquired) => *acquired, QuotaVerificationDecision::RetryWithAnotherAccount => { - rate_limited_count += 1; - last_rate_limit_error = Some(QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string()); + exhausted_accounts + .record_rate_limited(None, QUOTA_VERIFY_LIMIT_REACHED_MESSAGE); continue; } QuotaVerificationDecision::MaxAttemptsReached => { - let error = ResponseDispatchError::RateLimited { - count: rate_limited_count + 1, - upstream_error: QUOTA_VERIFY_LIMIT_REACHED_MESSAGE.to_string(), - }; + exhausted_accounts.record_rate_limited( + Some(&acquired_account_id), + QUOTA_VERIFY_LIMIT_REACHED_MESSAGE, + ); + let error = exhausted_accounts + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount); self.record_compact_dispatch_error( request_id, requested_model, started_at, - last_attempted_account_id.as_deref(), + Some(&acquired_account_id), &error, ) .await; @@ -1947,8 +1613,10 @@ impl ResponseDispatchService { return Ok(response.body); } Err(error) if is_rate_limit_upstream_error(&error) => { - rate_limited_count += 1; - last_rate_limit_error = Some(upstream_error_body(&error)); + exhausted_accounts.record_rate_limited( + Some(&release_account_id), + upstream_error_body(&error), + ); let cooldown_until = rate_limit_cooldown_until(&error, Utc::now()); self.account_pool .mark_quota_limited_until(&release_account_id, cooldown_until) @@ -1956,8 +1624,10 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_quota_exhausted_upstream_error(&error) => { - quota_exhausted_count += 1; - last_quota_error = Some(upstream_error_body(&error)); + exhausted_accounts.record_quota_exhausted( + Some(&release_account_id), + upstream_error_body(&error), + ); self.account_pool .set_status(&release_account_id, AccountStatus::QuotaExhausted) .await; @@ -1966,21 +1636,12 @@ impl ResponseDispatchService { Err(error) if is_auth_upstream_error(&error) => { let upstream_error = upstream_error_body(&error); let account_status = auth_failure_account_status(&error); - match account_status { - AccountStatus::Disabled => { - disabled_count += 1; - last_disabled_auth_error = Some(upstream_error); - } - AccountStatus::Banned => { - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error); - } - _ => { - expired_count += 1; - last_auth_error = Some(upstream_error); - } - } + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + account_status, + upstream_error, + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, account_status) .await; @@ -1992,18 +1653,20 @@ impl ResponseDispatchService { excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_challenge_upstream_error(&error) => { - cloudflare_challenge_count += 1; - last_cloudflare_challenge_error = - Some(cloudflare_challenge_error_message().to_string()); + exhausted_accounts.record_cloudflare_challenge( + Some(&release_account_id), + cloudflare_challenge_error_message(), + ); self.cloudflare .apply_challenge(self.account_pool.as_ref(), &release_account_id) .await; excluded_account_ids.push(release_account_id); } Err(error) if is_cloudflare_path_block_upstream_error(&error) => { - cloudflare_path_block_count += 1; - last_cloudflare_path_block_error = - Some(cloudflare_path_block_error_message().to_string()); + exhausted_accounts.record_cloudflare_path_blocked( + Some(&release_account_id), + cloudflare_path_block_error_message(), + ); self.cloudflare .apply_path_block(self.account_pool.as_ref(), &release_account_id) .await; @@ -2011,11 +1674,10 @@ impl ResponseDispatchService { } Err(error) if is_model_unsupported_upstream_error(&error) => { let upstream_error = upstream_error_body(&error); - if model_unsupported_retry_used { - let error = ResponseDispatchError::ModelUnsupported { - count: model_unsupported_count + 1, - upstream_error, - }; + if let Some(exhausted) = + exhausted_accounts.model_unsupported_retry_exhausted(upstream_error.clone()) + { + let error = ResponseDispatchError::from_exhausted_account(exhausted); self.record_compact_dispatch_error( request_id, requested_model, @@ -2026,15 +1688,17 @@ impl ResponseDispatchService { .await; return Err(error); } - model_unsupported_count += 1; - last_model_unsupported_error = Some(upstream_error); - model_unsupported_retry_used = true; + exhausted_accounts + .record_model_unsupported(Some(&release_account_id), upstream_error); excluded_account_ids.push(release_account_id); } Err(error) if is_banned_upstream_error(&error) => { - banned_count += 1; - last_banned_status_code = Some(upstream_error_http_status(&error)); - last_banned_auth_error = Some(upstream_error_body(&error)); + exhausted_accounts.record_auth_failure( + Some(&release_account_id), + AccountStatus::Banned, + upstream_error_body(&error), + Some(upstream_error_http_status(&error)), + ); self.account_pool .set_status(&release_account_id, AccountStatus::Banned) .await; diff --git a/backend/src/proxy/dispatch/session_affinity.rs b/backend/src/proxy/dispatch/session_affinity.rs index b4cec4210..86676aefd 100644 --- a/backend/src/proxy/dispatch/session_affinity.rs +++ b/backend/src/proxy/dispatch/session_affinity.rs @@ -125,14 +125,6 @@ impl SessionAffinityMap { self.entries.remove(response_id).is_some() } - pub fn len(&self) -> usize { - self.entries.len() - } - - pub fn is_empty(&self) -> bool { - self.entries.is_empty() - } - fn active_entry(&self, response_id: &str, now: DateTime) -> Option<&SessionAffinityEntry> { self.entries .get(response_id) diff --git a/backend/src/proxy/openai/chat.rs b/backend/src/proxy/openai/chat.rs index f05ed05b9..c0bd5adfb 100644 --- a/backend/src/proxy/openai/chat.rs +++ b/backend/src/proxy/openai/chat.rs @@ -863,11 +863,12 @@ fn response_format_text(response_format: Option) -> PreparedResponseForma pub fn chat_completion_from_codex_sse( body: &str, model: &str, - _include_reasoning: bool, + include_reasoning: bool, tuple_schema: Option<&Value>, ) -> Result, ChatStreamTranslationError> { let events = parse_sse_events(body)?; let mut output_text = String::new(); + let mut reasoning_text = String::new(); let mut function_call_items: BTreeMap = BTreeMap::new(); let mut finished_call_ids: BTreeSet = BTreeSet::new(); let mut tool_calls: Vec = Vec::new(); @@ -904,6 +905,13 @@ pub fn chat_completion_from_codex_sse( } } } + Some("response.reasoning_summary_text.delta" | "response.reasoning_text.delta") + if include_reasoning => + { + if let Some(delta) = value.get("delta").and_then(Value::as_str) { + reasoning_text.push_str(delta); + } + } Some("response.output_item.added") => { if let Some((item_id, call_id, name)) = function_call_item(value.get("item")) { function_call_items.insert( @@ -957,23 +965,26 @@ pub fn chat_completion_from_codex_sse( let mut choices = Vec::new(); let finish_reason = openai_finish_reason_for_response(Some(&response), !tool_calls.is_empty()); + let mut message = serde_json::Map::new(); + message.insert("role".to_string(), Value::String("assistant".to_string())); + message.insert("content".to_string(), Value::String(output_text_value)); + if include_reasoning && !reasoning_text.is_empty() { + message.insert( + "reasoning_content".to_string(), + Value::String(reasoning_text), + ); + } if !tool_calls.is_empty() { + message.insert("tool_calls".to_string(), Value::Array(tool_calls)); choices.push(json!({ "index": 0, - "message": { - "role": "assistant", - "content": output_text_value, - "tool_calls": tool_calls, - }, + "message": Value::Object(message), "finish_reason": finish_reason, })); } else { choices.push(json!({ "index": 0, - "message": { - "role": "assistant", - "content": output_text_value, - }, + "message": Value::Object(message), "finish_reason": finish_reason, })); } diff --git a/backend/src/proxy/openai/errors.rs b/backend/src/proxy/openai/errors.rs index d956834a6..a529db4db 100644 --- a/backend/src/proxy/openai/errors.rs +++ b/backend/src/proxy/openai/errors.rs @@ -8,7 +8,10 @@ use axum::{ use serde_json::{json, Value}; use crate::proxy::{ - dispatch::{chat::ChatDispatchError, responses::errors::ResponseDispatchError}, + dispatch::{ + chat::ChatDispatchError, errors::DispatchFailureClass, + responses::errors::ResponseDispatchError, + }, openai::responses::response_failed_sse_event as encode_response_failed_sse_event, }; @@ -113,7 +116,7 @@ pub fn chat_stream_dispatch_openai_error(error: &ResponseDispatchError) -> OpenA pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Response { match error { - ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => { + ResponseDispatchError::NoActiveAccount => { responses_no_available_accounts_response().into_response() } error => response_dispatch_openai_error_response(&error), @@ -122,7 +125,7 @@ pub fn responses_dispatch_error_response(error: ResponseDispatchError) -> Respon pub fn responses_compact_dispatch_error_response(error: ResponseDispatchError) -> Response { match error { - ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => { + ResponseDispatchError::NoActiveAccount => { responses_no_available_accounts_response().into_response() } error => { @@ -224,77 +227,21 @@ fn invalid_openai_request_response(message: &str) -> (StatusCode, Json) { fn chat_dispatch_openai_error(error: &ChatDispatchError) -> OpenAiErrorDetails { let status = chat_dispatch_status(error); - let (message, kind) = match error { - ChatDispatchError::NoActiveAccount | ChatDispatchError::AccountStore => ( - NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned(), - UPSTREAM_UNAVAILABLE_ERROR, - ), - ChatDispatchError::Upstream(_) => ( - UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(), - UPSTREAM_ERROR, - ), - ChatDispatchError::QuotaExhausted { - count, - upstream_error, - } => ( - exhausted_dispatch_message(*count, "quota-exhausted", upstream_error, false), - INSUFFICIENT_QUOTA_ERROR, - ), - ChatDispatchError::RateLimited { - count, - upstream_error, - } => ( - exhausted_dispatch_message(*count, "rate-limited", upstream_error, false), - RATE_LIMIT_ERROR, - ), - ChatDispatchError::Expired { - count, - upstream_error, - } => ( - exhausted_dispatch_message(*count, "expired", upstream_error, false), - AUTHENTICATION_ERROR, - ), - ChatDispatchError::Disabled { - count, - upstream_error, - } => ( - exhausted_dispatch_message(*count, "disabled", upstream_error, false), - AUTHENTICATION_ERROR, - ), - ChatDispatchError::Banned { - count, - upstream_error, - .. - } => ( - exhausted_dispatch_message(*count, "banned", upstream_error, false), - AUTHENTICATION_ERROR, - ), - ChatDispatchError::CloudflareChallenge { - count, - upstream_error, - } => ( - exhausted_dispatch_message(*count, "cloudflare-challenge", upstream_error, false), - UPSTREAM_ERROR, - ), - ChatDispatchError::CloudflarePathBlocked { - count, - upstream_error, - } => ( - exhausted_dispatch_message(*count, "cloudflare-path-block", upstream_error, false), - UPSTREAM_ERROR, - ), - ChatDispatchError::ModelUnsupported { - count, - upstream_error, - } => ( - exhausted_dispatch_message(*count, "model-unsupported", upstream_error, false), - MODEL_NOT_FOUND_ERROR, - ), - ChatDispatchError::InvalidSse(_) | ChatDispatchError::EmptyUpstreamResponse => ( - INVALID_UPSTREAM_CODEX_RESPONSE_MESSAGE.to_owned(), - INVALID_UPSTREAM_RESPONSE_ERROR, - ), + let metadata = error.metadata(); + let message = if let Some(exhausted) = error.exhausted_account() { + exhausted_dispatch_message( + exhausted.count, + exhausted.kind.message_reason(), + exhausted.upstream_error, + false, + ) + } else { + dispatch_failure_message( + metadata.failure_class, + DispatchFailureMessageStyle::Standard, + ) }; + let kind = dispatch_failure_openai_error_kind(metadata.failure_class, status); OpenAiErrorDetails { status, @@ -321,27 +268,10 @@ fn response_dispatch_openai_error(error: &ResponseDispatchError) -> OpenAiErrorD } fn response_dispatch_error_kind(error: &ResponseDispatchError) -> OpenAiErrorKind { - match error { - ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => { - UPSTREAM_UNAVAILABLE_ERROR - } - ResponseDispatchError::QuotaExhausted { .. } => INSUFFICIENT_QUOTA_ERROR, - ResponseDispatchError::RateLimited { .. } => RATE_LIMIT_ERROR, - ResponseDispatchError::Expired { .. } - | ResponseDispatchError::Disabled { .. } - | ResponseDispatchError::Banned { .. } => AUTHENTICATION_ERROR, - ResponseDispatchError::ModelUnsupported { .. } => MODEL_NOT_FOUND_ERROR, - ResponseDispatchError::InvalidSse(_) - | ResponseDispatchError::MissingCompleted - | ResponseDispatchError::EmptyUpstreamResponse => INVALID_UPSTREAM_RESPONSE_ERROR, - ResponseDispatchError::Failed(_) => responses_error_kind_for_status(status_from_u16( - error.http_status_code(), - StatusCode::BAD_GATEWAY, - )), - ResponseDispatchError::Upstream(_) - | ResponseDispatchError::CloudflareChallenge { .. } - | ResponseDispatchError::CloudflarePathBlocked { .. } => UPSTREAM_ERROR, - } + dispatch_failure_openai_error_kind( + error.metadata().failure_class, + status_from_u16(error.http_status_code(), StatusCode::BAD_GATEWAY), + ) } pub fn response_dispatch_http_error( @@ -395,60 +325,75 @@ fn response_dispatch_message( error: &ResponseDispatchError, style: ResponseDispatchMessageStyle, ) -> String { - match error { - ResponseDispatchError::NoActiveAccount | ResponseDispatchError::AccountStore => { - NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned() - } - ResponseDispatchError::Upstream(_) => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(), - ResponseDispatchError::QuotaExhausted { - count, - upstream_error, - } => exhausted_dispatch_message(*count, "quota-exhausted", upstream_error, false), - ResponseDispatchError::RateLimited { - count, - upstream_error, - } => exhausted_dispatch_message(*count, "rate-limited", upstream_error, false), - ResponseDispatchError::Expired { - count, - upstream_error, - } => exhausted_dispatch_message(*count, "expired", upstream_error, false), - ResponseDispatchError::Disabled { - count, - upstream_error, - } => exhausted_dispatch_message(*count, "disabled", upstream_error, false), - ResponseDispatchError::Banned { - count, - upstream_error, - .. - } => exhausted_dispatch_message(*count, "banned", upstream_error, false), - ResponseDispatchError::CloudflareChallenge { - count, - upstream_error, - } => exhausted_dispatch_message(*count, "cloudflare-challenge", upstream_error, false), - ResponseDispatchError::CloudflarePathBlocked { - count, - upstream_error, - } => exhausted_dispatch_message( - *count, - "cloudflare-path-block", - upstream_error, - matches!(style, ResponseDispatchMessageStyle::ResponsesStream), - ), - ResponseDispatchError::ModelUnsupported { - count, - upstream_error, - } => exhausted_dispatch_message( - *count, - "model-unsupported", - upstream_error, - matches!(style, ResponseDispatchMessageStyle::ResponsesStream), - ), - ResponseDispatchError::InvalidSse(_) - | ResponseDispatchError::MissingCompleted - | ResponseDispatchError::EmptyUpstreamResponse => { + let metadata = error.metadata(); + if let Some(exhausted) = error.exhausted_account() { + return exhausted_dispatch_message( + exhausted.count, + exhausted.kind.message_reason(), + exhausted.upstream_error, + matches!( + (metadata.failure_class, style), + ( + DispatchFailureClass::CloudflarePathBlocked + | DispatchFailureClass::ModelUnsupported, + ResponseDispatchMessageStyle::ResponsesStream, + ) + ), + ); + } + dispatch_failure_message( + metadata.failure_class, + match style { + ResponseDispatchMessageStyle::Standard => DispatchFailureMessageStyle::Standard, + ResponseDispatchMessageStyle::ResponsesStream => { + DispatchFailureMessageStyle::ResponsesStream + } + }, + ) +} + +#[derive(Clone, Copy)] +enum DispatchFailureMessageStyle { + Standard, + ResponsesStream, +} + +fn dispatch_failure_message( + failure_class: DispatchFailureClass, + _style: DispatchFailureMessageStyle, +) -> String { + match failure_class { + DispatchFailureClass::NoAvailableAccounts => NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned(), + DispatchFailureClass::Upstream => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(), + DispatchFailureClass::InvalidSse + | DispatchFailureClass::MissingCompleted + | DispatchFailureClass::EmptyUpstreamResponse => { INVALID_UPSTREAM_CODEX_RESPONSE_MESSAGE.to_owned() } - ResponseDispatchError::Failed(_) => UPSTREAM_CODEX_RESPONSE_FAILED_MESSAGE.to_owned(), + DispatchFailureClass::ResponseFailed => UPSTREAM_CODEX_RESPONSE_FAILED_MESSAGE.to_owned(), + _ => UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(), + } +} + +fn dispatch_failure_openai_error_kind( + failure_class: DispatchFailureClass, + status: StatusCode, +) -> OpenAiErrorKind { + match failure_class { + DispatchFailureClass::NoAvailableAccounts => UPSTREAM_UNAVAILABLE_ERROR, + DispatchFailureClass::QuotaExhausted => INSUFFICIENT_QUOTA_ERROR, + DispatchFailureClass::RateLimited => RATE_LIMIT_ERROR, + DispatchFailureClass::Expired + | DispatchFailureClass::Disabled + | DispatchFailureClass::Banned => AUTHENTICATION_ERROR, + DispatchFailureClass::ModelUnsupported => MODEL_NOT_FOUND_ERROR, + DispatchFailureClass::InvalidSse + | DispatchFailureClass::MissingCompleted + | DispatchFailureClass::EmptyUpstreamResponse => INVALID_UPSTREAM_RESPONSE_ERROR, + DispatchFailureClass::ResponseFailed => responses_error_kind_for_status(status), + DispatchFailureClass::Upstream + | DispatchFailureClass::CloudflareChallenge + | DispatchFailureClass::CloudflarePathBlocked => UPSTREAM_ERROR, } } diff --git a/backend/src/runtime/bootstrap.rs b/backend/src/runtime/bootstrap.rs index 9778567e1..093dcb9b7 100644 --- a/backend/src/runtime/bootstrap.rs +++ b/backend/src/runtime/bootstrap.rs @@ -113,7 +113,7 @@ async fn build_application( let task_coordinator = TaskCoordinator::start(&runtime_config, &services); let trusted_proxies = TrustedProxyConfig::from_entries(&config.server.trusted_proxies)?; - let app_state = AppState { config, services }; + let app_state = AppState { services }; Ok(( crate::http::router::router_with_trusted_proxies(trusted_proxies).with_state(app_state), diff --git a/backend/src/runtime/state.rs b/backend/src/runtime/state.rs index 4389d2ea0..5f09d3a83 100644 --- a/backend/src/runtime/state.rs +++ b/backend/src/runtime/state.rs @@ -6,7 +6,6 @@ use crate::runtime::services::Services; /// HTTP handler 共享的运行时状态。 #[derive(Clone)] pub struct AppState { - pub config: AppConfig, pub services: Services, } diff --git a/backend/src/runtime/tasks/cleanup.rs b/backend/src/runtime/tasks/cleanup.rs index d3ffe9773..05ae4c7e4 100644 --- a/backend/src/runtime/tasks/cleanup.rs +++ b/backend/src/runtime/tasks/cleanup.rs @@ -1,12 +1,14 @@ //! 通用过期数据清理任务骨架。 -use std::{fmt::Display, future::Future, time::Duration}; +use std::{fmt::Display, future::Future}; use chrono::{DateTime, Utc}; -use tokio::time::interval; use tracing::{debug, info, warn}; -use super::coordinator::SchedulerHandle; +use super::{ + coordinator::SchedulerHandle, + periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner}, +}; /// 支持按时间清理过期记录的存储。 pub(crate) trait ExpiredCleanupStore: Send + Sync + 'static { @@ -47,39 +49,12 @@ where } pub(crate) fn start(self) -> SchedulerHandle { - let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); - - tokio::spawn(async move { - let mut ticker = interval(Duration::from_secs(self.interval_secs)); - info!( - interval_secs = self.interval_secs, - "{}", self.messages.started - ); - - loop { - tokio::select! { - _ = ticker.tick() => { - match self.cleanup_once().await { - Ok(deleted) if deleted > 0 => { - info!(deleted, "{}", self.messages.deleted); - } - Ok(_) => { - debug!("{}", self.messages.empty); - } - Err(error) => { - warn!(error = %error, "{}", self.messages.failed); - } - } - } - _ = shutdown_rx.recv() => { - info!("{}", self.messages.stopped); - break; - } - } - } - }); - - SchedulerHandle::new(shutdown_tx) + let config = PeriodicTaskConfig::new( + self.interval_secs, + self.messages.started, + self.messages.stopped, + ); + spawn_periodic_task(self, config) } pub(crate) async fn cleanup_once(&self) -> Result { @@ -90,3 +65,24 @@ where self.store.delete_expired_at(now).await } } + +impl PeriodicTaskRunner for CleanupTask +where + S: ExpiredCleanupStore, +{ + fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> { + Box::pin(async move { + match self.cleanup_once().await { + Ok(deleted) if deleted > 0 => { + info!(deleted, "{}", self.messages.deleted); + } + Ok(_) => { + debug!("{}", self.messages.empty); + } + Err(error) => { + warn!(error = %error, "{}", self.messages.failed); + } + } + }) + } +} diff --git a/backend/src/runtime/tasks/coordinator.rs b/backend/src/runtime/tasks/coordinator.rs index 005ef33c3..999f9da6c 100644 --- a/backend/src/runtime/tasks/coordinator.rs +++ b/backend/src/runtime/tasks/coordinator.rs @@ -1,6 +1,6 @@ //! 后台任务协调器。 -use std::{path::PathBuf, sync::Arc}; +use std::{path::PathBuf, sync::Arc, time::Duration}; use tokio::task::JoinHandle; @@ -120,22 +120,21 @@ impl TaskCoordinator { /// 后台任务关闭句柄。 pub enum SchedulerHandle { /// 通过 channel 发送关闭信号。 - Channel(tokio::sync::mpsc::Sender<()>), - /// 直接持有 tokio 任务句柄。 - JoinHandle(JoinHandle<()>), + Channel { + shutdown_tx: tokio::sync::mpsc::Sender<()>, + handle: JoinHandle<()>, + }, /// 关闭上游 WebSocket 连接池。 WebSocketPool(Arc), } impl SchedulerHandle { /// 使用关闭 channel 构造句柄。 - pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self { - Self::Channel(shutdown_tx) - } - - /// 使用 `JoinHandle` 构造句柄。 - pub(crate) fn from_join_handle(handle: JoinHandle<()>) -> Self { - Self::JoinHandle(handle) + pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>, handle: JoinHandle<()>) -> Self { + Self::Channel { + shutdown_tx, + handle, + } } /// 使用 WebSocket 连接池构造句柄。 @@ -146,11 +145,12 @@ impl SchedulerHandle { /// 关闭任务。 pub async fn shutdown(self) { match self { - Self::Channel(shutdown_tx) => { + Self::Channel { + shutdown_tx, + handle, + } => { let _ = shutdown_tx.send(()).await; - } - Self::JoinHandle(handle) => { - handle.abort(); + wait_for_task_shutdown(handle).await; } Self::WebSocketPool(pool) => { pool.shutdown().await; @@ -158,3 +158,32 @@ impl SchedulerHandle { } } } + +const TASK_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); + +async fn wait_for_task_shutdown(handle: JoinHandle<()>) { + let mut handle = handle; + let timeout = tokio::time::sleep(TASK_SHUTDOWN_TIMEOUT); + tokio::pin!(timeout); + + tokio::select! { + result = &mut handle => { + log_task_shutdown_result(result); + } + () = &mut timeout => { + tracing::warn!(timeout_secs = TASK_SHUTDOWN_TIMEOUT.as_secs(), "等待后台任务关闭超时"); + handle.abort(); + log_task_shutdown_result(handle.await); + } + } +} + +fn log_task_shutdown_result(result: Result<(), tokio::task::JoinError>) { + match result { + Ok(()) => {} + Err(error) if error.is_cancelled() => {} + Err(error) => { + tracing::warn!(error = %error, "后台任务关闭时异常退出"); + } + } +} diff --git a/backend/src/runtime/tasks/fingerprint_update.rs b/backend/src/runtime/tasks/fingerprint_update.rs index 9dbef918f..1ae718b1f 100644 --- a/backend/src/runtime/tasks/fingerprint_update.rs +++ b/backend/src/runtime/tasks/fingerprint_update.rs @@ -2,9 +2,14 @@ use std::path::PathBuf; -use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker}; +use tracing::warn; -use super::coordinator::SchedulerHandle; +use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker, APPCAST_POLL_INTERVAL}; + +use super::{ + coordinator::SchedulerHandle, + periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner}, +}; /// Codex Desktop 官方 appcast 地址。 pub const CODEX_DESKTOP_APPCAST_URL: &str = @@ -28,5 +33,35 @@ pub fn start_fingerprint_update_task( current_version, current_build, ); - SchedulerHandle::from_join_handle(update_checker.start_background_checker()) + spawn_periodic_task( + FingerprintUpdateTask { + checker: update_checker, + first_tick: true, + }, + PeriodicTaskConfig::new( + APPCAST_POLL_INTERVAL.as_secs(), + "fingerprint 后台版本检查器已启动", + "fingerprint 后台版本检查器已关闭", + ), + ) +} + +struct FingerprintUpdateTask { + checker: UpdateChecker, + first_tick: bool, +} + +impl PeriodicTaskRunner for FingerprintUpdateTask { + fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> { + Box::pin(async move { + let first_tick = std::mem::take(&mut self.first_tick); + if let Err(error) = self.checker.check_and_apply_update().await { + if first_tick { + warn!(error = %error, "fingerprint 首次版本检查失败"); + } else { + warn!(error = %error, "fingerprint 定期版本检查失败"); + } + } + }) + } } diff --git a/backend/src/runtime/tasks/mod.rs b/backend/src/runtime/tasks/mod.rs index 2929d45b7..62395806f 100644 --- a/backend/src/runtime/tasks/mod.rs +++ b/backend/src/runtime/tasks/mod.rs @@ -5,6 +5,7 @@ pub mod cookie_cleanup; pub mod coordinator; pub mod fingerprint_update; pub mod model_refresh; +mod periodic; pub mod quota_refresh; pub mod session_affinity_cleanup; pub mod session_cleanup; diff --git a/backend/src/runtime/tasks/model_refresh.rs b/backend/src/runtime/tasks/model_refresh.rs index 758707071..392451a73 100644 --- a/backend/src/runtime/tasks/model_refresh.rs +++ b/backend/src/runtime/tasks/model_refresh.rs @@ -3,7 +3,8 @@ use std::{sync::Arc, time::Duration}; use chrono::Utc; -use tokio::time::{interval, sleep}; +use tokio::sync::mpsc; +use tokio::time::sleep; use tracing::{info, warn}; use crate::upstream::accounts::pool::RuntimeAccountPoolService; @@ -11,7 +12,10 @@ use crate::upstream::models::service::{ ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError, }; -use super::coordinator::SchedulerHandle; +use super::{ + coordinator::SchedulerHandle, + periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup}, +}; /// 模型刷新任务接线器。 pub struct ModelRefreshTask { @@ -54,85 +58,13 @@ impl ModelRefreshTask { /// 启动后台刷新任务。 pub fn start(self) -> SchedulerHandle { - let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); - - tokio::spawn(async move { - info!("模型刷新任务已启动"); - - sleep(Duration::from_millis(self.initial_delay_ms)).await; - - let mut attempt = 0; - let mut has_fetched_once = false; - - while attempt < self.max_retries { - tokio::select! { - _ = shutdown_rx.recv() => { - info!("模型刷新任务在首次拉取前关闭"); - return; - } - () = async {} => {} - } - - match self.refresh_once().await { - Ok(result) => { - info!( - refreshed_plans = result.refreshed_plans, - model_count = result.model_count, - "首次模型刷新成功" - ); - has_fetched_once = true; - break; - } - Err(error) => { - attempt += 1; - if matches!(error, ModelServiceError::NoAccounts) { - warn!("首次模型刷新跳过：没有可用账户，将进入周期重试"); - break; - } - - warn!( - attempt, - max_retries = self.max_retries, - error = ?error, - "首次模型刷新失败" - ); - - if attempt < self.max_retries { - tokio::select! { - _ = shutdown_rx.recv() => { - info!("模型刷新任务在重试等待期间关闭"); - return; - } - () = sleep(Duration::from_millis(self.retry_delay_ms)) => {} - } - } - } - } - } - - if !has_fetched_once { - warn!("模型刷新任务首次尝试全部失败，切换到周期刷新"); - } - - let mut ticker = interval(Duration::from_secs(self.refresh_interval_secs)); - ticker.tick().await; - - loop { - tokio::select! { - _ = ticker.tick() => { - if let Err(error) = self.refresh_once().await { - warn!(error = ?error, "刷新模型列表失败"); - } - } - _ = shutdown_rx.recv() => { - info!("模型刷新任务正在关闭"); - break; - } - } - } - }); - - SchedulerHandle::new(shutdown_tx) + let config = PeriodicTaskConfig::new( + self.refresh_interval_secs, + "模型刷新任务已启动", + "模型刷新任务正在关闭", + ) + .wait_first_interval(); + spawn_periodic_task(self, config) } /// 执行一次模型刷新并同步账号池模型计划约束。 @@ -167,3 +99,75 @@ impl ModelRefreshTask { Ok(result) } } + +impl PeriodicTaskRunner for ModelRefreshTask { + fn before_loop<'a>( + &'a mut self, + shutdown_rx: &'a mut mpsc::Receiver<()>, + ) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> { + Box::pin(async move { + tokio::select! { + _ = shutdown_rx.recv() => { + info!("模型刷新任务在首次拉取前关闭"); + return PeriodicTaskStartup::Stop; + } + () = sleep(Duration::from_millis(self.initial_delay_ms)) => {} + } + + let mut attempt = 0; + let mut has_fetched_once = false; + + while attempt < self.max_retries { + match self.refresh_once().await { + Ok(result) => { + info!( + refreshed_plans = result.refreshed_plans, + model_count = result.model_count, + "首次模型刷新成功" + ); + has_fetched_once = true; + break; + } + Err(error) => { + attempt += 1; + if matches!(error, ModelServiceError::NoAccounts) { + warn!("首次模型刷新跳过：没有可用账户，将进入周期重试"); + break; + } + + warn!( + attempt, + max_retries = self.max_retries, + error = ?error, + "首次模型刷新失败" + ); + + if attempt < self.max_retries { + tokio::select! { + _ = shutdown_rx.recv() => { + info!("模型刷新任务在重试等待期间关闭"); + return PeriodicTaskStartup::Stop; + } + () = sleep(Duration::from_millis(self.retry_delay_ms)) => {} + } + } + } + } + } + + if !has_fetched_once { + warn!("模型刷新任务首次尝试全部失败，切换到周期刷新"); + } + + PeriodicTaskStartup::Continue + }) + } + + fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> { + Box::pin(async move { + if let Err(error) = self.refresh_once().await { + warn!(error = ?error, "刷新模型列表失败"); + } + }) + } +} diff --git a/backend/src/runtime/tasks/periodic.rs b/backend/src/runtime/tasks/periodic.rs new file mode 100644 index 000000000..002f3b597 --- /dev/null +++ b/backend/src/runtime/tasks/periodic.rs @@ -0,0 +1,105 @@ +//! 周期性后台任务运行器。 + +use std::{future::Future, pin::Pin, time::Duration}; + +use tokio::sync::mpsc; +use tokio::time::interval; +use tracing::info; + +use super::coordinator::SchedulerHandle; + +pub(crate) type TaskFuture<'a, T> = Pin + Send + 'a>>; + +/// 周期任务进入循环前的决策。 +pub(crate) enum PeriodicTaskStartup { + /// 继续进入周期循环。 + Continue, + /// 任务已自行结束，不再进入周期循环。 + Stop, +} + +/// 周期任务调度配置。 +pub(crate) struct PeriodicTaskConfig { + interval_secs: u64, + started_message: &'static str, + stopped_message: &'static str, + run_immediately: bool, +} + +impl PeriodicTaskConfig { + pub(crate) fn new( + interval_secs: u64, + started_message: &'static str, + stopped_message: &'static str, + ) -> Self { + Self { + interval_secs, + started_message, + stopped_message, + run_immediately: true, + } + } + + pub(crate) fn wait_first_interval(mut self) -> Self { + self.run_immediately = false; + self + } +} + +/// 周期任务行为。 +pub(crate) trait PeriodicTaskRunner: Send + 'static { + fn before_loop<'a>( + &'a mut self, + _shutdown_rx: &'a mut mpsc::Receiver<()>, + ) -> TaskFuture<'a, PeriodicTaskStartup> { + Box::pin(std::future::ready(PeriodicTaskStartup::Continue)) + } + + fn tick(&mut self) -> TaskFuture<'_, ()>; + + fn shutdown(&mut self) -> TaskFuture<'_, ()> { + Box::pin(std::future::ready(())) + } +} + +/// 启动周期任务，并统一处理关闭信号与 JoinHandle。 +pub(crate) fn spawn_periodic_task(mut task: T, config: PeriodicTaskConfig) -> SchedulerHandle +where + T: PeriodicTaskRunner, +{ + let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1); + + let handle = tokio::spawn(async move { + info!( + interval_secs = config.interval_secs, + "{}", config.started_message + ); + + if matches!( + task.before_loop(&mut shutdown_rx).await, + PeriodicTaskStartup::Stop + ) { + return; + } + + let mut ticker = interval(Duration::from_secs(config.interval_secs)); + if !config.run_immediately { + ticker.tick().await; + } + + loop { + tokio::select! { + _ = ticker.tick() => { + task.tick().await; + } + _ = shutdown_rx.recv() => { + task.shutdown().await; + info!("{}", config.stopped_message); + break; + } + } + } + }); + + SchedulerHandle::new(shutdown_tx, handle) +} diff --git a/backend/src/runtime/tasks/quota_refresh.rs b/backend/src/runtime/tasks/quota_refresh.rs index 61ea76dd3..3402d44c6 100644 --- a/backend/src/runtime/tasks/quota_refresh.rs +++ b/backend/src/runtime/tasks/quota_refresh.rs @@ -1,8 +1,8 @@ //! 配额刷新后台任务接线器。 -use std::{collections::HashMap, sync::Arc, time::Duration}; +use std::{collections::HashMap, sync::Arc}; -use tokio::time::interval; +use tokio::time::Instant; use tracing::{debug, info, warn}; use crate::{ @@ -15,12 +15,16 @@ use crate::{ upstream::transport::CodexBackendClient, }; -use super::coordinator::SchedulerHandle; +use super::{ + coordinator::SchedulerHandle, + periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner}, +}; /// 主动配额刷新后台任务。 pub struct QuotaRefreshTask { service: RuntimeQuotaRefreshService, interval_secs: u64, + last_refreshed: HashMap, } impl QuotaRefreshTask { @@ -38,6 +42,7 @@ impl QuotaRefreshTask { min_refresh_interval_secs, ), interval_secs, + last_refreshed: HashMap::new(), } } @@ -61,42 +66,39 @@ impl QuotaRefreshTask { /// 启动后台刷新任务。 pub fn start(self) -> SchedulerHandle { - let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); + let config = PeriodicTaskConfig::new( + self.interval_secs, + "quota 刷新任务已启动", + "quota 刷新任务已关闭", + ); + spawn_periodic_task(self, config) + } +} - tokio::spawn(async move { - let mut ticker = interval(Duration::from_secs(self.interval_secs)); - let mut last_refreshed = HashMap::new(); - info!(interval_secs = self.interval_secs, "quota 刷新任务已启动"); - - loop { - tokio::select! { - _ = ticker.tick() => { - match self.service.refresh_locked_accounts(&mut last_refreshed).await { - Ok(summary) if summary.refreshed > 0 => { - log_refreshed_summary(summary); - } - Ok(summary) => { - debug!( - candidates = summary.candidates, - skipped_recent = summary.skipped_recent, - failed = summary.failed, - "没有需要刷新的 quota 锁定或待验证账号" - ); - } - Err(error) => { - warn!(error = %error, "quota 刷新任务失败"); - } - } - } - _ = shutdown_rx.recv() => { - info!("quota 刷新任务已关闭"); - break; - } +impl PeriodicTaskRunner for QuotaRefreshTask { + fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> { + Box::pin(async move { + match self + .service + .refresh_locked_accounts(&mut self.last_refreshed) + .await + { + Ok(summary) if summary.refreshed > 0 => { + log_refreshed_summary(summary); + } + Ok(summary) => { + debug!( + candidates = summary.candidates, + skipped_recent = summary.skipped_recent, + failed = summary.failed, + "没有需要刷新的 quota 锁定或待验证账号" + ); + } + Err(error) => { + warn!(error = %error, "quota 刷新任务失败"); } } - }); - - SchedulerHandle::new(shutdown_tx) + }) } } diff --git a/backend/src/runtime/tasks/token_refresh.rs b/backend/src/runtime/tasks/token_refresh.rs index 0b024b535..b1aef1e44 100644 --- a/backend/src/runtime/tasks/token_refresh.rs +++ b/backend/src/runtime/tasks/token_refresh.rs @@ -1,8 +1,5 @@ //! 令牌刷新后台任务接线器。 -use std::time::Duration; - -use tokio::time::interval; use tracing::{debug, info, warn}; use crate::upstream::accounts::{ @@ -12,7 +9,10 @@ use crate::upstream::accounts::{ }, }; -use super::coordinator::SchedulerHandle; +use super::{ + coordinator::SchedulerHandle, + periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner}, +}; const DEFAULT_INTERVAL_SECS: u64 = 60; @@ -53,46 +53,48 @@ where /// 启动后台任务。 pub fn start(self) -> SchedulerHandle { - let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1); + let config = PeriodicTaskConfig::new( + self.interval_secs, + "token 刷新任务已启动", + "token 刷新任务已关闭", + ); + spawn_periodic_task(self, config) + } +} - tokio::spawn(async move { - let mut ticker = interval(Duration::from_secs(self.interval_secs)); - info!(interval_secs = self.interval_secs, "token 刷新任务已启动"); - - loop { - tokio::select! { - _ = ticker.tick() => { - match self.service.schedule_account_timers_once().await { - Ok(summary) if summary.changed() > 0 => { - info!( - scheduled = summary.scheduled, - immediate = summary.immediate, - recovery_scheduled = summary.recovery_scheduled, - replaced = summary.replaced, - "token 刷新定时器已调度" - ); - } - Ok(summary) => { - debug!( - scanned = summary.scanned, - skipped = summary.skipped, - "没有需要调度刷新的 token" - ); - } - Err(error) => { - warn!(error = %error, "token 刷新定时器调度失败"); - } - } - } - _ = shutdown_rx.recv() => { - self.service.clear_scheduled_timers().await; - info!("token 刷新任务已关闭"); - break; - } +impl PeriodicTaskRunner for TokenRefreshTask +where + C: TokenRefresher, +{ + fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> { + Box::pin(async move { + match self.service.schedule_account_timers_once().await { + Ok(summary) if summary.changed() > 0 => { + info!( + scheduled = summary.scheduled, + immediate = summary.immediate, + recovery_scheduled = summary.recovery_scheduled, + replaced = summary.replaced, + "token 刷新定时器已调度" + ); + } + Ok(summary) => { + debug!( + scanned = summary.scanned, + skipped = summary.skipped, + "没有需要调度刷新的 token" + ); + } + Err(error) => { + warn!(error = %error, "token 刷新定时器调度失败"); } } - }); + }) + } - SchedulerHandle::new(shutdown_tx) + fn shutdown(&mut self) -> super::periodic::TaskFuture<'_, ()> { + Box::pin(async move { + self.service.clear_scheduled_timers().await; + }) } } diff --git a/backend/src/upstream/accounts/cookies.rs b/backend/src/upstream/accounts/cookies.rs index c9b4b1f50..9752ee167 100644 --- a/backend/src/upstream/accounts/cookies.rs +++ b/backend/src/upstream/accounts/cookies.rs @@ -8,26 +8,6 @@ use thiserror::Error; use tokio::sync::RwLock; use uuid::Uuid; -/// 单个 Cookie 条目。 -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct CookieEntry { - /// Cookie 名称。 - pub name: String, - /// Cookie 值。 - pub value: String, -} - -/// 将 Cookie 条目组装为 HTTP Cookie 头。 -pub fn cookie_header(cookies: &[CookieEntry]) -> Option { - let header = cookies - .iter() - .filter(|cookie| !cookie.name.trim().is_empty()) - .map(|cookie| format!("{}={}", cookie.name.trim(), cookie.value)) - .collect::>() - .join("; "); - (!header.is_empty()).then_some(header) -} - const DEFAULT_COOKIE_DOMAIN: &str = "chatgpt.com"; const CAPTURABLE_COOKIES: &[&str] = &["cf_clearance"]; diff --git a/backend/src/upstream/accounts/importing.rs b/backend/src/upstream/accounts/importing.rs index 98f72d0a1..dfdb9c319 100644 --- a/backend/src/upstream/accounts/importing.rs +++ b/backend/src/upstream/accounts/importing.rs @@ -141,61 +141,55 @@ fn explicit_account_import_source( fn parse_cpr_payload(payload: &Value) -> Result { Ok(ParsedAccountImport { source: AccountImportSource::Cpr, - entries: parse_cpr_entries(payload)?, + entries: parse_account_entries( + payload, + cpr_account_import_entry_from_value, + Some(ACCOUNT_IMPORT_CONTAINER_KEYS), + )?, }) } -fn parse_cpr_entries(payload: &Value) -> Result, &'static str> { - if let Some(accounts) = payload.get("accounts") { - ensure_account_import_keys(payload, ACCOUNT_IMPORT_CONTAINER_KEYS)?; - let accounts = accounts.as_array().ok_or("no importable accounts")?; - return parse_entries(accounts, cpr_account_import_entry_from_value); - } - if let Some(accounts) = payload.as_array() { - return parse_entries(accounts, cpr_account_import_entry_from_value); - } - cpr_account_import_entry_from_value(payload).map(option_to_vec) -} - fn parse_sub2api_payload(payload: &Value) -> Result { Ok(ParsedAccountImport { source: AccountImportSource::Sub2Api, - entries: parse_sub2api_entries(payload)?, + entries: parse_account_entries(payload, sub2api_account_import_entry_from_value, None)?, }) } -fn parse_sub2api_entries(payload: &Value) -> Result, &'static str> { - if let Some(accounts) = payload.get("accounts") { - let accounts = accounts.as_array().ok_or("no importable accounts")?; - return parse_entries(accounts, sub2api_account_import_entry_from_value); - } - if let Some(accounts) = payload.as_array() { - return parse_entries(accounts, sub2api_account_import_entry_from_value); - } - sub2api_account_import_entry_from_value(payload).map(option_to_vec) -} - fn parse_cli_proxy_api_payload(payload: &Value) -> Result { Ok(ParsedAccountImport { source: AccountImportSource::CliProxyApi, - entries: parse_cli_proxy_api_entries(payload)?, + entries: parse_account_entries( + payload, + cli_proxy_api_account_import_entry_from_value, + None, + )?, }) } -fn parse_cli_proxy_api_entries(payload: &Value) -> Result, &'static str> { +type AccountImportEntryParser = fn(&Value) -> Result, &'static str>; + +fn parse_account_entries( + payload: &Value, + parser: AccountImportEntryParser, + container_keys: Option<&[&str]>, +) -> Result, &'static str> { if let Some(accounts) = payload.get("accounts") { + if let Some(container_keys) = container_keys { + ensure_account_import_keys(payload, container_keys)?; + } let accounts = accounts.as_array().ok_or("no importable accounts")?; - return parse_entries(accounts, cli_proxy_api_account_import_entry_from_value); + return parse_entries(accounts, parser); } if let Some(accounts) = payload.as_array() { - return parse_entries(accounts, cli_proxy_api_account_import_entry_from_value); + return parse_entries(accounts, parser); } - cli_proxy_api_account_import_entry_from_value(payload).map(option_to_vec) + parser(payload).map(option_to_vec) } fn parse_entries( accounts: &[Value], - parser: fn(&Value) -> Result, &'static str>, + parser: AccountImportEntryParser, ) -> Result, &'static str> { let mut entries = Vec::new(); for account in accounts { diff --git a/backend/src/upstream/accounts/model.rs b/backend/src/upstream/accounts/model.rs index 81154e74d..5e5389eef 100644 --- a/backend/src/upstream/accounts/model.rs +++ b/backend/src/upstream/accounts/model.rs @@ -20,16 +20,23 @@ pub enum AccountStatus { Banned, } +impl AccountStatus { + /// 返回账号状态的持久化/API 字符串。 + pub fn as_str(self) -> &'static str { + match self { + Self::Active => "active", + Self::Expired => "expired", + Self::QuotaExhausted => "quota_exhausted", + Self::Refreshing => "refreshing", + Self::Disabled => "disabled", + Self::Banned => "banned", + } + } +} + impl std::fmt::Display for AccountStatus { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Active => write!(f, "active"), - Self::Expired => write!(f, "expired"), - Self::QuotaExhausted => write!(f, "quota_exhausted"), - Self::Refreshing => write!(f, "refreshing"), - Self::Disabled => write!(f, "disabled"), - Self::Banned => write!(f, "banned"), - } + f.write_str(self.as_str()) } } diff --git a/backend/src/upstream/accounts/pool.rs b/backend/src/upstream/accounts/pool.rs index c0706ab3e..2d0b947b3 100644 --- a/backend/src/upstream/accounts/pool.rs +++ b/backend/src/upstream/accounts/pool.rs @@ -24,18 +24,6 @@ use crate::upstream::protocol::events::{ parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage, }; -/// Token usage statistics. -#[derive(Debug, Clone, Default, serde::Serialize)] -pub struct TokenUsage { - pub input_tokens: i64, - pub output_tokens: i64, - pub cached_tokens: i64, - pub reasoning_tokens: i64, - pub total_tokens: i64, - pub image_input_tokens: i64, - pub image_output_tokens: i64, -} - /// 账号轮转策略。 #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RotationStrategy { @@ -1288,6 +1276,9 @@ impl RuntimeAccountPoolService { let limit_reached = quota_snapshot_limit_reached(quota); let reset_at = quota_snapshot_reset_at(quota); let cooldown_until = limit_reached.then_some(reset_at).flatten(); + let limit_window_seconds = reset_at.and_then(|_| { + crate::upstream::accounts::quota::quota_snapshot_limit_window_seconds(quota) + }); let quota_json = quota.to_string(); let persisted = match self .store @@ -1304,15 +1295,8 @@ impl RuntimeAccountPoolService { false } }; - let in_memory = - self.pool - .lock() - .await - .apply_quota_state(account_id, limit_reached, cooldown_until); if let Some(reset_at) = reset_at { - let limit_window_seconds = - crate::upstream::accounts::quota::quota_snapshot_limit_window_seconds(quota); if let Err(error) = self .store .sync_rate_limit_window(account_id, reset_at, limit_window_seconds) @@ -1324,11 +1308,12 @@ impl RuntimeAccountPoolService { "failed to persist verified quota window" ); } - self.pool.lock().await.sync_rate_limit_window( - account_id, - reset_at, - limit_window_seconds, - ); + } + + let mut pool = self.pool.lock().await; + let in_memory = pool.apply_quota_state(account_id, limit_reached, cooldown_until); + if let Some(reset_at) = reset_at { + pool.sync_rate_limit_window(account_id, reset_at, limit_window_seconds); } persisted || in_memory diff --git a/backend/src/upstream/accounts/store/mod.rs b/backend/src/upstream/accounts/store/mod.rs index 1ee19e786..064078bfa 100644 --- a/backend/src/upstream/accounts/store/mod.rs +++ b/backend/src/upstream/accounts/store/mod.rs @@ -21,7 +21,7 @@ mod rows; use queries::*; use rows::{ count_account_metadata, get_pool_account, list_pool_accounts, map_account_store_error, - metadata_from_row, model_usage_from_row, optional_positive_i64_to_u64, optional_update_value, + metadata_from_row, optional_positive_i64_to_u64, optional_update_value, push_account_metadata_search, quota_plan_type, quota_snapshot_from_row, sqlite_usage_delta, status_to_db, stored_account_from_row, to_page, u64_to_i64_saturating, }; @@ -196,27 +196,6 @@ impl AccountMetadataUpdate { } } -/// 账号模型维度用量记录。 -#[derive(Debug, Clone)] -pub struct AccountModelUsageRecord { - /// 账号 ID。 - pub account_id: String, - /// 模型 ID。 - pub model: String, - /// 历史请求总数。 - pub request_count: i64, - /// 历史错误数。 - pub error_count: i64, - /// 累计输入 token。 - pub input_tokens: i64, - /// 累计输出 token。 - pub output_tokens: i64, - /// 累计缓存 token。 - pub cached_tokens: i64, - /// 最近使用时间。 - pub last_used_at: Option>, -} - /// 账号配额快照。 #[derive(Debug, Clone)] pub struct AccountQuotaSnapshot { @@ -322,13 +301,6 @@ pub trait AccountStore: Send + Sync + 'static { /// 读取账号当前配额 JSON。 async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult>; - /// 更新账号当前配额 JSON。 - async fn update_quota_json( - &self, - account_id: &str, - quota_json: &str, - ) -> AccountStoreResult; - /// 应用已经验证过的账号配额快照。 async fn apply_quota_snapshot( &self, @@ -871,14 +843,6 @@ impl SqliteAccountStore { Ok(()) } - /// 列出模型维度用量。 - pub async fn list_model_usage(&self) -> SqliteAccountStoreResult> { - let rows = sqlx::query(LIST_MODEL_USAGE_SQL) - .fetch_all(&self.pool) - .await?; - rows.iter().map(model_usage_from_row).collect() - } - /// 读取配额快照列表。 pub async fn list_quota_snapshots( &self, @@ -1121,16 +1085,6 @@ impl AccountStore for SqliteAccountStore { .map_err(|error| map_account_store_error(&error)) } - async fn update_quota_json( - &self, - account_id: &str, - quota_json: &str, - ) -> AccountStoreResult { - SqliteAccountStore::update_quota_json(self, account_id, quota_json) - .await - .map_err(|error| map_account_store_error(&error)) - } - async fn apply_quota_snapshot( &self, account_id: &str, diff --git a/backend/src/upstream/accounts/store/queries.rs b/backend/src/upstream/accounts/store/queries.rs index 400cc5dac..01b7b91d0 100644 --- a/backend/src/upstream/accounts/store/queries.rs +++ b/backend/src/upstream/accounts/store/queries.rs @@ -327,19 +327,6 @@ on conflict(account_id, model) do update set cached_tokens = cached_tokens + excluded.cached_tokens, last_used_at = excluded.last_used_at"; -pub(super) const LIST_MODEL_USAGE_SQL: &str = r" -select - account_id, - model, - request_count, - error_count, - input_tokens, - output_tokens, - cached_tokens, - last_used_at -from account_model_usage -order by account_id asc, request_count desc, last_used_at desc, model asc"; - pub(super) const LIST_QUOTA_SNAPSHOTS_SQL: &str = r" select id, diff --git a/backend/src/upstream/accounts/store/rows.rs b/backend/src/upstream/accounts/store/rows.rs index ab9d63f0c..a520e763c 100644 --- a/backend/src/upstream/accounts/store/rows.rs +++ b/backend/src/upstream/accounts/store/rows.rs @@ -17,9 +17,8 @@ use crate::{ use super::{ queries::{GET_POOL_ACCOUNT_SQL, LIST_POOL_ACCOUNTS_SQL}, - AccountModelUsageRecord, AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore, - SqliteAccountStoreError, SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata, - UsageDelta, + AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore, SqliteAccountStoreError, + SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata, UsageDelta, }; pub(super) async fn list_pool_accounts( @@ -230,23 +229,6 @@ pub(super) fn quota_plan_type(quota_json: &str) -> Option { .map(ToString::to_string) } -pub(super) fn model_usage_from_row( - row: &sqlx::sqlite::SqliteRow, -) -> SqliteAccountStoreResult { - Ok(AccountModelUsageRecord { - account_id: row.get("account_id"), - model: row.get("model"), - request_count: row.get("request_count"), - error_count: row.get("error_count"), - input_tokens: row.get("input_tokens"), - output_tokens: row.get("output_tokens"), - cached_tokens: row.get("cached_tokens"), - last_used_at: parse_optional_rfc3339( - row.get::, _>("last_used_at").as_deref(), - )?, - }) -} - pub(super) fn quota_snapshot_from_row( row: &sqlx::sqlite::SqliteRow, ) -> SqliteAccountStoreResult { @@ -261,14 +243,7 @@ pub(super) fn quota_snapshot_from_row( } pub(super) fn status_to_db(status: AccountStatus) -> &'static str { - match status { - AccountStatus::Active => "active", - AccountStatus::Expired => "expired", - AccountStatus::QuotaExhausted => "quota_exhausted", - AccountStatus::Refreshing => "refreshing", - AccountStatus::Disabled => "disabled", - AccountStatus::Banned => "banned", - } + status.as_str() } pub(super) fn optional_update_value(value: &Option>) -> Option<&str> { diff --git a/backend/src/upstream/accounts/token_refresh/mod.rs b/backend/src/upstream/accounts/token_refresh/mod.rs index 13ae602d3..5c5682a04 100644 --- a/backend/src/upstream/accounts/token_refresh/mod.rs +++ b/backend/src/upstream/accounts/token_refresh/mod.rs @@ -259,9 +259,6 @@ pub enum RefreshFailure { /// 刷新令牌无效或已过期。 #[error("refresh token is invalid or expired")] InvalidGrant, - /// 账号配额耗尽。 - #[error("account quota is exhausted")] - QuotaExhausted, /// 账号被上游封禁。 #[error("account is banned")] Banned, @@ -350,7 +347,6 @@ where match self.client.refresh(refresh_token).await { Ok(token_pair) => Ok(apply_token_pair(account, token_pair)), - Err(RefreshFailure::QuotaExhausted) => Err(RefreshError::Transport), Err(RefreshFailure::RetryableTransport) => Err(RefreshError::RetryableTransport), Err(RefreshFailure::Transport) => Err(RefreshError::Transport), Err(error) => Ok(apply_refresh_failure(account, error)), @@ -411,7 +407,6 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco let mut updated = account.clone(); updated.status = match failure { RefreshFailure::InvalidGrant => AccountStatus::Expired, - RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted, RefreshFailure::Banned => AccountStatus::Banned, RefreshFailure::RetryableTransport => AccountStatus::Active, RefreshFailure::Transport => AccountStatus::Active, diff --git a/backend/src/upstream/fingerprint.rs b/backend/src/upstream/fingerprint.rs index 614bb502a..0a3392427 100644 --- a/backend/src/upstream/fingerprint.rs +++ b/backend/src/upstream/fingerprint.rs @@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize}; use sqlx::{Row, SqlitePool}; use thiserror::Error; use tokio::sync::Mutex; -use tracing::{info, warn}; +use tracing::info; use uuid::Uuid; /// 运行时当前指纹槽位。 @@ -17,7 +17,7 @@ pub const CURRENT_FINGERPRINT_ID: &str = "current"; const AUTO_UPDATE_SOURCE: &str = "auto_update"; const CONFIG_SEED_SOURCE: &str = "config_seed"; const APPCAST_TIMEOUT: Duration = Duration::from_secs(30); -const POLL_INTERVAL: Duration = Duration::from_hours(72); +pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_hours(72); type AppcastFields = (Option, Option, Option); @@ -442,29 +442,6 @@ impl UpdateChecker { Ok(true) } - /// 启动后台轮询任务。 - pub fn start_background_checker(self) -> tokio::task::JoinHandle<()> { - tokio::spawn(async move { - let mut ticker = tokio::time::interval(POLL_INTERVAL); - - info!( - interval_secs = POLL_INTERVAL.as_secs(), - "fingerprint 后台版本检查器已启动" - ); - - if let Err(error) = self.check_and_apply_update().await { - warn!(error = %error, "fingerprint 首次版本检查失败"); - } - - loop { - ticker.tick().await; - if let Err(error) = self.check_and_apply_update().await { - warn!(error = %error, "fingerprint 定期版本检查失败"); - } - } - }) - } - async fn fetch_appcast(&self) -> Result { let response = self .client diff --git a/backend/src/upstream/protocol/websocket.rs b/backend/src/upstream/protocol/websocket.rs index 17d70c93c..fc92c5583 100644 --- a/backend/src/upstream/protocol/websocket.rs +++ b/backend/src/upstream/protocol/websocket.rs @@ -48,17 +48,6 @@ pub struct WebSocketAuditArtifact { pub opening: Option, /// 首个 `response.create` payload 快照。 pub payload: Option, - /// 失败快照。 - pub error: Option, -} - -/// WebSocket audit error snapshot. -#[derive(Debug, Clone, Serialize, PartialEq, Eq)] -pub struct WebSocketAuditErrorSnapshot { - /// 错误分类。 - pub classification: String, - /// 错误消息。 - pub message: String, } /// 从 WebSocket 错误帧推导出的上游错误分类。 @@ -978,7 +967,6 @@ pub fn websocket_audit_artifact_from_attempt( fallback_allowed: http_sse_fallback_allowed(request), opening: Some(opening), payload: Some(payload), - error: None, } } diff --git a/backend/src/upstream/transport/client.rs b/backend/src/upstream/transport/client.rs index 81b0c7631..0c37cecb9 100644 --- a/backend/src/upstream/transport/client.rs +++ b/backend/src/upstream/transport/client.rs @@ -24,7 +24,7 @@ use uuid::Uuid; use crate::upstream::fingerprint::Fingerprint; use crate::upstream::models::backend_entry::BackendModelEntry; -use crate::upstream::protocol::events::extract_sse_usage; +use crate::upstream::protocol::events::{extract_sse_usage, retry_after_seconds_from_body}; use crate::upstream::protocol::responses::{ http_sse_fallback_allowed, transport_for_request, CodexCompactRequest, CodexResponsesRequest, CodexTransport, @@ -1012,18 +1012,6 @@ pub(super) fn retry_after_seconds(headers: &HeaderMap, body: Option<&str>) -> Op .or_else(|| body.and_then(retry_after_seconds_from_body)) } -pub(super) fn retry_after_seconds_from_body(body: &str) -> Option { - let value: Value = serde_json::from_str(body).ok()?; - value - .get("retry_after_seconds") - .or_else(|| value.get("retry_after")) - .and_then(|v| { - v.as_u64() - .or_else(|| v.as_str().and_then(|s| s.parse::().ok())) - }) - .filter(|seconds| *seconds > 0) -} - pub(super) fn truncate_for_error(body: &str) -> String { body.chars().take(200).collect() } diff --git a/backend/src/upstream/transport/usage.rs b/backend/src/upstream/transport/usage.rs index 3ad10cbd7..33c5f1878 100644 --- a/backend/src/upstream/transport/usage.rs +++ b/backend/src/upstream/transport/usage.rs @@ -1,11 +1,10 @@ use reqwest::StatusCode; use serde_json::Value; +use crate::upstream::protocol::events::retry_after_seconds_from_body; + use super::{ - client::{ - read_capped_error_body, retry_after_seconds, retry_after_seconds_from_body, - truncate_for_error, - }, + client::{read_capped_error_body, retry_after_seconds, truncate_for_error}, endpoints::usage_endpoint_urls, CodexBackendClient, CodexClientError, CodexClientResult, CodexRequestContext, }; diff --git a/backend/src/upstream/transport/websocket.rs b/backend/src/upstream/transport/websocket.rs index c07d88dee..0f320e136 100644 --- a/backend/src/upstream/transport/websocket.rs +++ b/backend/src/upstream/transport/websocket.rs @@ -46,6 +46,8 @@ const REDACTED_HEADER_VALUE: &str = ""; const CODEX_RESPONSES_PATH: &str = "/codex/responses"; const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits"; const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20); +const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60); +const WEBSOCKET_STREAM_BUFFER: usize = 16; /// WebSocket audit artifact 输出目录环境变量。 pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR"; @@ -110,8 +112,6 @@ pub struct CodexWebSocketExchange { pub set_cookie_headers: Vec, /// 上游握手响应里的限流头。 pub rate_limit_headers: Vec<(String, String)>, - /// 打开握手响应状态码。 - pub handshake_status: u16, /// 首个有效上游 WebSocket 事件到达代理的耗时。 pub first_token_ms: Option, /// WebSocket 连接池决策。 @@ -132,8 +132,6 @@ pub struct CodexWebSocketStreamingExchange { pub rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, /// 上游内部 metadata 事件里的动态 turn state。 pub turn_state_update: CodexWebSocketTurnStateUpdate, - /// 打开握手响应状态码。 - pub handshake_status: u16, /// WebSocket 连接池决策。 pub pool_decision: Option, } @@ -678,18 +676,18 @@ async fn first_reused_stream_text( websocket: &mut CodexWsStream, ) -> Result { loop { - let message = - match next_websocket_message(websocket, Some(WEBSOCKET_RECEIVE_IDLE_TIMEOUT)).await { - Ok(Some(message)) => message, - Ok(None) => { - return Err( - CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { - message: "websocket closed".to_string(), - }, - ); - } - Err(error) => return Err(reused_connection_died_before_first_frame(&error)), - }; + let message = match next_websocket_message(websocket, WEBSOCKET_RECEIVE_IDLE_TIMEOUT).await + { + Ok(Some(message)) => message, + Ok(None) => { + return Err( + CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame { + message: "websocket closed".to_string(), + }, + ); + } + Err(error) => return Err(reused_connection_died_before_first_frame(&error)), + }; match message { Message::Text(text) => return Ok(text.to_string()), Message::Close(_) => { @@ -738,18 +736,16 @@ async fn collect_websocket_response( let mut first_token_ms = None; loop { - let message = match next_websocket_message( - &mut websocket, - receive_timeout_before_first_response_frame(saw_response_frame), - ) - .await - { - Ok(message) => message, - Err(error) if reused_connection && !saw_response_frame => { - return Err(reused_connection_died_before_first_frame(&error)); - } - Err(error) => return Err(error), - }; + let message = + match next_websocket_message(&mut websocket, receive_idle_timeout(saw_response_frame)) + .await + { + Ok(message) => message, + Err(error) if reused_connection && !saw_response_frame => { + return Err(reused_connection_died_before_first_frame(&error)); + } + Err(error) => return Err(error), + }; let Some(message) = message else { break; }; @@ -808,7 +804,6 @@ async fn collect_websocket_response( turn_state: metadata.turn_state.clone(), set_cookie_headers: metadata.set_cookie_headers.clone(), rate_limit_headers: metadata.rate_limit_headers.clone(), - handshake_status: metadata.handshake_status, first_token_ms, pool_decision: None, }; @@ -842,7 +837,6 @@ fn websocket_connection_metadata( turn_state: response_meta::turn_state(response.headers()), set_cookie_headers: response_meta::set_cookie_headers(response.headers()), rate_limit_headers: response_meta::rate_limit_headers(response.headers()), - handshake_status: response.status().as_u16(), } } @@ -868,7 +862,7 @@ fn stream_websocket_response( let rate_limit_header_updates_for_task = Arc::clone(&rate_limit_header_updates); let turn_state_update = Arc::new(Mutex::new(metadata.turn_state.clone())); let turn_state_update_for_task = Arc::clone(&turn_state_update); - let (tx, rx) = mpsc::unbounded(); + let (tx, rx) = mpsc::channel(WEBSOCKET_STREAM_BUFFER); tokio::spawn(async move { forward_websocket_response_stream( websocket, @@ -889,7 +883,6 @@ fn stream_websocket_response( rate_limit_headers: response_metadata.rate_limit_headers, rate_limit_header_updates, turn_state_update, - handshake_status: response_metadata.handshake_status, pool_decision: None, } } @@ -901,7 +894,7 @@ async fn forward_websocket_response_stream( prefetched_text: Option, rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates, turn_state_update: CodexWebSocketTurnStateUpdate, - tx: mpsc::UnboundedSender>, + mut tx: mpsc::Sender>, ) { let mut pool_return = pool_return; let mut prefetched_text = prefetched_text; @@ -912,14 +905,14 @@ async fn forward_websocket_response_stream( } else { let message = match next_websocket_message( &mut websocket, - receive_timeout_before_first_response_frame(saw_response_frame), + receive_idle_timeout(saw_response_frame), ) .await { Ok(message) => message, Err(error) => { discard_stream_websocket(websocket, pool_return).await; - let _ = tx.unbounded_send(Err(error)); + let _ = tx.send(Err(error)).await; return; } }; @@ -930,14 +923,16 @@ async fn forward_websocket_response_stream( Message::Text(text) => text.to_string(), Message::Binary(_) => { discard_stream_websocket(websocket, pool_return).await; - let _ = - tx.unbounded_send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent)); + let _ = tx + .send(Err(CodexWebSocketExchangeError::UnexpectedBinaryEvent)) + .await; return; } Message::Close(_) => { discard_stream_websocket(websocket, pool_return).await; - let _ = - tx.unbounded_send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)); + let _ = tx + .send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)) + .await; return; } _ => continue, @@ -947,26 +942,32 @@ async fn forward_websocket_response_stream( discard_stream_websocket(websocket, pool_return).await; let retry_after_seconds = retry_after_seconds_from_wrapped_error_headers(&raw) .or_else(|| events::retry_after_seconds_from_body(&raw)); - let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::Upstream { - status_code: classified.status_code, - retry_after_seconds, - body: raw, - set_cookie_headers: Vec::new(), - })); + let _ = tx + .send(Err(CodexWebSocketExchangeError::Upstream { + status_code: classified.status_code, + retry_after_seconds, + body: raw, + set_cookie_headers: Vec::new(), + })) + .await; return; } if let Some(reason) = websocket_incomplete_response_reason(&raw) { discard_stream_websocket(websocket, pool_return).await; - let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::IncompleteResponse { - reason, - })); + let _ = tx + .send(Err(CodexWebSocketExchangeError::IncompleteResponse { + reason, + })) + .await; return; } if let Some(message) = websocket_response_completed_parse_error(&raw) { discard_stream_websocket(websocket, pool_return).await; - let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::InvalidCompletedResponse { - message, - })); + let _ = tx + .send(Err(CodexWebSocketExchangeError::InvalidCompletedResponse { + message, + })) + .await; return; } if let Some(headers) = websocket_rate_limit_event_headers(&raw) { @@ -984,7 +985,7 @@ async fn forward_websocket_response_stream( continue; }; let terminal = event.as_deref().is_some_and(is_terminal_websocket_event); - if tx.unbounded_send(Ok(Bytes::from(frame))).is_err() { + if tx.send(Ok(Bytes::from(frame))).await.is_err() { discard_stream_websocket(websocket, pool_return).await; return; } @@ -996,25 +997,23 @@ async fn forward_websocket_response_stream( } discard_stream_websocket(websocket, pool_return).await; - let _ = tx.unbounded_send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)); + let _ = tx + .send(Err(CodexWebSocketExchangeError::ClosedBeforeTerminal)) + .await; } -fn receive_timeout_before_first_response_frame(saw_response_frame: bool) -> Option { +fn receive_idle_timeout(saw_response_frame: bool) -> Duration { if saw_response_frame { - None + WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT } else { - Some(WEBSOCKET_RECEIVE_IDLE_TIMEOUT) + WEBSOCKET_RECEIVE_IDLE_TIMEOUT } } async fn next_websocket_message( websocket: &mut CodexWsStream, - receive_timeout: Option, + receive_timeout: Duration, ) -> Result, CodexWebSocketExchangeError> { - let Some(receive_timeout) = receive_timeout else { - return websocket.next().await.transpose().map_err(Into::into); - }; - match timeout(receive_timeout, websocket.next()).await { Ok(message) => message.transpose().map_err(Into::into), Err(_) => Err(CodexWebSocketExchangeError::ReceiveIdleTimeout { diff --git a/backend/src/upstream/transport/websocket_pool.rs b/backend/src/upstream/transport/websocket_pool.rs index d0d9ac7f5..c235587d1 100644 --- a/backend/src/upstream/transport/websocket_pool.rs +++ b/backend/src/upstream/transport/websocket_pool.rs @@ -401,7 +401,6 @@ pub(crate) struct CodexWebSocketConnectionMetadata { pub(crate) turn_state: Option, pub(crate) set_cookie_headers: Vec, pub(crate) rate_limit_headers: Vec<(String, String)>, - pub(crate) handshake_status: u16, } pub(crate) struct PooledWebSocketConnection { diff --git a/backend/tests/admin/accounts/importing.rs b/backend/tests/admin/accounts/importing.rs index f38e2d272..128fd27fd 100644 --- a/backend/tests/admin/accounts/importing.rs +++ b/backend/tests/admin/accounts/importing.rs @@ -32,7 +32,6 @@ async fn admin_accounts_import_should_store_cpr_account_tokens() { let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); @@ -642,7 +641,6 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); @@ -859,7 +857,6 @@ async fn admin_accounts_import_should_require_admin_session_cookie() { let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); diff --git a/backend/tests/admin/accounts/lifecycle.rs b/backend/tests/admin/accounts/lifecycle.rs index 6fcdc5570..eb3637e3c 100644 --- a/backend/tests/admin/accounts/lifecycle.rs +++ b/backend/tests/admin/accounts/lifecycle.rs @@ -25,7 +25,6 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() { let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); @@ -768,7 +767,6 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() { let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state.clone()); @@ -812,7 +810,6 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); diff --git a/backend/tests/admin/accounts/list.rs b/backend/tests/admin/accounts/list.rs index 4ed5523a9..d883c9476 100644 --- a/backend/tests/admin/accounts/list.rs +++ b/backend/tests/admin/accounts/list.rs @@ -29,7 +29,6 @@ async fn admin_accounts_list_should_not_expose_account_tokens() { let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); diff --git a/backend/tests/admin/accounts/mod.rs b/backend/tests/admin/accounts/mod.rs index 24c434ced..062c6bd05 100644 --- a/backend/tests/admin/accounts/mod.rs +++ b/backend/tests/admin/accounts/mod.rs @@ -172,7 +172,6 @@ async fn admin_accounts_test_app_with_overrides( let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state.clone()); diff --git a/backend/tests/admin/auth/session.rs b/backend/tests/admin/auth/session.rs index 3a2155639..f069a7f32 100644 --- a/backend/tests/admin/auth/session.rs +++ b/backend/tests/admin/auth/session.rs @@ -205,7 +205,6 @@ async fn admin_login_test_app( fingerprint, )); let state = codex_proxy_rs::runtime::state::AppState { - config, services: (*services).clone(), }; ( diff --git a/backend/tests/admin/keys/mod.rs b/backend/tests/admin/keys/mod.rs index 8d799f8b6..fb4cff736 100644 --- a/backend/tests/admin/keys/mod.rs +++ b/backend/tests/admin/keys/mod.rs @@ -43,7 +43,6 @@ async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, tempfile::Te let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; ( diff --git a/backend/tests/admin/monitoring/dashboard.rs b/backend/tests/admin/monitoring/dashboard.rs index 4673a02e9..5d04e03be 100644 --- a/backend/tests/admin/monitoring/dashboard.rs +++ b/backend/tests/admin/monitoring/dashboard.rs @@ -397,7 +397,6 @@ async fn dashboard_test_app( }; let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; ( diff --git a/backend/tests/admin/monitoring/usage_record_routes.rs b/backend/tests/admin/monitoring/usage_record_routes.rs index f49bcc2a9..bde6a9dd6 100644 --- a/backend/tests/admin/monitoring/usage_record_routes.rs +++ b/backend/tests/admin/monitoring/usage_record_routes.rs @@ -885,7 +885,6 @@ async fn admin_usage_records_test_app_with_config( let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; ( diff --git a/backend/tests/admin/settings/routes.rs b/backend/tests/admin/settings/routes.rs index c8a8253f2..e272d3722 100644 --- a/backend/tests/admin/settings/routes.rs +++ b/backend/tests/admin/settings/routes.rs @@ -283,7 +283,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() { let services = Services::new(&config, stores, fingerprint); let services = std::sync::Arc::new(services); let state = AppState { - config: config.clone(), services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); @@ -417,7 +416,6 @@ async fn admin_settings_update_should_apply_runtime_services() { .await .unwrap(); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); @@ -526,7 +524,6 @@ async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() { let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); @@ -569,7 +566,6 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp let fingerprint = crate::support::fingerprint::test_fingerprint(); let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint)); let state = AppState { - config, services: (*services).clone(), }; ( diff --git a/backend/tests/admin/system/mod.rs b/backend/tests/admin/system/mod.rs index f3b0b42d3..109ceb4f6 100644 --- a/backend/tests/admin/system/mod.rs +++ b/backend/tests/admin/system/mod.rs @@ -1220,7 +1220,6 @@ async fn admin_system_test_app(db_name: &str) -> (axum::Router, tempfile::TempDi let stores = stores(pool); let services = Arc::new(Services::new(&config, stores, test_fingerprint())); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); diff --git a/backend/tests/http/web_assets/mod.rs b/backend/tests/http/web_assets/mod.rs index 3fb936aa0..206159519 100644 --- a/backend/tests/http/web_assets/mod.rs +++ b/backend/tests/http/web_assets/mod.rs @@ -203,7 +203,6 @@ async fn router_with_dist_and_database(dist: &Path, database_url: &str) -> (Rout fingerprint, )); let state = codex_proxy_rs::runtime::state::AppState { - config, services: (*services).clone(), }; ( diff --git a/backend/tests/proxy/dispatch/chat_upstream/mod.rs b/backend/tests/proxy/dispatch/chat_upstream/mod.rs index 97cc7a9a4..1292c77cb 100644 --- a/backend/tests/proxy/dispatch/chat_upstream/mod.rs +++ b/backend/tests/proxy/dispatch/chat_upstream/mod.rs @@ -283,7 +283,7 @@ async fn test_app_state_with_pool_installation_id_and_usage_record_options( .initialize_hot_path_state() .await .expect("hot path state should initialize"); - AppState { config, services } + AppState { services } } async fn test_app_with_account(base_url: String) -> (axum::Router, String, tempfile::TempDir) { diff --git a/backend/tests/proxy/openai/models_routes.rs b/backend/tests/proxy/openai/models_routes.rs index 7ff54d8fa..b3b99fb7f 100644 --- a/backend/tests/proxy/openai/models_routes.rs +++ b/backend/tests/proxy/openai/models_routes.rs @@ -170,7 +170,6 @@ async fn build_test_app( .await .expect("hot path state should initialize"); let state = AppState { - config, services: (*services).clone(), }; let app = codex_proxy_rs::http::router::router().with_state(state); diff --git a/backend/tests/proxy/openai/responses_routes.rs b/backend/tests/proxy/openai/responses_routes.rs index 15b49b4a7..95de034c2 100644 --- a/backend/tests/proxy/openai/responses_routes.rs +++ b/backend/tests/proxy/openai/responses_routes.rs @@ -154,7 +154,6 @@ async fn test_app_with_client_api_key() -> (axum::Router, String, tempfile::Temp .await .expect("hot path state should initialize"); let state = AppState { - config, services: (*services).clone(), }; ( diff --git a/backend/tests/runtime/tasks/mod.rs b/backend/tests/runtime/tasks/mod.rs index ed781131e..c302d0824 100644 --- a/backend/tests/runtime/tasks/mod.rs +++ b/backend/tests/runtime/tasks/mod.rs @@ -79,14 +79,6 @@ impl AccountStore for FakeAccountStore { Ok(None) } - async fn update_quota_json( - &self, - _account_id: &str, - _quota_json: &str, - ) -> AccountStoreResult { - Ok(false) - } - async fn apply_quota_snapshot( &self, _account_id: &str, diff --git a/backend/tests/runtime/tasks/model_refresh.rs b/backend/tests/runtime/tasks/model_refresh.rs index 89451d84c..90c11119c 100644 --- a/backend/tests/runtime/tasks/model_refresh.rs +++ b/backend/tests/runtime/tasks/model_refresh.rs @@ -192,14 +192,6 @@ impl AccountStore for PlanAccountStore { Ok(None) } - async fn update_quota_json( - &self, - _account_id: &str, - _quota_json: &str, - ) -> AccountStoreResult { - Ok(false) - } - async fn apply_quota_snapshot( &self, _account_id: &str, diff --git a/backend/tests/upstream/accounts/account_pool/selection.rs b/backend/tests/upstream/accounts/account_pool/selection.rs index cfaf5d206..a0ad33771 100644 --- a/backend/tests/upstream/accounts/account_pool/selection.rs +++ b/backend/tests/upstream/accounts/account_pool/selection.rs @@ -281,7 +281,7 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex .acquire_with(&AccountAcquireRequest::new("gpt-5.5", now)) .await .unwrap(); - let stored: ( + type StoredQuotaResetState = ( String, i64, i64, @@ -293,7 +293,20 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex i64, Option, Option, - ) = sqlx::query_as( + ); + let ( + stored_status, + stored_quota_limit_reached, + stored_quota_verify_required, + stored_quota_cooldown_until, + stored_cloudflare_cooldown_until, + stored_window_request_count, + stored_window_input_tokens, + stored_window_output_tokens, + stored_window_cached_tokens, + stored_window_reset_at, + stored_limit_window_seconds, + ): StoredQuotaResetState = sqlx::query_as( r#" select a.status, @@ -316,22 +329,29 @@ async fn runtime_account_pool_should_persist_quota_window_reset_when_cooldown_ex .fetch_one(&pool) .await .unwrap(); - let persisted_reset_at = stored - .9 + let persisted_reset_at = stored_window_reset_at .as_deref() .and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok()) .map(|value| value.with_timezone(&Utc)); assert_eq!(acquired.account.id, "acct_quota_reset"); assert!(acquired.account.quota_verify_required); - assert_eq!(stored.0, "active"); - assert_eq!(stored.1, 0); - assert_eq!(stored.2, 1); - assert!(stored.3.is_none()); - assert!(stored.4.is_none()); - assert_eq!((stored.5, stored.6, stored.7, stored.8), (0, 0, 0, 0)); + assert_eq!(stored_status, "active"); + assert_eq!(stored_quota_limit_reached, 0); + assert_eq!(stored_quota_verify_required, 1); + assert!(stored_quota_cooldown_until.is_none()); + assert!(stored_cloudflare_cooldown_until.is_none()); + assert_eq!( + ( + stored_window_request_count, + stored_window_input_tokens, + stored_window_output_tokens, + stored_window_cached_tokens, + ), + (0, 0, 0, 0) + ); assert!(persisted_reset_at.is_some_and(|reset_at| reset_at > now)); - assert_eq!(stored.10, Some(60)); + assert_eq!(stored_limit_window_seconds, Some(60)); } #[test] diff --git a/backend/tests/upstream/protocol/codex_websocket.rs b/backend/tests/upstream/protocol/codex_websocket.rs index cacf0a5c0..2540ea0d9 100644 --- a/backend/tests/upstream/protocol/codex_websocket.rs +++ b/backend/tests/upstream/protocol/codex_websocket.rs @@ -154,7 +154,6 @@ fn codex_websocket_audit_artifact_should_record_transport_opening_and_payload() artifact.payload.expect("payload").body["input"], "" ); - assert_eq!(artifact.error, None); } #[test] diff --git a/backend/tests/upstream/protocol/openai_chat.rs b/backend/tests/upstream/protocol/openai_chat.rs index 47c633858..3508f1df3 100644 --- a/backend/tests/upstream/protocol/openai_chat.rs +++ b/backend/tests/upstream/protocol/openai_chat.rs @@ -162,6 +162,35 @@ fn chat_completion_from_codex_sse_should_convert_completed_response() { assert_eq!(response["choices"][0]["message"]["content"], "hello"); } +#[test] +fn chat_completion_from_codex_sse_should_include_reasoning_when_requested() { + let body = include_str!("../../fixtures/responses/http_sse/chat_reasoning_text_completed.sse"); + + let response = codex_proxy_rs::proxy::openai::chat::chat_completion_from_codex_sse( + body, "gpt-5.5", true, None, + ) + .expect("conversion should succeed") + .expect("completed response"); + + assert_eq!( + response["choices"][0]["message"]["reasoning_content"], + "I considered the context." + ); +} + +#[test] +fn chat_completion_from_codex_sse_should_omit_reasoning_when_not_requested() { + let body = include_str!("../../fixtures/responses/http_sse/chat_reasoning_text_completed.sse"); + + let response = codex_proxy_rs::proxy::openai::chat::chat_completion_from_codex_sse( + body, "gpt-5.5", false, None, + ) + .expect("conversion should succeed") + .expect("completed response"); + + assert!(response["choices"][0]["message"]["reasoning_content"].is_null()); +} + #[test] fn chat_completion_from_codex_sse_should_convert_incomplete_response() { let body = format!( diff --git a/backend/tests/upstream/transport/http_client.rs b/backend/tests/upstream/transport/http_client.rs index 78b975486..923c7ff83 100644 --- a/backend/tests/upstream/transport/http_client.rs +++ b/backend/tests/upstream/transport/http_client.rs @@ -145,6 +145,65 @@ async fn codex_backend_client_should_cap_non_success_error_body_at_one_mib() { assert_eq!(body.len(), 1024 * 1024); } +#[tokio::test] +async fn codex_backend_client_should_parse_retry_after_from_rate_limit_error_body() { + let server = wiremock::MockServer::start().await; + wiremock::Mock::given(wiremock::matchers::method("POST")) + .and(wiremock::matchers::path("/codex/responses")) + .respond_with(wiremock::ResponseTemplate::new(429).set_body_json(json!({ + "error": { + "code": "rate_limit_exceeded", + "message": "Rate limit exceeded, try again in 12s" + } + }))) + .mount(&server) + .await; + let client = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + server.uri(), + crate::support::fingerprint::test_fingerprint(), + ); + let mut request = + codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "", + Vec::new(), + ); + request.force_http_sse = true; + + let result = client + .create_response( + &request, + CodexRequestContext { + access_token: "access-token", + account_id: Some("chatgpt-account"), + request_id: "req_http_retry_after_body", + 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, + }, + ) + .await; + + let Err(CodexClientError::Upstream { + status, + retry_after_seconds, + .. + }) = result + else { + panic!("expected upstream error"); + }; + assert_eq!(status, reqwest::StatusCode::TOO_MANY_REQUESTS); + assert_eq!(retry_after_seconds, Some(12)); +} + #[tokio::test] async fn build_reqwest_client_should_reuse_cached_connection_pool() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/backend/tests/upstream/transport/websocket.rs b/backend/tests/upstream/transport/websocket.rs index 8bf6fb2f8..edc44c4d7 100644 --- a/backend/tests/upstream/transport/websocket.rs +++ b/backend/tests/upstream/transport/websocket.rs @@ -16,7 +16,6 @@ async fn websocket_audit_artifact_should_require_explicit_directory() { ..OpeningAuditSnapshot::default() }), payload: None, - error: None, }; let disabled = write_websocket_audit_artifact_for_dir(None, &artifact) @@ -1182,6 +1181,72 @@ async fn codex_backend_client_stream_should_wait_for_terminal_after_active_webso .contains("resp_after_active_gap")); } +#[tokio::test(start_paused = true)] +async fn codex_backend_client_stream_should_timeout_when_active_websocket_stalls() { + 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 _message = websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "response.output_text.delta", + "delta": "partial" + }) + .to_string() + .into(), + )) + .await + .unwrap(); + futures::future::pending::<()>().await; + }); + let mut request = + codex_proxy_rs::upstream::protocol::responses::CodexResponsesRequest::new_http_sse( + "gpt-5.5", + "be brief", + Vec::new(), + ); + request.previous_response_id = Some("resp_previous".to_string()); + request.force_http_sse = false; + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + crate::support::fingerprint::test_fingerprint(), + ); + + let mut response = backend + .create_response_stream( + &request, + request_context("req_stream_active_stall", Some("chatgpt-account")), + ) + .await + .expect("websocket stream should open"); + let first = response + .body + .next() + .await + .expect("stream should yield partial frame") + .expect("partial frame should be valid"); + assert!(std::str::from_utf8(&first).unwrap().contains("partial")); + tokio::time::advance(Duration::from_secs(5 * 60)).await; + tokio::task::yield_now().await; + let error = response + .body + .next() + .await + .expect("stream should yield idle timeout after active stall") + .expect_err("active stall should time out"); + server.abort(); + + assert!(matches!( + error, + CodexClientError::WebSocket(CodexWebSocketExchangeError::ReceiveIdleTimeout { timeout }) + if timeout == Duration::from_secs(5 * 60) + )); +} + #[tokio::test] async fn codex_backend_client_should_use_websocket_when_previous_response_id_is_present() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); diff --git a/frontend/src/views/usage/components/UsageCostCell.vue b/frontend/src/views/usage/components/UsageCostCell.vue index 8bbb36068..87e126421 100644 --- a/frontend/src/views/usage/components/UsageCostCell.vue +++ b/frontend/src/views/usage/components/UsageCostCell.vue @@ -29,8 +29,8 @@ const billingItems = computed(() => { return [ { label: '服务档位', value: details.serviceTierDisplay, tone: 'info' }, { label: '倍率', value: details.multiplierDisplay, tone: 'info' }, + { label: '总费用', value: details.totalCostDisplay, tone: 'success' }, { label: '原始', value: details.originalCostDisplay, tone: 'default' }, - { label: '计费', value: details.billedCostDisplay, tone: 'success' }, ] }) diff --git a/frontend/src/views/usage/components/UsageRecordDetailModal.vue b/frontend/src/views/usage/components/UsageRecordDetailModal.vue index 73b8303fc..076063ce0 100644 --- a/frontend/src/views/usage/components/UsageRecordDetailModal.vue +++ b/frontend/src/views/usage/components/UsageRecordDetailModal.vue @@ -104,7 +104,6 @@ const costItems = computed(() => { { label: '输入', value: details.inputCostDisplay, mono: true }, { label: '输出', value: details.outputCostDisplay, mono: true }, { label: '缓存读取', value: details.cacheReadCostDisplay, mono: true }, - { label: '计费', value: details.billedCostDisplay, mono: true }, { label: '原始', value: details.originalCostDisplay, mono: true }, { label: '输入单价', value: details.inputPriceDisplay, mono: true }, { label: '输出单价', value: details.outputPriceDisplay, mono: true },