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:
liame-oai
2026-10-29 00:03:28 +00:00
committed by copyberry
parent f9c10053fb
commit f276b82185
11 changed files with 364 additions and 133 deletions
@@ -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)) => {