code-mode: add and verify stdio session provider

codex-rs/Cargo.lock

@@ -2507,12 +2507,25 @@ dependencies = [
  "v8",
 ]
 
+[[package]]
+name = "codex-code-mode-client"
+version = "0.0.0"
+dependencies = [
+ "codex-code-mode-protocol",
+ "codex-protocol",
+ "tokio",
+ "tokio-util",
+ "tracing",
+]
+
 [[package]]
 name = "codex-code-mode-host"
 version = "0.0.0"
 dependencies = [
  "codex-code-mode",
+ "codex-code-mode-client",
  "codex-code-mode-protocol",
+ "codex-protocol",
  "codex-utils-cargo-bin",
  "pretty_assertions",
  "serde_json",

codex-rs/Cargo.toml

@@ -21,6 +21,7 @@ members = [
     "install-context",
     "codex-backend-openapi-models",
     "code-mode",
+    "code-mode-client",
     "code-mode-host",
     "code-mode-protocol",
     "codex-home",
@@ -160,6 +161,7 @@ codex-cloud-config = { path = "cloud-config" }
 codex-cloud-tasks-client = { path = "cloud-tasks-client" }
 codex-cloud-tasks-mock-client = { path = "cloud-tasks-mock-client" }
 codex-code-mode = { path = "code-mode" }
+codex-code-mode-client = { path = "code-mode-client" }
 codex-code-mode-protocol = { path = "code-mode-protocol" }
 codex-home = { path = "codex-home" }
 codex-config = { path = "config" }

codex-rs/code-mode-client/BUILD.bazel

@@ -0,0 +1,6 @@
+load("//:defs.bzl", "codex_rust_crate")
+
+codex_rust_crate(
+    name = "code-mode-client",
+    crate_name = "codex_code_mode_client",
+)

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

@@ -0,0 +1,20 @@
+[package]
+name = "codex-code-mode-client"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[lib]
+name = "codex_code_mode_client"
+path = "src/lib.rs"
+doctest = false
+
+[lints]
+workspace = true
+
+[dependencies]
+codex-code-mode-protocol = { workspace = true }
+codex-protocol = { workspace = true }
+tokio = { workspace = true, features = ["io-util", "process", "rt", "sync"] }
+tokio-util = { workspace = true, features = ["rt"] }
+tracing = { workspace = true }

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

@@ -0,0 +1,481 @@
+use std::collections::HashMap;
+use std::future::Future;
+use std::process::Stdio;
+use std::sync::Arc;
+use std::sync::PoisonError;
+use std::sync::Weak;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
+
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::wire;
+use tokio::io::AsyncBufReadExt;
+use tokio::io::BufReader;
+use tokio::process::Command;
+use tokio::sync::Mutex;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
+use tracing::debug;
+use tracing::warn;
+
+use crate::CodeModeHostCommand;
+use crate::convert;
+
+const IPC_CHANNEL_CAPACITY: usize = 128;
+
+#[derive(Clone, Debug)]
+pub(super) enum RequestError {
+    Host(wire::Error),
+    Transport(String),
+}
+
+impl RequestError {
+    pub(super) fn into_message(self) -> String {
+        match self {
+            Self::Host(error) => convert::wire_error_message(&error),
+            Self::Transport(message) => message,
+        }
+    }
+}
+
+pub(super) struct Connection {
+    state: Arc<ConnectionState>,
+    cancellation: CancellationToken,
+}
+
+impl Connection {
+    pub(super) async fn spawn(command: &CodeModeHostCommand) -> Result<Self, String> {
+        let mut child = host_process(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}",
+                    command.program.display()
+                )
+            })?;
+        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 stderr = child.stderr.take();
+        let (outgoing_tx, mut outgoing_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY);
+        let cancellation = CancellationToken::new();
+        let state = Arc::new(ConnectionState::new(outgoing_tx));
+
+        let writer_state = Arc::downgrade(&state);
+        let writer_cancellation = cancellation.clone();
+        spawn_critical_connection_task("writer", Arc::downgrade(&state), async move {
+            let mut stdin = stdin;
+            loop {
+                tokio::select! {
+                    _ = writer_cancellation.cancelled() => break,
+                    message = outgoing_rx.recv() => {
+                        let Some(message) = message else {
+                            break;
+                        };
+                        if let Err(err) = wire::write_frame(&mut stdin, &message).await {
+                            if let Some(state) = writer_state.upgrade() {
+                                state
+                                    .fail(format!(
+                                        "failed to write code-mode host message: {err}"
+                                    ))
+                                    .await;
+                            }
+                            break;
+                        }
+                    }
+                }
+            }
+        });
+
+        let reader_cancellation = cancellation.clone();
+        spawn_critical_connection_task(
+            "reader",
+            Arc::downgrade(&state),
+            drive_reader(stdout, Arc::downgrade(&state), reader_cancellation),
+        );
+
+        if let Some(stderr) = stderr {
+            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 supervisor_state = Arc::downgrade(&state);
+        let supervisor_cancellation = cancellation.clone();
+        spawn_critical_connection_task("supervisor", Arc::downgrade(&state), async move {
+            tokio::select! {
+                result = child.wait() => {
+                    let reason = match result {
+                        Ok(status) => format!("code-mode host exited with status {status}"),
+                        Err(err) => format!("failed waiting for code-mode host: {err}"),
+                    };
+                    if let Some(state) = supervisor_state.upgrade() {
+                        state.fail(reason).await;
+                    }
+                }
+                _ = supervisor_cancellation.cancelled() => {
+                    let _ = child.start_kill();
+                    let _ = child.wait().await;
+                    if let Some(state) = supervisor_state.upgrade() {
+                        state
+                            .fail("code-mode host connection closed".to_string())
+                            .await;
+                    }
+                }
+            }
+        });
+
+        Ok(Self {
+            state,
+            cancellation,
+        })
+    }
+
+    pub(super) fn is_alive(&self) -> bool {
+        self.state.alive.load(Ordering::Acquire)
+    }
+
+    pub(super) async fn request(
+        &self,
+        request: wire::HostRequest,
+    ) -> Result<wire::HostResponse, RequestError> {
+        let id = self.state.next_request_id.fetch_add(1, Ordering::Relaxed);
+        let (response_tx, response_rx) = oneshot::channel();
+        self.state.pending.lock().await.insert(id, response_tx);
+        let mut pending_request = PendingRequest::new(Arc::clone(&self.state), id);
+        if let Err(err) = self
+            .state
+            .send(wire::ClientMessage::Request { id, request })
+            .await
+        {
+            pending_request.cleanup().await;
+            return Err(RequestError::Transport(err));
+        }
+        pending_request.mark_request_sent();
+        let response = response_rx
+            .await
+            .map_err(|_| RequestError::Transport(self.state.failure_message()))?;
+        pending_request.disarm();
+        response
+    }
+
+    pub(super) async fn register_delegate(
+        &self,
+        session_id: wire::SessionId,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+    ) {
+        self.state
+            .delegates
+            .lock()
+            .await
+            .insert(session_id, delegate);
+    }
+
+    pub(super) async fn remove_delegate(&self, session_id: wire::SessionId) {
+        self.state.delegates.lock().await.remove(&session_id);
+    }
+}
+
+impl Drop for Connection {
+    fn drop(&mut self) {
+        self.cancellation.cancel();
+    }
+}
+
+struct ConnectionState {
+    outgoing_tx: mpsc::Sender<wire::ClientMessage>,
+    pending:
+        Mutex<HashMap<wire::RequestId, oneshot::Sender<Result<wire::HostResponse, RequestError>>>>,
+    delegates: Mutex<HashMap<wire::SessionId, Arc<dyn CodeModeSessionDelegate>>>,
+    callback_cancellations: Mutex<HashMap<wire::CallbackId, CancellationToken>>,
+    next_request_id: AtomicU64,
+    alive: AtomicBool,
+    failure: std::sync::Mutex<Option<String>>,
+}
+
+impl ConnectionState {
+    fn new(outgoing_tx: mpsc::Sender<wire::ClientMessage>) -> Self {
+        Self {
+            outgoing_tx,
+            pending: Mutex::new(HashMap::new()),
+            delegates: Mutex::new(HashMap::new()),
+            callback_cancellations: Mutex::new(HashMap::new()),
+            next_request_id: AtomicU64::new(1),
+            alive: AtomicBool::new(true),
+            failure: std::sync::Mutex::new(None),
+        }
+    }
+
+    async fn send(&self, message: wire::ClientMessage) -> Result<(), String> {
+        if !self.alive.load(Ordering::Acquire) {
+            return Err(self.failure_message());
+        }
+        self.outgoing_tx
+            .send(message)
+            .await
+            .map_err(|_| self.failure_message())
+    }
+
+    fn failure_message(&self) -> String {
+        self.failure
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .clone()
+            .unwrap_or_else(|| "code-mode host connection closed".to_string())
+    }
+
+    async fn fail(&self, reason: String) {
+        if !self.alive.swap(false, Ordering::AcqRel) {
+            return;
+        }
+        warn!(%reason, "code-mode host connection failed");
+        *self.failure.lock().unwrap_or_else(PoisonError::into_inner) = Some(reason.clone());
+        for (_, sender) in self.pending.lock().await.drain() {
+            let _ = sender.send(Err(RequestError::Transport(reason.clone())));
+        }
+        for (_, cancellation) in self.callback_cancellations.lock().await.drain() {
+            cancellation.cancel();
+        }
+        self.delegates.lock().await.clear();
+    }
+}
+
+struct PendingRequest {
+    cleanup: Option<PendingRequestCleanup>,
+}
+
+struct PendingRequestCleanup {
+    state: Arc<ConnectionState>,
+    id: wire::RequestId,
+    request_sent: bool,
+}
+
+impl PendingRequest {
+    fn new(state: Arc<ConnectionState>, id: wire::RequestId) -> Self {
+        Self {
+            cleanup: Some(PendingRequestCleanup {
+                state,
+                id,
+                request_sent: false,
+            }),
+        }
+    }
+
+    fn mark_request_sent(&mut self) {
+        if let Some(cleanup) = self.cleanup.as_mut() {
+            cleanup.request_sent = true;
+        }
+    }
+
+    fn disarm(&mut self) {
+        self.cleanup = None;
+    }
+
+    async fn cleanup(&mut self) {
+        if let Some(task) = self.spawn_cleanup() {
+            let _ = task.await;
+        }
+    }
+
+    fn spawn_cleanup(&mut self) -> Option<tokio::task::JoinHandle<()>> {
+        let handle = tokio::runtime::Handle::try_current().ok()?;
+        let cleanup = self.cleanup.take()?;
+        Some(handle.spawn(cleanup.run()))
+    }
+}
+
+impl Drop for PendingRequest {
+    fn drop(&mut self) {
+        std::mem::drop(self.spawn_cleanup());
+    }
+}
+
+impl PendingRequestCleanup {
+    async fn run(self) {
+        self.state.pending.lock().await.remove(&self.id);
+        if self.request_sent {
+            let _ = self
+                .state
+                .send(wire::ClientMessage::CancelRequest { id: self.id })
+                .await;
+        }
+    }
+}
+
+async fn drive_reader(
+    mut stdout: tokio::process::ChildStdout,
+    state: Weak<ConnectionState>,
+    cancellation: CancellationToken,
+) {
+    loop {
+        let message = tokio::select! {
+            _ = cancellation.cancelled() => return,
+            result = wire::read_frame::<_, wire::HostMessage>(&mut stdout) => result,
+        };
+        match message {
+            Ok(Some(message)) => {
+                let Some(state) = state.upgrade() else {
+                    return;
+                };
+                handle_host_message(state, message).await;
+            }
+            Ok(None) => {
+                if let Some(state) = state.upgrade() {
+                    state
+                        .fail("code-mode host closed its stdout".to_string())
+                        .await;
+                }
+                return;
+            }
+            Err(err) => {
+                if let Some(state) = state.upgrade() {
+                    state
+                        .fail(format!("failed to read code-mode host message: {err}"))
+                        .await;
+                }
+                return;
+            }
+        }
+    }
+}
+
+async fn handle_host_message(state: Arc<ConnectionState>, message: wire::HostMessage) {
+    match message {
+        wire::HostMessage::Response { id, result } => {
+            if let Some(sender) = state.pending.lock().await.remove(&id) {
+                let result = match result {
+                    wire::WireResult::Ok { value } => Ok(value),
+                    wire::WireResult::Err { error } => Err(RequestError::Host(error)),
+                };
+                let _ = sender.send(result);
+            }
+        }
+        wire::HostMessage::CellClosed {
+            session_id,
+            cell_id,
+        } => {
+            if let Some(delegate) = state.delegates.lock().await.get(&session_id).cloned() {
+                delegate.cell_closed(&convert::protocol_cell_id(&cell_id));
+            }
+        }
+        wire::HostMessage::CallbackRequest {
+            id,
+            session_id,
+            request,
+        } => {
+            start_callback(state, id, session_id, request).await;
+        }
+        wire::HostMessage::CancelCallback { id } => {
+            if let Some(cancellation) = state.callback_cancellations.lock().await.remove(&id) {
+                cancellation.cancel();
+            }
+        }
+    }
+}
+
+async fn start_callback(
+    state: Arc<ConnectionState>,
+    id: wire::CallbackId,
+    session_id: wire::SessionId,
+    request: wire::CallbackRequest,
+) {
+    let delegate = state.delegates.lock().await.get(&session_id).cloned();
+    let Some(delegate) = delegate else {
+        let response = match request {
+            wire::CallbackRequest::InvokeTool { .. } => wire::CallbackResponse::ToolError {
+                error_text: format!("code-mode session {session_id} not found"),
+            },
+            wire::CallbackRequest::Notify { .. } => wire::CallbackResponse::NotificationError {
+                error_text: format!("code-mode session {session_id} not found"),
+            },
+        };
+        let _ = state
+            .send(wire::ClientMessage::CallbackResponse { id, response })
+            .await;
+        return;
+    };
+
+    let cancellation = CancellationToken::new();
+    state
+        .callback_cancellations
+        .lock()
+        .await
+        .insert(id, cancellation.clone());
+    tokio::spawn(async move {
+        let response = match request {
+            wire::CallbackRequest::InvokeTool { invocation } => delegate
+                .invoke_tool(convert::nested_tool_call(invocation), cancellation)
+                .await
+                .map_or_else(
+                    |error_text| wire::CallbackResponse::ToolError { error_text },
+                    |result| wire::CallbackResponse::ToolResult { result },
+                ),
+            wire::CallbackRequest::Notify {
+                call_id,
+                cell_id,
+                text,
+            } => delegate
+                .notify(
+                    call_id,
+                    convert::protocol_cell_id(&cell_id),
+                    text,
+                    cancellation,
+                )
+                .await
+                .map_or_else(
+                    |error_text| wire::CallbackResponse::NotificationError { error_text },
+                    |()| wire::CallbackResponse::NotificationDelivered,
+                ),
+        };
+        state.callback_cancellations.lock().await.remove(&id);
+        let _ = state
+            .send(wire::ClientMessage::CallbackResponse { id, response })
+            .await;
+    });
+}
+
+fn spawn_critical_connection_task(
+    description: &'static str,
+    state: Weak<ConnectionState>,
+    future: impl Future<Output = ()> + Send + 'static,
+) {
+    let task = tokio::spawn(future);
+    tokio::spawn(async move {
+        if let Err(err) = task.await
+            && let Some(state) = state.upgrade()
+        {
+            state
+                .fail(format!("code-mode host {description} task failed: {err}"))
+                .await;
+        }
+    });
+}
+
+fn host_process(command: &CodeModeHostCommand) -> Command {
+    let mut process = Command::new(&command.program);
+    process.args(&command.args);
+    #[cfg(unix)]
+    process.process_group(0);
+    process
+}

codex-rs/code-mode-client/src/convert.rs

@@ -0,0 +1,144 @@
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeNestedToolCall;
+use codex_code_mode_protocol::CodeModeToolKind;
+use codex_code_mode_protocol::CreateCellRequest;
+use codex_code_mode_protocol::FunctionCallOutputContentItem;
+use codex_code_mode_protocol::ImageDetail;
+use codex_code_mode_protocol::RuntimeResponse;
+use codex_code_mode_protocol::ToolDefinition;
+use codex_code_mode_protocol::wire;
+
+pub(super) fn create_cell_request(request: CreateCellRequest) -> wire::CreateCellRequest {
+    wire::CreateCellRequest {
+        tool_call_id: request.tool_call_id,
+        enabled_tools: request
+            .enabled_tools
+            .into_iter()
+            .map(tool_definition)
+            .collect(),
+        source: request.source,
+    }
+}
+
+pub(super) fn wire_cell_id(cell_id: &CellId) -> wire::CellId {
+    wire::CellId::new(cell_id.as_str())
+}
+
+pub(super) fn protocol_cell_id(cell_id: &wire::CellId) -> CellId {
+    CellId::new(cell_id.as_str().to_string())
+}
+
+pub(super) fn runtime_response(
+    cell_id: &CellId,
+    event: wire::CellEvent,
+) -> Result<RuntimeResponse, String> {
+    match event {
+        wire::CellEvent::Yielded { content_items } => Ok(RuntimeResponse::Yielded {
+            cell_id: cell_id.clone(),
+            content_items: content_items.into_iter().map(output_item).collect(),
+        }),
+        wire::CellEvent::Completed {
+            content_items,
+            error_text,
+        } => Ok(RuntimeResponse::Result {
+            cell_id: cell_id.clone(),
+            content_items: content_items.into_iter().map(output_item).collect(),
+            error_text,
+        }),
+        wire::CellEvent::Terminated { content_items } => Ok(RuntimeResponse::Terminated {
+            cell_id: cell_id.clone(),
+            content_items: content_items.into_iter().map(output_item).collect(),
+        }),
+        wire::CellEvent::Pending { .. } => {
+            Err("code-mode host returned an unexpected pending frontier".to_string())
+        }
+    }
+}
+
+pub(super) fn missing_cell_response(cell_id: CellId) -> RuntimeResponse {
+    RuntimeResponse::Result {
+        error_text: Some(format!("exec cell {cell_id} not found")),
+        cell_id,
+        content_items: Vec::new(),
+    }
+}
+
+pub(super) fn nested_tool_call(invocation: wire::NestedToolCall) -> CodeModeNestedToolCall {
+    CodeModeNestedToolCall {
+        cell_id: protocol_cell_id(&invocation.cell_id),
+        runtime_tool_call_id: invocation.runtime_tool_call_id,
+        tool_name: codex_protocol::ToolName {
+            name: invocation.tool_name.name,
+            namespace: invocation.tool_name.namespace,
+        },
+        tool_kind: match invocation.tool_kind {
+            wire::ToolKind::Function => CodeModeToolKind::Function,
+            wire::ToolKind::Freeform => CodeModeToolKind::Freeform,
+        },
+        input: invocation.input,
+    }
+}
+
+pub(super) fn wire_error_message(error: &wire::Error) -> String {
+    match error {
+        wire::Error::MissingSession { session_id } => {
+            format!("code-mode session {session_id} not found")
+        }
+        wire::Error::ShuttingDown => "code mode session is shutting down".to_string(),
+        wire::Error::DuplicateCell { cell_id } => {
+            format!("exec cell {} already exists", cell_id.as_str())
+        }
+        wire::Error::MissingCell { cell_id } => {
+            format!("exec cell {} not found", cell_id.as_str())
+        }
+        wire::Error::ClosedCell { cell_id } => {
+            format!("exec cell {} closed unexpectedly", cell_id.as_str())
+        }
+        wire::Error::BusyObserver { cell_id } => {
+            format!(
+                "exec cell {} already has an active observer",
+                cell_id.as_str()
+            )
+        }
+        wire::Error::AlreadyTerminating { cell_id } => {
+            format!("exec cell {} is already terminating", cell_id.as_str())
+        }
+        wire::Error::Cancelled => "code-mode request was cancelled".to_string(),
+        wire::Error::Runtime { message }
+        | wire::Error::InvalidRequest { message }
+        | wire::Error::CallbackFailed { message }
+        | wire::Error::Internal { message } => message.clone(),
+    }
+}
+
+fn tool_definition(definition: ToolDefinition) -> wire::ToolDefinition {
+    wire::ToolDefinition {
+        name: definition.name,
+        tool_name: wire::ToolName {
+            name: definition.tool_name.name,
+            namespace: definition.tool_name.namespace,
+        },
+        description: definition.description,
+        kind: match definition.kind {
+            CodeModeToolKind::Function => wire::ToolKind::Function,
+            CodeModeToolKind::Freeform => wire::ToolKind::Freeform,
+        },
+    }
+}
+
+fn output_item(item: wire::OutputItem) -> FunctionCallOutputContentItem {
+    match item {
+        wire::OutputItem::Text { text } => FunctionCallOutputContentItem::InputText { text },
+        wire::OutputItem::Image { image_url, detail } => {
+            FunctionCallOutputContentItem::InputImage {
+                image_url,
+                detail: detail.map(|detail| match detail {
+                    wire::ImageDetail::Auto => ImageDetail::Auto,
+                    wire::ImageDetail::Low => ImageDetail::Low,
+                    wire::ImageDetail::High => ImageDetail::High,
+                    wire::ImageDetail::Original => ImageDetail::Original,
+                }),
+            }
+        }
+    }
+}

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

@@ -0,0 +1,263 @@
+use std::path::PathBuf;
+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::CellOutcome;
+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::CreateCellRequest;
+use codex_code_mode_protocol::ObserveRequest;
+use codex_code_mode_protocol::wire;
+use tokio::sync::Semaphore;
+
+mod connection;
+mod convert;
+
+use connection::Connection;
+use connection::RequestError;
+
+const CODE_MODE_HOST_PATH_ENV: &str = "CODEX_CODE_MODE_HOST_PATH";
+
+#[derive(Clone, Debug)]
+pub struct CodeModeHostCommand {
+    pub program: PathBuf,
+    pub args: Vec<String>,
+}
+
+impl Default for CodeModeHostCommand {
+    fn default() -> Self {
+        Self {
+            program: default_host_program(),
+            args: Vec::new(),
+        }
+    }
+}
+
+pub struct StdioCodeModeSessionProvider {
+    command: CodeModeHostCommand,
+    connection: std::sync::Mutex<Option<Arc<Connection>>>,
+    spawn_permit: Semaphore,
+}
+
+impl StdioCodeModeSessionProvider {
+    pub fn new(command: CodeModeHostCommand) -> Self {
+        Self {
+            command,
+            connection: std::sync::Mutex::new(None),
+            spawn_permit: Semaphore::new(/*permits*/ 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 connection = Arc::new(Connection::spawn(&self.command).await?);
+        *self
+            .connection
+            .lock()
+            .map_err(|_| "code-mode host connection lock poisoned".to_string())? =
+            Some(Arc::clone(&connection));
+        Ok(connection)
+    }
+
+    fn live_connection(&self) -> Result<Option<Arc<Connection>>, String> {
+        let connection = self
+            .connection
+            .lock()
+            .map_err(|_| "code-mode host connection lock poisoned".to_string())?;
+        Ok(connection
+            .as_ref()
+            .filter(|connection| connection.is_alive())
+            .cloned())
+    }
+}
+
+impl Default for StdioCodeModeSessionProvider {
+    fn default() -> Self {
+        Self::new(CodeModeHostCommand::default())
+    }
+}
+
+impl CodeModeSessionProvider for StdioCodeModeSessionProvider {
+    fn create_session<'a>(
+        &'a self,
+        delegate: Arc<dyn CodeModeSessionDelegate>,
+    ) -> CodeModeSessionProviderFuture<'a> {
+        Box::pin(async move {
+            let connection = self.connection().await?;
+            let response = connection
+                .request(wire::HostRequest::CreateSession)
+                .await
+                .map_err(RequestError::into_message)?;
+            let wire::HostResponse::SessionCreated { session_id } = response else {
+                return Err(
+                    "code-mode host returned an invalid create-session response".to_string()
+                );
+            };
+            connection.register_delegate(session_id, delegate).await;
+            let session: Arc<dyn CodeModeSession> = Arc::new(StdioCodeModeSession {
+                connection,
+                session_id,
+                shutdown: AtomicBool::new(false),
+            });
+            Ok(session)
+        })
+    }
+}
+
+struct StdioCodeModeSession {
+    connection: Arc<Connection>,
+    session_id: wire::SessionId,
+    shutdown: AtomicBool,
+}
+
+impl CodeModeSession for StdioCodeModeSession {
+    fn is_alive(&self) -> bool {
+        !self.shutdown.load(Ordering::Acquire) && self.connection.is_alive()
+    }
+
+    fn create_cell<'a>(
+        &'a self,
+        request: CreateCellRequest,
+    ) -> CodeModeSessionResultFuture<'a, CellId> {
+        Box::pin(async move {
+            self.ensure_active()?;
+            let response = self
+                .connection
+                .request(wire::HostRequest::CreateCell {
+                    session_id: self.session_id,
+                    request: convert::create_cell_request(request),
+                })
+                .await
+                .map_err(RequestError::into_message)?;
+            match response {
+                wire::HostResponse::CellCreated { cell_id } => {
+                    Ok(convert::protocol_cell_id(&cell_id))
+                }
+                _ => Err("code-mode host returned an invalid create-cell response".to_string()),
+            }
+        })
+    }
+
+    fn observe<'a>(
+        &'a self,
+        request: ObserveRequest,
+    ) -> CodeModeSessionResultFuture<'a, CellOutcome> {
+        Box::pin(async move {
+            self.ensure_active()?;
+            let cell_id = request.cell_id;
+            let response = self
+                .connection
+                .request(wire::HostRequest::Observe {
+                    session_id: self.session_id,
+                    cell_id: convert::wire_cell_id(&cell_id),
+                    mode: wire::ObserveMode::YieldAfter {
+                        duration_ms: request.yield_time_ms,
+                    },
+                })
+                .await;
+            match response {
+                Ok(wire::HostResponse::Observed { event }) => Ok(CellOutcome::LiveCell(
+                    convert::runtime_response(&cell_id, event)?,
+                )),
+                Ok(_) => Err("code-mode host returned an invalid observe response".to_string()),
+                Err(RequestError::Host(wire::Error::MissingCell { .. })) => Ok(
+                    CellOutcome::MissingCell(convert::missing_cell_response(cell_id)),
+                ),
+                Err(error) => Err(error.into_message()),
+            }
+        })
+    }
+
+    fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, CellOutcome> {
+        Box::pin(async move {
+            self.ensure_active()?;
+            let response = self
+                .connection
+                .request(wire::HostRequest::Terminate {
+                    session_id: self.session_id,
+                    cell_id: convert::wire_cell_id(&cell_id),
+                })
+                .await;
+            match response {
+                Ok(wire::HostResponse::Observed { event }) => Ok(CellOutcome::LiveCell(
+                    convert::runtime_response(&cell_id, event)?,
+                )),
+                Ok(_) => Err("code-mode host returned an invalid terminate response".to_string()),
+                Err(RequestError::Host(wire::Error::MissingCell { .. })) => Ok(
+                    CellOutcome::MissingCell(convert::missing_cell_response(cell_id)),
+                ),
+                Err(error) => Err(error.into_message()),
+            }
+        })
+    }
+
+    fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> {
+        Box::pin(async move {
+            if self.shutdown.swap(true, Ordering::AcqRel) {
+                return Ok(());
+            }
+            let result = self
+                .connection
+                .request(wire::HostRequest::ShutdownSession {
+                    session_id: self.session_id,
+                })
+                .await;
+            self.connection.remove_delegate(self.session_id).await;
+            match result.map_err(RequestError::into_message)? {
+                wire::HostResponse::SessionShutdown => Ok(()),
+                _ => Err("code-mode host returned an invalid shutdown response".to_string()),
+            }
+        })
+    }
+}
+
+impl StdioCodeModeSession {
+    fn ensure_active(&self) -> Result<(), String> {
+        if self.shutdown.load(Ordering::Acquire) {
+            Err("code mode session is shutting down".to_string())
+        } else {
+            Ok(())
+        }
+    }
+}
+
+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 = "tests.rs"]
+mod tests;

codex-rs/code-mode-client/src/tests.rs

@@ -0,0 +1,53 @@
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use codex_code_mode_protocol::CellId;
+use codex_code_mode_protocol::CodeModeNestedToolCall;
+use codex_code_mode_protocol::CodeModeSessionDelegate;
+use codex_code_mode_protocol::CodeModeSessionProvider;
+use codex_code_mode_protocol::NotificationFuture;
+use codex_code_mode_protocol::ToolInvocationFuture;
+use tokio_util::sync::CancellationToken;
+
+use super::CodeModeHostCommand;
+use super::StdioCodeModeSessionProvider;
+
+struct TestDelegate;
+
+impl CodeModeSessionDelegate for TestDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        _invocation: CodeModeNestedToolCall,
+        _cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        Box::pin(async { Err("unexpected tool invocation".to_string()) })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        _call_id: String,
+        _cell_id: CellId,
+        _text: String,
+        _cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        Box::pin(async { Ok(()) })
+    }
+
+    fn cell_closed(&self, _cell_id: &CellId) {}
+}
+
+#[tokio::test]
+async fn create_session_reports_host_spawn_errors() {
+    let provider = StdioCodeModeSessionProvider::new(CodeModeHostCommand {
+        program: PathBuf::from("codex-code-mode-host-does-not-exist"),
+        args: Vec::new(),
+    });
+
+    let error = provider
+        .create_session(Arc::new(TestDelegate))
+        .await
+        .err()
+        .expect("session creation should fail");
+
+    assert!(error.contains("failed to spawn code-mode host"));
+}

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

@@ -19,5 +19,7 @@ tokio = { workspace = true, features = ["io-std", "io-util", "macros", "process"
 tokio-util = { workspace = true, features = ["rt"] }
 
 [dev-dependencies]
+codex-code-mode-client = { workspace = true }
+codex-protocol = { workspace = true }
 codex-utils-cargo-bin = { workspace = true }
 pretty_assertions = { workspace = true }

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

@@ -0,0 +1,422 @@
+#![allow(clippy::expect_used)]
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use codex_code_mode::CellId;
+use codex_code_mode::CellOutcome;
+use codex_code_mode::CodeModeNestedToolCall;
+use codex_code_mode::CodeModeSession;
+use codex_code_mode::CodeModeSessionDelegate;
+use codex_code_mode::CodeModeSessionProvider;
+use codex_code_mode::CodeModeToolKind;
+use codex_code_mode::CreateCellRequest;
+use codex_code_mode::FunctionCallOutputContentItem;
+use codex_code_mode::InProcessCodeModeSessionProvider;
+use codex_code_mode::NotificationFuture;
+use codex_code_mode::ObserveRequest;
+use codex_code_mode::RuntimeResponse;
+use codex_code_mode::ToolDefinition;
+use codex_code_mode::ToolInvocationFuture;
+use codex_code_mode_client::CodeModeHostCommand;
+use codex_code_mode_client::StdioCodeModeSessionProvider;
+use codex_protocol::ToolName;
+use pretty_assertions::assert_eq;
+use tokio::sync::mpsc;
+use tokio_util::sync::CancellationToken;
+
+struct ProviderCase {
+    name: &'static str,
+    provider: Arc<dyn CodeModeSessionProvider>,
+}
+
+fn provider_cases() -> Vec<ProviderCase> {
+    let host_program =
+        codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("resolve host binary");
+    vec![
+        ProviderCase {
+            name: "in-process",
+            provider: Arc::new(InProcessCodeModeSessionProvider),
+        },
+        ProviderCase {
+            name: "stdio",
+            provider: Arc::new(StdioCodeModeSessionProvider::new(CodeModeHostCommand {
+                program: host_program,
+                args: Vec::new(),
+            })),
+        },
+    ]
+}
+
+#[derive(Debug, PartialEq)]
+enum DelegateEvent {
+    NotificationStarted,
+    ToolStarted,
+    CellClosed(CellId),
+}
+
+struct BlockingDelegate {
+    events_tx: mpsc::UnboundedSender<DelegateEvent>,
+}
+
+impl BlockingDelegate {
+    fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
+        let (events_tx, events_rx) = mpsc::unbounded_channel();
+        (Arc::new(Self { events_tx }), events_rx)
+    }
+}
+
+impl CodeModeSessionDelegate for BlockingDelegate {
+    fn invoke_tool<'a>(
+        &'a self,
+        _invocation: CodeModeNestedToolCall,
+        cancellation_token: CancellationToken,
+    ) -> ToolInvocationFuture<'a> {
+        Box::pin(async move {
+            let _ = self.events_tx.send(DelegateEvent::ToolStarted);
+            let _ = cancellation_token;
+            std::future::pending().await
+        })
+    }
+
+    fn notify<'a>(
+        &'a self,
+        _call_id: String,
+        _cell_id: CellId,
+        _text: String,
+        cancellation_token: CancellationToken,
+    ) -> NotificationFuture<'a> {
+        Box::pin(async move {
+            let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
+            let _ = cancellation_token;
+            std::future::pending().await
+        })
+    }
+
+    fn cell_closed(&self, cell_id: &CellId) {
+        let _ = self
+            .events_tx
+            .send(DelegateEvent::CellClosed(cell_id.clone()));
+    }
+}
+
+fn create_cell_request(source: &str) -> CreateCellRequest {
+    CreateCellRequest {
+        tool_call_id: "call-1".to_string(),
+        enabled_tools: Vec::new(),
+        source: source.to_string(),
+    }
+}
+
+fn blocking_tool() -> ToolDefinition {
+    ToolDefinition {
+        name: "block".to_string(),
+        tool_name: ToolName::plain("block"),
+        description: String::new(),
+        kind: CodeModeToolKind::Function,
+        input_schema: None,
+        output_schema: None,
+    }
+}
+
+async fn create_session(
+    provider: &ProviderCase,
+    delegate: Arc<dyn CodeModeSessionDelegate>,
+) -> Arc<dyn CodeModeSession> {
+    provider
+        .provider
+        .create_session(delegate)
+        .await
+        .unwrap_or_else(|err| panic!("{} session creation failed: {err}", provider.name))
+}
+
+async fn observe(
+    session: &Arc<dyn CodeModeSession>,
+    cell_id: &CellId,
+    yield_time_ms: u64,
+) -> Result<CellOutcome, String> {
+    session
+        .observe(ObserveRequest {
+            cell_id: cell_id.clone(),
+            yield_time_ms,
+        })
+        .await
+}
+
+async fn next_event(
+    provider_name: &str,
+    events_rx: &mut mpsc::UnboundedReceiver<DelegateEvent>,
+) -> DelegateEvent {
+    tokio::time::timeout(Duration::from_secs(2), events_rx.recv())
+        .await
+        .unwrap_or_else(|_| panic!("{provider_name} delegate event timeout"))
+        .unwrap_or_else(|| panic!("{provider_name} delegate event channel closed"))
+}
+
+#[tokio::test]
+async fn providers_yield_and_resume() {
+    for provider in provider_cases() {
+        let (delegate, _) = BlockingDelegate::new();
+        let session = create_session(&provider, delegate).await;
+        let cell_id = session
+            .create_cell(create_cell_request(
+                r#"text("before"); yield_control(); text("after");"#,
+            ))
+            .await
+            .unwrap_or_else(|err| panic!("{} create failed: {err}", provider.name));
+
+        assert_eq!(
+            observe(&session, &cell_id, 60_000).await.unwrap(),
+            CellOutcome::LiveCell(RuntimeResponse::Yielded {
+                cell_id: cell_id.clone(),
+                content_items: vec![FunctionCallOutputContentItem::InputText {
+                    text: "before".to_string(),
+                }],
+            }),
+            "{}",
+            provider.name
+        );
+        assert_eq!(
+            observe(&session, &cell_id, 1).await.unwrap(),
+            CellOutcome::LiveCell(RuntimeResponse::Result {
+                cell_id: cell_id.clone(),
+                content_items: vec![FunctionCallOutputContentItem::InputText {
+                    text: "after".to_string(),
+                }],
+                error_text: None,
+            }),
+            "{}",
+            provider.name
+        );
+        session.shutdown().await.unwrap();
+    }
+}
+
+#[tokio::test]
+async fn providers_preserve_natural_completion_when_termination_arrives_late() {
+    for provider in provider_cases() {
+        let (delegate, _) = BlockingDelegate::new();
+        let session = create_session(&provider, delegate).await;
+        let cell_id = session
+            .create_cell(create_cell_request(
+                r#"yield_control(); store("finished", true); text("done");"#,
+            ))
+            .await
+            .unwrap();
+        assert!(matches!(
+            observe(&session, &cell_id, 60_000).await.unwrap(),
+            CellOutcome::LiveCell(RuntimeResponse::Yielded { .. })
+        ));
+
+        tokio::time::timeout(Duration::from_secs(2), async {
+            loop {
+                let probe_id = session
+                    .create_cell(create_cell_request(r#"text(String(load("finished")));"#))
+                    .await
+                    .unwrap();
+                let response = observe(&session, &probe_id, 60_000).await.unwrap();
+                let CellOutcome::LiveCell(RuntimeResponse::Result { content_items, .. }) = response
+                else {
+                    panic!("{} stored-value probe did not complete", provider.name);
+                };
+                if content_items
+                    == vec![FunctionCallOutputContentItem::InputText {
+                        text: "true".to_string(),
+                    }]
+                {
+                    break;
+                }
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .unwrap_or_else(|_| panic!("{} completion probe timed out", provider.name));
+
+        assert_eq!(
+            session.terminate(cell_id.clone()).await.unwrap(),
+            CellOutcome::LiveCell(RuntimeResponse::Result {
+                cell_id,
+                content_items: vec![FunctionCallOutputContentItem::InputText {
+                    text: "done".to_string(),
+                }],
+                error_text: None,
+            }),
+            "{}",
+            provider.name
+        );
+        session.shutdown().await.unwrap();
+    }
+}
+
+#[tokio::test]
+async fn providers_cancel_callbacks_before_termination_returns() {
+    for provider in provider_cases() {
+        let (delegate, mut events_rx) = BlockingDelegate::new();
+        let session = create_session(&provider, delegate).await;
+        let cell_id = session
+            .create_cell(create_cell_request(
+                r#"notify("pending"); await new Promise(() => {});"#,
+            ))
+            .await
+            .unwrap();
+        let observing_session = Arc::clone(&session);
+        let observed_cell_id = cell_id.clone();
+        let observer =
+            tokio::spawn(
+                async move { observe(&observing_session, &observed_cell_id, 60_000).await },
+            );
+
+        assert_eq!(
+            next_event(provider.name, &mut events_rx).await,
+            DelegateEvent::NotificationStarted
+        );
+        let terminated = CellOutcome::LiveCell(RuntimeResponse::Terminated {
+            cell_id: cell_id.clone(),
+            content_items: Vec::new(),
+        });
+        assert_eq!(
+            session.terminate(cell_id).await.unwrap(),
+            terminated,
+            "{}",
+            provider.name
+        );
+        assert_eq!(
+            observer.await.unwrap().unwrap(),
+            terminated,
+            "{}",
+            provider.name
+        );
+        session.shutdown().await.unwrap();
+    }
+}
+
+#[tokio::test]
+async fn providers_reject_a_second_observer_without_displacing_the_first() {
+    for provider in provider_cases() {
+        let (delegate, _) = BlockingDelegate::new();
+        let session = create_session(&provider, delegate).await;
+        let cell_id = session
+            .create_cell(create_cell_request("await new Promise(() => {});"))
+            .await
+            .unwrap();
+        assert!(matches!(
+            observe(&session, &cell_id, 1).await.unwrap(),
+            CellOutcome::LiveCell(RuntimeResponse::Yielded { .. })
+        ));
+
+        let mut first_observer = session.observe(ObserveRequest {
+            cell_id: cell_id.clone(),
+            yield_time_ms: 60_000,
+        });
+        assert!(
+            tokio::time::timeout(Duration::from_millis(20), &mut first_observer)
+                .await
+                .is_err(),
+            "{} first observer returned unexpectedly",
+            provider.name
+        );
+        assert_eq!(
+            observe(&session, &cell_id, 60_000).await.unwrap_err(),
+            format!("exec cell {cell_id} already has an active observer"),
+            "{}",
+            provider.name
+        );
+
+        let terminated = CellOutcome::LiveCell(RuntimeResponse::Terminated {
+            cell_id: cell_id.clone(),
+            content_items: Vec::new(),
+        });
+        assert_eq!(
+            session.terminate(cell_id).await.unwrap(),
+            terminated,
+            "{}",
+            provider.name
+        );
+        assert_eq!(
+            first_observer.await.unwrap(),
+            terminated,
+            "{}",
+            provider.name
+        );
+        session.shutdown().await.unwrap();
+    }
+}
+
+#[tokio::test]
+async fn providers_cleanup_callbacks_before_natural_completion_returns() {
+    for provider in provider_cases() {
+        let (delegate, mut events_rx) = BlockingDelegate::new();
+        let session = create_session(&provider, delegate).await;
+        let mut request = create_cell_request(
+            r#"
+tools.block({});
+await new Promise((resolve) => setTimeout(resolve, 0));
+text("done");
+"#,
+        );
+        request.enabled_tools.push(blocking_tool());
+        let cell_id = session.create_cell(request).await.unwrap();
+
+        assert_eq!(
+            next_event(provider.name, &mut events_rx).await,
+            DelegateEvent::ToolStarted
+        );
+        assert_eq!(
+            observe(&session, &cell_id, 60_000).await.unwrap(),
+            CellOutcome::LiveCell(RuntimeResponse::Result {
+                cell_id,
+                content_items: vec![FunctionCallOutputContentItem::InputText {
+                    text: "done".to_string(),
+                }],
+                error_text: None,
+            }),
+            "{}",
+            provider.name
+        );
+        session.shutdown().await.unwrap();
+    }
+}
+
+#[tokio::test]
+async fn providers_report_background_completion_to_the_delegate() {
+    for provider in provider_cases() {
+        let (delegate, mut events_rx) = BlockingDelegate::new();
+        let session = create_session(&provider, delegate).await;
+        let cell_id = session
+            .create_cell(create_cell_request(concat!(
+                "await new Promise(resolve => setTimeout(resolve, 100));",
+                "text('done');",
+            )))
+            .await
+            .unwrap();
+
+        assert_eq!(
+            observe(&session, &cell_id, 1).await.unwrap(),
+            CellOutcome::LiveCell(RuntimeResponse::Yielded {
+                cell_id: cell_id.clone(),
+                content_items: Vec::new(),
+            }),
+            "{}",
+            provider.name
+        );
+        assert_eq!(
+            next_event(provider.name, &mut events_rx).await,
+            DelegateEvent::CellClosed(cell_id.clone()),
+            "{}",
+            provider.name
+        );
+        assert_eq!(
+            observe(&session, &cell_id, 60_000).await.unwrap(),
+            CellOutcome::LiveCell(RuntimeResponse::Result {
+                cell_id,
+                content_items: vec![FunctionCallOutputContentItem::InputText {
+                    text: "done".to_string(),
+                }],
+                error_text: None,
+            }),
+            "{}",
+            provider.name
+        );
+        session.shutdown().await.unwrap();
+    }
+}