[codex] add process-owned code-mode session client (#30112)

## Summary

- add `ProcessOwnedCodeModeSessionProvider` and logical session
generation/rebinding state
- add the supervised child-process connection, reader/writer tasks, and
driver state machine
- make dropped execute/wait/open callers cancellation-safe with explicit
ownership handoff and durable cleanup
- validate cell/delegate lifecycle state and reject invalid protocol
transitions
- add end-to-end stdio coverage for delegates, cancellation, frame
limits, child loss, stale generations, replacement, and long-lived
sessions

## Why

This final stage exposes the process-owned client only after the wire
protocol, host-safe runtime, and standalone host are independently in
place. Transport failure is fail-stop: the client closes local state,
cancels callbacks, reaps the child, and lazily rebuilds a fresh host
generation rather than transactionally recovering the old connection.

## Stack

This is **4 of 4** in the process-owned code-mode session stack.

- Depends on #30111
- Full stack: #30108 → #30110 → #30111 → this PR

## Validation

- `just test -p codex-code-mode -p codex-code-mode-host` — 86 passed
- `just fix -p codex-code-mode`
- `just fix -p codex-code-mode-host`
- `just bazel-lock-update`
- `just bazel-lock-check`
- `bazel test //codex-rs/code-mode:code-mode-unit-tests
//codex-rs/code-mode-host:code-mode-host-unit-tests
//codex-rs/code-mode-host:code-mode-host-stdio-test
//codex-rs/code-mode-protocol:code-mode-protocol-unit-tests` — 4/4
passed
- `just fmt`

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

@@ -1,4 +1,5 @@
 mod cell_actor;
+mod remote_session;
 mod runtime;
 mod service;
 mod session_runtime;
@@ -6,6 +7,8 @@ mod session_runtime;
 pub(crate) type TaskFailureHandler = std::sync::Arc<dyn Fn(String) + Send + Sync>;
 
 pub use codex_code_mode_protocol::*;
+pub use remote_session::ProcessOwnedCodeModeSession;
+pub use remote_session::ProcessOwnedCodeModeSessionProvider;
 pub use service::InProcessCodeModeSession;
 pub use service::InProcessCodeModeSessionProvider;
 pub use service::NoopCodeModeSessionDelegate;

codex-rs/code-mode/src/remote_session.rs

@@ -0,0 +1,496 @@
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::sync::Mutex as StdMutex;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeSession;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::CodeModeSessionProvider;
+use codex_code_mode_protocol::CodeModeSessionProviderFuture;
+use codex_code_mode_protocol::CodeModeSessionResultFuture;
+use codex_code_mode_protocol::ExecuteRequest;
+use codex_code_mode_protocol::StartedCell;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_code_mode_protocol::WaitRequest;
+use codex_code_mode_protocol::host::SessionId;
+use tokio::sync::Semaphore;
+use tokio::sync::watch;
+
+use self::connection::Connection;
+use self::connection::RemoteSession;
+use self::connection::SessionCleanup;
+use crate::NoopCodeModeSessionDelegate;
+
+mod connection;
+
+const CODE_MODE_HOST_PATH_ENV: &str = "CODEX_CODE_MODE_HOST_PATH";
+
+type ShutdownResultReceiver = watch::Receiver<Option<Result<(), String>>>;
+
+/// Creates code-mode sessions backed by one lazily spawned process host.
+pub struct ProcessOwnedCodeModeSessionProvider {
+    host_program: PathBuf,
+    process_host: StdMutex<Option<Arc<OwnedProcessHost>>>,
+}
+
+impl ProcessOwnedCodeModeSessionProvider {
+    pub fn with_host_program(host_program: PathBuf) -> Self {
+        Self {
+            host_program,
+            process_host: StdMutex::new(None),
+        }
+    }
+
+    fn process_host(&self) -> Arc<OwnedProcessHost> {
+        let mut process_host = self
+            .process_host
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if let Some(process_host) = process_host.as_ref() {
+            return Arc::clone(process_host);
+        }
+
+        let new_process_host = Arc::new(OwnedProcessHost::new(self.host_program.clone()));
+        *process_host = Some(Arc::clone(&new_process_host));
+        new_process_host
+    }
+}
+
+impl Default for ProcessOwnedCodeModeSessionProvider {
+    fn default() -> Self {
+        Self::with_host_program(default_host_program())
+    }
+}
+
+impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider {
+    fn create_session<'a>(
+        &'a self,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+    ) -> CodeModeSessionProviderFuture<'a> {
+        let session = ProcessOwnedCodeModeSession::with_process_host(delegate, self.process_host());
+        Box::pin(async move {
+            session.connection().await?;
+            let session: Arc<dyn CodeModeSession> = Arc::new(session);
+            Ok(session)
+        })
+    }
+}
+
+struct OwnedProcessHost {
+    host_program: PathBuf,
+    connection: StdMutex<Option<Arc<Connection>>>,
+    spawn_permit: Semaphore,
+    next_session_id: AtomicU64,
+}
+
+impl OwnedProcessHost {
+    fn new(host_program: PathBuf) -> Self {
+        Self {
+            host_program,
+            connection: StdMutex::new(None),
+            spawn_permit: Semaphore::new(/*permits*/ 1),
+            next_session_id: AtomicU64::new(1),
+        }
+    }
+
+    async fn connection(&self) -> Result<Arc<Connection>, String> {
+        if let Some(connection) = self.live_connection() {
+            return Ok(connection);
+        }
+
+        let _spawn_permit = self
+            .spawn_permit
+            .acquire()
+            .await
+            .map_err(|_| "code-mode host spawn coordinator closed".to_string())?;
+        if let Some(connection) = self.live_connection() {
+            return Ok(connection);
+        }
+        let new_connection = Arc::new(Connection::spawn(&self.host_program).await?);
+        *self
+            .connection
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Arc::clone(&new_connection));
+        Ok(new_connection)
+    }
+
+    fn live_connection(&self) -> Option<Arc<Connection>> {
+        self.connection
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .as_ref()
+            .filter(|connection| connection.is_alive())
+            .cloned()
+    }
+
+    fn allocate_session_id(&self) -> SessionId {
+        let value = self.next_session_id.fetch_add(1, Ordering::Relaxed);
+        match SessionId::new(format!("session-{value}")) {
+            Ok(session_id) => session_id,
+            Err(_) => unreachable!("a generated code-mode session ID is nonempty"),
+        }
+    }
+}
+
+enum SessionState {
+    New,
+    Opening {
+        remote: RemoteSession,
+        result_rx: watch::Receiver<Option<Result<SessionBinding, String>>>,
+    },
+    Open(SessionBinding),
+    Closing,
+    Closed,
+}
+
+#[derive(Clone)]
+struct SessionBinding {
+    connection: Arc<Connection>,
+    remote: RemoteSession,
+    cleanup: SessionCleanup,
+}
+
+struct SessionInner {
+    process_host: Arc<OwnedProcessHost>,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
+    state: StdMutex<SessionState>,
+    next_generation: AtomicU64,
+    shutdown_requested: AtomicBool,
+    shutdown_result: StdMutex<Option<ShutdownResultReceiver>>,
+    retired_cleanups: StdMutex<Vec<SessionCleanup>>,
+}
+
+/// A logical code-mode session assigned to a process-owned host.
+pub struct ProcessOwnedCodeModeSession {
+    inner: Arc<SessionInner>,
+}
+
+impl ProcessOwnedCodeModeSession {
+    pub fn new() -> Self {
+        Self::with_process_host(
+            Arc::new(NoopCodeModeSessionDelegate),
+            Arc::new(OwnedProcessHost::new(default_host_program())),
+        )
+    }
+
+    fn with_process_host(
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+        process_host: Arc<OwnedProcessHost>,
+    ) -> Self {
+        Self {
+            inner: Arc::new(SessionInner {
+                process_host,
+                delegate,
+                state: StdMutex::new(SessionState::New),
+                next_generation: AtomicU64::new(1),
+                shutdown_requested: AtomicBool::new(false),
+                shutdown_result: StdMutex::new(None),
+                retired_cleanups: StdMutex::new(Vec::new()),
+            }),
+        }
+    }
+
+    async fn connection(&self) -> Result<SessionBinding, String> {
+        self.inner.connection().await
+    }
+
+    pub async fn execute(&self, request: ExecuteRequest) -> Result<StartedCell, String> {
+        let binding = self.connection().await?;
+        binding.connection.execute(binding.remote, request).await
+    }
+
+    pub async fn wait(&self, request: WaitRequest) -> Result<WaitOutcome, String> {
+        let binding = self.connection().await?;
+        binding.connection.wait(binding.remote, request).await
+    }
+
+    pub async fn terminate(&self, cell_id: CellId) -> Result<WaitOutcome, String> {
+        let binding = self.connection().await?;
+        binding.connection.terminate(binding.remote, cell_id).await
+    }
+
+    pub async fn shutdown(&self) -> Result<(), String> {
+        wait_for_watch(self.inner.request_shutdown()).await
+    }
+}
+
+impl SessionInner {
+    async fn connection(self: &Arc<Self>) -> Result<SessionBinding, String> {
+        loop {
+            if self.shutdown_requested.load(Ordering::Acquire) {
+                return Err("code mode session is shutting down".to_string());
+            }
+            let (result_rx, start) = {
+                let mut state = self
+                    .state
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner);
+                match &*state {
+                    SessionState::New => {
+                        let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
+                        let remote = RemoteSession {
+                            id: self.process_host.allocate_session_id(),
+                            generation,
+                        };
+                        let (result_tx, result_rx) = watch::channel(None);
+                        *state = SessionState::Opening {
+                            remote: remote.clone(),
+                            result_rx: result_rx.clone(),
+                        };
+                        (result_rx, Some((remote, result_tx)))
+                    }
+                    SessionState::Opening { result_rx, .. } => (result_rx.clone(), None),
+                    SessionState::Open(binding) if binding.connection.is_alive() => {
+                        return Ok(binding.clone());
+                    }
+                    SessionState::Open(binding) => {
+                        self.retain_cleanup(binding.cleanup.clone());
+                        *state = SessionState::New;
+                        continue;
+                    }
+                    SessionState::Closing | SessionState::Closed => {
+                        return Err("code mode session is shutting down".to_string());
+                    }
+                }
+            };
+            if let Some((remote, result_tx)) = start {
+                let inner = Arc::clone(self);
+                tokio::spawn(async move {
+                    inner.open(remote, result_tx).await;
+                });
+            }
+            return wait_for_watch(result_rx).await;
+        }
+    }
+
+    async fn open(
+        self: Arc<Self>,
+        remote: RemoteSession,
+        result_tx: watch::Sender<Option<Result<SessionBinding, String>>>,
+    ) {
+        let result = match self.process_host.connection().await {
+            Ok(connection) => {
+                let cleanup = connection
+                    .open_session(remote.clone(), Arc::clone(&self.delegate))
+                    .await;
+                cleanup.map(|cleanup| SessionBinding {
+                    connection,
+                    remote: remote.clone(),
+                    cleanup,
+                })
+            }
+            Err(err) => Err(err),
+        };
+        {
+            let mut state = self
+                .state
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            if matches!(
+                &*state,
+                SessionState::Opening {
+                    remote: opening_remote,
+                    ..
+                } if opening_remote == &remote
+            ) {
+                *state = match &result {
+                    Ok(binding) => SessionState::Open(binding.clone()),
+                    Err(_) => SessionState::New,
+                };
+            }
+        }
+        result_tx.send_replace(Some(result));
+    }
+
+    fn request_shutdown(self: &Arc<Self>) -> ShutdownResultReceiver {
+        self.shutdown_requested.store(true, Ordering::Release);
+        let mut shutdown_result = self
+            .shutdown_result
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if let Some(result_rx) = shutdown_result.as_ref() {
+            return result_rx.clone();
+        }
+        let (result_tx, result_rx) = watch::channel(None);
+        *shutdown_result = Some(result_rx.clone());
+        let inner = Arc::clone(self);
+        tokio::spawn(async move {
+            let result = inner.drive_shutdown().await;
+            result_tx.send_replace(Some(result));
+        });
+        result_rx
+    }
+
+    async fn drive_shutdown(self: &Arc<Self>) -> Result<(), String> {
+        loop {
+            let action = {
+                let mut state = self
+                    .state
+                    .lock()
+                    .unwrap_or_else(std::sync::PoisonError::into_inner);
+                match &*state {
+                    SessionState::New => {
+                        *state = SessionState::Closed;
+                        ShutdownAction::Finish
+                    }
+                    SessionState::Opening { result_rx, .. } => {
+                        ShutdownAction::WaitForOpen(result_rx.clone())
+                    }
+                    SessionState::Open(binding) if !binding.connection.is_alive() => {
+                        let cleanup = binding.cleanup.clone();
+                        *state = SessionState::Closing;
+                        ShutdownAction::WaitForSessionCleanup(cleanup)
+                    }
+                    SessionState::Open(binding) => {
+                        let binding = binding.clone();
+                        *state = SessionState::Closing;
+                        ShutdownAction::Close(binding)
+                    }
+                    SessionState::Closing => {
+                        return Err("code-mode session shutdown driver entered twice".to_string());
+                    }
+                    SessionState::Closed => return Ok(()),
+                }
+            };
+            match action {
+                ShutdownAction::WaitForOpen(result_rx) => {
+                    let _ = wait_for_watch(result_rx).await;
+                }
+                ShutdownAction::Finish => {
+                    self.wait_for_retired_cleanups().await;
+                    return Ok(());
+                }
+                ShutdownAction::WaitForSessionCleanup(cleanup) => {
+                    cleanup.wait().await;
+                    self.wait_for_retired_cleanups().await;
+                    *self
+                        .state
+                        .lock()
+                        .unwrap_or_else(std::sync::PoisonError::into_inner) = SessionState::Closed;
+                    return Ok(());
+                }
+                ShutdownAction::Close(binding) => {
+                    let result = binding.connection.shutdown_session(binding.remote).await;
+                    if result.is_err() && !binding.connection.is_alive() {
+                        binding.cleanup.wait().await;
+                    }
+                    self.wait_for_retired_cleanups().await;
+                    *self
+                        .state
+                        .lock()
+                        .unwrap_or_else(std::sync::PoisonError::into_inner) = SessionState::Closed;
+                    return result;
+                }
+            }
+        }
+    }
+
+    fn retain_cleanup(&self, cleanup: SessionCleanup) {
+        let mut retired = self
+            .retired_cleanups
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        retired.retain(|cleanup| !cleanup.is_complete());
+        if !cleanup.is_complete() {
+            retired.push(cleanup);
+        }
+    }
+
+    async fn wait_for_retired_cleanups(&self) {
+        let retired = std::mem::take(
+            &mut *self
+                .retired_cleanups
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner),
+        );
+        for cleanup in retired {
+            cleanup.wait().await;
+        }
+    }
+}
+
+enum ShutdownAction {
+    WaitForOpen(watch::Receiver<Option<Result<SessionBinding, String>>>),
+    Finish,
+    WaitForSessionCleanup(SessionCleanup),
+    Close(SessionBinding),
+}
+
+async fn wait_for_watch<T>(
+    mut result_rx: watch::Receiver<Option<Result<T, String>>>,
+) -> Result<T, String>
+where
+    T: Clone,
+{
+    loop {
+        if let Some(result) = result_rx.borrow().clone() {
+            return result;
+        }
+        result_rx
+            .changed()
+            .await
+            .map_err(|_| "code-mode session transition stopped".to_string())?;
+    }
+}
+
+impl Drop for ProcessOwnedCodeModeSession {
+    fn drop(&mut self) {
+        if tokio::runtime::Handle::try_current().is_ok() {
+            self.inner.request_shutdown();
+        }
+    }
+}
+
+impl Default for ProcessOwnedCodeModeSession {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl CodeModeSession for ProcessOwnedCodeModeSession {
+    fn execute<'a>(
+        &'a self,
+        request: ExecuteRequest,
+    ) -> CodeModeSessionResultFuture<'a, StartedCell> {
+        Box::pin(ProcessOwnedCodeModeSession::execute(self, request))
+    }
+
+    fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
+        Box::pin(ProcessOwnedCodeModeSession::wait(self, request))
+    }
+
+    fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
+        Box::pin(ProcessOwnedCodeModeSession::terminate(self, cell_id))
+    }
+
+    fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
+        Box::pin(ProcessOwnedCodeModeSession::shutdown(self))
+    }
+}
+
+fn default_host_program() -> PathBuf {
+    if let Some(path) = std::env::var_os(CODE_MODE_HOST_PATH_ENV) {
+        return PathBuf::from(path);
+    }
+    let executable_name = if cfg!(windows) {
+        "codex-code-mode-host.exe"
+    } else {
+        "codex-code-mode-host"
+    };
+    if let Ok(current_exe) = std::env::current_exe()
+        && let Some(parent) = current_exe.parent()
+    {
+        let sibling = parent.join(executable_name);
+        if sibling.is_file() {
+            return sibling;
+        }
+    }
+    PathBuf::from(executable_name)
+}
+
+#[cfg(test)]
+#[path = "remote_session_tests.rs"]
+mod tests;

codex-rs/code-mode/src/remote_session/connection.rs

@@ -0,0 +1,458 @@
+use std::path::Path;
+use std::process::Stdio;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+use std::time::Duration;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::ExecuteRequest;
+use codex_code_mode_protocol::StartedCell;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_code_mode_protocol::WaitRequest;
+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::EncodedFrame;
+use codex_code_mode_protocol::host::FramedReader;
+use codex_code_mode_protocol::host::FramedWriter;
+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::SupportedProtocolVersions;
+use tokio::io::AsyncBufReadExt;
+use tokio::io::BufReader;
+use tokio::process::Child;
+use tokio::process::Command;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio::task::JoinHandle;
+use tokio_util::sync::CancellationToken;
+use tracing::debug;
+use tracing::warn;
+
+use self::driver::ConnectionDriver;
+use self::driver::DriverCommand;
+use self::driver::DriverEvent;
+use self::driver::DriverLifecycle;
+pub(super) use self::driver::RemoteSession;
+pub(super) use self::driver::SessionCleanup;
+use self::reader::drive_reader;
+
+mod driver;
+mod reader;
+
+const IPC_CHANNEL_CAPACITY: usize = 128;
+const HOST_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
+
+pub(super) struct Connection {
+    command_tx: mpsc::Sender<DriverCommand>,
+    execute_claim_tx: mpsc::UnboundedSender<RequestId>,
+    alive: Arc<AtomicBool>,
+    failure: Arc<std::sync::Mutex<Option<String>>>,
+    cancellation: CancellationToken,
+}
+
+struct CallerCancellation {
+    token: CancellationToken,
+    armed: bool,
+}
+
+struct ConnectionSupervisor {
+    child: Child,
+    event_tx: mpsc::Sender<DriverEvent>,
+    cancellation: CancellationToken,
+    alive: Arc<AtomicBool>,
+    failure: Arc<std::sync::Mutex<Option<String>>>,
+    driver_task: JoinHandle<()>,
+    reader_task: JoinHandle<Result<(), String>>,
+    writer_task: JoinHandle<Result<(), String>>,
+}
+
+impl CallerCancellation {
+    fn new() -> Self {
+        Self {
+            token: CancellationToken::new(),
+            armed: true,
+        }
+    }
+
+    fn token(&self) -> CancellationToken {
+        self.token.clone()
+    }
+
+    fn disarm(mut self) {
+        self.armed = false;
+    }
+}
+
+impl Drop for CallerCancellation {
+    fn drop(&mut self) {
+        if self.armed {
+            self.token.cancel();
+        }
+    }
+}
+
+impl Connection {
+    pub(super) async fn spawn(host_program: &Path) -> Result<Self, String> {
+        let mut command = Command::new(host_program);
+        #[cfg(unix)]
+        command.process_group(0);
+        let mut child = command
+            .stdin(Stdio::piped())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::piped())
+            .kill_on_drop(true)
+            .spawn()
+            .map_err(|err| {
+                format!(
+                    "failed to spawn code-mode host {}: {err}",
+                    host_program.display()
+                )
+            })?;
+
+        if let Some(stderr) = child.stderr.take() {
+            tokio::spawn(async move {
+                let mut lines = BufReader::new(stderr).lines();
+                loop {
+                    match lines.next_line().await {
+                        Ok(Some(line)) => debug!("code-mode host stderr: {line}"),
+                        Ok(None) => break,
+                        Err(err) => {
+                            warn!("failed to read code-mode host stderr: {err}");
+                            break;
+                        }
+                    }
+                }
+            });
+        }
+
+        let stdin = child
+            .stdin
+            .take()
+            .ok_or_else(|| "spawned code-mode host has no stdin".to_string())?;
+        let stdout = child
+            .stdout
+            .take()
+            .ok_or_else(|| "spawned code-mode host has no stdout".to_string())?;
+        let mut reader = FramedReader::new(stdout);
+        let mut writer = FramedWriter::new(stdin);
+        let handshake = async {
+            let hello = ClientHello::new(
+                SupportedProtocolVersions::try_new([ProtocolVersion::V1])
+                    .map_err(|err| err.to_string())?,
+                CapabilitySet::empty(),
+                CapabilitySet::empty(),
+            )
+            .map_err(|err| err.to_string())?;
+            writer
+                .write(&ClientToHost::ClientHello(hello))
+                .await
+                .map_err(|err| format!("failed to write code-mode host hello: {err}"))?;
+            match reader
+                .read::<HostToClient>()
+                .await
+                .map_err(|err| format!("failed to read code-mode host hello: {err}"))?
+            {
+                Some(HostToClient::HostHello(hello))
+                    if hello.selected_version() == ProtocolVersion::V1 =>
+                {
+                    Ok(())
+                }
+                Some(HostToClient::HandshakeRejected { reason }) => {
+                    Err(format!("code-mode host rejected the handshake: {reason:?}"))
+                }
+                Some(message) => Err(format!(
+                    "code-mode host returned an invalid handshake response: {message:?}"
+                )),
+                None => Err("code-mode host exited during handshake".to_string()),
+            }
+        };
+        let handshake_result = match tokio::time::timeout(HOST_HANDSHAKE_TIMEOUT, handshake).await {
+            Ok(result) => result,
+            Err(_) => {
+                kill_and_reap(&mut child).await;
+                return Err("timed out negotiating with the code-mode host".to_string());
+            }
+        };
+        if let Err(err) = handshake_result {
+            kill_and_reap(&mut child).await;
+            return Err(err);
+        }
+
+        let (command_tx, command_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY);
+        let (event_tx, event_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY);
+        let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<EncodedFrame>(IPC_CHANNEL_CAPACITY);
+        let cancellation = CancellationToken::new();
+        let alive = Arc::new(AtomicBool::new(true));
+        let failure = Arc::new(std::sync::Mutex::new(None));
+
+        let writer_cancellation = cancellation.clone();
+        let writer_task = tokio::spawn(async move {
+            loop {
+                tokio::select! {
+                    _ = writer_cancellation.cancelled() => return Ok(()),
+                    frame = outgoing_rx.recv() => {
+                        let Some(frame) = frame else {
+                            return Err("code-mode host outgoing stream closed".to_string());
+                        };
+                        if let Err(err) = writer.write_frame(&frame).await {
+                            return Err(format!("failed to write code-mode host message: {err}"));
+                        }
+                    }
+                }
+            }
+        });
+
+        let reader_events = event_tx.clone();
+        let reader_cancellation = cancellation.clone();
+        let reader_task =
+            tokio::spawn(
+                async move { drive_reader(reader, reader_events, reader_cancellation).await },
+            );
+
+        let (driver, execute_claim_tx) = ConnectionDriver::new(
+            command_rx,
+            event_rx,
+            event_tx.clone(),
+            outgoing_tx,
+            DriverLifecycle {
+                alive: Arc::clone(&alive),
+                failure: Arc::clone(&failure),
+                cancellation: cancellation.clone(),
+            },
+        );
+        let driver_task = tokio::spawn(driver.run());
+        tokio::spawn(
+            ConnectionSupervisor {
+                child,
+                event_tx,
+                cancellation: cancellation.clone(),
+                alive: Arc::clone(&alive),
+                failure: Arc::clone(&failure),
+                driver_task,
+                reader_task,
+                writer_task,
+            }
+            .run(),
+        );
+
+        Ok(Self {
+            command_tx,
+            execute_claim_tx,
+            alive,
+            failure,
+            cancellation,
+        })
+    }
+
+    pub(super) fn is_alive(&self) -> bool {
+        if self.command_tx.is_closed() {
+            mark_connection_dead(
+                &self.alive,
+                &self.failure,
+                "code-mode connection driver closed".to_string(),
+            );
+        }
+        self.alive.load(Ordering::Acquire)
+    }
+
+    pub(super) async fn open_session(
+        &self,
+        session: RemoteSession,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+    ) -> Result<SessionCleanup, String> {
+        let cleanup = SessionCleanup::new();
+        let cancellation = CallerCancellation::new();
+        let (response_tx, response_rx) = oneshot::channel();
+        self.send(DriverCommand::OpenSession {
+            session,
+            delegate,
+            cleanup: cleanup.clone(),
+            caller_cancellation: cancellation.token(),
+            response_tx,
+        })
+        .await?;
+        let result = self.receive(response_rx).await;
+        cancellation.disarm();
+        result?;
+        Ok(cleanup)
+    }
+
+    pub(super) async fn execute(
+        &self,
+        session: RemoteSession,
+        request: ExecuteRequest,
+    ) -> Result<StartedCell, String> {
+        let cancellation = CallerCancellation::new();
+        let (response_tx, response_rx) = oneshot::channel();
+        self.send(DriverCommand::Execute {
+            session,
+            request,
+            caller_cancellation: cancellation.token(),
+            response_tx,
+        })
+        .await?;
+        let delivered = match self.receive(response_rx).await {
+            Ok(delivered) => delivered,
+            Err(err) => {
+                cancellation.disarm();
+                return Err(err);
+            }
+        };
+        self.execute_claim_tx
+            .send(delivered.request_id)
+            .map_err(|_| self.failure_message())?;
+        cancellation.disarm();
+        Ok(delivered.started)
+    }
+
+    pub(super) async fn wait(
+        &self,
+        session: RemoteSession,
+        request: WaitRequest,
+    ) -> Result<WaitOutcome, String> {
+        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;
+        cancellation.disarm();
+        result
+    }
+
+    pub(super) async fn terminate(
+        &self,
+        session: RemoteSession,
+        cell_id: CellId,
+    ) -> Result<WaitOutcome, String> {
+        let (response_tx, response_rx) = oneshot::channel();
+        self.send(DriverCommand::Terminate {
+            session,
+            cell_id,
+            response_tx,
+        })
+        .await?;
+        self.receive(response_rx).await
+    }
+
+    pub(super) async fn shutdown_session(&self, session: RemoteSession) -> Result<(), String> {
+        let (response_tx, response_rx) = oneshot::channel();
+        self.send(DriverCommand::ShutdownSession {
+            session,
+            response_tx,
+        })
+        .await?;
+        self.receive(response_rx).await
+    }
+
+    async fn send(&self, command: DriverCommand) -> Result<(), String> {
+        if !self.is_alive() {
+            return Err(self.failure_message());
+        }
+        self.command_tx
+            .send(command)
+            .await
+            .map_err(|_| self.failure_message())
+    }
+
+    async fn receive<T>(
+        &self,
+        response_rx: oneshot::Receiver<Result<T, String>>,
+    ) -> Result<T, String> {
+        response_rx.await.map_err(|_| self.failure_message())?
+    }
+
+    fn failure_message(&self) -> String {
+        self.failure
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .clone()
+            .unwrap_or_else(|| "code-mode host connection closed".to_string())
+    }
+}
+
+impl Drop for Connection {
+    fn drop(&mut self) {
+        mark_connection_dead(
+            &self.alive,
+            &self.failure,
+            "code-mode host connection closed".to_string(),
+        );
+        self.cancellation.cancel();
+    }
+}
+
+impl ConnectionSupervisor {
+    async fn run(mut self) {
+        let mut child_exited = false;
+        let reason = tokio::select! {
+            biased;
+            _ = self.cancellation.cancelled() => failure_message(&self.failure),
+            result = &mut self.driver_task => match result {
+                Ok(()) => "code-mode connection driver exited unexpectedly".to_string(),
+                Err(err) => format!("code-mode connection driver task failed: {err}"),
+            },
+            result = &mut self.reader_task => task_failure("reader", result),
+            result = &mut self.writer_task => task_failure("writer", result),
+            result = self.child.wait() => {
+                child_exited = true;
+                match result {
+                    Ok(status) => format!("code-mode host exited with status {status}"),
+                    Err(err) => format!("failed waiting for code-mode host: {err}"),
+                }
+            }
+        };
+        mark_connection_dead(&self.alive, &self.failure, reason.clone());
+        let _ = self.event_tx.try_send(DriverEvent::Failed(reason));
+        self.cancellation.cancel();
+        if !child_exited {
+            kill_and_reap(&mut self.child).await;
+        }
+    }
+}
+
+fn task_failure(
+    task_name: &str,
+    result: Result<Result<(), String>, tokio::task::JoinError>,
+) -> String {
+    match result {
+        Ok(Ok(())) => format!("code-mode connection {task_name} exited unexpectedly"),
+        Ok(Err(err)) => err,
+        Err(err) => format!("code-mode connection {task_name} task failed: {err}"),
+    }
+}
+
+fn mark_connection_dead(
+    alive: &AtomicBool,
+    failure: &std::sync::Mutex<Option<String>>,
+    reason: String,
+) {
+    alive.store(false, Ordering::Release);
+    let mut failure = failure
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner);
+    if failure.is_none() {
+        *failure = Some(reason);
+    }
+}
+
+fn failure_message(failure: &std::sync::Mutex<Option<String>>) -> String {
+    failure
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner)
+        .clone()
+        .unwrap_or_else(|| "code-mode host connection closed".to_string())
+}
+
+async fn kill_and_reap(child: &mut Child) {
+    let _ = child.start_kill();
+    let _ = child.wait().await;
+}

codex-rs/code-mode/src/remote_session/connection/driver.rs

@@ -0,0 +1,179 @@
+use std::panic::AssertUnwindSafe;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::host::EncodedFrame;
+use codex_code_mode_protocol::host::RequestId;
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+
+pub(in crate::remote_session) use self::cleanup::SessionCleanup;
+use self::delegate_runtime::DelegateRuntime;
+use self::request_tracker::RequestTracker;
+use self::session_registry::SessionRegistry;
+pub(super) use self::types::DriverCommand;
+pub(super) use self::types::DriverEvent;
+pub(in crate::remote_session) use self::types::RemoteSession;
+
+mod cell_ids;
+mod cleanup;
+mod commands;
+mod delegate_runtime;
+mod request_tracker;
+mod responses;
+mod session_registry;
+mod types;
+
+pub(super) struct DriverLifecycle {
+    pub(super) alive: Arc<AtomicBool>,
+    pub(super) failure: Arc<std::sync::Mutex<Option<String>>>,
+    pub(super) cancellation: CancellationToken,
+}
+
+pub(super) struct ConnectionDriver {
+    command_rx: mpsc::Receiver<DriverCommand>,
+    event_rx: mpsc::Receiver<DriverEvent>,
+    event_tx: mpsc::Sender<DriverEvent>,
+    execute_claim_rx: mpsc::UnboundedReceiver<RequestId>,
+    outgoing_tx: mpsc::Sender<EncodedFrame>,
+    requests: RequestTracker,
+    sessions: SessionRegistry,
+    delegates: DelegateRuntime,
+    alive: Arc<AtomicBool>,
+    failure: Arc<std::sync::Mutex<Option<String>>>,
+    cancellation: CancellationToken,
+    failed: bool,
+}
+
+impl ConnectionDriver {
+    pub(super) fn new(
+        command_rx: mpsc::Receiver<DriverCommand>,
+        event_rx: mpsc::Receiver<DriverEvent>,
+        event_tx: mpsc::Sender<DriverEvent>,
+        outgoing_tx: mpsc::Sender<EncodedFrame>,
+        lifecycle: DriverLifecycle,
+    ) -> (Self, mpsc::UnboundedSender<RequestId>) {
+        let (execute_claim_tx, execute_claim_rx) = mpsc::unbounded_channel();
+        (
+            Self {
+                command_rx,
+                event_rx,
+                event_tx: event_tx.clone(),
+                execute_claim_rx,
+                outgoing_tx,
+                requests: RequestTracker::new(),
+                sessions: SessionRegistry::new(),
+                delegates: DelegateRuntime::new(event_tx),
+                alive: lifecycle.alive,
+                failure: lifecycle.failure,
+                cancellation: lifecycle.cancellation,
+                failed: false,
+            },
+            execute_claim_tx,
+        )
+    }
+
+    pub(super) async fn run(mut self) {
+        loop {
+            tokio::select! {
+                biased;
+                _ = self.cancellation.cancelled() => {
+                    self.fail("code-mode host connection closed".to_string());
+                    return;
+                }
+                event = self.event_rx.recv() => {
+                    let Some(event) = event else {
+                        self.fail("code-mode host event stream closed".to_string());
+                        return;
+                    };
+                    if !self.cancel_dropped_callers() || !self.handle_event(event) {
+                        return;
+                    }
+                }
+                claim = self.execute_claim_rx.recv() => {
+                    let Some(request_id) = claim else {
+                        self.fail("code-mode execute claim stream closed".to_string());
+                        return;
+                    };
+                    self.requests.claim_execute(request_id);
+                }
+                command = self.command_rx.recv() => {
+                    let Some(command) = command else {
+                        self.fail("code-mode host command stream closed".to_string());
+                        return;
+                    };
+                    if !self.cancel_dropped_callers() || !self.handle_command(command) {
+                        return;
+                    }
+                }
+            }
+        }
+    }
+
+    fn handle_event(&mut self, event: DriverEvent) -> bool {
+        let keep_running = match event {
+            DriverEvent::HostMessage(message) => self.handle_host_message(message),
+            DriverEvent::DelegateCompleted { id, result } => self.complete_delegate(id, result),
+            DriverEvent::RequestCancelled(id) => self.cancel_request(id),
+            DriverEvent::Failed(reason) => {
+                self.fail(reason);
+                false
+            }
+        };
+        if keep_running {
+            self.flush_deferred_waits()
+        } else {
+            false
+        }
+    }
+
+    fn queue_frame(&mut self, frame: EncodedFrame) -> bool {
+        match self.outgoing_tx.try_send(frame) {
+            Ok(()) => true,
+            Err(mpsc::error::TrySendError::Full(_)) => {
+                self.fail("code-mode host outgoing queue is full".to_string());
+                false
+            }
+            Err(mpsc::error::TrySendError::Closed(_)) => {
+                self.fail("code-mode host writer closed".to_string());
+                false
+            }
+        }
+    }
+
+    fn fail(&mut self, reason: String) {
+        if self.failed {
+            return;
+        }
+        self.failed = true;
+        self.alive.store(false, Ordering::Release);
+        let reason = {
+            let mut failure = self
+                .failure
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            failure.get_or_insert(reason).clone()
+        };
+        self.requests.fail_all(&reason);
+        let failed_sessions = self.sessions.drain();
+        self.delegates.fail_all(failed_sessions);
+        self.cancellation.cancel();
+    }
+}
+
+impl Drop for ConnectionDriver {
+    fn drop(&mut self) {
+        self.fail("code-mode connection driver stopped unexpectedly".to_string());
+    }
+}
+
+fn notify_cell_closed(delegate: &Arc<dyn CodeModeSessionDelegate>, cell_id: &CellId) {
+    let _ = std::panic::catch_unwind(AssertUnwindSafe(|| delegate.cell_closed(cell_id)));
+}
+
+#[cfg(test)]
+#[path = "driver_tests.rs"]
+mod tests;

codex-rs/code-mode/src/remote_session/connection/driver/cell_ids.rs

@@ -0,0 +1,111 @@
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::RuntimeResponse;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_code_mode_protocol::WaitRequest;
+use codex_code_mode_protocol::host::WireCellId;
+use codex_code_mode_protocol::host::WireRuntimeResponse;
+use codex_code_mode_protocol::host::WireWaitOutcome;
+use codex_code_mode_protocol::host::WireWaitRequest;
+
+use super::RemoteSession;
+
+pub(super) fn public_cell_id(generation: u64, cell_id: &WireCellId) -> CellId {
+    if generation == 1 {
+        CellId::new(cell_id.as_str().to_string())
+    } else {
+        CellId::new(format!("g{generation}:{}", cell_id.as_str()))
+    }
+}
+
+pub(super) fn public_cell_id_from_protocol(generation: u64, cell_id: &CellId) -> CellId {
+    public_cell_id(generation, &WireCellId::new(cell_id.as_str()))
+}
+
+pub(super) fn remote_cell_id(
+    session: &RemoteSession,
+    cell_id: &CellId,
+) -> Result<WireCellId, String> {
+    if session.generation == 1 {
+        if cell_id.as_str().starts_with('g') && cell_id.as_str().contains(':') {
+            return Err(format!(
+                "cell {cell_id} belongs to a stale code-mode host generation"
+            ));
+        }
+        return Ok(WireCellId::new(cell_id.as_str()));
+    }
+    let prefix = format!("g{}:", session.generation);
+    let Some(remote_id) = cell_id.as_str().strip_prefix(&prefix) else {
+        return Err(format!(
+            "cell {cell_id} belongs to a stale code-mode host generation"
+        ));
+    };
+    Ok(WireCellId::new(remote_id))
+}
+
+pub(super) fn remote_wait_request(
+    session: &RemoteSession,
+    request: WaitRequest,
+) -> Result<WireWaitRequest, String> {
+    Ok(WireWaitRequest {
+        cell_id: remote_cell_id(session, &request.cell_id)?,
+        yield_time_ms: request.yield_time_ms,
+    })
+}
+
+pub(super) fn public_runtime_response(
+    generation: u64,
+    response: RuntimeResponse,
+) -> RuntimeResponse {
+    match response {
+        RuntimeResponse::Yielded {
+            cell_id,
+            content_items,
+        } => RuntimeResponse::Yielded {
+            cell_id: public_cell_id_from_protocol(generation, &cell_id),
+            content_items,
+        },
+        RuntimeResponse::Terminated {
+            cell_id,
+            content_items,
+        } => RuntimeResponse::Terminated {
+            cell_id: public_cell_id_from_protocol(generation, &cell_id),
+            content_items,
+        },
+        RuntimeResponse::Result {
+            cell_id,
+            content_items,
+            error_text,
+        } => RuntimeResponse::Result {
+            cell_id: public_cell_id_from_protocol(generation, &cell_id),
+            content_items,
+            error_text,
+        },
+    }
+}
+
+pub(super) fn public_wait_outcome(generation: u64, outcome: WaitOutcome) -> WaitOutcome {
+    match outcome {
+        WaitOutcome::LiveCell(response) => {
+            WaitOutcome::LiveCell(public_runtime_response(generation, response))
+        }
+        WaitOutcome::MissingCell(response) => {
+            WaitOutcome::MissingCell(public_runtime_response(generation, response))
+        }
+    }
+}
+
+pub(super) fn runtime_response_cell_id(response: &WireRuntimeResponse) -> &WireCellId {
+    match response {
+        WireRuntimeResponse::Yielded { cell_id, .. }
+        | WireRuntimeResponse::Terminated { cell_id, .. }
+        | WireRuntimeResponse::Result { cell_id, .. } => cell_id,
+    }
+}
+
+pub(super) fn wait_outcome_cell_id(outcome: &WireWaitOutcome) -> &WireCellId {
+    match outcome {
+        WireWaitOutcome::LiveCell(response) | WireWaitOutcome::MissingCell(response) => {
+            runtime_response_cell_id(response)
+        }
+    }
+}

codex-rs/code-mode/src/remote_session/connection/driver/cleanup.rs

@@ -0,0 +1,40 @@
+use std::sync::Arc;
+
+use tokio_util::sync::CancellationToken;
+
+use super::notify_cell_closed;
+use super::session_registry::CellOwner;
+
+struct CleanupInner {
+    complete: CancellationToken,
+}
+
+#[derive(Clone)]
+pub(in crate::remote_session) struct SessionCleanup {
+    inner: Arc<CleanupInner>,
+}
+
+impl SessionCleanup {
+    pub(in crate::remote_session) fn new() -> Self {
+        Self {
+            inner: Arc::new(CleanupInner {
+                complete: CancellationToken::new(),
+            }),
+        }
+    }
+
+    pub(super) fn fail(&self, cells: Vec<CellOwner>) {
+        for owner in cells {
+            notify_cell_closed(&owner.delegate, &owner.cell_id);
+        }
+        self.inner.complete.cancel();
+    }
+
+    pub(in crate::remote_session) async fn wait(&self) {
+        self.inner.complete.cancelled().await;
+    }
+
+    pub(in crate::remote_session) fn is_complete(&self) -> bool {
+        self.inner.complete.is_cancelled()
+    }
+}

codex-rs/code-mode/src/remote_session/connection/driver/commands.rs

@@ -0,0 +1,298 @@
+use std::sync::Arc;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::ExecuteRequest;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_code_mode_protocol::WaitRequest;
+use codex_code_mode_protocol::host::ClientToHost;
+use codex_code_mode_protocol::host::EncodedFrame;
+use codex_code_mode_protocol::host::HostRequest;
+use codex_code_mode_protocol::host::WireWaitRequest;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+
+use super::ConnectionDriver;
+use super::cell_ids::remote_cell_id;
+use super::cell_ids::remote_wait_request;
+use super::types::CancellableRequest;
+use super::types::DeferredWait;
+use super::types::DeliveredExecute;
+use super::types::DriverCommand;
+use super::types::PendingRequest;
+use super::types::RemoteSession;
+
+impl ConnectionDriver {
+    pub(super) fn handle_command(&mut self, command: DriverCommand) -> bool {
+        match command {
+            DriverCommand::OpenSession {
+                session,
+                delegate,
+                cleanup,
+                caller_cancellation,
+                response_tx,
+            } => self.open_session(session, delegate, cleanup, caller_cancellation, response_tx),
+            DriverCommand::Execute {
+                session,
+                request,
+                caller_cancellation,
+                response_tx,
+            } => self.execute(session, request, caller_cancellation, response_tx),
+            DriverCommand::Wait {
+                session,
+                request,
+                caller_cancellation,
+                response_tx,
+            } => self.wait(session, request, caller_cancellation, response_tx),
+            DriverCommand::Terminate {
+                session,
+                cell_id,
+                response_tx,
+            } => self.terminate(session, cell_id, response_tx),
+            DriverCommand::ShutdownSession {
+                session,
+                response_tx,
+            } => self.shutdown_session(session, response_tx),
+        }
+    }
+
+    fn open_session(
+        &mut self,
+        session: RemoteSession,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+        cleanup: super::cleanup::SessionCleanup,
+        caller_cancellation: CancellationToken,
+        response_tx: oneshot::Sender<Result<(), String>>,
+    ) -> bool {
+        if self.sessions.contains(&session.id) || self.requests.contains_pending_open(&session) {
+            let _ = response_tx.send(Err(format!(
+                "code-mode session {} is already open",
+                session.id
+            )));
+            return true;
+        }
+        let request_id = match self.requests.allocate_id() {
+            Ok(id) => id,
+            Err(err) => {
+                let _ = response_tx.send(Err(err));
+                return false;
+            }
+        };
+        let message = ClientToHost::Request {
+            id: request_id,
+            request: HostRequest::OpenSession {
+                session_id: session.id.clone(),
+            },
+        };
+        let frame = match EncodedFrame::encode(&message) {
+            Ok(frame) => frame,
+            Err(err) => {
+                let _ = response_tx.send(Err(format!(
+                    "failed to encode code-mode open-session request: {err}"
+                )));
+                return true;
+            }
+        };
+        let cancellation = CancellableRequest::new(caller_cancellation);
+        self.requests.insert_pending(
+            request_id,
+            PendingRequest::OpenSession {
+                session,
+                delegate,
+                cleanup,
+                cancellation,
+                response_tx,
+            },
+            &self.event_tx,
+        );
+        self.queue_frame(frame)
+    }
+
+    fn execute(
+        &mut self,
+        session: RemoteSession,
+        request: ExecuteRequest,
+        caller_cancellation: CancellationToken,
+        response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
+    ) -> bool {
+        if let Err(err) = self.sessions.require_ready(&session) {
+            let _ = response_tx.send(Err(err));
+            return true;
+        }
+        let request = match request.try_into() {
+            Ok(request) => request,
+            Err(err) => {
+                let _ = response_tx.send(Err(format!(
+                    "failed to encode code-mode execute request: {err}"
+                )));
+                return true;
+            }
+        };
+        let request_id = match self.requests.allocate_id() {
+            Ok(id) => id,
+            Err(err) => {
+                let _ = response_tx.send(Err(err));
+                return false;
+            }
+        };
+        let message = ClientToHost::Request {
+            id: request_id,
+            request: HostRequest::Execute {
+                session_id: session.id.clone(),
+                request,
+            },
+        };
+        let frame = match EncodedFrame::encode(&message) {
+            Ok(frame) => frame,
+            Err(err) => {
+                let _ = response_tx.send(Err(format!(
+                    "code-mode execute request exceeds the IPC frame limit: {err}"
+                )));
+                return true;
+            }
+        };
+        let (initial_response_tx, initial_response_rx) = oneshot::channel();
+        let cancellation = CancellableRequest::new(caller_cancellation);
+        self.requests.insert_pending(
+            request_id,
+            PendingRequest::Execute {
+                session,
+                response_tx,
+                initial_response_tx,
+                initial_response_rx,
+                cancellation,
+            },
+            &self.event_tx,
+        );
+        self.queue_frame(frame)
+    }
+
+    fn wait(
+        &mut self,
+        session: RemoteSession,
+        request: WaitRequest,
+        caller_cancellation: CancellationToken,
+        response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
+    ) -> bool {
+        if let Err(err) = self.sessions.require_ready(&session) {
+            let _ = response_tx.send(Err(err));
+            return true;
+        }
+        let request = match remote_wait_request(&session, request) {
+            Ok(request) => request,
+            Err(err) => {
+                let _ = response_tx.send(Err(err));
+                return true;
+            }
+        };
+        if self.requests.has_cancelled_wait(&session, &request.cell_id) {
+            self.requests.push_deferred_wait(DeferredWait {
+                session,
+                request,
+                caller_cancellation,
+                response_tx,
+            });
+            return true;
+        }
+        self.start_wait(session, request, caller_cancellation, response_tx)
+    }
+
+    pub(super) fn start_wait(
+        &mut self,
+        session: RemoteSession,
+        request: WireWaitRequest,
+        caller_cancellation: CancellationToken,
+        response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
+    ) -> bool {
+        let cell_id = request.cell_id.clone();
+        self.send_request(
+            HostRequest::Wait {
+                session_id: session.id.clone(),
+                request,
+            },
+            PendingRequest::Wait {
+                session,
+                cell_id,
+                cancellation: CancellableRequest::new(caller_cancellation),
+                response_tx,
+            },
+        )
+    }
+
+    fn terminate(
+        &mut self,
+        session: RemoteSession,
+        cell_id: CellId,
+        response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
+    ) -> bool {
+        if let Err(err) = self.sessions.require_ready(&session) {
+            let _ = response_tx.send(Err(err));
+            return true;
+        }
+        let cell_id = match remote_cell_id(&session, &cell_id) {
+            Ok(cell_id) => cell_id,
+            Err(err) => {
+                let _ = response_tx.send(Err(err));
+                return true;
+            }
+        };
+        let pending_cell_id = cell_id.clone();
+        self.send_request(
+            HostRequest::Terminate {
+                session_id: session.id.clone(),
+                cell_id,
+            },
+            PendingRequest::Terminate {
+                session,
+                cell_id: pending_cell_id,
+                response_tx,
+            },
+        )
+    }
+
+    fn shutdown_session(
+        &mut self,
+        session: RemoteSession,
+        response_tx: oneshot::Sender<Result<(), String>>,
+    ) -> bool {
+        if let Err(err) = self.sessions.begin_shutdown(&session) {
+            let _ = response_tx.send(Err(err));
+            return true;
+        }
+        self.send_request(
+            HostRequest::ShutdownSession {
+                session_id: session.id.clone(),
+            },
+            PendingRequest::ShutdownSession {
+                session,
+                response_tx,
+            },
+        )
+    }
+
+    pub(super) fn send_request(&mut self, request: HostRequest, pending: PendingRequest) -> bool {
+        let request_id = match self.requests.allocate_id() {
+            Ok(id) => id,
+            Err(err) => {
+                pending.fail(err);
+                return false;
+            }
+        };
+        let message = ClientToHost::Request {
+            id: request_id,
+            request,
+        };
+        let frame = match EncodedFrame::encode(&message) {
+            Ok(frame) => frame,
+            Err(err) => {
+                pending.fail(format!(
+                    "code-mode request exceeds the IPC frame limit: {err}"
+                ));
+                return true;
+            }
+        };
+        self.requests
+            .insert_pending(request_id, pending, &self.event_tx);
+        self.queue_frame(frame)
+    }
+}

codex-rs/code-mode/src/remote_session/connection/driver/delegate_runtime.rs

@@ -0,0 +1,339 @@
+//! Client-side delegate task and closure lifecycle.
+//!
+//! Cancellation revokes the task's completion path before removing its active-call state. The
+//! delegate future may finish later, but it can no longer send a response or affect cell closure.
+
+use std::collections::HashMap;
+use std::collections::HashSet;
+use std::collections::VecDeque;
+
+use codex_code_mode_protocol::CodeModeNestedToolCall;
+use codex_code_mode_protocol::host::ClientToHost;
+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::EncodedFrame;
+use codex_code_mode_protocol::host::SessionId;
+use codex_code_mode_protocol::host::WireCellId;
+use codex_code_mode_protocol::host::WireResult;
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+
+use super::ConnectionDriver;
+use super::notify_cell_closed;
+use super::session_registry::CellOwner;
+use super::session_registry::DelegateTarget;
+use super::session_registry::FailedSession;
+use super::types::DriverEvent;
+
+const MAX_RECENT_DELEGATE_REQUEST_IDS: usize = 4096;
+
+#[derive(Clone, Eq, Hash, PartialEq)]
+struct CellKey {
+    session_id: codex_code_mode_protocol::host::SessionId,
+    cell_id: codex_code_mode_protocol::CellId,
+}
+
+impl CellKey {
+    fn for_owner(owner: &CellOwner) -> Self {
+        Self {
+            session_id: owner.session_id.clone(),
+            cell_id: owner.cell_id.clone(),
+        }
+    }
+}
+
+struct DelegateCall {
+    cell: CellKey,
+    cancellation: CancellationToken,
+    completion_stop: CancellationToken,
+}
+
+impl DelegateCall {
+    fn revoke(&self) {
+        self.cancellation.cancel();
+        self.completion_stop.cancel();
+    }
+}
+
+enum DelegateTask {
+    InvokeTool(CodeModeNestedToolCall),
+    Notify {
+        call_id: String,
+        cell_id: codex_code_mode_protocol::CellId,
+        text: String,
+    },
+}
+
+pub(super) struct DelegateEffects {
+    pub(super) response: Option<(DelegateRequestId, Result<DelegateResponse, String>)>,
+    pub(super) closed_cells: Vec<CellOwner>,
+}
+
+impl DelegateEffects {
+    fn empty() -> Self {
+        Self {
+            response: None,
+            closed_cells: Vec::new(),
+        }
+    }
+
+    fn append(&mut self, mut other: Self) {
+        debug_assert!(self.response.is_none());
+        self.response = other.response.take();
+        self.closed_cells.append(&mut other.closed_cells);
+    }
+}
+
+pub(super) struct DelegateRuntime {
+    calls: HashMap<DelegateRequestId, DelegateCall>,
+    seen_requests: HashSet<DelegateRequestId>,
+    request_order: VecDeque<DelegateRequestId>,
+    event_tx: mpsc::Sender<DriverEvent>,
+}
+
+impl DelegateRuntime {
+    pub(super) fn new(event_tx: mpsc::Sender<DriverEvent>) -> Self {
+        Self {
+            calls: HashMap::new(),
+            seen_requests: HashSet::new(),
+            request_order: VecDeque::new(),
+            event_tx,
+        }
+    }
+
+    pub(super) fn start(
+        &mut self,
+        id: DelegateRequestId,
+        target: DelegateTarget,
+        request: DelegateRequest,
+    ) -> Result<(), String> {
+        if self.calls.contains_key(&id) || self.seen_requests.contains(&id) {
+            return Err(format!("duplicate code-mode delegate request ID {id:?}"));
+        }
+        self.remember_request(id);
+        let cancellation = CancellationToken::new();
+        let task_request = match request {
+            DelegateRequest::InvokeTool { invocation } => {
+                let mut invocation: CodeModeNestedToolCall = invocation.into();
+                invocation.cell_id = target.cell_id.clone();
+                DelegateTask::InvokeTool(invocation)
+            }
+            DelegateRequest::Notify {
+                call_id,
+                cell_id: _,
+                text,
+            } => DelegateTask::Notify {
+                call_id,
+                cell_id: target.cell_id.clone(),
+                text,
+            },
+        };
+        let delegate = target.delegate;
+        let task_cancellation = cancellation.clone();
+        let delegate_task = tokio::spawn(async move {
+            match task_request {
+                DelegateTask::InvokeTool(invocation) => delegate
+                    .invoke_tool(invocation, task_cancellation)
+                    .await
+                    .map(|result| DelegateResponse::ToolResult { result }),
+                DelegateTask::Notify {
+                    call_id,
+                    cell_id,
+                    text,
+                } => delegate
+                    .notify(call_id, cell_id, text, task_cancellation)
+                    .await
+                    .map(|()| DelegateResponse::NotificationDelivered),
+            }
+        });
+        let completion_stop = CancellationToken::new();
+        self.calls.insert(
+            id,
+            DelegateCall {
+                cell: CellKey {
+                    session_id: target.session_id,
+                    cell_id: target.cell_id,
+                },
+                cancellation,
+                completion_stop: completion_stop.clone(),
+            },
+        );
+        let event_tx = self.event_tx.clone();
+        tokio::spawn(async move {
+            let result = tokio::select! {
+                biased;
+                _ = completion_stop.cancelled() => return,
+                result = delegate_task => match result {
+                    Ok(result) => result,
+                    Err(err) => Err(format!("code-mode delegate task failed: {err}")),
+                },
+            };
+            tokio::select! {
+                biased;
+                _ = completion_stop.cancelled() => {}
+                _ = event_tx.send(DriverEvent::DelegateCompleted { id, result }) => {}
+            }
+        });
+        Ok(())
+    }
+
+    pub(super) fn cancel(&mut self, id: DelegateRequestId) {
+        if let Some(call) = self.calls.remove(&id) {
+            call.revoke();
+        }
+    }
+
+    pub(super) fn complete(
+        &mut self,
+        id: DelegateRequestId,
+        result: Result<DelegateResponse, String>,
+    ) -> DelegateEffects {
+        if self.calls.remove(&id).is_none() {
+            return DelegateEffects::empty();
+        }
+        let mut effects = DelegateEffects::empty();
+        effects.response = Some((id, result));
+        effects
+    }
+
+    pub(super) fn close_cell(&mut self, owner: CellOwner) -> DelegateEffects {
+        let key = CellKey::for_owner(&owner);
+        self.calls.retain(|_, call| {
+            if call.cell != key {
+                return true;
+            }
+            call.revoke();
+            false
+        });
+        let mut effects = DelegateEffects::empty();
+        effects.closed_cells.push(owner);
+        effects
+    }
+
+    pub(super) fn close_cells(&mut self, owners: Vec<CellOwner>) -> DelegateEffects {
+        let mut effects = DelegateEffects::empty();
+        for owner in owners {
+            effects.append(self.close_cell(owner));
+        }
+        effects
+    }
+
+    pub(super) fn fail_all(&mut self, failed_sessions: Vec<FailedSession>) {
+        for (_, call) in self.calls.drain() {
+            call.revoke();
+        }
+        for session in failed_sessions {
+            session.cleanup.fail(session.cells);
+        }
+    }
+
+    fn remember_request(&mut self, id: DelegateRequestId) {
+        self.seen_requests.insert(id);
+        self.request_order.push_back(id);
+        while self.request_order.len() > MAX_RECENT_DELEGATE_REQUEST_IDS {
+            if let Some(expired) = self.request_order.pop_front() {
+                self.seen_requests.remove(&expired);
+            }
+        }
+    }
+}
+
+impl ConnectionDriver {
+    pub(super) fn start_delegate(
+        &mut self,
+        id: DelegateRequestId,
+        session_id: SessionId,
+        request: DelegateRequest,
+    ) -> bool {
+        let wire_cell_id = match &request {
+            DelegateRequest::InvokeTool { invocation } => &invocation.cell_id,
+            DelegateRequest::Notify { cell_id, .. } => cell_id,
+        };
+        let target = match self.sessions.delegate_target(&session_id, wire_cell_id) {
+            Ok(target) => target,
+            Err(err) => {
+                self.fail(err);
+                return false;
+            }
+        };
+        if let Err(err) = self.delegates.start(id, target, request) {
+            self.fail(err);
+            return false;
+        }
+        true
+    }
+
+    pub(super) fn complete_delegate(
+        &mut self,
+        id: DelegateRequestId,
+        result: Result<DelegateResponse, String>,
+    ) -> bool {
+        let effects = self.delegates.complete(id, result);
+        self.apply_delegate_effects(effects)
+    }
+
+    fn send_delegate_response(
+        &mut self,
+        id: DelegateRequestId,
+        result: Result<DelegateResponse, String>,
+    ) -> bool {
+        let message = ClientToHost::DelegateResponse {
+            id,
+            result: WireResult::from_result(result),
+        };
+        let frame = match EncodedFrame::encode(&message) {
+            Ok(frame) => frame,
+            Err(err) => {
+                let fallback = ClientToHost::DelegateResponse {
+                    id,
+                    result: WireResult::Err {
+                        message: format!(
+                            "code-mode delegate response exceeds the IPC frame limit: {err}"
+                        ),
+                    },
+                };
+                match EncodedFrame::encode(&fallback) {
+                    Ok(frame) => frame,
+                    Err(fallback_err) => {
+                        self.fail(format!(
+                            "failed to encode code-mode delegate error response: {fallback_err}"
+                        ));
+                        return false;
+                    }
+                }
+            }
+        };
+        self.queue_frame(frame)
+    }
+
+    pub(super) fn close_cell(&mut self, session_id: SessionId, cell_id: WireCellId) -> bool {
+        let owner = match self.sessions.remove_cell(&session_id, &cell_id) {
+            Ok(owner) => owner,
+            Err(err) => {
+                self.fail(err);
+                return false;
+            }
+        };
+        let effects = self.delegates.close_cell(owner);
+        self.apply_delegate_effects(effects)
+    }
+
+    pub(super) fn close_session_locally(&mut self, session_id: &SessionId) -> DelegateEffects {
+        self.requests.remove_unclaimed_for_session(session_id);
+        let owners = self.sessions.remove_session(session_id);
+        self.delegates.close_cells(owners)
+    }
+
+    pub(super) fn apply_delegate_effects(&mut self, effects: DelegateEffects) -> bool {
+        if let Some((id, result)) = effects.response
+            && !self.send_delegate_response(id, result)
+        {
+            return false;
+        }
+        for closed in effects.closed_cells {
+            notify_cell_closed(&closed.delegate, &closed.cell_id);
+        }
+        true
+    }
+}

codex-rs/code-mode/src/remote_session/connection/driver/request_tracker.rs

@@ -0,0 +1,186 @@
+use std::collections::HashMap;
+use std::collections::VecDeque;
+
+use codex_code_mode_protocol::host::RequestId;
+use codex_code_mode_protocol::host::SessionId;
+use codex_code_mode_protocol::host::WireCellId;
+use tokio::sync::mpsc;
+
+use super::types::DeferredWait;
+use super::types::DriverEvent;
+use super::types::InitialResponse;
+use super::types::PendingRequest;
+use super::types::RemoteSession;
+use super::types::UnclaimedExecute;
+
+pub(super) enum CancellationAction {
+    Send(RequestId),
+    Terminate {
+        request_id: RequestId,
+        execute: UnclaimedExecute,
+    },
+}
+
+pub(super) struct RequestTracker {
+    pending: HashMap<RequestId, PendingRequest>,
+    unclaimed_executes: HashMap<RequestId, UnclaimedExecute>,
+    initial_responses: HashMap<RequestId, InitialResponse>,
+    deferred_waits: VecDeque<DeferredWait>,
+    next_request_id: i64,
+}
+
+impl RequestTracker {
+    pub(super) fn new() -> Self {
+        Self {
+            pending: HashMap::new(),
+            unclaimed_executes: HashMap::new(),
+            initial_responses: HashMap::new(),
+            deferred_waits: VecDeque::new(),
+            next_request_id: 1,
+        }
+    }
+
+    pub(super) fn contains_pending_open(&self, session: &RemoteSession) -> bool {
+        self.pending.values().any(|pending| {
+            matches!(
+                pending,
+                PendingRequest::OpenSession {
+                    session: pending_session,
+                    ..
+                } if pending_session.id == session.id
+            )
+        })
+    }
+
+    pub(super) fn allocate_id(&mut self) -> Result<RequestId, String> {
+        let id = self.next_request_id;
+        self.next_request_id = self
+            .next_request_id
+            .checked_add(1)
+            .ok_or_else(|| "code-mode host request ID space exhausted".to_string())?;
+        Ok(RequestId::new(id))
+    }
+
+    pub(super) fn insert_pending(
+        &mut self,
+        id: RequestId,
+        pending: PendingRequest,
+        event_tx: &mpsc::Sender<DriverEvent>,
+    ) {
+        self.pending.insert(id, pending);
+        if let Some(cancellation) = self
+            .pending
+            .get_mut(&id)
+            .and_then(PendingRequest::cancellation_mut)
+        {
+            cancellation.spawn_watcher(id, event_tx.clone());
+        }
+    }
+
+    pub(super) fn remove_pending(&mut self, id: RequestId) -> Option<PendingRequest> {
+        self.pending.remove(&id)
+    }
+
+    pub(super) fn insert_initial_response(&mut self, id: RequestId, response: InitialResponse) {
+        self.initial_responses.insert(id, response);
+    }
+
+    pub(super) fn remove_initial_response(&mut self, id: RequestId) -> Option<InitialResponse> {
+        self.initial_responses.remove(&id)
+    }
+
+    pub(super) fn insert_unclaimed_execute(&mut self, id: RequestId, execute: UnclaimedExecute) {
+        self.unclaimed_executes.insert(id, execute);
+    }
+
+    pub(super) fn claim_execute(&mut self, id: RequestId) {
+        self.unclaimed_executes.remove(&id);
+    }
+
+    pub(super) fn collect_cancellations(&mut self) -> Vec<CancellationAction> {
+        let mut actions = self
+            .pending
+            .iter_mut()
+            .filter_map(|(id, pending)| {
+                let cancellation = pending.cancellation_mut()?;
+                (cancellation.is_cancelled() && cancellation.mark_reported())
+                    .then_some(CancellationAction::Send(*id))
+            })
+            .collect::<Vec<_>>();
+        actions.extend(
+            self.unclaimed_executes
+                .extract_if(|_, execute| {
+                    execute.cancellation.is_cancelled() && execute.cancellation.mark_reported()
+                })
+                .map(|(request_id, execute)| CancellationAction::Terminate {
+                    request_id,
+                    execute,
+                }),
+        );
+        actions
+    }
+
+    pub(super) fn mark_cancelled(&mut self, id: RequestId) -> Option<CancellationAction> {
+        if let Some(cancellation) = self
+            .pending
+            .get_mut(&id)
+            .and_then(PendingRequest::cancellation_mut)
+        {
+            return cancellation
+                .mark_reported()
+                .then_some(CancellationAction::Send(id));
+        }
+        let execute = self.unclaimed_executes.get_mut(&id)?;
+        if !execute.cancellation.mark_reported() {
+            return None;
+        }
+        self.unclaimed_executes
+            .remove(&id)
+            .map(|execute| CancellationAction::Terminate {
+                request_id: id,
+                execute,
+            })
+    }
+
+    pub(super) fn has_cancelled_wait(&self, session: &RemoteSession, cell_id: &WireCellId) -> bool {
+        self.pending.values().any(|pending| {
+            matches!(
+                pending,
+                PendingRequest::Wait {
+                    session: pending_session,
+                    cell_id: pending_cell_id,
+                    cancellation,
+                    ..
+                } if pending_session == session
+                    && pending_cell_id == cell_id
+                    && cancellation.is_cancelled()
+            )
+        })
+    }
+
+    pub(super) fn push_deferred_wait(&mut self, wait: DeferredWait) {
+        self.deferred_waits.push_back(wait);
+    }
+
+    pub(super) fn take_deferred_waits(&mut self) -> VecDeque<DeferredWait> {
+        std::mem::take(&mut self.deferred_waits)
+    }
+
+    pub(super) fn remove_unclaimed_for_session(&mut self, session_id: &SessionId) {
+        self.unclaimed_executes
+            .retain(|_, execute| &execute.session.id != session_id);
+    }
+
+    pub(super) fn fail_all(&mut self, reason: &str) {
+        for (_, pending) in self.pending.drain() {
+            pending.fail(reason.to_string());
+        }
+        self.unclaimed_executes.clear();
+        for (_, initial) in self.initial_responses.drain() {
+            let _ = initial.response_tx.send(Err(reason.to_string()));
+        }
+        for wait in self.deferred_waits.drain(..) {
+            let _ = wait.response_tx.send(Err(reason.to_string()));
+        }
+    }
+}

codex-rs/code-mode/src/remote_session/connection/driver/responses.rs

@@ -0,0 +1,397 @@
+use codex_code_mode_protocol::StartedCell;
+use codex_code_mode_protocol::host::ClientToHost;
+use codex_code_mode_protocol::host::EncodedFrame;
+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::RequestId;
+use codex_code_mode_protocol::host::WireCellId;
+use tokio::sync::oneshot;
+
+use super::ConnectionDriver;
+use super::cell_ids::public_runtime_response;
+use super::cell_ids::public_wait_outcome;
+use super::cell_ids::runtime_response_cell_id;
+use super::cell_ids::wait_outcome_cell_id;
+use super::request_tracker::CancellationAction;
+use super::session_registry::CellAdmissionError;
+use super::types::DeliveredExecute;
+use super::types::InitialResponse;
+use super::types::PendingRequest;
+use super::types::RemoteSession;
+use super::types::UnclaimedExecute;
+
+impl ConnectionDriver {
+    pub(super) fn flush_deferred_waits(&mut self) -> bool {
+        let mut deferred = self.requests.take_deferred_waits();
+        while let Some(wait) = deferred.pop_front() {
+            if wait.caller_cancellation.is_cancelled() {
+                let _ = wait
+                    .response_tx
+                    .send(Err("code-mode request cancelled".to_string()));
+                continue;
+            }
+            if self
+                .requests
+                .has_cancelled_wait(&wait.session, &wait.request.cell_id)
+            {
+                self.requests.push_deferred_wait(wait);
+                continue;
+            }
+            if !self.start_wait(
+                wait.session,
+                wait.request,
+                wait.caller_cancellation,
+                wait.response_tx,
+            ) {
+                for wait in deferred {
+                    let _ = wait
+                        .response_tx
+                        .send(Err("code-mode host connection closed".to_string()));
+                }
+                return false;
+            }
+        }
+        true
+    }
+
+    pub(super) fn handle_host_message(&mut self, message: HostToClient) -> bool {
+        match message {
+            HostToClient::Response { id, result } => {
+                self.complete_request(id, result.into_result())
+            }
+            HostToClient::InitialResponse { id, result } => {
+                self.complete_initial_response(id, result.into_result())
+            }
+            HostToClient::DelegateRequest {
+                id,
+                session_id,
+                request,
+            } => self.start_delegate(id, session_id, request),
+            HostToClient::CancelDelegateRequest { id } => {
+                self.delegates.cancel(id);
+                true
+            }
+            HostToClient::CellClosed {
+                session_id,
+                cell_id,
+            } => self.close_cell(session_id, cell_id),
+            HostToClient::HostHello(_) | HostToClient::HandshakeRejected { .. } => {
+                self.fail("code-mode host sent a second handshake response".to_string());
+                false
+            }
+        }
+    }
+
+    fn complete_request(&mut self, id: RequestId, result: Result<HostResponse, String>) -> bool {
+        let Some(pending) = self.requests.remove_pending(id) else {
+            self.fail(format!("code-mode host returned unknown request ID {id:?}"));
+            return false;
+        };
+        match pending {
+            PendingRequest::OpenSession {
+                session,
+                delegate,
+                cleanup,
+                cancellation,
+                response_tx,
+            } => match result {
+                Ok(HostResponse::SessionReady { session_id }) if session_id == session.id => {
+                    let abandoned = cancellation.is_cancelled() || response_tx.is_closed();
+                    self.sessions
+                        .insert_ready(session.clone(), delegate, cleanup);
+                    if abandoned || response_tx.send(Ok(())).is_err() {
+                        return self.shutdown_abandoned_session(session);
+                    }
+                }
+                Ok(_) => {
+                    let reason =
+                        "code-mode host returned an invalid open-session response".to_string();
+                    let _ = response_tx.send(Err(reason.clone()));
+                    self.fail(reason);
+                    return false;
+                }
+                Err(err) => {
+                    let _ = response_tx.send(Err(err));
+                }
+            },
+            PendingRequest::Execute {
+                session,
+                response_tx,
+                initial_response_tx,
+                initial_response_rx,
+                cancellation,
+            } => match result {
+                Ok(HostResponse::ExecutionStarted { cell_id }) => {
+                    // The host owns a checked, never-reused ID sequence. Retain only live
+                    // IDs so client memory scales with concurrency, not session lifetime.
+                    let remote_cell_id = cell_id.clone();
+                    let public_id = match self.sessions.admit_cell(&session, cell_id) {
+                        Ok(public_id) => public_id,
+                        Err(CellAdmissionError::MissingSession) => {
+                            let _ = response_tx
+                                .send(Err("code-mode session closed during execute".to_string()));
+                            return true;
+                        }
+                        Err(CellAdmissionError::DuplicateCell) => {
+                            let reason = format!(
+                                "code-mode host reused live cell {} in session {}",
+                                remote_cell_id.as_str(),
+                                session.id
+                            );
+                            let _ = response_tx.send(Err(reason.clone()));
+                            self.fail(reason);
+                            return false;
+                        }
+                    };
+                    self.requests.insert_initial_response(
+                        id,
+                        InitialResponse {
+                            generation: session.generation,
+                            cell_id: remote_cell_id.clone(),
+                            response_tx: initial_response_tx,
+                        },
+                    );
+                    let started = StartedCell::from_result_receiver(public_id, initial_response_rx);
+                    if cancellation.is_cancelled() || response_tx.is_closed() {
+                        return self.terminate_abandoned_cell(session, remote_cell_id);
+                    }
+                    let delivered = DeliveredExecute {
+                        request_id: id,
+                        started,
+                    };
+                    if response_tx.send(Ok(delivered)).is_err() {
+                        return self.terminate_abandoned_cell(session, remote_cell_id);
+                    }
+                    self.requests.insert_unclaimed_execute(
+                        id,
+                        UnclaimedExecute {
+                            session,
+                            cell_id: remote_cell_id,
+                            cancellation,
+                        },
+                    );
+                }
+                Ok(_) => {
+                    let reason = "code-mode host returned an invalid execute response".to_string();
+                    let _ = response_tx.send(Err(reason.clone()));
+                    self.fail(reason);
+                    return false;
+                }
+                Err(err) => {
+                    let _ = response_tx.send(Err(err));
+                }
+            },
+            PendingRequest::Wait {
+                session,
+                cell_id,
+                cancellation: _,
+                response_tx,
+            } => {
+                let result = match result {
+                    Ok(HostResponse::WaitCompleted { outcome }) => {
+                        if wait_outcome_cell_id(&outcome) != &cell_id {
+                            let reason = format!(
+                                "code-mode host returned cell {} for request targeting {}",
+                                wait_outcome_cell_id(&outcome).as_str(),
+                                cell_id.as_str()
+                            );
+                            let _ = response_tx.send(Err(reason.clone()));
+                            self.fail(reason);
+                            return false;
+                        }
+                        Ok(public_wait_outcome(session.generation, outcome.into()))
+                    }
+                    Ok(_) => {
+                        let reason = "code-mode host returned an invalid cell response".to_string();
+                        let _ = response_tx.send(Err(reason.clone()));
+                        self.fail(reason);
+                        return false;
+                    }
+                    Err(err) => Err(err),
+                };
+                let _ = response_tx.send(result);
+            }
+            PendingRequest::Terminate {
+                session,
+                cell_id,
+                response_tx,
+            } => {
+                let result = match result {
+                    Ok(HostResponse::WaitCompleted { outcome }) => {
+                        if wait_outcome_cell_id(&outcome) != &cell_id {
+                            let reason = format!(
+                                "code-mode host returned cell {} for request targeting {}",
+                                wait_outcome_cell_id(&outcome).as_str(),
+                                cell_id.as_str()
+                            );
+                            let _ = response_tx.send(Err(reason.clone()));
+                            self.fail(reason);
+                            return false;
+                        }
+                        public_wait_outcome(session.generation, outcome.into())
+                    }
+                    Ok(_) => {
+                        let reason = "code-mode host returned an invalid cell response".to_string();
+                        let _ = response_tx.send(Err(reason.clone()));
+                        self.fail(reason);
+                        return false;
+                    }
+                    Err(err) => {
+                        let _ = response_tx.send(Err(err));
+                        return true;
+                    }
+                };
+                let _ = response_tx.send(Ok(result));
+            }
+            PendingRequest::ShutdownSession {
+                session,
+                response_tx,
+            } => match result {
+                Ok(HostResponse::SessionClosed { session_id }) if session_id == session.id => {
+                    let effects = self.close_session_locally(&session.id);
+                    if !self.apply_delegate_effects(effects) {
+                        return false;
+                    }
+                    let _ = response_tx.send(Ok(()));
+                }
+                Ok(_) => {
+                    let err = "code-mode host returned an invalid shutdown response".to_string();
+                    let _ = response_tx.send(Err(err.clone()));
+                    self.fail(err);
+                    return false;
+                }
+                Err(err) => {
+                    let _ = response_tx.send(Err(err.clone()));
+                    self.fail(err);
+                    return false;
+                }
+            },
+        }
+        true
+    }
+
+    pub(super) fn cancel_dropped_callers(&mut self) -> bool {
+        for action in self.requests.collect_cancellations() {
+            if !self.apply_cancellation(action) {
+                return false;
+            }
+        }
+        true
+    }
+
+    pub(super) fn cancel_request(&mut self, id: RequestId) -> bool {
+        self.requests
+            .mark_cancelled(id)
+            .is_none_or(|action| self.apply_cancellation(action))
+    }
+
+    fn apply_cancellation(&mut self, action: CancellationAction) -> bool {
+        match action {
+            CancellationAction::Send(id) => self.send_cancel_request(id),
+            CancellationAction::Terminate {
+                request_id,
+                execute,
+            } => {
+                if !self.send_cancel_request(request_id) {
+                    return false;
+                }
+                self.terminate_abandoned_cell(execute.session, execute.cell_id)
+            }
+        }
+    }
+
+    fn send_cancel_request(&mut self, id: RequestId) -> bool {
+        let frame = match EncodedFrame::encode(&ClientToHost::CancelRequest { id }) {
+            Ok(frame) => frame,
+            Err(err) => {
+                self.fail(format!(
+                    "failed to encode code-mode cancellation request: {err}"
+                ));
+                return false;
+            }
+        };
+        self.queue_frame(frame)
+    }
+
+    fn shutdown_abandoned_session(&mut self, session: RemoteSession) -> bool {
+        let Some(should_shutdown) = self.sessions.begin_abandoned_shutdown(&session.id) else {
+            self.fail(format!(
+                "code-mode host committed abandoned session {} without local state",
+                session.id
+            ));
+            return false;
+        };
+        if !should_shutdown {
+            return true;
+        }
+        let (response_tx, response_rx) = oneshot::channel();
+        drop(response_rx);
+        self.send_request(
+            HostRequest::ShutdownSession {
+                session_id: session.id.clone(),
+            },
+            PendingRequest::ShutdownSession {
+                session,
+                response_tx,
+            },
+        )
+    }
+
+    fn terminate_abandoned_cell(&mut self, session: RemoteSession, cell_id: WireCellId) -> bool {
+        let Some(is_closing) = self.sessions.is_closing(&session.id) else {
+            self.fail(format!(
+                "code-mode host admitted an abandoned cell in unknown session {}",
+                session.id
+            ));
+            return false;
+        };
+        if is_closing {
+            return true;
+        }
+        let (response_tx, response_rx) = oneshot::channel();
+        drop(response_rx);
+        self.send_request(
+            HostRequest::Terminate {
+                session_id: session.id.clone(),
+                cell_id: cell_id.clone(),
+            },
+            PendingRequest::Terminate {
+                session,
+                cell_id,
+                response_tx,
+            },
+        )
+    }
+
+    fn complete_initial_response(
+        &mut self,
+        id: RequestId,
+        result: Result<codex_code_mode_protocol::host::WireRuntimeResponse, String>,
+    ) -> bool {
+        let Some(initial) = self.requests.remove_initial_response(id) else {
+            self.fail(format!(
+                "code-mode host returned initial response for unknown request ID {id:?}"
+            ));
+            return false;
+        };
+        let response = match result {
+            Ok(response) if runtime_response_cell_id(&response) == &initial.cell_id => {
+                Ok(public_runtime_response(initial.generation, response.into()))
+            }
+            Ok(response) => {
+                let reason = format!(
+                    "code-mode host returned initial response for cell {} instead of {}",
+                    runtime_response_cell_id(&response).as_str(),
+                    initial.cell_id.as_str()
+                );
+                let _ = initial.response_tx.send(Err(reason.clone()));
+                self.fail(reason);
+                return false;
+            }
+            Err(err) => Err(err),
+        };
+        let _ = initial.response_tx.send(response);
+        true
+    }
+}

codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs

@@ -0,0 +1,222 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::host::SessionId;
+use codex_code_mode_protocol::host::WireCellId;
+
+use super::cell_ids::public_cell_id;
+use super::cleanup::SessionCleanup;
+use super::types::RemoteSession;
+
+pub(super) struct CellOwner {
+    pub(super) session_id: SessionId,
+    pub(super) cell_id: CellId,
+    pub(super) delegate: Arc<dyn CodeModeSessionDelegate>,
+}
+
+pub(super) struct DelegateTarget {
+    pub(super) session_id: SessionId,
+    pub(super) cell_id: CellId,
+    pub(super) delegate: Arc<dyn CodeModeSessionDelegate>,
+}
+
+pub(super) struct FailedSession {
+    pub(super) cleanup: SessionCleanup,
+    pub(super) cells: Vec<CellOwner>,
+}
+
+pub(super) enum CellAdmissionError {
+    MissingSession,
+    DuplicateCell,
+}
+
+struct SessionRecord {
+    remote: RemoteSession,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
+    cleanup: SessionCleanup,
+    phase: SessionPhase,
+    cells: HashMap<WireCellId, CellId>,
+}
+
+#[derive(Clone, Copy, Eq, PartialEq)]
+enum SessionPhase {
+    Ready,
+    Closing,
+}
+
+pub(super) struct SessionRegistry {
+    records: HashMap<SessionId, SessionRecord>,
+}
+
+impl SessionRegistry {
+    pub(super) fn new() -> Self {
+        Self {
+            records: HashMap::new(),
+        }
+    }
+
+    pub(super) fn contains(&self, session_id: &SessionId) -> bool {
+        self.records.contains_key(session_id)
+    }
+
+    pub(super) fn insert_ready(
+        &mut self,
+        session: RemoteSession,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+        cleanup: SessionCleanup,
+    ) {
+        self.records.insert(
+            session.id.clone(),
+            SessionRecord {
+                remote: session,
+                delegate,
+                cleanup,
+                phase: SessionPhase::Ready,
+                cells: HashMap::new(),
+            },
+        );
+    }
+
+    pub(super) fn require_ready(&self, session: &RemoteSession) -> Result<(), String> {
+        let record = self
+            .records
+            .get(&session.id)
+            .ok_or_else(|| format!("unknown code-mode session {}", session.id))?;
+        if record.remote != *session {
+            return Err("stale code-mode session generation".to_string());
+        }
+        if record.phase != SessionPhase::Ready {
+            return Err("code-mode session is shutting down".to_string());
+        }
+        Ok(())
+    }
+
+    pub(super) fn begin_shutdown(&mut self, session: &RemoteSession) -> Result<(), String> {
+        let record = self
+            .records
+            .get_mut(&session.id)
+            .ok_or_else(|| format!("unknown code-mode session {}", session.id))?;
+        if record.remote != *session {
+            return Err("stale code-mode session generation".to_string());
+        }
+        if record.phase == SessionPhase::Closing {
+            return Err("code-mode session is already closing".to_string());
+        }
+        record.phase = SessionPhase::Closing;
+        Ok(())
+    }
+
+    pub(super) fn begin_abandoned_shutdown(&mut self, session_id: &SessionId) -> Option<bool> {
+        let record = self.records.get_mut(session_id)?;
+        if record.phase == SessionPhase::Closing {
+            return Some(false);
+        }
+        record.phase = SessionPhase::Closing;
+        Some(true)
+    }
+
+    pub(super) fn is_closing(&self, session_id: &SessionId) -> Option<bool> {
+        self.records
+            .get(session_id)
+            .map(|record| record.phase == SessionPhase::Closing)
+    }
+
+    pub(super) fn admit_cell(
+        &mut self,
+        session: &RemoteSession,
+        cell_id: WireCellId,
+    ) -> Result<CellId, CellAdmissionError> {
+        let Some(record) = self.records.get_mut(&session.id) else {
+            return Err(CellAdmissionError::MissingSession);
+        };
+        if record.cells.contains_key(&cell_id) {
+            return Err(CellAdmissionError::DuplicateCell);
+        }
+        let public_id = public_cell_id(session.generation, &cell_id);
+        record.cells.insert(cell_id, public_id.clone());
+        Ok(public_id)
+    }
+
+    pub(super) fn delegate_target(
+        &self,
+        session_id: &SessionId,
+        cell_id: &WireCellId,
+    ) -> Result<DelegateTarget, String> {
+        let session = self
+            .records
+            .get(session_id)
+            .ok_or_else(|| format!("code-mode host delegated for unknown session {session_id}"))?;
+        let public_id = session.cells.get(cell_id).cloned().ok_or_else(|| {
+            format!(
+                "code-mode host delegated for unknown cell {} in session {session_id}",
+                cell_id.as_str()
+            )
+        })?;
+        Ok(DelegateTarget {
+            session_id: session_id.clone(),
+            cell_id: public_id,
+            delegate: Arc::clone(&session.delegate),
+        })
+    }
+
+    pub(super) fn remove_cell(
+        &mut self,
+        session_id: &SessionId,
+        cell_id: &WireCellId,
+    ) -> Result<CellOwner, String> {
+        let session = self.records.get_mut(session_id).ok_or_else(|| {
+            format!(
+                "code-mode host closed cell {} in unknown session {session_id}",
+                cell_id.as_str()
+            )
+        })?;
+        let public_id = session
+            .cells
+            .remove(cell_id)
+            .ok_or_else(|| format!("code-mode host closed unknown cell in session {session_id}"))?;
+        Ok(CellOwner {
+            session_id: session_id.clone(),
+            cell_id: public_id,
+            delegate: Arc::clone(&session.delegate),
+        })
+    }
+
+    pub(super) fn remove_session(&mut self, session_id: &SessionId) -> Vec<CellOwner> {
+        let Some(session) = self.records.remove(session_id) else {
+            return Vec::new();
+        };
+        session
+            .cells
+            .into_values()
+            .map(|cell_id| CellOwner {
+                session_id: session_id.clone(),
+                cell_id,
+                delegate: Arc::clone(&session.delegate),
+            })
+            .collect()
+    }
+
+    pub(super) fn drain(&mut self) -> Vec<FailedSession> {
+        let sessions = std::mem::take(&mut self.records);
+        sessions
+            .into_iter()
+            .map(|(session_id, session)| {
+                let cells = session
+                    .cells
+                    .into_values()
+                    .map(|cell_id| CellOwner {
+                        session_id: session_id.clone(),
+                        cell_id,
+                        delegate: Arc::clone(&session.delegate),
+                    })
+                    .collect();
+                FailedSession {
+                    cleanup: session.cleanup,
+                    cells,
+                }
+            })
+            .collect()
+    }
+}

codex-rs/code-mode/src/remote_session/connection/driver/types.rs

@@ -0,0 +1,196 @@
+use std::sync::Arc;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::ExecuteRequest;
+use codex_code_mode_protocol::RuntimeResponse;
+use codex_code_mode_protocol::StartedCell;
+use codex_code_mode_protocol::WaitOutcome;
+use codex_code_mode_protocol::WaitRequest;
+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::WireCellId;
+use codex_code_mode_protocol::host::WireWaitRequest;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+
+use super::cleanup::SessionCleanup;
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub(in crate::remote_session) struct RemoteSession {
+    pub(in crate::remote_session) id: SessionId,
+    pub(in crate::remote_session) generation: u64,
+}
+
+pub(in crate::remote_session::connection) enum DriverCommand {
+    OpenSession {
+        session: RemoteSession,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+        cleanup: SessionCleanup,
+        caller_cancellation: CancellationToken,
+        response_tx: oneshot::Sender<Result<(), String>>,
+    },
+    Execute {
+        session: RemoteSession,
+        request: ExecuteRequest,
+        caller_cancellation: CancellationToken,
+        response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
+    },
+    Wait {
+        session: RemoteSession,
+        request: WaitRequest,
+        caller_cancellation: CancellationToken,
+        response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
+    },
+    Terminate {
+        session: RemoteSession,
+        cell_id: CellId,
+        response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
+    },
+    ShutdownSession {
+        session: RemoteSession,
+        response_tx: oneshot::Sender<Result<(), String>>,
+    },
+}
+
+pub(in crate::remote_session::connection) enum DriverEvent {
+    HostMessage(HostToClient),
+    DelegateCompleted {
+        id: DelegateRequestId,
+        result: Result<DelegateResponse, String>,
+    },
+    RequestCancelled(RequestId),
+    Failed(String),
+}
+
+pub(super) struct CancellableRequest {
+    caller_cancellation: CancellationToken,
+    watcher_stop: CancellationToken,
+    reported: bool,
+}
+
+impl CancellableRequest {
+    pub(super) fn new(caller_cancellation: CancellationToken) -> Self {
+        Self {
+            caller_cancellation,
+            watcher_stop: CancellationToken::new(),
+            reported: false,
+        }
+    }
+
+    pub(super) fn is_cancelled(&self) -> bool {
+        self.caller_cancellation.is_cancelled()
+    }
+
+    pub(super) fn mark_reported(&mut self) -> bool {
+        if self.reported {
+            return false;
+        }
+        self.reported = true;
+        true
+    }
+
+    pub(super) fn spawn_watcher(&self, id: RequestId, event_tx: mpsc::Sender<DriverEvent>) {
+        let caller_cancellation = self.caller_cancellation.clone();
+        let watcher_stop = self.watcher_stop.clone();
+        tokio::spawn(async move {
+            tokio::select! {
+                _ = caller_cancellation.cancelled() => {
+                    let _ = event_tx.send(DriverEvent::RequestCancelled(id)).await;
+                }
+                _ = watcher_stop.cancelled() => {}
+            }
+        });
+    }
+}
+
+impl Drop for CancellableRequest {
+    fn drop(&mut self) {
+        self.watcher_stop.cancel();
+    }
+}
+
+pub(super) struct InitialResponse {
+    pub(super) generation: u64,
+    pub(super) cell_id: WireCellId,
+    pub(super) response_tx: oneshot::Sender<Result<RuntimeResponse, String>>,
+}
+
+pub(in crate::remote_session::connection) struct DeliveredExecute {
+    pub(in crate::remote_session::connection) request_id: RequestId,
+    pub(in crate::remote_session::connection) started: StartedCell,
+}
+
+pub(super) struct UnclaimedExecute {
+    pub(super) session: RemoteSession,
+    pub(super) cell_id: WireCellId,
+    pub(super) cancellation: CancellableRequest,
+}
+
+pub(super) enum PendingRequest {
+    OpenSession {
+        session: RemoteSession,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+        cleanup: SessionCleanup,
+        cancellation: CancellableRequest,
+        response_tx: oneshot::Sender<Result<(), String>>,
+    },
+    Execute {
+        session: RemoteSession,
+        response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
+        initial_response_tx: oneshot::Sender<Result<RuntimeResponse, String>>,
+        initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
+        cancellation: CancellableRequest,
+    },
+    Wait {
+        session: RemoteSession,
+        cell_id: WireCellId,
+        cancellation: CancellableRequest,
+        response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
+    },
+    Terminate {
+        session: RemoteSession,
+        cell_id: WireCellId,
+        response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
+    },
+    ShutdownSession {
+        session: RemoteSession,
+        response_tx: oneshot::Sender<Result<(), String>>,
+    },
+}
+
+pub(super) struct DeferredWait {
+    pub(super) session: RemoteSession,
+    pub(super) request: WireWaitRequest,
+    pub(super) caller_cancellation: CancellationToken,
+    pub(super) response_tx: oneshot::Sender<Result<WaitOutcome, String>>,
+}
+
+impl PendingRequest {
+    pub(super) fn cancellation_mut(&mut self) -> Option<&mut CancellableRequest> {
+        match self {
+            Self::OpenSession { cancellation, .. }
+            | Self::Execute { cancellation, .. }
+            | Self::Wait { cancellation, .. } => Some(cancellation),
+            Self::Terminate { .. } | Self::ShutdownSession { .. } => None,
+        }
+    }
+
+    pub(super) fn fail(self, reason: String) {
+        match self {
+            Self::OpenSession { response_tx, .. } | Self::ShutdownSession { response_tx, .. } => {
+                let _ = response_tx.send(Err(reason));
+            }
+            Self::Execute { response_tx, .. } => {
+                let _ = response_tx.send(Err(reason));
+            }
+            Self::Wait { response_tx, .. } | Self::Terminate { response_tx, .. } => {
+                let _ = response_tx.send(Err(reason));
+            }
+        }
+    }
+}

codex-rs/code-mode/src/remote_session/connection/reader.rs

@@ -0,0 +1,29 @@
+use codex_code_mode_protocol::host::FramedReader;
+use codex_code_mode_protocol::host::HostToClient;
+use tokio::process::ChildStdout;
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+
+use super::driver::DriverEvent;
+
+pub(super) async fn drive_reader(
+    mut reader: FramedReader<ChildStdout>,
+    events: mpsc::Sender<DriverEvent>,
+    cancellation: CancellationToken,
+) -> Result<(), String> {
+    loop {
+        let message = tokio::select! {
+            _ = cancellation.cancelled() => return Ok(()),
+            result = reader.read::<HostToClient>() => result,
+        };
+        let message = match message {
+            Ok(Some(message)) => message,
+            Ok(None) => return Err("code-mode host closed its stdout".to_string()),
+            Err(err) => return Err(format!("failed to read code-mode host message: {err}")),
+        };
+        events
+            .send(DriverEvent::HostMessage(message))
+            .await
+            .map_err(|_| "code-mode connection driver closed".to_string())?;
+    }
+}

codex-rs/code-mode/src/remote_session_tests.rs

@@ -0,0 +1,52 @@
+use std::sync::Arc;
+
+use codex_code_mode_protocol::CodeModeSessionProvider;
+
+use super::ProcessOwnedCodeModeSession;
+use super::ProcessOwnedCodeModeSessionProvider;
+use crate::NoopCodeModeSessionDelegate;
+
+#[test]
+fn provider_reuses_its_live_process_host() {
+    let provider = ProcessOwnedCodeModeSessionProvider::default();
+
+    let first = provider.process_host();
+    let second = provider.process_host();
+
+    assert!(Arc::ptr_eq(&first, &second));
+}
+
+#[tokio::test]
+async fn provider_reports_host_spawn_failure() {
+    let provider = ProcessOwnedCodeModeSessionProvider::with_host_program(
+        "codex-code-mode-host-does-not-exist".into(),
+    );
+
+    let error = provider
+        .create_session(Arc::new(NoopCodeModeSessionDelegate))
+        .await
+        .err()
+        .expect("session creation should fail");
+
+    assert!(error.contains("failed to spawn code-mode host"));
+}
+
+#[tokio::test]
+async fn shutdown_before_open_does_not_spawn_the_host() {
+    let session = ProcessOwnedCodeModeSession::new();
+
+    session.shutdown().await.expect("shutdown session");
+    let error = session
+        .execute(codex_code_mode_protocol::ExecuteRequest {
+            tool_call_id: "call-1".to_string(),
+            enabled_tools: Vec::new(),
+            source: "text('unreachable')".to_string(),
+            yield_time_ms: None,
+            max_output_tokens: None,
+        })
+        .await
+        .err()
+        .expect("shutdown session should reject execution");
+
+    assert_eq!(error, "code mode session is shutting down");
+}