From 84a6cbe9441a879264ad333b6aad68daad1aed1f Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 18 Mar 2026 14:45:23 -0700 Subject: [PATCH] Add exec-server process RPC implementation Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 202 ++++++++- .../exec-server/src/client/local_backend.rs | 57 +++ codex-rs/exec-server/src/client_api.rs | 10 + codex-rs/exec-server/src/lib.rs | 12 + codex-rs/exec-server/src/protocol.rs | 144 +++++++ codex-rs/exec-server/src/server/handler.rs | 407 +++++++++++++++++- .../exec-server/src/server/handler/tests.rs | 74 ++++ codex-rs/exec-server/src/server/registry.rs | 34 +- codex-rs/exec-server/tests/stdio_smoke.rs | 206 ++++++--- 9 files changed, 1069 insertions(+), 77 deletions(-) create mode 100644 codex-rs/exec-server/src/server/handler/tests.rs diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 52c7da824d..652c9d01d9 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -1,24 +1,46 @@ use std::sync::Arc; use std::time::Duration; +use codex_app_server_protocol::JSONRPCNotification; +use serde_json::Value; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; +use tokio::sync::broadcast; use tokio::sync::mpsc; use tokio::time::timeout; use tokio_tungstenite::connect_async; +use tracing::debug; use tracing::warn; use crate::client_api::ExecServerClientConnectOptions; +use crate::client_api::ExecServerEvent; use crate::client_api::RemoteExecServerConnectArgs; use crate::connection::JsonRpcConnection; +use crate::protocol::EXEC_EXITED_METHOD; +use crate::protocol::EXEC_METHOD; +use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; +use crate::protocol::EXEC_READ_METHOD; +use crate::protocol::EXEC_TERMINATE_METHOD; +use crate::protocol::EXEC_WRITE_METHOD; +use crate::protocol::ExecExitedNotification; +use crate::protocol::ExecOutputDeltaNotification; +use crate::protocol::ExecParams; +use crate::protocol::ExecResponse; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; use crate::protocol::InitializeResponse; +use crate::protocol::ReadParams; +use crate::protocol::ReadResponse; +use crate::protocol::TerminateParams; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteParams; +use crate::protocol::WriteResponse; use crate::rpc::RpcCallError; use crate::rpc::RpcClient; use crate::rpc::RpcClientEvent; use crate::rpc::RpcNotificationSender; +use crate::rpc::RpcServerOutboundMessage; mod local_backend; use local_backend::LocalBackend; @@ -78,6 +100,7 @@ impl ClientBackend { struct Inner { backend: ClientBackend, + events_tx: broadcast::Sender, reader_task: tokio::task::JoinHandle<()>, } @@ -128,14 +151,32 @@ impl ExecServerClient { pub async fn connect_in_process( options: ExecServerClientConnectOptions, ) -> Result { - let (outgoing_tx, _outgoing_rx) = mpsc::channel(1); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(256); let backend = LocalBackend::new(crate::server::ExecServerHandler::new( RpcNotificationSender::new(outgoing_tx), )); - let inner = Arc::new(Inner { - backend: ClientBackend::InProcess(backend), - reader_task: tokio::spawn(async {}), + let inner = Arc::new_cyclic(|weak| { + let weak = weak.clone(); + let reader_task = tokio::spawn(async move { + while let Some(message) = outgoing_rx.recv().await { + if let Some(inner) = weak.upgrade() + && let Err(err) = handle_in_process_outbound_message(&inner, message).await + { + warn!( + "in-process exec-server client closing after unexpected response: {err}" + ); + return; + } + } + }); + + Inner { + backend: ClientBackend::InProcess(backend), + events_tx: broadcast::channel(256).0, + reader_task, + } }); + let client = Self { inner }; client.initialize(options).await?; Ok(client) @@ -183,6 +224,10 @@ impl ExecServerClient { .await } + pub fn event_receiver(&self) -> broadcast::Receiver { + self.inner.events_tx.subscribe() + } + pub async fn initialize( &self, options: ExecServerClientConnectOptions, @@ -213,36 +258,111 @@ impl ExecServerClient { })? } + pub async fn exec(&self, params: ExecParams) -> Result { + if let Some(backend) = self.inner.backend.as_local() { + return backend.exec(params).await; + } + let Some(remote) = self.inner.backend.as_remote() else { + return Err(ExecServerError::Protocol( + "remote backend missing during exec".to_string(), + )); + }; + remote.call(EXEC_METHOD, ¶ms).await.map_err(Into::into) + } + + pub async fn read(&self, params: ReadParams) -> Result { + if let Some(backend) = self.inner.backend.as_local() { + return backend.exec_read(params).await; + } + let Some(remote) = self.inner.backend.as_remote() else { + return Err(ExecServerError::Protocol( + "remote backend missing during read".to_string(), + )); + }; + remote + .call(EXEC_READ_METHOD, ¶ms) + .await + .map_err(Into::into) + } + + pub async fn write( + &self, + process_id: &str, + chunk: Vec, + ) -> Result { + let params = WriteParams { + process_id: process_id.to_string(), + chunk: chunk.into(), + }; + if let Some(backend) = self.inner.backend.as_local() { + return backend.exec_write(params).await; + } + let Some(remote) = self.inner.backend.as_remote() else { + return Err(ExecServerError::Protocol( + "remote backend missing during write".to_string(), + )); + }; + remote + .call(EXEC_WRITE_METHOD, ¶ms) + .await + .map_err(Into::into) + } + + pub async fn terminate(&self, process_id: &str) -> Result { + let params = TerminateParams { + process_id: process_id.to_string(), + }; + if let Some(backend) = self.inner.backend.as_local() { + return backend.terminate(params).await; + } + let Some(remote) = self.inner.backend.as_remote() else { + return Err(ExecServerError::Protocol( + "remote backend missing during terminate".to_string(), + )); + }; + remote + .call(EXEC_TERMINATE_METHOD, ¶ms) + .await + .map_err(Into::into) + } + async fn connect( connection: JsonRpcConnection, options: ExecServerClientConnectOptions, ) -> Result { let (rpc_client, mut events_rx) = RpcClient::new(connection); - let reader_task = tokio::spawn(async move { - while let Some(event) = events_rx.recv().await { - match event { - RpcClientEvent::Notification(notification) => { - warn!( - "ignoring unexpected exec-server notification during stub phase: {}", - notification.method - ); - } - RpcClientEvent::Disconnected { reason } => { - if let Some(reason) = reason { - warn!("exec-server client transport disconnected: {reason}"); + let inner = Arc::new_cyclic(|weak| { + let weak = weak.clone(); + let reader_task = tokio::spawn(async move { + while let Some(event) = events_rx.recv().await { + match event { + RpcClientEvent::Notification(notification) => { + if let Some(inner) = weak.upgrade() + && let Err(err) = + handle_server_notification(&inner, notification).await + { + warn!("exec-server client closing after protocol error: {err}"); + return; + } + } + RpcClientEvent::Disconnected { reason } => { + if let Some(reason) = reason { + warn!("exec-server client transport disconnected: {reason}"); + } + return; } - return; } } + }); + + Inner { + backend: ClientBackend::Remote(rpc_client), + events_tx: broadcast::channel(256).0, + reader_task, } }); - let client = Self { - inner: Arc::new(Inner { - backend: ClientBackend::Remote(rpc_client), - reader_task, - }), - }; + let client = Self { inner }; client.initialize(options).await?; Ok(client) } @@ -270,3 +390,39 @@ impl From for ExecServerError { } } } + +async fn handle_in_process_outbound_message( + inner: &Arc, + message: RpcServerOutboundMessage, +) -> Result<(), ExecServerError> { + match message { + RpcServerOutboundMessage::Response { .. } | RpcServerOutboundMessage::Error { .. } => Err( + ExecServerError::Protocol("unexpected in-process RPC response".to_string()), + ), + RpcServerOutboundMessage::Notification(notification) => { + handle_server_notification(inner, notification).await + } + } +} + +async fn handle_server_notification( + inner: &Arc, + notification: JSONRPCNotification, +) -> Result<(), ExecServerError> { + match notification.method.as_str() { + EXEC_OUTPUT_DELTA_METHOD => { + let params: ExecOutputDeltaNotification = + serde_json::from_value(notification.params.unwrap_or(Value::Null))?; + let _ = inner.events_tx.send(ExecServerEvent::OutputDelta(params)); + } + EXEC_EXITED_METHOD => { + let params: ExecExitedNotification = + serde_json::from_value(notification.params.unwrap_or(Value::Null))?; + let _ = inner.events_tx.send(ExecServerEvent::Exited(params)); + } + other => { + debug!("ignoring unknown exec-server notification: {other}"); + } + } + Ok(()) +} diff --git a/codex-rs/exec-server/src/client/local_backend.rs b/codex-rs/exec-server/src/client/local_backend.rs index 8f9a2481f8..16b16d3b10 100644 --- a/codex-rs/exec-server/src/client/local_backend.rs +++ b/codex-rs/exec-server/src/client/local_backend.rs @@ -1,6 +1,14 @@ use std::sync::Arc; +use crate::protocol::ExecParams; +use crate::protocol::ExecResponse; use crate::protocol::InitializeResponse; +use crate::protocol::ReadParams; +use crate::protocol::ReadResponse; +use crate::protocol::TerminateParams; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteParams; +use crate::protocol::WriteResponse; use crate::server::ExecServerHandler; use super::ExecServerError; @@ -35,4 +43,53 @@ impl LocalBackend { .initialized() .map_err(ExecServerError::Protocol) } + + pub(super) async fn exec(&self, params: ExecParams) -> Result { + self.handler + .exec(params) + .await + .map_err(|error| ExecServerError::Server { + code: error.code, + message: error.message, + }) + } + + pub(super) async fn exec_read( + &self, + params: ReadParams, + ) -> Result { + self.handler + .exec_read(params) + .await + .map_err(|error| ExecServerError::Server { + code: error.code, + message: error.message, + }) + } + + pub(super) async fn exec_write( + &self, + params: WriteParams, + ) -> Result { + self.handler + .exec_write(params) + .await + .map_err(|error| ExecServerError::Server { + code: error.code, + message: error.message, + }) + } + + pub(super) async fn terminate( + &self, + params: TerminateParams, + ) -> Result { + self.handler + .terminate(params) + .await + .map_err(|error| ExecServerError::Server { + code: error.code, + message: error.message, + }) + } } diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index 6e89763416..962d3ba364 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -1,5 +1,8 @@ use std::time::Duration; +use crate::protocol::ExecExitedNotification; +use crate::protocol::ExecOutputDeltaNotification; + /// Connection options for any exec-server client transport. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ExecServerClientConnectOptions { @@ -15,3 +18,10 @@ pub struct RemoteExecServerConnectArgs { pub connect_timeout: Duration, pub initialize_timeout: Duration, } + +/// Connection-level server events. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecServerEvent { + OutputDelta(ExecOutputDeltaNotification), + Exited(ExecExitedNotification), +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index e204d6e084..13e910351e 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -9,12 +9,24 @@ mod server; pub use client::ExecServerClient; pub use client::ExecServerError; pub use client_api::ExecServerClientConnectOptions; +pub use client_api::ExecServerEvent; pub use client_api::RemoteExecServerConnectArgs; pub use local::ExecServerLaunchCommand; pub use local::SpawnedExecServer; pub use local::spawn_local_exec_server; +pub use protocol::ExecExitedNotification; +pub use protocol::ExecOutputDeltaNotification; +pub use protocol::ExecOutputStream; +pub use protocol::ExecParams; +pub use protocol::ExecResponse; pub use protocol::InitializeParams; pub use protocol::InitializeResponse; +pub use protocol::ReadParams; +pub use protocol::ReadResponse; +pub use protocol::TerminateParams; +pub use protocol::TerminateResponse; +pub use protocol::WriteParams; +pub use protocol::WriteResponse; pub use server::ExecServerTransport; pub use server::ExecServerTransportParseError; pub use server::run_main; diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs index a365dc33ea..7ed8e20ae4 100644 --- a/codex-rs/exec-server/src/protocol.rs +++ b/codex-rs/exec-server/src/protocol.rs @@ -1,10 +1,36 @@ +use std::collections::HashMap; +use std::path::PathBuf; + +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; use serde::Deserialize; use serde::Serialize; pub const INITIALIZE_METHOD: &str = "initialize"; pub const INITIALIZED_METHOD: &str = "initialized"; +pub const EXEC_METHOD: &str = "process/start"; +pub const EXEC_READ_METHOD: &str = "process/read"; +pub const EXEC_WRITE_METHOD: &str = "process/write"; +pub const EXEC_TERMINATE_METHOD: &str = "process/terminate"; +pub const EXEC_OUTPUT_DELTA_METHOD: &str = "process/output"; +pub const EXEC_EXITED_METHOD: &str = "process/exited"; pub const PROTOCOL_VERSION: &str = "exec-server.v0"; +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct ByteChunk(#[serde(with = "base64_bytes")] pub Vec); + +impl ByteChunk { + pub fn into_inner(self) -> Vec { + self.0 + } +} + +impl From> for ByteChunk { + fn from(value: Vec) -> Self { + Self(value) + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct InitializeParams { @@ -16,3 +42,121 @@ pub struct InitializeParams { pub struct InitializeResponse { pub protocol_version: String, } + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExecParams { + /// Client-chosen logical process handle scoped to this connection/session. + /// This is a protocol key, not an OS pid. + pub process_id: String, + pub argv: Vec, + pub cwd: PathBuf, + pub env: HashMap, + pub tty: bool, + pub arg0: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExecResponse { + pub process_id: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ReadParams { + pub process_id: String, + pub after_seq: Option, + pub max_bytes: Option, + pub wait_ms: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProcessOutputChunk { + pub seq: u64, + pub stream: ExecOutputStream, + pub chunk: ByteChunk, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ReadResponse { + pub chunks: Vec, + pub next_seq: u64, + pub exited: bool, + pub exit_code: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WriteParams { + pub process_id: String, + pub chunk: ByteChunk, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct WriteResponse { + pub accepted: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TerminateParams { + pub process_id: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct TerminateResponse { + pub running: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum ExecOutputStream { + Stdout, + Stderr, + Pty, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExecOutputDeltaNotification { + pub process_id: String, + pub stream: ExecOutputStream, + pub chunk: ByteChunk, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExecExitedNotification { + pub process_id: String, + pub exit_code: i32, +} + +mod base64_bytes { + use super::BASE64_STANDARD; + use base64::Engine as _; + use serde::Deserialize; + use serde::Deserializer; + use serde::Serializer; + + pub fn serialize(bytes: &[u8], serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str(&BASE64_STANDARD.encode(bytes)) + } + + pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + where + D: Deserializer<'de>, + { + let encoded = String::deserialize(deserializer)?; + BASE64_STANDARD + .decode(encoded) + .map_err(serde::de::Error::custom) + } +} diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index f10fa8ddd9..40a13cf1f9 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -1,14 +1,67 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +use std::time::Duration; use codex_app_server_protocol::JSONRPCErrorError; +use codex_utils_pty::ExecCommandSession; +use codex_utils_pty::TerminalSize; +use tokio::sync::Mutex; +use tokio::sync::Notify; +use tracing::warn; +use crate::protocol::ExecExitedNotification; +use crate::protocol::ExecOutputDeltaNotification; +use crate::protocol::ExecOutputStream; +use crate::protocol::ExecParams; +use crate::protocol::ExecResponse; use crate::protocol::InitializeResponse; use crate::protocol::PROTOCOL_VERSION; +use crate::protocol::ProcessOutputChunk; +use crate::protocol::ReadParams; +use crate::protocol::ReadResponse; +use crate::protocol::TerminateParams; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteParams; +use crate::protocol::WriteResponse; use crate::rpc::RpcNotificationSender; +use crate::rpc::internal_error; +use crate::rpc::invalid_params; +use crate::rpc::invalid_request; + +const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024; +#[cfg(test)] +const EXITED_PROCESS_RETENTION: Duration = Duration::from_millis(25); +#[cfg(not(test))] +const EXITED_PROCESS_RETENTION: Duration = Duration::from_secs(30); + +#[derive(Clone)] +struct RetainedOutputChunk { + seq: u64, + stream: ExecOutputStream, + chunk: Vec, +} + +struct RunningProcess { + session: ExecCommandSession, + tty: bool, + output: VecDeque, + retained_bytes: usize, + next_seq: u64, + exit_code: Option, + output_notify: Arc, +} + +enum ProcessEntry { + Starting, + Running(Box), +} pub(crate) struct ExecServerHandler { - _notifications: RpcNotificationSender, + notifications: RpcNotificationSender, + processes: Arc>>, initialize_requested: AtomicBool, initialized: AtomicBool, } @@ -16,17 +69,32 @@ pub(crate) struct ExecServerHandler { impl ExecServerHandler { pub(crate) fn new(notifications: RpcNotificationSender) -> Self { Self { - _notifications: notifications, + notifications, + processes: Arc::new(Mutex::new(HashMap::new())), initialize_requested: AtomicBool::new(false), initialized: AtomicBool::new(false), } } - pub(crate) async fn shutdown(&self) {} + pub(crate) async fn shutdown(&self) { + let remaining = { + let mut processes = self.processes.lock().await; + processes + .drain() + .filter_map(|(_, process)| match process { + ProcessEntry::Starting => None, + ProcessEntry::Running(process) => Some(process), + }) + .collect::>() + }; + for process in remaining { + process.session.terminate(); + } + } pub(crate) fn initialize(&self) -> Result { if self.initialize_requested.swap(true, Ordering::SeqCst) { - return Err(crate::rpc::invalid_request( + return Err(invalid_request( "initialize may only be sent once per connection".to_string(), )); } @@ -42,4 +110,335 @@ impl ExecServerHandler { self.initialized.store(true, Ordering::SeqCst); Ok(()) } + + fn require_initialized(&self) -> Result<(), JSONRPCErrorError> { + if !self.initialize_requested.load(Ordering::SeqCst) { + return Err(invalid_request( + "client must call initialize before using exec methods".to_string(), + )); + } + if !self.initialized.load(Ordering::SeqCst) { + return Err(invalid_request( + "client must send initialized before using exec methods".to_string(), + )); + } + Ok(()) + } + + pub(crate) async fn exec(&self, params: ExecParams) -> Result { + self.require_initialized()?; + let process_id = params.process_id.clone(); + + let (program, args) = params + .argv + .split_first() + .ok_or_else(|| invalid_params("argv must not be empty".to_string()))?; + + { + let mut process_map = self.processes.lock().await; + if process_map.contains_key(&process_id) { + return Err(invalid_request(format!( + "process {process_id} already exists" + ))); + } + process_map.insert(process_id.clone(), ProcessEntry::Starting); + } + + let spawned_result = if params.tty { + codex_utils_pty::spawn_pty_process( + program, + args, + params.cwd.as_path(), + ¶ms.env, + ¶ms.arg0, + TerminalSize::default(), + ) + .await + } else { + codex_utils_pty::spawn_pipe_process_no_stdin( + program, + args, + params.cwd.as_path(), + ¶ms.env, + ¶ms.arg0, + ) + .await + }; + let spawned = match spawned_result { + Ok(spawned) => spawned, + Err(err) => { + let mut process_map = self.processes.lock().await; + if matches!(process_map.get(&process_id), Some(ProcessEntry::Starting)) { + process_map.remove(&process_id); + } + return Err(internal_error(err.to_string())); + } + }; + + let output_notify = Arc::new(Notify::new()); + { + let mut process_map = self.processes.lock().await; + process_map.insert( + process_id.clone(), + ProcessEntry::Running(Box::new(RunningProcess { + session: spawned.session, + tty: params.tty, + output: VecDeque::new(), + retained_bytes: 0, + next_seq: 1, + exit_code: None, + output_notify: Arc::clone(&output_notify), + })), + ); + } + + tokio::spawn(stream_output( + process_id.clone(), + if params.tty { + ExecOutputStream::Pty + } else { + ExecOutputStream::Stdout + }, + spawned.stdout_rx, + self.notifications.clone(), + Arc::clone(&self.processes), + Arc::clone(&output_notify), + )); + tokio::spawn(stream_output( + process_id.clone(), + if params.tty { + ExecOutputStream::Pty + } else { + ExecOutputStream::Stderr + }, + spawned.stderr_rx, + self.notifications.clone(), + Arc::clone(&self.processes), + Arc::clone(&output_notify), + )); + tokio::spawn(watch_exit( + process_id.clone(), + spawned.exit_rx, + self.notifications.clone(), + Arc::clone(&self.processes), + output_notify, + )); + + Ok(ExecResponse { process_id }) + } + + pub(crate) async fn exec_read( + &self, + params: ReadParams, + ) -> Result { + self.require_initialized()?; + let after_seq = params.after_seq.unwrap_or(0); + let max_bytes = params.max_bytes.unwrap_or(usize::MAX); + let wait = Duration::from_millis(params.wait_ms.unwrap_or(0)); + let deadline = tokio::time::Instant::now() + wait; + + loop { + let (response, output_notify) = { + let process_map = self.processes.lock().await; + let process = process_map.get(¶ms.process_id).ok_or_else(|| { + invalid_request(format!("unknown process id {}", params.process_id)) + })?; + let ProcessEntry::Running(process) = process else { + return Err(invalid_request(format!( + "process id {} is starting", + params.process_id + ))); + }; + + let mut chunks = Vec::new(); + let mut total_bytes = 0; + let mut next_seq = process.next_seq; + for retained in process.output.iter().filter(|chunk| chunk.seq > after_seq) { + let chunk_len = retained.chunk.len(); + if !chunks.is_empty() && total_bytes + chunk_len > max_bytes { + break; + } + total_bytes += chunk_len; + chunks.push(ProcessOutputChunk { + seq: retained.seq, + stream: retained.stream, + chunk: retained.chunk.clone().into(), + }); + next_seq = retained.seq + 1; + if total_bytes >= max_bytes { + break; + } + } + + ( + ReadResponse { + chunks, + next_seq, + exited: process.exit_code.is_some(), + exit_code: process.exit_code, + }, + Arc::clone(&process.output_notify), + ) + }; + + if !response.chunks.is_empty() + || response.exited + || tokio::time::Instant::now() >= deadline + { + return Ok(response); + } + + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Ok(response); + } + let _ = tokio::time::timeout(remaining, output_notify.notified()).await; + } + } + + pub(crate) async fn exec_write( + &self, + params: WriteParams, + ) -> Result { + self.require_initialized()?; + let writer_tx = { + let process_map = self.processes.lock().await; + let process = process_map.get(¶ms.process_id).ok_or_else(|| { + invalid_request(format!("unknown process id {}", params.process_id)) + })?; + let ProcessEntry::Running(process) = process else { + return Err(invalid_request(format!( + "process id {} is starting", + params.process_id + ))); + }; + if !process.tty { + return Err(invalid_request(format!( + "stdin is closed for process {}", + params.process_id + ))); + } + process.session.writer_sender() + }; + + writer_tx + .send(params.chunk.into_inner()) + .await + .map_err(|_| internal_error("failed to write to process stdin".to_string()))?; + + Ok(WriteResponse { accepted: true }) + } + + pub(crate) async fn terminate( + &self, + params: TerminateParams, + ) -> Result { + self.require_initialized()?; + let running = { + let process_map = self.processes.lock().await; + match process_map.get(¶ms.process_id) { + Some(ProcessEntry::Running(process)) => { + process.session.terminate(); + true + } + Some(ProcessEntry::Starting) | None => false, + } + }; + + Ok(TerminateResponse { running }) + } } + +async fn stream_output( + process_id: String, + stream: ExecOutputStream, + mut receiver: tokio::sync::mpsc::Receiver>, + notifications: RpcNotificationSender, + processes: Arc>>, + output_notify: Arc, +) { + while let Some(chunk) = receiver.recv().await { + let notification = { + let mut processes = processes.lock().await; + let Some(entry) = processes.get_mut(&process_id) else { + break; + }; + let ProcessEntry::Running(process) = entry else { + break; + }; + let seq = process.next_seq; + process.next_seq += 1; + process.retained_bytes += chunk.len(); + process.output.push_back(RetainedOutputChunk { + seq, + stream, + chunk: chunk.clone(), + }); + while process.retained_bytes > RETAINED_OUTPUT_BYTES_PER_PROCESS { + let Some(evicted) = process.output.pop_front() else { + break; + }; + process.retained_bytes = process.retained_bytes.saturating_sub(evicted.chunk.len()); + warn!( + "retained output cap exceeded for process {process_id}; dropping oldest output" + ); + } + ExecOutputDeltaNotification { + process_id: process_id.clone(), + stream, + chunk: chunk.into(), + } + }; + output_notify.notify_waiters(); + + if notifications + .notify(crate::protocol::EXEC_OUTPUT_DELTA_METHOD, ¬ification) + .await + .is_err() + { + break; + } + } +} + +async fn watch_exit( + process_id: String, + exit_rx: tokio::sync::oneshot::Receiver, + notifications: RpcNotificationSender, + processes: Arc>>, + output_notify: Arc, +) { + let exit_code = exit_rx.await.unwrap_or(-1); + { + let mut processes = processes.lock().await; + if let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) { + process.exit_code = Some(exit_code); + } + } + output_notify.notify_waiters(); + if notifications + .notify( + crate::protocol::EXEC_EXITED_METHOD, + &ExecExitedNotification { + process_id: process_id.clone(), + exit_code, + }, + ) + .await + .is_err() + { + return; + } + + tokio::time::sleep(EXITED_PROCESS_RETENTION).await; + let mut processes = processes.lock().await; + if matches!( + processes.get(&process_id), + Some(ProcessEntry::Running(process)) if process.exit_code == Some(exit_code) + ) { + processes.remove(&process_id); + } +} + +#[cfg(test)] +mod tests; diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs new file mode 100644 index 0000000000..9682eb1f66 --- /dev/null +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -0,0 +1,74 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +use pretty_assertions::assert_eq; +use tokio::sync::mpsc; + +use super::ExecServerHandler; +use crate::protocol::ExecParams; +use crate::protocol::InitializeResponse; +use crate::protocol::PROTOCOL_VERSION; +use crate::rpc::RpcNotificationSender; + +fn exec_params(process_id: &str) -> ExecParams { + let mut env = HashMap::new(); + if let Some(path) = std::env::var_os("PATH") { + env.insert("PATH".to_string(), path.to_string_lossy().into_owned()); + } + ExecParams { + process_id: process_id.to_string(), + argv: vec![ + "bash".to_string(), + "-lc".to_string(), + "sleep 0.1".to_string(), + ], + cwd: std::env::current_dir().expect("cwd"), + env, + tty: false, + arg0: None, + } +} + +async fn initialized_handler() -> Arc { + let (outgoing_tx, _outgoing_rx) = mpsc::channel(16); + let handler = Arc::new(ExecServerHandler::new(RpcNotificationSender::new( + outgoing_tx, + ))); + assert_eq!( + handler.initialize().expect("initialize"), + InitializeResponse { + protocol_version: PROTOCOL_VERSION.to_string(), + } + ); + handler.initialized().expect("initialized"); + handler +} + +#[tokio::test] +async fn duplicate_process_ids_allow_only_one_successful_start() { + let handler = initialized_handler().await; + let first_handler = Arc::clone(&handler); + let second_handler = Arc::clone(&handler); + + let (first, second) = tokio::join!( + first_handler.exec(exec_params("proc-1")), + second_handler.exec(exec_params("proc-1")), + ); + + let (successes, failures): (Vec<_>, Vec<_>) = + [first, second].into_iter().partition(Result::is_ok); + assert_eq!(successes.len(), 1); + assert_eq!(failures.len(), 1); + + let error = failures + .into_iter() + .next() + .expect("one failed request") + .expect_err("expected duplicate process error"); + assert_eq!(error.code, -32600); + assert_eq!(error.message, "process proc-1 already exists"); + + tokio::time::sleep(Duration::from_millis(150)).await; + handler.shutdown().await; +} diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 100fe4a81a..50efdddd71 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -1,8 +1,16 @@ use std::sync::Arc; +use crate::protocol::EXEC_METHOD; +use crate::protocol::EXEC_READ_METHOD; +use crate::protocol::EXEC_TERMINATE_METHOD; +use crate::protocol::EXEC_WRITE_METHOD; +use crate::protocol::ExecParams; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; +use crate::protocol::ReadParams; +use crate::protocol::TerminateParams; +use crate::protocol::WriteParams; use crate::rpc::RpcRouter; use crate::server::ExecServerHandler; @@ -10,11 +18,35 @@ pub(crate) fn build_router() -> RpcRouter { let mut router = RpcRouter::new(); router.request( INITIALIZE_METHOD, - |handler: Arc, _params: InitializeParams| async move { handler.initialize() }, + |handler: Arc, _params: InitializeParams| async move { + handler.initialize() + }, ); router.notification( INITIALIZED_METHOD, |handler: Arc, (): ()| async move { handler.initialized() }, ); + router.request( + EXEC_METHOD, + |handler: Arc, params: ExecParams| async move { handler.exec(params).await }, + ); + router.request( + EXEC_READ_METHOD, + |handler: Arc, params: ReadParams| async move { + handler.exec_read(params).await + }, + ); + router.request( + EXEC_WRITE_METHOD, + |handler: Arc, params: WriteParams| async move { + handler.exec_write(params).await + }, + ); + router.request( + EXEC_TERMINATE_METHOD, + |handler: Arc, params: TerminateParams| async move { + handler.terminate(params).await + }, + ); router } diff --git a/codex-rs/exec-server/tests/stdio_smoke.rs b/codex-rs/exec-server/tests/stdio_smoke.rs index 6a89318133..c08d7f3c9b 100644 --- a/codex-rs/exec-server/tests/stdio_smoke.rs +++ b/codex-rs/exec-server/tests/stdio_smoke.rs @@ -3,19 +3,29 @@ use std::process::Stdio; use std::time::Duration; +use anyhow::Context; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; +use codex_exec_server::ExecOutputStream; +use codex_exec_server::ExecParams; +use codex_exec_server::ExecServerClient; +use codex_exec_server::ExecServerClientConnectOptions; +use codex_exec_server::ExecServerEvent; +use codex_exec_server::ExecServerLaunchCommand; use codex_exec_server::InitializeParams; use codex_exec_server::InitializeResponse; +use codex_exec_server::RemoteExecServerConnectArgs; +use codex_exec_server::spawn_local_exec_server; use codex_utils_cargo_bin::cargo_bin; use pretty_assertions::assert_eq; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::process::Command; +use tokio::sync::broadcast; use tokio::time::timeout; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -66,62 +76,160 @@ async fn exec_server_accepts_initialize_over_stdio() -> anyhow::Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn exec_server_stubs_process_start_over_stdio() -> anyhow::Result<()> { - let binary = cargo_bin("codex-exec-server")?; - let mut child = Command::new(binary); - child.stdin(Stdio::piped()); - child.stdout(Stdio::piped()); - child.stderr(Stdio::inherit()); - let mut child = child.spawn()?; +async fn exec_server_client_streams_output_and_accepts_writes() -> anyhow::Result<()> { + let mut env = std::collections::HashMap::new(); + if let Some(path) = std::env::var_os("PATH") { + env.insert("PATH".to_string(), path.to_string_lossy().into_owned()); + } - let mut stdin = child.stdin.take().expect("stdin"); - let stdout = child.stdout.take().expect("stdout"); - let mut stdout = BufReader::new(stdout).lines(); - - let initialize = JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(1), - method: "initialize".to_string(), - params: Some(serde_json::to_value(InitializeParams { + let server = spawn_local_exec_server( + ExecServerLaunchCommand { + program: cargo_bin("codex-exec-server")?, + args: Vec::new(), + }, + ExecServerClientConnectOptions { client_name: "exec-server-test".to_string(), - })?), - trace: None, - }); - stdin - .write_all(format!("{}\n", serde_json::to_string(&initialize)?).as_bytes()) - .await?; - let _ = timeout(Duration::from_secs(5), stdout.next_line()).await??; + initialize_timeout: Duration::from_secs(5), + }, + ) + .await?; - let exec = JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(2), - method: "process/start".to_string(), - params: Some(serde_json::json!({ - "processId": "proc-1", - "argv": ["true"], - "cwd": std::env::current_dir()?, - "env": {}, - "tty": false, - "arg0": null - })), - trace: None, - }); - stdin - .write_all(format!("{}\n", serde_json::to_string(&exec)?).as_bytes()) + let client = server.client(); + let mut events = client.event_receiver(); + let response = client + .exec(ExecParams { + process_id: "proc-1".to_string(), + argv: vec![ + "bash".to_string(), + "-lc".to_string(), + "printf 'ready\\n'; while IFS= read -r line; do printf 'echo:%s\\n' \"$line\"; done" + .to_string(), + ], + cwd: std::env::current_dir()?, + env, + tty: true, + arg0: None, + }) .await?; + let process_id = response.process_id; - let response_line = timeout(Duration::from_secs(5), stdout.next_line()).await??; - let response_line = response_line.expect("exec response line"); - let response: JSONRPCMessage = serde_json::from_str(&response_line)?; - let JSONRPCMessage::Error(codex_app_server_protocol::JSONRPCError { id, error }) = response - else { - panic!("expected process/start stub error"); - }; - assert_eq!(id, RequestId::Integer(2)); - assert_eq!(error.code, -32601); - assert_eq!( - error.message, - "exec-server stub does not implement `process/start` yet" + let (stream, ready_output) = recv_until_contains(&mut events, &process_id, "ready").await?; + assert_eq!(stream, ExecOutputStream::Pty); + assert!( + ready_output.contains("ready"), + "expected initial ready output" ); + client.write(&process_id, b"hello\n".to_vec()).await?; + + let (stream, echoed_output) = + recv_until_contains(&mut events, &process_id, "echo:hello").await?; + assert_eq!(stream, ExecOutputStream::Pty); + assert!( + echoed_output.contains("echo:hello"), + "expected echoed output" + ); + + client.terminate(&process_id).await?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exec_server_client_connects_over_websocket() -> anyhow::Result<()> { + let mut env = std::collections::HashMap::new(); + if let Some(path) = std::env::var_os("PATH") { + env.insert("PATH".to_string(), path.to_string_lossy().into_owned()); + } + + let binary = cargo_bin("codex-exec-server")?; + let mut child = Command::new(binary); + child.args(["--listen", "ws://127.0.0.1:0"]); + child.stdin(Stdio::null()); + child.stdout(Stdio::null()); + child.stderr(Stdio::piped()); + let mut child = child.spawn()?; + let stderr = child.stderr.take().expect("stderr"); + let mut stderr_lines = BufReader::new(stderr).lines(); + let websocket_url = read_websocket_url(&mut stderr_lines).await?; + + let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: "exec-server-test".to_string(), + connect_timeout: Duration::from_secs(5), + initialize_timeout: Duration::from_secs(5), + }) + .await?; + + let mut events = client.event_receiver(); + let response = client + .exec(ExecParams { + process_id: "proc-1".to_string(), + argv: vec![ + "bash".to_string(), + "-lc".to_string(), + "printf 'ready\\n'; while IFS= read -r line; do printf 'echo:%s\\n' \"$line\"; done" + .to_string(), + ], + cwd: std::env::current_dir()?, + env, + tty: true, + arg0: None, + }) + .await?; + let process_id = response.process_id; + + let (stream, ready_output) = recv_until_contains(&mut events, &process_id, "ready").await?; + assert_eq!(stream, ExecOutputStream::Pty); + assert!( + ready_output.contains("ready"), + "expected initial ready output" + ); + + client.write(&process_id, b"hello\n".to_vec()).await?; + + let (stream, echoed_output) = + recv_until_contains(&mut events, &process_id, "echo:hello").await?; + assert_eq!(stream, ExecOutputStream::Pty); + assert!( + echoed_output.contains("echo:hello"), + "expected echoed output" + ); + + client.terminate(&process_id).await?; child.start_kill()?; Ok(()) } + +async fn read_websocket_url(lines: &mut tokio::io::Lines>) -> anyhow::Result +where + R: tokio::io::AsyncRead + Unpin, +{ + let line = timeout(Duration::from_secs(5), lines.next_line()).await??; + let line = line.context("missing websocket startup banner")?; + let websocket_url = line + .split_whitespace() + .find(|part| part.starts_with("ws://")) + .context("missing websocket URL in startup banner")?; + Ok(websocket_url.to_string()) +} + +async fn recv_until_contains( + events: &mut broadcast::Receiver, + process_id: &str, + needle: &str, +) -> anyhow::Result<(ExecOutputStream, String)> { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + let mut collected = String::new(); + loop { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + let event = timeout(remaining, events.recv()).await??; + if let ExecServerEvent::OutputDelta(output_event) = event + && output_event.process_id == process_id + { + collected.push_str(&String::from_utf8_lossy(&output_event.chunk.into_inner())); + if collected.contains(needle) { + return Ok((output_event.stream, collected)); + } + } + } +}