Cover remote code mode runtime behavior

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",
 ]

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 }

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());
     }
 }

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)),
                 }

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<HashMap<DelegateRequestId, oneshot::Sender<Result<DelegateResponse, String>>>>,
     next_request_id: AtomicI64,
     disconnected: CancellationToken,
+    cell_routes: StdMutex<HashMap<(SessionId, CellId), CellRoute>>,
+}
+
+enum CellRoute {
+    Pending(Vec<CellMessage>),
+    Active(mpsc::UnboundedSender<CellMessage>),
+}
+
+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<Self>,
+        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<HostPeer>,
+    key: (SessionId, CellId),
+    request_id: RequestId,
+    started: StartedCell,
+    mut messages_rx: mpsc::UnboundedReceiver<CellMessage>,
+) {
+    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 {

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::<HostToClient>().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::<HostToClient>().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::<HostToClient>()
+            .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::<HostToClient>()
+            .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::<HostToClient>()
+        .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::<HostToClient>()
+            .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::<HostToClient>()
+            .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<libc::pid_t> {
+    std::fs::read_to_string(path)
+        .expect("read host pid log")
+        .lines()
+        .map(|line| line.parse().expect("parse host pid"))
+        .collect()
+}

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<dyn CodeModeSessionDelegate>) -> 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<Option<Arc<Connection>>, 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()),
         }
     }