mirror of
https://github.com/openai/codex.git
synced 2026-10-31 19:28:49 +00:00
Allow upstream interruption when canceling Responses Lite turns (#57864)
## What changed - Add the disabled-by-default `interrupt_on_cancel` feature to interrupt upstream inference when a Responses Lite turn is canceled, with a one-second cleanup deadline. - Detach canceled WebSocket responses so replacement requests can open a new connection without waiting for cleanup. - Add an optional `interrupt_on_cancel` override for thread samples, inheriting the parent setting when omitted. ## Testing Add coverage for cancellation during inference and steering, replacement turns proceeding before cleanup finishes, and prompt cancellation of model response interceptors. GitOrigin-RevId: 501368746fc04fd3fdbd45c1f34c75ce57f4e955
This commit is contained in:
@@ -338,54 +338,70 @@ impl ResponsesWebsocketConnection {
|
||||
if tx_event.is_closed() {
|
||||
return;
|
||||
}
|
||||
let Some(mut ws_stream) = guard.take() else {
|
||||
let _ = tx_event
|
||||
.send(Err(ApiError::Stream(
|
||||
"websocket connection is closed".to_string(),
|
||||
)))
|
||||
.await;
|
||||
return;
|
||||
};
|
||||
let mut completion = None;
|
||||
let result = {
|
||||
let Some(ws_stream) = guard.as_mut() else {
|
||||
let _ = tx_event
|
||||
.send(Err(ApiError::Stream(
|
||||
"websocket connection is closed".to_string(),
|
||||
)))
|
||||
.await;
|
||||
return;
|
||||
};
|
||||
let response = run_websocket_response_stream(
|
||||
&mut ws_stream,
|
||||
tx_event.clone(),
|
||||
request_text,
|
||||
idle_timeout,
|
||||
telemetry,
|
||||
turn_state.as_deref(),
|
||||
&timing_log_context,
|
||||
interrupt.clone(),
|
||||
&mut completion,
|
||||
);
|
||||
tokio::pin!(response);
|
||||
|
||||
tokio::select! {
|
||||
biased;
|
||||
result = run_websocket_response_stream(
|
||||
ws_stream,
|
||||
tx_event.clone(),
|
||||
request_text,
|
||||
idle_timeout,
|
||||
telemetry,
|
||||
turn_state.as_deref(),
|
||||
&timing_log_context,
|
||||
interrupt.clone(),
|
||||
&mut completion,
|
||||
) => result,
|
||||
result = &mut response => result,
|
||||
() = &mut interrupt_timeout => Err(ApiError::Stream(
|
||||
"timed out draining interrupted response".into(),
|
||||
)),
|
||||
_ = tx_event.closed() => Err(ApiError::Stream(
|
||||
"response event consumer dropped".to_string(),
|
||||
)),
|
||||
() = tx_event.closed() => {
|
||||
// Release the lock so later requests can open a new connection
|
||||
// while this task finishes cleanup on the detached socket.
|
||||
drop(guard);
|
||||
if matches!(
|
||||
interrupt.clone().now_or_never(),
|
||||
Some(ResponseInterrupt::Cancel { .. })
|
||||
) {
|
||||
tokio::select! {
|
||||
_ = &mut response => {}
|
||||
() = &mut interrupt_timeout => {}
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if let Err(err) = result {
|
||||
// A terminal stream error should reach the caller immediately. Waiting for a
|
||||
// graceful close handshake here can stall indefinitely and mask the error.
|
||||
let failed_stream = guard.take();
|
||||
drop(guard);
|
||||
drop(failed_stream);
|
||||
drop(ws_stream);
|
||||
// Failed cancellation retires the socket while preserving any received
|
||||
// completion and its usage for the abandoned classifier's cleanup.
|
||||
let event = completion.map_or_else(|| Err(err), Ok);
|
||||
let _ = tx_event.send(event).await;
|
||||
} else if let Some(completed) = completion {
|
||||
// The handshake is finished. Consumer backpressure must not let the
|
||||
// interruption deadline discard completion or its usage accounting.
|
||||
drop(guard);
|
||||
let _ = tx_event.send(Ok(completed)).await;
|
||||
} else {
|
||||
*guard = Some(ws_stream);
|
||||
if let Some(completed) = completion {
|
||||
// The handshake is finished. Consumer backpressure must not let the
|
||||
// interruption deadline discard completion or its usage accounting.
|
||||
drop(guard);
|
||||
let _ = tx_event.send(Ok(completed)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
.instrument(current_span),
|
||||
@@ -925,11 +941,7 @@ async fn run_websocket_response_stream(
|
||||
last_server_model = Some(model);
|
||||
}
|
||||
for metadata_event in metadata_events {
|
||||
if tx_event.send(Ok(metadata_event)).await.is_err() {
|
||||
return Err(ApiError::Stream(
|
||||
"response event consumer dropped".to_string(),
|
||||
));
|
||||
}
|
||||
let _ = tx_event.send(Ok(metadata_event)).await;
|
||||
}
|
||||
match event {
|
||||
Ok(Some(event)) => {
|
||||
|
||||
Reference in New Issue
Block a user