Stabilize websocket response.failed error delivery (#14017)

## What changed
- Drop failed websocket connections immediately after a terminal stream
error instead of awaiting a graceful close handshake before forwarding
the error to the caller.
- Keep the success path and the closed-connection guard behavior
unchanged.

## Why this fixes the flake
- The failing integration test waits for the second websocket stream to
surface the model error before issuing a follow-up request.
- On slower runners, the old error path awaited
`ws_stream.close().await` before sending the error downstream. If that
close handshake stalled, the test kept waiting for an error that had
already happened server-side and nextest timed it out.
- Dropping the failed websocket immediately makes the terminal error
observable right away and marks the session closed so the next request
reconnects cleanly instead of depending on a best-effort close
handshake.

## Code or test?
- This is a production logic fix in `codex-api`. The existing websocket
integration test already exercises the regression path.
This commit is contained in:
Ahmed Ibrahim
2026-03-10 17:59:41 -07:00
committed by Michael Bolin
parent 285b3a5143
commit c8446d7cf3
5 changed files with 102 additions and 32 deletions

View File

@@ -416,6 +416,11 @@ pub struct WebSocketConnectionConfig {
/// Tests use this to force websocket setup into an in-flight state so first-turn warmup paths
/// can be exercised deterministically.
pub accept_delay: Option<Duration>,
/// Whether the server should send a websocket close frame after all scripted responses.
///
/// Tests can disable this to simulate a peer that surfaces a terminal event but never
/// completes the close handshake.
pub close_after_requests: bool,
}
pub struct WebSocketTestServer {
@@ -1168,6 +1173,7 @@ pub async fn start_websocket_server(connections: Vec<Vec<Vec<Value>>>) -> WebSoc
requests,
response_headers: Vec::new(),
accept_delay: None,
close_after_requests: true,
})
.collect();
start_websocket_server_with_headers(connections).await
@@ -1261,6 +1267,7 @@ pub async fn start_websocket_server_with_headers(
log.push(Vec::new());
log.len() - 1
};
let close_after_requests = connection.close_after_requests;
for request_events in connection.requests {
let Some(Ok(message)) = ws_stream.next().await else {
break;
@@ -1324,7 +1331,12 @@ pub async fn start_websocket_server_with_headers(
}
}
let _ = ws_stream.close(None).await;
if close_after_requests {
let _ = ws_stream.close(None).await;
} else {
let _ = shutdown_rx.await;
return;
}
if connections.lock().unwrap().is_empty() {
return;