diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 935ece0229..8d0d929ad0 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -21,6 +21,7 @@ use std::sync::Arc; use std::task::Context; use std::task::Poll; use tokio::sync::mpsc; +use tokio::sync::oneshot; pub const WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY: &str = "ws_request_header_traceparent"; pub const WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY: &str = "ws_request_header_tracestate"; @@ -413,6 +414,8 @@ pub struct ResponseStream { pub rx_event: mpsc::Receiver>, /// Server-assigned `x-request-id` response header, when present. pub upstream_request_id: Option, + /// Requests a graceful interrupt. Keep consuming events through completion. + pub interrupt: Option>, } impl Stream for ResponseStream { diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index f887edd56c..6e47262737 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -16,6 +16,7 @@ use codex_client::TransportError; use codex_http_client::HttpClientFactory; use codex_websocket_client::WebSocketConnection; use codex_websocket_client::WebSocketConnector; +use futures::FutureExt; use futures::SinkExt; use futures::StreamExt; use http::HeaderMap; @@ -243,6 +244,7 @@ impl ResponsesWebsocketConnection { ) -> Result { let (tx_event, rx_event) = mpsc::channel::>(1600); + let (tx_interrupt, rx_interrupt) = oneshot::channel(); let stream = Arc::clone(&self.stream); let idle_timeout = self.idle_timeout; let server_reasoning_included = self.server_reasoning_included; @@ -318,6 +320,7 @@ impl ResponsesWebsocketConnection { telemetry, turn_state.as_deref(), &timing_log_context, + rx_interrupt, ) => result, _ = tx_event.closed() => Err(ApiError::Stream( "response event consumer dropped".to_string(), @@ -340,6 +343,7 @@ impl ResponsesWebsocketConnection { Ok(ResponseStream { rx_event, upstream_request_id: None, + interrupt: Some(tx_interrupt), }) } } @@ -659,6 +663,7 @@ fn map_wrapped_websocket_error_event( })) } +#[allow(clippy::too_many_arguments)] async fn run_websocket_response_stream( ws_stream: &mut WsStream, tx_event: mpsc::Sender>, @@ -667,6 +672,7 @@ async fn run_websocket_response_stream( telemetry: Option>, turn_state: Option<&OnceLock>, timing_log_context: &ResponsesWebsocketTimingLogContext, + interrupt: oneshot::Receiver<()>, ) -> Result<(), ApiError> { let mut last_server_model: Option = None; let mut safety_buffering_treatment = SafetyBufferingTreatment::default(); @@ -679,11 +685,31 @@ async fn run_websocket_response_stream( ) .await?; + // A response owns its interrupt, and create must be sent before interrupt. + let mut interrupt = interrupt.fuse(); + let mut response_id = None; loop { let poll_start = Instant::now(); - let response = tokio::time::timeout(idle_timeout, ws_stream.next()) - .await - .map_err(|_| ApiError::Stream("idle timeout waiting for websocket".into())); + let response = tokio::select! { + response = tokio::time::timeout(idle_timeout, ws_stream.next()) => { + response.map_err(|_| ApiError::Stream("idle timeout waiting for websocket".into())) + } + Ok(()) = &mut interrupt, if response_id.is_some() => { + send_websocket_request( + ws_stream, + serde_json::json!({ + "type": "response.interrupt", + "response_id": response_id, + "mode": "discard_partial_items", + }).to_string(), + idle_timeout, + /*telemetry*/ None, + timing_log_context.connection_reused, + ) + .await?; + continue; + } + }; if let Some(t) = telemetry.as_ref() { t.on_ws_event(&response, poll_start.elapsed()); } @@ -793,6 +819,9 @@ async fn run_websocket_response_stream( } match process_responses_event(event) { Ok(Some(event)) => { + if let ResponseEvent::Created { response_id: id } = &event { + response_id.clone_from(id); + } let is_completed = matches!(event, ResponseEvent::Completed { .. }); let _ = tx_event.send(Ok(event)).await; if is_completed { diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index b01b8ad46b..c01b13f1fe 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -99,6 +99,7 @@ pub fn spawn_response_stream( ResponseStream { rx_event, upstream_request_id, + interrupt: None, } } @@ -414,18 +415,21 @@ pub fn process_responses_event( event.response, ))); } - "response.incomplete" => { - let reason = event.response.as_ref().and_then(|response| { - response - .get("incomplete_details") - .and_then(|details| details.get("reason")) - .and_then(Value::as_str) - }); - let reason = reason.unwrap_or("unknown"); - let message = format!("Incomplete response returned, reason: {reason}"); - return Err(ResponsesEventError::Api(ApiError::Stream(message))); - } - "response.completed" => { + "response.completed" | "response.incomplete" => { + let interrupted = event.kind == "response.incomplete"; + if interrupted { + let reason = event.response.as_ref().and_then(|response| { + response + .get("incomplete_details") + .and_then(|details| details.get("reason")) + .and_then(Value::as_str) + }); + let reason = reason.unwrap_or("unknown"); + if reason != "interrupted" { + let message = format!("Incomplete response returned, reason: {reason}"); + return Err(ResponsesEventError::Api(ApiError::Stream(message))); + } + } if let Some(resp_val) = event.response { let metadata = resp_val .get("usage") @@ -440,7 +444,11 @@ pub fn process_responses_event( response_id: resp.id, token_usage: resp.usage.map(Into::into), usage_metadata: resp.usage_metadata, - end_turn: resp.end_turn, + end_turn: if interrupted { + Some(false) + } else { + resp.end_turn + }, })); } Err(err) => { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index a182bef7e6..f02cba4d6c 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -2263,11 +2263,6 @@ impl ModelClientSession { } } - /// Drops the cached WebSocket connection and its continuation state. - pub(crate) fn drop_connection(&mut self) { - self.websocket_session.reset(Some("other")); - } - /// Permanently disables WebSockets for this Codex session and resets WebSocket state. /// /// This is used after exhausting the provider retry budget, to force subsequent requests onto @@ -2347,27 +2342,23 @@ const RESPONSE_STREAM_CHANNEL_CAPACITY: usize = 1600; const STREAM_DROPPED_REASON: &str = "response stream dropped before provider terminal event"; fn map_response_stream( - api_stream: codex_api::ResponseStream, + mut api_stream: codex_api::ResponseStream, session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, provider: SharedModelProvider, interceptors: Vec>, ) -> (ResponseStream, oneshot::Receiver) { - let codex_api::ResponseStream { - rx_event, - upstream_request_id, - } = api_stream; - let api_stream = codex_api::ResponseStream { - rx_event, - upstream_request_id: None, - }; - map_response_events( + let upstream_request_id = api_stream.upstream_request_id.take(); + let interrupt = api_stream.interrupt.take(); + let (mut stream, last_response) = map_response_events( upstream_request_id, crate::model_request::intercept_stream(Box::pin(api_stream), interceptors), session_telemetry, inference_trace_attempt, provider, - ) + ); + stream.interrupt = interrupt; + (stream, last_response) } fn map_response_events( @@ -2514,6 +2505,7 @@ where ( ResponseStream { rx_event, + interrupt: None, consumer_dropped: consumer_dropped_for_stream, }, rx_last_response, diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index 96cd7eab01..4e76ac5905 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -15,6 +15,7 @@ use std::sync::Arc; use std::task::Context; use std::task::Poll; use tokio::sync::mpsc; +use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; /// API request payload for a single model turn @@ -117,6 +118,7 @@ fn normalize_image_detail(detail: &mut Option, model_info: &ModelIn pub struct ResponseStream { pub(crate) rx_event: mpsc::Receiver>, + pub(crate) interrupt: Option>, /// Signals the mapper task that the consumer stopped polling before the /// provider stream reached its own terminal event. pub(crate) consumer_dropped: CancellationToken, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index ac3f943ed1..07d57f5536 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -2097,6 +2097,7 @@ async fn intercepted_output_reaches_trace_and_websocket_bookkeeping() -> anyhow: codex_api::ResponseStream { rx_event, upstream_request_id: None, + interrupt: None, }, test_session_telemetry(), attempt, diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index ff3b80573a..fcc6064f18 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -845,6 +845,7 @@ mod tests { drop(tx_event); ResponseStream { rx_event, + interrupt: None, consumer_dropped: CancellationToken::new(), } } diff --git a/codex-rs/core/src/model_request_tests.rs b/codex-rs/core/src/model_request_tests.rs index 82da1f65fe..84ab5014d1 100644 --- a/codex-rs/core/src/model_request_tests.rs +++ b/codex-rs/core/src/model_request_tests.rs @@ -31,6 +31,7 @@ async fn interceptors_forward_before_completion_in_order_and_cancel_upstream() { Box::pin( ResponseStream { rx_event, + interrupt: None, consumer_dropped, } .map(|event| { diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index ce389bd303..b88b343c1d 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2532,11 +2532,11 @@ async fn try_run_sampling_request( .features .enabled(Feature::ConcurrentReasoningSummaries) && turn_context.provider.info().is_openai(); - let preempt = step_context.preempt.clone().unwrap_or_default(); + let mut preempt = step_context.preempt.clone().unwrap_or_default(); let effort = sess .reasoning_effort_for_request(&step_context.settings, super::RequestEffortUsage::Sampling) .await; - let stream = client_session + let mut stream = client_session .stream( prompt, &step_context.settings.model_info, @@ -2548,21 +2548,11 @@ async fn try_run_sampling_request( &inference_trace, ) .instrument(trace_span!("stream_request")) - .or_cancel(&preempt) .or_cancel(&cancellation_token) - .await?; + .await??; if cancellation_token.is_cancelled() { return Err(CodexErr::TurnAborted); } - if preempt.is_cancelled() { - drop(stream); - client_session.drop_connection(); - return Ok(SamplingRequestResult { - needs_follow_up: true, - last_agent_message: None, - }); - } - let mut stream = stream??; let mut in_flight: FuturesOrdered> = FuturesOrdered::new(); let mut needs_follow_up = false; let mut last_agent_message: Option = None; @@ -2620,9 +2610,17 @@ async fn try_run_sampling_request( let event = match event { Ok(Ok(event)) => event, Ok(Err(_)) => { + if let Some(interrupt) = stream.interrupt.take() { + if step_context.settings.model_info.use_responses_lite { + let _ = interrupt.send(()); + } + // Drain the response normally before reusing its connection and history. + preempt = CancellationToken::new(); + needs_follow_up = true; + continue; + } // TODO: Reconcile any response item already being presented to the client. drop(stream); - client_session.drop_connection(); break Ok(SamplingRequestResult { needs_follow_up: true, last_agent_message, diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index fdef849440..0df8b4a20b 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -1131,28 +1131,53 @@ async fn injected_response_item_reopens_turn_after_final_answer() { server.shutdown().await; } +#[test_case(false; "completed_reasoning")] +#[test_case(true; "discarded_reasoning")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn steer_reconnects_websocket_and_sends_full_history() -> anyhow::Result<()> { +async fn steer_interrupts_and_drains_websocket(discard_partial: bool) -> anyhow::Result<()> { core_test_support::skip_if_no_network!(Ok(())); - let server = responses::start_websocket_server(vec![ + let mut interrupted = ev_completed_with_tokens("resp-interrupted", /*total_tokens*/ 42); + interrupted["type"] = json!("response.incomplete"); + interrupted["response"]["status"] = json!("incomplete"); + interrupted["response"]["incomplete_details"] = json!({"reason": "interrupted"}); + let server = responses::start_websocket_server(vec![vec![ + vec![ev_response_created("warm-1"), ev_completed("warm-1")], vec![ - vec![ev_response_created("warm-1"), ev_completed("warm-1")], - vec![ - ev_response_created("resp-interrupted"), - ev_reasoning_item_added("reason-1", &["thinking"]), - ], - // Keep the first connection open until the client closes it. - vec![], + ev_response_created("resp-interrupted"), + ev_reasoning_item_added("reason-1", &["thinking"]), ], - vec![vec![ + // Finish the active response only after receiving response.interrupt. + vec![ + json!({ + "type": "response.interrupt.accepted", + "response_id": "resp-interrupted", + "sequence_number": 3, + }), + if discard_partial { + json!({ + "type": "response.output_item.interrupted", + "response_id": "resp-interrupted", + "item_id": "reason-1", + "output_index": 0, + "sequence_number": 4, + }) + } else { + ev_reasoning_item("reason-1", &["thinking"], &[]) + }, + interrupted, + ], + vec![ ev_response_created("resp-follow-up"), ev_completed("resp-follow-up"), - ]], - ]) + ], + ]]) .await; let test = test_codex() .with_model("gpt-5.4") + .with_model_info_override("gpt-5.4", |model_info| { + model_info.use_responses_lite = true; + }) .with_config(|config| { config .features @@ -1174,27 +1199,34 @@ async fn steer_reconnects_websocket_and_sends_full_history() -> anyhow::Result<( steer_user_input(codex, "second prompt").await; let follow_up = tokio::time::timeout( std::time::Duration::from_secs(10), - server.wait_for_request(/*connection_index*/ 1, /*request_index*/ 0), + server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 3), ) .await - .expect("steer should close the first socket and open a new one") + .expect("steer should interrupt, drain, and reuse the first socket") .body_json(); - assert!(follow_up.get("previous_response_id").is_none()); + assert_eq!( + server.single_connection()[2].body_json(), + json!({"type": "response.interrupt", "response_id": "resp-interrupted", "mode": "discard_partial_items"}), + ); + // Incremental continuation is only possible when the drained output made it into history. + assert_eq!(follow_up["previous_response_id"], "resp-interrupted"); let prompts = message_input_texts(&follow_up, "user") .into_iter() .filter(|text| text == "first prompt" || text == "second prompt") .collect::>(); - assert_eq!(prompts, vec!["first prompt", "second prompt"]); + assert_eq!(prompts, vec!["second prompt"]); wait_for_event(codex, |event| { - assert!(!matches!( + matches!( event, - EventMsg::TurnAborted(_) | EventMsg::StreamError(_) - )); - matches!(event, EventMsg::TurnComplete(completed) if completed.error.is_none()) + EventMsg::TokenCount(count) + if count.info.as_ref().is_some_and(|info| info.last_token_usage.total_tokens == 42) + ) }) .await; + wait_for_successful_turn(codex).await; - assert_eq!(server.connections().len(), 2); + assert_eq!(server.connections().len(), 1); + codex.shutdown_and_wait().await?; server.shutdown().await; Ok(()) }