From d2f43f4edc29bada481082116567c1e1cab1b3a0 Mon Sep 17 00:00:00 2001 From: jiahao6635 Date: Thu, 24 Sep 2026 15:17:00 +0800 Subject: [PATCH] fix(openai): preserve WebSocket pump exit diagnostics Distinguish pump termination and stream EOF from actual upstream Close frames without changing send-state or replay-safety classification. Cover Pong timeout diagnostics on fresh and reused connections. Refs #249 Co-authored-by: Codex (cherry picked from commit 5a89bbc26be66f85c91dd9b1e7b84a45d621b82d) --- .../providers/openai/src/provider/failure.rs | 23 +++- .../openai/src/transport/client_sse.rs | 11 +- .../openai/src/transport/websocket/error.rs | 7 + .../src/transport/websocket/exchange/io.rs | 1 + .../transport/websocket/exchange/stream.rs | 48 ++++--- .../openai/tests/provider/contract.rs | 128 ++++++++++++++++++ 6 files changed, 193 insertions(+), 25 deletions(-) diff --git a/backend/crates/providers/openai/src/provider/failure.rs b/backend/crates/providers/openai/src/provider/failure.rs index 2bf4a32a6..168356380 100644 --- a/backend/crates/providers/openai/src/provider/failure.rs +++ b/backend/crates/providers/openai/src/provider/failure.rs @@ -1090,8 +1090,9 @@ fn websocket_diagnostic(error: &CodexWebSocketExchangeError) -> ProviderDiagnost } CodexWebSocketExchangeError::InvalidSse(_) => ("decode", "invalid_sse"), CodexWebSocketExchangeError::UnexpectedBinaryEvent => ("decode", "unexpected_binary_event"), - CodexWebSocketExchangeError::ClosedBeforeTerminal(_) => { - ("receive", "closed_before_terminal") + CodexWebSocketExchangeError::ClosedBeforeTerminal(_) => ("receive", "upstream_close"), + CodexWebSocketExchangeError::StreamEndedBeforeTerminal { reason, .. } => { + ("receive", *reason) } CodexWebSocketExchangeError::ConnectionLimitReached(_) => { ("upstream", "connection_limit_reached") @@ -1188,6 +1189,22 @@ fn websocket_diagnostic_message(error: &CodexWebSocketExchangeError) -> Provider CodexWebSocketExchangeError::UnexpectedBinaryEvent => { "OpenAI WebSocket returned an unexpected binary event".to_owned() } + CodexWebSocketExchangeError::StreamEndedBeforeTerminal { + reason, + timeout, + last_event_type, + } => { + let mut message = + format!("OpenAI WebSocket stream ended before a terminal response ({reason})"); + if let Some(timeout) = timeout { + message.push_str(&format!("; local keepalive timeout after {timeout:?}")); + } + if let Some(last_event_type) = last_event_type { + message.push_str("; last event type: "); + message.push_str(last_event_type); + } + message + } CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => { "Reused OpenAI WebSocket died before the first upstream event".to_owned() } @@ -1459,6 +1476,7 @@ pub(super) fn websocket_send_state(error: &CodexWebSocketExchangeError) -> Upstr | CodexWebSocketExchangeError::PostSendAmbiguous { .. } | CodexWebSocketExchangeError::SendTimeout { .. } | CodexWebSocketExchangeError::ClosedBeforeTerminal(_) + | CodexWebSocketExchangeError::StreamEndedBeforeTerminal { .. } | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => { UpstreamSendState::Ambiguous @@ -1489,6 +1507,7 @@ pub(super) fn websocket_error_kind(error: &CodexWebSocketExchangeError) -> Provi | CodexWebSocketExchangeError::Connect(_) | CodexWebSocketExchangeError::PostSendAmbiguous { .. } | CodexWebSocketExchangeError::ClosedBeforeTerminal(_) + | CodexWebSocketExchangeError::StreamEndedBeforeTerminal { .. } | CodexWebSocketExchangeError::ReusedConnectionDiedBeforeFirstEvent { .. } => { ProviderErrorKind::Transport } diff --git a/backend/crates/providers/openai/src/transport/client_sse.rs b/backend/crates/providers/openai/src/transport/client_sse.rs index 437da2348..1c24a763c 100644 --- a/backend/crates/providers/openai/src/transport/client_sse.rs +++ b/backend/crates/providers/openai/src/transport/client_sse.rs @@ -607,12 +607,11 @@ async fn await_websocket_delivery_boundary( } Some(Err(error)) => return Err(error), None => { - return Err(CodexWebSocketExchangeError::closed_before_terminal_on( - exchange.websocket_connection_id, - None, - None, - None, - )); + return Err(CodexWebSocketExchangeError::StreamEndedBeforeTerminal { + reason: "stream_eof", + timeout: None, + last_event_type: None, + }); } } } diff --git a/backend/crates/providers/openai/src/transport/websocket/error.rs b/backend/crates/providers/openai/src/transport/websocket/error.rs index d779ebe2d..6db5dbaae 100644 --- a/backend/crates/providers/openai/src/transport/websocket/error.rs +++ b/backend/crates/providers/openai/src/transport/websocket/error.rs @@ -73,6 +73,13 @@ pub enum CodexWebSocketExchangeError { /// 上游在 terminal 事件前关闭。 #[error("{0}")] ClosedBeforeTerminal(CodexWebSocketCloseError), + /// 未收到 Close 帧即结束，保留 pump 的安全原因码与本地保活时限。 + #[error("websocket stream ended before terminal event ({reason})")] + StreamEndedBeforeTerminal { + reason: &'static str, + timeout: Option, + last_event_type: Option, + }, /// 上游在指定时间内没有发送任何事件。 #[error("websocket receive idle timeout after {timeout:?}")] ReceiveIdleTimeout { 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 02ef44f96..9ff34b642 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/io.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/io.rs @@ -25,6 +25,7 @@ pub(super) fn reused_stream_receive_error( ) -> CodexWebSocketExchangeError { match error { CodexWebSocketExchangeError::ClosedBeforeTerminal(_) + | CodexWebSocketExchangeError::StreamEndedBeforeTerminal { .. } | CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } | CodexWebSocketExchangeError::Transport(_) => { reused_connection_died_before_first_event(error) 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 a240db532..f31073189 100644 --- a/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs +++ b/backend/crates/providers/openai/src/transport/websocket/exchange/stream.rs @@ -219,10 +219,6 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { } }; let Some(message) = message else { - trace.record( - "upstream.eof", - json!({"lastEventType": last_event_type, "terminalSeen": false}), - ); break; }; let raw = match message { @@ -364,20 +360,37 @@ async fn forward_websocket_response_stream(state: WebSocketStreamForwardState) { } let websocket_connection_id = websocket.connection_id(); - let upstream_close = websocket - .exit_reason() - .and_then(|reason| reason.upstream_close().cloned()) - .map(|close| close.with_connection_id(websocket_connection_id)); - let error = match upstream_close { - Some(close) => CodexWebSocketExchangeError::ClosedBeforeTerminal( - close.with_last_event_type(last_event_type), - ), - None => CodexWebSocketExchangeError::closed_before_terminal_on( - websocket_connection_id, - None, - None, + let exit_reason = websocket.exit_reason(); + trace.record( + "upstream.eof", + json!({ + "lastEventType": last_event_type, + "terminalSeen": false, + "failureReason": exit_reason.as_ref().map(PumpExitReason::as_str), + }), + ); + let error = match exit_reason.as_ref() { + Some(PumpExitReason::UpstreamCloseFrame { close }) => { + CodexWebSocketExchangeError::closed_before_terminal_on( + websocket_connection_id, + close.as_ref().and_then(|close| close.code()), + close + .as_ref() + .and_then(|close| close.reason().map(str::to_owned)), + last_event_type, + ) + } + reason => CodexWebSocketExchangeError::StreamEndedBeforeTerminal { + reason: reason.map_or("stream_eof", PumpExitReason::as_str), + timeout: match reason { + Some( + PumpExitReason::PongTimeout { timeout } + | PumpExitReason::LivenessTimeout { timeout }, + ) => Some(*timeout), + _ => None, + }, last_event_type, - ), + }, }; let observation = connection_observation(&websocket, &error); discard_stream_websocket( @@ -478,6 +491,7 @@ fn exchange_exit_reason(error: &CodexWebSocketExchangeError) -> &'static str { "upstream_close" } } + CodexWebSocketExchangeError::StreamEndedBeforeTerminal { reason, .. } => reason, CodexWebSocketExchangeError::ReceiveIdleTimeout { .. } => "receive_idle_timeout", CodexWebSocketExchangeError::UnexpectedBinaryEvent => "unexpected_binary_event", _ => "exchange_failure", diff --git a/backend/crates/providers/openai/tests/provider/contract.rs b/backend/crates/providers/openai/tests/provider/contract.rs index e9d997ce2..604f43943 100644 --- a/backend/crates/providers/openai/tests/provider/contract.rs +++ b/backend/crates/providers/openai/tests/provider/contract.rs @@ -4092,6 +4092,130 @@ async fn websocket_idle_timeout_diagnosis_survives_ambiguous_send_wrapping() { ); } +#[tokio::test] +async fn websocket_pong_timeout_diagnosis_survives_ambiguous_send_wrapping() { + for reuse in [false, true] { + const ACCOUNT_ID: &str = "acct_abrupt_disconnect"; + let store = Arc::new(MemoryAccountStore::default()); + create_account(&store, ACCOUNT_ID).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let (ping_seen_tx, ping_seen_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_codex_test_websocket(stream).await; + websocket.next().await.unwrap().unwrap(); + if reuse { + websocket + .send(Message::Text( + json!({ + "type": "response.completed", + "response": { + "id": "resp_warmup", + "model": "gpt-5.4", + "status": "completed", + "output": [], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} + } + }) + .to_string() + .into(), + )) + .await + .unwrap(); + websocket.next().await.unwrap().unwrap(); + } + for event in [ + json!({"type": "response.created", "response": {"id": "resp_pong", "model": "gpt-5.4"}}), + json!({"type": "response.output_text.delta", "delta": "private-response-body"}), + ] { + websocket + .send(Message::Text(event.to_string().into())) + .await + .unwrap(); + } + std::assert_matches!(websocket.next().await.unwrap().unwrap(), Message::Ping(_)); + ping_seen_tx.send(()).unwrap(); + // 停止 poll，避免 tungstenite 自动回 Pong，复现本地保活超时。 + futures::future::pending::<()>().await; + }); + let operation = Operation::Generate(generate_with_persisted_session_context( + ACCOUNT_ID, + "conversation-pong-timeout", + "session-pong-timeout", + "turn-pong-timeout", + )); + let provider = provider_with_base_url(&store, base_url); + if reuse { + let mut warmup = Arc::clone(&provider) + .execute( + planned_request("openai", operation.clone()), + context("req_pong_warmup", CancellationToken::new()), + ) + .await + .unwrap(); + while let Some(event) = warmup.next().await { + event.unwrap(); + } + } + let mut stream = provider + .execute( + planned_request("openai", operation), + context("req_pong_timeout", CancellationToken::new()), + ) + .await + .unwrap(); + let mut pool_kind = None; + loop { + let event = stream.next().await.unwrap().unwrap(); + if let Some(observation) = event.response_observation() { + pool_kind = observation.websocket_pool().or(pool_kind); + } + if event.has_client_event() { + break; + } + } + assert_eq!( + pool_kind, + Some(if reuse { + WebSocketPoolKind::Reuse + } else { + WebSocketPoolKind::New + }) + ); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(25)).await; + ping_seen_rx.await.unwrap(); + tokio::time::advance(Duration::from_secs(31)).await; + let error = loop { + match stream.next().await { + Some(Ok(_)) => {} + Some(Err(error)) => break error, + None => panic!("Pong timeout must surface a provider error"), + } + }; + server.abort(); + tokio::time::resume(); + assert_eq!(error.send_state(), UpstreamSendState::Ambiguous); + assert!(!error.replay_is_safe()); + assert_eq!(error.pre_delivery_retry(), None); + assert_eq!(error.kind(), ProviderErrorKind::Transport); + assert_eq!( + error.connection_observation().unwrap().exit_reason(), + "pong_timeout" + ); + let diagnostic = error.diagnostic().unwrap(); + assert_eq!(diagnostic.code(), Some("pong_timeout")); + assert_eq!(diagnostic.stage(), Some("receive")); + assert_eq!( + diagnostic.as_str(), + "OpenAI WebSocket stream ended before a terminal response (pong_timeout); local keepalive timeout after 30s; last event type: response.output_text.delta", + ); + assert!(error.raw_upstream_error().is_none()); + assert!(error.client_visible_upstream_error().is_none()); + } +} + #[tokio::test] async fn ambiguous_websocket_close_reconnects_and_retains_native_continuation() { const ACCOUNT_ID: &str = "acct_websocket_close"; @@ -4290,6 +4414,10 @@ async fn ambiguous_websocket_close_reconnects_and_retains_native_continuation() error.upstream_code().map(|code| code.as_str()), Some("websocket_close_1000") ); + assert_eq!( + error.diagnostic().and_then(|diagnostic| diagnostic.code()), + Some("upstream_close") + ); assert_eq!( error.diagnostic().map(|diagnostic| diagnostic.as_str()), Some(