From 4cb5caac100c5edd4ec99853769c9f438e2c8360 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 23 Jun 2026 22:19:50 -0700 Subject: [PATCH] Cover remote code mode runtime behavior --- codex-rs/Cargo.lock | 2 + codex-rs/code-mode-host/Cargo.toml | 4 + codex-rs/code-mode-host/src/delegate.rs | 7 +- codex-rs/code-mode-host/src/lib.rs | 7 +- codex-rs/code-mode-host/src/peer.rs | 157 ++++++++- codex-rs/code-mode-host/tests/stdio.rs | 386 +++++++++++++++++++++++ codex-rs/code-mode/src/remote_session.rs | 20 +- 7 files changed, 558 insertions(+), 25 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index dfe4a7fa67..9b8bbbce49 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2520,8 +2520,10 @@ dependencies = [ "codex-code-mode-protocol", "codex-protocol", "codex-utils-cargo-bin", + "libc", "pretty_assertions", "serde_json", + "tempfile", "tokio", "tokio-util", ] diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 2be4713410..66e9ba11a2 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -28,3 +28,7 @@ codex-protocol = { workspace = true } codex-utils-cargo-bin = { workspace = true } pretty_assertions = { workspace = true } serde_json = { workspace = true } +tempfile = { workspace = true } + +[target.'cfg(unix)'.dev-dependencies] +libc = { workspace = true } diff --git a/codex-rs/code-mode-host/src/delegate.rs b/codex-rs/code-mode-host/src/delegate.rs index b4fee25c0a..50d45e7de7 100644 --- a/codex-rs/code-mode-host/src/delegate.rs +++ b/codex-rs/code-mode-host/src/delegate.rs @@ -7,7 +7,6 @@ use codex_code_mode_protocol::NotificationFuture; use codex_code_mode_protocol::ToolInvocationFuture; use codex_code_mode_protocol::host::DelegateRequest; use codex_code_mode_protocol::host::DelegateResponse; -use codex_code_mode_protocol::host::HostToClient; use codex_code_mode_protocol::host::SessionId; use tokio_util::sync::CancellationToken; @@ -80,9 +79,7 @@ impl CodeModeSessionDelegate for RemoteDelegate { } fn cell_closed(&self, cell_id: &CellId) { - self.peer.send(HostToClient::CellClosed { - session_id: self.session_id.clone(), - cell_id: cell_id.into(), - }); + self.peer + .close_cell(self.session_id.clone(), cell_id.clone()); } } diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index 93f1299605..48453345ba 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -227,12 +227,7 @@ impl HostState { cell_id: cell_id.into(), }), ); - self.peer.send(HostToClient::InitialResponse { - id: request_id, - result: WireResult::from_result( - started.initial_response().await.map(Into::into), - ), - }); + self.peer.start_cell(session_id, request_id, started); } Err(err) => self.respond(request_id, Err(err)), } diff --git a/codex-rs/code-mode-host/src/peer.rs b/codex-rs/code-mode-host/src/peer.rs index 582e01a518..7e237323f5 100644 --- a/codex-rs/code-mode-host/src/peer.rs +++ b/codex-rs/code-mode-host/src/peer.rs @@ -1,13 +1,19 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::PoisonError; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::StartedCell; use codex_code_mode_protocol::host::DelegateRequest; use codex_code_mode_protocol::host::DelegateRequestId; use codex_code_mode_protocol::host::DelegateResponse; use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::RequestId; use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::WireResult; use tokio::sync::Mutex; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -18,6 +24,20 @@ pub(super) struct HostPeer { pending: Mutex>>>, next_request_id: AtomicI64, disconnected: CancellationToken, + cell_routes: StdMutex>, +} + +enum CellRoute { + Pending(Vec), + Active(mpsc::UnboundedSender), +} + +enum CellMessage { + Delegate { + id: DelegateRequestId, + message: HostToClient, + }, + Closed, } impl HostPeer { @@ -27,6 +47,7 @@ impl HostPeer { pending: Mutex::new(HashMap::new()), next_request_id: AtomicI64::new(1), disconnected: CancellationToken::new(), + cell_routes: StdMutex::new(HashMap::new()), } } @@ -47,15 +68,21 @@ impl HostPeer { let (response_tx, response_rx) = oneshot::channel(); self.pending.lock().await.insert(id, response_tx); let mut pending = PendingDelegateRequest::new(Arc::clone(self), id); - if !self.send(HostToClient::DelegateRequest { - id, - session_id, - request, - }) { - self.pending.lock().await.remove(&id); - pending.disarm(); - return Err("code-mode client connection closed".to_string()); - } + let cell_id = match &request { + DelegateRequest::InvokeTool { invocation } => invocation.cell_id.clone(), + DelegateRequest::Notify { cell_id, .. } => cell_id.clone(), + }; + self.route_cell_message( + (session_id.clone(), cell_id.into()), + CellMessage::Delegate { + id, + message: HostToClient::DelegateRequest { + id, + session_id, + request, + }, + }, + ); tokio::select! { response = response_rx => { @@ -92,6 +119,118 @@ impl HostPeer { pub(super) fn disconnect(&self) { self.disconnected.cancel(); } + + pub(super) fn start_cell( + self: &Arc, + session_id: SessionId, + request_id: RequestId, + started: StartedCell, + ) { + let key = (session_id, started.cell_id.clone()); + let (messages_tx, messages_rx) = mpsc::unbounded_channel(); + let pending = self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .insert(key.clone(), CellRoute::Active(messages_tx.clone())); + if let Some(CellRoute::Pending(messages)) = pending { + for message in messages { + let _ = messages_tx.send(message); + } + } + let peer = Arc::clone(self); + tokio::spawn(async move { + drive_cell(peer, key, request_id, started, messages_rx).await; + }); + } + + pub(super) fn close_cell(&self, session_id: SessionId, cell_id: CellId) { + self.route_cell_message((session_id, cell_id), CellMessage::Closed); + } + + fn route_cell_message(&self, key: (SessionId, CellId), message: CellMessage) { + use std::collections::hash_map::Entry; + + match self + .cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .entry(key) + { + Entry::Occupied(mut entry) => match entry.get_mut() { + CellRoute::Pending(messages) => messages.push(message), + CellRoute::Active(sender) => { + let _ = sender.send(message); + } + }, + Entry::Vacant(entry) => { + entry.insert(CellRoute::Pending(vec![message])); + } + } + } + + async fn send_delegate_if_pending(&self, id: DelegateRequestId, message: HostToClient) { + if self.pending.lock().await.contains_key(&id) { + self.send(message); + } + } +} + +async fn drive_cell( + peer: Arc, + key: (SessionId, CellId), + request_id: RequestId, + started: StartedCell, + mut messages_rx: mpsc::UnboundedReceiver, +) { + let initial_response = started.initial_response(); + tokio::pin!(initial_response); + let closed = loop { + tokio::select! { + biased; + result = &mut initial_response => { + peer.send(HostToClient::InitialResponse { + id: request_id, + result: WireResult::from_result(result.map(Into::into)), + }); + break false; + } + message = messages_rx.recv() => match message { + Some(CellMessage::Delegate { id, message }) => { + peer.send_delegate_if_pending(id, message).await; + } + Some(CellMessage::Closed) | None => break true, + }, + _ = peer.disconnected.cancelled() => return, + } + }; + + if closed { + peer.send(HostToClient::InitialResponse { + id: request_id, + result: WireResult::from_result(initial_response.await.map(Into::into)), + }); + } else { + loop { + tokio::select! { + message = messages_rx.recv() => match message { + Some(CellMessage::Delegate { id, message }) => { + peer.send_delegate_if_pending(id, message).await; + } + Some(CellMessage::Closed) | None => break, + }, + _ = peer.disconnected.cancelled() => return, + } + } + } + peer.send(HostToClient::CellClosed { + session_id: key.0.clone(), + cell_id: key.1.clone().into(), + }); + peer.cell_routes + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(&key); } struct PendingDelegateRequest { diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs index 4c2c5d1718..1a636def32 100644 --- a/codex-rs/code-mode-host/tests/stdio.rs +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -1,7 +1,9 @@ #![allow(clippy::expect_used)] +use std::process::Stdio; use std::sync::Arc; use std::sync::Mutex; +use std::time::Duration; use codex_code_mode::CellId; use codex_code_mode::CodeModeNestedToolCall; @@ -18,9 +20,26 @@ use codex_code_mode::ToolDefinition; use codex_code_mode::ToolInvocationFuture; use codex_code_mode::WaitOutcome; use codex_code_mode::WaitRequest; +use codex_code_mode::host::CapabilitySet; +use codex_code_mode::host::ClientHello; +use codex_code_mode::host::ClientToHost; +use codex_code_mode::host::DelegateRequest; +use codex_code_mode::host::DelegateResponse; +use codex_code_mode::host::FramedReader; +use codex_code_mode::host::FramedWriter; +use codex_code_mode::host::HostHello; +use codex_code_mode::host::HostRequest; +use codex_code_mode::host::HostResponse; +use codex_code_mode::host::HostToClient; +use codex_code_mode::host::ProtocolVersion; +use codex_code_mode::host::RequestId; +use codex_code_mode::host::SessionId; +use codex_code_mode::host::SupportedProtocolVersions; +use codex_code_mode::host::WireResult; use codex_protocol::ToolName; use pretty_assertions::assert_eq; use serde_json::json; +use tokio::process::Command; use tokio_util::sync::CancellationToken; #[derive(Default)] @@ -69,6 +88,10 @@ fn cell_id(value: &str) -> CellId { CellId::new(value.to_string()) } +fn request_id(value: i64) -> RequestId { + RequestId::new(value) +} + fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { tool_call_id: "call-1".to_string(), @@ -190,3 +213,366 @@ text(result.value); vec![cell_id("1"), cell_id("2"), cell_id("3")] ); } + +#[tokio::test] +async fn cpu_bound_output_loop_can_be_terminated() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let session = provider + .create_session(Arc::new(RecordingDelegate::default())) + .await + .expect("create remote session"); + let mut request = execute_request( + r#" +for (;;) { + text("lots of output"); +} +"#, + ); + request.yield_time_ms = Some(1); + let started = session.execute(request).await.expect("start output loop"); + let running_cell_id = started.cell_id.clone(); + let initial_response = tokio::time::timeout(Duration::from_secs(5), async { + started.initial_response().await + }) + .await + .expect("initial response timeout") + .expect("initial response"); + let RuntimeResponse::Yielded { + cell_id: yielded_cell_id, + content_items, + } = initial_response + else { + panic!("expected the output loop to yield"); + }; + assert_eq!(yielded_cell_id, running_cell_id); + assert!(!content_items.is_empty()); + assert!(content_items.iter().all(|item| { + item == &FunctionCallOutputContentItem::InputText { + text: "lots of output".to_string(), + } + })); + + let termination = tokio::time::timeout( + Duration::from_secs(5), + session.terminate(running_cell_id.clone()), + ) + .await + .expect("termination timeout") + .expect("terminate output loop"); + let WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: terminated_cell_id, + content_items, + }) = termination + else { + panic!("expected the output loop to terminate"); + }; + assert_eq!(terminated_cell_id, running_cell_id); + assert!(content_items.iter().all(|item| { + item == &FunctionCallOutputContentItem::InputText { + text: "lots of output".to_string(), + } + })); + session.shutdown().await.expect("shutdown remote session"); +} + +#[tokio::test] +async fn explicit_yield_frame_precedes_notification_and_terminal_output_drops_timers() { + let mut child = Command::new( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .kill_on_drop(true) + .spawn() + .expect("spawn host"); + let mut writer = FramedWriter::new(child.stdin.take().expect("host stdin")); + let mut reader = FramedReader::new(child.stdout.take().expect("host stdout")); + writer + .write(&ClientToHost::ClientHello( + ClientHello::new( + SupportedProtocolVersions::try_new([ProtocolVersion::V1]) + .expect("supported versions"), + CapabilitySet::empty(), + CapabilitySet::empty(), + ) + .expect("client hello"), + )) + .await + .expect("write client hello"); + assert_eq!( + reader.read::().await.expect("host hello"), + Some(HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + CapabilitySet::empty(), + ))) + ); + + let session_id = SessionId::new("session-1").expect("session ID"); + writer + .write(&ClientToHost::Request { + id: request_id(1), + request: HostRequest::OpenSession { + session_id: session_id.clone(), + }, + }) + .await + .expect("open session"); + assert_eq!( + reader.read::().await.expect("session ready"), + Some(HostToClient::Response { + id: request_id(1), + result: WireResult::Ok { + value: HostResponse::SessionReady { + session_id: session_id.clone(), + }, + }, + }) + ); + + writer + .write(&ClientToHost::Request { + id: request_id(2), + request: HostRequest::Execute { + session_id: session_id.clone(), + request: execute_request( + r#" +text("hello"); +yield_control(); +text("world"); +notify("this is important"); +setTimeout(() => { text("should never emit"); }, 60000); +"#, + ) + .try_into() + .expect("wire execute request"), + }, + }) + .await + .expect("execute request"); + assert_eq!( + reader + .read::() + .await + .expect("execution started"), + Some(HostToClient::Response { + id: request_id(2), + result: WireResult::Ok { + value: HostResponse::ExecutionStarted { + cell_id: cell_id("1").into(), + }, + }, + }) + ); + assert_eq!( + reader + .read::() + .await + .expect("initial response"), + Some(HostToClient::InitialResponse { + id: request_id(2), + result: WireResult::Ok { + value: RuntimeResponse::Yielded { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "hello".to_string(), + }], + } + .into(), + }, + }) + ); + + let delegate_request_id = match reader + .read::() + .await + .expect("notification request") + { + Some(HostToClient::DelegateRequest { + id, + session_id: callback_session_id, + request: + DelegateRequest::Notify { + call_id, + cell_id: callback_cell_id, + text, + }, + }) => { + assert_eq!(callback_session_id, session_id); + assert_eq!(call_id, "call-1"); + assert_eq!(callback_cell_id, cell_id("1").into()); + assert_eq!(text, "this is important"); + id + } + message => panic!("unexpected notification frame: {message:?}"), + }; + writer + .write(&ClientToHost::DelegateResponse { + id: delegate_request_id, + result: WireResult::Ok { + value: DelegateResponse::NotificationDelivered, + }, + }) + .await + .expect("notification response"); + writer + .write(&ClientToHost::Request { + id: request_id(3), + request: HostRequest::Wait { + session_id: session_id.clone(), + request: WaitRequest { + cell_id: cell_id("1"), + yield_time_ms: 60_000, + } + .into(), + }, + }) + .await + .expect("wait request"); + + let terminal = loop { + match reader + .read::() + .await + .expect("terminal response") + { + Some(HostToClient::Response { id, result }) if id == request_id(3) => break result, + Some(HostToClient::CellClosed { + session_id: closed_session_id, + cell_id: closed_cell_id, + }) => { + assert_eq!(closed_session_id, session_id); + assert_eq!(closed_cell_id, cell_id("1").into()); + } + message => panic!("unexpected terminal frame: {message:?}"), + } + }; + assert_eq!( + terminal, + WireResult::Ok { + value: HostResponse::WaitCompleted { + outcome: WaitOutcome::LiveCell(RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "world".to_string(), + }], + error_text: None, + }) + .into(), + }, + } + ); + + writer + .write(&ClientToHost::Request { + id: request_id(4), + request: HostRequest::ShutdownSession { + session_id: session_id.clone(), + }, + }) + .await + .expect("shutdown request"); + loop { + match reader + .read::() + .await + .expect("shutdown response") + { + Some(HostToClient::Response { id, result }) if id == request_id(4) => { + assert_eq!( + result, + WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session_id.clone(), + }, + } + ); + break; + } + Some(HostToClient::CellClosed { .. }) => {} + message => panic!("unexpected shutdown frame: {message:?}"), + } + } + drop(writer); + let status = tokio::time::timeout(Duration::from_secs(5), child.wait()) + .await + .expect("host exit timeout") + .expect("wait for host"); + assert!(status.success(), "host exited with {status}"); +} + +#[cfg(unix)] +#[tokio::test] +async fn crashed_host_fails_in_flight_exec_and_next_exec_respawns() { + use std::os::unix::fs::PermissionsExt; + + let host_binary = + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"); + let temp_dir = tempfile::tempdir().expect("temp dir"); + let pid_log = temp_dir.path().join("host-pids"); + let wrapper = temp_dir.path().join("code-mode-host-wrapper"); + let script = format!( + "#!/bin/sh\nprintf '%s\\n' \"$$\" >> {}\nexec {}\n", + shell_quote(&pid_log), + shell_quote(&host_binary), + ); + std::fs::write(&wrapper, script).expect("write host wrapper"); + let mut permissions = std::fs::metadata(&wrapper) + .expect("host wrapper metadata") + .permissions(); + permissions.set_mode(/*mode*/ 0o755); + std::fs::set_permissions(&wrapper, permissions).expect("make host wrapper executable"); + + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(wrapper); + let session = provider + .create_session(Arc::new(RecordingDelegate::default())) + .await + .expect("create remote session"); + let mut request = execute_request("await new Promise(() => {});"); + request.yield_time_ms = Some(60_000); + let started = session.execute(request).await.expect("start execution"); + let first_pid = recorded_pids(&pid_log)[0]; + + // SAFETY: `first_pid` was written by the wrapper immediately before it + // replaced itself with the host process owned by this test. + assert_eq!(unsafe { libc::kill(first_pid, libc::SIGKILL) }, 0); + let error = tokio::time::timeout(Duration::from_secs(5), started.initial_response()) + .await + .expect("host crash propagation timeout") + .expect_err("crashed host should fail the in-flight execution"); + assert!( + error.contains("code-mode host"), + "unexpected error: {error}" + ); + + assert_eq!( + execute(&session, execute_request(r#"text("recovered");"#)).await, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "recovered".to_string(), + }], + error_text: None, + } + ); + let pids = recorded_pids(&pid_log); + assert_eq!(pids.len(), 2); + assert_ne!(pids[0], pids[1]); + session.shutdown().await.expect("shutdown remote session"); +} + +#[cfg(unix)] +fn shell_quote(path: &std::path::Path) -> String { + format!("'{}'", path.to_string_lossy().replace('\'', "'\"'\"'")) +} + +#[cfg(unix)] +fn recorded_pids(path: &std::path::Path) -> Vec { + std::fs::read_to_string(path) + .expect("read host pid log") + .lines() + .map(|line| line.parse().expect("parse host pid")) + .collect() +} diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs index f5eb443809..83c64a0575 100644 --- a/codex-rs/code-mode/src/remote_session.rs +++ b/codex-rs/code-mode/src/remote_session.rs @@ -146,8 +146,12 @@ pub struct ProcessOwnedCodeModeSession { impl ProcessOwnedCodeModeSession { pub fn new() -> Self { + Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate)) + } + + pub fn with_delegate(delegate: Arc) -> Self { Self::with_process_host( - Arc::new(NoopCodeModeSessionDelegate), + delegate, Arc::new(OwnedProcessHost::new(default_host_program())), ) } @@ -192,13 +196,19 @@ impl ProcessOwnedCodeModeSession { } fn current_connection(&self) -> Result>, String> { - match &*self + let mut state = self .state .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - { + .unwrap_or_else(std::sync::PoisonError::into_inner); + match &*state { SessionState::New => Ok(None), - SessionState::Open(connection) => Ok(Some(Arc::clone(connection))), + SessionState::Open(connection) if connection.is_alive() => { + Ok(Some(Arc::clone(connection))) + } + SessionState::Open(_) => { + *state = SessionState::New; + Ok(None) + } SessionState::Shutdown => Err("code mode session is shutting down".to_string()), } }