diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs index 30310f3e83..d08c8fbbc4 100644 --- a/codex-rs/code-mode/src/remote_session/connection.rs +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -1,4 +1,5 @@ use std::fmt; +use std::future::Future; use std::io; use std::path::Path; use std::path::PathBuf; @@ -62,6 +63,8 @@ mod transport; const IPC_CHANNEL_CAPACITY: usize = 128; const HOST_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); +// TODO(anp) make this timeout configurable if 60 seconds is insufficient. +const DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT: Duration = Duration::from_secs(60); const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::(); // Host spawn errors become model-visible tool output. Bound configured paths // while preserving the executable-bearing suffix needed to diagnose failures. @@ -526,16 +529,23 @@ impl Connection { session: RemoteSession, request: WaitRequest, ) -> Result { + // Account for the runtime's one-second yield grace separately from transport. + let runtime_timeout = + Duration::from_millis(request.yield_time_ms).saturating_add(Duration::from_secs(1)); let cancellation = CallerCancellation::new(); let (response_tx, response_rx) = oneshot::channel(); - self.send(DriverCommand::Wait { - session, - request, - caller_cancellation: cancellation.token(), - response_tx, - }) - .await?; - let result = self.receive(response_rx).await; + let result = self + .with_transport_deadline(runtime_timeout, "wait", async { + self.send(DriverCommand::Wait { + session, + request, + caller_cancellation: cancellation.token(), + response_tx, + }) + .await?; + self.receive(response_rx).await + }) + .await; cancellation.disarm(); result } @@ -546,13 +556,16 @@ impl Connection { cell_id: CellId, ) -> Result { let (response_tx, response_rx) = oneshot::channel(); - self.send(DriverCommand::Terminate { - session, - cell_id, - response_tx, + self.with_transport_deadline(Duration::ZERO, "terminate", async { + self.send(DriverCommand::Terminate { + session, + cell_id, + response_tx, + }) + .await?; + self.receive(response_rx).await }) - .await?; - self.receive(response_rx).await + .await } pub(super) async fn shutdown_session(&self, session: RemoteSession) -> Result<(), String> { @@ -565,6 +578,26 @@ impl Connection { self.receive(response_rx).await } + async fn with_transport_deadline( + &self, + runtime_timeout: Duration, + request_type: &str, + request: impl Future>, + ) -> Result { + let deadline = runtime_timeout.saturating_add(DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT); + match tokio::time::timeout(deadline, request).await { + Ok(result) => result, + Err(_) => { + warn!(request_type, "code-mode host request exceeded its deadline"); + let reason = + format!("code-mode host timed out waiting for {request_type} response"); + mark_connection_dead(&self.alive, &self.failure, reason.clone()); + self.cancellation.cancel(); + Err(reason) + } + } + } + async fn send(&self, command: DriverCommand) -> Result<(), String> { if !self.is_alive() { return Err(self.failure_message()); diff --git a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs index bb33002ec2..2dc84eeeec 100644 --- a/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs +++ b/codex-rs/code-mode/src/remote_session/connection/driver_tests.rs @@ -32,6 +32,8 @@ use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; +use super::super::Connection; +use super::super::DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT; use super::ConnectionDriver; use super::DriverCommand; use super::DriverEvent; @@ -46,6 +48,7 @@ struct DriverHarness { outgoing_rx: mpsc::Receiver, cancellation: CancellationToken, alive: Arc, + failure: Arc>>, driver_task: tokio::task::JoinHandle<()>, } @@ -56,6 +59,7 @@ impl DriverHarness { let (outgoing_tx, outgoing_rx) = mpsc::channel(/*max_capacity*/ 16); let cancellation = CancellationToken::new(); let alive = Arc::new(AtomicBool::new(true)); + let failure = Arc::new(StdMutex::new(None)); let (driver, execute_claim_tx) = ConnectionDriver::new( command_rx, event_rx, @@ -63,7 +67,7 @@ impl DriverHarness { outgoing_tx, DriverLifecycle { alive: Arc::clone(&alive), - failure: Arc::new(StdMutex::new(None)), + failure: Arc::clone(&failure), cancellation: cancellation.clone(), }, ); @@ -75,6 +79,7 @@ impl DriverHarness { outgoing_rx, cancellation, alive, + failure, driver_task, } } @@ -647,6 +652,7 @@ async fn concurrent_large_delegate_results_do_not_disconnect_a_backpressured_bul let (bulk_tx, mut bulk_rx) = mpsc::channel(MAX_PENDING_DELEGATE_CALLS); let cancellation = CancellationToken::new(); let alive = Arc::new(AtomicBool::new(true)); + let failure = Arc::new(StdMutex::new(None)); let (driver, execute_claim_tx) = ConnectionDriver::new( command_rx, event_rx, @@ -654,7 +660,7 @@ async fn concurrent_large_delegate_results_do_not_disconnect_a_backpressured_bul outgoing_tx, DriverLifecycle { alive: Arc::clone(&alive), - failure: Arc::new(StdMutex::new(None)), + failure: Arc::clone(&failure), cancellation: cancellation.clone(), }, ); @@ -666,6 +672,7 @@ async fn concurrent_large_delegate_results_do_not_disconnect_a_backpressured_bul outgoing_rx, cancellation, alive, + failure, driver_task, }; let session = remote_session(); @@ -1328,6 +1335,90 @@ async fn remote_wait_accepts_durations_longer_than_five_minutes() { ); } +/// A stalled remote wait must expire even when its request never reaches the host. +#[tokio::test(start_paused = true)] +async fn queued_remote_wait_times_out_and_invalidates_the_connection() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + harness + .open(session.clone(), Arc::new(RecordingDelegate::default())) + .await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let connection = Connection { + command_tx: harness.command_tx.clone(), + execute_claim_tx: harness.execute_claim_tx.clone(), + alive: Arc::clone(&harness.alive), + failure: Arc::clone(&harness.failure), + cancellation: harness.cancellation.clone(), + }; + let response = tokio::spawn(async move { + let result = connection + .wait( + session, + WaitRequest { + cell_id: CellId::new("1".to_string()), + yield_time_ms: 1, + }, + ) + .await; + (connection, result) + }); + while harness.outgoing_rx.is_empty() { + tokio::task::yield_now().await; + } + tokio::time::advance(DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT + Duration::from_secs(2)).await; + + let (connection, result) = response.await.expect("wait task"); + assert_eq!( + result, + Err("code-mode host timed out waiting for wait response".to_string()) + ); + assert!(!harness.alive.load(Ordering::Acquire)); + assert!(harness.cancellation.is_cancelled()); + drop(connection); +} + +/// A stalled termination must expire and invalidate the same connection as a stalled wait. +#[tokio::test(start_paused = true)] +async fn queued_remote_termination_times_out_and_invalidates_the_connection() { + let mut harness = DriverHarness::start(); + let session = remote_session(); + harness + .open(session.clone(), Arc::new(RecordingDelegate::default())) + .await; + let _started = harness + .start_cell(session.clone(), /*request_id*/ 2, "1") + .await; + let connection = Connection { + command_tx: harness.command_tx.clone(), + execute_claim_tx: harness.execute_claim_tx.clone(), + alive: Arc::clone(&harness.alive), + failure: Arc::clone(&harness.failure), + cancellation: harness.cancellation.clone(), + }; + let response = tokio::spawn(async move { + let result = connection + .terminate(session, CellId::new("1".to_string())) + .await; + (connection, result) + }); + while harness.outgoing_rx.is_empty() { + tokio::task::yield_now().await; + } + tokio::time::advance(DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT + Duration::from_secs(1)).await; + + let (connection, result) = response.await.expect("termination task"); + assert_eq!( + result, + Err("code-mode host timed out waiting for terminate response".to_string()) + ); + assert!(!harness.alive.load(Ordering::Acquire)); + assert!(harness.cancellation.is_cancelled()); + drop(connection); +} + #[tokio::test] async fn cancelled_wait_is_retired_before_next_wait_is_sent() { let mut harness = DriverHarness::start(); diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 5ed90d2b62..1056b2f664 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -21,6 +21,7 @@ use codex_protocol::dynamic_tools::DynamicToolNamespaceTool; use codex_protocol::dynamic_tools::DynamicToolResponse; use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::models::PermissionProfile; +use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ToolMode; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; @@ -2053,6 +2054,140 @@ try { Ok(()) } +/// A stalled host wait must return its timeout to the model and reconnect for the next exec. +#[tokio::test(flavor = "current_thread")] +async fn code_mode_wait_timeout_reconnects_on_next_exec() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let mut builder = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(|config| { + config + .features + .enable(Feature::CodeMode) + .expect("code mode should be enabled"); + }); + let test = builder.build_with_auto_env(&server).await?; + + let first_turn = responses::mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_custom_tool_call( + "call-1", + "exec", + "yield_control(); await new Promise(() => {});", + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_assistant_message("msg-1", "waiting"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + + test.submit_turn("start a stalled code-mode cell").await?; + let first_request = first_turn + .last_request() + .expect("initial exec should be returned to the model"); + let first_items = custom_tool_output_items(&first_request, "call-1"); + let cell_id = extract_running_cell_id(text_item(&first_items, /*index*/ 0)); + + let timeout_completion = responses::mount_function_call_agent_response( + &server, + "call-2", + &serde_json::to_string(&serde_json::json!({ + "cell_id": cell_id, + "yield_time_ms": 60_000, + }))?, + "wait", + ) + .await + .completion; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "wait for the stalled cell".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_event_match(&test.codex, |event| match event { + EventMsg::RawResponseItem(raw) => match &raw.item { + ResponseItem::FunctionCall { call_id, .. } if call_id == "call-2" => Some(()), + _ => None, + }, + _ => None, + }) + .await; + + tokio::time::pause(); + for _ in 0..130 { + tokio::time::advance(Duration::from_secs(1)).await; + if timeout_completion + .function_call_output_text("call-2") + .is_some() + { + break; + } + } + tokio::time::resume(); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let timeout_output = timeout_completion + .function_call_output_text("call-2") + .expect("timed-out wait should be returned to the model"); + assert!( + timeout_output.contains("code-mode host timed out waiting for wait response"), + "unexpected wait output: {timeout_output}" + ); + + let reconnect_turn = responses::mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-5"), + ev_custom_tool_call( + "call-3", + "exec", + r#"text("reconnected"); yield_control(); await new Promise(() => {});"#, + ), + ev_completed("resp-5"), + ]), + sse(vec![ + ev_assistant_message("msg-3", "reconnected"), + ev_completed("resp-6"), + ]), + ], + ) + .await; + + test.submit_turn("run a cell after the timeout").await?; + let reconnect_request = reconnect_turn + .last_request() + .expect("replacement exec should be returned to the model"); + let reconnect_items = custom_tool_output_items(&reconnect_request, "call-3"); + assert_eq!( + extract_running_cell_id(text_item(&reconnect_items, /*index*/ 0)), + "g2:1" + ); + assert_eq!(text_item(&reconnect_items, /*index*/ 1), "reconnected"); + + Ok(()) +} + #[cfg_attr(windows, ignore = "no exec_command on Windows")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_can_yield_and_resume_with_wait() -> Result<()> {