Add process-owned code mode session

codex-rs/code-mode-host/Cargo.toml

@@ -8,5 +8,23 @@ license.workspace = true
 name = "codex-code-mode-host"
 path = "src/main.rs"
 
+[lib]
+doctest = false
+name = "codex_code_mode_host"
+path = "src/lib.rs"
+
 [lints]
 workspace = true
+
+[dependencies]
+anyhow = { workspace = true }
+codex-code-mode = { workspace = true }
+codex-code-mode-protocol = { workspace = true }
+tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt"] }
+tokio-util = { workspace = true, features = ["rt"] }
+
+[dev-dependencies]
+codex-protocol = { workspace = true }
+codex-utils-cargo-bin = { workspace = true }
+pretty_assertions = { workspace = true }
+serde_json = { workspace = true }

codex-rs/code-mode-host/src/delegate.rs

@@ -0,0 +1,88 @@
+use std::sync::Arc;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeNestedToolCall;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+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;
+
+use crate::peer::HostPeer;
+
+pub(super) struct RemoteDelegate {
+    session_id: SessionId,
+    peer: Arc<HostPeer>,
+}
+
+impl RemoteDelegate {
+    pub(super) fn new(session_id: SessionId, peer: Arc<HostPeer>) -> Self {
+        Self { session_id, peer }
+    }
+}
+
+impl CodeModeSessionDelegate for RemoteDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        invocation: CodeModeNestedToolCall,
+        cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        Box::pin(async move {
+            match self
+                .peer
+                .call(
+                    self.session_id.clone(),
+                    DelegateRequest::InvokeTool {
+                        invocation: invocation.into(),
+                    },
+                    cancellation_token,
+                )
+                .await?
+            {
+                DelegateResponse::ToolResult { result } => Ok(result),
+                DelegateResponse::NotificationDelivered => {
+                    Err("code-mode client returned an invalid tool result".to_string())
+                }
+            }
+        })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        call_id: String,
+        cell_id: CellId,
+        text: String,
+        cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        Box::pin(async move {
+            match self
+                .peer
+                .call(
+                    self.session_id.clone(),
+                    DelegateRequest::Notify {
+                        call_id,
+                        cell_id: cell_id.into(),
+                        text,
+                    },
+                    cancellation_token,
+                )
+                .await?
+            {
+                DelegateResponse::NotificationDelivered => Ok(()),
+                DelegateResponse::ToolResult { .. } => {
+                    Err("code-mode client returned an invalid notification result".to_string())
+                }
+            }
+        })
+    }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        self.peer.send(HostToClient::CellClosed {
+            session_id: self.session_id.clone(),
+            cell_id: cell_id.into(),
+        });
+    }
+}

codex-rs/code-mode-host/src/host_tests.rs

@@ -0,0 +1,257 @@
+use codex_code_mode_protocol::host::Capability;
+use codex_code_mode_protocol::host::CapabilitySet;
+use codex_code_mode_protocol::host::ClientHello;
+use codex_code_mode_protocol::host::ClientToHost;
+use codex_code_mode_protocol::host::FramedReader;
+use codex_code_mode_protocol::host::FramedWriter;
+use codex_code_mode_protocol::host::HandshakeRejectReason;
+use codex_code_mode_protocol::host::HostHello;
+use codex_code_mode_protocol::host::HostRequest;
+use codex_code_mode_protocol::host::HostResponse;
+use codex_code_mode_protocol::host::HostToClient;
+use codex_code_mode_protocol::host::ProtocolVersion;
+use codex_code_mode_protocol::host::RequestId;
+use codex_code_mode_protocol::host::SessionId;
+use codex_code_mode_protocol::host::SupportedProtocolVersions;
+use codex_code_mode_protocol::host::WireResult;
+use pretty_assertions::assert_eq;
+
+use super::run;
+
+fn client_hello(
+    versions: impl IntoIterator<Item = ProtocolVersion>,
+    required_capabilities: CapabilitySet,
+) -> ClientToHost {
+    ClientToHost::ClientHello(
+        ClientHello::new(
+            SupportedProtocolVersions::try_new(versions).expect("supported versions"),
+            required_capabilities,
+            CapabilitySet::empty(),
+        )
+        .expect("client hello"),
+    )
+}
+
+fn session_id(value: &str) -> SessionId {
+    SessionId::new(value).expect("session ID")
+}
+
+fn request_id(value: i64) -> RequestId {
+    RequestId::new(value)
+}
+
+#[tokio::test]
+async fn handshake_and_multiple_session_lifecycles_are_ordered() {
+    let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096);
+    let (host_reader, host_writer) = tokio::io::split(host_stream);
+    let (client_reader, client_writer) = tokio::io::split(client_stream);
+    let host = tokio::spawn(run(host_reader, host_writer));
+    let mut reader = FramedReader::new(client_reader);
+    let mut writer = FramedWriter::new(client_writer);
+
+    writer
+        .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty()))
+        .await
+        .expect("write hello");
+    assert_eq!(
+        reader.read::<HostToClient>().await.expect("read hello"),
+        Some(HostToClient::HostHello(HostHello::new(
+            ProtocolVersion::V1,
+            CapabilitySet::empty(),
+        )))
+    );
+
+    for (request_id, id) in [(request_id(1), "session-1"), (request_id(2), "session-2")] {
+        writer
+            .write(&ClientToHost::Request {
+                id: request_id,
+                request: HostRequest::OpenSession {
+                    session_id: session_id(id),
+                },
+            })
+            .await
+            .expect("open session");
+        assert_eq!(
+            reader.read::<HostToClient>().await.expect("session ready"),
+            Some(HostToClient::Response {
+                id: request_id,
+                result: WireResult::Ok {
+                    value: HostResponse::SessionReady {
+                        session_id: session_id(id),
+                    },
+                },
+            })
+        );
+    }
+
+    for (request_id, id) in [(request_id(3), "session-1"), (request_id(4), "session-2")] {
+        writer
+            .write(&ClientToHost::Request {
+                id: request_id,
+                request: HostRequest::ShutdownSession {
+                    session_id: session_id(id),
+                },
+            })
+            .await
+            .expect("shutdown session");
+        assert_eq!(
+            reader.read::<HostToClient>().await.expect("session closed"),
+            Some(HostToClient::Response {
+                id: request_id,
+                result: WireResult::Ok {
+                    value: HostResponse::SessionClosed {
+                        session_id: session_id(id),
+                    },
+                },
+            })
+        );
+    }
+
+    drop(writer);
+    drop(reader);
+    host.await.expect("host task").expect("host connection");
+}
+
+#[tokio::test]
+async fn incompatible_or_invalid_handshake_is_rejected() {
+    let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
+    let (host_reader, host_writer) = tokio::io::split(host_stream);
+    let (client_reader, client_writer) = tokio::io::split(client_stream);
+    let host = tokio::spawn(run(host_reader, host_writer));
+    let mut reader = FramedReader::new(client_reader);
+    let mut writer = FramedWriter::new(client_writer);
+    let version_two = ProtocolVersion::new(/*value*/ 2).expect("protocol version");
+
+    writer
+        .write(&client_hello([version_two], CapabilitySet::empty()))
+        .await
+        .expect("write hello");
+    assert_eq!(
+        reader.read::<HostToClient>().await.expect("rejection"),
+        Some(HostToClient::HandshakeRejected {
+            reason: HandshakeRejectReason::NoCompatibleVersion {
+                supported_versions: SupportedProtocolVersions::try_new([ProtocolVersion::V1])
+                    .expect("host versions"),
+            },
+        })
+    );
+    host.await.expect("host task").expect("host connection");
+
+    let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
+    let (host_reader, host_writer) = tokio::io::split(host_stream);
+    let (client_reader, client_writer) = tokio::io::split(client_stream);
+    let host = tokio::spawn(run(host_reader, host_writer));
+    let mut reader = FramedReader::new(client_reader);
+    let mut writer = FramedWriter::new(client_writer);
+    writer
+        .write(&ClientToHost::Request {
+            id: request_id(1),
+            request: HostRequest::OpenSession {
+                session_id: session_id("session-1"),
+            },
+        })
+        .await
+        .expect("write invalid first message");
+    assert_eq!(
+        reader.read::<HostToClient>().await.expect("rejection"),
+        Some(HostToClient::HandshakeRejected {
+            reason: HandshakeRejectReason::InvalidHello {
+                message: "first message must be connection/hello".to_string(),
+            },
+        })
+    );
+    host.await.expect("host task").expect("host connection");
+}
+
+#[tokio::test]
+async fn unsupported_required_capability_is_rejected() {
+    let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024);
+    let (host_reader, host_writer) = tokio::io::split(host_stream);
+    let (client_reader, client_writer) = tokio::io::split(client_stream);
+    let host = tokio::spawn(run(host_reader, host_writer));
+    let mut reader = FramedReader::new(client_reader);
+    let mut writer = FramedWriter::new(client_writer);
+    let capability = Capability::new("required").expect("capability");
+
+    writer
+        .write(&client_hello(
+            [ProtocolVersion::V1],
+            CapabilitySet::try_new([capability.clone()]).expect("capabilities"),
+        ))
+        .await
+        .expect("write hello");
+    assert_eq!(
+        reader.read::<HostToClient>().await.expect("rejection"),
+        Some(HostToClient::HandshakeRejected {
+            reason: HandshakeRejectReason::MissingRequiredCapability { capability },
+        })
+    );
+    host.await.expect("host task").expect("host connection");
+}
+
+#[tokio::test]
+async fn session_id_cannot_be_reused_after_shutdown() {
+    let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 2048);
+    let (host_reader, host_writer) = tokio::io::split(host_stream);
+    let (client_reader, client_writer) = tokio::io::split(client_stream);
+    let host = tokio::spawn(run(host_reader, host_writer));
+    let mut reader = FramedReader::new(client_reader);
+    let mut writer = FramedWriter::new(client_writer);
+    writer
+        .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty()))
+        .await
+        .expect("write hello");
+    reader
+        .read::<HostToClient>()
+        .await
+        .expect("read hello")
+        .expect("host hello");
+
+    let id = session_id("session-1");
+    for (request_id, request) in [
+        (
+            request_id(1),
+            HostRequest::OpenSession {
+                session_id: id.clone(),
+            },
+        ),
+        (
+            request_id(2),
+            HostRequest::ShutdownSession {
+                session_id: id.clone(),
+            },
+        ),
+    ] {
+        writer
+            .write(&ClientToHost::Request {
+                id: request_id,
+                request,
+            })
+            .await
+            .expect("session request");
+        reader
+            .read::<HostToClient>()
+            .await
+            .expect("session response")
+            .expect("session response message");
+    }
+    writer
+        .write(&ClientToHost::Request {
+            id: request_id(3),
+            request: HostRequest::OpenSession { session_id: id },
+        })
+        .await
+        .expect("reuse session ID");
+    assert_eq!(
+        reader.read::<HostToClient>().await.expect("reuse response"),
+        Some(HostToClient::Response {
+            id: request_id(3),
+            result: WireResult::Err {
+                message: "code-mode session ID `session-1` was reused".to_string(),
+            },
+        })
+    );
+    drop(writer);
+    drop(reader);
+    host.await.expect("host task").expect("host connection");
+}

codex-rs/code-mode-host/src/lib.rs

@@ -0,0 +1,345 @@
+use std::collections::HashMap;
+use std::collections::HashSet;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::PoisonError;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+
+use anyhow::Context;
+use anyhow::Result;
+use codex_code_mode::InProcessCodeModeSession;
+use codex_code_mode_protocol::host::CapabilitySet;
+use codex_code_mode_protocol::host::ClientToHost;
+use codex_code_mode_protocol::host::FramedReader;
+use codex_code_mode_protocol::host::FramedWriter;
+use codex_code_mode_protocol::host::HandshakeRejectReason;
+use codex_code_mode_protocol::host::HostHello;
+use codex_code_mode_protocol::host::HostRequest;
+use codex_code_mode_protocol::host::HostResponse;
+use codex_code_mode_protocol::host::HostToClient;
+use codex_code_mode_protocol::host::ProtocolVersion;
+use codex_code_mode_protocol::host::RequestId;
+use codex_code_mode_protocol::host::SessionId;
+use codex_code_mode_protocol::host::SupportedProtocolVersions;
+use codex_code_mode_protocol::host::WireResult;
+use tokio::io::AsyncRead;
+use tokio::io::AsyncWrite;
+use tokio::sync::mpsc;
+use tokio_util::task::TaskTracker;
+
+use self::delegate::RemoteDelegate;
+use self::peer::HostPeer;
+
+mod delegate;
+mod peer;
+
+/// Runs one code-mode host connection over the process standard streams.
+pub async fn run_stdio() -> Result<()> {
+    run(tokio::io::stdin(), tokio::io::stdout()).await
+}
+
+/// Runs one code-mode host connection over an ordered input/output pair.
+async fn run<R, W>(reader: R, writer: W) -> Result<()>
+where
+    R: AsyncRead + Send + Unpin + 'static,
+    W: AsyncWrite + Send + Unpin + 'static,
+{
+    let mut reader = FramedReader::new(reader);
+    let mut writer = FramedWriter::new(writer);
+    if !negotiate(&mut reader, &mut writer).await? {
+        return Ok(());
+    }
+
+    let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel();
+    let peer = Arc::new(HostPeer::new(outgoing_tx));
+    let state = Arc::new(HostState {
+        sessions: Mutex::new(HashMap::new()),
+        seen_session_ids: Mutex::new(HashSet::new()),
+        request_tasks: TaskTracker::new(),
+        closing: AtomicBool::new(false),
+        peer: Arc::clone(&peer),
+    });
+    let writer_task = tokio::spawn(async move {
+        while let Some(message) = outgoing_rx.recv().await {
+            writer
+                .write(&message)
+                .await
+                .context("failed to write code-mode host message")?;
+        }
+        Ok::<(), anyhow::Error>(())
+    });
+
+    let input_result = async {
+        while let Some(message) = reader
+            .read::<ClientToHost>()
+            .await
+            .context("failed to read code-mode client message")?
+        {
+            match message {
+                ClientToHost::ClientHello(_) => {
+                    anyhow::bail!("received a second code-mode client hello");
+                }
+                ClientToHost::Request { id, request } => {
+                    state.spawn_request(id, request);
+                }
+                ClientToHost::DelegateResponse { id, result } => {
+                    peer.complete(id, result.into_result()).await;
+                }
+            }
+        }
+        Ok::<(), anyhow::Error>(())
+    }
+    .await;
+
+    peer.disconnect();
+    state.disconnect().await;
+    drop(state);
+    drop(peer);
+    let writer_result = writer_task.await.context("code-mode writer task failed")?;
+    input_result?;
+    writer_result
+}
+
+async fn negotiate<R, W>(reader: &mut FramedReader<R>, writer: &mut FramedWriter<W>) -> Result<bool>
+where
+    R: AsyncRead + Unpin,
+    W: AsyncWrite + Unpin,
+{
+    let Some(first_message) = reader
+        .read::<ClientToHost>()
+        .await
+        .context("failed to read code-mode client hello")?
+    else {
+        return Ok(false);
+    };
+    let ClientToHost::ClientHello(client_hello) = first_message else {
+        writer
+            .write(&HostToClient::HandshakeRejected {
+                reason: HandshakeRejectReason::InvalidHello {
+                    message: "first message must be connection/hello".to_string(),
+                },
+            })
+            .await
+            .context("failed to reject invalid code-mode client hello")?;
+        return Ok(false);
+    };
+
+    let supported_versions = SupportedProtocolVersions::try_new([ProtocolVersion::V1])?;
+    if !client_hello
+        .supported_versions()
+        .contains(ProtocolVersion::V1)
+    {
+        writer
+            .write(&HostToClient::HandshakeRejected {
+                reason: HandshakeRejectReason::NoCompatibleVersion { supported_versions },
+            })
+            .await
+            .context("failed to reject incompatible code-mode client")?;
+        return Ok(false);
+    }
+
+    let host_capabilities = CapabilitySet::empty();
+    if let Some(capability) = client_hello
+        .required_capabilities()
+        .iter()
+        .find(|capability| !host_capabilities.contains(capability))
+    {
+        writer
+            .write(&HostToClient::HandshakeRejected {
+                reason: HandshakeRejectReason::MissingRequiredCapability {
+                    capability: capability.clone(),
+                },
+            })
+            .await
+            .context("failed to reject unsupported code-mode capability")?;
+        return Ok(false);
+    }
+
+    writer
+        .write(&HostToClient::HostHello(HostHello::new(
+            ProtocolVersion::V1,
+            host_capabilities,
+        )))
+        .await
+        .context("failed to write code-mode host hello")?;
+    Ok(true)
+}
+
+struct HostState {
+    sessions: Mutex<HashMap<SessionId, Arc<InProcessCodeModeSession>>>,
+    seen_session_ids: Mutex<HashSet<SessionId>>,
+    request_tasks: TaskTracker,
+    closing: AtomicBool,
+    peer: Arc<HostPeer>,
+}
+
+impl HostState {
+    fn spawn_request(self: &Arc<Self>, request_id: RequestId, request: HostRequest) {
+        let state = Arc::clone(self);
+        self.request_tasks.spawn(async move {
+            state.handle_request(request_id, request).await;
+        });
+    }
+
+    async fn handle_request(&self, request_id: RequestId, request: HostRequest) {
+        if self.closing.load(Ordering::Acquire) {
+            self.respond(
+                request_id,
+                Err("code-mode host is shutting down".to_string()),
+            );
+            return;
+        }
+        match request {
+            HostRequest::OpenSession { session_id } => {
+                let result = self
+                    .open_session(session_id.clone())
+                    .map(|()| HostResponse::SessionReady { session_id });
+                self.respond(request_id, result);
+            }
+            HostRequest::Execute {
+                session_id,
+                request,
+            } => {
+                let request = match request.try_into() {
+                    Ok(request) => request,
+                    Err(err) => {
+                        self.respond(
+                            request_id,
+                            Err(format!("invalid code-mode execute request: {err}")),
+                        );
+                        return;
+                    }
+                };
+                let result = match self.session(&session_id) {
+                    Ok(session) => session.execute(request).await,
+                    Err(err) => {
+                        self.respond(request_id, Err(err));
+                        return;
+                    }
+                };
+                match result {
+                    Ok(started) => {
+                        let cell_id = started.cell_id.clone();
+                        self.respond(
+                            request_id,
+                            Ok(HostResponse::ExecutionStarted {
+                                cell_id: cell_id.into(),
+                            }),
+                        );
+                        self.peer.send(HostToClient::InitialResponse {
+                            id: request_id,
+                            result: WireResult::from_result(
+                                started.initial_response().await.map(Into::into),
+                            ),
+                        });
+                    }
+                    Err(err) => self.respond(request_id, Err(err)),
+                }
+            }
+            HostRequest::Wait {
+                session_id,
+                request,
+            } => {
+                let result = match self.session(&session_id) {
+                    Ok(session) => session.wait(request.into()).await.map(|outcome| {
+                        HostResponse::WaitCompleted {
+                            outcome: outcome.into(),
+                        }
+                    }),
+                    Err(err) => Err(err),
+                };
+                self.respond(request_id, result);
+            }
+            HostRequest::Terminate {
+                session_id,
+                cell_id,
+            } => {
+                let result = match self.session(&session_id) {
+                    Ok(session) => session.terminate(cell_id.into()).await.map(|outcome| {
+                        HostResponse::WaitCompleted {
+                            outcome: outcome.into(),
+                        }
+                    }),
+                    Err(err) => Err(err),
+                };
+                self.respond(request_id, result);
+            }
+            HostRequest::ShutdownSession { session_id } => {
+                let session = self
+                    .sessions
+                    .lock()
+                    .unwrap_or_else(PoisonError::into_inner)
+                    .remove(&session_id);
+                let result = match session {
+                    Some(session) => session
+                        .shutdown()
+                        .await
+                        .map(|()| HostResponse::SessionClosed { session_id }),
+                    None => Err(format!("unknown code-mode session {session_id}")),
+                };
+                self.respond(request_id, result);
+            }
+        }
+    }
+
+    fn open_session(&self, session_id: SessionId) -> Result<(), String> {
+        let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner);
+        if self.closing.load(Ordering::Acquire) {
+            return Err("code-mode host is shutting down".to_string());
+        }
+        if !self
+            .seen_session_ids
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .insert(session_id.clone())
+        {
+            return Err(format!("code-mode session ID `{session_id}` was reused"));
+        }
+        let delegate = Arc::new(RemoteDelegate::new(
+            session_id.clone(),
+            Arc::clone(&self.peer),
+        ));
+        sessions.insert(
+            session_id,
+            Arc::new(InProcessCodeModeSession::with_delegate(delegate)),
+        );
+        Ok(())
+    }
+
+    fn session(&self, session_id: &SessionId) -> Result<Arc<InProcessCodeModeSession>, String> {
+        self.sessions
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .get(session_id)
+            .cloned()
+            .ok_or_else(|| format!("unknown code-mode session {session_id}"))
+    }
+
+    fn respond(&self, id: RequestId, result: Result<HostResponse, String>) {
+        self.peer.send(HostToClient::Response {
+            id,
+            result: WireResult::from_result(result),
+        });
+    }
+
+    async fn disconnect(&self) {
+        self.closing.store(true, Ordering::Release);
+        let sessions = self
+            .sessions
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .drain()
+            .map(|(_, session)| session)
+            .collect::<Vec<_>>();
+        for session in sessions {
+            let _ = session.shutdown().await;
+        }
+        self.request_tasks.close();
+        self.request_tasks.wait().await;
+    }
+}
+
+#[cfg(test)]
+#[path = "host_tests.rs"]
+mod tests;

codex-rs/code-mode-host/src/main.rs

@@ -1 +1,4 @@
-fn main() {}
+#[tokio::main(flavor = "current_thread")]
+async fn main() -> anyhow::Result<()> {
+    codex_code_mode_host::run_stdio().await
+}

codex-rs/code-mode-host/src/peer.rs

@@ -0,0 +1,124 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::sync::atomic::AtomicI64;
+use std::sync::atomic::Ordering;
+
+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::SessionId;
+use tokio::sync::Mutex;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+
+pub(super) struct HostPeer {
+    outgoing_tx: mpsc::UnboundedSender<HostToClient>,
+    pending: Mutex<HashMap<DelegateRequestId, oneshot::Sender<Result<DelegateResponse, String>>>>,
+    next_request_id: AtomicI64,
+    disconnected: CancellationToken,
+}
+
+impl HostPeer {
+    pub(super) fn new(outgoing_tx: mpsc::UnboundedSender<HostToClient>) -> Self {
+        Self {
+            outgoing_tx,
+            pending: Mutex::new(HashMap::new()),
+            next_request_id: AtomicI64::new(1),
+            disconnected: CancellationToken::new(),
+        }
+    }
+
+    pub(super) fn send(&self, message: HostToClient) -> bool {
+        self.outgoing_tx.send(message).is_ok()
+    }
+
+    pub(super) async fn call(
+        self: &Arc<Self>,
+        session_id: SessionId,
+        request: DelegateRequest,
+        cancellation_token: CancellationToken,
+    ) -> Result<DelegateResponse, String> {
+        if self.disconnected.is_cancelled() {
+            return Err("code-mode client connection closed".to_string());
+        }
+        let id = DelegateRequestId::new(self.next_request_id.fetch_add(1, Ordering::Relaxed));
+        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());
+        }
+
+        tokio::select! {
+            response = response_rx => {
+                pending.disarm();
+                response.map_err(|_| {
+                    "code-mode client closed before returning delegate output".to_string()
+                })?
+            }
+            _ = cancellation_token.cancelled() => {
+                if self.pending.lock().await.remove(&id).is_some() {
+                    self.send(HostToClient::CancelDelegateRequest { id });
+                }
+                pending.disarm();
+                Err("code mode delegate request cancelled".to_string())
+            }
+            _ = self.disconnected.cancelled() => {
+                self.pending.lock().await.remove(&id);
+                pending.disarm();
+                Err("code-mode client connection closed".to_string())
+            }
+        }
+    }
+
+    pub(super) async fn complete(
+        &self,
+        id: DelegateRequestId,
+        response: Result<DelegateResponse, String>,
+    ) {
+        if let Some(sender) = self.pending.lock().await.remove(&id) {
+            let _ = sender.send(response);
+        }
+    }
+
+    pub(super) fn disconnect(&self) {
+        self.disconnected.cancel();
+    }
+}
+
+struct PendingDelegateRequest {
+    peer: Arc<HostPeer>,
+    id: Option<DelegateRequestId>,
+}
+
+impl PendingDelegateRequest {
+    fn new(peer: Arc<HostPeer>, id: DelegateRequestId) -> Self {
+        Self { peer, id: Some(id) }
+    }
+
+    fn disarm(&mut self) {
+        self.id = None;
+    }
+}
+
+impl Drop for PendingDelegateRequest {
+    fn drop(&mut self) {
+        let Some(id) = self.id.take() else {
+            return;
+        };
+        let peer = Arc::clone(&self.peer);
+        tokio::spawn(async move {
+            if peer.pending.lock().await.remove(&id).is_some() {
+                peer.send(HostToClient::CancelDelegateRequest { id });
+            }
+        });
+    }
+}

codex-rs/code-mode-host/tests/stdio.rs

@@ -0,0 +1,192 @@
+#![allow(clippy::expect_used)]
+
+use std::sync::Arc;
+use std::sync::Mutex;
+
+use codex_code_mode::CellId;
+use codex_code_mode::CodeModeNestedToolCall;
+use codex_code_mode::CodeModeSession;
+use codex_code_mode::CodeModeSessionDelegate;
+use codex_code_mode::CodeModeSessionProvider;
+use codex_code_mode::CodeModeToolKind;
+use codex_code_mode::ExecuteRequest;
+use codex_code_mode::FunctionCallOutputContentItem;
+use codex_code_mode::NotificationFuture;
+use codex_code_mode::ProcessOwnedCodeModeSessionProvider;
+use codex_code_mode::RuntimeResponse;
+use codex_code_mode::ToolDefinition;
+use codex_code_mode::ToolInvocationFuture;
+use codex_code_mode::WaitOutcome;
+use codex_code_mode::WaitRequest;
+use codex_protocol::ToolName;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+use tokio_util::sync::CancellationToken;
+
+#[derive(Default)]
+struct RecordingDelegate {
+    invocations: Mutex<Vec<CodeModeNestedToolCall>>,
+    notifications: Mutex<Vec<(String, CellId, String)>>,
+    closed_cells: Mutex<Vec<CellId>>,
+}
+
+impl CodeModeSessionDelegate for RecordingDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        invocation: CodeModeNestedToolCall,
+        _cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        self.invocations
+            .lock()
+            .expect("invocations lock")
+            .push(invocation);
+        Box::pin(async { Ok(json!({ "value": "output" })) })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        call_id: String,
+        cell_id: CellId,
+        text: String,
+        _cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        self.notifications
+            .lock()
+            .expect("notifications lock")
+            .push((call_id, cell_id, text));
+        Box::pin(async { Ok(()) })
+    }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        self.closed_cells
+            .lock()
+            .expect("closed cells lock")
+            .push(cell_id.clone());
+    }
+}
+
+fn cell_id(value: &str) -> CellId {
+    CellId::new(value.to_string())
+}
+
+fn execute_request(source: &str) -> ExecuteRequest {
+    ExecuteRequest {
+        tool_call_id: "call-1".to_string(),
+        enabled_tools: Vec::new(),
+        source: source.to_string(),
+        yield_time_ms: None,
+        max_output_tokens: None,
+    }
+}
+
+async fn execute(session: &Arc<dyn CodeModeSession>, request: ExecuteRequest) -> RuntimeResponse {
+    session
+        .execute(request)
+        .await
+        .expect("start execution")
+        .initial_response()
+        .await
+        .expect("initial response")
+}
+
+#[tokio::test]
+async fn remote_session_persists_values_forwards_delegates_and_controls_cells() {
+    let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(
+        codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"),
+    );
+    let delegate = Arc::new(RecordingDelegate::default());
+    let session = provider
+        .create_session(delegate.clone())
+        .await
+        .expect("create remote session");
+
+    assert_eq!(
+        execute(&session, execute_request(r#"store("key", "persisted");"#),).await,
+        RuntimeResponse::Result {
+            cell_id: cell_id("1"),
+            content_items: Vec::new(),
+            error_text: None,
+        }
+    );
+
+    let mut callback_request = execute_request(
+        r#"
+const result = await tools.echo({ value: String(load("key")) });
+notify("notice");
+text(result.value);
+"#,
+    );
+    callback_request.tool_call_id = "call-2".to_string();
+    callback_request.enabled_tools = vec![ToolDefinition {
+        name: "echo".to_string(),
+        tool_name: ToolName::plain("echo"),
+        description: String::new(),
+        kind: CodeModeToolKind::Function,
+        input_schema: None,
+        output_schema: None,
+    }];
+    assert_eq!(
+        execute(&session, callback_request).await,
+        RuntimeResponse::Result {
+            cell_id: cell_id("2"),
+            content_items: vec![FunctionCallOutputContentItem::InputText {
+                text: "output".to_string(),
+            }],
+            error_text: None,
+        }
+    );
+    assert_eq!(
+        *delegate.invocations.lock().expect("invocations lock"),
+        vec![CodeModeNestedToolCall {
+            cell_id: cell_id("2"),
+            runtime_tool_call_id: "tool-1".to_string(),
+            tool_name: ToolName::plain("echo"),
+            tool_kind: CodeModeToolKind::Function,
+            input: Some(json!({ "value": "persisted" })),
+        }]
+    );
+    assert_eq!(
+        *delegate.notifications.lock().expect("notifications lock"),
+        vec![("call-2".to_string(), cell_id("2"), "notice".to_string())]
+    );
+
+    let mut pending_request = execute_request("await new Promise(() => {});");
+    pending_request.tool_call_id = "call-3".to_string();
+    pending_request.yield_time_ms = Some(1);
+    assert_eq!(
+        execute(&session, pending_request).await,
+        RuntimeResponse::Yielded {
+            cell_id: cell_id("3"),
+            content_items: Vec::new(),
+        }
+    );
+    assert_eq!(
+        session
+            .wait(WaitRequest {
+                cell_id: cell_id("3"),
+                yield_time_ms: 1,
+            })
+            .await
+            .expect("wait for cell"),
+        WaitOutcome::LiveCell(RuntimeResponse::Yielded {
+            cell_id: cell_id("3"),
+            content_items: Vec::new(),
+        })
+    );
+    assert_eq!(
+        session
+            .terminate(cell_id("3"))
+            .await
+            .expect("terminate cell"),
+        WaitOutcome::LiveCell(RuntimeResponse::Terminated {
+            cell_id: cell_id("3"),
+            content_items: Vec::new(),
+        })
+    );
+
+    session.shutdown().await.expect("shutdown remote session");
+    assert_eq!(
+        *delegate.closed_cells.lock().expect("closed cells lock"),
+        vec![cell_id("1"), cell_id("2"), cell_id("3")]
+    );
+}