diff --git a/backend/src/api/admin/accounts_routes/probe_routes.rs b/backend/src/api/admin/accounts_routes/probe_routes.rs index ef261a5bf..2a634859e 100644 --- a/backend/src/api/admin/accounts_routes/probe_routes.rs +++ b/backend/src/api/admin/accounts_routes/probe_routes.rs @@ -60,3 +60,29 @@ pub async fn account_models( Err(error) => Err(account_error(&error)), } } + +/// `POST /api/admin/accounts/models` +pub async fn refresh_account_models( + State(state): State, + _auth: AdminAuth, + Json(payload): Json, +) -> Result { + let models = state + .services + .admin_accounts + .refresh_account_models(&payload.id) + .await + .map_err(|error| account_error(&error))?; + Ok(AdminResponse::new( + StatusCode::OK, + AdminEnvelope::ok(AccountModelsData { + models: models + .into_iter() + .map(|model| AccountModelData { + id: model.id, + label: model.label, + }) + .collect(), + }), + )) +} diff --git a/backend/src/api/admin/accounts_routes/query.rs b/backend/src/api/admin/accounts_routes/query.rs index 2219d7215..2a6991cf2 100644 --- a/backend/src/api/admin/accounts_routes/query.rs +++ b/backend/src/api/admin/accounts_routes/query.rs @@ -513,7 +513,9 @@ pub(super) fn account_error(error: &AccountManageError) -> AdminError { | AccountManageError::OAuthCallbackInvalid | AccountManageError::OAuthStateMismatch | AccountManageError::NoValidCookies => AdminError::bad_request(error.to_string()), - AccountManageError::OAuthCodeExchange(_) => AdminError::bad_gateway(error.to_string()), + AccountManageError::OAuthCodeExchange(_) | AccountManageError::RefreshModels(_) => { + AdminError::bad_gateway(error.to_string()) + } AccountManageError::NotFound => account_not_found(), AccountManageError::Inactive(_) => AdminError::conflict(error.to_string()), _ => AdminError::internal(error.to_string()), diff --git a/backend/src/api/admin/router.rs b/backend/src/api/admin/router.rs index 0da724519..fcb31bc6f 100644 --- a/backend/src/api/admin/router.rs +++ b/backend/src/api/admin/router.rs @@ -14,7 +14,8 @@ use super::{ accounts_routes::{ account_models, account_quota, accounts, batch_delete_accounts, create_account, export_accounts, health_check_accounts, import_accounts, oauth_authorize_account, - oauth_exchange_account, refresh_account, test_account_connection, update_account, + oauth_exchange_account, refresh_account, refresh_account_models, test_account_connection, + update_account, }, auth_routes::{login, logout, session_status}, dashboard_routes::{dashboard_summary, dashboard_trend}, @@ -72,7 +73,10 @@ pub fn router() -> Router { post(oauth_exchange_account), ) .route("/api/admin/accounts/test", get(test_account_connection)) - .route("/api/admin/accounts/models", get(account_models)) + .route( + "/api/admin/accounts/models", + get(account_models).post(refresh_account_models), + ) .route("/api/admin/accounts/delete", post(batch_delete_accounts)) .route("/api/admin/accounts/update", post(update_account)) .route("/api/admin/accounts/refresh", post(refresh_account)) diff --git a/backend/src/bootstrap/tasks/model_refresh.rs b/backend/src/bootstrap/tasks/model_refresh.rs index fa874c9f4..a3282ae61 100644 --- a/backend/src/bootstrap/tasks/model_refresh.rs +++ b/backend/src/bootstrap/tasks/model_refresh.rs @@ -13,23 +13,18 @@ use crate::models::service::{ ModelRefreshPlanAccount, ModelRefreshResult, ModelService, ModelServiceError, }; -use super::{ - coordinator::SchedulerHandle, - periodic::{PeriodicTaskConfig, PeriodicTaskRunner, PeriodicTaskStartup, spawn_periodic_task}, -}; +use super::coordinator::SchedulerHandle; /// 模型刷新任务接线器。 pub struct ModelRefreshTask { model_service: Arc, account_pool: Arc, - refresh_interval_secs: u64, initial_delay_ms: u64, retry_delay_ms: u64, max_retries: u32, account_pseudonymizer: Arc, } -const DEFAULT_REFRESH_INTERVAL_SECS: u64 = 3600; const INITIAL_DELAY_MS: u64 = 1000; const RETRY_DELAY_MS: u64 = 10_000; const MAX_RETRIES: u32 = 12; @@ -44,7 +39,6 @@ impl ModelRefreshTask { Self { model_service, account_pool, - refresh_interval_secs: DEFAULT_REFRESH_INTERVAL_SECS, initial_delay_ms: INITIAL_DELAY_MS, retry_delay_ms: RETRY_DELAY_MS, max_retries: MAX_RETRIES, @@ -52,15 +46,13 @@ impl ModelRefreshTask { } } - /// 启动后台刷新任务。 + /// 启动一次性首次刷新任务。 pub fn start(self) -> SchedulerHandle { - let config = PeriodicTaskConfig::new( - self.refresh_interval_secs, - "模型刷新任务已启动", - "模型刷新任务正在关闭", - ) - .wait_first_interval(); - spawn_periodic_task(self, config) + let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1); + let handle = tokio::spawn(async move { + self.run_initial_refresh(&mut shutdown_rx).await; + }); + SchedulerHandle::new(shutdown_tx, handle) } /// 在 Responses 响应声明模型目录 ETag 变化时异步刷新。 @@ -115,84 +107,52 @@ impl ModelRefreshTask { .await; 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)) => {} + async fn run_initial_refresh(&self, shutdown_rx: &mut mpsc::Receiver<()>) { + tokio::select! { + _ = shutdown_rx.recv() => { + info!("模型刷新任务在首次拉取前关闭"); + return; } + () = sleep(Duration::from_millis(self.initial_delay_ms)) => {} + } - let mut attempt = 0; - let mut has_fetched_once = false; - let mut deferred_until_accounts_available = 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) { - info!("首次模型刷新跳过:没有可用账户,将进入周期重试"); - deferred_until_accounts_available = true; - 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 && !deferred_until_accounts_available { - warn!("模型刷新任务首次尝试全部失败,切换到周期刷新"); - } - - PeriodicTaskStartup::Continue - }) - } - - fn tick(&mut self) -> super::periodic::TaskFuture<'_, ()> { - Box::pin(async move { + let mut attempt = 0; + while attempt < self.max_retries { match self.refresh_once().await { - Ok(_) => {} + Ok(result) => { + info!( + refreshed_plans = result.refreshed_plans, + model_count = result.model_count, + "首次模型刷新成功" + ); + return; + } Err(ModelServiceError::NoAccounts) => { - info!("模型刷新跳过:没有可用账户"); + info!("首次模型刷新跳过:没有可用账户,后续由 ETag 或手动操作触发"); + return; } Err(error) => { - warn!(error = ?error, "刷新模型列表失败"); + attempt += 1; + 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)) => {} + } + } } } - }) + } + warn!("模型刷新任务首次尝试全部失败,等待 ETag 或手动操作再次触发"); } } diff --git a/backend/src/bootstrap/tasks/periodic.rs b/backend/src/bootstrap/tasks/periodic.rs index 002f3b597..ff31cd0eb 100644 --- a/backend/src/bootstrap/tasks/periodic.rs +++ b/backend/src/bootstrap/tasks/periodic.rs @@ -10,20 +10,11 @@ use super::coordinator::SchedulerHandle; pub(crate) type TaskFuture<'a, T> = Pin + Send + 'a>>; -/// 周期任务进入循环前的决策。 -pub(crate) enum PeriodicTaskStartup { - /// 继续进入周期循环。 - Continue, - /// 任务已自行结束,不再进入周期循环。 - Stop, -} - /// 周期任务调度配置。 pub(crate) struct PeriodicTaskConfig { interval_secs: u64, started_message: &'static str, stopped_message: &'static str, - run_immediately: bool, } impl PeriodicTaskConfig { @@ -36,25 +27,12 @@ impl PeriodicTaskConfig { 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<'_, ()> { @@ -75,17 +53,7 @@ where "{}", 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! { diff --git a/backend/src/dispatch/affinity/identity.rs b/backend/src/dispatch/affinity/identity.rs index a4070b65d..64e70489a 100644 --- a/backend/src/dispatch/affinity/identity.rs +++ b/backend/src/dispatch/affinity/identity.rs @@ -191,12 +191,6 @@ impl AccountIdentityService { AccountScopedRequest { request, identity } } - pub(in crate::dispatch) fn scope_auxiliary(&self, account_id: &str) -> AccountScopedIdentity { - AccountScopedIdentity { - installation_id: self.pseudonymizer.installation_id(account_id), - } - } - fn scope_turn_metadata( raw: &str, identity: &AccountScopedIdentity, diff --git a/backend/src/dispatch/controllers/account_failure.rs b/backend/src/dispatch/controllers/account_failure.rs index 33e52ebb8..7e6b5c65a 100644 --- a/backend/src/dispatch/controllers/account_failure.rs +++ b/backend/src/dispatch/controllers/account_failure.rs @@ -119,14 +119,20 @@ impl AccountFailureController { if code.contains("invalid_request") || code.contains("not_found") || code.contains("context_window") + || code == "context_length_exceeded" || code.contains("invalid_prompt") || code.contains("bad_request") { return Some(ClientFailure::new(failure.clone(), 400, false)); } - if code.contains("server_overloaded") { + if matches!(code.as_str(), "server_is_overloaded" | "slow_down") + || code.contains("server_overloaded") + { return Some(ClientFailure::new(failure.clone(), 503, true)); } + if code == "usage_not_included" { + return Some(ClientFailure::new(failure.clone(), 403, false)); + } classify_response_failure(failure).map(|classified| { let status = match classified.kind { AccountFailureKind::ModelUnsupported => 400, diff --git a/backend/src/dispatch/controllers/mod.rs b/backend/src/dispatch/controllers/mod.rs index 85c81ac88..21fc2106c 100644 --- a/backend/src/dispatch/controllers/mod.rs +++ b/backend/src/dispatch/controllers/mod.rs @@ -19,7 +19,10 @@ use crate::{ AccountIdentityScope, AccountIdentityService, AccountScopedRequest, SessionAffinityService, }, - errors::{ClientFailure, ResponseDispatchError, upstream_error_set_cookie_headers}, + errors::{ + ClientFailure, ResponseDispatchError, upstream_error_rate_limit_headers, + upstream_error_set_cookie_headers, + }, lifecycle::trace::{ResponseDispatchAttempt, ResponseDispatchTrace}, }, fleet::{account::Account, pool::AccountLease}, @@ -271,22 +274,11 @@ pub(in crate::dispatch) struct ControllerEnter { pub excluded_account_ids: BTreeSet, } -pub(in crate::dispatch) struct AttemptAccountPreparationContext<'a> { - pub codex: &'a CodexBackendClient, - pub account_identity: &'a AccountIdentityService, - pub request_id: &'a str, -} - pub(in crate::dispatch) struct PreparedAttemptAccount { pub lease: AccountLease, pub cookie_header: Option, } -pub(in crate::dispatch) enum AttemptAccountPreparation { - Ready(Box), - Rejected, -} - pub(in crate::dispatch) struct AttemptRoutePreparationContext<'a> { pub account_identity: &'a AccountIdentityService, pub account_id: &'a str, @@ -300,38 +292,16 @@ pub(in crate::dispatch) enum AttemptRoutePreparation { impl ControllerSet { pub(in crate::dispatch) async fn prepare_attempt_account( &self, - context: AttemptAccountPreparationContext<'_>, acquired: AccountLease, - ) -> AttemptAccountPreparation { - let usage_cookie = self - .cloudflare - .cookie_header_for_request(&acquired.account.id, "/codex/usage") - .await; - let acquired = match QuotaController::enter( - quota::QuotaEnterContext { - account_pool: &self.account_pool, - codex: context.codex, - cookie_header: usage_cookie.as_deref(), - account_identity: context.account_identity, - request_id: context.request_id, - }, - acquired, - ) - .await - { - quota::QuotaEnterOutcome::Ready(acquired) => *acquired, - quota::QuotaEnterOutcome::LimitReached => { - return AttemptAccountPreparation::Rejected; - } - }; + ) -> PreparedAttemptAccount { let cookie_header = self .cloudflare .cookie_header_for_request(&acquired.account.id, "/codex/responses") .await; - AttemptAccountPreparation::Ready(Box::new(PreparedAttemptAccount { + PreparedAttemptAccount { lease: acquired, cookie_header, - })) + } } pub(in crate::dispatch) fn prepare_attempt_route( @@ -396,19 +366,24 @@ impl ControllerSet { &self, context: AttemptUpstreamErrorContext<'_>, ) { - let _ = best_effort_controller_io( + let cookie = best_effort_controller_io( "cookie.attempt_error", self.cloudflare.capture_set_cookie_headers( &context.account.id, upstream_error_set_cookie_headers(context.error), ), - ) - .await; - let _ = best_effort_controller_io( + ); + let rate_limit_headers = upstream_error_rate_limit_headers(context.error); + let quota = best_effort_controller_io( + "quota.attempt_error_headers", + self.account_pool + .sync_passive_rate_limit_headers(context.account, &rate_limit_headers), + ); + let telemetry = best_effort_controller_io( "telemetry.attempt_error", telemetry::TelemetryController::observe_upstream_error(&self.recorder, &context), - ) - .await; + ); + let _ = tokio::join!(cookie, quota, telemetry); } pub(in crate::dispatch) async fn observe_dispatch_error( @@ -944,9 +919,6 @@ impl ControllerSet { AttemptObservationKind::NoCandidate { .. } => { AttemptDecision::Return(AttemptReturnKind::Observed) } - AttemptObservationKind::CandidatePreparationRejected => { - unreachable!("candidate preparation controller owns it") - } AttemptObservationKind::RoutePreparationRejected { message } => { AttemptDecision::Return(AttemptReturnKind::RouteUnavailable { message: message.clone(), diff --git a/backend/src/dispatch/controllers/quota.rs b/backend/src/dispatch/controllers/quota.rs index 07ceb8286..af53b2a28 100644 --- a/backend/src/dispatch/controllers/quota.rs +++ b/backend/src/dispatch/controllers/quota.rs @@ -1,139 +1,29 @@ -//! 请求前配额验证规则 owner。 - -use serde_json::Value; +//! 配额失败规则 owner。 use crate::{ dispatch::{ - affinity::AccountIdentityService, controllers::ControllerFailureFact, errors::ClientFailure, failure::exhaustion::{AccountExhaustionRecord, ExhaustedAccountKind}, - lifecycle::contract::{ - AttemptDecision, AttemptObservation, AttemptObservationKind, AttemptReturnKind, - }, + lifecycle::contract::{AttemptDecision, AttemptObservation, AttemptReturnKind}, }, fleet::{ account_failure::{ AccountFailureKind, AccountStateEffect, apply_account_state_effect_immediately, classify_response_failure, classify_upstream_failure, }, - pool::{AccountLease, AccountPoolService}, - quota::{ - quota_from_usage, quota_snapshot_limit_reached, quota_snapshot_limit_window_seconds, - quota_snapshot_reset_at, - }, + pool::AccountPoolService, }, upstream::openai::{ - failure::UpstreamFailureFacts, - protocol::responses::ResponsesSseFailure, - transport::{CodexBackendClient, CodexRequestContext}, + failure::UpstreamFailureFacts, protocol::responses::ResponsesSseFailure, + transport::CodexBackendClient, }, }; -const LIMIT_REACHED_MESSAGE: &str = "Upstream usage quota still reports limit_reached"; - -pub(super) enum QuotaEnterOutcome { - Ready(Box), - LimitReached, -} - -pub(super) struct QuotaEnterContext<'a> { - pub account_pool: &'a AccountPoolService, - pub codex: &'a CodexBackendClient, - pub account_identity: &'a AccountIdentityService, - pub request_id: &'a str, - pub cookie_header: Option<&'a str>, -} - pub(super) struct QuotaController; impl QuotaController { - pub(super) async fn enter( - context: QuotaEnterContext<'_>, - acquired: AccountLease, - ) -> QuotaEnterOutcome { - if !acquired.account.quota_verify_required { - return QuotaEnterOutcome::Ready(Box::new(acquired)); - } - - let account_id = acquired.account.id.clone(); - let identity = context.account_identity.scope_auxiliary(&account_id); - let usage = context - .codex - .fetch_usage(CodexRequestContext { - access_token: &acquired.account.access_token, - account_id: acquired.account.account_id.as_deref(), - request_id: context.request_id, - turn_state: None, - turn_metadata: None, - beta_features: None, - include_timing_metrics: None, - version: None, - codex_window_id: None, - parent_thread_id: None, - cookie_header: context.cookie_header, - installation_id: Some(&identity.installation_id), - session_id: None, - thread_id: None, - client_request_id: None, - turn_id: None, - }) - .await; - - let raw = match usage { - Ok(raw) => raw, - Err(error) => { - tracing::warn!( - request_id = %context.request_id, - account_id = %account_id, - quota_verify_required = true, - quota_verify_result = "upstream_error", - retry_with_another_account = false, - error = %error, - "Failed to verify stale quota state before upstream request" - ); - return QuotaEnterOutcome::Ready(Box::new(acquired)); - } - }; - - let quota = quota_from_usage(&raw); - context - .account_pool - .apply_quota_snapshot(&account_id, "a) - .await; - if quota_snapshot_limit_reached("a) { - acquired.release_without_usage().await; - tracing::info!( - request_id = %context.request_id, - account_id = %account_id, - quota_verify_required = true, - quota_verify_result = "limit_reached", - retry_with_another_account = true, - "Quota verification reported exhausted account before upstream request" - ); - return QuotaEnterOutcome::LimitReached; - } - - QuotaEnterOutcome::Ready(Box::new(acquired_with_verified_quota(acquired, "a))) - } - pub(super) fn classify(observation: &AttemptObservation) -> Option { - if matches!( - observation.kind, - AttemptObservationKind::CandidatePreparationRejected - ) { - return observation - .account - .as_ref() - .map(|account| ClassifiedQuotaFailure { - exhaustion: AccountExhaustionRecord::new( - account.id.clone(), - ExhaustedAccountKind::RateLimited, - LIMIT_REACHED_MESSAGE, - ), - effect: None, - }); - } classify(observation) } @@ -229,19 +119,3 @@ fn classified_quota_failure( effect: classified.effect, }) } - -fn acquired_with_verified_quota(mut acquired: AccountLease, quota: &Value) -> AccountLease { - let limit_reached = quota_snapshot_limit_reached(quota); - acquired.account.quota_verify_required = false; - acquired.account.quota_limit_reached = limit_reached; - acquired.account.quota_cooldown_until = limit_reached - .then_some(quota_snapshot_reset_at(quota)) - .flatten(); - if let Some(reset_at) = quota_snapshot_reset_at(quota) { - acquired.account.window_reset_at = Some(reset_at); - if let Some(limit_window_seconds) = quota_snapshot_limit_window_seconds(quota) { - acquired.account.limit_window_seconds = Some(limit_window_seconds); - } - } - acquired -} diff --git a/backend/src/dispatch/errors.rs b/backend/src/dispatch/errors.rs index f4b28c200..e548ddfd4 100644 --- a/backend/src/dispatch/errors.rs +++ b/backend/src/dispatch/errors.rs @@ -9,8 +9,13 @@ use crate::{ sse::sse_failure_error_body, }, upstream::openai::{ - protocol::{responses::ResponsesSseFailure, sse::SseError}, - transport::{CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics}, + protocol::{ + events::is_rate_limit_header_name, responses::ResponsesSseFailure, sse::SseError, + }, + transport::{ + CodexBackendTransport, CodexClientError, CodexUpstreamDiagnostics, + websocket::CodexWebSocketExchangeError, + }, }, }; @@ -238,6 +243,21 @@ pub(crate) fn upstream_error_set_cookie_headers(error: &CodexClientError) -> &[S } } +pub(crate) fn upstream_error_rate_limit_headers(error: &CodexClientError) -> Vec<(String, String)> { + match error { + CodexClientError::Upstream { + rate_limit_headers, .. + } => rate_limit_headers.clone(), + CodexClientError::WebSocket(CodexWebSocketExchangeError::Upstream(upstream)) => upstream + .headers + .iter() + .filter(|(name, _)| is_rate_limit_header_name(name)) + .cloned() + .collect(), + _ => Vec::new(), + } +} + pub(crate) fn upstream_error_http_status(error: &CodexClientError) -> u16 { match error { CodexClientError::Upstream { status, .. } => status.as_u16(), diff --git a/backend/src/dispatch/lifecycle/attempt.rs b/backend/src/dispatch/lifecycle/attempt.rs index e86d58970..f024ce16c 100644 --- a/backend/src/dispatch/lifecycle/attempt.rs +++ b/backend/src/dispatch/lifecycle/attempt.rs @@ -8,9 +8,8 @@ use crate::{ dispatch::{ affinity::AccountIdentityService, controllers::{ - AttemptAccountPreparation, AttemptAccountPreparationContext, AttemptRoutePreparation, - AttemptRoutePreparationContext, AttemptUpstreamErrorContext, ControllerRequestScope, - ControllerSet, + AttemptRoutePreparation, AttemptRoutePreparationContext, AttemptUpstreamErrorContext, + ControllerRequestScope, ControllerSet, }, errors::{ClientFailure, ResponseDispatchError}, failure::exhaustion::AccountExhaustionTracker, @@ -191,25 +190,11 @@ impl<'a> AttemptRunner<'a> { .controller_scope .as_mut() .ok_or(AttemptContractError::Terminal)?; - let prepared_account = match self + let prepared_account = self .dependencies .controllers - .prepare_attempt_account( - AttemptAccountPreparationContext { - codex: self.dependencies.codex, - account_identity: self.dependencies.account_identity, - request_id: self.dependencies.request_id, - }, - acquired, - ) - .await - { - AttemptAccountPreparation::Ready(prepared) => *prepared, - AttemptAccountPreparation::Rejected => { - self.pending = Some(PendingAttempt::CandidatePreparationRejected { account }); - return self.step(); - } - }; + .prepare_attempt_account(acquired) + .await; let acquired = prepared_account.lease; let cookie_header = prepared_account.cookie_header; @@ -769,9 +754,6 @@ fn observation_kind( kind: *kind, } } - PendingAttempt::CandidatePreparationRejected { .. } => { - AttemptObservationKind::CandidatePreparationRejected - } PendingAttempt::RoutePreparationRejected { message, .. } => { AttemptObservationKind::RoutePreparationRejected { message: message.clone(), @@ -840,17 +822,15 @@ async fn rejection_error( }) } AttemptReturnKind::Observed => match pending { - PendingAttempt::NoCandidate | PendingAttempt::CandidatePreparationRejected { .. } => { - Ok(on_exhaustion.map_or_else( - || { - exhaustion - .last_exhausted() - .map(ResponseDispatchError::from_exhausted_account) - .unwrap_or(ResponseDispatchError::NoActiveAccount) - }, - ResponseDispatchError::Failed, - )) - } + PendingAttempt::NoCandidate => Ok(on_exhaustion.map_or_else( + || { + exhaustion + .last_exhausted() + .map(ResponseDispatchError::from_exhausted_account) + .unwrap_or(ResponseDispatchError::NoActiveAccount) + }, + ResponseDispatchError::Failed, + )), PendingAttempt::RoutePreparationRejected { message, .. } => { Ok(ResponseDispatchError::HistoryUnavailable { upstream_error: message, @@ -902,8 +882,7 @@ fn pending_account(pending: &PendingAttempt) -> Option<&str> { fn pending_account_value(pending: &PendingAttempt) -> Option<&Account> { match pending { PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } => None, - PendingAttempt::CandidatePreparationRejected { account } - | PendingAttempt::RoutePreparationRejected { account, .. } + PendingAttempt::RoutePreparationRejected { account, .. } | PendingAttempt::UpstreamFailure { account, .. } | PendingAttempt::ProtocolFailure { account, .. } => Some(account), PendingAttempt::CompleteResponse(complete) => Some(&complete.account), @@ -919,7 +898,6 @@ fn pending_attempt(pending: &PendingAttempt) -> Option<&super::trace::ResponseDi PendingAttempt::StreamResponse(stream) => Some(&stream.attempt), PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } - | PendingAttempt::CandidatePreparationRejected { .. } | PendingAttempt::RoutePreparationRejected { .. } => None, } } @@ -936,7 +914,6 @@ fn pending_attempt_request(pending: &PendingAttempt) -> Option<&CodexResponsesRe PendingAttempt::StreamResponse(stream) => Some(&stream.attempt_request), PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } - | PendingAttempt::CandidatePreparationRejected { .. } | PendingAttempt::RoutePreparationRejected { .. } => None, } } @@ -951,7 +928,6 @@ fn pending_transport( PendingAttempt::StreamResponse(stream) => Some(stream.response.transport), PendingAttempt::NoCandidate | PendingAttempt::PinnedCandidateUnavailable { .. } - | PendingAttempt::CandidatePreparationRejected { .. } | PendingAttempt::RoutePreparationRejected { .. } => None, } } @@ -973,7 +949,6 @@ fn invalid(decision: &'static str, pending: &PendingAttempt) -> AttemptContractE observation: match pending { PendingAttempt::NoCandidate => "no_candidate", PendingAttempt::PinnedCandidateUnavailable { .. } => "pinned_candidate_unavailable", - PendingAttempt::CandidatePreparationRejected { .. } => "candidate_preparation_rejected", PendingAttempt::RoutePreparationRejected { .. } => "route_preparation_rejected", PendingAttempt::UpstreamFailure { .. } => "upstream_failure", PendingAttempt::ProtocolFailure { .. } => "protocol_failure", diff --git a/backend/src/dispatch/lifecycle/contract.rs b/backend/src/dispatch/lifecycle/contract.rs index c12b4e853..8fb296a20 100644 --- a/backend/src/dispatch/lifecycle/contract.rs +++ b/backend/src/dispatch/lifecycle/contract.rs @@ -93,7 +93,6 @@ pub(in crate::dispatch) enum AttemptObservationKind { ledger: CandidateLedgerFacts, last_exhausted: Option, }, - CandidatePreparationRejected, RoutePreparationRejected { message: String, }, @@ -257,9 +256,6 @@ pub(super) enum PendingAttempt { account_id: String, kind: PinnedCandidateAcquireFailureKind, }, - CandidatePreparationRejected { - account: Account, - }, RoutePreparationRejected { account: Account, message: String, diff --git a/backend/src/fleet/account_failure.rs b/backend/src/fleet/account_failure.rs index 200801955..a848fb3fb 100644 --- a/backend/src/fleet/account_failure.rs +++ b/backend/src/fleet/account_failure.rs @@ -53,6 +53,8 @@ pub fn classify_upstream_failure(facts: &UpstreamFailureFacts) -> Option { status_code: Option, code: Option<&'a str>, error_type: Option<&'a str>, + identity_authorization_error: Option<&'a str>, + identity_error_code: Option<&'a str>, message: &'a str, body: &'a str, retry_after_seconds: Option, @@ -136,6 +142,8 @@ struct FailureFields<'a> { fn classify_failure(fields: FailureFields<'_>) -> Option { let code = fields.code.unwrap_or_default(); let error_type = fields.error_type.unwrap_or_default(); + let identity_error_code = fields.identity_error_code.unwrap_or_default(); + let identity_authorization_error = fields.identity_authorization_error.unwrap_or_default(); if [code, fields.message, fields.body] .into_iter() @@ -143,14 +151,13 @@ fn classify_failure(fields: FailureFields<'_>) -> Option Optio || message.contains("organization has been disabled") || message.contains("workspace has been deactivated") || message.contains("deactivated_workspace") - || code == "forbidden" - || error_type == "permission_error" { return Some(AccountStatus::Banned); } @@ -310,8 +315,3 @@ fn is_model_unsupported(value: &str) -> bool { || value.contains("not_supported") || value.contains("not_available"))) } - -fn is_html(value: &str) -> bool { - let value = value.trim_start().to_ascii_lowercase(); - value.starts_with(">(); if models.is_empty() { return Err(AccountManageError::NoModels); @@ -88,7 +87,32 @@ impl AccountManageService { Ok(models) } - async fn refresh_account_plan_models(&self, account: &StoredAccount, plan_type: &str) { + pub async fn refresh_account_models( + &self, + account_id: &str, + ) -> Result, AccountManageError> { + let account = self + .store + .get(account_id) + .await + .map_err(|_| AccountManageError::Inspect)? + .ok_or(AccountManageError::NotFound)?; + let plan_type = account_plan_type(&account); + self.refresh_account_plan_models(&account, &plan_type) + .await?; + let models = self.models.catalog().await.models_for_plan(&plan_type); + let models = models.iter().map(account_model_option).collect::>(); + if models.is_empty() { + return Err(AccountManageError::NoModels); + } + Ok(models) + } + + async fn refresh_account_plan_models( + &self, + account: &StoredAccount, + plan_type: &str, + ) -> Result { let request_id = uuid::Uuid::new_v4().to_string(); let plan_account = ModelRefreshPlanAccount { plan_type: plan_type.to_string(), @@ -96,18 +120,20 @@ impl AccountManageService { account_id: account.account_id.clone(), installation_id: self.account_pseudonymizer.installation_id(&account.id), }; - if let Err(error) = self + let result = self .models - .refresh_backend_models(&[plan_account], &request_id) + .refresh_selected_plan_models(&[plan_account], &request_id) .await - { - tracing::warn!( - account_id = %account.id, - plan_type, - error = %error, - "Failed to refresh account plan models" - ); - } + .map_err(|error| AccountManageError::RefreshModels(error.to_string()))?; + let routing = self + .models + .model_plan_routing() + .await + .map_err(|error| AccountManageError::RefreshModels(error.to_string()))?; + self.account_pool + .apply_model_plan_routing(routing.allowlist, routing.fetched_plan_types) + .await; + Ok(result) } pub async fn test_connection_stream( diff --git a/backend/src/fleet/manage/types.rs b/backend/src/fleet/manage/types.rs index 75629b4d9..f775ab568 100644 --- a/backend/src/fleet/manage/types.rs +++ b/backend/src/fleet/manage/types.rs @@ -86,6 +86,8 @@ pub enum AccountManageError { FetchQuota(String), #[error("no available models returned")] NoModels, + #[error("failed to refresh account models: {0}")] + RefreshModels(String), #[error("invalid status: {0}")] InvalidStatus(String), #[error("label must be 64 characters or fewer")] diff --git a/backend/src/models/service.rs b/backend/src/models/service.rs index 5c17217b0..9652d0a1c 100644 --- a/backend/src/models/service.rs +++ b/backend/src/models/service.rs @@ -189,6 +189,26 @@ impl ModelService { &self, plan_accounts: &[ModelRefreshPlanAccount], request_id: &str, + ) -> Result { + self.refresh_backend_models_inner(plan_accounts, request_id, false) + .await + } + + /// 刷新指定计划,并保留未参与本次刷新的计划快照。 + pub async fn refresh_selected_plan_models( + &self, + plan_accounts: &[ModelRefreshPlanAccount], + request_id: &str, + ) -> Result { + self.refresh_backend_models_inner(plan_accounts, request_id, true) + .await + } + + async fn refresh_backend_models_inner( + &self, + plan_accounts: &[ModelRefreshPlanAccount], + request_id: &str, + preserve_other_plans: bool, ) -> Result { let _refresh_guard = self.refresh_lock.lock().await; let store = self @@ -211,7 +231,11 @@ impl ModelService { .into_iter() .map(|snapshot| (snapshot.plan_type.clone(), snapshot)) .collect::>(); - let mut next_snapshots = BTreeMap::new(); + let mut next_snapshots = if preserve_other_plans { + previous_snapshots.clone() + } else { + BTreeMap::new() + }; let mut result = ModelRefreshResult { refreshed_plans: 0, diff --git a/backend/src/upstream/openai/failure.rs b/backend/src/upstream/openai/failure.rs index 9462f3ed6..50ab20262 100644 --- a/backend/src/upstream/openai/failure.rs +++ b/backend/src/upstream/openai/failure.rs @@ -30,6 +30,8 @@ pub struct UpstreamFailureFacts { pub status_code: Option, pub code: Option, pub error_type: Option, + pub identity_authorization_error: Option, + pub identity_error_code: Option, pub message: String, pub body: String, pub retry_after_seconds: Option, @@ -48,6 +50,8 @@ pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts status_code: error.status().map(|status| status.as_u16()), code: None, error_type: None, + identity_authorization_error: None, + identity_error_code: None, message: error.to_string(), body: error.to_string(), retry_after_seconds: None, @@ -57,6 +61,7 @@ pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts body, retry_after_seconds, transport, + diagnostics, .. } => { let (code, error_type, message) = error_fields(body); @@ -72,6 +77,8 @@ pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts status_code: Some(status.as_u16()), code, error_type, + identity_authorization_error: diagnostics.identity_authorization_error.clone(), + identity_error_code: diagnostics.identity_error_code.clone(), message: message.unwrap_or_else(|| body.clone()), body: body.clone(), retry_after_seconds: *retry_after_seconds, @@ -87,6 +94,11 @@ pub fn upstream_failure_facts(error: &CodexClientError) -> UpstreamFailureFacts status_code: Some(upstream.status_code), code: upstream.code.clone().or(body_code), error_type: upstream.error_type.clone().or(body_type), + identity_authorization_error: upstream + .diagnostics + .identity_authorization_error + .clone(), + identity_error_code: upstream.diagnostics.identity_error_code.clone(), message: upstream .message .clone() @@ -141,6 +153,8 @@ fn simple_facts(kind: UpstreamFailureKind, error: &CodexClientError) -> Upstream status_code: None, code: None, error_type: None, + identity_authorization_error: None, + identity_error_code: None, body: message.clone(), message, retry_after_seconds: None, diff --git a/backend/src/upstream/openai/protocol/events.rs b/backend/src/upstream/openai/protocol/events.rs index b2c231fec..2690ddb73 100644 --- a/backend/src/upstream/openai/protocol/events.rs +++ b/backend/src/upstream/openai/protocol/events.rs @@ -238,6 +238,8 @@ pub struct CreditsSnapshot { pub struct ParsedRateLimits { /// 以标准化 limit ID 为键的全部计量项。 pub limits: BTreeMap, + /// 当前响应对应的计量项 ID。 + pub active_limit: Option, pub credits: Option, pub plan_type: Option, pub promo_message: Option, @@ -251,10 +253,15 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option>(); + if let Some(active_limit) = &active_limit { + limit_ids.insert(active_limit.clone()); + } limit_ids.insert("codex".to_string()); let limits = limit_ids .into_iter() @@ -270,6 +277,7 @@ pub fn parse_rate_limit_headers(headers: &[(String, String)]) -> Option Option bool { "x-codex-credits-has-credits" | "x-codex-credits-unlimited" | "x-codex-credits-balance" + | "x-codex-active-limit" | "x-codex-plan-type" | "x-codex-promo-message" | "x-codex-rate-limit-reached-type" @@ -343,10 +353,11 @@ pub fn parse_rate_limits_event(value: &Value) -> Option { let mut limits = BTreeMap::new(); if let Some(details) = details { - limits.insert(limit_id, details); + limits.insert(limit_id.clone(), details); } Some(ParsedRateLimits { limits, + active_limit: Some(limit_id), credits, plan_type, promo_message: value @@ -430,9 +441,19 @@ pub fn rate_limit_quota( let spend_control = existing_quota .and_then(|quota| quota.get("spend_control").cloned()) .unwrap_or(Value::Null); + let active_limit = rate_limits.active_limit.as_ref().map_or_else( + || { + existing_quota + .and_then(|quota| quota.get("active_limit")) + .cloned() + .unwrap_or(Value::Null) + }, + |limit| Value::String(limit.clone()), + ); json!({ "plan_type": rate_limits.plan_type.as_deref().or(plan_type).unwrap_or("unknown"), + "active_limit": active_limit, "snapshots": snapshots, "monthly_limit": monthly_limit, "credits": credits, @@ -445,6 +466,9 @@ pub fn rate_limit_quota( /// 将限流状态转换回内部传输头键值对。 pub fn rate_limits_to_header_pairs(rate_limits: &ParsedRateLimits) -> Vec<(String, String)> { let mut headers = Vec::new(); + if let Some(active_limit) = &rate_limits.active_limit { + headers.push(("x-codex-active-limit".to_string(), active_limit.clone())); + } for details in rate_limits.limits.values() { let prefix = format!("x-{}", details.limit_id.replace('_', "-")); push_window_headers(&mut headers, &format!("{prefix}-primary"), details.primary); diff --git a/backend/src/upstream/openai/transport/client.rs b/backend/src/upstream/openai/transport/client.rs index c27b39fb5..b2b069e99 100644 --- a/backend/src/upstream/openai/transport/client.rs +++ b/backend/src/upstream/openai/transport/client.rs @@ -139,9 +139,11 @@ pub enum CodexClientError { /// 推导出的重试秒数。 retry_after_seconds: Option, /// 上游诊断元数据。 - diagnostics: CodexUpstreamDiagnostics, + diagnostics: Box, /// 上游透传的 `set-cookie` 列表。 set_cookie_headers: Vec, + /// 上游错误响应携带的限流头。 + rate_limit_headers: Vec<(String, String)>, /// 实际收到该上游响应的 transport。 transport: CodexBackendTransport, }, @@ -616,13 +618,22 @@ fn websocket_exchange_error_to_client_error( match error { CodexWebSocketExchangeError::Upstream(upstream) => { let upstream = *upstream; + let rate_limit_headers = upstream + .headers + .iter() + .filter(|(name, _)| { + crate::upstream::openai::protocol::events::is_rate_limit_header_name(name) + }) + .cloned() + .collect(); CodexClientError::Upstream { status: StatusCode::from_u16(upstream.status_code) .unwrap_or(StatusCode::BAD_GATEWAY), body: upstream.body, retry_after_seconds: upstream.retry_after_seconds, - diagnostics: upstream.diagnostics, + diagnostics: Box::new(upstream.diagnostics), set_cookie_headers: upstream.set_cookie_headers, + rate_limit_headers, transport: CodexBackendTransport::WebSocket, } } diff --git a/backend/src/upstream/openai/transport/client_sse.rs b/backend/src/upstream/openai/transport/client_sse.rs index 449eff6a6..77d69785b 100644 --- a/backend/src/upstream/openai/transport/client_sse.rs +++ b/backend/src/upstream/openai/transport/client_sse.rs @@ -84,8 +84,9 @@ impl CodexBackendClient { retry_after_seconds: retry_after_seconds .or_else(|| retry_after_seconds_from_body(&body)), body, - diagnostics, + diagnostics: Box::new(diagnostics), set_cookie_headers, + rate_limit_headers, transport: CodexBackendTransport::HttpSse, }); } @@ -166,8 +167,9 @@ impl CodexBackendClient { retry_after_seconds: retry_after_seconds .or_else(|| retry_after_seconds_from_body(&body)), body, - diagnostics, + diagnostics: Box::new(diagnostics), set_cookie_headers, + rate_limit_headers, transport: CodexBackendTransport::HttpSse, }); } @@ -539,8 +541,9 @@ impl CodexBackendClient { retry_after_seconds: retry_after_seconds .or_else(|| retry_after_seconds_from_body(&body)), body, - diagnostics, + diagnostics: Box::new(diagnostics), set_cookie_headers, + rate_limit_headers: Vec::new(), transport: CodexBackendTransport::HttpSse, }); } @@ -552,10 +555,11 @@ impl CodexBackendClient { "model catalog response is not valid JSON: {error}; body: {}", truncate_for_error(&body) ), - diagnostics: CodexUpstreamDiagnostics::with_status( + diagnostics: Box::new(CodexUpstreamDiagnostics::with_status( StatusCode::BAD_GATEWAY.as_u16(), - ), + )), set_cookie_headers, + rate_limit_headers: Vec::new(), transport: CodexBackendTransport::HttpSse, })?; let models = extract_model_entries(&parsed); @@ -567,8 +571,11 @@ impl CodexBackendClient { status: StatusCode::BAD_GATEWAY, retry_after_seconds: None, body: "model catalog response contains no models".to_string(), - diagnostics: CodexUpstreamDiagnostics::with_status(StatusCode::BAD_GATEWAY.as_u16()), + diagnostics: Box::new(CodexUpstreamDiagnostics::with_status( + StatusCode::BAD_GATEWAY.as_u16(), + )), set_cookie_headers: Vec::new(), + rate_limit_headers: Vec::new(), transport: CodexBackendTransport::HttpSse, }) } diff --git a/backend/src/upstream/openai/transport/diagnostics.rs b/backend/src/upstream/openai/transport/diagnostics.rs index b057b894c..4d211b5d1 100644 --- a/backend/src/upstream/openai/transport/diagnostics.rs +++ b/backend/src/upstream/openai/transport/diagnostics.rs @@ -1,9 +1,12 @@ //! Upstream response diagnostics captured at the transport boundary. +use base64::{Engine as _, engine::general_purpose::STANDARD}; use reqwest::header::HeaderMap; +use serde_json::Value; const UPSTREAM_REQUEST_ID_HEADERS: &[&str] = &[ "x-request-id", + "x-oai-request-id", "x-openai-request-id", "openai-request-id", "request-id", @@ -11,16 +14,21 @@ const UPSTREAM_REQUEST_ID_HEADERS: &[&str] = &[ const UPSTREAM_TRACE_HEADERS: &[&str] = &[ "x-request-id", + "x-oai-request-id", "x-openai-request-id", "openai-request-id", "request-id", "cf-ray", ]; +const IDENTITY_AUTHORIZATION_ERROR_HEADER: &str = "x-openai-authorization-error"; +const IDENTITY_ERROR_JSON_HEADER: &str = "x-error-json"; #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct CodexUpstreamDiagnostics { pub status_code: Option, pub request_id: Option, + pub identity_authorization_error: Option, + pub identity_error_code: Option, pub trace_headers: Vec<(String, String)>, } @@ -29,6 +37,12 @@ impl CodexUpstreamDiagnostics { Self { status_code, request_id: first_header(headers, UPSTREAM_REQUEST_ID_HEADERS), + identity_authorization_error: header_value( + headers, + IDENTITY_AUTHORIZATION_ERROR_HEADER, + ), + identity_error_code: header_value(headers, IDENTITY_ERROR_JSON_HEADER) + .and_then(|encoded| decode_identity_error_code(&encoded)), trace_headers: trace_headers(headers), } } @@ -48,15 +62,24 @@ impl CodexUpstreamDiagnostics { .iter() .filter_map(|name| pair_value(headers, name).map(|value| ((*name).to_string(), value))) .collect(); + let identity_authorization_error = pair_value(headers, IDENTITY_AUTHORIZATION_ERROR_HEADER); + let identity_error_code = pair_value(headers, IDENTITY_ERROR_JSON_HEADER) + .and_then(|encoded| decode_identity_error_code(&encoded)); Self { status_code, request_id, + identity_authorization_error, + identity_error_code, trace_headers, } } pub fn is_empty(&self) -> bool { - self.status_code.is_none() && self.request_id.is_none() && self.trace_headers.is_empty() + self.status_code.is_none() + && self.request_id.is_none() + && self.identity_authorization_error.is_none() + && self.identity_error_code.is_none() + && self.trace_headers.is_empty() } pub fn cf_ray(&self) -> Option<&str> { @@ -95,3 +118,14 @@ fn header_value(headers: &HeaderMap, name: &str) -> Option { .filter(|value| !value.is_empty()) .map(ToString::to_string) } + +fn decode_identity_error_code(encoded: &str) -> Option { + let decoded = STANDARD.decode(encoded).ok()?; + serde_json::from_slice::(&decoded) + .ok()? + .pointer("/error/code")? + .as_str() + .map(str::trim) + .filter(|code| !code.is_empty()) + .map(ToString::to_string) +} diff --git a/backend/src/upstream/openai/transport/usage.rs b/backend/src/upstream/openai/transport/usage.rs index 1682835f9..3ce18519c 100644 --- a/backend/src/upstream/openai/transport/usage.rs +++ b/backend/src/upstream/openai/transport/usage.rs @@ -34,8 +34,9 @@ impl CodexBackendClient { retry_after_seconds: retry_after_seconds .or_else(|| retry_after_seconds_from_body(&body)), body, - diagnostics, + diagnostics: Box::new(diagnostics), set_cookie_headers: Vec::new(), + rate_limit_headers: Vec::new(), transport: super::client::CodexBackendTransport::HttpSse, }); } @@ -54,8 +55,11 @@ impl CodexBackendClient { || "usage endpoint is unavailable".to_string(), |body| format!("invalid usage response: {}", truncate_for_error(&body)), ), - diagnostics: CodexUpstreamDiagnostics::with_status(StatusCode::BAD_GATEWAY.as_u16()), + diagnostics: Box::new(CodexUpstreamDiagnostics::with_status( + StatusCode::BAD_GATEWAY.as_u16(), + )), set_cookie_headers: Vec::new(), + rate_limit_headers: Vec::new(), transport: super::client::CodexBackendTransport::HttpSse, }) } diff --git a/backend/tests/api/admin/accounts_routes/probe_routes.rs b/backend/tests/api/admin/accounts_routes/probe_routes.rs index 6e8fa7752..4b7234929 100644 --- a/backend/tests/api/admin/accounts_routes/probe_routes.rs +++ b/backend/tests/api/admin/accounts_routes/probe_routes.rs @@ -85,27 +85,94 @@ async fn account_models_should_return_redis_snapshot_models() { } #[tokio::test] -async fn account_models_should_fetch_missing_plan_snapshot_with_account_token() { +async fn refresh_account_models_should_fetch_persist_and_preserve_other_plans() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/backend-api/codex/models")) + .and(header("authorization", "Bearer access-test")) + .and(header("chatgpt-account-id", "chatgpt-test")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "models": [{ + "slug": "gpt-plus-refreshed", + "display_name": "GPT Plus Refreshed" + }] + }))) + .expect(1) + .mount(&server) + .await; + let (app, state, pool, _dir) = admin_accounts_test_app_with_api_base_url( + "admin-account-models-refresh", + 92, + format!("{}/backend-api", server.uri()), + ) + .await; + seed_test_account(&pool).await; + let model_snapshots = + codex_proxy_rs::models::store::RedisModelSnapshotStore::new(state.redis.clone()); + model_snapshots + .replace_plan_snapshots(&[ + ModelPlanSnapshot::from_backend_entries( + "plus", + vec![BackendModelEntry { + id: Some("gpt-plus-old".to_string()), + name: Some("GPT Plus Old".to_string()), + ..BackendModelEntry::default() + }], + ), + ModelPlanSnapshot::from_backend_entries( + "team", + vec![BackendModelEntry { + id: Some("gpt-team-existing".to_string()), + name: Some("GPT Team Existing".to_string()), + ..BackendModelEntry::default() + }], + ), + ]) + .await + .unwrap(); + state + .services + .models + .reload_from_store() + .await + .expect("model catalog should reload from seeded snapshots"); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/api/admin/accounts/models") + .header("cookie", "cpr_admin_session=session_1") + .header("content-type", "application/json") + .body(Body::from(json!({ "id": "acct_test" }).to_string())) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let body = response_json(response).await; + assert_eq!(body["data"]["models"][0]["id"], "gpt-plus-refreshed"); + let stored = model_snapshots.list_plan_snapshots().await.unwrap(); + assert!(stored.iter().any(|snapshot| { + snapshot.plan_type == "plus" + && snapshot + .models + .iter() + .any(|model| model.id == "gpt-plus-refreshed") + })); + assert!(stored.iter().any(|snapshot| snapshot.plan_type == "team")); +} + +#[tokio::test] +async fn account_models_should_not_fetch_missing_plan_snapshot() { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/backend-api/codex/models")) .and(header("authorization", "Bearer access-plan-b")) .and(header("chatgpt-account-id", "chatgpt-plan-b")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "models": [ - { - "slug": "gpt-plan-b-live", - "display_name": "GPT Plan B Live", - "description": "Live model", - "is_default": true, - "supported_reasoning_efforts": [{"reasoning_effort": "medium", "description": "medium"}], - "default_reasoning_effort": "medium", - "input_modalities": ["text"], - "output_modalities": ["text"], - "supports_personality": false - } - ] - }))) + .respond_with(ResponseTemplate::new(500)) + .expect(0) .mount(&server) .await; let (app, state, pool, _dir) = admin_accounts_test_app_with_api_base_url( @@ -153,15 +220,12 @@ async fn account_models_should_fetch_missing_plan_snapshot_with_account_token() .unwrap(); let requests = server.received_requests().await.unwrap(); - assert_eq!(response.status(), StatusCode::OK); - assert_eq!(requests.len(), 1, "missing plan should be fetched upstream"); - let body = response_json(response).await; - assert_eq!(body["data"]["models"][0]["id"], "gpt-plan-b-live"); - assert_eq!(body["data"]["models"][0]["label"], "GPT Plan B Live"); + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert!(requests.is_empty(), "model snapshot lookup must stay local"); let model_snapshots = codex_proxy_rs::models::store::RedisModelSnapshotStore::new(state.redis.clone()); let stored = model_snapshots.list_plan_snapshots().await.unwrap(); - assert!(stored.iter().any(|snapshot| snapshot.plan_type == "plan-b")); + assert!(!stored.iter().any(|snapshot| snapshot.plan_type == "plan-b")); } #[tokio::test] diff --git a/backend/tests/dispatch/service/responses_http.rs b/backend/tests/dispatch/service/responses_http.rs index 122c24a2a..2603face5 100644 --- a/backend/tests/dispatch/service/responses_http.rs +++ b/backend/tests/dispatch/service/responses_http.rs @@ -44,7 +44,7 @@ async fn responses_json_should_preserve_upstream_client_failure() { let body = response_json(response).await; assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); - assert_eq!(body["error"]["code"], "server_overloaded"); + assert_eq!(body["error"]["code"], "server_is_overloaded"); assert_eq!( body["error"]["message"], "upstream is temporarily overloaded" @@ -97,11 +97,25 @@ async fn responses_sse_should_preserve_upstream_client_failure() { assert_eq!(status, StatusCode::OK); assert!(content_type.starts_with("text/event-stream")); - assert!(body.contains(r#""code":"server_overloaded""#)); + assert!(body.contains(r#""code":"server_is_overloaded""#)); assert!(body.contains("upstream is temporarily overloaded")); assert!(body.ends_with("data: [DONE]\n\n")); } +#[tokio::test] +async fn responses_context_length_exceeded_should_return_bad_request() { + let status = complete_response_failure_status("context_length_exceeded").await; + + assert_eq!(status, StatusCode::BAD_REQUEST); +} + +#[tokio::test] +async fn responses_usage_not_included_should_return_forbidden() { + let status = complete_response_failure_status("usage_not_included").await; + + assert_eq!(status, StatusCode::FORBIDDEN); +} + #[tokio::test] async fn responses_should_reject_invalid_json_without_upstream_request() { let server = MockServer::start().await; @@ -128,6 +142,53 @@ async fn responses_should_reject_invalid_json_without_upstream_request() { assert!(requests.is_empty()); } +async fn complete_response_failure_status(code: &str) -> StatusCode { + let server = MockServer::start().await; + let body = format!( + "event: response.failed\ndata: {}\n\n", + json!({ + "response": { + "id": "resp_failed", + "object": "response", + "status": "failed", + "error": {"code": code, "message": "upstream rejected request"} + } + }) + ); + Mock::given(method("POST")) + .and(path("/codex/responses")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(body), + ) + .expect(1) + .mount(&server) + .await; + let (app, api_key, _dir) = test_app_with_account(server.uri()).await; + + app.oneshot( + Request::builder() + .method("POST") + .uri("/v1/responses") + .header("authorization", format!("Bearer {api_key}")) + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-5.5", + "input": [], + "stream": false, + "use_websocket": false + }) + .to_string(), + )) + .unwrap(), + ) + .await + .unwrap() + .status() +} + #[tokio::test] async fn responses_should_reject_non_object_json_without_upstream_request() { let server = MockServer::start().await; diff --git a/backend/tests/dispatch/service/responses_recovery.rs b/backend/tests/dispatch/service/responses_recovery.rs index 181dbe615..1c26a150a 100644 --- a/backend/tests/dispatch/service/responses_recovery.rs +++ b/backend/tests/dispatch/service/responses_recovery.rs @@ -414,12 +414,17 @@ async fn responses_should_return_no_available_accounts_error_when_no_accounts_ar #[tokio::test] async fn responses_should_fallback_to_next_account_after_rate_limit() { let server = MockServer::start().await; + let reset_at = Utc::now().timestamp() + 120; Mock::given(method("POST")) .and(path("/codex/responses")) .and(header("authorization", "Bearer access-primary")) .respond_with( ResponseTemplate::new(429) .insert_header("retry-after", "120") + .insert_header("x-codex-active-limit", "codex") + .insert_header("x-codex-primary-used-percent", "100") + .insert_header("x-codex-primary-window-minutes", "5") + .insert_header("x-codex-primary-reset-at", reset_at.to_string()) .set_body_json(json!({ "error": { "message": "rate limited", @@ -484,6 +489,13 @@ async fn responses_should_fallback_to_next_account_after_rate_limit() { .unwrap(); assert_eq!(primary_usage, (1, 0, 0)); assert_eq!(secondary_usage, (1, 5, 2)); + let quota: (Value,) = sqlx::query_as("select quota_json from accounts where id = $1") + .bind("acct_primary") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(quota.0["active_limit"], "codex"); + assert_eq!(quota.0["snapshots"][0]["primary"]["reset_at"], reset_at); } #[tokio::test] @@ -856,7 +868,7 @@ async fn responses_stream_should_mark_banned_after_402_deactivated_workspace_and } #[tokio::test] -async fn responses_should_mark_banned_after_403_and_fallback() { +async fn responses_should_return_generic_403_without_banning_or_fallback() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/codex/responses")) @@ -875,7 +887,7 @@ async fn responses_should_mark_banned_after_403_and_fallback() { .insert_header("content-type", "text/event-stream") .set_body_string(RESPONSES_AFTER_403_SSE), ) - .expect(1) + .expect(0) .mount(&server) .await; @@ -908,9 +920,9 @@ async fn responses_should_mark_banned_after_403_and_fallback() { .await .unwrap(); - assert_eq!(status, StatusCode::OK); - assert_eq!(body["id"], "resp_after_403"); - assert_eq!(account_status.0, "banned"); + assert_eq!(status, StatusCode::FORBIDDEN); + assert_eq!(body["error"]["message"], "request forbidden"); + assert_eq!(account_status.0, "active"); } #[tokio::test] @@ -2152,7 +2164,7 @@ async fn responses_stream_should_return_auth_error_when_401_fallback_is_exhauste } #[tokio::test] -async fn responses_stream_should_mark_banned_after_403_and_fallback() { +async fn responses_stream_should_return_generic_403_without_banning_or_fallback() { let server = MockServer::start().await; Mock::given(method("POST")) .and(path("/codex/responses")) @@ -2171,7 +2183,7 @@ async fn responses_stream_should_mark_banned_after_403_and_fallback() { .insert_header("content-type", "text/event-stream") .set_body_string(RESPONSES_STREAM_AFTER_403_SSE), ) - .expect(1) + .expect(0) .mount(&server) .await; @@ -2210,10 +2222,10 @@ async fn responses_stream_should_mark_banned_after_403_and_fallback() { .await .unwrap(); - assert_eq!(status, StatusCode::OK); - assert!(content_type.starts_with("text/event-stream")); - assert!(body.contains("resp_stream_after_403")); - assert_eq!(account_status.0, "banned"); + assert_eq!(status, StatusCode::FORBIDDEN); + assert!(content_type.starts_with("application/json")); + assert!(body.contains("request forbidden")); + assert_eq!(account_status.0, "active"); } #[tokio::test] diff --git a/backend/tests/dispatch/service/usage_logging.rs b/backend/tests/dispatch/service/usage_logging.rs index e75f92bc4..fcb137587 100644 --- a/backend/tests/dispatch/service/usage_logging.rs +++ b/backend/tests/dispatch/service/usage_logging.rs @@ -399,6 +399,7 @@ async fn responses_should_passively_cache_rate_limit_headers() { .respond_with( ResponseTemplate::new(200) .insert_header("content-type", "text/event-stream") + .insert_header("x-codex-active-limit", "codex") .insert_header("x-codex-primary-used-percent", "100") .insert_header("x-codex-primary-window-minutes", "5") .insert_header("x-codex-primary-reset-at", reset_at.to_string()) @@ -440,6 +441,10 @@ async fn responses_should_passively_cache_rate_limit_headers() { .unwrap(); assert_eq!(response.status(), StatusCode::OK); + let upstream_requests = server.received_requests().await.unwrap(); + assert_eq!(upstream_requests.len(), 1); + assert_eq!(upstream_requests[0].method.as_str(), "POST"); + assert_eq!(upstream_requests[0].url.path(), "/codex/responses"); let stored: ( Value, String, @@ -454,6 +459,7 @@ async fn responses_should_passively_cache_rate_limit_headers() { .await .unwrap(); let quota = stored.0; + assert_eq!(quota["active_limit"], "codex"); assert_eq!(quota["snapshots"][0]["source"], "core"); assert_eq!(quota["snapshots"][0]["primary"]["limit_reached"], true); assert_eq!(quota["snapshots"][0]["primary"]["reset_at"], reset_at); diff --git a/backend/tests/fixtures/responses/http_sse/failed_server_overloaded.sse b/backend/tests/fixtures/responses/http_sse/failed_server_overloaded.sse index 28637c599..72970adb2 100644 --- a/backend/tests/fixtures/responses/http_sse/failed_server_overloaded.sse +++ b/backend/tests/fixtures/responses/http_sse/failed_server_overloaded.sse @@ -1,2 +1,2 @@ event: response.failed -data: {"response":{"id":"resp_server_overloaded","object":"response","status":"failed","error":{"code":"server_overloaded","message":"upstream is temporarily overloaded"}}} +data: {"response":{"id":"resp_server_overloaded","object":"response","status":"failed","error":{"code":"server_is_overloaded","message":"upstream is temporarily overloaded"}}} diff --git a/backend/tests/fleet/account_failure.rs b/backend/tests/fleet/account_failure.rs index 73dcb6e19..0efacc081 100644 --- a/backend/tests/fleet/account_failure.rs +++ b/backend/tests/fleet/account_failure.rs @@ -60,6 +60,38 @@ fn html_forbidden_page_should_not_ban_account() { assert!(classify_client_failure(&error).is_none()); } +#[test] +fn generic_forbidden_should_not_ban_account() { + let error = upstream_error(403, r#"{"error":{"message":"request forbidden"}}"#); + + assert!(classify_client_failure(&error).is_none()); +} + +#[test] +fn generic_permission_error_should_not_ban_account() { + let error = upstream_error( + 403, + r#"{"error":{"code":"forbidden","type":"permission_error"}}"#, + ); + + assert!(classify_client_failure(&error).is_none()); +} + +#[test] +fn identity_error_header_should_expire_account() { + let mut diagnostics = CodexUpstreamDiagnostics::with_status(403); + diagnostics.identity_error_code = Some("token_expired".to_string()); + let error = upstream_error_with_diagnostics( + 403, + r#"{"error":{"message":"request forbidden"}}"#, + diagnostics, + ); + + let failure = classify_client_failure(&error).expect("identity failure should be classified"); + + assert_eq!(failure.kind, AccountFailureKind::Expired); +} + fn classify( status: u16, body: &str, @@ -68,12 +100,21 @@ fn classify( } fn upstream_error(status: u16, body: &str) -> CodexClientError { + upstream_error_with_diagnostics(status, body, CodexUpstreamDiagnostics::with_status(status)) +} + +fn upstream_error_with_diagnostics( + status: u16, + body: &str, + diagnostics: CodexUpstreamDiagnostics, +) -> CodexClientError { CodexClientError::Upstream { status: StatusCode::from_u16(status).unwrap(), body: body.to_string(), retry_after_seconds: None, - diagnostics: CodexUpstreamDiagnostics::with_status(status), + diagnostics: Box::new(diagnostics), set_cookie_headers: Vec::new(), + rate_limit_headers: Vec::new(), transport: CodexBackendTransport::HttpSse, } } diff --git a/backend/tests/upstream/openai/protocol/events.rs b/backend/tests/upstream/openai/protocol/events.rs index f42dc9bcd..50ed43b77 100644 --- a/backend/tests/upstream/openai/protocol/events.rs +++ b/backend/tests/upstream/openai/protocol/events.rs @@ -348,6 +348,10 @@ fn parse_rate_limit_headers_should_extract_arbitrary_limits_and_account_metadata "x-codex-other-primary-used-percent".to_string(), "63.5".to_string(), ), + ( + "x-codex-active-limit".to_string(), + "codex-other".to_string(), + ), ( "x-codex-other-primary-window-minutes".to_string(), "1440".to_string(), @@ -371,6 +375,7 @@ fn parse_rate_limit_headers_should_extract_arbitrary_limits_and_account_metadata ]; let parsed = parse_rate_limit_headers(&headers).expect("rate limits should parse"); + assert_eq!(parsed.active_limit.as_deref(), Some("codex_other")); let other = parsed.limits.get("codex_other").expect("dynamic limit"); assert_eq!(other.limit_name.as_deref(), Some("Codex Other")); assert_eq!( diff --git a/backend/tests/upstream/openai/transport/diagnostics.rs b/backend/tests/upstream/openai/transport/diagnostics.rs index b93691aad..cbd7c6ebb 100644 --- a/backend/tests/upstream/openai/transport/diagnostics.rs +++ b/backend/tests/upstream/openai/transport/diagnostics.rs @@ -1,3 +1,4 @@ +use base64::{Engine as _, engine::general_purpose::STANDARD}; use codex_proxy_rs::upstream::openai::transport::CodexUpstreamDiagnostics; use reqwest::header::{HeaderMap, HeaderValue}; @@ -24,3 +25,45 @@ fn diagnostics_should_not_treat_cf_ray_as_request_id() { assert_eq!(diagnostics.request_id, None); assert_eq!(diagnostics.cf_ray(), Some("ray_1")); } + +#[test] +fn diagnostics_should_extract_oai_request_id() { + let mut headers = HeaderMap::new(); + headers.insert("x-oai-request-id", HeaderValue::from_static("oai_req_1")); + + let diagnostics = CodexUpstreamDiagnostics::from_headers(Some(500), &headers); + + assert_eq!(diagnostics.request_id.as_deref(), Some("oai_req_1")); +} + +#[test] +fn diagnostics_should_extract_identity_error_facts() { + let encoded = STANDARD.encode(r#"{"error":{"code":"token_expired"}}"#); + let mut headers = HeaderMap::new(); + headers.insert( + "x-openai-authorization-error", + HeaderValue::from_static("authorization failed"), + ); + headers.insert("x-error-json", HeaderValue::from_str(&encoded).unwrap()); + + let diagnostics = CodexUpstreamDiagnostics::from_headers(Some(403), &headers); + + assert_eq!( + diagnostics.identity_authorization_error.as_deref(), + Some("authorization failed") + ); + assert_eq!( + diagnostics.identity_error_code.as_deref(), + Some("token_expired") + ); +} + +#[test] +fn diagnostics_should_ignore_malformed_identity_error_json() { + let mut headers = HeaderMap::new(); + headers.insert("x-error-json", HeaderValue::from_static("not-base64")); + + let diagnostics = CodexUpstreamDiagnostics::from_headers(Some(403), &headers); + + assert_eq!(diagnostics.identity_error_code, None); +} diff --git a/frontend/src/api/modules/accounts.ts b/frontend/src/api/modules/accounts.ts index 4a8691fcb..5c925cd94 100644 --- a/frontend/src/api/modules/accounts.ts +++ b/frontend/src/api/modules/accounts.ts @@ -86,3 +86,11 @@ export function getAccountModels(data: any) { }, }) } + +export function refreshAccountModels(data: any) { + return request({ + url: '/api/admin/accounts/models', + method: 'POST', + data, + }) +} diff --git a/frontend/src/views/accounts/components/AccountConnectionTestModal.vue b/frontend/src/views/accounts/components/AccountConnectionTestModal.vue index b979d8e10..58d5c2e4e 100644 --- a/frontend/src/views/accounts/components/AccountConnectionTestModal.vue +++ b/frontend/src/views/accounts/components/AccountConnectionTestModal.vue @@ -1,4 +1,6 @@