From 09982d7221ef332e8b9d4409a429b0971eec38d6 Mon Sep 17 00:00:00 2001 From: Zyy Date: Mon, 31 Aug 2026 19:13:46 +0800 Subject: [PATCH] fix(openai): align websocket retry fallback Co-authored-by: Codex --- .../gateway-core/src/engine/coordinator.rs | 140 +++-- backend/crates/gateway-core/src/engine/mod.rs | 2 + backend/crates/gateway-core/src/error.rs | 39 +- .../gateway-core/tests/engine/coordinator.rs | 60 +++ backend/crates/providers/openai/src/config.rs | 19 + .../openai/src/credential/affinity.rs | 3 + .../providers/openai/src/credential/mod.rs | 3 +- .../openai/src/credential/selector.rs | 7 +- backend/crates/providers/openai/src/lib.rs | 2 + .../crates/providers/openai/src/provider.rs | 307 ++++++++++- .../providers/openai/src/session_transport.rs | 106 ++++ .../providers/openai/src/transport/client.rs | 1 + .../openai/src/transport/client_sse.rs | 134 +++-- .../openai/src/transport/websocket/error.rs | 27 +- .../transport/websocket/exchange/reducer.rs | 41 +- .../transport/websocket/exchange/stream.rs | 52 +- .../crates/providers/openai/tests/config.rs | 18 +- .../openai/tests/provider/contract.rs | 483 +++++++++++++++--- .../openai/tests/transport/websocket.rs | 36 +- 19 files changed, 1242 insertions(+), 238 deletions(-) create mode 100644 backend/crates/providers/openai/src/session_transport.rs diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index d7fceeb48..a7f2e833e 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -162,11 +162,12 @@ where account_state_owner, cancellation, attempts: 0, + routing_attempts: 0, candidate_index, excluded_accounts: BTreeSet::new(), credential_recovery_attempted_accounts: BTreeSet::new(), recovery_account: None, - transport_fallback_account: None, + transport_recovery: None, current: None, send_state_watermark: UpstreamSendState::NotSent, downstream_committed_at: None, @@ -217,6 +218,13 @@ struct FailureFinalization { retry_after_ms: Option, } +#[derive(Debug, Clone)] +struct PendingTransportRecovery { + account: crate::engine::credential::ProviderAccountId, + transport: AttemptTransport, + delay: Duration, +} + /// API 可逐事件消费的 Core 执行会话。 /// /// API 只能提交下游 delivery 边界；账号重试、断流终结与 @@ -238,7 +246,10 @@ pub struct ResponseExecutionSession { continuation_attempt: ContinuationAttempt, account_state_owner: Option, cancellation: CancellationToken, + /// 所有实际上游 attempt 数；包含同账号传输重试，作为持久化序号。 attempts: u32, + /// 路由预算只统计正常选号/账号恢复，不被 Provider-owned 传输预算消耗。 + routing_attempts: u32, candidate_index: usize, excluded_accounts: BTreeSet, credential_recovery_attempted_accounts: BTreeSet, @@ -246,9 +257,9 @@ pub struct ResponseExecutionSession { /// attempt 建立时即被消费，后续可重试错误仍可换号消耗剩余重试预算。 /// 与 `required_account`（外部指定、贯穿整个请求）语义不同，不可合并。 recovery_account: Option, - /// 上游传输在提交边界前失败后的一次性同账号备用传输钉选。 - /// 与凭据恢复分开保存，避免把 transport fallback 误记为 OAuth 恢复。 - transport_fallback_account: Option, + /// 上游传输在提交边界前失败后的一次性同账号重试或备用传输钉选。 + /// 与凭据恢复分开保存，避免把 transport recovery 误记为 OAuth 恢复。 + transport_recovery: Option, current: Option, /// 请求级发送状态水位；跨 attempt 单调不降，终态写回不得低于此档。 send_state_watermark: UpstreamSendState, @@ -590,9 +601,25 @@ where } async fn prepare_attempt(&mut self) -> Result, EngineError> { - if self.attempts >= self.plan.max_attempts().get() { + let transport_recovery = self.transport_recovery.take(); + if transport_recovery.is_none() && self.routing_attempts >= self.plan.max_attempts().get() { return Err(EngineError::EmptyRoutingPlan); } + if let Some(recovery) = transport_recovery.as_ref() + && !recovery.delay.is_zero() + { + match poll_retry_delay(recovery.delay, self.cancellation.clone(), self.deadline).await { + RetryDelayBoundary::Elapsed => {} + RetryDelayBoundary::Cancelled => { + self.finish_interruption(EngineError::Cancelled).await?; + return Err(EngineError::Cancelled); + } + RetryDelayBoundary::Deadline => { + self.finish_interruption(EngineError::Deadline).await?; + return Err(EngineError::Deadline); + } + } + } let Some(candidate) = self.plan.candidates().get(self.candidate_index).cloned() else { let error = GatewayError::new( GatewayErrorKind::NoAvailableProvider, @@ -616,23 +643,20 @@ where .ok_or(EngineError::EmptyRoutingPlan)?; // 请求局部恢复钉选在此被一次性消费，只绑定本次 replay attempt； // 外部 required_account 每次 attempt 都重新生效。 - let (pinned_account, attempt_transport) = match &self.account_selection { - AccountSelection::Diagnostic(account) => { - (Some(account.clone()), AttemptTransport::Default) - } - AccountSelection::Scheduled(_) => { - let transport_fallback_account = self.transport_fallback_account.take(); - let transport = if transport_fallback_account.is_some() { - AttemptTransport::Fallback - } else { - AttemptTransport::Default - }; - ( - transport_fallback_account - .or_else(|| self.recovery_account.take()) + let is_transport_recovery = transport_recovery.is_some(); + let (pinned_account, attempt_transport) = if let Some(recovery) = transport_recovery { + (Some(recovery.account), recovery.transport) + } else { + match &self.account_selection { + AccountSelection::Diagnostic(account) => { + (Some(account.clone()), AttemptTransport::Default) + } + AccountSelection::Scheduled(_) => ( + self.recovery_account + .take() .or_else(|| self.account_selection.required_account().cloned()), - transport, - ) + AttemptTransport::Default, + ), } }; let account_context = match &self.account_selection { @@ -881,6 +905,9 @@ where } } self.attempts = next_attempt.get(); + if !is_transport_recovery { + self.routing_attempts = self.routing_attempts.saturating_add(1); + } self.current = Some(CurrentAttempt { stream, metadata, @@ -1034,35 +1061,50 @@ where attempt_send_state, provider_proved_replay_safe, ); - let pre_delivery_retry = self.account_selection.required_account().is_none() + let account_rotation_retry = self.account_selection.required_account().is_none() && self.continuation_attempt == ContinuationAttempt::None && self.downstream_committed_at.is_none() && !self.delivery_pending - && error.allows_pre_delivery_retry() - && self.attempts < self.plan.max_attempts().get(); - let transport_fallback_retry = pre_delivery_retry && matches!( error.pre_delivery_retry(), - Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback) - ); + Some(crate::error::PreDeliveryRetry::AccountRotation) + ) + && self.routing_attempts < self.plan.max_attempts().get(); + let transport_recovery = match error.pre_delivery_retry() { + Some(crate::error::PreDeliveryRetry::SameAccountTransportRetry { + retry_index, + delay, + }) if self.downstream_committed_at.is_none() && !self.delivery_pending => { + Some((AttemptTransport::Retry(retry_index), delay)) + } + Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback) + if self.downstream_committed_at.is_none() && !self.delivery_pending => + { + Some((AttemptTransport::Fallback, Duration::ZERO)) + } + _ => None, + }; let ordinary_retry = self.account_selection.required_account().is_none() && self.continuation_attempt == ContinuationAttempt::None && self.downstream_committed_at.is_none() && !self.delivery_pending && attempt_send_state != UpstreamSendState::Ambiguous && provider_proved_replay_safe - && self.attempts < self.plan.max_attempts().get(); + && self.routing_attempts < self.plan.max_attempts().get(); let same_account_retry = error.retries_same_account() && provider_proved_replay_safe && self.downstream_committed_at.is_none() && !self.delivery_pending && attempt_send_state != UpstreamSendState::Ambiguous - && self.attempts < self.plan.max_attempts().get() + && self.routing_attempts < self.plan.max_attempts().get() && !self .credential_recovery_attempted_accounts .contains(current.metadata.provider_account_id()); - let retryable = - continuation_retry || same_account_retry || ordinary_retry || pre_delivery_retry; + let retryable = continuation_retry + || same_account_retry + || ordinary_retry + || account_rotation_retry + || transport_recovery.is_some(); if retryable { // 普通 clone 只保留稳定事实；原始 wire/HTTP response 由 request-local @@ -1075,9 +1117,12 @@ where // 只钉住紧随其后的 replay attempt；replay 再遇可重试错误时， // ordinary/continuation 重试门不受影响，仍可换号。 self.recovery_account = Some(account); - } else if transport_fallback_retry { - self.transport_fallback_account = - Some(current.metadata.provider_account_id().clone()); + } else if let Some((transport, delay)) = transport_recovery { + self.transport_recovery = Some(PendingTransportRecovery { + account: current.metadata.provider_account_id().clone(), + transport, + delay, + }); } else if !continuation_retry { self.excluded_accounts .insert(current.metadata.provider_account_id().clone()); @@ -1155,7 +1200,7 @@ where || self.delivery_pending || send_state == UpstreamSendState::Ambiguous || !provider_proved_replay_safe - || self.attempts >= self.plan.max_attempts().get() + || self.routing_attempts >= self.plan.max_attempts().get() || self .operation .provider_session_state(current.metadata.provider().as_str()) @@ -1636,6 +1681,31 @@ enum ProviderBoundary { Deadline, } +enum RetryDelayBoundary { + Elapsed, + Cancelled, + Deadline, +} + +async fn poll_retry_delay( + delay: Duration, + cancellation: CancellationToken, + deadline: SystemTime, +) -> RetryDelayBoundary { + let Ok(remaining) = deadline.duration_since(SystemTime::now()) else { + return RetryDelayBoundary::Deadline; + }; + let retry_delay = Delay::new(delay).fuse(); + let cancelled = cancellation.cancelled().fuse(); + let timeout = Delay::new(remaining).fuse(); + pin_mut!(retry_delay, cancelled, timeout); + select_biased! { + _ = cancelled => RetryDelayBoundary::Cancelled, + _ = timeout => RetryDelayBoundary::Deadline, + _ = retry_delay => RetryDelayBoundary::Elapsed, + } +} + async fn poll_provider( provider: Arc, request: ProviderRequest, diff --git a/backend/crates/gateway-core/src/engine/mod.rs b/backend/crates/gateway-core/src/engine/mod.rs index 0e10502d9..ae07f4124 100644 --- a/backend/crates/gateway-core/src/engine/mod.rs +++ b/backend/crates/gateway-core/src/engine/mod.rs @@ -112,6 +112,8 @@ pub enum AttemptTransport { /// 使用 Provider 的默认传输策略。 #[default] Default, + /// 固定账号重试 Provider 的首选传输；序号由 Provider 的独立预算驱动。 + Retry(NonZeroU32), /// 使用 Provider 定义的备用传输。 Fallback, } diff --git a/backend/crates/gateway-core/src/error.rs b/backend/crates/gateway-core/src/error.rs index 0f67828e7..db7998481 100644 --- a/backend/crates/gateway-core/src/error.rs +++ b/backend/crates/gateway-core/src/error.rs @@ -1,6 +1,7 @@ //! 网关核心使用的稳定错误分类。 use std::fmt; +use std::num::NonZeroU32; use std::time::Duration; use bytes::Bytes; @@ -202,6 +203,13 @@ pub enum ContinuationFailure { pub enum PreDeliveryRetry { /// 排除本次账号，按普通调度策略重新选号。 AccountRotation, + /// 固定本次账号，并按 Provider 给出的序号重试当前传输。 + SameAccountTransportRetry { + /// Provider-owned 传输重试序号，从 1 开始。 + retry_index: NonZeroU32, + /// 发起下一次传输尝试前的退避时长。 + delay: Duration, + }, /// 固定本次账号，并要求 Provider 使用备用传输。 SameAccountTransportFallback, } @@ -475,10 +483,39 @@ impl ProviderError { /// 下游已提交后的隐藏重放。 #[must_use] pub const fn with_pre_delivery_transport_fallback(mut self) -> Self { - self.pre_delivery_retry = Some(PreDeliveryRetry::SameAccountTransportFallback); + self.set_pre_delivery_transport_fallback(); self } + /// 原地设置同账号备用传输恢复，保留当前错误携带的原客户端响应。 + pub const fn set_pre_delivery_transport_fallback(&mut self) { + self.pre_delivery_retry = Some(PreDeliveryRetry::SameAccountTransportFallback); + } + + /// 允许 Core 在首个客户端事件前固定原账号并重试当前 Provider 传输。 + /// + /// 重试预算和退避策略由 Provider 所属协议定义；Core 只负责同账号钉选、 + /// deadline/取消边界以及下游 commit barrier。 + #[must_use] + pub const fn with_pre_delivery_transport_retry( + mut self, + retry_index: NonZeroU32, + delay: Duration, + ) -> Self { + self.set_pre_delivery_transport_retry(retry_index, delay); + self + } + + /// 原地设置同账号当前传输重试，保留当前错误携带的原客户端响应。 + pub const fn set_pre_delivery_transport_retry( + &mut self, + retry_index: NonZeroU32, + delay: Duration, + ) { + self.pre_delivery_retry = + Some(PreDeliveryRetry::SameAccountTransportRetry { retry_index, delay }); + } + /// 要求 Core 在账号凭据已恢复后仅对原账号重放一次。 #[must_use] pub const fn with_same_account_retry(mut self) -> Self { diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index 5ac141324..4590f79d5 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -2797,6 +2797,66 @@ fn pre_delivery_transport_fallback_retries_the_same_account_with_fallback_transp assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded); } +#[test] +fn provider_owned_transport_retries_keep_the_same_account_until_fallback() { + let operation = generate_operation(); + let route_plan = plan(&operation); + let retry_one = NonZeroU32::new(1).expect("non-zero retry index"); + let retry_two = NonZeroU32::new(2).expect("non-zero retry index"); + let transport_error = + || ProviderError::new(ProviderErrorKind::Transport, UpstreamSendState::Ambiguous); + let (coordinator, store, provider) = coordinator(vec![ + Script::Stream { + account_id: "acct_first", + items: vec![Err( + transport_error().with_pre_delivery_transport_retry(retry_one, Duration::ZERO) + )], + }, + Script::Stream { + account_id: "acct_first", + items: vec![Err( + transport_error().with_pre_delivery_transport_retry(retry_two, Duration::ZERO) + )], + }, + Script::Stream { + account_id: "acct_first", + items: vec![Err(transport_error().with_pre_delivery_transport_fallback())], + }, + Script::Stream { + account_id: "acct_first", + items: complete_stream(None), + }, + ]); + + let mut session = block_on(coordinator.start( + model_request(&operation, SystemTime::now() + Duration::from_secs(30)), + operation, + route_plan, + None, + None, + CancellationToken::new(), + )) + .expect("start execution"); + block_on(session.collect_uncommitted()).expect("fallback succeeds"); + block_on(session.commit_downstream(Some(200))).expect("commit winning response"); + + let original = ProviderAccountId::new("acct_first").expect("account id"); + let contexts = provider.contexts.lock().expect("contexts lock"); + assert_eq!(contexts.len(), 4); + assert_eq!(contexts[0].transport(), AttemptTransport::Default); + assert_eq!(contexts[1].transport(), AttemptTransport::Retry(retry_one)); + assert_eq!(contexts[2].transport(), AttemptTransport::Retry(retry_two)); + assert_eq!(contexts[3].transport(), AttemptTransport::Fallback); + for context in contexts.iter().skip(1) { + assert_eq!(context.required_account(), Some(&original)); + assert!(!context.excluded_accounts().contains(&original)); + } + let state = store.state.lock().expect("store lock"); + assert_eq!(state.intermediate_failures, 3); + assert_eq!(state.finalizations[0].attempt_count, 4); + assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded); +} + #[test] fn final_capacity_exhaustion_returns_the_last_retryable_upstream_failure() { let operation = generate_operation(); diff --git a/backend/crates/providers/openai/src/config.rs b/backend/crates/providers/openai/src/config.rs index 2332353e0..2aea60481 100644 --- a/backend/crates/providers/openai/src/config.rs +++ b/backend/crates/providers/openai/src/config.rs @@ -18,6 +18,15 @@ use crate::{ }, }; +/// 本服务的流式请求默认重试次数。 +pub const DEFAULT_STREAM_MAX_RETRIES: u64 = 3; +/// 防止错误配置产生无界隐藏重放；与官方 Codex 的硬上限一致。 +pub const MAX_STREAM_MAX_RETRIES: u64 = 100; + +const fn default_stream_max_retries() -> u64 { + DEFAULT_STREAM_MAX_RETRIES +} + /// OpenAI Provider 唯一启动配置。 #[derive(Debug, Clone, Deserialize, PartialEq, Eq)] #[serde(deny_unknown_fields)] @@ -30,6 +39,8 @@ pub struct OpenAiConfig { pub quota: CodexQuotaSettings, #[serde(default)] pub auth: CodexAuthSettings, + #[serde(default = "default_stream_max_retries")] + pub stream_max_retries: u64, pub wire_profile: CodexWireProfileConfig, #[serde(skip)] identity_secret_path: PathBuf, @@ -88,6 +99,13 @@ impl OpenAiConfig { self.auth.refresh_enabled } + /// 返回经官方同款硬上限约束后的上游流重试预算。 + #[must_use] + pub fn stream_max_retries(&self) -> u32 { + u32::try_from(self.stream_max_retries.min(MAX_STREAM_MAX_RETRIES)) + .unwrap_or(MAX_STREAM_MAX_RETRIES as u32) + } + pub(crate) fn session_identity( &self, ) -> Result { @@ -102,6 +120,7 @@ impl Default for OpenAiConfig { ws_pool: CodexWebSocketPoolSettings::default(), quota: CodexQuotaSettings::default(), auth: CodexAuthSettings::default(), + stream_max_retries: DEFAULT_STREAM_MAX_RETRIES, wire_profile: CodexWireProfileConfig::default(), identity_secret_path: PathBuf::new(), } diff --git a/backend/crates/providers/openai/src/credential/affinity.rs b/backend/crates/providers/openai/src/credential/affinity.rs index 706aa7e32..d0b95c10b 100644 --- a/backend/crates/providers/openai/src/credential/affinity.rs +++ b/backend/crates/providers/openai/src/credential/affinity.rs @@ -1,5 +1,7 @@ //! OpenAI 根会话锚点到 Store 不透明账号亲和键及诊断上下文的单向派生。 +use std::time::Duration; + use gateway_core::policy::ClientApiKeyId; use gateway_core::provider_ports::ProviderSessionAffinityKey; use sha2::{Digest, Sha256}; @@ -8,6 +10,7 @@ use crate::transport::protocol::responses::CodexResponsesRequest; use crate::transport::request::derive_conversation_anchor; const AFFINITY_KEY_HASH_LENGTH: usize = 12; +pub(crate) const CODEX_ROOT_SESSION_TTL: Duration = Duration::from_secs(4 * 60 * 60); /// 一次请求派生出的账号亲和键及其结构化日志上下文。 pub(crate) struct CodexSessionAffinity { diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index 5aac12c66..59f795b42 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -16,7 +16,8 @@ pub mod token_client; mod types; pub(crate) use affinity::{ - CodexSessionAffinity, derive_codex_cyber_policy_session_key, derive_codex_session_affinity, + CODEX_ROOT_SESSION_TTL, CodexSessionAffinity, derive_codex_cyber_policy_session_key, + derive_codex_session_affinity, }; pub(crate) use oauth::oauth_owner_ref; pub(crate) use types::parse_access_token_expiration; diff --git a/backend/crates/providers/openai/src/credential/selector.rs b/backend/crates/providers/openai/src/credential/selector.rs index a84af2eff..b45bd0622 100644 --- a/backend/crates/providers/openai/src/credential/selector.rs +++ b/backend/crates/providers/openai/src/credential/selector.rs @@ -22,7 +22,7 @@ use secrecy::ExposeSecret; use thiserror::Error; use url::Url; -use super::affinity::CodexSessionAffinity; +use super::affinity::{CODEX_ROOT_SESSION_TTL, CodexSessionAffinity}; use super::catalog::CodexCredentialCatalogService; use super::cookie::CodexCookiePolicy; use super::quota::CodexCredentialQuotaService; @@ -41,7 +41,6 @@ const CLOUDFLARE_CHALLENGE_BACKOFF: [Duration; 4] = [ Duration::from_secs(120), ]; const CLOUDFLARE_PATH_BLOCK_THRESHOLD: u32 = 3; -const SESSION_AFFINITY_TTL: Duration = Duration::from_secs(4 * 60 * 60); const SESSION_AFFINITY_TIMEOUT: Duration = Duration::from_millis(100); const CYBER_POLICY_SESSION_TTL: Duration = Duration::from_secs(60 * 60); @@ -681,7 +680,7 @@ impl CodexCredentialSelector { &self.provider_kind, key, selected_account_id, - SESSION_AFFINITY_TTL, + CODEX_ROOT_SESSION_TTL, ), ) .await @@ -955,7 +954,7 @@ impl CodexCredentialSelector { key, expected_affinity_account_id, account.id(), - SESSION_AFFINITY_TTL, + CODEX_ROOT_SESSION_TTL, ), ) .await diff --git a/backend/crates/providers/openai/src/lib.rs b/backend/crates/providers/openai/src/lib.rs index 3435a2e0b..db574db54 100644 --- a/backend/crates/providers/openai/src/lib.rs +++ b/backend/crates/providers/openai/src/lib.rs @@ -3,6 +3,7 @@ mod admin; pub mod config; mod provider; +mod session_transport; use std::sync::Arc; @@ -138,6 +139,7 @@ pub async fn initialize( profile.clone(), config.base_url().to_owned(), Arc::clone(&websocket_pool), + config.stream_max_retries(), ) .map_err(OpenAiInitializeError::Provider)? .with_session_identity(session_identity), diff --git a/backend/crates/providers/openai/src/provider.rs b/backend/crates/providers/openai/src/provider.rs index c6e086c10..f779096fa 100644 --- a/backend/crates/providers/openai/src/provider.rs +++ b/backend/crates/providers/openai/src/provider.rs @@ -2,6 +2,7 @@ use std::collections::BTreeSet; use std::fmt; +use std::num::NonZeroU32; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime}; @@ -57,6 +58,7 @@ use crate::credential::{ SelectCodexCredential, SelectCodexProviderEndpointCredential, derive_codex_cyber_policy_session_key, derive_codex_session_affinity, }; +use crate::session_transport::CodexSessionTransportFallbacks; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, }; @@ -71,8 +73,9 @@ use crate::transport::profile::{ }; use crate::transport::protocol::responses::{ CodexResponsesRequest, PREVIOUS_RESPONSE_NOT_FOUND_CODE, PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE, - PreviousResponseScope, ResponseEventSignals, transport_requirement, + PreviousResponseScope, ResponseEventSignals, TransportRequirement, transport_requirement, }; +use crate::transport::protocol::websocket::WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE; use crate::transport::request::{ CodexRequestEncodeError, RequestAccountScope, encode_generate_request, scope_request_to_account, }; @@ -126,6 +129,8 @@ pub struct CodexProvider { image_generations_url: Url, image_edits_url: Url, session_identity: Option, + session_transport_fallbacks: CodexSessionTransportFallbacks, + stream_max_retries: u32, } impl CodexProvider { @@ -140,6 +145,7 @@ impl CodexProvider { profile: CodexWireProfileState, base_url: String, websocket_pool: Arc, + stream_max_retries: u32, ) -> Result { let responses_url = Url::parse(&endpoint_url(&base_url, CODEX_RESPONSES_PATH)) .map_err(|_| CodexProviderConfigError::InvalidBaseUrl)?; @@ -160,6 +166,8 @@ impl CodexProvider { image_generations_url, image_edits_url, session_identity: None, + session_transport_fallbacks: CodexSessionTransportFallbacks::default(), + stream_max_retries, }) } @@ -288,11 +296,6 @@ impl Provider for CodexProvider { derive_codex_session_affinity(&upstream_request, context.client_api_key_ref()); let cyber_policy_session_key = derive_codex_cyber_policy_session_key(&upstream_request, context.client_api_key_ref()); - let transport = match context.transport() { - AttemptTransport::Default => selected_transport(&upstream_request), - AttemptTransport::Fallback => CodexProviderTransport::HttpOnly, - }; - apply_transport(&mut upstream_request, transport); let selection_started_at = Instant::now(); let lease = self @@ -388,6 +391,39 @@ impl Provider for CodexProvider { lease.installation_id(), account_scope, ); + let requirement = transport_requirement(&upstream_request); + let requested_transport = selected_transport(&upstream_request); + let sticky_http_fallback = context.transport() == AttemptTransport::Default + && requested_transport == CodexProviderTransport::PreferWebSocket + && !requirement.requires_websocket() + && session_affinity.as_ref().is_some_and(|affinity| { + self.session_transport_fallbacks + .is_http_only(affinity.key()) + }); + let transport = if requirement.requires_websocket() { + CodexProviderTransport::PreferWebSocket + } else { + match context.transport() { + AttemptTransport::Default if sticky_http_fallback => { + CodexProviderTransport::HttpOnly + } + AttemptTransport::Default => requested_transport, + AttemptTransport::Retry(_) => requested_transport, + AttemptTransport::Fallback => CodexProviderTransport::HttpOnly, + } + }; + if sticky_http_fallback { + tracing::info!( + request_id = %context.request_id(), + attempt_index = context.attempt_index().get(), + session_affinity_key_hash = session_affinity + .as_ref() + .map_or("", CodexSessionAffinity::key_hash), + transport_requirement = requirement.as_str(), + "Reusing session-sticky OpenAI upstream HTTP fallback" + ); + } + apply_transport(&mut upstream_request, transport); let metadata = ProviderCallMetadata::new( provider_kind, upstream_model.clone(), @@ -411,7 +447,14 @@ impl Provider for CodexProvider { continuation_scope: None, }); let allows_account_state_mutation = lease.allows_account_state_mutation(); + let session_affinity_key_hash = session_affinity + .as_ref() + .map(|affinity| affinity.key_hash().to_owned()); let session_affinity_key = session_affinity.map(CodexSessionAffinity::into_key); + let websocket_retry_count = match context.transport() { + AttemptTransport::Retry(retry_index) => retry_index.get(), + AttemptTransport::Default | AttemptTransport::Fallback => 0, + }; let events = cold_response_stream(ColdResponse { client: self.client.clone(), response_origin: self.responses_url.clone(), @@ -425,6 +468,10 @@ impl Provider for CodexProvider { lease: Arc::clone(&lease), output_started_at, session_affinity_key, + session_affinity_key_hash, + session_transport_fallbacks: self.session_transport_fallbacks.clone(), + websocket_retry_count, + stream_max_retries: self.stream_max_retries, session_capture, }); let stream = ProviderStream::new(metadata, events, lease); @@ -522,6 +569,10 @@ struct ColdResponse { lease: Arc, output_started_at: Instant, session_affinity_key: Option, + session_affinity_key_hash: Option, + session_transport_fallbacks: CodexSessionTransportFallbacks, + websocket_retry_count: u32, + stream_max_retries: u32, session_capture: Option, } @@ -663,7 +714,7 @@ async fn create_response_attempt( biased; _ = cancellation.cancelled() => Err(CodexHandshakeAttemptError::Cancelled), _ = tokio::time::sleep(handshake_deadline) => Err(CodexHandshakeAttemptError::Timeout), - response = client.create_response_stream_with_pool_account( + response = client.create_response_stream_with_deferred_websocket_recovery( request, request_context, Some(account_id), @@ -839,6 +890,10 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { lease, output_started_at, session_affinity_key, + session_affinity_key_hash, + session_transport_fallbacks, + websocket_retry_count, + stream_max_retries, mut session_capture, } = response; Box::pin(async_stream::try_stream! { @@ -870,6 +925,7 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { lease.escape_reason(), lease.account_switch(), ); + let request_transport_requirement = transport_requirement(&request); let response = create_response_attempt( &client, &request, @@ -887,6 +943,14 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { &cancellation, ) .await; + let websocket_failure_policy = match &response { + Err(CodexHandshakeAttemptError::Client(error)) + if transport_policy == CodexProviderTransport::PreferWebSocket => + { + websocket_client_failure_policy(error) + } + _ => None, + }; if let Err(CodexHandshakeAttemptError::Client(error)) = &response { log_client_upstream_error( UpstreamErrorLogContext::new(&context, &active_account, None), @@ -897,6 +961,23 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { let response = match response { Ok(response) => response, Err(mut failure) => { + if let Some(policy) = websocket_failure_policy { + apply_websocket_recovery_policy( + &mut failure, + WebSocketRecoveryContext { + policy, + requirement: request_transport_requirement, + retry_count: websocket_retry_count, + max_retries: stream_max_retries, + request_id: context.request_id().as_str(), + attempt_index: context.attempt_index().get(), + account_id: active_account.id().as_str(), + session_affinity_key: session_affinity_key.as_ref(), + session_affinity_key_hash: session_affinity_key_hash.as_deref(), + session_transport_fallbacks: &session_transport_fallbacks, + }, + ); + } if let Some(observation) = failure.observation.take() { yield ProviderEvent::observation(observation); } @@ -1029,12 +1110,25 @@ fn cold_response_stream(response: ColdResponse) -> EventStream { continue; } Err(mut failure) => { - if failure.error.allows_pre_delivery_retry() + if failure.websocket_transport_retryable && response_transport == CodexBackendTransport::WebSocket - && transport_requirement(&request).allows_pre_delivery_http_fallback() && !pre_commit_events.is_committed() { - failure.error = failure.error.with_pre_delivery_transport_fallback(); + apply_websocket_recovery_policy( + &mut failure, + WebSocketRecoveryContext { + policy: WebSocketFailurePolicy::Budgeted, + requirement: request_transport_requirement, + retry_count: websocket_retry_count, + max_retries: stream_max_retries, + request_id: context.request_id().as_str(), + attempt_index: context.attempt_index().get(), + account_id: active_account.id().as_str(), + session_affinity_key: session_affinity_key.as_ref(), + session_affinity_key_hash: session_affinity_key_hash.as_deref(), + session_transport_fallbacks: &session_transport_fallbacks, + }, + ); } let updates = take_rate_limit_updates(rate_limit_updates.as_ref()).await; if !updates.is_empty() { @@ -1645,7 +1739,11 @@ fn codex_response_observation( WebSocketPoolKind::New }); } - if let Some(status_code) = diagnostics.status_code { + // WebSocket opening 的 101 是成功升级事实，不是业务请求的失败 HTTP 状态。 + // opening 明确拒绝仍由 `codex_error_observation` 保存真实上游状态。 + if transport != CodexBackendTransport::WebSocket + && let Some(status_code) = diagnostics.status_code + { observation = observation.with_status_code(status_code); } if let Some(request_id) = diagnostics.request_id.as_deref() { @@ -2084,6 +2182,7 @@ fn map_selection_error(error: CredentialSelectionError) -> ProviderError { struct MappedProviderFailure { error: ProviderError, + websocket_transport_retryable: bool, account_failure: Option, /// 原始上游错误描述，仅在凭据错误状态下持久化。 error_message: Option, @@ -2098,6 +2197,7 @@ impl MappedProviderFailure { fn plain(error: ProviderError) -> Self { Self { error, + websocket_transport_retryable: false, account_failure: None, error_message: None, cyber_policy_failure: false, @@ -2311,7 +2411,7 @@ fn log_client_upstream_error(context: UpstreamErrorLogContext<'_>, error: &Codex .websocket_connection_id .map(|connection_id| connection_id.to_string()) .unwrap_or_default(); - let upstream_error_raw = close.reason().unwrap_or_default(); + let close_reason_bytes = close.reason().map_or(0, str::len); tracing::warn!( request_id = %context.request_id, account_id = %context.account_id, @@ -2321,8 +2421,11 @@ fn log_client_upstream_error(context: UpstreamErrorLogContext<'_>, error: &Codex upstream_error_kind = "websocket_close", upstream_close_code = close.code().unwrap_or_default(), upstream_close_code_present = close.code().is_some(), - upstream_error_raw, - upstream_error_raw_bytes = upstream_error_raw.len(), + upstream_last_event_type = close.last_event_type().unwrap_or_default(), + upstream_last_event_type_present = close.last_event_type().is_some(), + upstream_terminal_seen = false, + upstream_close_reason_bytes = close_reason_bytes, + upstream_close_reason_present = close.reason().is_some(), upstream_error_raw_present = close.reason().is_some(), "OpenAI upstream WebSocket closed before terminal event" ); @@ -2480,6 +2583,151 @@ fn map_handshake_error(error: CodexClientError) -> MappedProviderFailure { map_client_error(error, UpstreamSendState::Ambiguous, true) } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum WebSocketFailurePolicy { + Budgeted, + ImmediateFallback, +} + +struct WebSocketRecoveryContext<'a> { + policy: WebSocketFailurePolicy, + requirement: TransportRequirement, + retry_count: u32, + max_retries: u32, + request_id: &'a str, + attempt_index: u32, + account_id: &'a str, + session_affinity_key: Option<&'a ProviderSessionAffinityKey>, + session_affinity_key_hash: Option<&'a str>, + session_transport_fallbacks: &'a CodexSessionTransportFallbacks, +} + +fn websocket_client_failure_policy(error: &CodexClientError) -> Option { + match error { + CodexClientError::Upstream { + status, + transport: CodexBackendTransport::WebSocket, + .. + } if *status == reqwest::StatusCode::UPGRADE_REQUIRED => { + Some(WebSocketFailurePolicy::ImmediateFallback) + } + CodexClientError::Upstream { + status, + transport: CodexBackendTransport::WebSocket, + .. + } if status.is_server_error() + || matches!( + *status, + reqwest::StatusCode::REQUEST_TIMEOUT + | reqwest::StatusCode::CONFLICT + | reqwest::StatusCode::TOO_EARLY + | reqwest::StatusCode::TOO_MANY_REQUESTS + ) => + { + Some(WebSocketFailurePolicy::Budgeted) + } + CodexClientError::WebSocket(error) + if !matches!(error, CodexWebSocketExchangeError::InvalidRequest(_)) => + { + Some(WebSocketFailurePolicy::Budgeted) + } + _ => None, + } +} + +fn apply_websocket_recovery_policy( + failure: &mut MappedProviderFailure, + context: WebSocketRecoveryContext<'_>, +) { + let recovery_is_safe = match failure.error.send_state() { + UpstreamSendState::NotSent => context.requirement.allows_pre_send_http_fallback(), + UpstreamSendState::Sent | UpstreamSendState::Ambiguous => { + context.requirement.allows_pre_delivery_http_fallback() + } + }; + if !recovery_is_safe { + return; + } + + // 传输恢复必须保持原账号可调度；最终 HTTP attempt 若仍失败，再按真实 HTTP + // 结果更新账号健康度，避免中间 WS 错误把同账号钉选提前冷却掉。 + failure.account_failure = None; + let fallback_now = context.policy == WebSocketFailurePolicy::ImmediateFallback + || context.retry_count >= context.max_retries; + if !fallback_now { + let Some(retry_index) = context.retry_count.checked_add(1).and_then(NonZeroU32::new) else { + return; + }; + let delay = failure + .error + .retry_after() + .unwrap_or_else(|| websocket_retry_backoff(retry_index)); + failure + .error + .set_pre_delivery_transport_retry(retry_index, delay); + tracing::warn!( + request_id = context.request_id, + attempt_index = context.attempt_index, + account_id = context.account_id, + websocket_failure_kind = failure.error.kind().as_str(), + websocket_failure_code = failure + .error + .upstream_code() + .map_or("", OpaqueUpstreamValue::as_str), + upstream_status_code = failure.error.upstream_status().unwrap_or_default(), + upstream_status_code_present = failure.error.upstream_status().is_some(), + websocket_retry_count = retry_index.get(), + websocket_max_retries = context.max_retries, + websocket_retry_delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX), + transport_requirement = context.requirement.as_str(), + "OpenAI upstream WebSocket failed; retrying the same account" + ); + return; + } + + failure.error.set_pre_delivery_transport_fallback(); + let sticky_state_transition = context + .session_affinity_key + .map(|key| context.session_transport_fallbacks.disable_websocket(key)); + tracing::warn!( + request_id = context.request_id, + attempt_index = context.attempt_index, + account_id = context.account_id, + websocket_failure_kind = failure.error.kind().as_str(), + websocket_failure_code = failure + .error + .upstream_code() + .map_or("", OpaqueUpstreamValue::as_str), + upstream_status_code = failure.error.upstream_status().unwrap_or_default(), + upstream_status_code_present = failure.error.upstream_status().is_some(), + websocket_retry_count = context.retry_count, + websocket_max_retries = context.max_retries, + websocket_fallback_reason = match context.policy { + WebSocketFailurePolicy::Budgeted => "retry_budget_exhausted", + WebSocketFailurePolicy::ImmediateFallback => "upgrade_required", + }, + session_affinity_present = context.session_affinity_key.is_some(), + session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""), + sticky_http_enabled = sticky_state_transition.is_some(), + sticky_state_transition = sticky_state_transition.unwrap_or(false), + "OpenAI upstream WebSocket disabled for the current session" + ); +} + +fn websocket_retry_backoff(retry_index: NonZeroU32) -> Duration { + const INITIAL_DELAY_MS: u64 = 200; + + let shift = retry_index.get().saturating_sub(1).min(63); + let base_ms = INITIAL_DELAY_MS.saturating_mul(1_u64.checked_shl(shift).unwrap_or(u64::MAX)); + let mut random = [0_u8; 2]; + let jitter_per_mille = if getrandom::fill(&mut random).is_ok() { + 900_u64 + u64::from(u16::from_le_bytes(random) % 200) + } else { + 1_000 + }; + Duration::from_millis(base_ms.saturating_mul(jitter_per_mille) / 1_000) +} + fn continuation_replay_required_error() -> ProviderError { ProviderError::new( ProviderErrorKind::InvalidRequest, @@ -2495,8 +2743,10 @@ fn continuation_replay_required_error() -> ProviderError { fn map_stream_error(error: CodexClientError) -> MappedProviderFailure { let allows_pre_delivery_retry = stream_transport_allows_pre_delivery_retry(&error); + let websocket_failure = error.transport() == Some(CodexBackendTransport::WebSocket); let mut failure = map_client_error(error, UpstreamSendState::Sent, false); - if allows_pre_delivery_retry { + failure.websocket_transport_retryable = allows_pre_delivery_retry && websocket_failure; + if allows_pre_delivery_retry && !websocket_failure { failure.error = failure.error.with_pre_delivery_retry(); } failure @@ -2507,15 +2757,11 @@ fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool CodexClientError::Http(_) | CodexClientError::HttpJson(_) | CodexClientError::StreamIdleTimeout { .. } => true, - CodexClientError::WebSocket(error) => matches!( + CodexClientError::WebSocket(error) => !matches!( error, - CodexWebSocketExchangeError::Transport(_) - | CodexWebSocketExchangeError::PostSendAmbiguous { .. } - | CodexWebSocketExchangeError::SendTimeout { .. } - | CodexWebSocketExchangeError::ClosedBeforeTerminal(_) - | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } - | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } - | CodexWebSocketExchangeError::InitialEventTimeout { .. } + CodexWebSocketExchangeError::InvalidRequest(_) + | CodexWebSocketExchangeError::Upstream(_) + | CodexWebSocketExchangeError::ContinuationUnavailable { .. } ), _ => false, } @@ -2617,6 +2863,11 @@ fn map_client_error( "websocket_close_{close_code}" ))); } + if matches!(error, CodexWebSocketExchangeError::ConnectionLimitReached) { + failure.error = failure.error.with_upstream_code(OpaqueUpstreamValue::new( + WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned(), + )); + } if let Some(client_visible_error) = client_visible_error { failure.error = failure .error @@ -2637,6 +2888,13 @@ fn map_client_error( fn websocket_client_visible_error( error: &CodexWebSocketExchangeError, ) -> Option { + if matches!(error, CodexWebSocketExchangeError::ConnectionLimitReached) { + return Some(ClientVisibleUpstreamError::new( + "websocket connection limit reached", + Some(WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned()), + Some("websocket_error".to_owned()), + )); + } let close = error.close_before_terminal()?; let message = close .reason() @@ -2711,6 +2969,7 @@ fn map_upstream_failure( } MappedProviderFailure { error, + websocket_transport_retryable: false, account_failure: account_failure( category, failure.retry_after_seconds, @@ -2822,6 +3081,7 @@ const fn websocket_send_state(error: &CodexWebSocketExchangeError) -> UpstreamSe | CodexWebSocketExchangeError::SharedConnectFailed | CodexWebSocketExchangeError::ContinuationUnavailable { .. } => UpstreamSendState::NotSent, CodexWebSocketExchangeError::Upstream(_) + | CodexWebSocketExchangeError::ConnectionLimitReached | CodexWebSocketExchangeError::InvalidSse(_) | CodexWebSocketExchangeError::UnexpectedBinaryEvent => UpstreamSendState::Sent, CodexWebSocketExchangeError::Transport(_) @@ -2851,6 +3111,7 @@ const fn websocket_error_kind(error: &CodexWebSocketExchangeError) -> ProviderEr ProviderErrorKind::Unavailable } CodexWebSocketExchangeError::Upstream(_) => ProviderErrorKind::Unavailable, + CodexWebSocketExchangeError::ConnectionLimitReached => ProviderErrorKind::RateLimited, CodexWebSocketExchangeError::Transport(_) | CodexWebSocketExchangeError::Connect(_) | CodexWebSocketExchangeError::PostSendAmbiguous { .. } diff --git a/backend/crates/providers/openai/src/session_transport.rs b/backend/crates/providers/openai/src/session_transport.rs new file mode 100644 index 000000000..f4ea2f34e --- /dev/null +++ b/backend/crates/providers/openai/src/session_transport.rs @@ -0,0 +1,106 @@ +//! OpenAI 根会话级上游传输降级状态。 + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::{Duration, Instant}; + +use gateway_core::provider_ports::ProviderSessionAffinityKey; + +use crate::credential::CODEX_ROOT_SESSION_TTL; + +/// Provider 是进程级单例；限制不可显式观测结束的客户端会话占用。 +const MAX_SESSION_HTTP_FALLBACKS: usize = 16_384; + +#[derive(Clone)] +pub(crate) struct CodexSessionTransportFallbacks { + state: Arc>, + ttl: Duration, + capacity: usize, +} + +#[derive(Default)] +struct SessionTransportFallbackState { + http_only: HashMap, +} + +impl Default for CodexSessionTransportFallbacks { + fn default() -> Self { + Self::with_limits(CODEX_ROOT_SESSION_TTL, MAX_SESSION_HTTP_FALLBACKS) + } +} + +impl CodexSessionTransportFallbacks { + fn with_limits(ttl: Duration, capacity: usize) -> Self { + debug_assert!(!ttl.is_zero()); + debug_assert!(capacity > 0); + Self { + state: Arc::new(Mutex::new(SessionTransportFallbackState::default())), + ttl, + capacity, + } + } + + /// 返回该根会话是否已粘性切到 HTTP，并为活跃会话续期。 + pub(crate) fn is_http_only(&self, key: &ProviderSessionAffinityKey) -> bool { + self.is_http_only_at(key, Instant::now()) + } + + /// 在当前 WebSocket 重试预算耗尽后，为根会话启用粘性 HTTP。 + /// + /// 返回 `true` 表示本次发生了禁用状态转换；已禁用会话只续期。 + pub(crate) fn disable_websocket(&self, key: &ProviderSessionAffinityKey) -> bool { + self.disable_websocket_at(key, Instant::now()) + } + + fn is_http_only_at(&self, key: &ProviderSessionAffinityKey, now: Instant) -> bool { + let mut state = self.lock(); + let Some(last_seen) = state.http_only.get_mut(key) else { + return false; + }; + if now.saturating_duration_since(*last_seen) >= self.ttl { + state.http_only.remove(key); + return false; + } + *last_seen = now; + true + } + + fn disable_websocket_at(&self, key: &ProviderSessionAffinityKey, now: Instant) -> bool { + let mut state = self.lock(); + let already_disabled = state + .http_only + .get(key) + .is_some_and(|last_seen| now.saturating_duration_since(*last_seen) < self.ttl); + if !already_disabled && !state.http_only.contains_key(key) { + self.make_room(&mut state, now); + } + state.http_only.insert(key.clone(), now); + !already_disabled + } + + fn make_room(&self, state: &mut SessionTransportFallbackState, now: Instant) { + if state.http_only.len() < self.capacity { + return; + } + state + .http_only + .retain(|_, last_seen| now.saturating_duration_since(*last_seen) < self.ttl); + if state.http_only.len() < self.capacity { + return; + } + let oldest = state + .http_only + .iter() + .min_by_key(|(_, last_seen)| **last_seen) + .map(|(key, _)| key.clone()); + if let Some(oldest) = oldest { + state.http_only.remove(&oldest); + } + } + + fn lock(&self) -> MutexGuard<'_, SessionTransportFallbackState> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} diff --git a/backend/crates/providers/openai/src/transport/client.rs b/backend/crates/providers/openai/src/transport/client.rs index 01e9db875..5badb40e4 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -599,6 +599,7 @@ pub(crate) struct PreparedResponseTransport { pub(super) requirement: TransportRequirement, pub(super) route: PreparedResponseRoute, pub(super) metrics: CodexTransportMetrics, + pub(super) defer_websocket_recovery: bool, } pub(super) enum PreparedResponseRoute { diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index dcb73a555..c69d509c1 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -182,6 +182,20 @@ impl CodexBackendClient { .await } + /// Provider 生产路径：保留所有 WebSocket 失败，交由统一重试预算处理。 + pub(crate) async fn create_response_stream_with_deferred_websocket_recovery( + &self, + request: &CodexResponsesRequest, + context: CodexRequestContext<'_>, + pool_account_id: Option<&str>, + ) -> CodexClientResult { + let prepared = self + .prepare_response_transport(request, context, pool_account_id, false, true) + .await?; + self.create_response_stream_with_prepared(request, context, prepared) + .await + } + /// 在发送 payload 前完成 transport 选择和可取消的 WebSocket opening。 #[doc(hidden)] pub(crate) async fn prepare_response_transport_with_pool_account( @@ -190,7 +204,7 @@ impl CodexBackendClient { context: CodexRequestContext<'_>, pool_account_id: Option<&str>, ) -> CodexClientResult { - self.prepare_response_transport(request, context, pool_account_id, false) + self.prepare_response_transport(request, context, pool_account_id, false, false) .await } @@ -202,12 +216,14 @@ impl CodexBackendClient { context: CodexRequestContext<'_>, pool_account_id: Option<&str>, force_fresh: bool, + defer_websocket_recovery: bool, ) -> CodexClientResult { let requirement = transport_requirement(request); if requirement == TransportRequirement::HttpRequired { return Ok(PreparedResponseTransport { requirement, route: PreparedResponseRoute::Http, + defer_websocket_recovery, metrics: CodexTransportMetrics { decision: Some(CodexTransportDecision::HttpRequired), ..CodexTransportMetrics::default() @@ -268,7 +284,8 @@ impl CodexBackendClient { let prepared = match prepared { Ok(prepared) => prepared, Err(error) - if requirement.allows_pre_send_http_fallback() + if (!defer_websocket_recovery || error.requires_immediate_pool_fallback()) + && requirement.allows_pre_send_http_fallback() && error.allows_pre_send_http_fallback() => { let decision = http_fallback_decision(&error); @@ -291,6 +308,7 @@ impl CodexBackendClient { transport_decision_wait_ms: Some(wait_ms), ..CodexTransportMetrics::default() }, + defer_websocket_recovery, }); } Err(error) => return Err(websocket_exchange_error_to_client_error(error)), @@ -328,6 +346,7 @@ impl CodexBackendClient { prepared, })), metrics, + defer_websocket_recovery, }) } @@ -342,6 +361,7 @@ impl CodexBackendClient { requirement, route, metrics, + defer_websocket_recovery, } = prepared; match route { PreparedResponseRoute::Http => self @@ -365,7 +385,8 @@ impl CodexBackendClient { { Ok(exchange) => exchange, Err(error) - if requirement.allows_pre_delivery_http_fallback() + if !defer_websocket_recovery + && requirement.allows_pre_delivery_http_fallback() && error.allows_pre_delivery_http_fallback() => { return self @@ -381,29 +402,20 @@ impl CodexBackendClient { } Err(error) => return Err(websocket_exchange_error_to_client_error(error)), }; - // 仅普通新链请求允许交付前降级/重试；延续类请求保持旧语义： - // 不进入交付边界，错误由 exchange 流原样上抛给下游。 if requirement.allows_pre_delivery_http_fallback() { let mut fresh_retry_used = false; loop { match await_websocket_delivery_boundary(&mut exchange).await { Ok(DeliveryBoundary::Ready) => break, - Ok(DeliveryBoundary::ConnectionLimitReached { - message, - raw_frame, - }) => { - let upstream_error_raw = String::from_utf8_lossy(&raw_frame); - tracing::warn!( - request_id = %context.request_id, - account_id = context.account_id.unwrap_or_default(), - websocket_connection_id = %exchange.websocket_connection_id, - upstream_transport = "websocket", - upstream_error_kind = "websocket_error_frame", - upstream_error_raw = %upstream_error_raw, - "OpenAI upstream returned a recoverable WebSocket error frame" - ); + Ok(DeliveryBoundary::ConnectionLimitReached) + if defer_websocket_recovery => + { + return Err(CodexClientError::WebSocket( + CodexWebSocketExchangeError::ConnectionLimitReached, + )); + } + Ok(DeliveryBoundary::ConnectionLimitReached) => { if fresh_retry_used { - // 新连接仍被限流：与官方重试预算耗尽一致，降级同账号 HTTP/SSE。 return self .http_fallback_before_delivery( request, @@ -411,20 +423,18 @@ impl CodexBackendClient { requirement, metrics, delivery_wait_started_at, - &message, + "websocket connection limit reached", ) - .await - ; + .await; } fresh_retry_used = true; tracing::warn!( request_id = %context.request_id, - %message, "WebSocket connection limit reached; retrying on a fresh connection" ); drop(exchange); match self - .prepare_response_transport(request, context, None, true) + .prepare_response_transport(request, context, None, true, false) .await { Ok(PreparedResponseTransport { @@ -493,19 +503,21 @@ impl CodexBackendClient { } } } + Err(error) if !defer_websocket_recovery + && error.allows_pre_delivery_http_fallback() => + { + return self + .fallback_to_http_before_websocket_delivery( + request, + context, + requirement, + metrics, + delivery_wait_started_at, + error, + ) + .await; + } Err(error) => { - if error.allows_pre_delivery_http_fallback() { - return self - .fallback_to_http_before_websocket_delivery( - request, - context, - requirement, - metrics, - delivery_wait_started_at, - error, - ) - .await; - } return Err(websocket_exchange_error_to_client_error( post_send_ambiguous(error), )); @@ -561,7 +573,6 @@ impl CodexBackendClient { delivery_wait_started_at: Instant, error: CodexWebSocketExchangeError, ) -> CodexClientResult { - log_raw_websocket_close(context, &error); self.http_fallback_before_delivery( request, context, @@ -573,7 +584,6 @@ impl CodexBackendClient { .await } - /// 交付前降级同账号 HTTP/SSE 的公共路径：统一标记 decision 与等待耗时。 async fn http_fallback_before_delivery( &self, request: &CodexResponsesRequest, @@ -829,14 +839,8 @@ impl CodexBackendClient { enum DeliveryBoundary { /// 已越过边界，可开始向下游投递。 Ready, - /// 首个可投递帧是上游连接寿命限制错误；该帧已放回流， - /// 调用方可选择丢弃并换新连接重试，或原样投递。 - ConnectionLimitReached { - /// 上游错误说明。 - message: String, - /// 从 WebSocket 文本事件构造的完整 SSE 错误帧。 - raw_frame: bytes::Bytes, - }, + /// 首个可投递帧是上游连接寿命限制错误。 + ConnectionLimitReached, } async fn await_websocket_delivery_boundary( @@ -847,19 +851,16 @@ async fn await_websocket_delivery_boundary( match exchange.body.next().await { Some(Ok(frame)) if is_websocket_lifecycle_prelude(&frame) => prelude.push(frame), Some(Ok(frame)) => { - let limit_message = websocket_connection_limit_message(&frame); - let raw_frame = limit_message.as_ref().map(|_| frame.clone()); + let connection_limit_reached = websocket_connection_limit_message(&frame).is_some(); prelude.push(frame); let remaining = std::mem::replace(&mut exchange.body, Box::pin(futures::stream::empty())); exchange.body = Box::pin(futures::stream::iter(prelude.into_iter().map(Ok)).chain(remaining)); - return Ok(match limit_message { - Some(message) => DeliveryBoundary::ConnectionLimitReached { - message, - raw_frame: raw_frame.unwrap_or_default(), - }, - None => DeliveryBoundary::Ready, + return Ok(if connection_limit_reached { + DeliveryBoundary::ConnectionLimitReached + } else { + DeliveryBoundary::Ready }); } Some(Err(error)) => return Err(error), @@ -868,34 +869,13 @@ async fn await_websocket_delivery_boundary( exchange.websocket_connection_id, None, None, + None, )); } } } } -fn log_raw_websocket_close(context: CodexRequestContext<'_>, error: &CodexWebSocketExchangeError) { - let Some(close) = error.close_before_terminal() else { - return; - }; - let websocket_connection_id = close - .connection_id() - .map(|connection_id| connection_id.to_string()) - .unwrap_or_default(); - let upstream_error_raw = close.reason().unwrap_or_default(); - tracing::warn!( - request_id = %context.request_id, - account_id = context.account_id.unwrap_or_default(), - websocket_connection_id, - upstream_transport = "websocket", - upstream_error_kind = "websocket_close", - upstream_close_code = ?close.code(), - upstream_error_raw, - upstream_error_raw_present = close.reason().is_some(), - "OpenAI upstream WebSocket closed before delivery" - ); -} - fn is_websocket_lifecycle_prelude(frame: &[u8]) -> bool { frame.starts_with(b"event: response.created\n") || frame.starts_with(b"event: response.in_progress\n") diff --git a/backend/crates/providers/openai/src/transport/websocket/error.rs b/backend/crates/providers/openai/src/transport/websocket/error.rs index 9f1474da5..c96ebb819 100644 --- a/backend/crates/providers/openai/src/transport/websocket/error.rs +++ b/backend/crates/providers/openai/src/transport/websocket/error.rs @@ -66,6 +66,9 @@ pub enum CodexWebSocketExchangeError { /// 上游 WebSocket 错误帧。 #[error("{0}")] Upstream(Box), + /// 上游要求结束当前 WebSocket 连接并在新连接上重试。 + #[error("websocket connection limit reached")] + ConnectionLimitReached, /// 请求依赖的连接本地 previous response 无法在当前连接满足。 #[error("websocket continuation unavailable: {reason}")] ContinuationUnavailable { @@ -108,6 +111,7 @@ pub struct CodexWebSocketCloseError { connection_id: Option, code: Option, reason: Option, + last_event_type: Option, } impl CodexWebSocketCloseError { @@ -116,6 +120,7 @@ impl CodexWebSocketCloseError { connection_id: None, code, reason, + last_event_type: None, } } @@ -124,6 +129,11 @@ impl CodexWebSocketCloseError { self } + pub(crate) fn with_last_event_type(mut self, event_type: Option) -> Self { + self.last_event_type = event_type; + self + } + /// 返回承载该 close frame 的 WebSocket 连接标识。 #[must_use] pub const fn connection_id(&self) -> Option { @@ -141,6 +151,12 @@ impl CodexWebSocketCloseError { pub fn reason(&self) -> Option<&str> { self.reason.as_deref() } + + /// 返回关闭前最后一条经过白名单校验的事件类型；不包含事件 payload。 + #[must_use] + pub fn last_event_type(&self) -> Option<&str> { + self.last_event_type.as_deref() + } } impl fmt::Debug for CodexWebSocketCloseError { @@ -150,6 +166,7 @@ impl fmt::Debug for CodexWebSocketCloseError { .field("connection_id", &self.connection_id) .field("code", &self.code) .field("reason", &self.reason.as_ref().map(|_| "")) + .field("last_event_type", &self.last_event_type) .finish() } } @@ -214,9 +231,12 @@ impl CodexWebSocketExchangeError { connection_id: Uuid, code: Option, reason: Option, + last_event_type: Option, ) -> Self { Self::ClosedBeforeTerminal( - CodexWebSocketCloseError::new(code, reason).with_connection_id(connection_id), + CodexWebSocketCloseError::new(code, reason) + .with_connection_id(connection_id) + .with_last_event_type(last_event_type), ) } @@ -269,6 +289,11 @@ impl CodexWebSocketExchangeError { ) } + /// 精确连接状态不可用属于本地池路由事实，不消耗上游 WS 重试预算。 + pub(in crate::transport) fn requires_immediate_pool_fallback(&self) -> bool { + matches!(self, Self::ContinuationUnavailable { .. }) + } + /// 首个业务事件交付前，普通 WebSocket 失败可切到同账号 HTTP。 pub(in crate::transport) fn allows_pre_delivery_http_fallback(&self) -> bool { match self { diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs index 0cfd7c4ed..f40589049 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/reducer.rs @@ -29,20 +29,32 @@ pub(super) enum ExchangeAction { Ignore, } +pub(super) struct ReducedWebSocketEvent { + pub(super) action: ExchangeAction, + pub(super) diagnostic_event_type: Option, +} + pub(super) fn reduce_websocket_event( raw: &str, metadata: &mut CodexWebSocketConnectionMetadata, continuation: &mut WebSocketContinuationState, -) -> Result { +) -> Result { // 每帧只解析一次 JSON，后续提取全部复用同一 Value； // 不可解析的帧不承载可路由的事件类型，忽略。 let Ok(value) = serde_json::from_str::(raw) else { - return Ok(ExchangeAction::Ignore); + return Ok(ReducedWebSocketEvent { + action: ExchangeAction::Ignore, + diagnostic_event_type: None, + }); }; + let diagnostic_event_type = diagnostic_event_type(websocket_event_type(&value)); if let Some(parsed) = events::parse_rate_limits_event(&value) { let headers = events::rate_limits_to_header_pairs(&parsed); metadata.rate_limit_headers.extend(headers); - return Ok(ExchangeAction::RateLimits(parsed)); + return Ok(ReducedWebSocketEvent { + action: ExchangeAction::RateLimits(parsed), + diagnostic_event_type, + }); } response_meta::merge_response_metadata( @@ -51,7 +63,10 @@ pub(super) fn reduce_websocket_event( ); if let Some(turn_state) = websocket_metadata_turn_state(&value) { metadata.turn_state = Some(turn_state.clone()); - return Ok(ExchangeAction::TurnState(turn_state)); + return Ok(ReducedWebSocketEvent { + action: ExchangeAction::TurnState(turn_state), + diagnostic_event_type, + }); } let event = websocket_event_type(&value); @@ -67,8 +82,24 @@ pub(super) fn reduce_websocket_event( Some("response.failed" | "error") => Some(WebSocketTerminalKind::Failed), _ => None, }; - Ok(match websocket_event_frame(&value, raw) { + let action = match websocket_event_frame(&value, raw) { Some(frame) => ExchangeAction::Forward { frame, terminal }, None => ExchangeAction::Ignore, + }; + Ok(ReducedWebSocketEvent { + action, + diagnostic_event_type, }) } + +fn diagnostic_event_type(event_type: Option<&str>) -> Option { + const MAX_EVENT_TYPE_BYTES: usize = 128; + + let event_type = event_type?; + (!event_type.is_empty() + && event_type.len() <= MAX_EVENT_TYPE_BYTES + && event_type.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || matches!(byte, b'.' | b'_' | b'-') + })) + .then(|| event_type.to_owned()) +} diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs index 1114729b1..35666c92e 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -143,6 +143,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { .map(|pool_return| std::mem::take(&mut pool_return.continuation)) .unwrap_or_default(); let mut saw_upstream_activity = false; + let mut last_event_type = None; loop { let message = tokio::select! { biased; @@ -219,6 +220,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { websocket_connection_id, frame.as_ref().map(|frame| u16::from(frame.code)), frame.map(|frame| frame.reason.to_string()), + last_event_type.clone(), ); let error = if reused_connection { reused_stream_receive_error(error) @@ -230,18 +232,8 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { } _ => continue, }; - let (frame, terminal) = match reduce_websocket_event(&raw, &mut metadata, &mut continuation) - { - Ok(ExchangeAction::RateLimits(rate_limits)) => { - rate_limit_updates.lock().await.push(rate_limits); - continue; - } - Ok(ExchangeAction::TurnState(turn_state)) => { - *turn_state_update.lock().await = Some(turn_state); - continue; - } - Ok(ExchangeAction::Forward { frame, terminal }) => (frame, terminal), - Ok(ExchangeAction::Ignore) => continue, + let reduced = match reduce_websocket_event(&raw, &mut metadata, &mut continuation) { + Ok(reduced) => reduced, Err(error) => { discard_stream_websocket( websocket, @@ -253,6 +245,21 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { return; } }; + if let Some(event_type) = reduced.diagnostic_event_type { + last_event_type = Some(event_type); + } + let (frame, terminal) = match reduced.action { + ExchangeAction::RateLimits(rate_limits) => { + rate_limit_updates.lock().await.push(rate_limits); + continue; + } + ExchangeAction::TurnState(turn_state) => { + *turn_state_update.lock().await = Some(turn_state); + continue; + } + ExchangeAction::Forward { frame, terminal } => (frame, terminal), + ExchangeAction::Ignore => continue, + }; if tx.send(Ok(Bytes::from(frame))).await.is_err() { discard_stream_websocket( websocket, @@ -300,16 +307,17 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { StreamWebSocketDiscardReason::UpstreamClosed, ) .await; - let error = upstream_close.map_or_else( - || { - CodexWebSocketExchangeError::closed_before_terminal_on( - websocket_connection_id, - None, - None, - ) - }, - CodexWebSocketExchangeError::ClosedBeforeTerminal, - ); + let error = match upstream_close { + Some(close) => CodexWebSocketExchangeError::ClosedBeforeTerminal( + close.with_last_event_type(last_event_type), + ), + None => CodexWebSocketExchangeError::closed_before_terminal_on( + websocket_connection_id, + None, + None, + last_event_type, + ), + }; let error = if reused_connection { reused_stream_receive_error(error) } else { diff --git a/backend/crates/providers/openai/tests/config.rs b/backend/crates/providers/openai/tests/config.rs index 873d87953..7b7a85fa0 100644 --- a/backend/crates/providers/openai/tests/config.rs +++ b/backend/crates/providers/openai/tests/config.rs @@ -1,7 +1,9 @@ use std::path::Path; use chrono::{TimeZone as _, Utc}; -use provider_openai::config::{CodexWireProfileConfig, OpenAiConfig}; +use provider_openai::config::{ + CodexWireProfileConfig, DEFAULT_STREAM_MAX_RETRIES, MAX_STREAM_MAX_RETRIES, OpenAiConfig, +}; #[test] fn openai_config_builds_the_audited_wire_profile() { @@ -95,6 +97,7 @@ fn openai_config_restricts_upstream_base_url_to_https_or_loopback_http() { #[test] fn openai_config_defaults_to_the_provider_owned_operating_values() { let config = OpenAiConfig::default(); + assert_eq!(DEFAULT_STREAM_MAX_RETRIES, 3); assert_eq!( ( @@ -107,6 +110,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() { config.auth.refresh_enabled, config.auth.oauth_client_id.as_str(), config.auth.oauth_token_endpoint.as_str(), + config.stream_max_retries(), ), ( "https://chatgpt.com/backend-api", @@ -118,6 +122,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() { true, "app_EMoamEEZ73f0CkXaXp7hrann", "https://auth.openai.com/oauth/token", + u32::try_from(DEFAULT_STREAM_MAX_RETRIES).expect("default retry budget fits u32"), ) ); assert_eq!( @@ -126,6 +131,17 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() { ); } +#[test] +fn openai_stream_retry_budget_uses_the_official_hard_cap() { + let mut config = OpenAiConfig::default(); + config.stream_max_retries = MAX_STREAM_MAX_RETRIES + 1; + + assert_eq!( + config.stream_max_retries(), + u32::try_from(MAX_STREAM_MAX_RETRIES).expect("retry cap fits u32") + ); +} + fn valid_config() -> OpenAiConfig { let mut config = OpenAiConfig::default(); config.wire_profile = CodexWireProfileConfig { diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 56653c7f8..4c2fde63e 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -34,6 +34,7 @@ use gateway_core::routing::{ ProviderModel, PublicModelId, RoutingContext, RuntimeAccount, RuntimeAccountDirectory, RuntimeSnapshot, UpstreamModelId, }; +use provider_openai::config::DEFAULT_STREAM_MAX_RETRIES; use provider_openai::credential::{ CodexCookiePolicy, CodexCredentialCatalogService, CodexCredentialQuotaService, CodexCredentialSelector, ImportCodexOAuthCredential, @@ -148,7 +149,26 @@ fn provider_with_affinity( } fn provider_with_base_url(store: &Arc, base_url: String) -> CodexProvider { - provider_with_affinity_and_base_url(store, Arc::new(MemorySessionAffinity::default()), base_url) + provider_with_base_url_and_retry_budget( + store, + base_url, + u32::try_from(DEFAULT_STREAM_MAX_RETRIES).expect("default retry budget fits u32"), + ) +} + +fn provider_with_base_url_and_retry_budget( + store: &Arc, + base_url: String, + stream_max_retries: u32, +) -> CodexProvider { + provider_and_quota_with_affinity_and_base_url_and_leases( + store, + Arc::new(MemorySessionAffinity::default()), + base_url, + Arc::new(TestLeaseCoordinator::default()), + stream_max_retries, + ) + .0 } fn provider_with_leases( @@ -187,6 +207,7 @@ fn provider_with_affinity_and_base_url_and_leases( session_affinity, base_url, leases, + u32::try_from(DEFAULT_STREAM_MAX_RETRIES).expect("default retry budget fits u32"), ) .0 } @@ -196,6 +217,7 @@ fn provider_and_quota_with_affinity_and_base_url_and_leases( session_affinity: Arc, base_url: String, leases: Arc, + stream_max_retries: u32, ) -> (CodexProvider, Arc) { let profile = wire_profile(); let http = reqwest::Client::builder() @@ -239,6 +261,7 @@ fn provider_and_quota_with_affinity_and_base_url_and_leases( profile, base_url, websocket_pool, + stream_max_retries, ) .expect("official OpenAI provider"); (provider, quota) @@ -769,21 +792,25 @@ async fn truncated_chunked_sse_server() -> (String, tokio::task::JoinHandle<()>) (base_url, server) } -async fn read_http_request(stream: &mut TcpStream) { +async fn capture_http_request(stream: &mut TcpStream) -> Vec { let mut request = Vec::new(); let mut buffer = [0_u8; 1024]; loop { let read = stream.read(&mut buffer).await.expect("read HTTP request"); if read == 0 { - return; + return request; } request.extend_from_slice(&buffer[..read]); if request.windows(4).any(|window| window == b"\r\n\r\n") { - return; + return request; } } } +async fn read_http_request(stream: &mut TcpStream) { + drop(capture_http_request(stream).await); +} + async fn write_http_chunk(stream: &mut TcpStream, body: &str) { stream .write_all(format!("{:X}\r\n", body.len()).as_bytes()) @@ -1131,7 +1158,7 @@ async fn selection_infrastructure_errors_have_a_distinct_classification() { } #[tokio::test] -async fn websocket_close_details_are_only_exposed_through_the_client_error() { +async fn websocket_close_after_delivery_preserves_details_without_enabling_sticky_http() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_websocket_close").await; let listener = TcpListener::bind("127.0.0.1:0") @@ -1170,11 +1197,49 @@ async fn websocket_close_details_are_only_exposed_through_the_client_error() { })) .await .expect("close WebSocket"); + + let (stream, _) = listener + .accept() + .await + .expect("accept next-turn WebSocket connection"); + let mut websocket = accept_codex_test_websocket(stream).await; + let _request = websocket + .next() + .await + .expect("next-turn WebSocket request") + .expect("valid next-turn WebSocket request"); + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_after_committed_close", + "model": "gpt-5.4", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2 + } + } + }) + .to_string() + .into(), + )) + .await + .expect("complete next-turn WebSocket request"); }); - let mut stream = provider_with_base_url(&store, base_url) + let provider = provider_with_base_url(&store, base_url); + let first_operation = Operation::Generate(generate_with_session_context( + "committed-websocket-session", + Some("thread-first"), + None, + )); + let mut stream = provider .execute( - planned_request("openai", generate_operation()), + planned_request("openai", first_operation), context("req_websocket_close", CancellationToken::new()), ) .await @@ -1186,13 +1251,16 @@ async fn websocket_close_details_are_only_exposed_through_the_client_error() { None => panic!("WebSocket close must surface a provider error"), } }; - server.await.expect("WebSocket server"); + drop(stream); assert!(!format!("{error:?}").contains("message too big")); assert!(!error.to_string().contains("message too big")); assert_eq!(error.kind(), ProviderErrorKind::Transport); assert_eq!(error.send_state(), UpstreamSendState::Ambiguous); - assert!(error.allows_pre_delivery_retry()); + assert!( + !error.allows_pre_delivery_retry(), + "a close after the delivery boundary must not request a hidden replay" + ); let detail = error .client_visible_upstream_error() .expect("WebSocket close detail"); @@ -1203,10 +1271,28 @@ async fn websocket_close_details_are_only_exposed_through_the_client_error() { error.upstream_code().map(|code| code.as_str()), Some("websocket_close_1009") ); + + let second_operation = Operation::Generate(generate_with_session_context( + "committed-websocket-session", + Some("thread-second"), + None, + )); + let mut next_stream = provider + .execute( + planned_request("openai", second_operation), + context("req_after_committed_close", CancellationToken::new()), + ) + .await + .expect("post-delivery close must not enable sticky HTTP fallback"); + assert_eq!(next_stream.metadata().transport().as_str(), "websocket"); + while let Some(event) = next_stream.next().await { + event.expect("next turn WebSocket response"); + } + server.await.expect("WebSocket server"); } #[tokio::test] -async fn websocket_normal_close_before_semantic_output_requests_same_account_http_fallback() { +async fn websocket_retry_budget_exhaustion_is_sticky_only_for_the_current_session() { let store = Arc::new(MemoryAccountStore::default()); create_account(&store, "acct_websocket_close").await; let listener = TcpListener::bind("127.0.0.1:0") @@ -1217,79 +1303,350 @@ async fn websocket_normal_close_before_semantic_output_requests_same_account_htt listener.local_addr().expect("listener address") ); let server = tokio::spawn(async move { + for retry_count in 0..=2 { + let (stream, _) = listener + .accept() + .await + .expect("accept WebSocket retry connection"); + let mut websocket = accept_codex_test_websocket(stream).await; + let _request = websocket + .next() + .await + .expect("WebSocket request") + .expect("valid WebSocket request"); + for event in [ + json!({ + "type": "response.created", + "response": {"id": format!("resp_structural_{retry_count}"), "model": "gpt-5.4"} + }), + json!({ + "type": "response.in_progress", + "response": {"id": format!("resp_structural_{retry_count}"), "model": "gpt-5.4"} + }), + ] { + websocket + .send(Message::Text(event.to_string().into())) + .await + .expect("send structural WebSocket event"); + } + if retry_count == 1 { + websocket + .send(Message::Text( + json!({ + "type": "error", + "status": 400, + "error": { + "type": "invalid_request_error", + "code": "websocket_connection_limit_reached", + "message": "create a new websocket connection" + } + }) + .to_string() + .into(), + )) + .await + .expect("send connection-limit error"); + continue; + } + websocket + .send(Message::Text( + json!({ + "type": "response.output_item.added", + "output_index": 0, + "item": { + "id": format!("msg_structural_{retry_count}"), + "type": "message", + "role": "assistant", + "status": "in_progress", + "content": [] + } + }) + .to_string() + .into(), + )) + .await + .expect("send structural output event"); + let close_code = match retry_count { + 0 => CloseCode::Normal, + _ => CloseCode::Error, + }; + websocket + .close(Some(CloseFrame { + code: close_code, + reason: "".into(), + })) + .await + .expect("close WebSocket retry attempt"); + } + + let (mut http, _) = listener.accept().await.expect("accept sticky HTTP request"); + let request = capture_http_request(&mut http).await; + let request_head = String::from_utf8_lossy(&request).to_ascii_lowercase(); + assert!(!request_head.contains("upgrade: websocket")); + http.write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{CAPTURE_COMPLETED_SSE}", + CAPTURE_COMPLETED_SSE.len() + ) + .as_bytes(), + ) + .await + .expect("write sticky HTTP response"); + let (stream, _) = listener .accept() .await - .expect("accept WebSocket connection"); + .expect("accept other-session WebSocket connection"); let mut websocket = accept_codex_test_websocket(stream).await; let _request = websocket .next() .await - .expect("WebSocket request") - .expect("valid WebSocket request"); - for event in [ - json!({ - "type": "response.created", - "response": {"id": "resp_structural", "model": "gpt-5.4"} - }), - json!({ - "type": "response.in_progress", - "response": {"id": "resp_structural", "model": "gpt-5.4"} - }), - json!({ - "type": "response.output_item.added", - "output_index": 0, - "item": { - "id": "msg_structural", - "type": "message", - "role": "assistant", - "status": "in_progress", - "content": [] - } - }), - ] { - websocket - .send(Message::Text(event.to_string().into())) - .await - .expect("send structural WebSocket event"); - } + .expect("other-session WebSocket request") + .expect("valid other-session WebSocket request"); websocket - .close(Some(CloseFrame { - code: CloseCode::Normal, - reason: "".into(), - })) + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_other_session", + "model": "gpt-5.4", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2 + } + } + }) + .to_string() + .into(), + )) .await - .expect("close WebSocket"); + .expect("complete other-session WebSocket request"); }); - let mut stream = provider_with_base_url(&store, base_url) + let provider = provider_with_base_url_and_retry_budget(&store, base_url, 2); + let mut saw_websocket_observation = false; + for retry_count in 0..=2 { + let operation = Operation::Generate(generate_with_session_context( + "sticky-websocket-session", + Some("thread-first"), + None, + )); + let attempt_context = match NonZeroU32::new(retry_count) { + Some(retry_index) => context( + &format!("req_websocket_normal_close_{retry_count}"), + CancellationToken::new(), + ) + .with_transport(AttemptTransport::Retry(retry_index)), + None => context("req_websocket_normal_close_0", CancellationToken::new()), + }; + let mut stream = provider + .execute(planned_request("openai", operation), attempt_context) + .await + .expect("prepare WebSocket provider stream"); + let error = loop { + match stream.next().await { + Some(Ok(event)) => { + if let Some(observation) = event.response_observation() { + saw_websocket_observation = true; + assert_eq!(observation.status_code(), None); + if let Some(provider_metadata) = observation.provider_metadata() { + let metadata: Value = serde_json::from_str(provider_metadata.as_json()) + .expect("OpenAI provider metadata JSON"); + assert_eq!( + metadata.get("upstreamStatus").and_then(Value::as_u64), + Some(101), + "successful upgrade remains provider-owned diagnostics" + ); + } + } + assert!( + !event.has_client_event(), + "structural events must remain behind the replay boundary" + ); + } + Some(Err(error)) => break error, + None => panic!("WebSocket close must surface a provider error"), + } + }; + drop(stream); + + if retry_count < 2 { + let Some(PreDeliveryRetry::SameAccountTransportRetry { retry_index, delay }) = + error.pre_delivery_retry() + else { + panic!("expected a same-account WebSocket retry: {error:?}"); + }; + assert_eq!(retry_index.get(), retry_count + 1); + assert!(!delay.is_zero()); + } else { + assert_eq!( + error.pre_delivery_retry(), + Some(PreDeliveryRetry::SameAccountTransportFallback), + "unexpected provider error at retry {retry_count}: {error:?}" + ); + } + let expected_close_code = match retry_count { + 0 => "websocket_close_1000", + 1 => "websocket_connection_limit_reached", + _ => "websocket_close_1011", + }; + assert_eq!( + error.upstream_code().map(|code| code.as_str()), + Some(expected_close_code) + ); + } + + assert!(saw_websocket_observation); + + let second_operation = Operation::Generate(generate_with_session_context( + "sticky-websocket-session", + Some("thread-second"), + None, + )); + let mut sticky_stream = provider .execute( - planned_request("openai", generate_operation()), - context("req_websocket_normal_close", CancellationToken::new()), + planned_request("openai", second_operation), + context("req_websocket_sticky_http", CancellationToken::new()), ) .await - .expect("prepare WebSocket provider stream"); - let error = loop { - match stream.next().await { - Some(Ok(event)) => assert!( - !event.has_client_event(), - "structural events must remain behind the replay boundary" + .expect("same session should prepare an HTTP stream"); + assert_eq!(sticky_stream.metadata().transport().as_str(), "http_sse"); + while let Some(event) = sticky_stream.next().await { + event.expect("sticky HTTP response"); + } + drop(sticky_stream); + + let other_operation = Operation::Generate(generate_with_session_context( + "other-websocket-session", + Some("thread-first"), + None, + )); + let mut other_stream = provider + .execute( + planned_request("openai", other_operation), + context("req_other_session_websocket", CancellationToken::new()), + ) + .await + .expect("another session should still prepare a WebSocket stream"); + assert_eq!(other_stream.metadata().transport().as_str(), "websocket"); + while let Some(event) = other_stream.next().await { + event.expect("other-session WebSocket response"); + } + drop(other_stream); + + let warmup_operation = Operation::Generate(GenerateRequest::from_protocol_payload( + ProtocolPayload::json_object( + "openai", + Map::from_iter([ + ("model".to_owned(), json!("gpt-5.4")), + ("input".to_owned(), json!("warm up")), + ("session_id".to_owned(), json!("sticky-websocket-session")), + ("generate".to_owned(), json!(false)), + ("store".to_owned(), json!(false)), + ]), + ) + .expect("warmup payload"), + )); + let warmup_stream = provider + .execute( + planned_request("openai", warmup_operation), + context( + "req_sticky_session_required_warmup", + CancellationToken::new(), ), + ) + .await + .expect("required warmup must bypass sticky HTTP fallback"); + assert_eq!(warmup_stream.metadata().transport().as_str(), "websocket"); + drop(warmup_stream); + server.await.expect("WebSocket and HTTP server"); +} + +#[tokio::test] +async fn websocket_upgrade_required_immediately_enables_session_http_fallback() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_websocket_close").await; + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind upstream listener"); + let base_url = format!( + "http://{}", + listener.local_addr().expect("listener address") + ); + let server = tokio::spawn(async move { + let (mut opening, _) = listener.accept().await.expect("accept WS opening"); + let request = capture_http_request(&mut opening).await; + assert!(String::from_utf8_lossy(&request).starts_with("GET /codex/responses")); + let body = r#"{"error":{"code":"upgrade_required","message":"use HTTP"}}"#; + opening + .write_all( + format!( + "HTTP/1.1 426 Upgrade Required\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .await + .expect("write 426 response"); + + let (mut http, _) = listener.accept().await.expect("accept sticky HTTP request"); + let request = capture_http_request(&mut http).await; + assert!(String::from_utf8_lossy(&request).starts_with("POST /codex/responses")); + http.write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{CAPTURE_COMPLETED_SSE}", + CAPTURE_COMPLETED_SSE.len() + ) + .as_bytes(), + ) + .await + .expect("write HTTP fallback response"); + }); + + let provider = provider_with_base_url(&store, base_url); + let operation = || { + Operation::Generate(generate_with_session_context( + "sticky-websocket-session", + Some("thread-first"), + None, + )) + }; + let mut first = provider + .execute( + planned_request("openai", operation()), + context("req_upgrade_required", CancellationToken::new()), + ) + .await + .expect("prepare WebSocket stream"); + let error = loop { + match first.next().await { + Some(Ok(_)) => {} Some(Err(error)) => break error, - None => panic!("WebSocket close must surface a provider error"), + None => panic!("426 opening must surface a fallback signal"), } }; - server.await.expect("WebSocket server"); - + assert_eq!(error.upstream_status(), Some(426)); assert_eq!( error.pre_delivery_retry(), - Some(PreDeliveryRetry::SameAccountTransportFallback), - "unexpected provider error: {error:?}" - ); - assert_eq!( - error.upstream_code().map(|code| code.as_str()), - Some("websocket_close_1000") + Some(PreDeliveryRetry::SameAccountTransportFallback) ); + drop(first); + + let mut second = provider + .execute( + planned_request("openai", operation()), + context("req_upgrade_required_next_turn", CancellationToken::new()), + ) + .await + .expect("same session should select HTTP"); + assert_eq!(second.metadata().transport().as_str(), "http_sse"); + while let Some(event) = second.next().await { + event.expect("HTTP fallback response"); + } + server.await.expect("upstream server"); } #[tokio::test] @@ -2230,6 +2587,7 @@ async fn affinity_quota_switch_should_clear_old_turn_state_without_a_provider_st Arc::clone(&affinity), server.uri(), Arc::new(TestLeaseCoordinator::default()), + u32::try_from(DEFAULT_STREAM_MAX_RETRIES).expect("default retry budget fits u32"), ); let session_id = "stable-affinity-switch-session"; let first_payload = ProtocolPayload::json_object( @@ -3585,6 +3943,7 @@ async fn quota_limited_account_diagnostic_uses_upstream() { Arc::new(MemorySessionAffinity::default()), server.uri(), Arc::new(TestLeaseCoordinator::default()), + u32::try_from(DEFAULT_STREAM_MAX_RETRIES).expect("default retry budget fits u32"), ); quota .prepare_scheduling(std::slice::from_ref(&account)) diff --git a/backend/crates/providers/openai/tests/transport/websocket.rs b/backend/crates/providers/openai/tests/transport/websocket.rs index c48cdeb53..83a55885c 100644 --- a/backend/crates/providers/openai/tests/transport/websocket.rs +++ b/backend/crates/providers/openai/tests/transport/websocket.rs @@ -1507,6 +1507,17 @@ async fn codex_backend_client_stream_should_error_when_websocket_closes_before_t )) .await .unwrap(); + websocket + .send(Message::Text( + json!({ + "type": "secret@example.com", + "payload": "must not become a diagnostic event type" + }) + .to_string() + .into(), + )) + .await + .unwrap(); websocket.close(None).await.unwrap(); }); let mut request = codex_request("gpt-5.5", "be brief", Vec::new()); @@ -1533,14 +1544,27 @@ async fn codex_backend_client_stream_should_error_when_websocket_closes_before_t .expect("stream should yield partial frame") .expect("partial frame should be valid"); assert!(std::str::from_utf8(&first).unwrap().contains("partial")); - let error = response - .body - .next() - .await - .expect("stream should yield close-before-terminal error") - .expect_err("close before terminal should be an error"); + let error = loop { + match response.body.next().await { + Some(Ok(_)) => {} + Some(Err(error)) => break error, + None => panic!("stream should yield close-before-terminal error"), + } + }; server.await.unwrap(); + let CodexClientError::WebSocket(CodexWebSocketExchangeError::PostSendAmbiguous { + source: Some(source), + .. + }) = &error + else { + panic!("close-before-terminal should preserve its typed source: {error:?}"); + }; + let CodexWebSocketExchangeError::ClosedBeforeTerminal(close) = source.as_ref() else { + panic!("post-send ambiguity should wrap a close-before-terminal error: {source:?}"); + }; + assert_eq!(close.last_event_type(), Some("response.output_text.delta")); + std::assert_matches!( error, CodexClientError::WebSocket(CodexWebSocketExchangeError::PostSendAmbiguous {