fix: clean up quota dispatch audit issues

Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
zyycn
2026-07-04 21:34:05 +08:00
co-authored by Codex
parent 50d9b6eef8
commit c0cd7ccbc1
66 changed files with 1989 additions and 1846 deletions
+40 -54
View File
@@ -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<AdminAccountRefreshResult> for AccountRefreshData {
Self {
id: result.id,
email: result.email,
previous_status: account_status_str(result.previous_status).to_string(),
previous_status: result.previous_status.as_str().to_string(),
result: account_refresh_outcome_str(result.outcome),
error: result.error,
duration_ms: result.duration_ms,
@@ -472,8 +472,8 @@ pub(crate) async fn accounts(
) -> Result<impl IntoResponse, AdminError> {
let limit = clamp_limit(params.page_size.or(params.limit).unwrap_or(50));
let use_numbered_page = params.page.is_some() || params.page_size.is_some();
let stats = account_list_stats(&state).await;
let summary = account_summary_data(&state, &stats).await;
let quota_by_account = quota_snapshots_by_account(&state).await;
let summary = account_summary_data(&state, &quota_by_account).await;
if use_numbered_page {
return match state
@@ -483,6 +483,7 @@ pub(crate) async fn accounts(
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await;
let page = crate::infra::json::NumberedPage {
items: page
.items
@@ -509,6 +510,7 @@ pub(crate) async fn accounts(
.await
{
Ok(page) => {
let stats = account_list_stats(&state, &page.items, &quota_by_account).await;
let page = Page {
items: page
.items
@@ -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<String, AdminAccountQuotaData>,
) -> AccountListStats {
let account_ids = accounts
.iter()
.map(|account| account.id.clone())
.collect::<Vec<_>>();
let usage_records = state
.services
.background_tasks
.accounts
.list_quota_snapshots()
.usage
.list_by_account_ids(&account_ids)
.await
.unwrap_or_default();
let 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<String, AdminAccountQuotaData>,
) -> 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<String, AdminAccountQuotaData> {
state
.services
.background_tasks
.accounts
.list_quota_snapshots()
.await
.unwrap_or_default()
.into_iter()
.map(|snapshot| {
(
snapshot.account_id,
quota_data(&snapshot.quota_json, snapshot.quota_fetched_at),
)
})
.collect()
}
async fn list_all_account_metadata(state: &AppState) -> Vec<AdminAccountMetadata> {
let mut page = 1;
let mut accounts = Vec::new();
@@ -934,26 +951,6 @@ fn account_summary_needs_attention(status: AccountStatus) -> bool {
)
}
async fn list_all_usage_records(state: &AppState) -> Vec<AdminUsageRecord> {
let mut cursor = None;
let mut records = Vec::new();
loop {
let Ok(page) = state
.services
.usage
.list(cursor, ACCOUNT_STATS_PAGE_LIMIT)
.await
else {
return Vec::new();
};
records.extend(page.items);
let Some(next_cursor) = page.next_cursor else {
return records;
};
cursor = Some(next_cursor);
}
}
async fn list_current_window_model_usage(
state: &AppState,
usage_records: &[AdminUsageRecord],
@@ -1202,14 +1199,3 @@ fn account_export_ids(value: Option<&str>) -> Vec<String> {
.map(ToString::to_string)
.collect()
}
fn account_status_str(status: crate::upstream::accounts::model::AccountStatus) -> &'static str {
match status {
crate::upstream::accounts::model::AccountStatus::Active => "active",
crate::upstream::accounts::model::AccountStatus::Expired => "expired",
crate::upstream::accounts::model::AccountStatus::QuotaExhausted => "quota_exhausted",
crate::upstream::accounts::model::AccountStatus::Refreshing => "refreshing",
crate::upstream::accounts::model::AccountStatus::Disabled => "disabled",
crate::upstream::accounts::model::AccountStatus::Banned => "banned",
}
}
@@ -113,12 +113,6 @@ pub enum AdminAccountError {
Inactive(AcctStatus),
}
#[derive(Debug, Clone)]
pub struct UpdatedAccountStatus {
pub id: String,
pub status: AcctStatus,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdminAccountRefreshOutcome {
Alive,
@@ -233,6 +227,14 @@ pub(super) struct ManualCreateTokens {
pub(super) access_token: String,
pub(super) refresh_token_for_new: Option<String>,
pub(super) refresh_token_for_existing: Option<String>,
pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims,
}
#[derive(Debug, Clone)]
pub(super) struct RefreshedAccountTokens {
pub(super) access_token: String,
pub(super) refresh_token: Option<String>,
pub(super) claims: crate::upstream::accounts::token_refresh::ManualAccountClaims,
}
#[derive(Debug, Clone)]
@@ -4,10 +4,7 @@ use super::{
contracts::{ExportedAccount, ExportedAccounts},
AdminAccountError, AdminAccountService,
};
use crate::{
infra::time::china_rfc3339_str,
upstream::accounts::{model::AccountStatus, store::StoredAccount},
};
use crate::{infra::time::china_rfc3339_str, upstream::accounts::store::StoredAccount};
impl AdminAccountService {
pub async fn export(&self, ids: Vec<String>) -> Result<ExportedAccounts, AdminAccountError> {
@@ -58,20 +55,9 @@ impl From<StoredAccount> for ExportedAccount {
access_token_expires_at: account
.access_token_expires_at
.map(|value| value.to_rfc3339()),
status: export_status_str(account.status).to_string(),
status: account.status.as_str().to_string(),
added_at: china_rfc3339_str(&account.added_at),
updated_at: china_rfc3339_str(&account.updated_at),
}
}
}
fn export_status_str(status: AccountStatus) -> &'static str {
match status {
AccountStatus::Active => "active",
AccountStatus::Expired => "expired",
AccountStatus::QuotaExhausted => "quota_exhausted",
AccountStatus::Refreshing => "refreshing",
AccountStatus::Disabled => "disabled",
AccountStatus::Banned => "banned",
}
}
+11 -38
View File
@@ -281,8 +281,7 @@ impl AdminAccountService {
token: Option<String>,
refresh_token: Option<String>,
) -> Result<Option<ResolvedImportTokens>, AdminAccountError> {
let mut refresh_token =
crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
let refresh_token = crate::upstream::accounts::importing::normalize_nonempty(refresh_token);
let Some(access_token) = crate::upstream::accounts::importing::normalize_nonempty(
token
.as_deref()
@@ -292,26 +291,12 @@ impl AdminAccountService {
return Ok(None);
};
let refreshed = self
.token_refresher
.refresh(&existing_refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(
&refreshed.access_token,
),
))
.ok_or(AdminAccountError::TokenRequired)?;
refresh_token = refreshed.refresh_token;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
.refresh_tokens_from_refresh_token(&existing_refresh_token)
.await?;
return Ok(Some(ResolvedImportTokens {
access_token,
refresh_token,
claims: Some(claims),
access_token: refreshed.access_token,
refresh_token: refreshed.refresh_token,
claims: Some(refreshed.claims),
}));
};
@@ -334,24 +319,12 @@ impl AdminAccountService {
}));
};
let refreshed = self
.token_refresher
.refresh(&existing_refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(&refreshed.access_token),
))
.ok_or(AdminAccountError::TokenRequired)?;
refresh_token = refreshed.refresh_token;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
.refresh_tokens_from_refresh_token(&existing_refresh_token)
.await?;
Ok(Some(ResolvedImportTokens {
access_token,
refresh_token,
claims: Some(claims),
access_token: refreshed.access_token,
refresh_token: refreshed.refresh_token,
claims: Some(refreshed.claims),
}))
}
async fn import_target_account_id(
+31 -86
View File
@@ -20,7 +20,7 @@ use super::{
contracts::{
parse_account_status, stored_to_admin_metadata, AdminAccountError, AdminAccountHealthCheck,
AdminAccountMetadata, AdminAccountRefreshOutcome, AdminAccountRefreshResult,
AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens, UpdatedAccountStatus,
AdminAccountUpdate, BatchDeleteAccounts, ManualCreateTokens,
},
AdminAccountService,
};
@@ -92,37 +92,32 @@ impl AdminAccountService {
.as_deref()
.map(crate::upstream::accounts::importing::normalize_bearer_token),
) {
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
ManualCreateTokens {
access_token,
refresh_token_for_new: provided_refresh_token.clone(),
refresh_token_for_existing: provided_refresh_token,
claims,
}
} else if let Some(refresh_token) = provided_refresh_token {
let token_pair = self
.token_refresher
.refresh(&refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(
&token_pair.access_token,
),
))
.ok_or(AdminAccountError::TokenRequired)?;
let refreshed = self
.refresh_tokens_from_refresh_token(&refresh_token)
.await?;
ManualCreateTokens {
access_token,
refresh_token_for_new: token_pair.refresh_token.clone(),
refresh_token_for_existing: token_pair.refresh_token,
access_token: refreshed.access_token,
refresh_token_for_new: refreshed.refresh_token.clone(),
refresh_token_for_existing: refreshed.refresh_token,
claims: refreshed.claims,
}
} else {
return Err(AdminAccountError::TokenRequired);
};
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&tokens.access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
let claims = tokens.claims;
let existing = if let Some(account_id) = claims.account_id.as_deref() {
self.store
.find_by_chatgpt_identity(account_id, claims.user_id.as_deref())
@@ -195,45 +190,6 @@ impl AdminAccountService {
.map(stored_to_admin_metadata)
.ok_or(AdminAccountError::NotFound)
}
pub async fn update_label(
&self,
account_id: &str,
label: Option<String>,
) -> Result<bool, AdminAccountError> {
if label.as_ref().is_some_and(|l| l.chars().count() > 64) {
return Err(AdminAccountError::LabelTooLong);
}
let updated = self
.store
.set_label(account_id, label)
.await
.map_err(|_| AdminAccountError::UpdateLabel)?;
if updated {
self.sync_account_pool_best_effort(account_id, "account label update")
.await;
}
Ok(updated)
}
pub async fn update_status(
&self,
account_id: &str,
status: &str,
) -> Result<Option<UpdatedAccountStatus>, AdminAccountError> {
let status = parse_account_status(status)?;
match self.store.set_status(account_id, status).await {
Ok(true) => {
self.sync_account_pool_best_effort(account_id, "account status update")
.await;
self.evict_account_websocket_pool(account_id).await;
Ok(Some(UpdatedAccountStatus {
id: account_id.to_string(),
status,
}))
}
Ok(false) => Ok(None),
Err(_) => Err(AdminAccountError::UpdateStatus),
}
}
pub async fn update_account(
&self,
account_id: &str,
@@ -336,39 +292,28 @@ impl AdminAccountService {
}
match self
.token_refresher
.refresh(refresh_token.expose_secret())
.refresh_tokens_from_refresh_token(refresh_token.expose_secret())
.await
{
Ok(tokens) => {
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(
&tokens.access_token,
),
))
.ok_or(AdminAccountError::TokenRequired)?;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
Ok(refreshed) => {
let updated = self
.store
.update_from_claims(
account_id,
crate::upstream::accounts::store::AccountClaimsUpdate {
email: claims.email,
account_id: claims.account_id.or(account.account_id),
user_id: claims.user_id,
plan_type: claims.plan_type,
access_token: SecretString::new(access_token.into()),
refresh_token: tokens
email: refreshed.claims.email,
account_id: refreshed.claims.account_id.or(account.account_id),
user_id: refreshed.claims.user_id,
plan_type: refreshed.claims.plan_type,
access_token: SecretString::new(refreshed.access_token.into()),
refresh_token: refreshed
.refresh_token
.map(|token| SecretString::new(token.into())),
access_token_expires_at: Some(claims.expires_at),
next_refresh_at: Some(
self.next_refresh_at_for_expires_at(account_id, claims.expires_at),
),
access_token_expires_at: Some(refreshed.claims.expires_at),
next_refresh_at: Some(self.next_refresh_at_for_expires_at(
account_id,
refreshed.claims.expires_at,
)),
status: crate::upstream::accounts::model::AccountStatus::Active,
},
)
@@ -387,7 +332,7 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
})
}
Err(failure) => {
Err(AdminAccountError::RefreshTokenExchange(failure)) => {
if manual_refresh_failure_is_permanent(&failure) {
let status = crate::upstream::accounts::model::AccountStatus::Expired;
let updated = self
@@ -422,6 +367,7 @@ impl AdminAccountService {
duration_ms: elapsed_millis_i64(started_at),
})
}
Err(error) => Err(error),
}
}
@@ -494,8 +440,7 @@ fn manual_refresh_failure_is_permanent(
match failure {
crate::upstream::accounts::token_refresh::RefreshFailure::InvalidGrant
| crate::upstream::accounts::token_refresh::RefreshFailure::Banned => true,
crate::upstream::accounts::token_refresh::RefreshFailure::QuotaExhausted
| crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport
crate::upstream::accounts::token_refresh::RefreshFailure::RetryableTransport
| crate::upstream::accounts::token_refresh::RefreshFailure::Transport => false,
}
}
+28 -2
View File
@@ -11,7 +11,7 @@ mod testing;
use std::sync::Arc as StdArc;
use chrono::DateTime;
use chrono::{DateTime, Utc};
use crate::{
upstream::accounts::{
@@ -27,7 +27,7 @@ use crate::{
pub use contracts::{
AdminAccountError, AdminAccountHealthCheck, AdminAccountMetadata, AdminAccountRefreshOutcome,
AdminAccountRefreshResult, AdminAccountUpdate, BatchDeleteAccounts, ExportedAccounts,
ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput, UpdatedAccountStatus,
ImportedAccounts, OAuthAuthorizeResult, OAuthExchangeInput,
};
#[derive(Clone)]
@@ -82,6 +82,32 @@ impl AdminAccountService {
)
}
async fn refresh_tokens_from_refresh_token(
&self,
refresh_token: &str,
) -> Result<contracts::RefreshedAccountTokens, AdminAccountError> {
let token_pair = self
.token_refresher
.refresh(refresh_token)
.await
.map_err(AdminAccountError::RefreshTokenExchange)?;
let access_token = crate::upstream::accounts::importing::normalize_nonempty(Some(
crate::upstream::accounts::importing::normalize_bearer_token(&token_pair.access_token),
))
.ok_or(AdminAccountError::TokenRequired)?;
let claims = crate::upstream::accounts::token_refresh::manual_account_claims(
&access_token,
Utc::now(),
)
.map_err(AdminAccountError::InvalidToken)?;
Ok(contracts::RefreshedAccountTokens {
access_token,
refresh_token: token_pair.refresh_token,
claims,
})
}
pub(crate) async fn sync_account_pool(
&self,
account_id: &str,
@@ -37,6 +37,18 @@ impl AdminUsageService {
})
}
/// 按账号 ID 批量读取账号用量。
pub async fn list_by_account_ids(
&self,
account_ids: &[String],
) -> Result<Vec<AdminUsageRecord>, AdminUsageError> {
self.store
.list_usage_by_account_ids(account_ids)
.await
.map(|records| records.into_iter().map(AdminUsageRecord::from).collect())
.map_err(|_| AdminUsageError::List)
}
/// 汇总账号用量。
pub async fn summary(&self) -> Result<AdminUsageSummary, AdminUsageError> {
self.store
@@ -245,6 +245,27 @@ impl SqliteUsageStore {
Ok(to_page(&rows, limit))
}
/// 按账号 ID 批量读取账号用量。
pub async fn list_usage_by_account_ids(
&self,
account_ids: &[String],
) -> SqliteUsageStoreResult<Vec<UsageListRecord>> {
if account_ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(LIST_USAGE_SELECT_SQL);
builder.push("\nwhere au.account_id in (");
let mut separated = builder.separated(", ");
for account_id in account_ids {
separated.push_bind(account_id);
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await?;
rows.iter().map(usage_list_from_row).collect()
}
/// 汇总账号用量。
pub async fn usage_summary(&self) -> SqliteUsageStoreResult<UsageSummary> {
let row = sqlx::query(USAGE_SUMMARY_SQL).fetch_one(&self.pool).await?;
@@ -125,7 +125,6 @@ pub(crate) struct UsageRecordCostDetailsData {
output_cost: f64,
cache_read_cost: f64,
total_cost: f64,
billed_cost: f64,
original_cost: f64,
input_price_per_mtoken: f64,
output_price_per_mtoken: f64,
@@ -137,7 +136,6 @@ pub(crate) struct UsageRecordCostDetailsData {
output_cost_display: String,
cache_read_cost_display: String,
total_cost_display: String,
billed_cost_display: String,
original_cost_display: String,
input_price_display: String,
output_price_display: String,
@@ -332,19 +330,7 @@ pub(crate) async fn usage_records_token_trend(
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
let filter = filter_from_query(query)?;
match state.services.usage_records.token_trend(filter).await {
Ok(points) => Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
points
.into_iter()
.map(UsageRecordTrendPointData::from)
.collect::<Vec<_>>(),
),
)),
Err(error) => Err(log_error(&error)),
}
usage_records_trend_response(state, query).await
}
/// `GET /api/admin/usage/records/insights/latency-trend`
@@ -353,8 +339,15 @@ pub(crate) async fn usage_records_latency_trend(
_auth: AdminAuth,
Query(query): Query<UsageRecordsQuery>,
) -> Result<impl IntoResponse, AdminError> {
usage_records_trend_response(state, query).await
}
async fn usage_records_trend_response(
state: AppState,
query: UsageRecordsQuery,
) -> Result<AdminResponse<AdminEnvelope<Vec<UsageRecordTrendPointData>>>, AdminError> {
let filter = filter_from_query(query)?;
match state.services.usage_records.latency_trend(filter).await {
match state.services.usage_records.trend(filter).await {
Ok(points) => Ok(AdminResponse::new(
StatusCode::OK,
AdminEnvelope::ok(
@@ -607,7 +600,6 @@ fn usage_cost_details(
output_cost: cost.output_cost,
cache_read_cost: cost.cache_read_cost,
total_cost: cost.total_cost,
billed_cost: cost.billed_cost,
original_cost: cost.original_cost,
input_price_per_mtoken: cost.input_price_per_mtoken,
output_price_per_mtoken: cost.output_price_per_mtoken,
@@ -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),
@@ -92,7 +92,6 @@ pub(crate) struct AdminUsageRecordCostDetails {
pub output_cost: f64,
pub cache_read_cost: f64,
pub total_cost: f64,
pub billed_cost: f64,
pub original_cost: f64,
pub input_price_per_mtoken: f64,
pub output_price_per_mtoken: f64,
@@ -237,19 +236,8 @@ impl AdminUsageRecordService {
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合 Token 趋势。
pub async fn token_trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
self.store
.trend(filter.into())
.await
.map_err(|_| AdminUsageRecordError::List)
}
/// 聚合延迟趋势。
pub async fn latency_trend(
/// 聚合使用记录趋势。
pub async fn trend(
&self,
filter: AdminUsageRecordFilter,
) -> Result<Vec<UsageRecordTrendPoint>, AdminUsageRecordError> {
@@ -343,7 +331,6 @@ pub(crate) fn usage_record_cost_details(
output_cost: breakdown.output_cost,
cache_read_cost: breakdown.cache_read_cost,
total_cost: breakdown.total_cost,
billed_cost: breakdown.total_cost,
original_cost,
input_price_per_mtoken: breakdown.input_price_per_mtoken,
output_price_per_mtoken: breakdown.output_price_per_mtoken,
+244 -222
View File
@@ -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<u16> = None;
let mut cloudflare_challenge_count = 0usize;
let mut last_cloudflare_challenge_error = None;
let mut cloudflare_path_block_count = 0usize;
let mut last_cloudflare_path_block_error = None;
let mut model_unsupported_count = 0usize;
let mut last_model_unsupported_error = None;
let mut model_unsupported_retry_used = false;
let mut 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<ExhaustedAccountRef<'_>> {
match self {
Self::QuotaExhausted {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::QuotaExhausted,
count: *count,
upstream_error,
}),
Self::RateLimited {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::RateLimited,
count: *count,
upstream_error,
}),
Self::Expired {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Expired,
count: *count,
upstream_error,
}),
Self::Disabled {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Disabled,
count: *count,
upstream_error,
}),
Self::Banned {
count,
upstream_error,
..
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Banned,
count: *count,
upstream_error,
}),
Self::CloudflareChallenge {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflareChallenge,
count: *count,
upstream_error,
}),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflarePathBlocked,
count: *count,
upstream_error,
}),
Self::ModelUnsupported {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::ModelUnsupported,
count: *count,
upstream_error,
}),
_ => None,
}
}
fn exhausted_metadata(
kind: ExhaustedAccountKind,
count: usize,
upstream_error: &str,
) -> DispatchErrorMetadata {
DispatchErrorMetadata::exhausted_ref(ExhaustedAccountRef {
kind,
count,
upstream_error,
})
}
}
// ====================================================================
@@ -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<String, Value>,
error: &ChatDispatchError,
) {
let (failure_class, exhausted_count, upstream_error, upstream_status) = match error {
ChatDispatchError::AccountStore => ("account_store", None, None, None),
ChatDispatchError::NoActiveAccount => ("no_available_accounts", None, None, None),
ChatDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
"quota_exhausted",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::RateLimited {
count,
upstream_error,
} => (
"rate_limited",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::Expired {
count,
upstream_error,
} => ("expired", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::Disabled {
count,
upstream_error,
} => ("disabled", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::Banned {
count,
upstream_error,
..
} => ("banned", Some(*count), Some(upstream_error.clone()), None),
ChatDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
"cloudflare_challenge",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
"cloudflare_path_blocked",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
"model_unsupported",
Some(*count),
Some(upstream_error.clone()),
None,
),
ChatDispatchError::Upstream(error) => {
let upstream_status = match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
};
(
"upstream",
None,
Some(upstream_error_body(error)),
upstream_status,
)
}
ChatDispatchError::InvalidSse(_) => ("invalid_sse", None, None, None),
ChatDispatchError::EmptyUpstreamResponse => ("empty_upstream_response", None, None, None),
};
object.insert(
"failureClass".to_string(),
Value::String(failure_class.to_string()),
);
if let Some(count) = exhausted_count {
object.insert("exhaustedCount".to_string(), serde_json::json!(count));
}
if let Some(error) = upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = upstream_status {
object.insert("upstreamStatus".to_string(), serde_json::json!(status));
}
}
+124
View File
@@ -1,12 +1,136 @@
//! Shared upstream error classification for dispatch routes.
use chrono::{DateTime, Duration, Utc};
use serde_json::{json, Map, Value};
use crate::{
proxy::dispatch::exhaustion::{ExhaustedAccountKind, ExhaustedAccountRef},
upstream::accounts::model::AccountStatus,
upstream::transport::{is_banned_auth_signal, CodexBackendTransport, CodexClientError},
};
#[derive(Clone, Debug)]
pub(crate) struct DispatchErrorMetadata {
pub failure_class: DispatchFailureClass,
pub exhausted_count: Option<usize>,
pub upstream_error: Option<String>,
pub upstream_status: Option<u16>,
}
#[derive(Clone, Copy, Debug)]
pub(crate) enum DispatchFailureClass {
NoAvailableAccounts,
QuotaExhausted,
RateLimited,
Expired,
Disabled,
Banned,
CloudflareChallenge,
CloudflarePathBlocked,
ModelUnsupported,
Upstream,
InvalidSse,
MissingCompleted,
EmptyUpstreamResponse,
ResponseFailed,
}
impl DispatchFailureClass {
pub(crate) fn as_str(self) -> &'static str {
match self {
Self::NoAvailableAccounts => "no_available_accounts",
Self::QuotaExhausted => "quota_exhausted",
Self::RateLimited => "rate_limited",
Self::Expired => "expired",
Self::Disabled => "disabled",
Self::Banned => "banned",
Self::CloudflareChallenge => "cloudflare_challenge",
Self::CloudflarePathBlocked => "cloudflare_path_blocked",
Self::ModelUnsupported => "model_unsupported",
Self::Upstream => "upstream",
Self::InvalidSse => "invalid_sse",
Self::MissingCompleted => "missing_completed",
Self::EmptyUpstreamResponse => "empty_upstream_response",
Self::ResponseFailed => "response_failed",
}
}
}
impl DispatchErrorMetadata {
pub(crate) fn no_available_accounts() -> Self {
Self {
failure_class: DispatchFailureClass::NoAvailableAccounts,
exhausted_count: None,
upstream_error: None,
upstream_status: None,
}
}
pub(crate) fn exhausted_ref(exhausted: ExhaustedAccountRef<'_>) -> Self {
Self {
failure_class: DispatchFailureClass::from(exhausted.kind),
exhausted_count: Some(exhausted.count),
upstream_error: Some(exhausted.upstream_error.to_string()),
upstream_status: None,
}
}
pub(crate) fn upstream(error: &CodexClientError) -> Self {
Self {
failure_class: DispatchFailureClass::Upstream,
exhausted_count: None,
upstream_error: Some(upstream_error_body(error)),
upstream_status: match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
},
}
}
pub(crate) fn simple(failure_class: DispatchFailureClass) -> Self {
Self {
failure_class,
exhausted_count: None,
upstream_error: None,
upstream_status: None,
}
}
}
impl From<ExhaustedAccountKind> for DispatchFailureClass {
fn from(kind: ExhaustedAccountKind) -> Self {
match kind {
ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted,
ExhaustedAccountKind::RateLimited => Self::RateLimited,
ExhaustedAccountKind::Expired => Self::Expired,
ExhaustedAccountKind::Disabled => Self::Disabled,
ExhaustedAccountKind::Banned => Self::Banned,
ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge,
ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked,
ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported,
}
}
}
pub(crate) fn insert_dispatch_error_metadata(
object: &mut Map<String, Value>,
metadata: DispatchErrorMetadata,
) {
object.insert(
"failureClass".to_string(),
Value::String(metadata.failure_class.as_str().to_string()),
);
if let Some(count) = metadata.exhausted_count {
object.insert("exhaustedCount".to_string(), json!(count));
}
if let Some(error) = metadata.upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = metadata.upstream_status {
object.insert("upstreamStatus".to_string(), json!(status));
}
}
pub(crate) fn is_rate_limit_upstream_error(error: &CodexClientError) -> bool {
matches!(
error,
+222
View File
@@ -0,0 +1,222 @@
//! 账号耗尽状态聚合。
use crate::upstream::accounts::model::AccountStatus;
#[derive(Clone, Copy, Debug)]
pub(crate) enum ExhaustedAccountKind {
QuotaExhausted,
RateLimited,
Expired,
Disabled,
Banned,
CloudflareChallenge,
CloudflarePathBlocked,
ModelUnsupported,
}
impl ExhaustedAccountKind {
pub(crate) fn message_reason(self) -> &'static str {
match self {
Self::QuotaExhausted => "quota-exhausted",
Self::RateLimited => "rate-limited",
Self::Expired => "expired",
Self::Disabled => "disabled",
Self::Banned => "banned",
Self::CloudflareChallenge => "cloudflare-challenge",
Self::CloudflarePathBlocked => "cloudflare-path-block",
Self::ModelUnsupported => "model-unsupported",
}
}
}
#[derive(Clone, Debug)]
pub(crate) struct ExhaustedAccount {
pub kind: ExhaustedAccountKind,
pub count: usize,
pub upstream_error: String,
pub status_code: Option<u16>,
}
#[derive(Clone, Copy, Debug)]
pub(crate) struct ExhaustedAccountRef<'a> {
pub kind: ExhaustedAccountKind,
pub count: usize,
pub upstream_error: &'a str,
}
#[derive(Default)]
pub(crate) struct AccountExhaustionTracker {
last_kind: Option<ExhaustedAccountKind>,
last_account_id: Option<String>,
quota_exhausted: ExhaustedCounter,
rate_limited: ExhaustedCounter,
expired: ExhaustedCounter,
disabled: ExhaustedCounter,
banned: ExhaustedCounter,
cloudflare_challenge: ExhaustedCounter,
cloudflare_path_blocked: ExhaustedCounter,
model_unsupported: ExhaustedCounter,
model_unsupported_retry_used: bool,
}
#[derive(Default)]
struct ExhaustedCounter {
count: usize,
upstream_error: Option<String>,
status_code: Option<u16>,
}
impl AccountExhaustionTracker {
pub(crate) fn last_account_id(&self) -> Option<&str> {
self.last_account_id.as_deref()
}
pub(crate) fn record_rate_limited(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::RateLimited,
upstream_error,
None,
);
}
pub(crate) fn record_quota_exhausted(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::QuotaExhausted,
upstream_error,
None,
);
}
pub(crate) fn record_auth_failure(
&mut self,
account_id: Option<&str>,
account_status: AccountStatus,
upstream_error: impl Into<String>,
status_code: Option<u16>,
) {
let kind = match account_status {
AccountStatus::Disabled => ExhaustedAccountKind::Disabled,
AccountStatus::Banned => ExhaustedAccountKind::Banned,
_ => ExhaustedAccountKind::Expired,
};
self.record(account_id, kind, upstream_error, status_code);
}
pub(crate) fn record_cloudflare_challenge(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::CloudflareChallenge,
upstream_error,
None,
);
}
pub(crate) fn record_cloudflare_path_blocked(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.record(
account_id,
ExhaustedAccountKind::CloudflarePathBlocked,
upstream_error,
None,
);
}
pub(crate) fn record_model_unsupported(
&mut self,
account_id: Option<&str>,
upstream_error: impl Into<String>,
) {
self.model_unsupported_retry_used = true;
self.record(
account_id,
ExhaustedAccountKind::ModelUnsupported,
upstream_error,
None,
);
}
pub(crate) fn model_unsupported_retry_exhausted(
&self,
upstream_error: impl Into<String>,
) -> Option<ExhaustedAccount> {
self.model_unsupported_retry_used.then(|| ExhaustedAccount {
kind: ExhaustedAccountKind::ModelUnsupported,
count: self.model_unsupported.count + 1,
upstream_error: upstream_error.into(),
status_code: None,
})
}
pub(crate) fn last_exhausted(&self) -> Option<ExhaustedAccount> {
let kind = self.last_kind?;
let counter = self.counter(kind);
Some(ExhaustedAccount {
kind,
count: counter.count,
upstream_error: counter.upstream_error.clone().unwrap_or_default(),
status_code: counter.status_code,
})
}
fn record(
&mut self,
account_id: Option<&str>,
kind: ExhaustedAccountKind,
upstream_error: impl Into<String>,
status_code: Option<u16>,
) {
self.last_kind = Some(kind);
if let Some(account_id) = account_id {
self.last_account_id = Some(account_id.to_string());
}
let counter = self.counter_mut(kind);
counter.count += 1;
counter.upstream_error = Some(upstream_error.into());
if let Some(status_code) = status_code {
counter.status_code = Some(status_code);
}
}
fn counter(&self, kind: ExhaustedAccountKind) -> &ExhaustedCounter {
match kind {
ExhaustedAccountKind::QuotaExhausted => &self.quota_exhausted,
ExhaustedAccountKind::RateLimited => &self.rate_limited,
ExhaustedAccountKind::Expired => &self.expired,
ExhaustedAccountKind::Disabled => &self.disabled,
ExhaustedAccountKind::Banned => &self.banned,
ExhaustedAccountKind::CloudflareChallenge => &self.cloudflare_challenge,
ExhaustedAccountKind::CloudflarePathBlocked => &self.cloudflare_path_blocked,
ExhaustedAccountKind::ModelUnsupported => &self.model_unsupported,
}
}
fn counter_mut(&mut self, kind: ExhaustedAccountKind) -> &mut ExhaustedCounter {
match kind {
ExhaustedAccountKind::QuotaExhausted => &mut self.quota_exhausted,
ExhaustedAccountKind::RateLimited => &mut self.rate_limited,
ExhaustedAccountKind::Expired => &mut self.expired,
ExhaustedAccountKind::Disabled => &mut self.disabled,
ExhaustedAccountKind::Banned => &mut self.banned,
ExhaustedAccountKind::CloudflareChallenge => &mut self.cloudflare_challenge,
ExhaustedAccountKind::CloudflarePathBlocked => &mut self.cloudflare_path_blocked,
ExhaustedAccountKind::ModelUnsupported => &mut self.model_unsupported,
}
}
}
+1
View File
@@ -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;
+202 -104
View File
@@ -1,8 +1,14 @@
use serde_json::{json, Value};
use serde_json::Value;
use thiserror::Error;
use crate::{
proxy::dispatch::errors::{upstream_error_body, upstream_error_http_status},
proxy::dispatch::{
errors::{
insert_dispatch_error_metadata, upstream_error_http_status, DispatchErrorMetadata,
DispatchFailureClass,
},
exhaustion::{ExhaustedAccount, ExhaustedAccountKind, ExhaustedAccountRef},
},
upstream::{
protocol::{responses::ResponsesSseFailure, sse::SseError},
transport::CodexClientError,
@@ -14,8 +20,6 @@ use super::sse_failure::{sse_failure_error_body, stream_failure_http_status};
/// Responses 调度错误。
#[derive(Debug, Error)]
pub enum ResponseDispatchError {
#[error("failed to list runtime accounts")]
AccountStore,
#[error("no active account is available")]
NoActiveAccount,
#[error("all accounts exhausted by quota")]
@@ -72,9 +76,47 @@ pub enum ResponseDispatchError {
}
impl ResponseDispatchError {
pub(crate) fn from_exhausted_account(exhausted: ExhaustedAccount) -> Self {
match exhausted.kind {
ExhaustedAccountKind::QuotaExhausted => Self::QuotaExhausted {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::RateLimited => Self::RateLimited {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Expired => Self::Expired {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Disabled => Self::Disabled {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::Banned => Self::Banned {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
status_code: exhausted.status_code.unwrap_or(403),
},
ExhaustedAccountKind::CloudflareChallenge => Self::CloudflareChallenge {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::CloudflarePathBlocked => Self::CloudflarePathBlocked {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
ExhaustedAccountKind::ModelUnsupported => Self::ModelUnsupported {
count: exhausted.count,
upstream_error: exhausted.upstream_error,
},
}
}
pub fn http_status_code(&self) -> u16 {
match self {
Self::NoActiveAccount | Self::AccountStore => 503,
Self::NoActiveAccount => 503,
Self::QuotaExhausted { .. } => 429,
Self::RateLimited { .. } => 429,
Self::Expired { .. } | Self::Disabled { .. } => 401,
@@ -89,6 +131,160 @@ impl ResponseDispatchError {
Self::Upstream(error) => upstream_error_http_status(error),
}
}
pub(crate) fn metadata(&self) -> DispatchErrorMetadata {
match self {
Self::NoActiveAccount => DispatchErrorMetadata::no_available_accounts(),
Self::QuotaExhausted {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::QuotaExhausted,
*count,
upstream_error,
),
Self::RateLimited {
count,
upstream_error,
} => {
Self::exhausted_metadata(ExhaustedAccountKind::RateLimited, *count, upstream_error)
}
Self::Expired {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Expired, *count, upstream_error),
Self::Disabled {
count,
upstream_error,
} => Self::exhausted_metadata(ExhaustedAccountKind::Disabled, *count, upstream_error),
Self::Banned {
count,
upstream_error,
..
} => Self::exhausted_metadata(ExhaustedAccountKind::Banned, *count, upstream_error),
Self::CloudflareChallenge {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflareChallenge,
*count,
upstream_error,
),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::CloudflarePathBlocked,
*count,
upstream_error,
),
Self::ModelUnsupported {
count,
upstream_error,
} => Self::exhausted_metadata(
ExhaustedAccountKind::ModelUnsupported,
*count,
upstream_error,
),
Self::Upstream(error) => DispatchErrorMetadata::upstream(error),
Self::InvalidSse(_) => DispatchErrorMetadata::simple(DispatchFailureClass::InvalidSse),
Self::MissingCompleted => {
DispatchErrorMetadata::simple(DispatchFailureClass::MissingCompleted)
}
Self::EmptyUpstreamResponse => {
DispatchErrorMetadata::simple(DispatchFailureClass::EmptyUpstreamResponse)
}
Self::Failed(failure) => DispatchErrorMetadata {
failure_class: DispatchFailureClass::ResponseFailed,
exhausted_count: None,
upstream_error: Some(sse_failure_error_body(failure)),
upstream_status: None,
},
}
}
pub(crate) fn exhausted_account(&self) -> Option<ExhaustedAccountRef<'_>> {
match self {
Self::QuotaExhausted {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::QuotaExhausted,
count: *count,
upstream_error,
}),
Self::RateLimited {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::RateLimited,
count: *count,
upstream_error,
}),
Self::Expired {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Expired,
count: *count,
upstream_error,
}),
Self::Disabled {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Disabled,
count: *count,
upstream_error,
}),
Self::Banned {
count,
upstream_error,
..
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::Banned,
count: *count,
upstream_error,
}),
Self::CloudflareChallenge {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflareChallenge,
count: *count,
upstream_error,
}),
Self::CloudflarePathBlocked {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::CloudflarePathBlocked,
count: *count,
upstream_error,
}),
Self::ModelUnsupported {
count,
upstream_error,
} => Some(ExhaustedAccountRef {
kind: ExhaustedAccountKind::ModelUnsupported,
count: *count,
upstream_error,
}),
_ => None,
}
}
fn exhausted_metadata(
kind: ExhaustedAccountKind,
count: usize,
upstream_error: &str,
) -> DispatchErrorMetadata {
DispatchErrorMetadata::exhausted_ref(ExhaustedAccountRef {
kind,
count,
upstream_error,
})
}
}
/// Responses live SSE body stream error.
@@ -132,103 +328,5 @@ pub(super) fn enrich_response_dispatch_error_metadata(
let Some(object) = metadata.as_object_mut() else {
return;
};
let (failure_class, exhausted_count, upstream_error, upstream_status) = match error {
ResponseDispatchError::AccountStore => ("account_store", None, None, None),
ResponseDispatchError::NoActiveAccount => ("no_available_accounts", None, None, None),
ResponseDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
"quota_exhausted",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::RateLimited {
count,
upstream_error,
} => (
"rate_limited",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::Expired {
count,
upstream_error,
} => ("expired", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::Disabled {
count,
upstream_error,
} => ("disabled", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::Banned {
count,
upstream_error,
..
} => ("banned", Some(*count), Some(upstream_error.clone()), None),
ResponseDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
"cloudflare_challenge",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
"cloudflare_path_blocked",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
"model_unsupported",
Some(*count),
Some(upstream_error.clone()),
None,
),
ResponseDispatchError::Upstream(error) => {
let upstream_status = match error {
CodexClientError::Upstream { status, .. } => Some(status.as_u16()),
_ => None,
};
(
"upstream",
None,
Some(upstream_error_body(error)),
upstream_status,
)
}
ResponseDispatchError::InvalidSse(_) => ("invalid_sse", None, None, None),
ResponseDispatchError::MissingCompleted => ("missing_completed", None, None, None),
ResponseDispatchError::EmptyUpstreamResponse => {
("empty_upstream_response", None, None, None)
}
ResponseDispatchError::Failed(failure) => (
"response_failed",
None,
Some(sse_failure_error_body(failure)),
None,
),
};
object.insert(
"failureClass".to_string(),
Value::String(failure_class.to_string()),
);
if let Some(count) = exhausted_count {
object.insert("exhaustedCount".to_string(), json!(count));
}
if let Some(error) = upstream_error {
object.insert("upstreamError".to_string(), Value::String(error));
}
if let Some(status) = upstream_status {
object.insert("upstreamStatus".to_string(), json!(status));
}
insert_dispatch_error_metadata(object, error.metadata());
}
File diff suppressed because it is too large Load Diff
@@ -125,14 +125,6 @@ impl SessionAffinityMap {
self.entries.remove(response_id).is_some()
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
fn active_entry(&self, response_id: &str, now: DateTime<Utc>) -> Option<&SessionAffinityEntry> {
self.entries
.get(response_id)
+21 -10
View File
@@ -863,11 +863,12 @@ fn response_format_text(response_format: Option<Value>) -> PreparedResponseForma
pub fn chat_completion_from_codex_sse(
body: &str,
model: &str,
_include_reasoning: bool,
include_reasoning: bool,
tuple_schema: Option<&Value>,
) -> Result<Option<Value>, 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<String, FunctionCallInfo> = BTreeMap::new();
let mut finished_call_ids: BTreeSet<String> = BTreeSet::new();
let mut tool_calls: Vec<Value> = 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,
}));
}
+91 -146
View File
@@ -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<Value>) {
fn chat_dispatch_openai_error(error: &ChatDispatchError) -> OpenAiErrorDetails {
let status = chat_dispatch_status(error);
let (message, kind) = match error {
ChatDispatchError::NoActiveAccount | ChatDispatchError::AccountStore => (
NO_ACTIVE_UPSTREAM_ACCOUNT_MESSAGE.to_owned(),
UPSTREAM_UNAVAILABLE_ERROR,
),
ChatDispatchError::Upstream(_) => (
UPSTREAM_CODEX_REQUEST_FAILED_MESSAGE.to_owned(),
UPSTREAM_ERROR,
),
ChatDispatchError::QuotaExhausted {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "quota-exhausted", upstream_error, false),
INSUFFICIENT_QUOTA_ERROR,
),
ChatDispatchError::RateLimited {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "rate-limited", upstream_error, false),
RATE_LIMIT_ERROR,
),
ChatDispatchError::Expired {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "expired", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::Disabled {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "disabled", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::Banned {
count,
upstream_error,
..
} => (
exhausted_dispatch_message(*count, "banned", upstream_error, false),
AUTHENTICATION_ERROR,
),
ChatDispatchError::CloudflareChallenge {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "cloudflare-challenge", upstream_error, false),
UPSTREAM_ERROR,
),
ChatDispatchError::CloudflarePathBlocked {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "cloudflare-path-block", upstream_error, false),
UPSTREAM_ERROR,
),
ChatDispatchError::ModelUnsupported {
count,
upstream_error,
} => (
exhausted_dispatch_message(*count, "model-unsupported", upstream_error, false),
MODEL_NOT_FOUND_ERROR,
),
ChatDispatchError::InvalidSse(_) | ChatDispatchError::EmptyUpstreamResponse => (
INVALID_UPSTREAM_CODEX_RESPONSE_MESSAGE.to_owned(),
INVALID_UPSTREAM_RESPONSE_ERROR,
),
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,
}
}
+1 -1
View File
@@ -113,7 +113,7 @@ async fn build_application(
let task_coordinator = TaskCoordinator::start(&runtime_config, &services);
let trusted_proxies = TrustedProxyConfig::from_entries(&config.server.trusted_proxies)?;
let app_state = AppState { config, services };
let app_state = AppState { services };
Ok((
crate::http::router::router_with_trusted_proxies(trusted_proxies).with_state(app_state),
-1
View File
@@ -6,7 +6,6 @@ use crate::runtime::services::Services;
/// HTTP handler 共享的运行时状态。
#[derive(Clone)]
pub struct AppState {
pub config: AppConfig,
pub services: Services,
}
+32 -36
View File
@@ -1,12 +1,14 @@
//! 通用过期数据清理任务骨架。
use std::{fmt::Display, future::Future, time::Duration};
use std::{fmt::Display, future::Future};
use chrono::{DateTime, Utc};
use tokio::time::interval;
use tracing::{debug, info, warn};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// 支持按时间清理过期记录的存储。
pub(crate) trait ExpiredCleanupStore: Send + Sync + 'static {
@@ -47,39 +49,12 @@ where
}
pub(crate) fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(
interval_secs = self.interval_secs,
"{}", self.messages.started
);
loop {
tokio::select! {
_ = ticker.tick() => {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "{}", self.messages.deleted);
}
Ok(_) => {
debug!("{}", self.messages.empty);
}
Err(error) => {
warn!(error = %error, "{}", self.messages.failed);
}
}
}
_ = shutdown_rx.recv() => {
info!("{}", self.messages.stopped);
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
let config = PeriodicTaskConfig::new(
self.interval_secs,
self.messages.started,
self.messages.stopped,
);
spawn_periodic_task(self, config)
}
pub(crate) async fn cleanup_once(&self) -> Result<u64, S::Error> {
@@ -90,3 +65,24 @@ where
self.store.delete_expired_at(now).await
}
}
impl<S> PeriodicTaskRunner for CleanupTask<S>
where
S: ExpiredCleanupStore,
{
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self.cleanup_once().await {
Ok(deleted) if deleted > 0 => {
info!(deleted, "{}", self.messages.deleted);
}
Ok(_) => {
debug!("{}", self.messages.empty);
}
Err(error) => {
warn!(error = %error, "{}", self.messages.failed);
}
}
})
}
}
+44 -15
View File
@@ -1,6 +1,6 @@
//! 后台任务协调器。
use std::{path::PathBuf, sync::Arc};
use std::{path::PathBuf, sync::Arc, time::Duration};
use tokio::task::JoinHandle;
@@ -120,22 +120,21 @@ impl TaskCoordinator {
/// 后台任务关闭句柄。
pub enum SchedulerHandle {
/// 通过 channel 发送关闭信号。
Channel(tokio::sync::mpsc::Sender<()>),
/// 直接持有 tokio 任务句柄。
JoinHandle(JoinHandle<()>),
Channel {
shutdown_tx: tokio::sync::mpsc::Sender<()>,
handle: JoinHandle<()>,
},
/// 关闭上游 WebSocket 连接池。
WebSocketPool(Arc<crate::upstream::transport::CodexWebSocketPool>),
}
impl SchedulerHandle {
/// 使用关闭 channel 构造句柄。
pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>) -> Self {
Self::Channel(shutdown_tx)
}
/// 使用 `JoinHandle` 构造句柄。
pub(crate) fn from_join_handle(handle: JoinHandle<()>) -> Self {
Self::JoinHandle(handle)
pub(crate) fn new(shutdown_tx: tokio::sync::mpsc::Sender<()>, handle: JoinHandle<()>) -> Self {
Self::Channel {
shutdown_tx,
handle,
}
}
/// 使用 WebSocket 连接池构造句柄。
@@ -146,11 +145,12 @@ impl SchedulerHandle {
/// 关闭任务。
pub async fn shutdown(self) {
match self {
Self::Channel(shutdown_tx) => {
Self::Channel {
shutdown_tx,
handle,
} => {
let _ = shutdown_tx.send(()).await;
}
Self::JoinHandle(handle) => {
handle.abort();
wait_for_task_shutdown(handle).await;
}
Self::WebSocketPool(pool) => {
pool.shutdown().await;
@@ -158,3 +158,32 @@ impl SchedulerHandle {
}
}
}
const TASK_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
async fn wait_for_task_shutdown(handle: JoinHandle<()>) {
let mut handle = handle;
let timeout = tokio::time::sleep(TASK_SHUTDOWN_TIMEOUT);
tokio::pin!(timeout);
tokio::select! {
result = &mut handle => {
log_task_shutdown_result(result);
}
() = &mut timeout => {
tracing::warn!(timeout_secs = TASK_SHUTDOWN_TIMEOUT.as_secs(), "等待后台任务关闭超时");
handle.abort();
log_task_shutdown_result(handle.await);
}
}
}
fn log_task_shutdown_result(result: Result<(), tokio::task::JoinError>) {
match result {
Ok(()) => {}
Err(error) if error.is_cancelled() => {}
Err(error) => {
tracing::warn!(error = %error, "后台任务关闭时异常退出");
}
}
}
@@ -2,9 +2,14 @@
use std::path::PathBuf;
use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker};
use tracing::warn;
use super::coordinator::SchedulerHandle;
use crate::upstream::fingerprint::{FingerprintRepository, UpdateChecker, APPCAST_POLL_INTERVAL};
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// Codex Desktop 官方 appcast 地址。
pub const CODEX_DESKTOP_APPCAST_URL: &str =
@@ -28,5 +33,35 @@ pub fn start_fingerprint_update_task(
current_version,
current_build,
);
SchedulerHandle::from_join_handle(update_checker.start_background_checker())
spawn_periodic_task(
FingerprintUpdateTask {
checker: update_checker,
first_tick: true,
},
PeriodicTaskConfig::new(
APPCAST_POLL_INTERVAL.as_secs(),
"fingerprint 后台版本检查器已启动",
"fingerprint 后台版本检查器已关闭",
),
)
}
struct FingerprintUpdateTask {
checker: UpdateChecker,
first_tick: bool,
}
impl PeriodicTaskRunner for FingerprintUpdateTask {
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
let first_tick = std::mem::take(&mut self.first_tick);
if let Err(error) = self.checker.check_and_apply_update().await {
if first_tick {
warn!(error = %error, "fingerprint 首次版本检查失败");
} else {
warn!(error = %error, "fingerprint 定期版本检查失败");
}
}
})
}
}
+1
View File
@@ -5,6 +5,7 @@ pub mod cookie_cleanup;
pub mod coordinator;
pub mod fingerprint_update;
pub mod model_refresh;
mod periodic;
pub mod quota_refresh;
pub mod session_affinity_cleanup;
pub mod session_cleanup;
+85 -81
View File
@@ -3,7 +3,8 @@
use std::{sync::Arc, time::Duration};
use chrono::Utc;
use tokio::time::{interval, sleep};
use tokio::sync::mpsc;
use tokio::time::sleep;
use tracing::{info, warn};
use crate::upstream::accounts::pool::RuntimeAccountPoolService;
@@ -11,7 +12,10 @@ use crate::upstream::models::service::{
ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError,
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup},
};
/// 模型刷新任务接线器。
pub struct ModelRefreshTask {
@@ -54,85 +58,13 @@ impl ModelRefreshTask {
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
tokio::spawn(async move {
info!("模型刷新任务已启动");
sleep(Duration::from_millis(self.initial_delay_ms)).await;
let mut attempt = 0;
let mut has_fetched_once = false;
while attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return;
}
() = async {} => {}
}
match self.refresh_once().await {
Ok(result) => {
info!(
refreshed_plans = result.refreshed_plans,
model_count = result.model_count,
"首次模型刷新成功"
);
has_fetched_once = true;
break;
}
Err(error) => {
attempt += 1;
if matches!(error, ModelServiceError::NoAccounts) {
warn!("首次模型刷新跳过:没有可用账户,将进入周期重试");
break;
}
warn!(
attempt,
max_retries = self.max_retries,
error = ?error,
"首次模型刷新失败"
);
if attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
}
}
}
}
if !has_fetched_once {
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
}
let mut ticker = interval(Duration::from_secs(self.refresh_interval_secs));
ticker.tick().await;
loop {
tokio::select! {
_ = ticker.tick() => {
if let Err(error) = self.refresh_once().await {
warn!(error = ?error, "刷新模型列表失败");
}
}
_ = shutdown_rx.recv() => {
info!("模型刷新任务正在关闭");
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx)
let config = PeriodicTaskConfig::new(
self.refresh_interval_secs,
"模型刷新任务已启动",
"模型刷新任务正在关闭",
)
.wait_first_interval();
spawn_periodic_task(self, config)
}
/// 执行一次模型刷新并同步账号池模型计划约束。
@@ -167,3 +99,75 @@ impl ModelRefreshTask {
Ok(result)
}
}
impl PeriodicTaskRunner for ModelRefreshTask {
fn before_loop<'a>(
&'a mut self,
shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> super::periodic::TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(async move {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在首次拉取前关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.initial_delay_ms)) => {}
}
let mut attempt = 0;
let mut has_fetched_once = false;
while attempt < self.max_retries {
match self.refresh_once().await {
Ok(result) => {
info!(
refreshed_plans = result.refreshed_plans,
model_count = result.model_count,
"首次模型刷新成功"
);
has_fetched_once = true;
break;
}
Err(error) => {
attempt += 1;
if matches!(error, ModelServiceError::NoAccounts) {
warn!("首次模型刷新跳过:没有可用账户,将进入周期重试");
break;
}
warn!(
attempt,
max_retries = self.max_retries,
error = ?error,
"首次模型刷新失败"
);
if attempt < self.max_retries {
tokio::select! {
_ = shutdown_rx.recv() => {
info!("模型刷新任务在重试等待期间关闭");
return PeriodicTaskStartup::Stop;
}
() = sleep(Duration::from_millis(self.retry_delay_ms)) => {}
}
}
}
}
}
if !has_fetched_once {
warn!("模型刷新任务首次尝试全部失败,切换到周期刷新");
}
PeriodicTaskStartup::Continue
})
}
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
if let Err(error) = self.refresh_once().await {
warn!(error = ?error, "刷新模型列表失败");
}
})
}
}
+105
View File
@@ -0,0 +1,105 @@
//! 周期性后台任务运行器。
use std::{future::Future, pin::Pin, time::Duration};
use tokio::sync::mpsc;
use tokio::time::interval;
use tracing::info;
use super::coordinator::SchedulerHandle;
pub(crate) type TaskFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
/// 周期任务进入循环前的决策。
pub(crate) enum PeriodicTaskStartup {
/// 继续进入周期循环。
Continue,
/// 任务已自行结束,不再进入周期循环。
Stop,
}
/// 周期任务调度配置。
pub(crate) struct PeriodicTaskConfig {
interval_secs: u64,
started_message: &'static str,
stopped_message: &'static str,
run_immediately: bool,
}
impl PeriodicTaskConfig {
pub(crate) fn new(
interval_secs: u64,
started_message: &'static str,
stopped_message: &'static str,
) -> Self {
Self {
interval_secs,
started_message,
stopped_message,
run_immediately: true,
}
}
pub(crate) fn wait_first_interval(mut self) -> Self {
self.run_immediately = false;
self
}
}
/// 周期任务行为。
pub(crate) trait PeriodicTaskRunner: Send + 'static {
fn before_loop<'a>(
&'a mut self,
_shutdown_rx: &'a mut mpsc::Receiver<()>,
) -> TaskFuture<'a, PeriodicTaskStartup> {
Box::pin(std::future::ready(PeriodicTaskStartup::Continue))
}
fn tick(&mut self) -> TaskFuture<'_, ()>;
fn shutdown(&mut self) -> TaskFuture<'_, ()> {
Box::pin(std::future::ready(()))
}
}
/// 启动周期任务,并统一处理关闭信号与 JoinHandle。
pub(crate) fn spawn_periodic_task<T>(mut task: T, config: PeriodicTaskConfig) -> SchedulerHandle
where
T: PeriodicTaskRunner,
{
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
let handle = tokio::spawn(async move {
info!(
interval_secs = config.interval_secs,
"{}", config.started_message
);
if matches!(
task.before_loop(&mut shutdown_rx).await,
PeriodicTaskStartup::Stop
) {
return;
}
let mut ticker = interval(Duration::from_secs(config.interval_secs));
if !config.run_immediately {
ticker.tick().await;
}
loop {
tokio::select! {
_ = ticker.tick() => {
task.tick().await;
}
_ = shutdown_rx.recv() => {
task.shutdown().await;
info!("{}", config.stopped_message);
break;
}
}
}
});
SchedulerHandle::new(shutdown_tx, handle)
}
+38 -36
View File
@@ -1,8 +1,8 @@
//! 配额刷新后台任务接线器。
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{collections::HashMap, sync::Arc};
use tokio::time::interval;
use tokio::time::Instant;
use tracing::{debug, info, warn};
use crate::{
@@ -15,12 +15,16 @@ use crate::{
upstream::transport::CodexBackendClient,
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
/// 主动配额刷新后台任务。
pub struct QuotaRefreshTask {
service: RuntimeQuotaRefreshService,
interval_secs: u64,
last_refreshed: HashMap<String, Instant>,
}
impl QuotaRefreshTask {
@@ -38,6 +42,7 @@ impl QuotaRefreshTask {
min_refresh_interval_secs,
),
interval_secs,
last_refreshed: HashMap::new(),
}
}
@@ -61,42 +66,39 @@ impl QuotaRefreshTask {
/// 启动后台刷新任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
let config = PeriodicTaskConfig::new(
self.interval_secs,
"quota 刷新任务已启动",
"quota 刷新任务已关闭",
);
spawn_periodic_task(self, config)
}
}
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
let mut last_refreshed = HashMap::new();
info!(interval_secs = self.interval_secs, "quota 刷新任务已启动");
loop {
tokio::select! {
_ = ticker.tick() => {
match self.service.refresh_locked_accounts(&mut last_refreshed).await {
Ok(summary) if summary.refreshed > 0 => {
log_refreshed_summary(summary);
}
Ok(summary) => {
debug!(
candidates = summary.candidates,
skipped_recent = summary.skipped_recent,
failed = summary.failed,
"没有需要刷新的 quota 锁定或待验证账号"
);
}
Err(error) => {
warn!(error = %error, "quota 刷新任务失败");
}
}
}
_ = shutdown_rx.recv() => {
info!("quota 刷新任务已关闭");
break;
}
impl PeriodicTaskRunner for QuotaRefreshTask {
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self
.service
.refresh_locked_accounts(&mut self.last_refreshed)
.await
{
Ok(summary) if summary.refreshed > 0 => {
log_refreshed_summary(summary);
}
Ok(summary) => {
debug!(
candidates = summary.candidates,
skipped_recent = summary.skipped_recent,
failed = summary.failed,
"没有需要刷新的 quota 锁定或待验证账号"
);
}
Err(error) => {
warn!(error = %error, "quota 刷新任务失败");
}
}
});
SchedulerHandle::new(shutdown_tx)
})
}
}
+43 -41
View File
@@ -1,8 +1,5 @@
//! 令牌刷新后台任务接线器。
use std::time::Duration;
use tokio::time::interval;
use tracing::{debug, info, warn};
use crate::upstream::accounts::{
@@ -12,7 +9,10 @@ use crate::upstream::accounts::{
},
};
use super::coordinator::SchedulerHandle;
use super::{
coordinator::SchedulerHandle,
periodic::{spawn_periodic_task, PeriodicTaskConfig, PeriodicTaskRunner},
};
const DEFAULT_INTERVAL_SECS: u64 = 60;
@@ -53,46 +53,48 @@ where
/// 启动后台任务。
pub fn start(self) -> SchedulerHandle {
let (shutdown_tx, mut shutdown_rx) = tokio::sync::mpsc::channel(1);
let config = PeriodicTaskConfig::new(
self.interval_secs,
"token 刷新任务已启动",
"token 刷新任务已关闭",
);
spawn_periodic_task(self, config)
}
}
tokio::spawn(async move {
let mut ticker = interval(Duration::from_secs(self.interval_secs));
info!(interval_secs = self.interval_secs, "token 刷新任务已启动");
loop {
tokio::select! {
_ = ticker.tick() => {
match self.service.schedule_account_timers_once().await {
Ok(summary) if summary.changed() > 0 => {
info!(
scheduled = summary.scheduled,
immediate = summary.immediate,
recovery_scheduled = summary.recovery_scheduled,
replaced = summary.replaced,
"token 刷新定时器已调度"
);
}
Ok(summary) => {
debug!(
scanned = summary.scanned,
skipped = summary.skipped,
"没有需要调度刷新的 token"
);
}
Err(error) => {
warn!(error = %error, "token 刷新定时器调度失败");
}
}
}
_ = shutdown_rx.recv() => {
self.service.clear_scheduled_timers().await;
info!("token 刷新任务已关闭");
break;
}
impl<C> PeriodicTaskRunner for TokenRefreshTask<C>
where
C: TokenRefresher,
{
fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
match self.service.schedule_account_timers_once().await {
Ok(summary) if summary.changed() > 0 => {
info!(
scheduled = summary.scheduled,
immediate = summary.immediate,
recovery_scheduled = summary.recovery_scheduled,
replaced = summary.replaced,
"token 刷新定时器已调度"
);
}
Ok(summary) => {
debug!(
scanned = summary.scanned,
skipped = summary.skipped,
"没有需要调度刷新的 token"
);
}
Err(error) => {
warn!(error = %error, "token 刷新定时器调度失败");
}
}
});
})
}
SchedulerHandle::new(shutdown_tx)
fn shutdown(&mut self) -> super::periodic::TaskFuture<'_, ()> {
Box::pin(async move {
self.service.clear_scheduled_timers().await;
})
}
}
-20
View File
@@ -8,26 +8,6 @@ use thiserror::Error;
use tokio::sync::RwLock;
use uuid::Uuid;
/// 单个 Cookie 条目。
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CookieEntry {
/// Cookie 名称。
pub name: String,
/// Cookie 值。
pub value: String,
}
/// 将 Cookie 条目组装为 HTTP Cookie 头。
pub fn cookie_header(cookies: &[CookieEntry]) -> Option<String> {
let header = cookies
.iter()
.filter(|cookie| !cookie.name.trim().is_empty())
.map(|cookie| format!("{}={}", cookie.name.trim(), cookie.value))
.collect::<Vec<_>>()
.join("; ");
(!header.is_empty()).then_some(header)
}
const DEFAULT_COOKIE_DOMAIN: &str = "chatgpt.com";
const CAPTURABLE_COOKIES: &[&str] = &["cf_clearance"];
+25 -31
View File
@@ -141,61 +141,55 @@ fn explicit_account_import_source(
fn parse_cpr_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::Cpr,
entries: parse_cpr_entries(payload)?,
entries: parse_account_entries(
payload,
cpr_account_import_entry_from_value,
Some(ACCOUNT_IMPORT_CONTAINER_KEYS),
)?,
})
}
fn parse_cpr_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
ensure_account_import_keys(payload, ACCOUNT_IMPORT_CONTAINER_KEYS)?;
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, cpr_account_import_entry_from_value);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, cpr_account_import_entry_from_value);
}
cpr_account_import_entry_from_value(payload).map(option_to_vec)
}
fn parse_sub2api_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::Sub2Api,
entries: parse_sub2api_entries(payload)?,
entries: parse_account_entries(payload, sub2api_account_import_entry_from_value, None)?,
})
}
fn parse_sub2api_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, sub2api_account_import_entry_from_value);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, sub2api_account_import_entry_from_value);
}
sub2api_account_import_entry_from_value(payload).map(option_to_vec)
}
fn parse_cli_proxy_api_payload(payload: &Value) -> Result<ParsedAccountImport, &'static str> {
Ok(ParsedAccountImport {
source: AccountImportSource::CliProxyApi,
entries: parse_cli_proxy_api_entries(payload)?,
entries: parse_account_entries(
payload,
cli_proxy_api_account_import_entry_from_value,
None,
)?,
})
}
fn parse_cli_proxy_api_entries(payload: &Value) -> Result<Vec<AccountImportEntry>, &'static str> {
type AccountImportEntryParser = fn(&Value) -> Result<Option<AccountImportEntry>, &'static str>;
fn parse_account_entries(
payload: &Value,
parser: AccountImportEntryParser,
container_keys: Option<&[&str]>,
) -> Result<Vec<AccountImportEntry>, &'static str> {
if let Some(accounts) = payload.get("accounts") {
if let Some(container_keys) = container_keys {
ensure_account_import_keys(payload, container_keys)?;
}
let accounts = accounts.as_array().ok_or("no importable accounts")?;
return parse_entries(accounts, cli_proxy_api_account_import_entry_from_value);
return parse_entries(accounts, parser);
}
if let Some(accounts) = payload.as_array() {
return parse_entries(accounts, cli_proxy_api_account_import_entry_from_value);
return parse_entries(accounts, parser);
}
cli_proxy_api_account_import_entry_from_value(payload).map(option_to_vec)
parser(payload).map(option_to_vec)
}
fn parse_entries(
accounts: &[Value],
parser: fn(&Value) -> Result<Option<AccountImportEntry>, &'static str>,
parser: AccountImportEntryParser,
) -> Result<Vec<AccountImportEntry>, &'static str> {
let mut entries = Vec::new();
for account in accounts {
+15 -8
View File
@@ -20,16 +20,23 @@ pub enum AccountStatus {
Banned,
}
impl AccountStatus {
/// 返回账号状态的持久化/API 字符串。
pub fn as_str(self) -> &'static str {
match self {
Self::Active => "active",
Self::Expired => "expired",
Self::QuotaExhausted => "quota_exhausted",
Self::Refreshing => "refreshing",
Self::Disabled => "disabled",
Self::Banned => "banned",
}
}
}
impl std::fmt::Display for AccountStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Active => write!(f, "active"),
Self::Expired => write!(f, "expired"),
Self::QuotaExhausted => write!(f, "quota_exhausted"),
Self::Refreshing => write!(f, "refreshing"),
Self::Disabled => write!(f, "disabled"),
Self::Banned => write!(f, "banned"),
}
f.write_str(self.as_str())
}
}
+9 -24
View File
@@ -24,18 +24,6 @@ use crate::upstream::protocol::events::{
parse_rate_limit_headers, rate_limit_quota, TokenUsage as CodexTokenUsage,
};
/// Token usage statistics.
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct TokenUsage {
pub input_tokens: i64,
pub output_tokens: i64,
pub cached_tokens: i64,
pub reasoning_tokens: i64,
pub total_tokens: i64,
pub image_input_tokens: i64,
pub image_output_tokens: i64,
}
/// 账号轮转策略。
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RotationStrategy {
@@ -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
+1 -47
View File
@@ -21,7 +21,7 @@ mod rows;
use queries::*;
use rows::{
count_account_metadata, get_pool_account, list_pool_accounts, map_account_store_error,
metadata_from_row, model_usage_from_row, optional_positive_i64_to_u64, optional_update_value,
metadata_from_row, optional_positive_i64_to_u64, optional_update_value,
push_account_metadata_search, quota_plan_type, quota_snapshot_from_row, sqlite_usage_delta,
status_to_db, stored_account_from_row, to_page, u64_to_i64_saturating,
};
@@ -196,27 +196,6 @@ impl AccountMetadataUpdate {
}
}
/// 账号模型维度用量记录。
#[derive(Debug, Clone)]
pub struct AccountModelUsageRecord {
/// 账号 ID。
pub account_id: String,
/// 模型 ID。
pub model: String,
/// 历史请求总数。
pub request_count: i64,
/// 历史错误数。
pub error_count: i64,
/// 累计输入 token。
pub input_tokens: i64,
/// 累计输出 token。
pub output_tokens: i64,
/// 累计缓存 token。
pub cached_tokens: i64,
/// 最近使用时间。
pub last_used_at: Option<DateTime<Utc>>,
}
/// 账号配额快照。
#[derive(Debug, Clone)]
pub struct AccountQuotaSnapshot {
@@ -322,13 +301,6 @@ pub trait AccountStore: Send + Sync + 'static {
/// 读取账号当前配额 JSON。
async fn get_quota_json(&self, account_id: &str) -> AccountStoreResult<Option<String>>;
/// 更新账号当前配额 JSON。
async fn update_quota_json(
&self,
account_id: &str,
quota_json: &str,
) -> AccountStoreResult<bool>;
/// 应用已经验证过的账号配额快照。
async fn apply_quota_snapshot(
&self,
@@ -871,14 +843,6 @@ impl SqliteAccountStore {
Ok(())
}
/// 列出模型维度用量。
pub async fn list_model_usage(&self) -> SqliteAccountStoreResult<Vec<AccountModelUsageRecord>> {
let rows = sqlx::query(LIST_MODEL_USAGE_SQL)
.fetch_all(&self.pool)
.await?;
rows.iter().map(model_usage_from_row).collect()
}
/// 读取配额快照列表。
pub async fn list_quota_snapshots(
&self,
@@ -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<bool> {
SqliteAccountStore::update_quota_json(self, account_id, quota_json)
.await
.map_err(|error| map_account_store_error(&error))
}
async fn apply_quota_snapshot(
&self,
account_id: &str,
@@ -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,
+3 -28
View File
@@ -17,9 +17,8 @@ use crate::{
use super::{
queries::{GET_POOL_ACCOUNT_SQL, LIST_POOL_ACCOUNTS_SQL},
AccountModelUsageRecord, AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore,
SqliteAccountStoreError, SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata,
UsageDelta,
AccountQuotaSnapshot, AccountStoreError, SqliteAccountStore, SqliteAccountStoreError,
SqliteAccountStoreResult, StoredAccount, StoredAccountMetadata, UsageDelta,
};
pub(super) async fn list_pool_accounts(
@@ -230,23 +229,6 @@ pub(super) fn quota_plan_type(quota_json: &str) -> Option<String> {
.map(ToString::to_string)
}
pub(super) fn model_usage_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountModelUsageRecord> {
Ok(AccountModelUsageRecord {
account_id: row.get("account_id"),
model: row.get("model"),
request_count: row.get("request_count"),
error_count: row.get("error_count"),
input_tokens: row.get("input_tokens"),
output_tokens: row.get("output_tokens"),
cached_tokens: row.get("cached_tokens"),
last_used_at: parse_optional_rfc3339(
row.get::<Option<String>, _>("last_used_at").as_deref(),
)?,
})
}
pub(super) fn quota_snapshot_from_row(
row: &sqlx::sqlite::SqliteRow,
) -> SqliteAccountStoreResult<AccountQuotaSnapshot> {
@@ -261,14 +243,7 @@ pub(super) fn quota_snapshot_from_row(
}
pub(super) fn status_to_db(status: AccountStatus) -> &'static str {
match status {
AccountStatus::Active => "active",
AccountStatus::Expired => "expired",
AccountStatus::QuotaExhausted => "quota_exhausted",
AccountStatus::Refreshing => "refreshing",
AccountStatus::Disabled => "disabled",
AccountStatus::Banned => "banned",
}
status.as_str()
}
pub(super) fn optional_update_value(value: &Option<Option<String>>) -> Option<&str> {
@@ -259,9 +259,6 @@ pub enum RefreshFailure {
/// 刷新令牌无效或已过期。
#[error("refresh token is invalid or expired")]
InvalidGrant,
/// 账号配额耗尽。
#[error("account quota is exhausted")]
QuotaExhausted,
/// 账号被上游封禁。
#[error("account is banned")]
Banned,
@@ -350,7 +347,6 @@ where
match self.client.refresh(refresh_token).await {
Ok(token_pair) => Ok(apply_token_pair(account, token_pair)),
Err(RefreshFailure::QuotaExhausted) => Err(RefreshError::Transport),
Err(RefreshFailure::RetryableTransport) => Err(RefreshError::RetryableTransport),
Err(RefreshFailure::Transport) => Err(RefreshError::Transport),
Err(error) => Ok(apply_refresh_failure(account, error)),
@@ -411,7 +407,6 @@ pub fn apply_refresh_failure(account: &Account, failure: RefreshFailure) -> Acco
let mut updated = account.clone();
updated.status = match failure {
RefreshFailure::InvalidGrant => AccountStatus::Expired,
RefreshFailure::QuotaExhausted => AccountStatus::QuotaExhausted,
RefreshFailure::Banned => AccountStatus::Banned,
RefreshFailure::RetryableTransport => AccountStatus::Active,
RefreshFailure::Transport => AccountStatus::Active,
+2 -25
View File
@@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize};
use sqlx::{Row, SqlitePool};
use thiserror::Error;
use tokio::sync::Mutex;
use tracing::{info, warn};
use tracing::info;
use uuid::Uuid;
/// 运行时当前指纹槽位。
@@ -17,7 +17,7 @@ pub const CURRENT_FINGERPRINT_ID: &str = "current";
const AUTO_UPDATE_SOURCE: &str = "auto_update";
const CONFIG_SEED_SOURCE: &str = "config_seed";
const APPCAST_TIMEOUT: Duration = Duration::from_secs(30);
const POLL_INTERVAL: Duration = Duration::from_hours(72);
pub const APPCAST_POLL_INTERVAL: Duration = Duration::from_hours(72);
type AppcastFields = (Option<String>, Option<String>, Option<String>);
@@ -442,29 +442,6 @@ impl UpdateChecker {
Ok(true)
}
/// 启动后台轮询任务。
pub fn start_background_checker(self) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move {
let mut ticker = tokio::time::interval(POLL_INTERVAL);
info!(
interval_secs = POLL_INTERVAL.as_secs(),
"fingerprint 后台版本检查器已启动"
);
if let Err(error) = self.check_and_apply_update().await {
warn!(error = %error, "fingerprint 首次版本检查失败");
}
loop {
ticker.tick().await;
if let Err(error) = self.check_and_apply_update().await {
warn!(error = %error, "fingerprint 定期版本检查失败");
}
}
})
}
async fn fetch_appcast(&self) -> Result<String, UpdateError> {
let response = self
.client
@@ -48,17 +48,6 @@ pub struct WebSocketAuditArtifact {
pub opening: Option<OpeningAuditSnapshot>,
/// 首个 `response.create` payload 快照。
pub payload: Option<PayloadAuditSnapshot>,
/// 失败快照。
pub error: Option<WebSocketAuditErrorSnapshot>,
}
/// WebSocket audit error snapshot.
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
pub struct WebSocketAuditErrorSnapshot {
/// 错误分类。
pub classification: String,
/// 错误消息。
pub message: String,
}
/// 从 WebSocket 错误帧推导出的上游错误分类。
@@ -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,
}
}
+1 -13
View File
@@ -24,7 +24,7 @@ use uuid::Uuid;
use crate::upstream::fingerprint::Fingerprint;
use crate::upstream::models::backend_entry::BackendModelEntry;
use crate::upstream::protocol::events::extract_sse_usage;
use crate::upstream::protocol::events::{extract_sse_usage, retry_after_seconds_from_body};
use crate::upstream::protocol::responses::{
http_sse_fallback_allowed, transport_for_request, CodexCompactRequest, CodexResponsesRequest,
CodexTransport,
@@ -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<u64> {
let value: Value = serde_json::from_str(body).ok()?;
value
.get("retry_after_seconds")
.or_else(|| value.get("retry_after"))
.and_then(|v| {
v.as_u64()
.or_else(|| v.as_str().and_then(|s| s.parse::<u64>().ok()))
})
.filter(|seconds| *seconds > 0)
}
pub(super) fn truncate_for_error(body: &str) -> String {
body.chars().take(200).collect()
}
+3 -4
View File
@@ -1,11 +1,10 @@
use reqwest::StatusCode;
use serde_json::Value;
use crate::upstream::protocol::events::retry_after_seconds_from_body;
use super::{
client::{
read_capped_error_body, retry_after_seconds, retry_after_seconds_from_body,
truncate_for_error,
},
client::{read_capped_error_body, retry_after_seconds, truncate_for_error},
endpoints::usage_endpoint_urls,
CodexBackendClient, CodexClientError, CodexClientResult, CodexRequestContext,
};
+60 -61
View File
@@ -46,6 +46,8 @@ const REDACTED_HEADER_VALUE: &str = "<redacted>";
const CODEX_RESPONSES_PATH: &str = "/codex/responses";
const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits";
const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20);
const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60);
const WEBSOCKET_STREAM_BUFFER: usize = 16;
/// WebSocket audit artifact 输出目录环境变量。
pub const WS_AUDIT_DIR_ENV: &str = "CODEX_PROXY_WS_AUDIT_DIR";
@@ -110,8 +112,6 @@ pub struct CodexWebSocketExchange {
pub set_cookie_headers: Vec<String>,
/// 上游握手响应里的限流头。
pub rate_limit_headers: Vec<(String, String)>,
/// 打开握手响应状态码。
pub handshake_status: u16,
/// 首个有效上游 WebSocket 事件到达代理的耗时。
pub first_token_ms: Option<i64>,
/// WebSocket 连接池决策。
@@ -132,8 +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<WebSocketPoolDecision>,
}
@@ -678,18 +676,18 @@ async fn first_reused_stream_text(
websocket: &mut CodexWsStream,
) -> Result<String, CodexWebSocketExchangeError> {
loop {
let message =
match next_websocket_message(websocket, Some(WEBSOCKET_RECEIVE_IDLE_TIMEOUT)).await {
Ok(Some(message)) => message,
Ok(None) => {
return Err(
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
message: "websocket closed".to_string(),
},
);
}
Err(error) => return Err(reused_connection_died_before_first_frame(&error)),
};
let message = match next_websocket_message(websocket, WEBSOCKET_RECEIVE_IDLE_TIMEOUT).await
{
Ok(Some(message)) => message,
Ok(None) => {
return Err(
CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstFrame {
message: "websocket closed".to_string(),
},
);
}
Err(error) => return Err(reused_connection_died_before_first_frame(&error)),
};
match message {
Message::Text(text) => return Ok(text.to_string()),
Message::Close(_) => {
@@ -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<String>,
rate_limit_header_updates: CodexWebSocketRateLimitHeaderUpdates,
turn_state_update: CodexWebSocketTurnStateUpdate,
tx: mpsc::UnboundedSender<Result<Bytes, CodexWebSocketExchangeError>>,
mut tx: mpsc::Sender<Result<Bytes, CodexWebSocketExchangeError>>,
) {
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<Duration> {
fn receive_idle_timeout(saw_response_frame: bool) -> Duration {
if saw_response_frame {
None
WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT
} else {
Some(WEBSOCKET_RECEIVE_IDLE_TIMEOUT)
WEBSOCKET_RECEIVE_IDLE_TIMEOUT
}
}
async fn next_websocket_message(
websocket: &mut CodexWsStream,
receive_timeout: Option<Duration>,
receive_timeout: Duration,
) -> Result<Option<Message>, CodexWebSocketExchangeError> {
let Some(receive_timeout) = receive_timeout else {
return websocket.next().await.transpose().map_err(Into::into);
};
match timeout(receive_timeout, websocket.next()).await {
Ok(message) => message.transpose().map_err(Into::into),
Err(_) => Err(CodexWebSocketExchangeError::ReceiveIdleTimeout {
@@ -401,7 +401,6 @@ pub(crate) struct CodexWebSocketConnectionMetadata {
pub(crate) turn_state: Option<String>,
pub(crate) set_cookie_headers: Vec<String>,
pub(crate) rate_limit_headers: Vec<(String, String)>,
pub(crate) handshake_status: u16,
}
pub(crate) struct PooledWebSocketConnection {
@@ -32,7 +32,6 @@ async fn admin_accounts_import_should_store_cpr_account_tokens() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -642,7 +641,6 @@ async fn admin_accounts_import_should_fetch_usage_to_complete_missing_plan_and_q
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -859,7 +857,6 @@ async fn admin_accounts_import_should_require_admin_session_cookie() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -25,7 +25,6 @@ async fn admin_accounts_lifecycle_should_update_and_delete_accounts() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -768,7 +767,6 @@ async fn admin_account_create_should_derive_claims_and_store_plain_tokens() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state.clone());
@@ -812,7 +810,6 @@ async fn admin_account_manual_create_should_reject_missing_invalid_expired_or_un
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
-1
View File
@@ -29,7 +29,6 @@ async fn admin_accounts_list_should_not_expose_account_tokens() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
-1
View File
@@ -172,7 +172,6 @@ async fn admin_accounts_test_app_with_overrides(
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state.clone());
-1
View File
@@ -205,7 +205,6 @@ async fn admin_login_test_app(
fingerprint,
));
let state = codex_proxy_rs::runtime::state::AppState {
config,
services: (*services).clone(),
};
(
-1
View File
@@ -43,7 +43,6 @@ async fn admin_client_key_test_app(db_name: &str) -> (axum::Router, tempfile::Te
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
(
@@ -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(),
};
(
@@ -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(),
};
(
-4
View File
@@ -283,7 +283,6 @@ async fn admin_settings_update_should_persist_runtime_settings_to_database() {
let services = Services::new(&config, stores, fingerprint);
let services = std::sync::Arc::new(services);
let state = AppState {
config: config.clone(),
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -417,7 +416,6 @@ async fn admin_settings_update_should_apply_runtime_services() {
.await
.unwrap();
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -526,7 +524,6 @@ async fn admin_settings_update_should_reject_unsupported_or_invalid_fields() {
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
@@ -569,7 +566,6 @@ async fn admin_settings_test_app(db_name: &str) -> (axum::Router, tempfile::Temp
let fingerprint = crate::support::fingerprint::test_fingerprint();
let services = std::sync::Arc::new(Services::new(&config, stores, fingerprint));
let state = AppState {
config,
services: (*services).clone(),
};
(
-1
View File
@@ -1220,7 +1220,6 @@ async fn admin_system_test_app(db_name: &str) -> (axum::Router, tempfile::TempDi
let stores = stores(pool);
let services = Arc::new(Services::new(&config, stores, test_fingerprint()));
let state = AppState {
config,
services: (*services).clone(),
};
let app = codex_proxy_rs::http::router::router().with_state(state);
-1
View File
@@ -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(),
};
(
@@ -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) {
@@ -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);
@@ -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(),
};
(
-8
View File
@@ -79,14 +79,6 @@ impl AccountStore for FakeAccountStore {
Ok(None)
}
async fn update_quota_json(
&self,
_account_id: &str,
_quota_json: &str,
) -> AccountStoreResult<bool> {
Ok(false)
}
async fn apply_quota_snapshot(
&self,
_account_id: &str,
@@ -192,14 +192,6 @@ impl AccountStore for PlanAccountStore {
Ok(None)
}
async fn update_quota_json(
&self,
_account_id: &str,
_quota_json: &str,
) -> AccountStoreResult<bool> {
Ok(false)
}
async fn apply_quota_snapshot(
&self,
_account_id: &str,
@@ -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<String>,
Option<i64>,
) = 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]
@@ -154,7 +154,6 @@ fn codex_websocket_audit_artifact_should_record_transport_opening_and_payload()
artifact.payload.expect("payload").body["input"],
"<redacted>"
);
assert_eq!(artifact.error, None);
}
#[test]
@@ -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!(
@@ -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();
+66 -1
View File
@@ -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();
@@ -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' },
]
})
@@ -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 },