From a12f41d4e87992b44ad9e694b6e650cd553331d0 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 16 Jun 2026 05:10:25 +0000 Subject: [PATCH] code-mode: add and verify stdio session provider --- codex-rs/Cargo.lock | 13 + codex-rs/Cargo.toml | 2 + codex-rs/code-mode-client/BUILD.bazel | 6 + codex-rs/code-mode-client/Cargo.toml | 20 + codex-rs/code-mode-client/src/connection.rs | 481 ++++++++++++++++++ codex-rs/code-mode-client/src/convert.rs | 144 ++++++ codex-rs/code-mode-client/src/lib.rs | 263 ++++++++++ codex-rs/code-mode-client/src/tests.rs | 53 ++ codex-rs/code-mode-host/Cargo.toml | 2 + .../code-mode-host/tests/provider_contract.rs | 422 +++++++++++++++ 10 files changed, 1406 insertions(+) create mode 100644 codex-rs/code-mode-client/BUILD.bazel create mode 100644 codex-rs/code-mode-client/Cargo.toml create mode 100644 codex-rs/code-mode-client/src/connection.rs create mode 100644 codex-rs/code-mode-client/src/convert.rs create mode 100644 codex-rs/code-mode-client/src/lib.rs create mode 100644 codex-rs/code-mode-client/src/tests.rs create mode 100644 codex-rs/code-mode-host/tests/provider_contract.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 2ef511cdfb..2dae185061 100644 --- a/codex-rs/Cargo.lock +++ b/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", diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 667bb7dfdf..69e32a893f 100644 --- a/codex-rs/Cargo.toml +++ b/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" } diff --git a/codex-rs/code-mode-client/BUILD.bazel b/codex-rs/code-mode-client/BUILD.bazel new file mode 100644 index 0000000000..071598cc2c --- /dev/null +++ b/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", +) diff --git a/codex-rs/code-mode-client/Cargo.toml b/codex-rs/code-mode-client/Cargo.toml new file mode 100644 index 0000000000..e1cc680400 --- /dev/null +++ b/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 } diff --git a/codex-rs/code-mode-client/src/connection.rs b/codex-rs/code-mode-client/src/connection.rs new file mode 100644 index 0000000000..155adac070 --- /dev/null +++ b/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, + cancellation: CancellationToken, +} + +impl Connection { + pub(super) async fn spawn(command: &CodeModeHostCommand) -> Result { + 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 { + 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, + ) { + 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, + pending: + Mutex>>>, + delegates: Mutex>>, + callback_cancellations: Mutex>, + next_request_id: AtomicU64, + alive: AtomicBool, + failure: std::sync::Mutex>, +} + +impl ConnectionState { + fn new(outgoing_tx: mpsc::Sender) -> 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, +} + +struct PendingRequestCleanup { + state: Arc, + id: wire::RequestId, + request_sent: bool, +} + +impl PendingRequest { + fn new(state: Arc, 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> { + 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, + 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, 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, + 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, + future: impl Future + 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 +} diff --git a/codex-rs/code-mode-client/src/convert.rs b/codex-rs/code-mode-client/src/convert.rs new file mode 100644 index 0000000000..7b346b75da --- /dev/null +++ b/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 { + 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, + }), + } + } + } +} diff --git a/codex-rs/code-mode-client/src/lib.rs b/codex-rs/code-mode-client/src/lib.rs new file mode 100644 index 0000000000..b526fac421 --- /dev/null +++ b/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, +} + +impl Default for CodeModeHostCommand { + fn default() -> Self { + Self { + program: default_host_program(), + args: Vec::new(), + } + } +} + +pub struct StdioCodeModeSessionProvider { + command: CodeModeHostCommand, + connection: std::sync::Mutex>>, + 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, 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>, 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, + ) -> 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 = Arc::new(StdioCodeModeSession { + connection, + session_id, + shutdown: AtomicBool::new(false), + }); + Ok(session) + }) + } +} + +struct StdioCodeModeSession { + connection: Arc, + 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; diff --git a/codex-rs/code-mode-client/src/tests.rs b/codex-rs/code-mode-client/src/tests.rs new file mode 100644 index 0000000000..e1476a3e0f --- /dev/null +++ b/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")); +} diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 6a34b8bb79..6aed00ba5a 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/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 } diff --git a/codex-rs/code-mode-host/tests/provider_contract.rs b/codex-rs/code-mode-host/tests/provider_contract.rs new file mode 100644 index 0000000000..b7f2766d8a --- /dev/null +++ b/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, +} + +fn provider_cases() -> Vec { + 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, +} + +impl BlockingDelegate { + fn new() -> (Arc, mpsc::UnboundedReceiver) { + 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, +) -> Arc { + provider + .provider + .create_session(delegate) + .await + .unwrap_or_else(|err| panic!("{} session creation failed: {err}", provider.name)) +} + +async fn observe( + session: &Arc, + cell_id: &CellId, + yield_time_ms: u64, +) -> Result { + 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 { + 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(); + } +}