diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 2c50f268f0..dfe4a7fa67 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2514,6 +2514,17 @@ dependencies = [ [[package]] name = "codex-code-mode-host" version = "0.0.0" +dependencies = [ + "anyhow", + "codex-code-mode", + "codex-code-mode-protocol", + "codex-protocol", + "codex-utils-cargo-bin", + "pretty_assertions", + "serde_json", + "tokio", + "tokio-util", +] [[package]] name = "codex-code-mode-protocol" diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index a2c384d010..2be4713410 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -8,5 +8,23 @@ license.workspace = true name = "codex-code-mode-host" path = "src/main.rs" +[lib] +doctest = false +name = "codex_code_mode_host" +path = "src/lib.rs" + [lints] workspace = true + +[dependencies] +anyhow = { workspace = true } +codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } +tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt"] } +tokio-util = { workspace = true, features = ["rt"] } + +[dev-dependencies] +codex-protocol = { workspace = true } +codex-utils-cargo-bin = { workspace = true } +pretty_assertions = { workspace = true } +serde_json = { workspace = true } diff --git a/codex-rs/code-mode-host/src/delegate.rs b/codex-rs/code-mode-host/src/delegate.rs new file mode 100644 index 0000000000..b4fee25c0a --- /dev/null +++ b/codex-rs/code-mode-host/src/delegate.rs @@ -0,0 +1,88 @@ +use std::sync::Arc; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::SessionId; +use tokio_util::sync::CancellationToken; + +use crate::peer::HostPeer; + +pub(super) struct RemoteDelegate { + session_id: SessionId, + peer: Arc, +} + +impl RemoteDelegate { + pub(super) fn new(session_id: SessionId, peer: Arc) -> Self { + Self { session_id, peer } + } +} + +impl CodeModeSessionDelegate for RemoteDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id.clone(), + DelegateRequest::InvokeTool { + invocation: invocation.into(), + }, + cancellation_token, + ) + .await? + { + DelegateResponse::ToolResult { result } => Ok(result), + DelegateResponse::NotificationDelivered => { + Err("code-mode client returned an invalid tool result".to_string()) + } + } + }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id.clone(), + DelegateRequest::Notify { + call_id, + cell_id: cell_id.into(), + text, + }, + cancellation_token, + ) + .await? + { + DelegateResponse::NotificationDelivered => Ok(()), + DelegateResponse::ToolResult { .. } => { + Err("code-mode client returned an invalid notification result".to_string()) + } + } + }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.peer.send(HostToClient::CellClosed { + session_id: self.session_id.clone(), + cell_id: cell_id.into(), + }); + } +} diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs new file mode 100644 index 0000000000..4521c8c3ad --- /dev/null +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -0,0 +1,257 @@ +use codex_code_mode_protocol::host::Capability; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientHello; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HandshakeRejectReason; +use codex_code_mode_protocol::host::HostHello; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use codex_code_mode_protocol::host::WireResult; +use pretty_assertions::assert_eq; + +use super::run; + +fn client_hello( + versions: impl IntoIterator, + required_capabilities: CapabilitySet, +) -> ClientToHost { + ClientToHost::ClientHello( + ClientHello::new( + SupportedProtocolVersions::try_new(versions).expect("supported versions"), + required_capabilities, + CapabilitySet::empty(), + ) + .expect("client hello"), + ) +} + +fn session_id(value: &str) -> SessionId { + SessionId::new(value).expect("session ID") +} + +fn request_id(value: i64) -> RequestId { + RequestId::new(value) +} + +#[tokio::test] +async fn handshake_and_multiple_session_lifecycles_are_ordered() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("read hello"), + Some(HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + CapabilitySet::empty(), + ))) + ); + + for (request_id, id) in [(request_id(1), "session-1"), (request_id(2), "session-2")] { + writer + .write(&ClientToHost::Request { + id: request_id, + request: HostRequest::OpenSession { + session_id: session_id(id), + }, + }) + .await + .expect("open session"); + assert_eq!( + reader.read::().await.expect("session ready"), + Some(HostToClient::Response { + id: request_id, + result: WireResult::Ok { + value: HostResponse::SessionReady { + session_id: session_id(id), + }, + }, + }) + ); + } + + for (request_id, id) in [(request_id(3), "session-1"), (request_id(4), "session-2")] { + writer + .write(&ClientToHost::Request { + id: request_id, + request: HostRequest::ShutdownSession { + session_id: session_id(id), + }, + }) + .await + .expect("shutdown session"); + assert_eq!( + reader.read::().await.expect("session closed"), + Some(HostToClient::Response { + id: request_id, + result: WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session_id(id), + }, + }, + }) + ); + } + + drop(writer); + drop(reader); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn incompatible_or_invalid_handshake_is_rejected() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + let version_two = ProtocolVersion::new(/*value*/ 2).expect("protocol version"); + + writer + .write(&client_hello([version_two], CapabilitySet::empty())) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::NoCompatibleVersion { + supported_versions: SupportedProtocolVersions::try_new([ProtocolVersion::V1]) + .expect("host versions"), + }, + }) + ); + host.await.expect("host task").expect("host connection"); + + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + writer + .write(&ClientToHost::Request { + id: request_id(1), + request: HostRequest::OpenSession { + session_id: session_id("session-1"), + }, + }) + .await + .expect("write invalid first message"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::InvalidHello { + message: "first message must be connection/hello".to_string(), + }, + }) + ); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn unsupported_required_capability_is_rejected() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + let capability = Capability::new("required").expect("capability"); + + writer + .write(&client_hello( + [ProtocolVersion::V1], + CapabilitySet::try_new([capability.clone()]).expect("capabilities"), + )) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::MissingRequiredCapability { capability }, + }) + ); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn session_id_cannot_be_reused_after_shutdown() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 2048); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write hello"); + reader + .read::() + .await + .expect("read hello") + .expect("host hello"); + + let id = session_id("session-1"); + for (request_id, request) in [ + ( + request_id(1), + HostRequest::OpenSession { + session_id: id.clone(), + }, + ), + ( + request_id(2), + HostRequest::ShutdownSession { + session_id: id.clone(), + }, + ), + ] { + writer + .write(&ClientToHost::Request { + id: request_id, + request, + }) + .await + .expect("session request"); + reader + .read::() + .await + .expect("session response") + .expect("session response message"); + } + writer + .write(&ClientToHost::Request { + id: request_id(3), + request: HostRequest::OpenSession { session_id: id }, + }) + .await + .expect("reuse session ID"); + assert_eq!( + reader.read::().await.expect("reuse response"), + Some(HostToClient::Response { + id: request_id(3), + result: WireResult::Err { + message: "code-mode session ID `session-1` was reused".to_string(), + }, + }) + ); + drop(writer); + drop(reader); + host.await.expect("host task").expect("host connection"); +} diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs new file mode 100644 index 0000000000..93f1299605 --- /dev/null +++ b/codex-rs/code-mode-host/src/lib.rs @@ -0,0 +1,345 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use anyhow::Context; +use anyhow::Result; +use codex_code_mode::InProcessCodeModeSession; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HandshakeRejectReason; +use codex_code_mode_protocol::host::HostHello; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use codex_code_mode_protocol::host::WireResult; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::sync::mpsc; +use tokio_util::task::TaskTracker; + +use self::delegate::RemoteDelegate; +use self::peer::HostPeer; + +mod delegate; +mod peer; + +/// Runs one code-mode host connection over the process standard streams. +pub async fn run_stdio() -> Result<()> { + run(tokio::io::stdin(), tokio::io::stdout()).await +} + +/// Runs one code-mode host connection over an ordered input/output pair. +async fn run(reader: R, writer: W) -> Result<()> +where + R: AsyncRead + Send + Unpin + 'static, + W: AsyncWrite + Send + Unpin + 'static, +{ + let mut reader = FramedReader::new(reader); + let mut writer = FramedWriter::new(writer); + if !negotiate(&mut reader, &mut writer).await? { + return Ok(()); + } + + let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel(); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = Arc::new(HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(HashSet::new()), + request_tasks: TaskTracker::new(), + closing: AtomicBool::new(false), + peer: Arc::clone(&peer), + }); + let writer_task = tokio::spawn(async move { + while let Some(message) = outgoing_rx.recv().await { + writer + .write(&message) + .await + .context("failed to write code-mode host message")?; + } + Ok::<(), anyhow::Error>(()) + }); + + let input_result = async { + while let Some(message) = reader + .read::() + .await + .context("failed to read code-mode client message")? + { + match message { + ClientToHost::ClientHello(_) => { + anyhow::bail!("received a second code-mode client hello"); + } + ClientToHost::Request { id, request } => { + state.spawn_request(id, request); + } + ClientToHost::DelegateResponse { id, result } => { + peer.complete(id, result.into_result()).await; + } + } + } + Ok::<(), anyhow::Error>(()) + } + .await; + + peer.disconnect(); + state.disconnect().await; + drop(state); + drop(peer); + let writer_result = writer_task.await.context("code-mode writer task failed")?; + input_result?; + writer_result +} + +async fn negotiate(reader: &mut FramedReader, writer: &mut FramedWriter) -> Result +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + let Some(first_message) = reader + .read::() + .await + .context("failed to read code-mode client hello")? + else { + return Ok(false); + }; + let ClientToHost::ClientHello(client_hello) = first_message else { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::InvalidHello { + message: "first message must be connection/hello".to_string(), + }, + }) + .await + .context("failed to reject invalid code-mode client hello")?; + return Ok(false); + }; + + let supported_versions = SupportedProtocolVersions::try_new([ProtocolVersion::V1])?; + if !client_hello + .supported_versions() + .contains(ProtocolVersion::V1) + { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::NoCompatibleVersion { supported_versions }, + }) + .await + .context("failed to reject incompatible code-mode client")?; + return Ok(false); + } + + let host_capabilities = CapabilitySet::empty(); + if let Some(capability) = client_hello + .required_capabilities() + .iter() + .find(|capability| !host_capabilities.contains(capability)) + { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::MissingRequiredCapability { + capability: capability.clone(), + }, + }) + .await + .context("failed to reject unsupported code-mode capability")?; + return Ok(false); + } + + writer + .write(&HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + host_capabilities, + ))) + .await + .context("failed to write code-mode host hello")?; + Ok(true) +} + +struct HostState { + sessions: Mutex>>, + seen_session_ids: Mutex>, + request_tasks: TaskTracker, + closing: AtomicBool, + peer: Arc, +} + +impl HostState { + fn spawn_request(self: &Arc, request_id: RequestId, request: HostRequest) { + let state = Arc::clone(self); + self.request_tasks.spawn(async move { + state.handle_request(request_id, request).await; + }); + } + + async fn handle_request(&self, request_id: RequestId, request: HostRequest) { + if self.closing.load(Ordering::Acquire) { + self.respond( + request_id, + Err("code-mode host is shutting down".to_string()), + ); + return; + } + match request { + HostRequest::OpenSession { session_id } => { + let result = self + .open_session(session_id.clone()) + .map(|()| HostResponse::SessionReady { session_id }); + self.respond(request_id, result); + } + HostRequest::Execute { + session_id, + request, + } => { + let request = match request.try_into() { + Ok(request) => request, + Err(err) => { + self.respond( + request_id, + Err(format!("invalid code-mode execute request: {err}")), + ); + return; + } + }; + let result = match self.session(&session_id) { + Ok(session) => session.execute(request).await, + Err(err) => { + self.respond(request_id, Err(err)); + return; + } + }; + match result { + Ok(started) => { + let cell_id = started.cell_id.clone(); + self.respond( + request_id, + Ok(HostResponse::ExecutionStarted { + cell_id: cell_id.into(), + }), + ); + self.peer.send(HostToClient::InitialResponse { + id: request_id, + result: WireResult::from_result( + started.initial_response().await.map(Into::into), + ), + }); + } + Err(err) => self.respond(request_id, Err(err)), + } + } + HostRequest::Wait { + session_id, + request, + } => { + let result = match self.session(&session_id) { + Ok(session) => session.wait(request.into()).await.map(|outcome| { + HostResponse::WaitCompleted { + outcome: outcome.into(), + } + }), + Err(err) => Err(err), + }; + self.respond(request_id, result); + } + HostRequest::Terminate { + session_id, + cell_id, + } => { + let result = match self.session(&session_id) { + Ok(session) => session.terminate(cell_id.into()).await.map(|outcome| { + HostResponse::WaitCompleted { + outcome: outcome.into(), + } + }), + Err(err) => Err(err), + }; + self.respond(request_id, result); + } + HostRequest::ShutdownSession { session_id } => { + let session = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(&session_id); + let result = match session { + Some(session) => session + .shutdown() + .await + .map(|()| HostResponse::SessionClosed { session_id }), + None => Err(format!("unknown code-mode session {session_id}")), + }; + self.respond(request_id, result); + } + } + } + + fn open_session(&self, session_id: SessionId) -> Result<(), String> { + let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner); + if self.closing.load(Ordering::Acquire) { + return Err("code-mode host is shutting down".to_string()); + } + if !self + .seen_session_ids + .lock() + .unwrap_or_else(PoisonError::into_inner) + .insert(session_id.clone()) + { + return Err(format!("code-mode session ID `{session_id}` was reused")); + } + let delegate = Arc::new(RemoteDelegate::new( + session_id.clone(), + Arc::clone(&self.peer), + )); + sessions.insert( + session_id, + Arc::new(InProcessCodeModeSession::with_delegate(delegate)), + ); + Ok(()) + } + + fn session(&self, session_id: &SessionId) -> Result, String> { + self.sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(session_id) + .cloned() + .ok_or_else(|| format!("unknown code-mode session {session_id}")) + } + + fn respond(&self, id: RequestId, result: Result) { + self.peer.send(HostToClient::Response { + id, + result: WireResult::from_result(result), + }); + } + + async fn disconnect(&self) { + self.closing.store(true, Ordering::Release); + let sessions = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .drain() + .map(|(_, session)| session) + .collect::>(); + for session in sessions { + let _ = session.shutdown().await; + } + self.request_tasks.close(); + self.request_tasks.wait().await; + } +} + +#[cfg(test)] +#[path = "host_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index f328e4d9d0..b215869c8a 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -1 +1,4 @@ -fn main() {} +#[tokio::main(flavor = "current_thread")] +async fn main() -> anyhow::Result<()> { + codex_code_mode_host::run_stdio().await +} diff --git a/codex-rs/code-mode-host/src/peer.rs b/codex-rs/code-mode-host/src/peer.rs new file mode 100644 index 0000000000..582e01a518 --- /dev/null +++ b/codex-rs/code-mode-host/src/peer.rs @@ -0,0 +1,124 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::SessionId; +use tokio::sync::Mutex; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +pub(super) struct HostPeer { + outgoing_tx: mpsc::UnboundedSender, + pending: Mutex>>>, + next_request_id: AtomicI64, + disconnected: CancellationToken, +} + +impl HostPeer { + pub(super) fn new(outgoing_tx: mpsc::UnboundedSender) -> Self { + Self { + outgoing_tx, + pending: Mutex::new(HashMap::new()), + next_request_id: AtomicI64::new(1), + disconnected: CancellationToken::new(), + } + } + + pub(super) fn send(&self, message: HostToClient) -> bool { + self.outgoing_tx.send(message).is_ok() + } + + pub(super) async fn call( + self: &Arc, + session_id: SessionId, + request: DelegateRequest, + cancellation_token: CancellationToken, + ) -> Result { + if self.disconnected.is_cancelled() { + return Err("code-mode client connection closed".to_string()); + } + let id = DelegateRequestId::new(self.next_request_id.fetch_add(1, Ordering::Relaxed)); + let (response_tx, response_rx) = oneshot::channel(); + self.pending.lock().await.insert(id, response_tx); + let mut pending = PendingDelegateRequest::new(Arc::clone(self), id); + if !self.send(HostToClient::DelegateRequest { + id, + session_id, + request, + }) { + self.pending.lock().await.remove(&id); + pending.disarm(); + return Err("code-mode client connection closed".to_string()); + } + + tokio::select! { + response = response_rx => { + pending.disarm(); + response.map_err(|_| { + "code-mode client closed before returning delegate output".to_string() + })? + } + _ = cancellation_token.cancelled() => { + if self.pending.lock().await.remove(&id).is_some() { + self.send(HostToClient::CancelDelegateRequest { id }); + } + pending.disarm(); + Err("code mode delegate request cancelled".to_string()) + } + _ = self.disconnected.cancelled() => { + self.pending.lock().await.remove(&id); + pending.disarm(); + Err("code-mode client connection closed".to_string()) + } + } + } + + pub(super) async fn complete( + &self, + id: DelegateRequestId, + response: Result, + ) { + if let Some(sender) = self.pending.lock().await.remove(&id) { + let _ = sender.send(response); + } + } + + pub(super) fn disconnect(&self) { + self.disconnected.cancel(); + } +} + +struct PendingDelegateRequest { + peer: Arc, + id: Option, +} + +impl PendingDelegateRequest { + fn new(peer: Arc, id: DelegateRequestId) -> Self { + Self { peer, id: Some(id) } + } + + fn disarm(&mut self) { + self.id = None; + } +} + +impl Drop for PendingDelegateRequest { + fn drop(&mut self) { + let Some(id) = self.id.take() else { + return; + }; + let peer = Arc::clone(&self.peer); + tokio::spawn(async move { + if peer.pending.lock().await.remove(&id).is_some() { + peer.send(HostToClient::CancelDelegateRequest { id }); + } + }); + } +} diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs new file mode 100644 index 0000000000..4c2c5d1718 --- /dev/null +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -0,0 +1,192 @@ +#![allow(clippy::expect_used)] + +use std::sync::Arc; +use std::sync::Mutex; + +use codex_code_mode::CellId; +use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeSession; +use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::CodeModeToolKind; +use codex_code_mode::ExecuteRequest; +use codex_code_mode::FunctionCallOutputContentItem; +use codex_code_mode::NotificationFuture; +use codex_code_mode::ProcessOwnedCodeModeSessionProvider; +use codex_code_mode::RuntimeResponse; +use codex_code_mode::ToolDefinition; +use codex_code_mode::ToolInvocationFuture; +use codex_code_mode::WaitOutcome; +use codex_code_mode::WaitRequest; +use codex_protocol::ToolName; +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio_util::sync::CancellationToken; + +#[derive(Default)] +struct RecordingDelegate { + invocations: Mutex>, + notifications: Mutex>, + closed_cells: Mutex>, +} + +impl CodeModeSessionDelegate for RecordingDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + _cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + self.invocations + .lock() + .expect("invocations lock") + .push(invocation); + Box::pin(async { Ok(json!({ "value": "output" })) }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + self.notifications + .lock() + .expect("notifications lock") + .push((call_id, cell_id, text)); + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.closed_cells + .lock() + .expect("closed cells lock") + .push(cell_id.clone()); + } +} + +fn cell_id(value: &str) -> CellId { + CellId::new(value.to_string()) +} + +fn execute_request(source: &str) -> ExecuteRequest { + ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: None, + max_output_tokens: None, + } +} + +async fn execute(session: &Arc, request: ExecuteRequest) -> RuntimeResponse { + session + .execute(request) + .await + .expect("start execution") + .initial_response() + .await + .expect("initial response") +} + +#[tokio::test] +async fn remote_session_persists_values_forwards_delegates_and_controls_cells() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let delegate = Arc::new(RecordingDelegate::default()); + let session = provider + .create_session(delegate.clone()) + .await + .expect("create remote session"); + + assert_eq!( + execute(&session, execute_request(r#"store("key", "persisted");"#),).await, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None, + } + ); + + let mut callback_request = execute_request( + r#" +const result = await tools.echo({ value: String(load("key")) }); +notify("notice"); +text(result.value); +"#, + ); + callback_request.tool_call_id = "call-2".to_string(); + callback_request.enabled_tools = vec![ToolDefinition { + name: "echo".to_string(), + tool_name: ToolName::plain("echo"), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }]; + assert_eq!( + execute(&session, callback_request).await, + RuntimeResponse::Result { + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "output".to_string(), + }], + error_text: None, + } + ); + assert_eq!( + *delegate.invocations.lock().expect("invocations lock"), + vec![CodeModeNestedToolCall { + cell_id: cell_id("2"), + runtime_tool_call_id: "tool-1".to_string(), + tool_name: ToolName::plain("echo"), + tool_kind: CodeModeToolKind::Function, + input: Some(json!({ "value": "persisted" })), + }] + ); + assert_eq!( + *delegate.notifications.lock().expect("notifications lock"), + vec![("call-2".to_string(), cell_id("2"), "notice".to_string())] + ); + + let mut pending_request = execute_request("await new Promise(() => {});"); + pending_request.tool_call_id = "call-3".to_string(); + pending_request.yield_time_ms = Some(1); + assert_eq!( + execute(&session, pending_request).await, + RuntimeResponse::Yielded { + cell_id: cell_id("3"), + content_items: Vec::new(), + } + ); + assert_eq!( + session + .wait(WaitRequest { + cell_id: cell_id("3"), + yield_time_ms: 1, + }) + .await + .expect("wait for cell"), + WaitOutcome::LiveCell(RuntimeResponse::Yielded { + cell_id: cell_id("3"), + content_items: Vec::new(), + }) + ); + assert_eq!( + session + .terminate(cell_id("3")) + .await + .expect("terminate cell"), + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("3"), + content_items: Vec::new(), + }) + ); + + session.shutdown().await.expect("shutdown remote session"); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![cell_id("1"), cell_id("2"), cell_id("3")] + ); +} diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 09f1d7380d..513628a507 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -20,7 +20,7 @@ codex-code-mode-protocol = { workspace = true } codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } serde_json = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tokio = { workspace = true, features = ["io-util", "macros", "process", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["rt"] } tracing = { workspace = true } v8 = { workspace = true } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 089bb6f78d..b186bd00e6 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -1,9 +1,12 @@ mod cell_actor; +mod remote_session; mod runtime; mod service; mod session_runtime; 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; diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs new file mode 100644 index 0000000000..f5eb443809 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session.rs @@ -0,0 +1,297 @@ +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +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 self::connection::Connection; +use crate::NoopCodeModeSessionDelegate; + +mod connection; + +const CODE_MODE_HOST_PATH_ENV: &str = "CODEX_CODE_MODE_HOST_PATH"; + +/// Creates code-mode sessions backed by one lazily spawned process host. +pub struct ProcessOwnedCodeModeSessionProvider { + host_program: PathBuf, + process_host: StdMutex>>, +} + +impl ProcessOwnedCodeModeSessionProvider { + pub fn with_host_program(host_program: PathBuf) -> Self { + Self { + host_program, + process_host: StdMutex::new(None), + } + } + + fn process_host(&self) -> Arc { + 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, + ) -> CodeModeSessionProviderFuture<'a> { + let session = ProcessOwnedCodeModeSession::with_process_host(delegate, self.process_host()); + Box::pin(async move { + session.connection().await?; + let session: Arc = Arc::new(session); + Ok(session) + }) + } +} + +struct OwnedProcessHost { + host_program: PathBuf, + connection: StdMutex>>, + 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, 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> { + 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, + Open(Arc), + Shutdown, +} + +/// A logical code-mode session assigned to a process-owned host. +pub struct ProcessOwnedCodeModeSession { + process_host: Arc, + session_id: SessionId, + delegate: Arc, + state: StdMutex, + transition_permit: Semaphore, +} + +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, + process_host: Arc, + ) -> Self { + let session_id = process_host.allocate_session_id(); + Self { + process_host, + session_id, + delegate, + state: StdMutex::new(SessionState::New), + transition_permit: Semaphore::new(/*permits*/ 1), + } + } + + async fn connection(&self) -> Result, String> { + if let Some(connection) = self.current_connection()? { + return Ok(connection); + } + + let _transition_permit = self + .transition_permit + .acquire() + .await + .map_err(|_| "code-mode session transition coordinator closed".to_string())?; + if let Some(connection) = self.current_connection()? { + return Ok(connection); + } + let connection = self.process_host.connection().await?; + connection + .open_session(self.session_id.clone(), Arc::clone(&self.delegate)) + .await?; + *self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + SessionState::Open(Arc::clone(&connection)); + Ok(connection) + } + + fn current_connection(&self) -> Result>, String> { + match &*self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + { + SessionState::New => Ok(None), + SessionState::Open(connection) => Ok(Some(Arc::clone(connection))), + SessionState::Shutdown => Err("code mode session is shutting down".to_string()), + } + } + + pub async fn execute(&self, request: ExecuteRequest) -> Result { + self.connection() + .await? + .execute(self.session_id.clone(), request) + .await + } + + pub async fn wait(&self, request: WaitRequest) -> Result { + self.connection() + .await? + .wait(self.session_id.clone(), request) + .await + } + + pub async fn terminate(&self, cell_id: CellId) -> Result { + self.connection() + .await? + .terminate(self.session_id.clone(), cell_id) + .await + } + + pub async fn shutdown(&self) -> Result<(), String> { + let transition_permit = self + .transition_permit + .acquire() + .await + .map_err(|_| "code-mode session transition coordinator closed".to_string())?; + let connection = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match std::mem::replace(&mut *state, SessionState::Shutdown) { + SessionState::Open(connection) => connection, + SessionState::New | SessionState::Shutdown => return Ok(()), + } + }; + drop(transition_permit); + connection.shutdown_session(self.session_id.clone()).await + } +} + +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; diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs new file mode 100644 index 0000000000..413f25fc21 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -0,0 +1,405 @@ +use std::collections::HashMap; +use std::path::Path; +use std::process::Stdio; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +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::CapabilitySet; +use codex_code_mode_protocol::host::ClientHello; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use 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 self::reader::drive_reader; + +mod reader; + +const IPC_CHANNEL_CAPACITY: usize = 128; + +pub(super) struct Connection { + state: Arc, + cancellation: CancellationToken, +} + +impl Connection { + pub(super) async fn spawn(host_program: &Path) -> Result { + 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() + ) + })?; + 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 mut reader = FramedReader::new(stdout); + let mut writer = FramedWriter::new(stdin); + + 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::() + .await + .map_err(|err| format!("failed to read code-mode host hello: {err}"))? + { + Some(HostToClient::HostHello(hello)) + if hello.selected_version() == ProtocolVersion::V1 => {} + Some(HostToClient::HandshakeRejected { reason }) => { + return Err(format!("code-mode host rejected the handshake: {reason:?}")); + } + Some(message) => { + return Err(format!( + "code-mode host returned an invalid handshake response: {message:?}" + )); + } + None => return Err("code-mode host exited during handshake".to_string()), + } + + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY); + let cancellation = CancellationToken::new(); + let state = Arc::new(ConnectionState::new(outgoing_tx, cancellation.clone())); + + let writer_state = Arc::downgrade(&state); + let writer_cancellation = cancellation.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + _ = writer_cancellation.cancelled() => break, + message = outgoing_rx.recv() => { + let Some(message) = message else { + break; + }; + if let Err(err) = writer.write(&message).await { + if let Some(state) = writer_state.upgrade() { + state.fail(format!( + "failed to write code-mode host message: {err}" + )).await; + } + break; + } + } + } + } + }); + + let reader_state = Arc::downgrade(&state); + let reader_cancellation = cancellation.clone(); + tokio::spawn(async move { + drive_reader(reader, reader_state, reader_cancellation).await; + }); + + 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(); + tokio::spawn(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; + } + } + }); + + Ok(Self { + state, + cancellation, + }) + } + + pub(super) fn is_alive(&self) -> bool { + self.state.alive.load(Ordering::Acquire) + } + + pub(super) async fn open_session( + &self, + session_id: SessionId, + delegate: Arc, + ) -> Result<(), String> { + self.state + .delegates + .lock() + .await + .insert(session_id.clone(), delegate); + let response = self + .request(HostRequest::OpenSession { + session_id: session_id.clone(), + }) + .await; + match response { + Ok(HostResponse::SessionReady { + session_id: ready_session_id, + }) if ready_session_id == session_id => Ok(()), + Ok(_) => { + self.state.delegates.lock().await.remove(&session_id); + Err("code-mode host returned an invalid open-session response".to_string()) + } + Err(err) => { + self.state.delegates.lock().await.remove(&session_id); + Err(err) + } + } + } + + pub(super) async fn execute( + &self, + session_id: SessionId, + request: ExecuteRequest, + ) -> Result { + let request = request + .try_into() + .map_err(|err| format!("failed to encode code-mode execute request: {err}"))?; + let id = RequestId::new(self.state.next_request_id.fetch_add(1, Ordering::Relaxed)); + let (response_tx, response_rx) = oneshot::channel(); + let (initial_tx, initial_rx) = oneshot::channel(); + self.state.pending.lock().await.insert(id, response_tx); + self.state + .initial_responses + .lock() + .await + .insert(id, initial_tx); + if let Err(err) = self + .state + .send(ClientToHost::Request { + id, + request: HostRequest::Execute { + session_id, + request, + }, + }) + .await + { + self.state.pending.lock().await.remove(&id); + self.state.initial_responses.lock().await.remove(&id); + return Err(err); + } + let response = match response_rx.await { + Ok(Ok(response)) => response, + Ok(Err(err)) => { + self.state.initial_responses.lock().await.remove(&id); + return Err(err); + } + Err(_) => { + self.state.initial_responses.lock().await.remove(&id); + return Err(self.state.failure_message()); + } + }; + match response { + HostResponse::ExecutionStarted { cell_id } => Ok(StartedCell::from_result_receiver( + cell_id.into(), + initial_rx, + )), + _ => { + self.state.initial_responses.lock().await.remove(&id); + Err("code-mode host returned an invalid execute response".to_string()) + } + } + } + + pub(super) async fn wait( + &self, + session_id: SessionId, + request: WaitRequest, + ) -> Result { + match self + .request(HostRequest::Wait { + session_id, + request: request.into(), + }) + .await? + { + HostResponse::WaitCompleted { outcome } => Ok(outcome.into()), + _ => Err("code-mode host returned an invalid wait response".to_string()), + } + } + + pub(super) async fn terminate( + &self, + session_id: SessionId, + cell_id: CellId, + ) -> Result { + match self + .request(HostRequest::Terminate { + session_id, + cell_id: cell_id.into(), + }) + .await? + { + HostResponse::WaitCompleted { outcome } => Ok(outcome.into()), + _ => Err("code-mode host returned an invalid terminate response".to_string()), + } + } + + pub(super) async fn shutdown_session(&self, session_id: SessionId) -> Result<(), String> { + let response = self + .request(HostRequest::ShutdownSession { + session_id: session_id.clone(), + }) + .await; + self.state.delegates.lock().await.remove(&session_id); + match response? { + HostResponse::SessionClosed { + session_id: closed_session_id, + } if closed_session_id == session_id => Ok(()), + _ => Err("code-mode host returned an invalid shutdown response".to_string()), + } + } + + async fn request(&self, request: HostRequest) -> Result { + let id = RequestId::new(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); + if let Err(err) = self.state.send(ClientToHost::Request { id, request }).await { + self.state.pending.lock().await.remove(&id); + return Err(err); + } + response_rx + .await + .map_err(|_| self.state.failure_message())? + } +} + +impl Drop for Connection { + fn drop(&mut self) { + self.cancellation.cancel(); + } +} + +struct ConnectionState { + outgoing_tx: mpsc::Sender, + pending: Mutex>>>, + initial_responses: Mutex>>>, + delegates: Mutex>>, + delegate_cancellations: Mutex>, + next_request_id: AtomicI64, + alive: AtomicBool, + failure: std::sync::Mutex>, + cancellation: CancellationToken, +} + +impl ConnectionState { + fn new(outgoing_tx: mpsc::Sender, cancellation: CancellationToken) -> Self { + Self { + outgoing_tx, + pending: Mutex::new(HashMap::new()), + initial_responses: Mutex::new(HashMap::new()), + delegates: Mutex::new(HashMap::new()), + delegate_cancellations: Mutex::new(HashMap::new()), + next_request_id: AtomicI64::new(1), + alive: AtomicBool::new(true), + failure: std::sync::Mutex::new(None), + cancellation, + } + } + + async fn send(&self, message: ClientToHost) -> 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() + .ok() + .and_then(|failure| failure.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; + } + if let Ok(mut failure) = self.failure.lock() { + *failure = Some(reason.clone()); + } + self.cancellation.cancel(); + for (_, sender) in self.pending.lock().await.drain() { + let _ = sender.send(Err(reason.clone())); + } + for (_, sender) in self.initial_responses.lock().await.drain() { + let _ = sender.send(Err(reason.clone())); + } + self.delegates.lock().await.clear(); + for (_, cancellation) in self.delegate_cancellations.lock().await.drain() { + cancellation.cancel(); + } + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/reader.rs b/codex-rs/code-mode/src/remote_session/connection/reader.rs new file mode 100644 index 0000000000..277a293024 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/reader.rs @@ -0,0 +1,131 @@ +use std::sync::Arc; +use std::sync::Weak; + +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::WireResult; +use tokio::process::ChildStdout; +use tokio_util::sync::CancellationToken; + +use super::ConnectionState; + +pub(super) async fn drive_reader( + mut reader: FramedReader, + state: Weak, + cancellation: CancellationToken, +) { + loop { + let message = tokio::select! { + _ = cancellation.cancelled() => return, + result = reader.read::() => 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: HostToClient) { + match message { + HostToClient::Response { id, result } => { + if let Some(sender) = state.pending.lock().await.remove(&id) { + let _ = sender.send(result.into_result()); + } + } + HostToClient::InitialResponse { id, result } => { + if let Some(sender) = state.initial_responses.lock().await.remove(&id) { + let _ = sender.send(result.into_result().map(Into::into)); + } + } + HostToClient::DelegateRequest { + id, + session_id, + request, + } => { + let delegate = state.delegates.lock().await.get(&session_id).cloned(); + let Some(delegate) = delegate else { + let _ = state + .send(ClientToHost::DelegateResponse { + id, + result: WireResult::Err { + message: format!("unknown code-mode session {session_id}"), + }, + }) + .await; + return; + }; + let cancellation = CancellationToken::new(); + state + .delegate_cancellations + .lock() + .await + .insert(id, cancellation.clone()); + tokio::spawn(async move { + let result = match request { + DelegateRequest::InvokeTool { invocation } => delegate + .invoke_tool(invocation.into(), cancellation) + .await + .map(|result| DelegateResponse::ToolResult { result }), + DelegateRequest::Notify { + call_id, + cell_id, + text, + } => delegate + .notify(call_id, cell_id.into(), text, cancellation) + .await + .map(|()| DelegateResponse::NotificationDelivered), + }; + state.delegate_cancellations.lock().await.remove(&id); + let _ = state + .send(ClientToHost::DelegateResponse { + id, + result: WireResult::from_result(result), + }) + .await; + }); + } + HostToClient::CancelDelegateRequest { id } => { + if let Some(cancellation) = state.delegate_cancellations.lock().await.remove(&id) { + cancellation.cancel(); + } + } + HostToClient::CellClosed { + session_id, + cell_id, + } => { + if let Some(delegate) = state.delegates.lock().await.get(&session_id).cloned() { + let cell_id = cell_id.into(); + delegate.cell_closed(&cell_id); + } + } + HostToClient::HostHello(_) | HostToClient::HandshakeRejected { .. } => { + state + .fail("code-mode host sent a second handshake response".to_string()) + .await; + } + } +} diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs new file mode 100644 index 0000000000..d971207553 --- /dev/null +++ b/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"); +} diff --git a/justfile b/justfile index 03f87391ac..d4a0e6566c 100644 --- a/justfile +++ b/justfile @@ -31,6 +31,10 @@ tui-with-exec-server *args: file-search *args: cargo run --bin codex-file-search -- {args} +# Run the standalone code-mode host from source. +code-mode-host *args: + cargo run --bin codex-code-mode-host -- {args} + # Build the CLI and run the app-server test client app-server-test-client *args: cargo build -p codex-cli @@ -107,6 +111,16 @@ bazel-codex *args: bazel-codex *args: bazel run //codex-rs/cli:codex --run_under='cd /d "{{ invocation_directory_native() }}" &&' -- @($args | Select-Object -Skip 1) +# Build and run the standalone code-mode host from source using Bazel. +[no-cd] +[unix] +bazel-code-mode-host *args: + bazel run //codex-rs/code-mode-host:codex-code-mode-host --run_under="cd $PWD &&" -- "$@" + +[windows] +bazel-code-mode-host *args: + bazel run //codex-rs/code-mode-host:codex-code-mode-host --run_under='cd /d "{{ invocation_directory_native() }}" &&' -- @($args | Select-Object -Skip 1) + [no-cd] bazel-lock-update: bazel mod deps --lockfile_mode=update