From f314ebb3a25669ee68d37d4909f144cfe1c0d29f Mon Sep 17 00:00:00 2001 From: Zyy Date: Sun, 2 Aug 2026 14:27:04 +0800 Subject: [PATCH] fix(openai): align wire profile os_type and harden quota 5xx backoff - Default wire profile os_type to macOS to match the audited official desktop user-agent - Retry quota usage 5xx rejections with 1s/2s exponential backoff, leaving 4xx to the existing quota state transitions - Classify 503 with Retry-After into a resettable cooldown instead of an unclassified transient warning - Add account_id and conversation hash to WebSocket pump exit logs so idle-connection resets can be attributed to an account Co-authored-by: Codex --- backend/crates/providers/openai/src/config.rs | 2 +- .../providers/openai/src/credential/quota.rs | 43 ++++++++++++++-- .../src/transport/websocket/coordinator.rs | 25 ++++++++-- .../src/transport/websocket/handshake.rs | 5 +- .../openai/src/transport/websocket/pump.rs | 41 +++++++++++++-- .../openai/tests/credential/quota.rs | 50 ++++++++++++++++++- .../openai/tests/transport/profile.rs | 6 +-- 7 files changed, 155 insertions(+), 17 deletions(-) diff --git a/backend/crates/providers/openai/src/config.rs b/backend/crates/providers/openai/src/config.rs index 6fe7c8f85..59b827497 100644 --- a/backend/crates/providers/openai/src/config.rs +++ b/backend/crates/providers/openai/src/config.rs @@ -295,7 +295,7 @@ impl Default for CodexWireProfileConfig { codex_version: "0.145.0".to_owned(), desktop_version: "26.721.31836".to_owned(), desktop_build: "5828".to_owned(), - os_type: "Mac OS".to_owned(), + os_type: "macOS".to_owned(), os_version: "15.7.1".to_owned(), arch: "arm64".to_owned(), terminal: "unknown".to_owned(), diff --git a/backend/crates/providers/openai/src/credential/quota.rs b/backend/crates/providers/openai/src/credential/quota.rs index aacbfb365..9c5807b3b 100644 --- a/backend/crates/providers/openai/src/credential/quota.rs +++ b/backend/crates/providers/openai/src/credential/quota.rs @@ -30,6 +30,9 @@ const MAX_RATE_LIMIT_COOLDOWN: Duration = Duration::from_secs(60 * 60); const QUOTA_SCHEDULING_TTL: Duration = Duration::from_secs(10 * 60); const QUOTA_HYDRATION_FAILURE_TTL: Duration = Duration::from_secs(5); const QUOTA_PERIODIC_REFRESH_MIN_INTERVAL: Duration = Duration::from_secs(30 * 60); +// 5xx 上游拒绝的短退避重试预算；指数退避 1s/2s，吞掉瞬时抖动（对齐 TS withRetry）。 +const QUOTA_FETCH_5XX_MAX_RETRIES: u32 = 2; +const QUOTA_FETCH_5XX_BASE_DELAY: Duration = Duration::from_secs(1); /// OpenAI Provider 主动额度刷新的调度策略。 /// @@ -1089,7 +1092,7 @@ impl CodexCredentialQuotaService { .prepare(account) .await .map_err(|_| CodexQuotaFetchError::InvalidCredential)?; - let mut result = fetch_usage_once(client, &prepared).await; + let mut result = fetch_usage_with_5xx_retry(client, &prepared).await; if let Err(CodexQuotaFetchAttemptError::Upstream(error)) = &result && let Some(recovered) = self .agent_identity @@ -1102,7 +1105,7 @@ impl CodexCredentialQuotaService { .map_err(|_| CodexQuotaFetchError::Recovery)? { prepared = recovered; - result = fetch_usage_once(client, &prepared).await; + result = fetch_usage_with_5xx_retry(client, &prepared).await; } match result { Ok(value) => Ok(FetchedCodexQuota { @@ -1143,6 +1146,39 @@ async fn fetch_usage_once( .map_err(CodexQuotaFetchAttemptError::Upstream) } +/// 对 5xx 上游拒绝做有限次指数退避重试（1s/2s），吞掉瞬时抖动。 +/// +/// 4xx（含 402/429）不重试：它们已经走额度状态转换，重试只会放大上游负载。 +async fn fetch_usage_with_5xx_retry( + client: &CodexBackendClient, + prepared: &PreparedCodexRuntimeCredential, +) -> Result { + let mut attempt = 0_u32; + loop { + let result = fetch_usage_once(client, prepared).await; + let retryable = match &result { + Ok(_) => false, + Err(CodexQuotaFetchAttemptError::Upstream(CodexClientError::Upstream { + status, + .. + })) => status.is_server_error(), + Err(_) => false, + }; + if !retryable || attempt >= QUOTA_FETCH_5XX_MAX_RETRIES { + return result; + } + attempt += 1; + let delay = QUOTA_FETCH_5XX_BASE_DELAY.saturating_mul(attempt); + tracing::warn!( + account_id = %prepared.account.id(), + retry_attempt = attempt, + retry_delay_ms = delay.as_millis(), + "OpenAI quota usage 5xx upstream rejection; retrying with backoff" + ); + tokio::time::sleep(delay).await; + } +} + fn eligible_periodic_quota_refresh(account: &ProviderAccount, now: SystemTime) -> bool { account.enabled() && access_token_is_current(account, now) } @@ -1354,7 +1390,8 @@ fn quota_state_transition( } match *status { StatusCode::PAYMENT_REQUIRED => Some(CodexQuotaStateTransition::Exhausted), - StatusCode::TOO_MANY_REQUESTS => { + // 429 与 503 都是"稍后重试"语义：优先尊重 Retry-After，缺失时用默认冷却。 + StatusCode::TOO_MANY_REQUESTS | StatusCode::SERVICE_UNAVAILABLE => { let duration = Duration::from_secs( retry_after_seconds.unwrap_or(DEFAULT_RATE_LIMIT_COOLDOWN.as_secs()), ) diff --git a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs index 435d8b0b5..696291840 100644 --- a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs +++ b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs @@ -24,7 +24,7 @@ use super::{ WebSocketPoolConnectLease, WebSocketPoolConnectOutcome, WebSocketPoolDecision, WebSocketPoolLease, }, - pump::{PumpKeepalive, PumpedWebSocket}, + pump::{PumpKeepalive, PumpLogContext, PumpedWebSocket}, }; pub(crate) const WEBSOCKET_FAST_PATH_BUDGET: Duration = Duration::from_millis(800); @@ -203,10 +203,12 @@ async fn prepare_unpooled_websocket( decision_started_at: Instant, ) -> Result { let permit = acquire_breaker_permit(breaker, origin_key)?; + let context = pump_log_context_from_connection(request.connection()); let connected = connect_with_budget( request.connection(), PumpKeepalive::disabled(), fast_path_budget, + context, ) .await; let (connection, connect_elapsed) = finish_breaker_attempt(permit, connected)?; @@ -302,6 +304,10 @@ fn start_pooled_websocket_connect( let cancellation = connect_lease.cancellation_token(); let fast_path_reporter = permit.fast_path_reporter(); let keepalive = pool.keepalive(); + let context = PumpLogContext::new( + Some(task_key.account_id().to_owned()), + Some(task_key.conversation_id_hash()), + ); let (sender, receiver) = oneshot::channel(); pool.spawn_connect_task(async move { let connected = tokio::select! { @@ -319,7 +325,7 @@ fn start_pooled_websocket_connect( ); return; } - result = connect_with_budget(&connection, keepalive, None) => result, + result = connect_with_budget(&connection, keepalive, None, context) => result, }; match finish_breaker_attempt(permit, connected) { Ok((connection, connect_elapsed)) => { @@ -392,17 +398,28 @@ fn duration_millis_u64(duration: Duration) -> u64 { .max(1) } +/// 非池化建连的 pump 日志上下文：从业务头提取账号归属（会话 hash 只有池 key 才有）。 +fn pump_log_context_from_connection(connection: &CodexWebSocketConnection) -> PumpLogContext { + let account_id = connection + .headers() + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("chatgpt-account-id")) + .map(|(_, value)| value.clone()); + PumpLogContext::new(account_id, None) +} + async fn connect_with_budget( connection: &CodexWebSocketConnection, keepalive: PumpKeepalive, fast_path_budget: Option, + context: PumpLogContext, ) -> Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError> { let started_at = Instant::now(); let connected = match fast_path_budget { - Some(budget) => timeout(budget, connect_pumped_websocket(connection, keepalive)) + Some(budget) => timeout(budget, connect_pumped_websocket(connection, keepalive, context)) .await .map_err(|_| CodexWebSocketExchangeError::FastPathTimeout { timeout: budget })?, - None => connect_pumped_websocket(connection, keepalive).await, + None => connect_pumped_websocket(connection, keepalive, context).await, }?; let (websocket, response) = connected; Ok(( diff --git a/backend/crates/providers/openai/src/transport/websocket/handshake.rs b/backend/crates/providers/openai/src/transport/websocket/handshake.rs index add2367ef..304673902 100644 --- a/backend/crates/providers/openai/src/transport/websocket/handshake.rs +++ b/backend/crates/providers/openai/src/transport/websocket/handshake.rs @@ -30,7 +30,7 @@ use super::{ error::CodexWebSocketExchangeError, model::{CodexWebSocketConnection, CodexWebSocketRequest, WebSocketContinuationRequirement}, pool::CodexWebSocketConnectionMetadata, - pump::{PumpKeepalive, PumpedWebSocket, RawWsStream}, + pump::{PumpKeepalive, PumpLogContext, PumpedWebSocket, RawWsStream}, }; const WEBSOCKET_EXTENSIONS: &str = "permessage-deflate; client_max_window_bits"; @@ -93,9 +93,10 @@ pub fn responses_websocket_endpoint(base_url: &str) -> String { pub(super) async fn connect_pumped_websocket( connection: &CodexWebSocketConnection, keepalive: PumpKeepalive, + context: PumpLogContext, ) -> Result<(PumpedWebSocket, WsResponse>>), CodexWebSocketExchangeError> { let (raw, response) = connect_websocket(connection).await?; - Ok((PumpedWebSocket::new(raw, keepalive), response)) + Ok((PumpedWebSocket::new(raw, keepalive, context), response)) } pub(super) async fn send_websocket_request( diff --git a/backend/crates/providers/openai/src/transport/websocket/pump.rs b/backend/crates/providers/openai/src/transport/websocket/pump.rs index 41a2b6259..2edcf4eb0 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pump.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pump.rs @@ -36,6 +36,22 @@ pub(crate) type RawWsStream = WebSocketStream>; const PUMP_COMMAND_BUFFER: usize = 32; +/// pump 日志的账号级关联上下文；空闲连接被上游重置等场景用于回溯归属。 +#[derive(Debug, Clone, Default)] +pub(crate) struct PumpLogContext { + pub(crate) account_id: Option, + pub(crate) conversation_id_hash: Option, +} + +impl PumpLogContext { + pub(crate) fn new(account_id: Option, conversation_id_hash: Option) -> Self { + Self { + account_id, + conversation_id_hash, + } + } +} + /// pump 保活策略。 #[derive(Debug, Clone, Copy)] pub(crate) struct PumpKeepalive { @@ -188,7 +204,7 @@ impl std::fmt::Debug for PumpedWebSocket { impl PumpedWebSocket { /// 用底层流启动一个 pump 任务并返回句柄。 - pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive) -> Self { + pub(crate) fn new(inner: RawWsStream, keepalive: PumpKeepalive, context: PumpLogContext) -> Self { let (tx_command, rx_command) = mpsc::channel::(PUMP_COMMAND_BUFFER); let (tx_message, rx_message) = mpsc::channel(PUMP_MESSAGE_BUFFER); let connection_id = Uuid::new_v4(); @@ -203,6 +219,7 @@ impl PumpedWebSocket { rx_command, tx_message, keepalive, + context, closed_for_task, exit_reason_for_task, ) @@ -273,12 +290,17 @@ impl Drop for PumpedWebSocket { } } +#[expect( + clippy::too_many_arguments, + reason = "pump owns one flat set of channel/state handles and log context" +)] async fn pump_loop( connection_id: Uuid, mut inner: RawWsStream, mut rx_command: mpsc::Receiver, tx_message: mpsc::Sender>, keepalive: PumpKeepalive, + context: PumpLogContext, closed: Arc, exit_reason: Arc>>, ) { @@ -429,7 +451,7 @@ async fn pump_loop( *exit_reason .lock() .expect("WebSocket pump exit reason lock poisoned") = Some(reason.clone()); - log_pump_exit(connection_id, &reason, backpressure_events); + log_pump_exit(connection_id, &context, &reason, backpressure_events); if reason.should_send_close() { let _ = inner.send(Message::Close(None)).await; @@ -510,11 +532,20 @@ async fn send_keepalive_ping( })) } -fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_events: u64) { +fn log_pump_exit( + connection_id: Uuid, + context: &PumpLogContext, + reason: &PumpExitReason, + backpressure_events: u64, +) { let detail = reason.detail().unwrap_or_default(); + let account_id = context.account_id.as_deref().unwrap_or("unknown"); + let conversation_id_hash = context.conversation_id_hash.as_deref().unwrap_or("unknown"); if reason.is_unexpected() { tracing::warn!( websocket_connection_id = %connection_id, + account_id = %account_id, + conversation_id_hash = %conversation_id_hash, pump_exit_reason = reason.as_str(), pump_exit_detail = detail, backpressure_events, @@ -523,6 +554,8 @@ fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_even } else if matches!(reason, PumpExitReason::UpstreamCloseFrame { .. }) { tracing::info!( websocket_connection_id = %connection_id, + account_id = %account_id, + conversation_id_hash = %conversation_id_hash, pump_exit_reason = reason.as_str(), pump_exit_detail = detail, backpressure_events, @@ -531,6 +564,8 @@ fn log_pump_exit(connection_id: Uuid, reason: &PumpExitReason, backpressure_even } else { tracing::debug!( websocket_connection_id = %connection_id, + account_id = %account_id, + conversation_id_hash = %conversation_id_hash, pump_exit_reason = reason.as_str(), pump_exit_detail = detail, backpressure_events, diff --git a/backend/crates/providers/openai/tests/credential/quota.rs b/backend/crates/providers/openai/tests/credential/quota.rs index 5fb30b43e..cb895cc1d 100644 --- a/backend/crates/providers/openai/tests/credential/quota.rs +++ b/backend/crates/providers/openai/tests/credential/quota.rs @@ -945,6 +945,53 @@ async fn periodic_quota_synchronization_does_not_downgrade_exhaustion_on_rate_li ); } +#[tokio::test] +async fn periodic_quota_synchronization_respects_retry_after_on_service_unavailable() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_periodic_503_retry_after").await; + let account = store + .account("acct_periodic_503_retry_after") + .expect("created account"); + store + .apply_state_change(AccountStateChange { + account_id: account.id().clone(), + expected_revision: account.revision(), + availability: AccountAvailability::QuotaExhausted, + reason: Some("quota_exhausted".to_owned()), + cooldown_until: None, + observed_at: SystemTime::now(), + }) + .await + .expect("mark account exhausted"); + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/codex/usage")) + // 5xx 短退避重试后仍失败，Retry-After 用于写冷却，不再落入"未分类"。 + .respond_with(ResponseTemplate::new(503).insert_header("Retry-After", "120")) + .mount(&server) + .await; + let service = quota_service_with_base_url( + &store, + reqwest::Client::builder() + .no_proxy() + .build() + .expect("client"), + server.uri(), + ); + + service.synchronize().await.expect("quota synchronization"); + + let account = store + .account("acct_periodic_503_retry_after") + .expect("persisted account"); + assert_eq!(account.availability(), AccountAvailability::QuotaExhausted); + assert!( + account + .cooldown_until() + .is_some_and(|until| until > SystemTime::now()) + ); +} + #[tokio::test] async fn periodic_quota_synchronization_skips_ordinary_cooldown_without_exhausted_quota() { let store = Arc::new(MemoryAccountStore::default()); @@ -1028,7 +1075,8 @@ async fn periodic_quota_synchronization_rechecks_exhausted_accounts_without_wait Mock::given(method("GET")) .and(path("/api/codex/usage")) .respond_with(ResponseTemplate::new(500)) - .expect(1) + // 5xx 短退避重试 1s/2s 后仍失败，最终计入 transient。 + .expect(3) .mount(&server) .await; let service = quota_service_with_base_url( diff --git a/backend/crates/providers/openai/tests/transport/profile.rs b/backend/crates/providers/openai/tests/transport/profile.rs index 79614581b..89bccdffc 100644 --- a/backend/crates/providers/openai/tests/transport/profile.rs +++ b/backend/crates/providers/openai/tests/transport/profile.rs @@ -66,7 +66,7 @@ fn wire_profile_should_generate_codex_core_user_agent() { codex_version: "0.144.2".to_owned(), desktop_version: "26.707.72221".to_owned(), desktop_build: "72221".to_owned(), - os_type: "Mac OS".to_owned(), + os_type: "macOS".to_owned(), os_version: "15.7.1".to_owned(), arch: "arm64".to_owned(), terminal: "unknown".to_owned(), @@ -78,7 +78,7 @@ fn wire_profile_should_generate_codex_core_user_agent() { assert_eq!( profile.user_agent(), - "Codex Desktop/0.144.2 (Mac OS 15.7.1; arm64) unknown" + "Codex Desktop/0.144.2 (macOS 15.7.1; arm64) unknown" ); } @@ -219,7 +219,7 @@ fn wire_profile() -> CodexWireProfile { codex_version: "0.144.2".to_owned(), desktop_version: "26.707.72221".to_owned(), desktop_build: "72221".to_owned(), - os_type: "Mac OS".to_owned(), + os_type: "macOS".to_owned(), os_version: "15.7.1".to_owned(), arch: "arm64".to_owned(), terminal: "unknown".to_owned(),