From 4220f4e09523eab261da162298b4ff59bd6edd88 Mon Sep 17 00:00:00 2001 From: Zyy Date: Thu, 3 Sep 2026 23:33:26 +0800 Subject: [PATCH] fix(openai): recover websocket sessions after ambiguous failures Co-authored-by: Codex --- .../providers/openai/src/credential/mod.rs | 4 +- .../openai/src/provider/execution.rs | 52 +- .../providers/openai/src/provider/failure.rs | 55 +- .../providers/openai/src/provider/mod.rs | 66 ++- .../providers/openai/src/session_transport.rs | 459 +++++++++++++-- .../openai/src/transport/client_sse.rs | 50 +- .../providers/openai/src/transport/mod.rs | 1 + .../src/transport/protocol/responses.rs | 4 +- .../openai/tests/provider/contract.rs | 539 ++++++++++++++++-- .../tests/transport/protocol/responses.rs | 8 +- 10 files changed, 1086 insertions(+), 152 deletions(-) diff --git a/backend/crates/providers/openai/src/credential/mod.rs b/backend/crates/providers/openai/src/credential/mod.rs index fc5ce3243..3daefc374 100644 --- a/backend/crates/providers/openai/src/credential/mod.rs +++ b/backend/crates/providers/openai/src/credential/mod.rs @@ -16,8 +16,8 @@ pub mod token_client; mod types; pub(crate) use affinity::{ - CODEX_ROOT_SESSION_TTL, CodexSessionAffinity, derive_codex_cyber_policy_session_key, - derive_codex_session_affinity, derive_previous_response_id_hash, + CodexSessionAffinity, derive_codex_cyber_policy_session_key, derive_codex_session_affinity, + derive_previous_response_id_hash, }; pub(crate) use oauth::oauth_owner_ref; pub(crate) use types::parse_access_token_expiration; diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index 13dfbe4e2..d5a7a54ad 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -147,6 +147,7 @@ pub(super) struct ColdResponse { pub(super) request: CodexResponsesRequest, pub(super) upstream_model: UpstreamModelId, pub(super) transport_policy: CodexProviderTransport, + pub(super) websocket_connection_preference: WebSocketConnectionPreference, pub(super) context: AttemptContext, pub(super) selector: Arc, pub(super) quota: Arc, @@ -155,7 +156,8 @@ pub(super) struct ColdResponse { pub(super) output_started_at: Instant, pub(super) session_affinity_key: Option, pub(super) session_affinity_key_hash: Option, - pub(super) session_transport_fallbacks: CodexSessionTransportFallbacks, + pub(super) session_transport_recovery: CodexSessionTransportRecovery, + pub(super) session_recovery_probe: Option, pub(super) websocket_retry_count: u32, pub(super) stream_max_retries: u32, pub(super) session_capture: Option, @@ -292,6 +294,7 @@ pub(super) async fn create_response_attempt( account_id: &str, deadline: SystemTime, cancellation: &CancellationToken, + connection_preference: WebSocketConnectionPreference, ) -> Result { let Some(handshake_deadline) = remaining(deadline) else { return Err(CodexHandshakeAttemptError::Timeout); @@ -304,6 +307,7 @@ pub(super) async fn create_response_attempt( request, request_context, Some(account_id), + connection_preference, ) => response.map_err(CodexHandshakeAttemptError::Client), } } @@ -472,6 +476,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { request, upstream_model, transport_policy, + websocket_connection_preference, context, selector, quota, @@ -480,7 +485,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { output_started_at, session_affinity_key, session_affinity_key_hash, - session_transport_fallbacks, + session_transport_recovery, + mut session_recovery_probe, websocket_retry_count, stream_max_retries, mut session_capture, @@ -530,6 +536,7 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { active_account.id().as_str(), context.deadline(), &cancellation, + websocket_connection_preference, ) .await; let websocket_failure_policy = match &response { @@ -563,7 +570,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { 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, + session_transport_recovery: &session_transport_recovery, + session_recovery_probe: &mut session_recovery_probe, }, ); } @@ -701,7 +709,10 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { Err(mut failure) => { if failure.websocket_transport_retryable && response_transport == CodexBackendTransport::WebSocket - && !pre_commit_events.is_committed() + && (matches!( + failure.error.send_state(), + UpstreamSendState::Sent | UpstreamSendState::Ambiguous + ) || !pre_commit_events.is_committed()) { apply_websocket_recovery_policy( &mut failure, @@ -715,7 +726,8 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { 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, + session_transport_recovery: &session_transport_recovery, + session_recovery_probe: &mut session_recovery_probe, }, ); } @@ -808,6 +820,21 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { .any(|event| matches!(event, GatewayEvent::Completed(_))); let terminal_changed = completed && observation_state.mark_completed(terminal_response_is_incomplete(&events)); + if response_transport == CodexBackendTransport::WebSocket + && completed + && terminal_failure.is_none() + && let Some(probe) = session_recovery_probe.take() + { + let session_recovery_cleared = probe.succeed(); + tracing::info!( + request_id = %context.request_id(), + attempt_index = context.attempt_index().get(), + session_affinity_key_hash = session_affinity_key_hash.as_deref().unwrap_or(""), + session_transport_action = "websocket_recovered", + session_recovery_cleared, + "OpenAI session fresh WebSocket probe reached a terminal response" + ); + } if allows_account_state_mutation && (completed || terminal_failure.is_some()) { synchronize_passive_quota( "a, @@ -939,6 +966,21 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { .any(|event| matches!(event, GatewayEvent::Completed(_))); let terminal_changed = completed && observation_state.mark_completed(terminal_response_is_incomplete(&events)); + if response_transport == CodexBackendTransport::WebSocket + && completed + && terminal_failure.is_none() + && let Some(probe) = session_recovery_probe.take() + { + let session_recovery_cleared = probe.succeed(); + tracing::info!( + request_id = %context.request_id(), + attempt_index = context.attempt_index().get(), + session_affinity_key_hash = session_affinity_key_hash.as_deref().unwrap_or(""), + session_transport_action = "websocket_recovered", + session_recovery_cleared, + "OpenAI session fresh WebSocket probe reached a terminal response" + ); + } if allows_account_state_mutation && completed && terminal_failure.is_none() { // 同上：尾部 finish() 也可能产出 completed，亲和记录必须先于任何下游 yield。 selector diff --git a/backend/crates/providers/openai/src/provider/failure.rs b/backend/crates/providers/openai/src/provider/failure.rs index f9eb0222c..c35004361 100644 --- a/backend/crates/providers/openai/src/provider/failure.rs +++ b/backend/crates/providers/openai/src/provider/failure.rs @@ -458,7 +458,8 @@ pub(super) struct WebSocketRecoveryContext<'a> { pub(super) account_id: &'a str, pub(super) session_affinity_key: Option<&'a ProviderSessionAffinityKey>, pub(super) session_affinity_key_hash: Option<&'a str>, - pub(super) session_transport_fallbacks: &'a CodexSessionTransportFallbacks, + pub(super) session_transport_recovery: &'a CodexSessionTransportRecovery, + pub(super) session_recovery_probe: &'a mut Option, } pub(super) fn websocket_client_failure_policy( @@ -509,10 +510,18 @@ pub(super) fn apply_websocket_recovery_policy( UpstreamSendState::Sent | UpstreamSendState::Ambiguous ) { // payload 已经发送或发送结果不确定时，代理不能用同一业务输入自动重放。 - // 只为客户端下一次完整请求记住 HTTP 偏好，当前请求保持终态失败。 - let sticky_state_transition = context - .session_affinity_key - .map(|key| context.session_transport_fallbacks.disable_websocket(key)); + // 当前请求保持终态失败；恢复状态只裁决客户端主动发起的下一次请求。 + let session_recovery_transition = match context.session_recovery_probe.take() { + Some(probe) => Some(probe.post_send_failed()), + None => context.session_affinity_key.map(|key| { + context + .session_transport_recovery + .record_post_send_failure(key) + }), + }; + let session_http_cooldown_ms = session_recovery_transition + .and_then(CodexSessionRecoveryTransition::cooldown) + .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)); tracing::warn!( request_id = context.request_id, attempt_index = context.attempt_index, @@ -527,8 +536,12 @@ pub(super) fn apply_websocket_recovery_policy( continuation_recovery_action = "client_replay_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), + session_transport_action = session_recovery_transition + .map_or("untracked", CodexSessionRecoveryTransition::action), + session_websocket_failure_count = session_recovery_transition + .map_or(0, CodexSessionRecoveryTransition::failure_count), + session_http_cooldown_ms = session_http_cooldown_ms.unwrap_or_default(), + session_http_cooldown_present = session_http_cooldown_ms.is_some(), "OpenAI upstream WebSocket failed after payload send; proxy replay was suppressed" ); return; @@ -576,9 +589,21 @@ pub(super) fn apply_websocket_recovery_policy( } failure.error.set_pre_delivery_transport_fallback(); - let sticky_state_transition = context - .session_affinity_key - .map(|key| context.session_transport_fallbacks.disable_websocket(key)); + let session_fallback = match context.policy { + WebSocketFailurePolicy::Budgeted => SessionWebSocketFallback::RetryBudgetExhausted, + WebSocketFailurePolicy::ImmediateFallback => SessionWebSocketFallback::UpgradeRequired, + }; + let session_recovery_transition = match context.session_recovery_probe.take() { + Some(probe) => Some(probe.fallback(session_fallback)), + None => context.session_affinity_key.map(|key| { + context + .session_transport_recovery + .record_websocket_fallback(key, session_fallback) + }), + }; + let session_http_cooldown_ms = session_recovery_transition + .and_then(CodexSessionRecoveryTransition::cooldown) + .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)); tracing::warn!( request_id = context.request_id, attempt_index = context.attempt_index, @@ -598,9 +623,13 @@ pub(super) fn apply_websocket_recovery_policy( }, 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" + session_transport_action = + session_recovery_transition.map_or("untracked", CodexSessionRecoveryTransition::action), + session_websocket_failure_count = + session_recovery_transition.map_or(0, CodexSessionRecoveryTransition::failure_count), + session_http_cooldown_ms = session_http_cooldown_ms.unwrap_or_default(), + session_http_cooldown_present = session_http_cooldown_ms.is_some(), + "OpenAI upstream WebSocket entered session HTTP cooldown" ); } diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 697f683cd..779522c6c 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -61,7 +61,10 @@ use crate::credential::{ derive_codex_cyber_policy_session_key, derive_codex_session_affinity, derive_previous_response_id_hash, }; -use crate::session_transport::CodexSessionTransportFallbacks; +use crate::session_transport::{ + CodexSessionRecoveryTransition, CodexSessionTransportDecision, CodexSessionTransportRecovery, + CodexSessionWebSocketProbe, SessionWebSocketFallback, +}; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, }; @@ -90,7 +93,8 @@ use crate::transport::{ CODEX_RESPONSES_PATH, CodexAccountSelectionTelemetry, CodexBackendClient, CodexBackendJsonResponse, CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexRateLimitUpdates, CodexRequestContext, CodexResponseMetadata, - CodexTransportMetrics, CodexUpstreamDiagnostics, CodexWebSocketPool, endpoint_url, + CodexTransportMetrics, CodexUpstreamDiagnostics, CodexWebSocketPool, + WebSocketConnectionPreference, endpoint_url, }; mod execution; @@ -145,7 +149,7 @@ pub struct CodexProvider { image_edits_url: Url, search_url: Url, session_identity: Option, - session_transport_fallbacks: CodexSessionTransportFallbacks, + session_transport_recovery: CodexSessionTransportRecovery, stream_max_retries: u32, } @@ -185,7 +189,7 @@ impl CodexProvider { image_edits_url, search_url, session_identity: None, - session_transport_fallbacks: CodexSessionTransportFallbacks::default(), + session_transport_recovery: CodexSessionTransportRecovery::default(), stream_max_retries, }) } @@ -442,26 +446,46 @@ impl Provider for CodexProvider { ); 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.allows_sticky_http_fallback() - && session_affinity.as_ref().is_some_and(|affinity| { - self.session_transport_fallbacks - .is_http_only(affinity.key()) - }); + let session_transport_decision = if matches!( + context.transport(), + AttemptTransport::Default | AttemptTransport::Retry(_) + ) && requested_transport + == CodexProviderTransport::PreferWebSocket + && requirement.allows_session_transport_recovery() + && let Some(affinity) = session_affinity.as_ref() + { + self.session_transport_recovery.decide(affinity.key()) + } else { + CodexSessionTransportDecision::Default + }; + let session_transport_action = session_transport_decision.action(); + let session_websocket_failure_count = session_transport_decision.failure_count(); + let session_http_cooldown_ms = session_transport_decision + .retry_after() + .map(|duration| u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)); + let mut websocket_connection_preference = WebSocketConnectionPreference::ReuseOrConnect; + let mut session_recovery_probe = None; let transport = if requirement.requires_websocket() { CodexProviderTransport::PreferWebSocket } else { match context.transport() { - AttemptTransport::Default if sticky_http_fallback => { - CodexProviderTransport::HttpOnly + AttemptTransport::Default | AttemptTransport::Retry(_) => { + match session_transport_decision { + CodexSessionTransportDecision::Default => requested_transport, + CodexSessionTransportDecision::FreshWebSocket { probe, .. } => { + websocket_connection_preference = WebSocketConnectionPreference::Fresh; + session_recovery_probe = Some(probe); + CodexProviderTransport::PreferWebSocket + } + CodexSessionTransportDecision::HttpSse { .. } => { + CodexProviderTransport::HttpOnly + } + } } - AttemptTransport::Default => requested_transport, - AttemptTransport::Retry(_) => requested_transport, AttemptTransport::Fallback => CodexProviderTransport::HttpOnly, } }; - if sticky_http_fallback { + if let Some(session_transport_action) = session_transport_action { tracing::info!( request_id = %context.request_id(), attempt_index = context.attempt_index().get(), @@ -469,7 +493,11 @@ impl Provider for CodexProvider { .as_ref() .map_or("", CodexSessionAffinity::key_hash), transport_requirement = requirement.as_str(), - "Reusing session-sticky OpenAI upstream HTTP fallback" + session_transport_action, + session_websocket_failure_count = session_websocket_failure_count.unwrap_or_default(), + session_http_cooldown_ms = session_http_cooldown_ms.unwrap_or_default(), + session_http_cooldown_present = session_http_cooldown_ms.is_some(), + "Applied OpenAI session transport recovery decision" ); } apply_transport(&mut upstream_request, transport); @@ -518,7 +546,9 @@ impl Provider for CodexProvider { output_started_at, session_affinity_key, session_affinity_key_hash, - session_transport_fallbacks: self.session_transport_fallbacks.clone(), + websocket_connection_preference, + session_transport_recovery: self.session_transport_recovery.clone(), + session_recovery_probe, websocket_retry_count, stream_max_retries: self.stream_max_retries, session_capture, diff --git a/backend/crates/providers/openai/src/session_transport.rs b/backend/crates/providers/openai/src/session_transport.rs index f4ea2f34e..05ee153e7 100644 --- a/backend/crates/providers/openai/src/session_transport.rs +++ b/backend/crates/providers/openai/src/session_transport.rs @@ -1,104 +1,457 @@ -//! OpenAI 根会话级上游传输降级状态。 +//! OpenAI 根会话级上游传输恢复状态。 use std::collections::HashMap; use std::sync::{Arc, Mutex, MutexGuard}; -use std::time::{Duration, Instant}; +use std::time::Duration; use gateway_core::provider_ports::ProviderSessionAffinityKey; +use tokio::time::Instant; +use uuid::Uuid; -use crate::credential::CODEX_ROOT_SESSION_TTL; - +/// 未成功恢复的失败链需要跨过最长冷却，避免 4 小时档在半开前被清理。 +const SESSION_RECOVERY_RETENTION: Duration = Duration::from_secs(8 * 60 * 60); +const HTTP_COOLDOWNS: [Duration; 4] = [ + Duration::from_secs(5 * 60), + Duration::from_secs(30 * 60), + Duration::from_secs(60 * 60), + Duration::from_secs(4 * 60 * 60), +]; /// Provider 是进程级单例；限制不可显式观测结束的客户端会话占用。 -const MAX_SESSION_HTTP_FALLBACKS: usize = 16_384; +const MAX_SESSION_RECOVERIES: usize = 16_384; #[derive(Clone)] -pub(crate) struct CodexSessionTransportFallbacks { - state: Arc>, - ttl: Duration, +pub(crate) struct CodexSessionTransportRecovery { + state: Arc>, + http_cooldowns: [Duration; 4], + retention: Duration, capacity: usize, } #[derive(Default)] -struct SessionTransportFallbackState { - http_only: HashMap, +struct SessionTransportRecoveryState { + sessions: HashMap, } -impl Default for CodexSessionTransportFallbacks { +enum SessionCircuit { + FreshPending { + failure_count: u8, + last_failure_at: Instant, + }, + Open { + failure_count: u8, + last_failure_at: Instant, + until: Instant, + }, + HalfOpen { + failure_count: u8, + last_failure_at: Instant, + probe_id: Uuid, + }, +} + +impl SessionCircuit { + const fn failure_count(&self) -> u8 { + match self { + Self::FreshPending { failure_count, .. } + | Self::Open { failure_count, .. } + | Self::HalfOpen { failure_count, .. } => *failure_count, + } + } + + const fn last_failure_at(&self) -> Instant { + match self { + Self::FreshPending { + last_failure_at, .. + } + | Self::Open { + last_failure_at, .. + } + | Self::HalfOpen { + last_failure_at, .. + } => *last_failure_at, + } + } +} + +impl Default for CodexSessionTransportRecovery { fn default() -> Self { - Self::with_limits(CODEX_ROOT_SESSION_TTL, MAX_SESSION_HTTP_FALLBACKS) + Self { + state: Arc::new(Mutex::new(SessionTransportRecoveryState::default())), + http_cooldowns: HTTP_COOLDOWNS, + retention: SESSION_RECOVERY_RETENTION, + capacity: MAX_SESSION_RECOVERIES, + } } } -impl CodexSessionTransportFallbacks { - fn with_limits(ttl: Duration, capacity: usize) -> Self { - debug_assert!(!ttl.is_zero()); - debug_assert!(capacity > 0); +#[derive(Default)] +pub(crate) enum CodexSessionTransportDecision { + #[default] + Default, + FreshWebSocket { + failure_count: u8, + probe: CodexSessionWebSocketProbe, + }, + HttpSse { + reason: SessionHttpFallbackReason, + failure_count: u8, + retry_after: Option, + }, +} + +impl CodexSessionTransportDecision { + pub(crate) const fn action(&self) -> Option<&'static str> { + match self { + Self::Default => None, + Self::FreshWebSocket { .. } => Some("fresh_websocket_probe"), + Self::HttpSse { + reason: SessionHttpFallbackReason::Cooldown, + .. + } => Some("http_cooldown"), + Self::HttpSse { + reason: SessionHttpFallbackReason::ProbeInFlight, + .. + } => Some("http_probe_in_flight"), + } + } + + pub(crate) const fn failure_count(&self) -> Option { + match self { + Self::Default => None, + Self::FreshWebSocket { failure_count, .. } | Self::HttpSse { failure_count, .. } => { + Some(*failure_count) + } + } + } + + pub(crate) const fn retry_after(&self) -> Option { + match self { + Self::HttpSse { retry_after, .. } => *retry_after, + Self::Default | Self::FreshWebSocket { .. } => None, + } + } +} + +#[derive(Clone, Copy)] +pub(crate) enum SessionWebSocketFallback { + RetryBudgetExhausted, + UpgradeRequired, +} + +impl SessionWebSocketFallback { + const fn minimum_failure_count(self) -> u8 { + match self { + Self::RetryBudgetExhausted => 2, + Self::UpgradeRequired => 5, + } + } +} + +#[derive(Clone, Copy)] +pub(crate) enum SessionHttpFallbackReason { + Cooldown, + ProbeInFlight, +} + +#[derive(Clone, Copy)] +pub(crate) struct CodexSessionRecoveryTransition { + failure_count: u8, + cooldown: Option, +} + +impl CodexSessionRecoveryTransition { + pub(crate) const fn action(self) -> &'static str { + if self.cooldown.is_some() { + "http_cooldown" + } else { + "fresh_websocket_next" + } + } + + pub(crate) const fn failure_count(self) -> u8 { + self.failure_count + } + + pub(crate) const fn cooldown(self) -> Option { + self.cooldown + } +} + +/// 一次根会话级 fresh WebSocket 许可。 +/// +/// 许可未得到明确结果就被释放时，恢复状态会重新允许下一次半开探测。 +pub(crate) struct CodexSessionWebSocketProbe { + recovery: CodexSessionTransportRecovery, + key: ProviderSessionAffinityKey, + probe_id: Uuid, + armed: bool, +} + +impl CodexSessionWebSocketProbe { + fn new( + recovery: CodexSessionTransportRecovery, + key: ProviderSessionAffinityKey, + probe_id: Uuid, + ) -> Self { Self { - state: Arc::new(Mutex::new(SessionTransportFallbackState::default())), - ttl, - capacity, + recovery, + key, + probe_id, + armed: true, } } - /// 返回该根会话是否已粘性切到 HTTP，并为活跃会话续期。 - pub(crate) fn is_http_only(&self, key: &ProviderSessionAffinityKey) -> bool { - self.is_http_only_at(key, Instant::now()) + pub(crate) fn succeed(mut self) -> bool { + let recovered = self.recovery.record_probe_success(&self.key, self.probe_id); + self.armed = false; + recovered } - /// 在当前 WebSocket 重试预算耗尽后，为根会话启用粘性 HTTP。 - /// - /// 返回 `true` 表示本次发生了禁用状态转换；已禁用会话只续期。 - pub(crate) fn disable_websocket(&self, key: &ProviderSessionAffinityKey) -> bool { - self.disable_websocket_at(key, Instant::now()) + pub(crate) fn post_send_failed(mut self) -> CodexSessionRecoveryTransition { + let transition = self.recovery.record_post_send_failure(&self.key); + self.armed = false; + transition } - fn is_http_only_at(&self, key: &ProviderSessionAffinityKey, now: Instant) -> bool { + pub(crate) fn fallback( + mut self, + fallback: SessionWebSocketFallback, + ) -> CodexSessionRecoveryTransition { + let transition = self.recovery.record_websocket_fallback(&self.key, fallback); + self.armed = false; + transition + } +} + +impl Drop for CodexSessionWebSocketProbe { + fn drop(&mut self) { + if self.armed { + self.recovery.cancel_probe(&self.key, self.probe_id); + } + } +} + +impl CodexSessionTransportRecovery { + pub(crate) fn decide(&self, key: &ProviderSessionAffinityKey) -> CodexSessionTransportDecision { + self.decide_at(key, Instant::now()) + } + + pub(crate) fn record_post_send_failure( + &self, + key: &ProviderSessionAffinityKey, + ) -> CodexSessionRecoveryTransition { + self.record_post_send_failure_at(key, Instant::now()) + } + + pub(crate) fn record_websocket_fallback( + &self, + key: &ProviderSessionAffinityKey, + fallback: SessionWebSocketFallback, + ) -> CodexSessionRecoveryTransition { + self.record_websocket_fallback_at(key, fallback, Instant::now()) + } + + fn decide_at( + &self, + key: &ProviderSessionAffinityKey, + now: Instant, + ) -> CodexSessionTransportDecision { let mut state = self.lock(); - let Some(last_seen) = state.http_only.get_mut(key) else { - return false; + let stale = state.sessions.get(key).is_some_and(|circuit| { + now.saturating_duration_since(circuit.last_failure_at()) >= self.retention + }); + if stale { + state.sessions.remove(key); + return CodexSessionTransportDecision::Default; + } + + let Some(circuit) = state.sessions.get_mut(key) else { + return CodexSessionTransportDecision::Default; }; - if now.saturating_duration_since(*last_seen) >= self.ttl { - state.http_only.remove(key); - return false; + match circuit { + SessionCircuit::Open { + failure_count, + until, + .. + } if now < *until => CodexSessionTransportDecision::HttpSse { + reason: SessionHttpFallbackReason::Cooldown, + failure_count: *failure_count, + retry_after: Some(until.saturating_duration_since(now)), + }, + SessionCircuit::HalfOpen { failure_count, .. } => { + CodexSessionTransportDecision::HttpSse { + reason: SessionHttpFallbackReason::ProbeInFlight, + failure_count: *failure_count, + retry_after: None, + } + } + SessionCircuit::FreshPending { + failure_count, + last_failure_at, + } + | SessionCircuit::Open { + failure_count, + last_failure_at, + .. + } => { + let failure_count = *failure_count; + let last_failure_at = *last_failure_at; + let probe_id = Uuid::new_v4(); + *circuit = SessionCircuit::HalfOpen { + failure_count, + last_failure_at, + probe_id, + }; + CodexSessionTransportDecision::FreshWebSocket { + failure_count, + probe: CodexSessionWebSocketProbe::new(self.clone(), key.clone(), probe_id), + } + } } - *last_seen = now; - true } - fn disable_websocket_at(&self, key: &ProviderSessionAffinityKey, now: Instant) -> bool { + fn record_post_send_failure_at( + &self, + key: &ProviderSessionAffinityKey, + now: Instant, + ) -> CodexSessionRecoveryTransition { 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) { + let failure_count = self.next_failure_count(state.sessions.get(key), now); + if !state.sessions.contains_key(key) { self.make_room(&mut state, now); } - state.http_only.insert(key.clone(), now); - !already_disabled + if failure_count == 1 { + state.sessions.insert( + key.clone(), + SessionCircuit::FreshPending { + failure_count, + last_failure_at: now, + }, + ); + return CodexSessionRecoveryTransition { + failure_count, + cooldown: None, + }; + } + + let cooldown = self.cooldown_for(failure_count); + state.sessions.insert( + key.clone(), + SessionCircuit::Open { + failure_count, + last_failure_at: now, + until: now + cooldown, + }, + ); + CodexSessionRecoveryTransition { + failure_count, + cooldown: Some(cooldown), + } } - fn make_room(&self, state: &mut SessionTransportFallbackState, now: Instant) { - if state.http_only.len() < self.capacity { + fn record_websocket_fallback_at( + &self, + key: &ProviderSessionAffinityKey, + fallback: SessionWebSocketFallback, + now: Instant, + ) -> CodexSessionRecoveryTransition { + let mut state = self.lock(); + let failure_count = self + .next_failure_count(state.sessions.get(key), now) + .max(fallback.minimum_failure_count()); + if !state.sessions.contains_key(key) { + self.make_room(&mut state, now); + } + let cooldown = self.cooldown_for(failure_count); + state.sessions.insert( + key.clone(), + SessionCircuit::Open { + failure_count, + last_failure_at: now, + until: now + cooldown, + }, + ); + CodexSessionRecoveryTransition { + failure_count, + cooldown: Some(cooldown), + } + } + + fn next_failure_count(&self, circuit: Option<&SessionCircuit>, now: Instant) -> u8 { + circuit.map_or(1, |circuit| { + if now.saturating_duration_since(circuit.last_failure_at()) < self.retention { + circuit.failure_count().saturating_add(1) + } else { + 1 + } + }) + } + + fn cooldown_for(&self, failure_count: u8) -> Duration { + let index = usize::from(failure_count.saturating_sub(2)).min(self.http_cooldowns.len() - 1); + self.http_cooldowns[index] + } + + fn record_probe_success(&self, key: &ProviderSessionAffinityKey, probe_id: Uuid) -> bool { + let mut state = self.lock(); + let owns_probe = matches!( + state.sessions.get(key), + Some(SessionCircuit::HalfOpen { + probe_id: active_probe_id, + .. + }) if *active_probe_id == probe_id + ); + if owns_probe { + state.sessions.remove(key); + } + owns_probe + } + + fn cancel_probe(&self, key: &ProviderSessionAffinityKey, probe_id: Uuid) { + let mut state = self.lock(); + let Some(SessionCircuit::HalfOpen { + failure_count, + last_failure_at, + probe_id: active_probe_id, + }) = state.sessions.get(key) + else { + return; + }; + if *active_probe_id != probe_id { return; } - state - .http_only - .retain(|_, last_seen| now.saturating_duration_since(*last_seen) < self.ttl); - if state.http_only.len() < self.capacity { + let failure_count = *failure_count; + let last_failure_at = *last_failure_at; + state.sessions.insert( + key.clone(), + SessionCircuit::FreshPending { + failure_count, + last_failure_at, + }, + ); + } + + fn make_room(&self, state: &mut SessionTransportRecoveryState, now: Instant) { + if state.sessions.len() < self.capacity { + return; + } + state.sessions.retain(|_, circuit| { + now.saturating_duration_since(circuit.last_failure_at()) < self.retention + }); + if state.sessions.len() < self.capacity { return; } let oldest = state - .http_only + .sessions .iter() - .min_by_key(|(_, last_seen)| **last_seen) + .min_by_key(|(_, circuit)| circuit.last_failure_at()) .map(|(key, _)| key.clone()); if let Some(oldest) = oldest { - state.http_only.remove(&oldest); + state.sessions.remove(&oldest); } } - fn lock(&self) -> MutexGuard<'_, SessionTransportFallbackState> { + fn lock(&self) -> MutexGuard<'_, SessionTransportRecoveryState> { self.state .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index c69d509c1..13b917af9 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -47,6 +47,16 @@ use crate::transport::{ use super::client::*; +/// 单次 Responses WebSocket opening 是否可以使用会话连接池。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(crate) enum WebSocketConnectionPreference { + /// 优先复用匹配的空闲连接，否则建立并登记一条新连接。 + #[default] + ReuseOrConnect, + /// 绕过连接池建立独立连接，用于根会话恢复探针。 + Fresh, +} + impl CodexBackendClient { /// 构造客户端。 pub fn new( @@ -188,9 +198,16 @@ impl CodexBackendClient { request: &CodexResponsesRequest, context: CodexRequestContext<'_>, pool_account_id: Option<&str>, + connection_preference: WebSocketConnectionPreference, ) -> CodexClientResult { let prepared = self - .prepare_response_transport(request, context, pool_account_id, false, true) + .prepare_response_transport( + request, + context, + pool_account_id, + connection_preference, + true, + ) .await?; self.create_response_stream_with_prepared(request, context, prepared) .await @@ -204,18 +221,24 @@ impl CodexBackendClient { context: CodexRequestContext<'_>, pool_account_id: Option<&str>, ) -> CodexClientResult { - self.prepare_response_transport(request, context, pool_account_id, false, false) - .await + self.prepare_response_transport( + request, + context, + pool_account_id, + WebSocketConnectionPreference::ReuseOrConnect, + false, + ) + .await } /// 与 [`Self::prepare_response_transport_with_pool_account`] 相同； - /// `force_fresh` 时绕过连接池，为连接寿命限制重试强制新建连接。 + /// fresh 偏好会绕过连接池，为受控恢复探针强制新建连接。 async fn prepare_response_transport( &self, request: &CodexResponsesRequest, context: CodexRequestContext<'_>, pool_account_id: Option<&str>, - force_fresh: bool, + connection_preference: WebSocketConnectionPreference, defer_websocket_recovery: bool, ) -> CodexClientResult { let requirement = transport_requirement(request); @@ -256,10 +279,11 @@ impl CodexBackendClient { let pool_key = self.websocket_pool_key(request, context, pool_account_id, &connection_profile); let pool_log_context = pool_key.as_ref().map(WebSocketPoolLogContext::from_key); - let pool = if force_fresh { - None - } else { - self.websocket_pool.as_deref().zip(pool_key) + let pool = match connection_preference { + WebSocketConnectionPreference::ReuseOrConnect => { + self.websocket_pool.as_deref().zip(pool_key) + } + WebSocketConnectionPreference::Fresh => None, }; let fast_path_budget = match requirement { TransportRequirement::PersistedContinuation | TransportRequirement::NewChain => { @@ -434,7 +458,13 @@ impl CodexBackendClient { ); drop(exchange); match self - .prepare_response_transport(request, context, None, true, false) + .prepare_response_transport( + request, + context, + None, + WebSocketConnectionPreference::Fresh, + false, + ) .await { Ok(PreparedResponseTransport { diff --git a/backend/crates/providers/openai/src/transport/mod.rs b/backend/crates/providers/openai/src/transport/mod.rs index 6a6604198..8bd26f19d 100644 --- a/backend/crates/providers/openai/src/transport/mod.rs +++ b/backend/crates/providers/openai/src/transport/mod.rs @@ -21,6 +21,7 @@ pub mod tls; pub mod usage; pub mod websocket; +pub(crate) use self::client_sse::WebSocketConnectionPreference; pub use self::{ canonical::{CodexCanonicalDecoder, CodexCanonicalError}, catalog::{ diff --git a/backend/crates/providers/openai/src/transport/protocol/responses.rs b/backend/crates/providers/openai/src/transport/protocol/responses.rs index 844bf2170..cdf452e22 100644 --- a/backend/crates/providers/openai/src/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/src/transport/protocol/responses.rs @@ -162,8 +162,8 @@ impl TransportRequirement { matches!(self, Self::NewChain) } - /// 会话级 sticky 状态是否可以把首选 WebSocket 改为 HTTP。 - pub fn allows_sticky_http_fallback(self) -> bool { + /// 会话级恢复状态是否可以为请求选择 fresh WebSocket 或临时 HTTP。 + pub fn allows_session_transport_recovery(self) -> bool { matches!(self, Self::NewChain) } diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index bf94b36be..c5ef76e49 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -24,7 +24,7 @@ use gateway_core::engine::{ use gateway_core::error::{ ContinuationFailure, ContinuationRecoveryDisposition, PreDeliveryRetry, ProviderErrorKind, }; -use gateway_core::event::GatewayEvent; +use gateway_core::event::{GatewayEvent, WebSocketPoolKind}; use gateway_core::lifecycle::CancellationToken; use gateway_core::operation::{ CapabilityRequirements, GenerateRequest, ImageRequest, ImageRequestKind, Operation, @@ -321,6 +321,25 @@ fn generate_with_session_context( ) } +fn generate_with_persisted_session_context( + account_id: &str, + conversation_id: &str, + session_id: &str, + thread_id: &str, +) -> GenerateRequest { + generate_with_session_context(session_id, Some(thread_id), None).with_provider_session_state( + ProviderSessionState::new( + "openai", + Map::from_iter([ + ("account_id".to_owned(), json!(account_id)), + ("conversation_id".to_owned(), json!(conversation_id)), + ("continuation_scope".to_owned(), json!("persisted")), + ]), + ) + .expect("provider session state"), + ) +} + fn http_generate_operation() -> Operation { let payload = ProtocolPayload::json_object( "openai", @@ -439,6 +458,7 @@ fn contract_account_scope() -> Arc { "acct_truncated_stream", "acct_usage_limit_request_path", "acct_websocket_close", + "acct_websocket_cooldown", "acct_websocket_busy_replay", "acct_websocket_turn_state", ] @@ -1295,9 +1315,13 @@ async fn selection_infrastructure_errors_have_a_distinct_classification() { } #[tokio::test] -async fn websocket_close_after_delivery_preserves_details_without_enabling_sticky_http() { +async fn websocket_close_after_delivery_preserves_details_and_forces_a_fresh_client_retry() { + const ACCOUNT_ID: &str = "acct_websocket_close"; + const CONVERSATION_ID: &str = "conversation-websocket-close-after-delivery"; + const SESSION_ID: &str = "committed-websocket-session"; + let store = Arc::new(MemoryAccountStore::default()); - create_account(&store, "acct_websocket_close").await; + create_account(&store, ACCOUNT_ID).await; let listener = TcpListener::bind("127.0.0.1:0") .await .expect("bind WebSocket listener"); @@ -1316,6 +1340,20 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick .await .expect("WebSocket request") .expect("valid WebSocket request"); + websocket + .send(Message::Text( + json!({ + "type": "response.created", + "response": { + "id": "resp_before_committed_close", + "model": "gpt-5.4" + } + }) + .to_string() + .into(), + )) + .await + .expect("start WebSocket response"); websocket .send(Message::Text( json!({ @@ -1345,6 +1383,20 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick .await .expect("next-turn WebSocket request") .expect("valid next-turn WebSocket request"); + websocket + .send(Message::Text( + json!({ + "type": "response.created", + "response": { + "id": "resp_after_committed_close", + "model": "gpt-5.4" + } + }) + .to_string() + .into(), + )) + .await + .expect("start next-turn WebSocket response"); websocket .send(Message::Text( json!({ @@ -1369,10 +1421,11 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick }); 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 first_operation = Operation::Generate(generate_with_persisted_session_context( + ACCOUNT_ID, + CONVERSATION_ID, + SESSION_ID, + "thread-first", )); let mut stream = provider .execute( @@ -1409,10 +1462,11 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick Some("websocket_close_1009") ); - let second_operation = Operation::Generate(generate_with_session_context( - "committed-websocket-session", - Some("thread-second"), - None, + let second_operation = Operation::Generate(generate_with_persisted_session_context( + ACCOUNT_ID, + CONVERSATION_ID, + SESSION_ID, + "thread-second", )); let mut next_stream = provider .execute( @@ -1420,18 +1474,29 @@ async fn websocket_close_after_delivery_preserves_details_without_enabling_stick context("req_after_committed_close", CancellationToken::new()), ) .await - .expect("post-delivery close must not enable sticky HTTP fallback"); + .expect("post-delivery close must allow a fresh WebSocket client retry"); assert_eq!(next_stream.metadata().transport().as_str(), "websocket"); + let mut retry_pool_observation = None; while let Some(event) = next_stream.next().await { - event.expect("next turn WebSocket response"); + let event = event.expect("next turn WebSocket response"); + if let Some(observation) = event.response_observation() + && observation.transport().as_str() == "websocket" + { + retry_pool_observation = Some(observation.websocket_pool()); + } } + assert_eq!(retry_pool_observation, Some(None)); server.await.expect("WebSocket server"); } #[tokio::test] -async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session() { +async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_chain() { + const ACCOUNT_ID: &str = "acct_websocket_close"; + const CONVERSATION_ID: &str = "conversation-websocket-ambiguous-close"; + const SESSION_ID: &str = "sticky-websocket-session"; + let store = Arc::new(MemoryAccountStore::default()); - create_account(&store, "acct_websocket_close").await; + create_account(&store, ACCOUNT_ID).await; let listener = TcpListener::bind("127.0.0.1:0") .await .expect("bind WebSocket listener"); @@ -1511,19 +1576,31 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session .await .expect("close WebSocket before terminal event"); - 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 fresh WebSocket probe"); + let mut websocket = accept_codex_test_websocket(stream).await; + let _request = websocket + .next() + .await + .expect("fresh WebSocket request") + .expect("valid fresh WebSocket request"); + websocket + .send(Message::Text( + json!({ + "type": "response.created", + "response": {"id": "resp_fresh_retry", "model": "gpt-5.4"} + }) + .to_string() + .into(), + )) + .await + .expect("start fresh WebSocket response"); + websocket + .send(completed_response("resp_fresh_retry")) + .await + .expect("complete fresh WebSocket request"); let (stream, _) = listener .accept() @@ -1568,10 +1645,11 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session let provider = provider_with_base_url_and_retry_budget(&store, base_url, 2); let mut saw_websocket_observation = false; - let operation = Operation::Generate(generate_with_session_context( - "sticky-websocket-session", - Some("thread-first"), - None, + let operation = Operation::Generate(generate_with_persisted_session_context( + ACCOUNT_ID, + CONVERSATION_ID, + SESSION_ID, + "thread-first", )); let mut stream = provider .execute( @@ -1640,23 +1718,31 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session assert_eq!(connection.exit_reason(), "normal_close"); assert!(connection.age_ms() >= connection.idle_ms()); - let second_operation = Operation::Generate(generate_with_session_context( - "sticky-websocket-session", - Some("thread-second"), - None, + let second_operation = Operation::Generate(generate_with_persisted_session_context( + ACCOUNT_ID, + CONVERSATION_ID, + SESSION_ID, + "thread-second", )); - let mut sticky_stream = provider + let mut fresh_stream = provider .execute( planned_request("openai", second_operation), - context("req_websocket_sticky_http", CancellationToken::new()), + context("req_websocket_fresh_retry", CancellationToken::new()), ) .await - .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"); + .expect("same session should prepare a fresh WebSocket stream"); + assert_eq!(fresh_stream.metadata().transport().as_str(), "websocket"); + let mut fresh_pool_observation = None; + while let Some(event) = fresh_stream.next().await { + let event = event.expect("fresh WebSocket response"); + if let Some(observation) = event.response_observation() + && observation.transport().as_str() == "websocket" + { + fresh_pool_observation = Some(observation.websocket_pool()); + } } - drop(sticky_stream); + assert_eq!(fresh_pool_observation, Some(None)); + drop(fresh_stream); let continuation_operation = Operation::Generate(generate_with_session_context( "sticky-websocket-session", @@ -1669,7 +1755,7 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session external_continuation_context("req_sticky_external_continuation"), ) .await - .expect("sticky state must not force a continuation onto HTTP"); + .expect("successful fresh probe must restore normal continuation routing"); assert_eq!( continuation_stream.metadata().transport().as_str(), "websocket" @@ -1719,12 +1805,375 @@ async fn ambiguous_websocket_close_only_sticks_new_chains_in_the_current_session ), ) .await - .expect("required warmup must bypass sticky HTTP fallback"); + .expect("required warmup must remain on WebSocket"); assert_eq!(warmup_stream.metadata().transport().as_str(), "websocket"); drop(warmup_stream); server.await.expect("WebSocket and HTTP server"); } +#[tokio::test] +async fn repeated_ambiguous_websocket_close_uses_steep_sse_cooldowns_and_one_half_open_probe() { + const ACCOUNT_ID: &str = "acct_websocket_cooldown"; + const CONVERSATION_ID: &str = "conversation-websocket-cooldown"; + + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, ACCOUNT_ID).await; + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket listener"); + let base_url = format!( + "http://{}", + listener.local_addr().expect("listener address") + ); + let server = tokio::spawn(async move { + let completed_response = |response_id: &str| { + Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": response_id, + "model": "gpt-5.4", + "status": "completed", + "output": [], + "usage": { + "input_tokens": 1, + "output_tokens": 1, + "total_tokens": 2 + } + } + }) + .to_string() + .into(), + ) + }; + + for response_id in ["resp_first_reset", "resp_fresh_probe_reset"] { + let (stream, _) = listener + .accept() + .await + .expect("accept failing WebSocket connection"); + let mut websocket = accept_codex_test_websocket(stream).await; + let _request = websocket + .next() + .await + .expect("WebSocket request") + .expect("valid WebSocket request"); + websocket + .send(Message::Text( + json!({ + "type": "response.created", + "response": {"id": response_id, "model": "gpt-5.4"} + }) + .to_string() + .into(), + )) + .await + .expect("send structural WebSocket event"); + websocket + .close(Some(CloseFrame { + code: CloseCode::Normal, + reason: "".into(), + })) + .await + .expect("close WebSocket before terminal event"); + } + + for response_id in [ + "resp_third_reset", + "resp_fourth_reset", + "resp_fifth_reset", + "resp_sixth_reset", + ] { + let (mut http, _) = listener + .accept() + .await + .expect("accept cooldown 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 cooldown HTTP response"); + + let (stream, _) = listener + .accept() + .await + .expect("accept failing half-open WebSocket connection"); + let mut websocket = accept_codex_test_websocket(stream).await; + let _request = websocket + .next() + .await + .expect("half-open WebSocket request") + .expect("valid half-open WebSocket request"); + websocket + .send(Message::Text( + json!({ + "type": "response.created", + "response": {"id": response_id, "model": "gpt-5.4"} + }) + .to_string() + .into(), + )) + .await + .expect("send half-open structural event"); + websocket + .close(Some(CloseFrame { + code: CloseCode::Normal, + reason: "".into(), + })) + .await + .expect("close half-open WebSocket before terminal event"); + } + + for _ in 0..2 { + let (mut http, _) = listener + .accept() + .await + .expect("accept final cooldown 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 final cooldown HTTP response"); + } + + for response_id in ["resp_half_open", "resp_recovered"] { + let (stream, _) = listener + .accept() + .await + .expect("accept recovery WebSocket connection"); + let mut websocket = accept_codex_test_websocket(stream).await; + let _request = websocket + .next() + .await + .expect("recovery WebSocket request") + .expect("valid recovery WebSocket request"); + websocket + .send(Message::Text( + json!({ + "type": "response.created", + "response": {"id": response_id, "model": "gpt-5.4"} + }) + .to_string() + .into(), + )) + .await + .expect("start recovery WebSocket response"); + websocket + .send(completed_response(response_id)) + .await + .expect("complete recovery WebSocket request"); + } + }); + + let provider = provider_with_base_url(&store, base_url); + let operation = |thread_id: &str| { + Operation::Generate(generate_with_persisted_session_context( + ACCOUNT_ID, + CONVERSATION_ID, + "websocket-cooldown-session", + thread_id, + )) + }; + + let mut first = provider + .execute( + planned_request("openai", operation("thread-first-reset")), + context("req_first_reset", CancellationToken::new()), + ) + .await + .expect("prepare first WebSocket stream"); + let first_error = loop { + match first.next().await { + Some(Ok(_)) => {} + Some(Err(error)) => break error, + None => panic!("first WebSocket reset must surface a provider error"), + } + }; + assert_eq!(first_error.send_state(), UpstreamSendState::Ambiguous); + assert_eq!(first_error.pre_delivery_retry(), None); + drop(first); + + let mut fresh = provider + .execute( + planned_request("openai", operation("thread-fresh-reset")), + context("req_fresh_reset", CancellationToken::new()), + ) + .await + .expect("prepare fresh WebSocket probe"); + assert_eq!(fresh.metadata().transport().as_str(), "websocket"); + let mut fresh_pool_observation = None; + let second_error = loop { + match fresh.next().await { + Some(Ok(event)) => { + if let Some(observation) = event.response_observation() + && observation.transport().as_str() == "websocket" + { + fresh_pool_observation = Some(observation.websocket_pool()); + } + } + Some(Err(error)) => break error, + None => panic!("fresh WebSocket reset must surface a provider error"), + } + }; + assert_eq!(fresh_pool_observation, Some(None)); + assert_eq!(second_error.send_state(), UpstreamSendState::Ambiguous); + assert_eq!(second_error.pre_delivery_retry(), None); + drop(fresh); + tokio::time::pause(); + + for (index, cooldown) in [ + Duration::from_secs(5 * 60), + Duration::from_secs(30 * 60), + Duration::from_secs(60 * 60), + Duration::from_secs(4 * 60 * 60), + ] + .into_iter() + .enumerate() + { + tokio::time::advance(cooldown - Duration::from_secs(60)).await; + let mut active_cooldown = provider + .execute( + planned_request( + "openai", + operation(&format!("thread-active-cooldown-{index}")), + ), + context( + &format!("req_active_cooldown_{index}"), + CancellationToken::new(), + ), + ) + .await + .expect("prepare active cooldown HTTP stream"); + assert_eq!(active_cooldown.metadata().transport().as_str(), "http_sse"); + while let Some(event) = active_cooldown.next().await { + event.expect("active cooldown HTTP response"); + } + + tokio::time::advance(Duration::from_secs(60)).await; + tokio::time::resume(); + let mut failed_probe = provider + .execute( + planned_request("openai", operation(&format!("thread-failed-probe-{index}"))), + context( + &format!("req_failed_probe_{index}"), + CancellationToken::new(), + ), + ) + .await + .expect("prepare half-open WebSocket probe"); + assert_eq!(failed_probe.metadata().transport().as_str(), "websocket"); + let mut failed_probe_pool_observation = None; + let probe_error = loop { + match failed_probe.next().await { + Some(Ok(event)) => { + if let Some(observation) = event.response_observation() + && observation.transport().as_str() == "websocket" + { + failed_probe_pool_observation = Some(observation.websocket_pool()); + } + } + Some(Err(error)) => break error, + None => panic!("half-open WebSocket reset must surface a provider error"), + } + }; + assert_eq!(failed_probe_pool_observation, Some(None)); + assert_eq!(probe_error.send_state(), UpstreamSendState::Ambiguous); + assert_eq!(probe_error.pre_delivery_retry(), None); + drop(failed_probe); + tokio::time::pause(); + } + + tokio::time::advance(Duration::from_secs(4 * 60 * 60 - 60)).await; + let mut capped_cooldown = provider + .execute( + planned_request("openai", operation("thread-capped-cooldown")), + context("req_capped_cooldown", CancellationToken::new()), + ) + .await + .expect("prepare capped cooldown HTTP stream"); + assert_eq!(capped_cooldown.metadata().transport().as_str(), "http_sse"); + while let Some(event) = capped_cooldown.next().await { + event.expect("capped cooldown HTTP response"); + } + + tokio::time::advance(Duration::from_secs(60)).await; + tokio::time::resume(); + let mut half_open = provider + .execute( + planned_request("openai", operation("thread-half-open")), + context("req_half_open", CancellationToken::new()), + ) + .await + .expect("prepare half-open WebSocket probe"); + assert_eq!(half_open.metadata().transport().as_str(), "websocket"); + + let mut concurrent = provider + .execute( + planned_request("openai", operation("thread-probe-busy")), + context("req_probe_busy", CancellationToken::new()), + ) + .await + .expect("prepare concurrent HTTP stream"); + assert_eq!(concurrent.metadata().transport().as_str(), "http_sse"); + while let Some(event) = concurrent.next().await { + event.expect("probe-busy HTTP response"); + } + + let mut half_open_pool_observation = None; + let mut half_open_completed = false; + while let Some(event) = half_open.next().await { + let event = event.expect("half-open WebSocket response"); + half_open_completed |= event + .canonical_facts() + .iter() + .any(|event| matches!(event, GatewayEvent::Completed(_))); + if let Some(observation) = event.response_observation() + && observation.transport().as_str() == "websocket" + { + half_open_pool_observation = Some(observation.websocket_pool()); + } + } + assert!(half_open_completed); + assert_eq!(half_open_pool_observation, Some(None)); + + let mut recovered = provider + .execute( + planned_request("openai", operation("thread-recovered")), + context("req_recovered", CancellationToken::new()), + ) + .await + .expect("prepare recovered WebSocket stream"); + assert_eq!(recovered.metadata().transport().as_str(), "websocket"); + let mut recovered_pool_observation = None; + while let Some(event) = recovered.next().await { + let event = event.expect("recovered WebSocket response"); + if let Some(observation) = event.response_observation() + && observation.transport().as_str() == "websocket" + { + recovered_pool_observation = Some(observation.websocket_pool()); + } + } + assert_eq!( + recovered_pool_observation, + Some(Some(WebSocketPoolKind::New)) + ); + server.await.expect("WebSocket and HTTP recovery server"); +} + #[tokio::test] async fn websocket_upgrade_required_immediately_enables_session_http_fallback() { let store = Arc::new(MemoryAccountStore::default()); diff --git a/backend/crates/providers/openai/tests/transport/protocol/responses.rs b/backend/crates/providers/openai/tests/transport/protocol/responses.rs index b58a03fba..7d5f91b54 100644 --- a/backend/crates/providers/openai/tests/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/tests/transport/protocol/responses.rs @@ -295,8 +295,8 @@ fn new_chain_should_allow_pre_delivery_http_fallback() { } #[test] -fn sticky_http_fallback_should_only_apply_to_new_chains() { - assert!(TransportRequirement::NewChain.allows_sticky_http_fallback()); +fn session_transport_recovery_should_only_apply_to_new_chains() { + assert!(TransportRequirement::NewChain.allows_session_transport_recovery()); for requirement in [ TransportRequirement::HttpRequired, @@ -306,8 +306,8 @@ fn sticky_http_fallback_should_only_apply_to_new_chains() { TransportRequirement::ExternalUnknown, ] { assert!( - !requirement.allows_sticky_http_fallback(), - "sticky HTTP must not override {}", + !requirement.allows_session_transport_recovery(), + "session recovery must not override {}", requirement.as_str() ); }