diff --git a/backend/crates/gateway-core/src/engine/coordinator.rs b/backend/crates/gateway-core/src/engine/coordinator.rs index 4bd5d9ad4..ff3ce8b45 100644 --- a/backend/crates/gateway-core/src/engine/coordinator.rs +++ b/backend/crates/gateway-core/src/engine/coordinator.rs @@ -223,8 +223,6 @@ struct CurrentAttempt { index: NonZeroU32, started_at: SystemTime, send_observed: bool, - /// Provider 已为当前 attempt 交付可供同请求恢复的不透明状态检查点。 - session_checkpoint_observed: bool, response_observation: Option, } @@ -596,7 +594,7 @@ where .wire_event(wire.protocol(), wire.event_type()); } if let Some(state) = event.session_update() { - self.observe_session_checkpoint(state); + self.observe_session_update(state); } if let Some(observation) = event.take_observation() { self.observe_response(observation); @@ -978,7 +976,6 @@ where index: next_attempt, started_at: SystemTime::now(), send_observed: false, - session_checkpoint_observed: false, response_observation: None, }); Ok(None) @@ -1004,14 +1001,13 @@ where self.mark_send_observed().await; } - fn observe_session_checkpoint(&mut self, state: &ProviderSessionState) { + fn observe_session_update(&mut self, state: &ProviderSessionState) { let Some(current) = self.current.as_mut() else { return; }; if state.provider() != current.metadata.provider().as_str() { return; } - current.session_checkpoint_observed = true; self.operation.set_provider_session_state(state.clone()); } @@ -1081,11 +1077,6 @@ where .await; } let provider_proved_replay_safe = provider_proved_replay_safe(&error); - // Ambiguous 仍是默认 fail-closed。唯一例外是 Provider 在同一 attempt 中先 - // 交付了不透明恢复检查点，并明确证明从该检查点重放安全；该例外只用于 - // 固定原账号的 transport recovery，不放开普通换号或 continuation 重放。 - let checkpointed_transport_replay = - current.session_checkpoint_observed && error.replay_is_safe(); let continuation_retry = self.prepare_continuation_retry( ¤t, &error, @@ -1108,16 +1099,14 @@ where delay, }) if self.downstream_committed_at.is_none() && !self.delivery_pending - && (attempt_send_state != UpstreamSendState::Ambiguous - || checkpointed_transport_replay) => + && attempt_send_state != UpstreamSendState::Ambiguous => { Some((AttemptTransport::Retry(retry_index), delay)) } Some(crate::error::PreDeliveryRetry::SameAccountTransportFallback) if self.downstream_committed_at.is_none() && !self.delivery_pending - && (attempt_send_state != UpstreamSendState::Ambiguous - || checkpointed_transport_replay) => + && attempt_send_state != UpstreamSendState::Ambiguous => { Some((AttemptTransport::Fallback, Duration::ZERO)) } diff --git a/backend/crates/gateway-core/tests/engine/coordinator.rs b/backend/crates/gateway-core/tests/engine/coordinator.rs index 982c68f61..a9bad8108 100644 --- a/backend/crates/gateway-core/tests/engine/coordinator.rs +++ b/backend/crates/gateway-core/tests/engine/coordinator.rs @@ -2909,7 +2909,7 @@ fn ambiguous_pre_delivery_retry_marker_does_not_rotate_account() { } #[test] -fn ambiguous_transport_retry_uses_provider_session_checkpoint_on_same_account() { +fn ambiguous_transport_retry_is_rejected_even_with_provider_session_state() { let operation = generate_operation(); let route_plan = plan(&operation); let checkpoint = ProviderSessionState::new( @@ -2920,7 +2920,7 @@ fn ambiguous_transport_retry_uses_provider_session_checkpoint_on_same_account() let mut checkpoint_event = ProviderEvent::observation(ProviderResponseObservation::new( UpstreamTransport::new("websocket").expect("transport"), )); - checkpoint_event.attach_session_update(checkpoint.clone()); + checkpoint_event.attach_session_update(checkpoint); let retry_index = NonZeroU32::new(1).expect("retry index"); let (coordinator, store, provider) = coordinator(vec![ Script::ObservedStream { @@ -2949,27 +2949,14 @@ fn ambiguous_transport_retry_uses_provider_session_checkpoint_on_same_account() CancellationToken::new(), )) .expect("start execution"); - block_on(session.collect_uncommitted()).expect("checkpointed retry 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(), 2); - assert_eq!( - contexts[1].transport(), - AttemptTransport::Retry(retry_index) - ); - assert_eq!(contexts[1].required_account(), Some(&original)); - drop(contexts); - let operations = provider.operations.lock().expect("operations lock"); - assert_eq!( - operations[1].provider_session_state("openai"), - Some(&checkpoint) - ); + let error = + block_on(session.collect_uncommitted()).expect_err("ambiguous send cannot be replayed"); + assert!(matches!(error, EngineError::Provider(_))); + assert_eq!(provider.contexts.lock().expect("contexts lock").len(), 1); let state = store.state.lock().expect("store lock"); - assert_eq!(state.intermediate_failures, 1); - assert_eq!(state.finalizations[0].attempt_count, 2); - assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Succeeded); + assert_eq!(state.intermediate_failures, 0); + assert_eq!(state.finalizations[0].attempt_count, 1); + assert_eq!(state.finalizations[0].outcome, ExecutionOutcome::Failed); } #[test] diff --git a/backend/crates/providers/openai/src/config.rs b/backend/crates/providers/openai/src/config.rs index 395f882ba..a4f7a4eeb 100644 --- a/backend/crates/providers/openai/src/config.rs +++ b/backend/crates/providers/openai/src/config.rs @@ -19,7 +19,7 @@ use crate::{ }; /// 本服务的流式请求默认重试次数。 -pub const DEFAULT_STREAM_MAX_RETRIES: u64 = 3; +pub const DEFAULT_STREAM_MAX_RETRIES: u64 = 5; /// 防止错误配置产生无界隐藏重放；与官方 Codex 的硬上限一致。 pub const MAX_STREAM_MAX_RETRIES: u64 = 100; @@ -175,7 +175,7 @@ pub struct CodexWebSocketPoolSettings { pub enabled: bool, pub max_age_ms: u64, pub max_connecting: usize, - pub initial_event_timeout_ms: u64, + pub stream_idle_timeout_ms: u64, } impl Default for CodexWebSocketPoolSettings { @@ -184,7 +184,7 @@ impl Default for CodexWebSocketPoolSettings { enabled: true, max_age_ms: 55 * 60 * 1000, max_connecting: 8, - initial_event_timeout_ms: 20_000, + stream_idle_timeout_ms: 300_000, } } } @@ -202,8 +202,8 @@ impl CodexWebSocketPoolSettings { enabled: self.enabled, max_age: Duration::from_millis(self.max_age_ms), max_connecting: self.max_connecting, - initial_event_timeout: (self.initial_event_timeout_ms != 0) - .then(|| Duration::from_millis(self.initial_event_timeout_ms)), + stream_idle_timeout: (self.stream_idle_timeout_ms != 0) + .then(|| Duration::from_millis(self.stream_idle_timeout_ms)), ..CodexWebSocketPoolConfig::default() } } diff --git a/backend/crates/providers/openai/src/provider/execution.rs b/backend/crates/providers/openai/src/provider/execution.rs index 8b6405c7f..ad44affe2 100644 --- a/backend/crates/providers/openai/src/provider/execution.rs +++ b/backend/crates/providers/openai/src/provider/execution.rs @@ -147,7 +147,6 @@ 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, @@ -157,7 +156,6 @@ pub(super) struct ColdResponse { pub(super) session_affinity_key: Option, pub(super) session_affinity_key_hash: Option, 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, @@ -186,13 +184,6 @@ pub(super) struct OpenAiSessionState { #[serde(default)] pub(super) client_turn_id: Option, pub(super) continuation_scope: OpenAiContinuationScope, - /// 该状态来自同一业务请求的 metadata 检查点，只允许被紧随其后的恢复 attempt 消费。 - #[serde(default, skip_serializing_if = "is_false")] - pub(super) retry_checkpoint: bool, -} - -const fn is_false(value: &bool) -> bool { - !*value } #[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -241,7 +232,6 @@ pub(super) fn encode_openai_session_state( fn encode_openai_session_capture( capture: &OpenAiSessionCapture, - retry_checkpoint: bool, ) -> Result { let Some(continuation_scope) = capture.continuation_scope else { return Err(provider_error( @@ -255,7 +245,6 @@ fn encode_openai_session_capture( turn_state: capture.turn_state.clone(), client_turn_id: capture.client_turn_id.clone(), continuation_scope, - retry_checkpoint, }) } @@ -272,7 +261,7 @@ pub(super) fn attach_openai_session_update( let Some(capture) = capture.take() else { return; }; - let Ok(update) = encode_openai_session_capture(&capture, false) else { + let Ok(update) = encode_openai_session_capture(&capture) else { return; }; events[terminal_index].attach_session_update(update); @@ -312,7 +301,6 @@ 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); @@ -321,11 +309,10 @@ pub(super) 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_deferred_websocket_recovery( + response = client.create_response_stream_with_pool_account( request, request_context, Some(account_id), - connection_preference, ) => response.map_err(CodexHandshakeAttemptError::Client), } } @@ -496,7 +483,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { request, upstream_model, transport_policy, - websocket_connection_preference, context, selector, quota, @@ -506,7 +492,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { session_affinity_key, session_affinity_key_hash, session_transport_recovery, - mut session_recovery_probe, websocket_retry_count, stream_max_retries, mut session_capture, @@ -557,7 +542,6 @@ 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 { @@ -584,7 +568,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { WebSocketRecoveryContext { policy, requirement: request_transport_requirement, - replay_checkpoint_available: false, retry_count: websocket_retry_count, max_retries: stream_max_retries, request_id: context.request_id().as_str(), @@ -593,7 +576,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { session_affinity_key: session_affinity_key.as_ref(), session_affinity_key_hash: session_affinity_key_hash.as_deref(), session_transport_recovery: &session_transport_recovery, - session_recovery_probe: &mut session_recovery_probe, }, ); } @@ -616,9 +598,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { Err(failure.error)?; return; } - let mut turn_state_checkpoint_observed = response.transport - == CodexBackendTransport::WebSocket - && response.turn_state.is_some(); if let Some(capture) = session_capture.as_mut() { capture.continuation_scope = Some(if capture.response_store { OpenAiContinuationScope::Persisted @@ -744,35 +723,11 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { &mut observation_state, ) .await; - turn_state_checkpoint_observed |= turn_state_merge.is_some(); - let retry_checkpoint_update = if turn_state_checkpoint_observed - && failure.websocket_checkpoint_retryable - { - session_capture - .as_ref() - .and_then(|capture| encode_openai_session_capture(capture, true).ok()) + let observation_event = if rate_limits_changed || turn_state_merge.is_some() { + observation_state.observation().map(ProviderEvent::observation) } else { None }; - let mut observation_event = if rate_limits_changed - || turn_state_merge.is_some() - || retry_checkpoint_update.is_some() - { - observation_state - .observation() - .map(ProviderEvent::observation) - } else { - None - }; - if let (Some(event), Some(update)) = - (observation_event.as_mut(), retry_checkpoint_update) - { - event.attach_session_update(update); - } - let replay_checkpoint_available = observation_event - .as_ref() - .and_then(ProviderEvent::session_update) - .is_some(); if failure.websocket_transport_retryable && response_transport == CodexBackendTransport::WebSocket && (matches!( @@ -785,7 +740,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { WebSocketRecoveryContext { policy: WebSocketFailurePolicy::Budgeted, requirement: request_transport_requirement, - replay_checkpoint_available, retry_count: websocket_retry_count, max_retries: stream_max_retries, request_id: context.request_id().as_str(), @@ -794,7 +748,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { session_affinity_key: session_affinity_key.as_ref(), session_affinity_key_hash: session_affinity_key_hash.as_deref(), session_transport_recovery: &session_transport_recovery, - session_recovery_probe: &mut session_recovery_probe, }, ); } @@ -829,7 +782,6 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { &mut observation_state, ) .await; - turn_state_checkpoint_observed |= turn_state_merge.is_some(); let turn_state_changed = turn_state_merge.unwrap_or(false); let first_event_changed = observation_state.observe_stream_chunk(&chunk, output_started_at); @@ -879,19 +831,10 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { 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() + && completed && terminal_failure.is_none() + && let Some(key) = session_affinity_key.as_ref() { - 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" - ); + session_transport_recovery.websocket_succeeded(key); } if allows_account_state_mutation && (completed || terminal_failure.is_some()) { synchronize_passive_quota( @@ -1022,19 +965,10 @@ pub(super) fn cold_response_stream(response: ColdResponse) -> EventStream { 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() + && completed && terminal_failure.is_none() + && let Some(key) = session_affinity_key.as_ref() { - 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" - ); + session_transport_recovery.websocket_succeeded(key); } if allows_account_state_mutation && completed && terminal_failure.is_none() { // 同上：尾部 finish() 也可能产出 completed，亲和记录必须先于任何下游 yield。 diff --git a/backend/crates/providers/openai/src/provider/failure.rs b/backend/crates/providers/openai/src/provider/failure.rs index 9289cc9a3..b11781931 100644 --- a/backend/crates/providers/openai/src/provider/failure.rs +++ b/backend/crates/providers/openai/src/provider/failure.rs @@ -39,8 +39,6 @@ pub fn openai_failure_affects_account_score(error: &ProviderError) -> bool { pub(super) struct MappedProviderFailure { pub(super) error: ProviderError, pub(super) websocket_transport_retryable: bool, - /// 兼容少数链路在 metadata 后、终态前关闭；仅在捕获 turn-state 时允许安全恢复。 - pub(super) websocket_checkpoint_retryable: bool, pub(super) account_failure: Option, /// 原始上游错误描述，仅在凭据错误状态下持久化。 pub(super) error_message: Option, @@ -56,7 +54,6 @@ impl MappedProviderFailure { Self { error, websocket_transport_retryable: false, - websocket_checkpoint_retryable: false, account_failure: None, error_message: None, cyber_policy_failure: false, @@ -454,8 +451,6 @@ pub(super) enum WebSocketFailurePolicy { pub(super) struct WebSocketRecoveryContext<'a> { pub(super) policy: WebSocketFailurePolicy, pub(super) requirement: TransportRequirement, - /// 本 attempt 已从 metadata 捕获可恢复同一请求的 turn-state 检查点。 - pub(super) replay_checkpoint_available: bool, pub(super) retry_count: u32, pub(super) max_retries: u32, pub(super) request_id: &'a str, @@ -464,7 +459,6 @@ pub(super) struct WebSocketRecoveryContext<'a> { pub(super) session_affinity_key: Option<&'a ProviderSessionAffinityKey>, pub(super) session_affinity_key_hash: Option<&'a str>, pub(super) session_transport_recovery: &'a CodexSessionTransportRecovery, - pub(super) session_recovery_probe: &'a mut Option, } pub(super) fn websocket_client_failure_policy( @@ -497,6 +491,9 @@ pub(super) fn websocket_client_failure_policy( if !matches!( error.classified(), CodexWebSocketExchangeError::InvalidRequest(_) + | CodexWebSocketExchangeError::ContinuationUnavailable { .. } + | CodexWebSocketExchangeError::OriginCircuitOpen + | CodexWebSocketExchangeError::OriginHalfOpenBusy ) => { Some(WebSocketFailurePolicy::Budgeted) @@ -509,25 +506,18 @@ pub(super) fn apply_websocket_recovery_policy( failure: &mut MappedProviderFailure, context: WebSocketRecoveryContext<'_>, ) { + let session_budget_exhausted = context.session_affinity_key.is_some_and(|key| { + context + .session_transport_recovery + .record_websocket_failure(key, context.max_retries) + }); let send_state = failure.error.send_state(); let post_send = matches!( send_state, UpstreamSendState::Sent | UpstreamSendState::Ambiguous ); - if post_send && !context.replay_checkpoint_available { - // payload 已经发送或发送结果不确定时，代理不能用同一业务输入自动重放。 - // 当前请求保持终态失败；恢复状态只裁决客户端主动发起的下一次请求。 - 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)); + if send_state == UpstreamSendState::Ambiguous || (post_send && !failure.error.replay_is_safe()) + { tracing::warn!( request_id = context.request_id, attempt_index = context.attempt_index, @@ -540,34 +530,27 @@ pub(super) fn apply_websocket_recovery_policy( upstream_send_state = ?send_state, transport_requirement = context.requirement.as_str(), continuation_recovery_action = "client_replay_required", + session_http_fallback = session_budget_exhausted, session_affinity_present = context.session_affinity_key.is_some(), session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""), - 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; } - let recovery_is_safe = context.requirement.allows_pre_send_http_fallback(); - if !recovery_is_safe { + if matches!( + context.requirement, + TransportRequirement::ExactWebSocketContinuation | TransportRequirement::HttpRequired + ) { return; } // 传输恢复必须保持原账号可调度；最终 HTTP attempt 若仍失败，再按真实 HTTP // 结果更新账号健康度，避免中间 WS 错误把同账号钉选提前冷却掉。 failure.account_failure = None; - if post_send { - // Core 对 Ambiguous 默认仍然 fail-closed；只有同一 attempt 的 session - // checkpoint 与这个 Provider proof 同时存在时，才允许同账号传输恢复。 - failure.error.set_replay_safe(); - } let fallback_now = context.policy == WebSocketFailurePolicy::ImmediateFallback - || context.retry_count >= context.max_retries; + || context.retry_count >= context.max_retries + || session_budget_exhausted; if !fallback_now { let Some(retry_index) = context.retry_count.checked_add(1).and_then(NonZeroU32::new) else { return; @@ -594,28 +577,18 @@ pub(super) fn apply_websocket_recovery_policy( 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(), - replay_checkpoint_available = context.replay_checkpoint_available, "OpenAI upstream WebSocket failed; retrying the same account" ); return; } + if !context.requirement.allows_pre_send_http_fallback() { + return; + } failure.error.set_pre_delivery_transport_fallback(); - 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)); + if let Some(key) = context.session_affinity_key { + context.session_transport_recovery.disable_websocket(key); + } tracing::warn!( request_id = context.request_id, attempt_index = context.attempt_index, @@ -633,16 +606,9 @@ pub(super) fn apply_websocket_recovery_policy( WebSocketFailurePolicy::Budgeted => "retry_budget_exhausted", WebSocketFailurePolicy::ImmediateFallback => "upgrade_required", }, - replay_checkpoint_available = context.replay_checkpoint_available, session_affinity_present = context.session_affinity_key.is_some(), session_affinity_key_hash = context.session_affinity_key_hash.unwrap_or(""), - 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" + "OpenAI upstream WebSocket disabled for this session" ); } @@ -686,31 +652,14 @@ pub(super) fn continuation_replay_required_error(reason: &'static str) -> Provid pub(super) 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 websocket_checkpoint_retryable = websocket_metadata_close_can_resume(&error); let mut failure = map_client_error(error, UpstreamSendState::Sent, false); failure.websocket_transport_retryable = allows_pre_delivery_retry && websocket_failure; - failure.websocket_checkpoint_retryable = - allows_pre_delivery_retry && websocket_checkpoint_retryable; if allows_pre_delivery_retry && !websocket_failure { failure.error = failure.error.with_pre_delivery_retry(); } failure } -fn websocket_metadata_close_can_resume(error: &CodexClientError) -> bool { - let CodexClientError::WebSocket(error) = error else { - return false; - }; - let Some(close) = error.close_before_terminal() else { - return false; - }; - close.code() == Some(1000) - && matches!( - close.last_event_type(), - Some("codex.response.metadata" | "response.metadata") - ) -} - pub(super) fn stream_transport_allows_pre_delivery_retry(error: &CodexClientError) -> bool { match error { CodexClientError::Http(_) @@ -865,9 +814,13 @@ pub(super) fn map_client_error( error.classified(), CodexWebSocketExchangeError::ConnectionLimitReached ) { - failure.error = failure.error.with_upstream_code(OpaqueUpstreamValue::new( - WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned(), - )); + failure.error = + failure + .error + .with_replay_safe() + .with_upstream_code(OpaqueUpstreamValue::new( + WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE.to_owned(), + )); } if let Some(client_visible_error) = client_visible_error { failure.error = failure @@ -963,9 +916,6 @@ fn websocket_diagnostic(error: &CodexWebSocketExchangeError) -> ProviderDiagnost CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => { "Reused OpenAI WebSocket died before the first upstream event".to_owned() } - CodexWebSocketExchangeError::InitialEventTimeout { timeout } => { - format!("OpenAI WebSocket first-event timeout after {timeout:?}") - } CodexWebSocketExchangeError::ClosedBeforeTerminal(_) => { unreachable!("close-before-terminal errors are handled before the variant match") } @@ -1098,7 +1048,6 @@ pub(super) fn map_upstream_failure( MappedProviderFailure { error, websocket_transport_retryable: false, - websocket_checkpoint_retryable: false, account_failure: account_failure( category, failure.retry_after_seconds, @@ -1217,8 +1166,9 @@ pub(super) fn websocket_send_state(error: &CodexWebSocketExchangeError) -> Upstr | CodexWebSocketExchangeError::SendTimeout { .. } | CodexWebSocketExchangeError::ClosedBeforeTerminal(_) | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } - | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } - | CodexWebSocketExchangeError::InitialEventTimeout { .. } => UpstreamSendState::Ambiguous, + | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => { + UpstreamSendState::Ambiguous + } CodexWebSocketExchangeError::ConnectionObserved { .. } => { unreachable!("classified websocket errors never retain observation wrappers") } @@ -1232,8 +1182,7 @@ pub(super) fn websocket_error_kind(error: &CodexWebSocketExchangeError) -> Provi | CodexWebSocketExchangeError::UnexpectedBinaryEvent => ProviderErrorKind::Protocol, CodexWebSocketExchangeError::ConnectTimeout { .. } | CodexWebSocketExchangeError::SendTimeout { .. } - | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } - | CodexWebSocketExchangeError::InitialEventTimeout { .. } => ProviderErrorKind::Timeout, + | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } => ProviderErrorKind::Timeout, CodexWebSocketExchangeError::OriginCircuitOpen | CodexWebSocketExchangeError::OriginHalfOpenBusy | CodexWebSocketExchangeError::SharedConnectFailed diff --git a/backend/crates/providers/openai/src/provider/mod.rs b/backend/crates/providers/openai/src/provider/mod.rs index 086609610..797fb4d60 100644 --- a/backend/crates/providers/openai/src/provider/mod.rs +++ b/backend/crates/providers/openai/src/provider/mod.rs @@ -61,10 +61,7 @@ use crate::credential::{ derive_codex_cyber_policy_session_key, derive_codex_session_affinity, derive_previous_response_id_hash, }; -use crate::session_transport::{ - CodexSessionRecoveryTransition, CodexSessionTransportDecision, CodexSessionTransportRecovery, - CodexSessionWebSocketProbe, SessionWebSocketFallback, -}; +use crate::session_transport::CodexSessionTransportRecovery; use crate::transport::canonical::{ CodexCanonicalDecoder, CodexCanonicalError, CodexCanonicalOutcome, }; @@ -94,7 +91,7 @@ use crate::transport::{ CodexBackendJsonResponse, CodexBackendStreamingResponse, CodexBackendTransport, CodexClientError, CodexRateLimitUpdates, CodexRequestContext, CodexResponseMetadata, CodexTransportMetrics, CodexTurnStateUpdate, CodexUpstreamDiagnostics, CodexWebSocketPool, - WebSocketConnectionPreference, endpoint_url, + endpoint_url, }; mod execution; @@ -321,14 +318,10 @@ impl Provider for CodexProvider { identity.prepare_local_conversation(&mut upstream_request); } if let Some(previous_session) = previous_session.as_ref() { - upstream_request.turn_state = if previous_session.retry_checkpoint - || same_client_turn( - previous_session.client_turn_id.as_deref(), - upstream_request.client_turn_id.as_deref(), - ) { - // 客户端可能携带 metadata-close 失败后才拿到的新状态；同一 - // turn 内显式回传值比最后一次成功响应保存的状态更新。同请求 - // 内部恢复由一次性 retry checkpoint 明确证明，不依赖客户端 turn ID。 + upstream_request.turn_state = if same_client_turn( + previous_session.client_turn_id.as_deref(), + upstream_request.client_turn_id.as_deref(), + ) { upstream_request .turn_state .take() @@ -453,60 +446,17 @@ impl Provider for CodexProvider { ); let requirement = transport_requirement(&upstream_request); let requested_transport = selected_transport(&upstream_request); - 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 session_http_fallback = requirement.allows_pre_send_http_fallback() + && session_affinity + .as_ref() + .is_some_and(|affinity| self.session_transport_recovery.uses_http(affinity.key())); let transport = if requirement.requires_websocket() { CodexProviderTransport::PreferWebSocket + } else if context.transport() == AttemptTransport::Fallback || session_http_fallback { + CodexProviderTransport::HttpOnly } else { - match context.transport() { - 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::Fallback => CodexProviderTransport::HttpOnly, - } + requested_transport }; - if let Some(session_transport_action) = session_transport_action { - 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(), - 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); let metadata = ProviderCallMetadata::new( provider_kind, @@ -553,9 +503,7 @@ impl Provider for CodexProvider { output_started_at, session_affinity_key, session_affinity_key_hash, - 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 05ee153e7..8a9d208f1 100644 --- a/backend/crates/providers/openai/src/session_transport.rs +++ b/backend/crates/providers/openai/src/session_transport.rs @@ -1,4 +1,4 @@ -//! OpenAI 根会话级上游传输恢复状态。 +//! 根会话的连续 WS 失败预算；耗尽后保持 HTTP，空闲过期后释放状态。 use std::collections::HashMap; use std::sync::{Arc, Mutex, MutexGuard}; @@ -6,453 +6,106 @@ use std::time::Duration; use gateway_core::provider_ports::ProviderSessionAffinityKey; use tokio::time::Instant; -use uuid::Uuid; -/// 未成功恢复的失败链需要跨过最长冷却，避免 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_RECOVERIES: usize = 16_384; +const SESSION_IDLE_RETENTION: Duration = Duration::from_secs(8 * 60 * 60); +const MAX_SESSION_TRANSPORTS: usize = 16_384; -#[derive(Clone)] +#[derive(Clone, Default)] pub(crate) struct CodexSessionTransportRecovery { - state: Arc>, - http_cooldowns: [Duration; 4], - retention: Duration, - capacity: usize, + sessions: Arc>>, } -#[derive(Default)] -struct SessionTransportRecoveryState { - sessions: HashMap, +struct SessionTransport { + state: TransportState, + last_used: Instant, } -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 { - state: Arc::new(Mutex::new(SessionTransportRecoveryState::default())), - http_cooldowns: HTTP_COOLDOWNS, - retention: SESSION_RECOVERY_RETENTION, - capacity: MAX_SESSION_RECOVERIES, - } - } -} - -#[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 { - recovery, - key, - probe_id, - armed: true, - } - } - - pub(crate) fn succeed(mut self) -> bool { - let recovered = self.recovery.record_probe_success(&self.key, self.probe_id); - self.armed = false; - recovered - } - - pub(crate) fn post_send_failed(mut self) -> CodexSessionRecoveryTransition { - let transition = self.recovery.record_post_send_failure(&self.key); - self.armed = false; - transition - } - - 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); - } - } +enum TransportState { + WebSocket { failures: u32 }, + Http, } 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 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; + pub(crate) fn uses_http(&self, key: &ProviderSessionAffinityKey) -> bool { + let now = Instant::now(); + let mut sessions = self.lock(); + let Some(session) = sessions.get_mut(key) else { + 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), - } - } + if now.saturating_duration_since(session.last_used) >= SESSION_IDLE_RETENTION { + sessions.remove(key); + return false; } + session.last_used = now; + matches!(session.state, TransportState::Http) } - fn record_post_send_failure_at( + /// 返回连续失败是否已耗尽首发加重试预算。 + pub(crate) fn record_websocket_failure( &self, key: &ProviderSessionAffinityKey, - now: Instant, - ) -> CodexSessionRecoveryTransition { - let mut state = self.lock(); - let failure_count = self.next_failure_count(state.sessions.get(key), now); - if !state.sessions.contains_key(key) { - self.make_room(&mut state, now); - } - 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 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 + max_retries: u32, + ) -> bool { + self.update(key, |state| match state { + TransportState::WebSocket { failures } if failures < max_retries => { + TransportState::WebSocket { + failures: failures + 1, + } } + TransportState::WebSocket { .. } | TransportState::Http => TransportState::Http, }) } - 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] + pub(crate) fn disable_websocket(&self, key: &ProviderSessionAffinityKey) { + self.update(key, |_| TransportState::Http); } - 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 + pub(crate) fn websocket_succeeded(&self, key: &ProviderSessionAffinityKey) { + self.lock().remove(key); } - 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; + fn update( + &self, + key: &ProviderSessionAffinityKey, + update: impl FnOnce(TransportState) -> TransportState, + ) -> bool { + let now = Instant::now(); + let mut sessions = self.lock(); + let previous = sessions + .remove(key) + .filter(|session| { + now.saturating_duration_since(session.last_used) < SESSION_IDLE_RETENTION + }) + .map_or(TransportState::WebSocket { failures: 0 }, |session| { + session.state + }); + if sessions.len() >= MAX_SESSION_TRANSPORTS { + sessions.retain(|_, session| { + now.saturating_duration_since(session.last_used) < SESSION_IDLE_RETENTION + }); + if sessions.len() >= MAX_SESSION_TRANSPORTS + && let Some(oldest) = sessions + .iter() + .min_by_key(|(_, session)| session.last_used) + .map(|(key, _)| key.clone()) + { + sessions.remove(&oldest); + } } - let failure_count = *failure_count; - let last_failure_at = *last_failure_at; - state.sessions.insert( + let state = update(previous); + let uses_http = matches!(state, TransportState::Http); + sessions.insert( key.clone(), - SessionCircuit::FreshPending { - failure_count, - last_failure_at, + SessionTransport { + state, + last_used: now, }, ); + uses_http } - 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 - .sessions - .iter() - .min_by_key(|(_, circuit)| circuit.last_failure_at()) - .map(|(key, _)| key.clone()); - if let Some(oldest) = oldest { - state.sessions.remove(&oldest); - } - } - - fn lock(&self) -> MutexGuard<'_, SessionTransportRecoveryState> { - self.state + fn lock(&self) -> MutexGuard<'_, HashMap> { + self.sessions .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 ceffb71bb..aea45187d 100644 --- a/backend/crates/providers/openai/src/transport/client.rs +++ b/backend/crates/providers/openai/src/transport/client.rs @@ -511,8 +511,6 @@ pub enum CodexTransportDecision { Http2WebSocketBudgetExhausted, Http2BreakerOpen, Http2PoolUnavailable, - Http2PreSendFailure, - Http2PreDeliveryFailure, } impl CodexTransportDecision { @@ -526,8 +524,6 @@ impl CodexTransportDecision { Self::Http2WebSocketBudgetExhausted => "http2_ws_budget_exhausted", Self::Http2BreakerOpen => "http2_breaker_open", Self::Http2PoolUnavailable => "http2_pool_unavailable", - Self::Http2PreSendFailure => "http2_ws_pre_send_failure", - Self::Http2PreDeliveryFailure => "http2_ws_pre_delivery_failure", } } } @@ -615,7 +611,6 @@ 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 { @@ -862,19 +857,18 @@ pub(super) fn websocket_success_decision( } } -pub(super) fn http_fallback_decision( +pub(super) fn local_http_fallback_decision( error: &CodexWebSocketExchangeError, -) -> CodexTransportDecision { +) -> Option { match error.classified() { CodexWebSocketExchangeError::OriginCircuitOpen | CodexWebSocketExchangeError::OriginHalfOpenBusy => { - CodexTransportDecision::Http2BreakerOpen + Some(CodexTransportDecision::Http2BreakerOpen) } - CodexWebSocketExchangeError::ContinuationUnavailable { .. } - | CodexWebSocketExchangeError::SharedConnectFailed => { - CodexTransportDecision::Http2PoolUnavailable + CodexWebSocketExchangeError::ContinuationUnavailable { .. } => { + Some(CodexTransportDecision::Http2PoolUnavailable) } - _ => CodexTransportDecision::Http2PreSendFailure, + _ => None, } } diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index 72f321ff0..7f20ad5c4 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -33,7 +33,7 @@ use crate::transport::{ response_meta, websocket::{ CodexWebSocketConnection, CodexWebSocketExchangeError, CodexWebSocketPool, - CodexWebSocketPoolKey, CodexWebSocketStreamingExchange, DEFAULT_INITIAL_EVENT_TIMEOUT, + CodexWebSocketPoolKey, CodexWebSocketStreamingExchange, DEFAULT_STREAM_IDLE_TIMEOUT, WEBSOCKET_FAST_PATH_BUDGET, WebSocketFastPath, WebSocketOriginBreaker, execute_prepared_response_create_request_stream, post_send_ambiguous, prepare_response_create_request_with_pool, websocket_audit_dir, @@ -43,16 +43,6 @@ 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( @@ -216,27 +206,6 @@ 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>, - connection_preference: WebSocketConnectionPreference, - ) -> CodexClientResult { - let prepared = self - .prepare_response_transport( - request, - context, - pool_account_id, - connection_preference, - 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( @@ -244,36 +213,18 @@ impl CodexBackendClient { request: &CodexResponsesRequest, context: CodexRequestContext<'_>, pool_account_id: Option<&str>, - ) -> CodexClientResult { - self.prepare_response_transport( - request, - context, - pool_account_id, - WebSocketConnectionPreference::ReuseOrConnect, - false, - ) - .await - } - - /// 与 [`Self::prepare_response_transport_with_pool_account`] 相同； - /// fresh 偏好会绕过连接池，为受控恢复探针强制新建连接。 - async fn prepare_response_transport( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - pool_account_id: Option<&str>, - connection_preference: WebSocketConnectionPreference, - defer_websocket_recovery: bool, ) -> CodexClientResult { let requirement = transport_requirement(request); - context.trace.cloned().unwrap_or_default().record("transport.preparing", serde_json::json!({ - "requirement": requirement.as_str(), "connectionPreference": format!("{connection_preference:?}"), - })); + context.trace.cloned().unwrap_or_default().record( + "transport.preparing", + serde_json::json!({ + "requirement": requirement.as_str(), + }), + ); if requirement == TransportRequirement::HttpRequired { return Ok(PreparedResponseTransport { requirement, route: PreparedResponseRoute::Http, - defer_websocket_recovery, metrics: CodexTransportMetrics { decision: Some(CodexTransportDecision::HttpRequired), ..CodexTransportMetrics::default() @@ -315,12 +266,7 @@ 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 = match connection_preference { - WebSocketConnectionPreference::ReuseOrConnect => { - self.websocket_pool.as_deref().zip(pool_key) - } - WebSocketConnectionPreference::Fresh => None, - }; + let pool = self.websocket_pool.as_deref().zip(pool_key); let fast_path_budget = match requirement { TransportRequirement::PersistedContinuation | TransportRequirement::NewChain => { Some(WEBSOCKET_FAST_PATH_BUDGET) @@ -338,7 +284,7 @@ impl CodexBackendClient { &self.websocket_origin_key, fast_path_budget, requirement.requires_websocket(), - Some(DEFAULT_INITIAL_EVENT_TIMEOUT), + Some(DEFAULT_STREAM_IDLE_TIMEOUT), ) .await; let prepared = match prepared { @@ -363,15 +309,12 @@ impl CodexBackendClient { transport_decision_wait_ms: Some(wait_ms), ..CodexTransportMetrics::default() }, - defer_websocket_recovery, }); } Err(error) - if (!defer_websocket_recovery || error.requires_immediate_pool_fallback()) - && requirement.allows_pre_send_http_fallback() - && error.allows_pre_send_http_fallback() => + if requirement.allows_pre_send_http_fallback() + && let Some(decision) = local_http_fallback_decision(&error) => { - let decision = http_fallback_decision(&error); let wait_ms = elapsed_duration_millis(prepare_started_at.elapsed()); context.trace.cloned().unwrap_or_default().record( "transport.fallback", serde_json::json!({ @@ -390,7 +333,6 @@ 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)), @@ -428,7 +370,6 @@ impl CodexBackendClient { prepared, })), metrics, - defer_websocket_recovery, }) } @@ -443,7 +384,6 @@ impl CodexBackendClient { requirement, route, metrics, - defer_websocket_recovery, } = prepared; context.trace.cloned().unwrap_or_default().record( "transport.selected", @@ -465,8 +405,7 @@ impl CodexBackendClient { request: websocket_request, prepared, } = *route; - let delivery_wait_started_at = Instant::now(); - let mut exchange = match execute_prepared_response_create_request_stream( + let mut exchange = execute_prepared_response_create_request_stream( &websocket_request, prepared, context @@ -476,161 +415,19 @@ impl CodexBackendClient { .exchange("websocket"), ) .await - { - Ok(exchange) => exchange, - Err(error) - if !defer_websocket_recovery - && requirement.allows_pre_delivery_http_fallback() - && 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) => return Err(websocket_exchange_error_to_client_error(error)), - }; - 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) - if defer_websocket_recovery => - { - return Err(CodexClientError::WebSocket( - CodexWebSocketExchangeError::ConnectionLimitReached, - )); - } - Ok(DeliveryBoundary::ConnectionLimitReached) => { - if fresh_retry_used { - return self - .http_fallback_before_delivery( - request, - context, - requirement, - metrics, - delivery_wait_started_at, - "websocket connection limit reached", - ) - .await; - } - fresh_retry_used = true; - context.trace.cloned().unwrap_or_default().record( - "transport.retry", - serde_json::json!({ - "reason": "websocket_connection_limit_reached", - "connectionPreference": "fresh", - }), - ); - drop(exchange); - match self - .prepare_response_transport( - request, - context, - None, - WebSocketConnectionPreference::Fresh, - false, - ) - .await - { - Ok(PreparedResponseTransport { - route: PreparedResponseRoute::WebSocket(route), - .. - }) => { - match execute_prepared_response_create_request_stream( - &route.request, - route.prepared, - context - .trace - .cloned() - .unwrap_or_default() - .exchange("websocket"), - ) - .await - { - Ok(next_exchange) => { - exchange = next_exchange; - continue; - } - 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; - } - Err(error) => { - return Err( - websocket_exchange_error_to_client_error( - post_send_ambiguous(error), - ), - ); - } - } - } - Ok(PreparedResponseTransport { - route: PreparedResponseRoute::Http, - .. - }) => { - return self - .create_response_stream_http_sse(request, context) - .await - .map(|mut response| { - merge_preparation_metrics( - &mut response.transport_metrics, - metrics, - ); - response - }); - } - Err(error) => { - return self - .http_fallback_before_delivery( - request, - context, - requirement, - metrics, - delivery_wait_started_at, - &error.to_string(), - ) - .await; - } - } - } - 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) => { - return Err(websocket_exchange_error_to_client_error( - post_send_ambiguous(error), - )); - } + .map_err(websocket_exchange_error_to_client_error)?; + if requirement.allows_connection_restart() { + match await_websocket_delivery_boundary(&mut exchange).await { + Ok(DeliveryBoundary::Ready) => {} + Ok(DeliveryBoundary::ConnectionLimitReached) => { + return Err(CodexClientError::WebSocket( + CodexWebSocketExchangeError::ConnectionLimitReached, + )); + } + Err(error) => { + return Err(websocket_exchange_error_to_client_error( + post_send_ambiguous(error), + )); } } } @@ -658,59 +455,6 @@ impl CodexBackendClient { } } - async fn fallback_to_http_before_websocket_delivery( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - requirement: TransportRequirement, - metrics: CodexTransportMetrics, - delivery_wait_started_at: Instant, - error: CodexWebSocketExchangeError, - ) -> CodexClientResult { - self.http_fallback_before_delivery( - request, - context, - requirement, - metrics, - delivery_wait_started_at, - &error.to_string(), - ) - .await - } - - async fn http_fallback_before_delivery( - &self, - request: &CodexResponsesRequest, - context: CodexRequestContext<'_>, - requirement: TransportRequirement, - mut metrics: CodexTransportMetrics, - delivery_wait_started_at: Instant, - detail: &str, - ) -> CodexClientResult { - let decision_wait_ms = metrics - .transport_decision_wait_ms - .unwrap_or_default() - .saturating_add(elapsed_duration_millis(delivery_wait_started_at.elapsed())); - metrics.decision = Some(CodexTransportDecision::Http2PreDeliveryFailure); - metrics.transport_decision_wait_ms = Some(decision_wait_ms); - context.trace.cloned().unwrap_or_default().record( - "transport.fallback", - serde_json::json!({ - "to": "http_sse", "reason": "websocket_pre_delivery_failure", - "requirement": requirement.as_str(), - "decision": CodexTransportDecision::Http2PreDeliveryFailure.as_str(), - "waitMs": decision_wait_ms, - "detail": diagnostic_json(&serde_json::json!({"message": detail})), - }), - ); - self.create_response_stream_http_sse(request, context) - .await - .map(|mut response| { - merge_preparation_metrics(&mut response.transport_metrics, metrics); - response - }) - } - fn websocket_pool_key( &self, request: &CodexResponsesRequest, diff --git a/backend/crates/providers/openai/src/transport/diagnostics.rs b/backend/crates/providers/openai/src/transport/diagnostics.rs index 96fc2e6c2..7ed350f66 100644 --- a/backend/crates/providers/openai/src/transport/diagnostics.rs +++ b/backend/crates/providers/openai/src/transport/diagnostics.rs @@ -1,10 +1,7 @@ //! 在 transport 边界采集的上游响应诊断信息。 use super::client::CodexClientVisibleUpstreamResponse; -use super::protocol::responses::{ - PREVIOUS_RESPONSE_NOT_FOUND_CODE, ResponsesSseFailure, - is_bare_invalid_previous_response_id_error, -}; +use super::protocol::responses::ResponsesSseFailure; use base64::{Engine as _, engine::general_purpose::STANDARD}; use chrono::{DateTime, Utc}; use reqwest::{StatusCode, header::HeaderMap}; @@ -146,7 +143,7 @@ impl CodexUpstreamFailure { rate_limit_headers: &[(String, String)], send_phase: CodexUpstreamSendPhase, ) -> Self { - let fields = ParsedUpstreamError::from_http_response(status, body); + let fields = ParsedUpstreamError::from_http_response(body); let category = classify_upstream_failure( UpstreamFailureSource::HttpResponse, Some(status), @@ -268,7 +265,7 @@ struct ParsedUpstreamError { } impl ParsedUpstreamError { - fn from_http_response(status: StatusCode, body: &str) -> Self { + fn from_http_response(body: &str) -> Self { let Ok(value) = serde_json::from_str::(body) else { return Self { code: None, @@ -294,16 +291,7 @@ impl ParsedUpstreamError { .or_else(|| value.get("message")) .and_then(Value::as_str) .and_then(non_empty_owned); - let code = if status == StatusCode::BAD_REQUEST - && is_bare_invalid_previous_response_id_error( - code_value, - error_type.as_deref(), - client_message.as_deref(), - ) { - Some(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()) - } else { - code_value.and_then(Value::as_str).and_then(non_empty_owned) - }; + let code = code_value.and_then(Value::as_str).and_then(non_empty_owned); let message = client_message .clone() .or_else(|| error.as_str().and_then(non_empty_owned)) diff --git a/backend/crates/providers/openai/src/transport/mod.rs b/backend/crates/providers/openai/src/transport/mod.rs index 1b9cd0c69..a0f3efa38 100644 --- a/backend/crates/providers/openai/src/transport/mod.rs +++ b/backend/crates/providers/openai/src/transport/mod.rs @@ -21,7 +21,6 @@ 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 18036033a..6abb4caf7 100644 --- a/backend/crates/providers/openai/src/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/src/transport/protocol/responses.rs @@ -15,19 +15,6 @@ pub(crate) const X_CODEX_TURN_STATE_CLIENT_METADATA_KEY: &str = "x-codex-turn-st /// 本地生成 history unavailable 错误时使用的官方提示文本。 pub(crate) const PREVIOUS_RESPONSE_NOT_FOUND_MESSAGE: &str = "Previous response was not found. Retrying the full request."; -const INVALID_PREVIOUS_RESPONSE_ID_MESSAGE: &str = "Invalid `previous_response_id`."; - -/// 精确匹配上游省略错误码的旧版 previous response 拒绝形状。 -pub(crate) fn is_bare_invalid_previous_response_id_error( - code: Option<&Value>, - error_type: Option<&str>, - message: Option<&str>, -) -> bool { - code.is_none_or(Value::is_null) - && error_type == Some("invalid_request_error") - && message == Some(INVALID_PREVIOUS_RESPONSE_ID_MESSAGE) -} - /// Codex Responses 上游请求体。 /// /// 发往上游的 Responses 请求。`body` 持有客户端原始 JSON object，逐字段（含顺序、 @@ -161,13 +148,8 @@ impl TransportRequirement { ) } - /// WebSocket 已发送 payload、但尚未交付业务事件时，是否允许同账号 HTTP 重放。 - pub fn allows_pre_delivery_http_fallback(self) -> bool { - matches!(self, Self::NewChain) - } - - /// 会话级恢复状态是否可以为请求选择 fresh WebSocket 或临时 HTTP。 - pub fn allows_session_transport_recovery(self) -> bool { + /// 无续接依赖的新请求可以在明确的连接级拒绝后重新建连。 + pub fn allows_connection_restart(self) -> bool { matches!(self, Self::NewChain) } diff --git a/backend/crates/providers/openai/src/transport/protocol/websocket.rs b/backend/crates/providers/openai/src/transport/protocol/websocket.rs index 7e0e8daae..703b81876 100644 --- a/backend/crates/providers/openai/src/transport/protocol/websocket.rs +++ b/backend/crates/providers/openai/src/transport/protocol/websocket.rs @@ -2,10 +2,7 @@ use gateway_protocol::openai::sse::encode_sse_event; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; -use crate::transport::protocol::responses::{ - CodexResponsesRequest, PREVIOUS_RESPONSE_NOT_FOUND_CODE, - is_bare_invalid_previous_response_id_error, transport_requirement, -}; +use crate::transport::protocol::responses::{CodexResponsesRequest, transport_requirement}; const REDACTED_PAYLOAD_VALUE: &str = ""; @@ -57,49 +54,13 @@ pub fn websocket_event_to_sse_frame(raw: &str) -> Option { websocket_event_frame(&value, raw) } -/// 同上，复用已解析的事件 JSON；`data` 段通常逐字节内嵌 `raw` 原文。 -/// -/// 仅剥离传输层内部帧：`codex.rate_limits` 与承载 turn_state 的 -/// `response.metadata` 不下发客户端。唯一的业务帧适配是为上游省略错误码的 -/// `Invalid previous_response_id` 补齐官方客户端用于完整历史重放的稳定错误码； -/// 原始 status、type 与 message 均保持不变。 +/// 复用已解析的 JSON，剥离内部帧并原样转发公开事件。 pub(crate) fn websocket_event_frame(value: &Value, raw: &str) -> Option { let event = websocket_event_type(value)?; if is_internal_websocket_event(event) || event == "response.metadata" { return None; } - let normalized = previous_response_not_found_event(value); - Some(encode_sse_event( - event, - normalized.as_deref().unwrap_or(raw), - )) -} - -fn previous_response_not_found_event(value: &Value) -> Option { - if websocket_event_type(value) != Some("error") - || value - .get("status") - .or_else(|| value.get("status_code")) - .and_then(Value::as_u64) - != Some(400) - { - return None; - } - let error = value.get("error")?.as_object()?; - if !is_bare_invalid_previous_response_id_error( - error.get("code"), - error.get("type").and_then(Value::as_str), - error.get("message").and_then(Value::as_str), - ) { - return None; - } - - let mut normalized = value.clone(); - normalized.get_mut("error")?.as_object_mut()?.insert( - "code".to_owned(), - Value::String(PREVIOUS_RESPONSE_NOT_FOUND_CODE.to_owned()), - ); - serde_json::to_string(&normalized).ok() + Some(encode_sse_event(event, raw)) } pub(crate) fn websocket_event_type(value: &Value) -> Option<&str> { diff --git a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs index 8dd6ab7a0..d8e4d67f8 100644 --- a/backend/crates/providers/openai/src/transport/websocket/coordinator.rs +++ b/backend/crates/providers/openai/src/transport/websocket/coordinator.rs @@ -1,8 +1,14 @@ //! Responses WebSocket pre-send/post-send 与 pool/breaker 编排。 -use std::time::{Duration, Instant}; +use std::{ + future::Future, + time::{Duration, Instant}, +}; -use tokio::{sync::oneshot, time::timeout}; +use tokio::{ + sync::oneshot, + time::{Instant as TokioInstant, timeout_at}, +}; use super::{ breaker::{ @@ -43,7 +49,7 @@ pub(crate) struct PreparedWebSocket { binding: PoolBinding, connect_elapsed: Option, decision_wait_elapsed: Duration, - initial_event_timeout: Option, + stream_idle_timeout: Option, } enum PoolBinding { @@ -111,7 +117,7 @@ pub(crate) async fn prepare_response_create_request_with_pool( origin_key: &str, fast_path_budget: Option, require_pool: bool, - fallback_initial_event_timeout: Option, + fallback_stream_idle_timeout: Option, ) -> Result, CodexWebSocketExchangeError> { let decision_started_at = Instant::now(); let Some((pool, key)) = pool else { @@ -125,7 +131,7 @@ pub(crate) async fn prepare_response_create_request_with_pool( breaker, origin_key, fast_path_budget, - fallback_initial_event_timeout, + fallback_stream_idle_timeout, decision_started_at, ) .await; @@ -159,7 +165,7 @@ pub(crate) async fn prepare_response_create_request_with_pool( }, connect_elapsed: None, decision_wait_elapsed: decision_started_at.elapsed(), - initial_event_timeout: pool.initial_event_timeout(), + stream_idle_timeout: pool.stream_idle_timeout(), })) } WebSocketPoolAcquire::Connect(connect_lease) => { @@ -209,10 +215,10 @@ async fn prepare_unpooled_websocket( breaker: &WebSocketOriginBreaker, origin_key: &str, fast_path_budget: Option, - initial_event_timeout: Option, + stream_idle_timeout: Option, decision_started_at: Instant, ) -> Result, CodexWebSocketExchangeError> { - let permit = acquire_breaker_permit(breaker, origin_key)?; + let permit = acquire_breaker_permit(breaker, origin_key, fast_path_budget.is_some())?; let context = pump_log_context_from_connection(request.connection()); let connected = connect_with_budget( request.connection(), @@ -224,7 +230,9 @@ async fn prepare_unpooled_websocket( let connected = match connected { Ok(WebSocketFastPath::Ready(connection)) => Ok(connection), Ok(WebSocketFastPath::Missed) => { - permit.fast_timeout(); + if let Some(permit) = permit { + permit.fast_timeout(); + } return Ok(WebSocketFastPath::Missed); } Err(error) => Err(error), @@ -235,7 +243,7 @@ async fn prepare_unpooled_websocket( binding: PoolBinding::Unpooled, connect_elapsed: Some(connect_elapsed), decision_wait_elapsed: decision_started_at.elapsed(), - initial_event_timeout, + stream_idle_timeout, })) } @@ -248,7 +256,7 @@ async fn prepare_pooled_websocket( fast_path_budget: Option, decision_started_at: Instant, ) -> Result, CodexWebSocketExchangeError> { - let permit = match acquire_breaker_permit(breaker, origin_key) { + let permit = match acquire_breaker_permit(breaker, origin_key, fast_path_budget.is_some()) { Ok(permit) => permit, Err(error) => { connect_lease.failed().await; @@ -271,7 +279,7 @@ async fn prepare_pooled_websocket( }, connect_elapsed: Some(handoff.connect_elapsed), decision_wait_elapsed: decision_started_at.elapsed(), - initial_event_timeout: pool.initial_event_timeout(), + stream_idle_timeout: pool.stream_idle_timeout(), })), WebSocketFastPath::Missed => Ok(WebSocketFastPath::Missed), } @@ -280,7 +288,7 @@ async fn prepare_pooled_websocket( struct PooledWebSocketConnectWaiter { started_at: tokio::time::Instant, receiver: oneshot::Receiver>, - fast_path_reporter: WebSocketOriginFastPathReporter, + fast_path_reporter: Option, } struct PooledWebSocketHandoff { @@ -294,19 +302,16 @@ impl PooledWebSocketConnectWaiter { &mut self, fast_path_budget: Option, ) -> Result, CodexWebSocketExchangeError> { - let received = match fast_path_budget { - Some(budget) => { - let remaining = budget.saturating_sub(self.started_at.elapsed()); - match timeout(remaining, &mut self.receiver).await { - Ok(received) => received, - Err(_) => { - self.fast_path_reporter.missed(); - return Ok(WebSocketFastPath::Missed); + let received = + match wait_for_fast_path(self.started_at, fast_path_budget, &mut self.receiver).await { + WebSocketFastPath::Ready(received) => received, + WebSocketFastPath::Missed => { + if let Some(reporter) = &self.fast_path_reporter { + reporter.missed(); } + return Ok(WebSocketFastPath::Missed); } - } - None => (&mut self.receiver).await, - }; + }; received .map_err(|_| CodexWebSocketExchangeError::SharedConnectFailed)? .map(WebSocketFastPath::Ready) @@ -317,12 +322,14 @@ fn start_pooled_websocket_connect( connection: CodexWebSocketConnection, pool: CodexWebSocketPool, connect_lease: WebSocketPoolConnectLease, - permit: WebSocketOriginBreakerPermit, + permit: Option, ) -> PooledWebSocketConnectWaiter { let task_key = connect_lease.key().clone(); let started_at = connect_lease.started_at(); let cancellation = connect_lease.cancellation_token(); - let fast_path_reporter = permit.fast_path_reporter(); + let fast_path_reporter = permit + .as_ref() + .map(WebSocketOriginBreakerPermit::fast_path_reporter); let keepalive = pool.keepalive(); let context = PumpLogContext::new( Some(task_key.account_id().to_owned()), @@ -333,7 +340,7 @@ fn start_pooled_websocket_connect( let connected = tokio::select! { biased; _ = cancellation.cancelled() => { - permit.cancel(); + if let Some(permit) = permit { permit.cancel(); } let _ = sender.send(Err(CodexWebSocketExchangeError::SharedConnectFailed)); connect_lease.failed().await; tracing::info!( @@ -450,21 +457,15 @@ async fn connect_with_budget( fast_path_budget: Option, context: PumpLogContext, ) -> Result, CodexWebSocketExchangeError> { - match fast_path_budget { - Some(budget) => { - let Ok(connected) = timeout( - budget, - connect_websocket_connection(connection, keepalive, context), - ) - .await - else { - return Ok(WebSocketFastPath::Missed); - }; - connected.map(WebSocketFastPath::Ready) - } - None => connect_websocket_connection(connection, keepalive, context) - .await - .map(WebSocketFastPath::Ready), + match wait_for_fast_path( + TokioInstant::now(), + fast_path_budget, + connect_websocket_connection(connection, keepalive, context), + ) + .await + { + WebSocketFastPath::Ready(connected) => connected.map(WebSocketFastPath::Ready), + WebSocketFastPath::Missed => Ok(WebSocketFastPath::Missed), } } @@ -489,9 +490,13 @@ async fn connect_websocket_connection( fn acquire_breaker_permit( breaker: &WebSocketOriginBreaker, origin_key: &str, -) -> Result { + allow_http_fallback: bool, +) -> Result, CodexWebSocketExchangeError> { + if !allow_http_fallback { + return Ok(None); + } match breaker.try_acquire(origin_key) { - WebSocketOriginBreakerDecision::Allowed(permit) => Ok(permit), + WebSocketOriginBreakerDecision::Allowed(permit) => Ok(Some(permit)), WebSocketOriginBreakerDecision::Open => Err(CodexWebSocketExchangeError::OriginCircuitOpen), WebSocketOriginBreakerDecision::HalfOpenBusy => { Err(CodexWebSocketExchangeError::OriginHalfOpenBusy) @@ -500,9 +505,12 @@ fn acquire_breaker_permit( } fn finish_breaker_attempt( - permit: WebSocketOriginBreakerPermit, + permit: Option, connected: Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError>, ) -> Result<(PooledWebSocketConnection, Duration), CodexWebSocketExchangeError> { + let Some(permit) = permit else { + return connected; + }; match connected { Ok(connection) => { permit.succeed(); @@ -524,15 +532,21 @@ async fn wait_for_shared_connect( waiter: pool::WebSocketPoolConnectWaiter, fast_path_budget: Option, ) -> WebSocketFastPath { - match fast_path_budget { - Some(budget) => { - let remaining = waiter.remaining_budget(budget); - match timeout(remaining, waiter.wait()).await { - Ok(outcome) => WebSocketFastPath::Ready(outcome), - Err(_) => WebSocketFastPath::Missed, - } - } - None => WebSocketFastPath::Ready(waiter.wait().await), + wait_for_fast_path(waiter.started_at(), fast_path_budget, waiter.wait()).await +} + +// 同一次 opening 的所有等待者共享截止时间；后来的请求不会重置 800ms 预算。 +async fn wait_for_fast_path( + started_at: TokioInstant, + budget: Option, + future: F, +) -> WebSocketFastPath { + let Some(budget) = budget else { + return WebSocketFastPath::Ready(future.await); + }; + match timeout_at(started_at + budget, future).await { + Ok(value) => WebSocketFastPath::Ready(value), + Err(_) => WebSocketFastPath::Missed, } } @@ -564,7 +578,7 @@ pub(crate) async fn execute_prepared_response_create_request_stream( let PreparedWebSocket { connection, binding, - initial_event_timeout, + stream_idle_timeout, .. } = prepared; let (lease, reused, pool_decision) = binding.into_parts(); @@ -607,7 +621,7 @@ pub(crate) async fn execute_prepared_response_create_request_stream( metadata, pool_return, reused, - initial_event_timeout, + stream_idle_timeout, trace, ); exchange.pool_decision = pool_decision; diff --git a/backend/crates/providers/openai/src/transport/websocket/error.rs b/backend/crates/providers/openai/src/transport/websocket/error.rs index 78cc60f1d..f73da7213 100644 --- a/backend/crates/providers/openai/src/transport/websocket/error.rs +++ b/backend/crates/providers/openai/src/transport/websocket/error.rs @@ -90,12 +90,6 @@ pub enum CodexWebSocketExchangeError { #[source] source: Option>, }, - /// 建连并发送后，上游在配置时间内没有产生任何事件。 - #[error("websocket first upstream event not received within {timeout:?}")] - InitialEventTimeout { - /// 首个上游事件超时时长。 - timeout: Duration, - }, /// 将一个已分类交互错误与物理连接生命周期快照绑定。 #[error("{source}")] ConnectionObserved { @@ -325,38 +319,4 @@ impl CodexWebSocketExchangeError { send_phase, })) } - - /// opening 阶段只有明确的 transport 可用性失败才能切到同账号 HTTP。 - pub(in crate::transport) fn allows_pre_send_http_fallback(&self) -> bool { - if let Self::ConnectionObserved { source, .. } = self { - return source.allows_pre_send_http_fallback(); - } - matches!( - self, - Self::Connect(_) - | Self::ConnectTimeout { .. } - | Self::OriginCircuitOpen - | Self::OriginHalfOpenBusy - | Self::SharedConnectFailed - | Self::ContinuationUnavailable { .. } - ) - } - - /// 精确连接状态不可用属于本地池路由事实，不消耗上游 WS 重试预算。 - pub(in crate::transport) fn requires_immediate_pool_fallback(&self) -> bool { - matches!(self.classified(), Self::ContinuationUnavailable { .. }) - } - - /// 首个业务事件交付前，普通 WebSocket 失败可切到同账号 HTTP。 - pub(in crate::transport) fn allows_pre_delivery_http_fallback(&self) -> bool { - match self { - Self::Upstream(_) => false, - Self::PostSendAmbiguous { - source: Some(source), - .. - } => source.allows_pre_delivery_http_fallback(), - Self::ConnectionObserved { source, .. } => source.allows_pre_delivery_http_fallback(), - _ => true, - } - } } diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/io.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/io.rs index 8f8d1f849..02ef44f96 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/io.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/io.rs @@ -6,23 +6,7 @@ use tokio::time::timeout; use tungstenite::Message; use super::super::pump::PumpedWebSocket; -use super::{ - CodexWebSocketExchangeError, WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT, - WEBSOCKET_RECEIVE_IDLE_TIMEOUT, -}; - -pub(super) fn receive_idle_timeout( - saw_upstream_activity: bool, - initial_event_timeout: Option, -) -> Duration { - if saw_upstream_activity { - WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT - } else { - initial_event_timeout - .filter(|timeout| !timeout.is_zero()) - .unwrap_or(WEBSOCKET_RECEIVE_IDLE_TIMEOUT) - } -} +use super::CodexWebSocketExchangeError; pub(super) async fn next_websocket_message( websocket: &mut PumpedWebSocket, diff --git a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs index 20dd4e255..c37fa7f9a 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/mod.rs @@ -4,7 +4,7 @@ mod io; mod reducer; mod stream; -use std::{pin::Pin, sync::Arc, time::Duration}; +use std::{pin::Pin, sync::Arc}; use bytes::Bytes; use futures::Stream; @@ -19,8 +19,6 @@ use crate::transport::{ use super::error::CodexWebSocketExchangeError; use super::pool::{CodexWebSocketConnectionMetadata, WebSocketPoolDecision}; -const WEBSOCKET_RECEIVE_IDLE_TIMEOUT: Duration = Duration::from_secs(20); -const WEBSOCKET_ACTIVE_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(5 * 60); const WEBSOCKET_STREAM_BUFFER: usize = 16; pub(super) use self::stream::{WebSocketStreamPoolReturn, stream_websocket_response}; 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 7098e3c7d..a167c4271 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -17,7 +17,7 @@ use super::super::{ PumpExitReason, PumpedWebSocket, WebSocketConnectionObservation, transport_metric_reason, }, }; -use super::io::{next_websocket_message, receive_idle_timeout, reused_stream_receive_error}; +use super::io::{next_websocket_message, reused_stream_receive_error}; use super::reducer::{ExchangeAction, WebSocketTerminalKind, reduce_websocket_event}; use super::{ CodexWebSocketExchangeError, CodexWebSocketRateLimitUpdates, CodexWebSocketStreamingExchange, @@ -62,7 +62,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response( metadata: CodexWebSocketConnectionMetadata, pool_return: Option, reused_connection: bool, - initial_event_timeout: Option, + stream_idle_timeout: Option, trace: TraceContext, ) -> CodexWebSocketStreamingExchange { let websocket_connection_id = websocket.connection_id(); @@ -85,7 +85,7 @@ pub(in crate::transport::websocket) fn stream_websocket_response( metadata, pool_return, reused_connection, - initial_event_timeout, + stream_idle_timeout, trace, shutdown, rate_limit_updates: rate_limit_updates_for_task, @@ -125,7 +125,7 @@ struct WebSocketStreamForwardState { metadata: CodexWebSocketConnectionMetadata, pool_return: Option, reused_connection: bool, - initial_event_timeout: Option, + stream_idle_timeout: Option, shutdown: CancellationToken, rate_limit_updates: CodexWebSocketRateLimitUpdates, turn_state_update: CodexWebSocketTurnStateUpdate, @@ -139,7 +139,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { mut metadata, pool_return, reused_connection, - initial_event_timeout, + stream_idle_timeout, shutdown, rate_limit_updates, turn_state_update, @@ -150,7 +150,6 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { .as_mut() .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! { @@ -177,20 +176,12 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { } message = next_websocket_message( &mut websocket, - receive_idle_timeout(saw_upstream_activity, initial_event_timeout), + stream_idle_timeout.filter(|timeout| !timeout.is_zero()).unwrap_or(super::super::pool::DEFAULT_STREAM_IDLE_TIMEOUT), ) => message, }; let message = match message { Ok(message) => message, Err(error) => { - let error = match error { - CodexWebSocketExchangeError::ReceiveIdleTimeout { timeout } - if !saw_upstream_activity => - { - CodexWebSocketExchangeError::InitialEventTimeout { timeout } - } - error => error, - }; trace.record( "upstream.read.failed", json!({ @@ -223,10 +214,7 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { break; }; let raw = match message { - tungstenite::Message::Text(text) => { - saw_upstream_activity = true; - text.to_string() - } + tungstenite::Message::Text(text) => text.to_string(), tungstenite::Message::Binary(bytes) => { trace.capture("upstream.binary", &bytes); let error = CodexWebSocketExchangeError::UnexpectedBinaryEvent; @@ -475,7 +463,6 @@ fn exchange_exit_reason(error: &CodexWebSocketExchangeError) -> &'static str { } } CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } => "receive_idle_timeout", - CodexWebSocketExchangeError::InitialEventTimeout { .. } => "initial_event_timeout", CodexWebSocketExchangeError::UnexpectedBinaryEvent => "unexpected_binary_event", _ => "exchange_failure", } diff --git a/backend/crates/providers/openai/src/transport/websocket/mod.rs b/backend/crates/providers/openai/src/transport/websocket/mod.rs index bb7e96e5c..b4f806f3b 100644 --- a/backend/crates/providers/openai/src/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/mod.rs @@ -41,5 +41,5 @@ pub(crate) use self::{ execute_prepared_response_create_request_stream, post_send_ambiguous, prepare_response_create_request_with_pool, }, - pool::DEFAULT_INITIAL_EVENT_TIMEOUT, + pool::DEFAULT_STREAM_IDLE_TIMEOUT, }; diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/lease.rs b/backend/crates/providers/openai/src/transport/websocket/pool/lease.rs index 203b9547a..cd26d8c04 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/lease.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/lease.rs @@ -1,7 +1,5 @@ //! WebSocket 连接池 reservation、handoff 与 lease 生命周期。 -use std::time::Duration; - use tokio::{ sync::{OwnedSemaphorePermit, watch}, time::Instant, @@ -40,8 +38,8 @@ pub(crate) struct WebSocketPoolConnectWaiter { } impl WebSocketPoolConnectWaiter { - pub(crate) fn remaining_budget(&self, budget: Duration) -> Duration { - budget.saturating_sub(self.started_at.elapsed()) + pub(crate) fn started_at(&self) -> Instant { + self.started_at } pub(crate) async fn wait(mut self) -> WebSocketPoolConnectOutcome { diff --git a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs index fcc931f2a..2e4e8e16d 100644 --- a/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs +++ b/backend/crates/providers/openai/src/transport/websocket/pool/mod.rs @@ -37,7 +37,7 @@ const DEFAULT_MAX_AGE: Duration = Duration::from_mins(55); const DEFAULT_MAINTENANCE_INTERVAL: Duration = Duration::from_secs(25); const DEFAULT_PING_INTERVAL: Duration = Duration::from_secs(25); const DEFAULT_PING_TIMEOUT: Duration = Duration::from_secs(5); -pub(crate) const DEFAULT_INITIAL_EVENT_TIMEOUT: Duration = Duration::from_secs(20); +pub(crate) const DEFAULT_STREAM_IDLE_TIMEOUT: Duration = Duration::from_secs(300); /// WebSocket 连接池。 #[derive(Clone)] @@ -73,8 +73,8 @@ pub struct CodexWebSocketPoolConfig { pub ping_timeout: Duration, /// idle socket 无活动多久后视为失活。 pub liveness_timeout: Option, - /// 建连并发送后首个上游事件到达前的超时；`None` 表示禁用。 - pub initial_event_timeout: Option, + /// 等待下一条上游消息的空闲超时；`None` 或零值使用默认 300 秒。 + pub stream_idle_timeout: Option, } impl Default for CodexWebSocketPoolConfig { @@ -87,9 +87,9 @@ impl Default for CodexWebSocketPoolConfig { ping_interval: Some(DEFAULT_PING_INTERVAL), ping_timeout: DEFAULT_PING_TIMEOUT, // idle 连接不设失活截断：靠 ping/pong 保活，只在 max_age（55 分钟） - // 或 ping 失败时关闭，最大化跨轮复用（对齐 Codex CLI 的长连接策略）。 + // 或 ping 失败时关闭，维持跨轮可复用连接。 liveness_timeout: None, - initial_event_timeout: Some(DEFAULT_INITIAL_EVENT_TIMEOUT), + stream_idle_timeout: Some(DEFAULT_STREAM_IDLE_TIMEOUT), } } } @@ -136,9 +136,9 @@ impl CodexWebSocketPool { self.config.keepalive() } - /// 建连并发送后首个上游事件到达前的超时；`None` 表示禁用。 - pub(crate) fn initial_event_timeout(&self) -> Option { - self.config.initial_event_timeout + /// 等待下一条上游消息的空闲超时；`None` 或零值使用默认 300 秒。 + pub(crate) fn stream_idle_timeout(&self) -> Option { + self.config.stream_idle_timeout } /// 注册由连接池生命周期托管的 opening 任务。 diff --git a/backend/crates/providers/openai/tests/config.rs b/backend/crates/providers/openai/tests/config.rs index c07ce58f2..7a20ffa9e 100644 --- a/backend/crates/providers/openai/tests/config.rs +++ b/backend/crates/providers/openai/tests/config.rs @@ -97,7 +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!(DEFAULT_STREAM_MAX_RETRIES, 5); assert_eq!( ( @@ -105,7 +105,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() { config.ws_pool.enabled, config.ws_pool.max_age_ms, config.ws_pool.max_connecting, - config.ws_pool.initial_event_timeout_ms, + config.ws_pool.stream_idle_timeout_ms, config.quota.refresh_interval_minutes, config.auth.refresh_enabled, config.auth.oauth_client_id.as_str(), @@ -117,7 +117,7 @@ fn openai_config_defaults_to_the_provider_owned_operating_values() { true, 3_300_000, 8, - 20_000, + 300_000, 15, true, "app_EMoamEEZ73f0CkXaXp7hrann", diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index 77dc24187..2e0bbab7f 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -459,7 +459,6 @@ fn contract_account_scope() -> Arc { "acct_truncated_stream", "acct_usage_limit_request_path", "acct_websocket_close", - "acct_websocket_cooldown", "acct_websocket_fast_path", "acct_websocket_busy_replay", "acct_websocket_metadata_close", @@ -666,7 +665,6 @@ async fn capture_turn_state_request( previous_turn_id: Option<&str>, current_turn_id: Option<&str>, client_turn_state: Option<&str>, - retry_checkpoint: bool, ) -> wiremock::Request { let account_id = "acct_session_affinity"; let store = Arc::new(MemoryAccountStore::default()); @@ -691,9 +689,6 @@ async fn capture_turn_state_request( if let Some(previous_turn_id) = previous_turn_id { session_state.insert("client_turn_id".to_owned(), json!(previous_turn_id)); } - if retry_checkpoint { - session_state.insert("retry_checkpoint".to_owned(), json!(true)); - } let mut protocol_context = Map::from_iter([("use_websocket".to_owned(), json!(false))]); if let Some(current_turn_id) = current_turn_id { protocol_context.insert("turn_id".to_owned(), json!(current_turn_id)); @@ -1326,7 +1321,7 @@ async fn selection_infrastructure_errors_have_a_distinct_classification() { } #[tokio::test] -async fn websocket_close_after_delivery_preserves_details_and_forces_a_fresh_client_retry() { +async fn websocket_close_after_delivery_preserves_details_and_reconnects_through_pool() { const ACCOUNT_ID: &str = "acct_websocket_close"; const CONVERSATION_ID: &str = "conversation-websocket-close-after-delivery"; const SESSION_ID: &str = "committed-websocket-session"; @@ -1496,7 +1491,7 @@ async fn websocket_close_after_delivery_preserves_details_and_forces_a_fresh_cli retry_pool_observation = Some(observation.websocket_pool()); } } - assert_eq!(retry_pool_observation, Some(None)); + assert_eq!(retry_pool_observation, Some(Some(WebSocketPoolKind::New))); server.await.expect("WebSocket server"); } @@ -1687,7 +1682,7 @@ async fn websocket_fast_path_miss_uses_http_and_keeps_background_preconnect() { } #[tokio::test] -async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_chain() { +async fn ambiguous_websocket_close_reconnects_and_retains_native_continuation() { const ACCOUNT_ID: &str = "acct_websocket_close"; const CONVERSATION_ID: &str = "conversation-websocket-ambiguous-close"; const SESSION_ID: &str = "sticky-websocket-session"; @@ -1776,7 +1771,7 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha let (stream, _) = listener .accept() .await - .expect("accept fresh WebSocket probe"); + .expect("accept new pooled WebSocket"); let mut websocket = accept_codex_test_websocket(stream).await; let _request = websocket .next() @@ -1799,11 +1794,6 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha .await .expect("complete fresh WebSocket request"); - let (stream, _) = listener - .accept() - .await - .expect("accept same-session continuation WebSocket connection"); - let mut websocket = accept_codex_test_websocket(stream).await; let request = websocket .next() .await @@ -1817,7 +1807,7 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha .expect("same-session continuation request JSON"); assert_eq!( payload.get("previous_response_id"), - Some(&json!("external-previous-response")) + Some(&json!("resp_fresh_retry")) ); websocket .send(completed_response("resp_external_continuation")) @@ -1930,29 +1920,48 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha .expect("same session should prepare a fresh WebSocket stream"); assert_eq!(fresh_stream.metadata().transport().as_str(), "websocket"); let mut fresh_pool_observation = None; + let mut recovered_session = None; while let Some(event) = fresh_stream.next().await { let event = event.expect("fresh WebSocket response"); + if let Some(update) = event.session_update() { + recovered_session = Some(update.clone()); + } if let Some(observation) = event.response_observation() && observation.transport().as_str() == "websocket" { fresh_pool_observation = Some(observation.websocket_pool()); } } - assert_eq!(fresh_pool_observation, Some(None)); + assert_eq!(fresh_pool_observation, Some(Some(WebSocketPoolKind::New))); + let recovered_session = recovered_session.expect("recovered session"); + assert_eq!( + recovered_session.payload().get("continuation_scope"), + Some(&json!("connection_local")) + ); drop(fresh_stream); - let continuation_operation = Operation::Generate(generate_with_session_context( - "sticky-websocket-session", - Some("thread-continuation"), - None, - )); + let continuation_operation = Operation::Generate( + generate_with_session_context( + "sticky-websocket-session", + Some("thread-continuation"), + None, + ) + .with_provider_session_state(recovered_session), + ); let mut continuation_stream = provider .execute( planned_request("openai", continuation_operation), - external_continuation_context("req_sticky_external_continuation"), + pinned_continuation_context( + "req_recovered_continuation", + ACCOUNT_ID, + "resp_fresh_retry", + "resp_fresh_retry", + 1, + ContinuationAttempt::Native, + ), ) .await - .expect("successful fresh probe must restore normal continuation routing"); + .expect("recovered pooled socket must support native continuation"); assert_eq!( continuation_stream.metadata().transport().as_str(), "websocket" @@ -2009,366 +2018,117 @@ async fn ambiguous_websocket_close_probes_a_fresh_websocket_for_the_next_new_cha } #[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"; - +async fn repeated_websocket_failures_exhaust_budget_then_use_http_and_report_http_failure() { let store = Arc::new(MemoryAccountStore::default()); - create_account(&store, ACCOUNT_ID).await; + create_account(&store, "acct_websocket_close").await; let listener = TcpListener::bind("127.0.0.1:0") .await - .expect("bind WebSocket listener"); + .expect("bind 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"); + for code in [1000, 1001, 1011] { + let (stream, _) = listener.accept().await.expect("accept websocket"); let mut websocket = accept_codex_test_websocket(stream).await; - let _request = websocket + 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"); + .expect("request") + .expect("valid frame"); websocket .close(Some(CloseFrame { - code: CloseCode::Normal, + code: CloseCode::from(code), reason: "".into(), })) .await - .expect("close WebSocket before terminal event"); + .expect("close websocket"); } - - 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"); + for status in [200, 503] { + let (mut http, _) = listener.accept().await.expect("accept HTTP fallback"); + assert!( + String::from_utf8_lossy(&capture_http_request(&mut http).await) + .starts_with("POST /codex/responses") + ); + let body = if status == 200 { + CAPTURE_COMPLETED_SSE + } else { + r#"{"error":{"code":"server_error","message":"unavailable"}}"# + }; + http.write_all(format!("HTTP/1.1 {status} Response\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", body.len()).as_bytes()).await.expect("HTTP response"); } }); - - 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 + let provider = provider_with_base_url_and_retry_budget(&store, base_url, 2); + for index in 0..3 { + let operation = Operation::Generate(generate_with_persisted_session_context( + "acct_websocket_close", + "conversation-repeated-close", + "repeated-close", + "turn", + )); + let mut stream = provider .execute( - planned_request( - "openai", - operation(&format!("thread-active-cooldown-{index}")), - ), + planned_request("openai", operation), context( - &format!("req_active_cooldown_{index}"), + &format!("req_repeated_close_{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()); - } - } + .expect("prepare stream"); + assert_eq!(stream.metadata().transport().as_str(), "websocket"); + let error = loop { + match stream.next().await { + Some(Ok(_)) => {} Some(Err(error)) => break error, - None => panic!("half-open WebSocket reset must surface a provider error"), + None => panic!("close must fail"), } }; - 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(); + assert_eq!(error.send_state(), UpstreamSendState::Ambiguous); + assert_eq!(error.pre_delivery_retry(), None); + assert!(!error.replay_is_safe()); } - - 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()); + for index in 0..2 { + let operation = Operation::Generate(generate_with_persisted_session_context( + "acct_websocket_close", + "conversation-repeated-close", + "repeated-close", + "turn", + )); + let mut stream = provider + .execute( + planned_request("openai", operation), + context( + &format!("req_http_after_ws_{index}"), + CancellationToken::new(), + ), + ) + .await + .expect("prepare HTTP"); + assert_eq!(stream.metadata().transport().as_str(), "http_sse"); + let mut failure = None; + while let Some(event) = stream.next().await { + if let Err(error) = event { + failure = Some(error); + break; + } + } + if index == 0 { + assert!(failure.is_none()); + } else { + let failure = failure.expect("HTTP failure is final transport result"); + assert_eq!(failure.upstream_status(), Some(503)); + assert!(!matches!( + failure.pre_delivery_retry(), + Some( + PreDeliveryRetry::SameAccountTransportRetry { .. } + | PreDeliveryRetry::SameAccountTransportFallback + ) + )); } } - 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"); + server.await.expect("server"); } #[tokio::test] @@ -2452,6 +2212,45 @@ async fn websocket_upgrade_required_immediately_enables_session_http_fallback() while let Some(event) = second.next().await { event.expect("HTTP fallback response"); } + tokio::time::pause(); + for index in 0..3 { + tokio::time::advance(Duration::from_secs(5 * 60 * 60)).await; + let stream = provider + .execute( + planned_request("openai", operation()), + context( + &format!("req_active_http_session_{index}"), + CancellationToken::new(), + ), + ) + .await + .expect("prepare active session"); + assert_eq!(stream.metadata().transport().as_str(), "http_sse"); + } + let unrelated = provider + .execute( + planned_request( + "openai", + Operation::Generate(generate_with_session_context( + "unrelated-session", + None, + None, + )), + ), + context("req_unrelated_session", CancellationToken::new()), + ) + .await + .expect("prepare unrelated session"); + assert_eq!(unrelated.metadata().transport().as_str(), "websocket"); + tokio::time::advance(Duration::from_secs(8 * 60 * 60)).await; + let expired = provider + .execute( + planned_request("openai", operation()), + context("req_expired_http_session", CancellationToken::new()), + ) + .await + .expect("prepare expired session"); + assert_eq!(expired.metadata().transport().as_str(), "websocket"); server.await.expect("upstream server"); } @@ -2572,7 +2371,7 @@ async fn websocket_turn_state_metadata_is_exposed_through_response_observation() } #[tokio::test] -async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() { +async fn websocket_turn_state_metadata_close_does_not_authorize_replay() { const ACCOUNT_ID: &str = "acct_websocket_metadata_close"; const TURN_STATE: &str = "turn-state-before-close"; @@ -2630,7 +2429,7 @@ async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() { .await .expect("prepare WebSocket provider stream"); let mut observed_turn_state = false; - let mut replay_checkpoint = None; + let mut session_update = None; let error = loop { match stream.next().await { Some(Ok(event)) => { @@ -2641,7 +2440,7 @@ async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() { }); } if let Some(update) = event.session_update() { - replay_checkpoint = Some(update.clone()); + session_update = Some(update.clone()); } } Some(Err(error)) => break error, @@ -2651,22 +2450,10 @@ async fn websocket_turn_state_metadata_close_requests_checkpointed_retry() { server.await.expect("WebSocket server"); assert!(observed_turn_state); - let replay_checkpoint = replay_checkpoint.expect("metadata close checkpoint"); - assert_eq!( - replay_checkpoint.payload().get("turn_state"), - Some(&json!(TURN_STATE)) - ); - assert_eq!( - replay_checkpoint.payload().get("retry_checkpoint"), - Some(&json!(true)) - ); + assert!(session_update.is_none()); assert_eq!(error.send_state(), UpstreamSendState::Ambiguous); - assert!(error.replay_is_safe()); - assert!(matches!( - error.pre_delivery_retry(), - Some(PreDeliveryRetry::SameAccountTransportRetry { retry_index, .. }) - if retry_index.get() == 1 - )); + assert!(!error.replay_is_safe()); + assert_eq!(error.pre_delivery_retry(), None); assert_eq!( error.upstream_code().map(|code| code.as_str()), Some("websocket_close_1000") @@ -3114,7 +2901,6 @@ async fn matching_turn_id_should_restore_previous_turn_state() { Some("turn-same"), Some("turn-same"), None, - false, ) .await; @@ -3124,17 +2910,6 @@ async fn matching_turn_id_should_restore_previous_turn_state() { ); } -#[tokio::test] -async fn retry_checkpoint_should_restore_turn_state_without_client_turn_id() { - let request = - capture_turn_state_request("req_restore_retry_checkpoint", None, None, None, true).await; - - assert_eq!( - captured_header_values(&request, "x-codex-turn-state"), - vec![b"previous-turn-state".to_vec()] - ); -} - #[tokio::test] async fn matching_turn_id_should_prefer_an_explicit_client_echo_over_saved_provider_state() { let request = capture_turn_state_request( @@ -3142,7 +2917,6 @@ async fn matching_turn_id_should_prefer_an_explicit_client_echo_over_saved_provi Some("turn-same"), Some("turn-same"), Some("client-turn-state"), - false, ) .await; @@ -3174,7 +2948,6 @@ async fn new_or_unidentified_turn_should_not_restore_previous_turn_state() { previous_turn_id, current_turn_id, client_turn_state, - false, ) .await; assert!(captured_header_values(&request, "x-codex-turn-state").is_empty()); @@ -5217,3 +4990,141 @@ async fn provider_routes_catalog_model_when_supported_in_api_is_false() { .is_some() ); } + +#[tokio::test] +async fn completed_websocket_response_resets_consecutive_failure_budget() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_websocket_close").await; + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!("http://{}", listener.local_addr().expect("address")); + let server = tokio::spawn(async move { + let mut active = None; + for succeeds in [false, true, false, false, true] { + if active.is_none() { + let (stream, _) = listener.accept().await.expect("accept WS"); + active = Some(accept_codex_test_websocket(stream).await); + } + let ws = active.as_mut().expect("active connection"); + ws.next().await.expect("request").expect("valid frame"); + if succeeds { + ws.send(Message::Text(json!({"type":"response.created","response":{"id":"resp_budget_reset","model":"gpt-5.4"}}).to_string().into())).await.expect("created response"); + ws.send(Message::Text(json!({"type":"response.completed","response":{"id":"resp_budget_reset","model":"gpt-5.4","status":"completed","output":[]}}).to_string().into())).await.expect("complete response"); + } else { + ws.close(None).await.expect("close connection"); + active = None; + } + } + }); + let provider = provider_with_base_url_and_retry_budget(&store, base_url, 2); + for (index, succeeds) in [false, true, false, false, true].into_iter().enumerate() { + let operation = Operation::Generate(generate_with_persisted_session_context( + "acct_websocket_close", + "conversation-budget-reset", + "budget-reset", + "turn", + )); + let mut stream = provider + .execute( + planned_request("openai", operation), + context( + &format!("req_reset_budget_{index}"), + CancellationToken::new(), + ), + ) + .await + .expect("prepare request"); + assert_eq!( + stream.metadata().transport().as_str(), + "websocket", + "request {index}" + ); + let mut failed = false; + let mut completed = false; + while let Some(event) = stream.next().await { + if let Ok(event) = &event { + completed |= event + .canonical_facts() + .iter() + .any(|fact| matches!(fact, GatewayEvent::Completed(_))); + } + if event.is_err() { + failed = true; + break; + } + } + assert_eq!(failed, !succeeds); + assert_eq!(completed, succeeds); + } + server.await.expect("server"); +} + +#[tokio::test] +async fn connection_limit_rejection_requests_one_provider_retry_and_reconnects_in_pool() { + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, "acct_websocket_close").await; + let listener = TcpListener::bind("127.0.0.1:0").await.expect("listener"); + let base_url = format!("http://{}", listener.local_addr().expect("address")); + let server = tokio::spawn(async move { + for first in [true, false] { + let (stream, _) = listener.accept().await.expect("accept WS"); + let mut ws = accept_codex_test_websocket(stream).await; + ws.next().await.expect("request").expect("valid frame"); + let event = if first { + json!({"type":"error","status":400,"error":{"code":"websocket_connection_limit_reached","type":"invalid_request_error","message":"connection expired"}}) + } else { + json!({"type":"response.completed","response":{"id":"resp_reconnected","model":"gpt-5.4","status":"completed","output":[]}}) + }; + if !first { + ws.send(Message::Text(json!({"type":"response.created","response":{"id":"resp_reconnected","model":"gpt-5.4"}}).to_string().into())).await.expect("created response"); + } + ws.send(Message::Text(event.to_string().into())) + .await + .expect("send response"); + } + }); + let provider = provider_with_base_url_and_retry_budget(&store, base_url, 1); + let operation = || { + Operation::Generate(generate_with_persisted_session_context( + "acct_websocket_close", + "conversation-rejected", + "rejected", + "turn", + )) + }; + let mut stream = provider + .execute( + planned_request("openai", operation()), + context("req_rejected", CancellationToken::new()), + ) + .await + .expect("prepare request"); + let error = loop { + match stream.next().await { + Some(Ok(_)) => {} + Some(Err(error)) => break error, + None => panic!("expected rejection"), + } + }; + assert_eq!(error.send_state(), UpstreamSendState::Sent); + assert!(error.replay_is_safe()); + assert!( + matches!(error.pre_delivery_retry(), Some(PreDeliveryRetry::SameAccountTransportRetry { retry_index, .. }) if retry_index.get() == 1) + ); + drop(stream); + let retry_context = context("req_rejected", CancellationToken::new()).with_transport( + AttemptTransport::Retry(NonZeroU32::new(1).expect("retry index")), + ); + let mut retry = provider + .execute(planned_request("openai", operation()), retry_context) + .await + .expect("prepare retry"); + let mut pooled = false; + while let Some(event) = retry.next().await { + let event = event.expect("retry succeeds"); + pooled |= event.response_observation().is_some_and(|observation| { + observation.websocket_pool() == Some(WebSocketPoolKind::New) + }); + } + assert!(pooled); + server.await.expect("server"); +} diff --git a/backend/crates/providers/openai/tests/transport/latency.rs b/backend/crates/providers/openai/tests/transport/latency.rs index da95b40b5..10fdf6c97 100644 --- a/backend/crates/providers/openai/tests/transport/latency.rs +++ b/backend/crates/providers/openai/tests/transport/latency.rs @@ -59,7 +59,7 @@ async fn record_websocket_origin_failures( let request_id = format!("req_{scenario}_{attempt}"); let error = backend .create_response( - &explicit_websocket_warmup_request(&conversation_id), + &new_chain_request(&conversation_id), request_context(&request_id, Some("chatgpt-account")), ) .await @@ -998,7 +998,7 @@ async fn concurrent_same_key_should_singleflight_websocket_opening() { } #[tokio::test] -async fn websocket_failure_before_first_delivery_should_use_same_account_http() { +async fn websocket_failure_before_first_delivery_should_not_replay_payload() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let server = tokio::spawn(async move { @@ -1061,18 +1061,12 @@ async fn websocket_failure_before_first_delivery_should_use_same_account_http() .await .unwrap(); - let (mut http, _) = listener.accept().await.unwrap(); - let request = read_http_request(&mut http).await; - assert!(request.starts_with("POST /codex/responses HTTP/1.1")); - assert_eq!( - read_header_value(&request, "authorization"), - Some("Bearer access-token") + assert!( + timeout(Duration::from_millis(100), listener.accept()) + .await + .is_err(), + "transport must not replay over HTTP or WS" ); - assert_eq!( - read_header_value(&request, "chatgpt-account-id"), - Some("chatgpt-account") - ); - write_completed_sse_response(&mut http).await; }); let backend = CodexBackendClient::new( reqwest::Client::builder().no_proxy().build().unwrap(), @@ -1081,22 +1075,16 @@ async fn websocket_failure_before_first_delivery_should_use_same_account_http() ) .with_websocket_pool(Arc::new(CodexWebSocketPool::default())); - let response = backend + let error = backend .create_response( &new_chain_request("conversation-pre-delivery-fallback"), request_context("req_pre_delivery_fallback", Some("chatgpt-account")), ) .await - .expect("pre-delivery WebSocket failure should use same-account HTTP"); + .expect_err("sent payload must not be replayed"); server.await.unwrap(); - assert_eq!(response.transport, CodexBackendTransport::HttpSse); - assert_eq!( - response.transport_metrics.decision, - Some(CodexTransportDecision::Http2PreDeliveryFailure) - ); - assert!(response.body.contains("resp_http_fallback")); - assert!(!response.body.contains("resp_abandoned_websocket")); + assert!(matches!(error, CodexClientError::WebSocket(_))); } #[tokio::test] @@ -1479,3 +1467,63 @@ async fn half_open_upstream_response_should_close_origin_breaker() { assert!(response.body.contains("resp_after_half_open_success")); pool.shutdown().await; } + +#[tokio::test] +async fn open_fast_path_breaker_should_allow_required_warmup_and_hot_reuse() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let threshold = WebSocketOriginBreakerConfig::default().failure_threshold; + let server = tokio::spawn(async move { + reject_websocket_openings(&listener, threshold).await; + let (mut http, _) = listener.accept().await.unwrap(); + assert!( + read_http_request(&mut http) + .await + .starts_with("POST /codex/responses") + ); + write_completed_sse_response(&mut http).await; + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + for id in ["resp_required_warmup", "resp_hot_after_breaker"] { + websocket.next().await.unwrap().unwrap(); + websocket + .send(Message::Text(completed_websocket_response(id, 2, 1).into())) + .await + .unwrap(); + } + }); + let pool = Arc::new(CodexWebSocketPool::default()); + let backend = CodexBackendClient::new( + reqwest::Client::builder().no_proxy().build().unwrap(), + format!("http://{addr}"), + test_wire_profile(), + ) + .with_websocket_pool(Arc::clone(&pool)); + record_websocket_origin_failures(&backend, threshold, "required-warmup").await; + let fallback = backend + .create_response( + &new_chain_request("ordinary-cold"), + request_context("req_cold_suppressed", Some("chatgpt-account")), + ) + .await + .expect("open fast path selects HTTP"); + assert_eq!(fallback.transport, CodexBackendTransport::HttpSse); + let warmup = backend + .create_response( + &explicit_websocket_warmup_request("required-warmup"), + request_context("req_required_warmup", Some("chatgpt-account")), + ) + .await + .expect("required WS bypasses optional fast path breaker"); + assert_eq!(warmup.transport, CodexBackendTransport::WebSocket); + let hot = backend + .create_response( + &new_chain_request("required-warmup"), + request_context("req_hot_after_breaker", Some("chatgpt-account")), + ) + .await + .expect("hot socket remains reusable"); + assert_eq!(hot.websocket_pool_decision.unwrap().kind(), "reuse"); + server.await.unwrap(); + pool.shutdown().await; +} diff --git a/backend/crates/providers/openai/tests/transport/mod.rs b/backend/crates/providers/openai/tests/transport/mod.rs index ca481211d..609b61465 100644 --- a/backend/crates/providers/openai/tests/transport/mod.rs +++ b/backend/crates/providers/openai/tests/transport/mod.rs @@ -506,7 +506,7 @@ fn websocket_pool_config_for_tests( ping_interval, ping_timeout: Duration::from_secs(1), liveness_timeout, - initial_event_timeout: None, + stream_idle_timeout: None, } } diff --git a/backend/crates/providers/openai/tests/transport/protocol/responses.rs b/backend/crates/providers/openai/tests/transport/protocol/responses.rs index 7d5f91b54..5eb2ac32b 100644 --- a/backend/crates/providers/openai/tests/transport/protocol/responses.rs +++ b/backend/crates/providers/openai/tests/transport/protocol/responses.rs @@ -289,30 +289,11 @@ fn new_chain_should_allow_pre_delivery_http_fallback() { request.use_websocket = true; assert!( - transport_requirement(&request).allows_pre_delivery_http_fallback(), - "new chains may follow the TS same-account fallback behavior" + transport_requirement(&request).allows_connection_restart(), + "new chains have no connection-bound state to lose on an explicit rejection" ); } -#[test] -fn session_transport_recovery_should_only_apply_to_new_chains() { - assert!(TransportRequirement::NewChain.allows_session_transport_recovery()); - - for requirement in [ - TransportRequirement::HttpRequired, - TransportRequirement::ExplicitWebSocketWarmup, - TransportRequirement::ExactWebSocketContinuation, - TransportRequirement::PersistedContinuation, - TransportRequirement::ExternalUnknown, - ] { - assert!( - !requirement.allows_session_transport_recovery(), - "session recovery must not override {}", - requirement.as_str() - ); - } -} - #[test] fn transport_should_mark_unknown_previous_response_as_external() { let mut request = codex_request("gpt-test", "be brief", Vec::new()); @@ -331,7 +312,7 @@ fn previous_response_should_forbid_pre_delivery_http_fallback() { request.set_previous_response_id(Some("resp_previous".to_owned())); assert!( - !transport_requirement(&request).allows_pre_delivery_http_fallback(), + !transport_requirement(&request).allows_connection_restart(), "a sent continuation must not be replayed on another transport" ); } diff --git a/backend/crates/providers/openai/tests/transport/protocol/websocket.rs b/backend/crates/providers/openai/tests/transport/protocol/websocket.rs index e834e3050..0947c2cd1 100644 --- a/backend/crates/providers/openai/tests/transport/protocol/websocket.rs +++ b/backend/crates/providers/openai/tests/transport/protocol/websocket.rs @@ -161,7 +161,7 @@ fn websocket_event_to_sse_should_forward_public_events_and_strip_internal_events } #[test] -fn websocket_event_to_sse_should_add_missing_previous_response_recovery_code() { +fn websocket_event_to_sse_should_preserve_errors_without_inventing_recovery_codes() { let event = json!({ "type": "error", "status": 400, @@ -174,7 +174,7 @@ fn websocket_event_to_sse_should_add_missing_previous_response_recovery_code() { let frame = websocket_event_to_sse_frame(&event).expect("public error event"); - assert!(frame.contains(r#""code":"previous_response_not_found""#)); + assert_eq!(frame, format!("event: error\ndata: {event}\n\n")); } #[test] diff --git a/backend/crates/providers/openai/tests/transport/websocket/mod.rs b/backend/crates/providers/openai/tests/transport/websocket/mod.rs index 39c8a558b..ce0344753 100644 --- a/backend/crates/providers/openai/tests/transport/websocket/mod.rs +++ b/backend/crates/providers/openai/tests/transport/websocket/mod.rs @@ -1392,12 +1392,12 @@ async fn codex_backend_client_should_timeout_when_upstream_is_silent() { message, .. }) - if message.contains("20s") + if message.contains("300s") ); } #[tokio::test] -async fn websocket_response_created_should_switch_to_active_stream_idle_timeout() { +async fn websocket_stream_should_allow_silence_below_idle_timeout() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); let (created_tx, created_rx) = tokio::sync::oneshot::channel(); @@ -1437,7 +1437,7 @@ async fn websocket_response_created_should_switch_to_active_stream_idle_timeout( let response = response_task .await .expect("websocket task should finish") - .expect("structural activity should disable the initial timeout"); + .expect("structural events and model output share the same idle timeout"); server.await.unwrap(); assert!(response.body.contains("event: response.created")); @@ -1445,60 +1445,32 @@ async fn websocket_response_created_should_switch_to_active_stream_idle_timeout( } #[tokio::test] -async fn websocket_connection_limit_error_should_retry_on_fresh_connection() { +async fn websocket_connection_limit_should_be_reported_without_transport_retry() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let (attempts_tx, mut attempts_rx) = tokio::sync::mpsc::channel(2); let server = tokio::spawn(async move { - for attempt in 0..2 { - let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = accept_codex_test_websocket(stream).await; - let _message = websocket.next().await.unwrap().unwrap(); - attempts_tx.send(attempt).await.unwrap(); - if attempt == 0 { - websocket - .send(Message::Text( - json!({ - "type": "error", - "status": 400, - "error": { - "type": "invalid_request_error", - "code": "websocket_connection_limit_reached", - "message": "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue." - } - }) - .to_string() - .into(), - )) - .await - .unwrap(); - } else { - websocket - .send(Message::Text( - completed_websocket_response("resp_retry_fresh", 3, 1).into(), - )) - .await - .unwrap(); - } - websocket.close(None).await.unwrap(); - } + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + websocket.next().await.unwrap().unwrap(); + websocket.send(Message::Text(json!({ + "type":"error", "status":400, + "error":{"type":"invalid_request_error", "code":"websocket_connection_limit_reached", "message":"connection expired"} + }).to_string().into())).await.unwrap(); + assert!( + timeout(Duration::from_millis(100), listener.accept()) + .await + .is_err() + ); }); let prepared = prepared_websocket_request(&format!("http://{addr}")); - let response = execute_response_create_request(&prepared) + let error = execute_response_create_request(&prepared) .await - .expect("connection limit should retry on a fresh connection"); + .expect_err("transport reports rejection"); + assert!(matches!( + error, + CodexClientError::WebSocket(CodexWebSocketExchangeError::ConnectionLimitReached) + )); server.await.unwrap(); - - let mut attempts = Vec::new(); - while let Some(attempt) = attempts_rx.recv().await { - attempts.push(attempt); - } - assert_eq!( - attempts, - vec![0, 1], - "client should retry once on a new connection" - ); - assert!(response.body.contains("resp_retry_fresh")); } #[tokio::test] diff --git a/backend/crates/providers/openai/tests/transport/websocket_pool.rs b/backend/crates/providers/openai/tests/transport/websocket_pool.rs index a1c01a13c..8d13cdd65 100644 --- a/backend/crates/providers/openai/tests/transport/websocket_pool.rs +++ b/backend/crates/providers/openai/tests/transport/websocket_pool.rs @@ -401,7 +401,7 @@ async fn codex_backend_client_stream_should_keep_fresh_socket_after_structural_a first_websocket.close(None).await.unwrap(); }); let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { - initial_event_timeout: Some(Duration::from_millis(30)), + stream_idle_timeout: Some(Duration::from_millis(200)), ..websocket_pool_config_for_tests(None, None, None) })); let backend = CodexBackendClient::new( @@ -490,7 +490,7 @@ async fn codex_backend_client_stream_should_keep_reused_socket_after_structural_ first_websocket.close(None).await.unwrap(); }); let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { - initial_event_timeout: Some(Duration::from_millis(30)), + stream_idle_timeout: Some(Duration::from_millis(200)), ..websocket_pool_config_for_tests(None, None, None) })); let backend = CodexBackendClient::new( @@ -548,7 +548,7 @@ async fn codex_backend_client_stream_should_use_http_when_connecting_limit_is_ex }); let pool = Arc::new(CodexWebSocketPool::with_config(CodexWebSocketPoolConfig { max_connecting: 0, - initial_event_timeout: Some(Duration::from_millis(30)), + stream_idle_timeout: Some(Duration::from_millis(200)), ..websocket_pool_config_for_tests(None, None, None) })); let backend = CodexBackendClient::new(