From 144c3593dbd3bc2ccdc6906f2afbc4cfc2f069b2 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 16 Mar 2026 23:10:50 +0000 Subject: [PATCH] Add codex-exec-server crate This adds the standalone exec-server stdio JSON-RPC crate and its smoke tests without wiring it into the CLI or unified-exec yet. Co-authored-by: Codex --- codex-rs/Cargo.lock | 17 + codex-rs/Cargo.toml | 1 + codex-rs/exec-server/Cargo.toml | 35 ++ .../exec-server/src/bin/codex-exec-server.rs | 7 + codex-rs/exec-server/src/client.rs | 497 ++++++++++++++++++ codex-rs/exec-server/src/lib.rs | 20 + codex-rs/exec-server/src/protocol.rs | 143 +++++ codex-rs/exec-server/src/server.rs | 420 +++++++++++++++ codex-rs/exec-server/tests/stdio_smoke.rs | 141 +++++ 9 files changed, 1281 insertions(+) create mode 100644 codex-rs/exec-server/Cargo.toml create mode 100644 codex-rs/exec-server/src/bin/codex-exec-server.rs create mode 100644 codex-rs/exec-server/src/client.rs create mode 100644 codex-rs/exec-server/src/lib.rs create mode 100644 codex-rs/exec-server/src/protocol.rs create mode 100644 codex-rs/exec-server/src/server.rs create mode 100644 codex-rs/exec-server/tests/stdio_smoke.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d6a13a3d4c..8aa93c9687 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2003,6 +2003,23 @@ dependencies = [ "wiremock", ] +[[package]] +name = "codex-exec-server" +version = "0.0.0" +dependencies = [ + "anyhow", + "base64 0.22.1", + "codex-app-server-protocol", + "codex-utils-cargo-bin", + "codex-utils-pty", + "pretty_assertions", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tracing", +] + [[package]] name = "codex-execpolicy" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 35ff64195e..7d4b8792b6 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -26,6 +26,7 @@ members = [ "hooks", "secrets", "exec", + "exec-server", "execpolicy", "execpolicy-legacy", "keyring-store", diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml new file mode 100644 index 0000000000..7b1ec13cef --- /dev/null +++ b/codex-rs/exec-server/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "codex-exec-server" +version.workspace = true +edition.workspace = true +license.workspace = true + +[[bin]] +name = "codex-exec-server" +path = "src/bin/codex-exec-server.rs" + +[lints] +workspace = true + +[dependencies] +base64 = { workspace = true } +codex-app-server-protocol = { workspace = true } +codex-utils-pty = { workspace = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true } +thiserror = { workspace = true } +tokio = { workspace = true, features = [ + "io-std", + "io-util", + "macros", + "process", + "rt-multi-thread", + "sync", + "time", +] } +tracing = { workspace = true } + +[dev-dependencies] +anyhow = { workspace = true } +codex-utils-cargo-bin = { workspace = true } +pretty_assertions = { workspace = true } diff --git a/codex-rs/exec-server/src/bin/codex-exec-server.rs b/codex-rs/exec-server/src/bin/codex-exec-server.rs new file mode 100644 index 0000000000..399167c1a9 --- /dev/null +++ b/codex-rs/exec-server/src/bin/codex-exec-server.rs @@ -0,0 +1,7 @@ +#[tokio::main] +async fn main() { + if let Err(err) = codex_exec_server::run_main().await { + eprintln!("{err}"); + std::process::exit(1); + } +} diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs new file mode 100644 index 0000000000..68a8de9b79 --- /dev/null +++ b/codex-rs/exec-server/src/client.rs @@ -0,0 +1,497 @@ +use std::collections::HashMap; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCErrorError; +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 serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::process::Child; +use tokio::process::Command; +use tokio::sync::Mutex; +use tokio::sync::broadcast; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; +use tracing::debug; +use tracing::warn; + +use crate::protocol::EXEC_EXITED_METHOD; +use crate::protocol::EXEC_METHOD; +use crate::protocol::EXEC_OUTPUT_DELTA_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::TerminateParams; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteParams; +use crate::protocol::WriteResponse; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecServerLaunchCommand { + pub program: PathBuf, + pub args: Vec, +} + +pub struct ExecServerProcess { + process_id: String, + output_rx: broadcast::Receiver>, + writer_tx: mpsc::Sender>, + status: Arc, + client: ExecServerClient, +} + +impl std::fmt::Debug for ExecServerProcess { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ExecServerProcess") + .field("process_id", &self.process_id) + .field("has_exited", &self.has_exited()) + .field("exit_code", &self.exit_code()) + .finish() + } +} + +impl ExecServerProcess { + pub fn writer_sender(&self) -> mpsc::Sender> { + self.writer_tx.clone() + } + + pub fn output_receiver(&self) -> broadcast::Receiver> { + self.output_rx.resubscribe() + } + + pub fn has_exited(&self) -> bool { + self.status.has_exited() + } + + pub fn exit_code(&self) -> Option { + self.status.exit_code() + } + + pub fn terminate(&self) { + self.status.mark_exited(None); + let client = self.client.clone(); + let process_id = self.process_id.clone(); + tokio::spawn(async move { + let _ = client.terminate_process(&process_id).await; + }); + } +} + +impl std::fmt::Debug for RemoteProcessStatus { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RemoteProcessStatus") + .field("exited", &self.has_exited()) + .field("exit_code", &self.exit_code()) + .finish() + } +} + +struct RemoteProcessStatus { + exited: AtomicBool, + exit_code: StdMutex>, +} + +impl RemoteProcessStatus { + fn new() -> Self { + Self { + exited: AtomicBool::new(false), + exit_code: StdMutex::new(None), + } + } + + fn has_exited(&self) -> bool { + self.exited.load(Ordering::SeqCst) + } + + fn exit_code(&self) -> Option { + self.exit_code.lock().ok().and_then(|guard| *guard) + } + + fn mark_exited(&self, exit_code: Option) { + self.exited.store(true, Ordering::SeqCst); + if let Ok(mut guard) = self.exit_code.lock() { + *guard = exit_code; + } + } +} + +struct RegisteredProcess { + output_tx: broadcast::Sender>, + status: Arc, +} + +struct Inner { + child: StdMutex>, + write_tx: mpsc::UnboundedSender, + pending: Mutex>>>, + processes: Mutex>, + next_request_id: AtomicI64, + reader_task: JoinHandle<()>, + writer_task: JoinHandle<()>, +} + +impl Drop for Inner { + fn drop(&mut self) { + self.reader_task.abort(); + self.writer_task.abort(); + if let Ok(mut child_guard) = self.child.lock() + && let Some(child) = child_guard.as_mut() + { + let _ = child.start_kill(); + } + } +} + +#[derive(Clone)] +pub struct ExecServerClient { + inner: Arc, +} + +#[derive(Debug, thiserror::Error)] +pub enum ExecServerError { + #[error("failed to spawn exec-server: {0}")] + Spawn(#[source] std::io::Error), + #[error("exec-server transport closed")] + Closed, + #[error("failed to serialize or deserialize exec-server JSON: {0}")] + Json(#[from] serde_json::Error), + #[error("exec-server protocol error: {0}")] + Protocol(String), + #[error("exec-server rejected request ({code}): {message}")] + Server { code: i64, message: String }, +} + +impl ExecServerClient { + pub async fn spawn(command: ExecServerLaunchCommand) -> Result { + let mut child = Command::new(&command.program); + child.args(&command.args); + child.stdin(Stdio::piped()); + child.stdout(Stdio::piped()); + child.stderr(Stdio::inherit()); + child.kill_on_drop(true); + + let mut child = child.spawn().map_err(ExecServerError::Spawn)?; + let stdin = child.stdin.take().ok_or_else(|| { + ExecServerError::Protocol("exec-server stdin was not captured".to_string()) + })?; + let stdout = child.stdout.take().ok_or_else(|| { + ExecServerError::Protocol("exec-server stdout was not captured".to_string()) + })?; + + let (write_tx, mut write_rx) = mpsc::unbounded_channel::(); + let writer_task = tokio::spawn(async move { + let mut stdin = stdin; + while let Some(message) = write_rx.recv().await { + let encoded = match serde_json::to_vec(&message) { + Ok(encoded) => encoded, + Err(err) => { + warn!("failed to encode exec-server message: {err}"); + break; + } + }; + if stdin.write_all(&encoded).await.is_err() { + break; + } + if stdin.write_all(b"\n").await.is_err() { + break; + } + if stdin.flush().await.is_err() { + break; + } + } + }); + + let pending = Mutex::new(HashMap::< + RequestId, + oneshot::Sender>, + >::new()); + let processes = Mutex::new(HashMap::::new()); + let inner = Arc::new_cyclic(move |weak| { + let weak = weak.clone(); + let reader_task = tokio::spawn(async move { + let mut lines = BufReader::new(stdout).lines(); + loop { + let Some(inner) = weak.upgrade() else { + break; + }; + let next_line = lines.next_line().await; + match next_line { + Ok(Some(line)) => { + if line.trim().is_empty() { + continue; + } + match serde_json::from_str::(&line) { + Ok(message) => { + if let Err(err) = handle_server_message(&inner, message).await { + warn!("failed to handle exec-server message: {err}"); + break; + } + } + Err(err) => { + warn!("failed to parse exec-server message: {err}"); + break; + } + } + } + Ok(None) => break, + Err(err) => { + warn!("failed to read exec-server stdout: {err}"); + break; + } + } + } + if let Some(inner) = weak.upgrade() { + handle_transport_shutdown(&inner).await; + } + }); + + Inner { + child: StdMutex::new(Some(child)), + write_tx, + pending, + processes, + next_request_id: AtomicI64::new(1), + reader_task, + writer_task, + } + }); + + let client = Self { inner }; + client.initialize().await?; + Ok(client) + } + + pub async fn start_process( + &self, + params: ExecParams, + ) -> Result { + let process_id = params.process_id.clone(); + let status = Arc::new(RemoteProcessStatus::new()); + let (output_tx, output_rx) = broadcast::channel(256); + self.inner.processes.lock().await.insert( + process_id.clone(), + RegisteredProcess { + output_tx, + status: Arc::clone(&status), + }, + ); + + let (writer_tx, mut writer_rx) = mpsc::channel::>(128); + let client = self.clone(); + let write_process_id = process_id.clone(); + tokio::spawn(async move { + while let Some(chunk) = writer_rx.recv().await { + let request = WriteParams { + process_id: write_process_id.clone(), + chunk: chunk.into(), + }; + if client.write_process(request).await.is_err() { + break; + } + } + }); + + let response = match self.request::<_, ExecResponse>(EXEC_METHOD, ¶ms).await { + Ok(response) => response, + Err(err) => { + self.inner.processes.lock().await.remove(&process_id); + return Err(err); + } + }; + + if let Some(exit_code) = response.exit_code { + status.mark_exited(Some(exit_code)); + } + + Ok(ExecServerProcess { + process_id, + output_rx, + writer_tx, + status, + client: self.clone(), + }) + } + + async fn initialize(&self) -> Result<(), ExecServerError> { + let _: InitializeResponse = self + .request( + INITIALIZE_METHOD, + &InitializeParams { + client_name: "codex-core".to_string(), + }, + ) + .await?; + self.notify(INITIALIZED_METHOD, &serde_json::json!({})) + .await + } + + async fn write_process(&self, params: WriteParams) -> Result { + self.request(EXEC_WRITE_METHOD, ¶ms).await + } + + async fn terminate_process( + &self, + process_id: &str, + ) -> Result { + self.request( + EXEC_TERMINATE_METHOD, + &TerminateParams { + process_id: process_id.to_string(), + }, + ) + .await + } + + async fn notify(&self, method: &str, params: &P) -> Result<(), ExecServerError> { + let params = serde_json::to_value(params)?; + self.inner + .write_tx + .send(JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(params), + })) + .map_err(|_| ExecServerError::Closed) + } + + async fn request(&self, method: &str, params: &P) -> Result + where + P: Serialize, + R: DeserializeOwned, + { + let request_id = + RequestId::Integer(self.inner.next_request_id.fetch_add(1, Ordering::SeqCst)); + let (response_tx, response_rx) = oneshot::channel(); + self.inner + .pending + .lock() + .await + .insert(request_id.clone(), response_tx); + + let params = serde_json::to_value(params)?; + let message = JSONRPCMessage::Request(JSONRPCRequest { + id: request_id.clone(), + method: method.to_string(), + params: Some(params), + trace: None, + }); + + if self.inner.write_tx.send(message).is_err() { + self.inner.pending.lock().await.remove(&request_id); + return Err(ExecServerError::Closed); + } + + let result = response_rx.await.map_err(|_| ExecServerError::Closed)?; + match result { + Ok(value) => serde_json::from_value(value).map_err(ExecServerError::from), + Err(error) => Err(ExecServerError::Server { + code: error.code, + message: error.message, + }), + } + } +} + +async fn handle_server_message( + inner: &Arc, + message: JSONRPCMessage, +) -> Result<(), ExecServerError> { + match message { + JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { + if let Some(tx) = inner.pending.lock().await.remove(&id) { + let _ = tx.send(Ok(result)); + } + } + JSONRPCMessage::Error(JSONRPCError { id, error }) => { + if let Some(tx) = inner.pending.lock().await.remove(&id) { + let _ = tx.send(Err(error)); + } + } + JSONRPCMessage::Notification(notification) => { + handle_server_notification(inner, notification).await?; + } + JSONRPCMessage::Request(request) => { + return Err(ExecServerError::Protocol(format!( + "unexpected exec-server request from child: {}", + request.method + ))); + } + } + + Ok(()) +} + +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 chunk = params.chunk.into_inner(); + let processes = inner.processes.lock().await; + if let Some(process) = processes.get(¶ms.process_id) { + let _ = process.output_tx.send(chunk); + } + } + EXEC_EXITED_METHOD => { + let params: ExecExitedNotification = + serde_json::from_value(notification.params.unwrap_or(Value::Null))?; + let mut processes = inner.processes.lock().await; + if let Some(process) = processes.remove(¶ms.process_id) { + process.status.mark_exited(Some(params.exit_code)); + } + } + other => { + debug!("ignoring unknown exec-server notification: {other}"); + } + } + Ok(()) +} + +async fn handle_transport_shutdown(inner: &Arc) { + let pending = { + let mut pending = inner.pending.lock().await; + pending.drain().map(|(_, tx)| tx).collect::>() + }; + for tx in pending { + let _ = tx.send(Err(JSONRPCErrorError { + code: -32000, + data: None, + message: "exec-server transport closed".to_string(), + })); + } + + let processes = { + let mut processes = inner.processes.lock().await; + processes + .drain() + .map(|(_, process)| process) + .collect::>() + }; + for process in processes { + process.status.mark_exited(None); + } +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs new file mode 100644 index 0000000000..5e915ce5ca --- /dev/null +++ b/codex-rs/exec-server/src/lib.rs @@ -0,0 +1,20 @@ +mod client; +mod protocol; +mod server; + +pub use client::ExecServerClient; +pub use client::ExecServerError; +pub use client::ExecServerLaunchCommand; +pub use client::ExecServerProcess; +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::TerminateParams; +pub use protocol::TerminateResponse; +pub use protocol::WriteParams; +pub use protocol::WriteResponse; +pub use server::run_main; diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs new file mode 100644 index 0000000000..dcec9a6d6a --- /dev/null +++ b/codex-rs/exec-server/src/protocol.rs @@ -0,0 +1,143 @@ +use std::collections::HashMap; +use std::path::PathBuf; + +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use codex_utils_pty::DEFAULT_OUTPUT_BYTES_CAP; +use serde::Deserialize; +use serde::Serialize; + +pub const INITIALIZE_METHOD: &str = "initialize"; +pub const INITIALIZED_METHOD: &str = "initialized"; +pub const EXEC_METHOD: &str = "command/exec"; +pub const EXEC_WRITE_METHOD: &str = "command/exec/write"; +pub const EXEC_TERMINATE_METHOD: &str = "command/exec/terminate"; +pub const EXEC_OUTPUT_DELTA_METHOD: &str = "command/exec/outputDelta"; +pub const EXEC_EXITED_METHOD: &str = "command/exec/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 { + pub client_name: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct InitializeResponse { + pub protocol_version: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExecParams { + pub process_id: String, + pub argv: Vec, + pub cwd: PathBuf, + pub env: HashMap, + pub tty: bool, + #[serde(default = "default_output_bytes_cap")] + pub output_bytes_cap: usize, + pub arg0: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExecResponse { + pub process_id: String, + pub running: bool, + pub exit_code: Option, + pub stdout: Option, + pub stderr: 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, +} + +#[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, +} + +fn default_output_bytes_cap() -> usize { + DEFAULT_OUTPUT_BYTES_CAP +} + +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.rs b/codex-rs/exec-server/src/server.rs new file mode 100644 index 0000000000..56d2206b0c --- /dev/null +++ b/codex-rs/exec-server/src/server.rs @@ -0,0 +1,420 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; + +use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCErrorError; +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_utils_pty::ExecCommandSession; +use codex_utils_pty::TerminalSize; +use serde::Serialize; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::io::BufWriter; +use tokio::sync::Mutex; + +use crate::protocol::EXEC_EXITED_METHOD; +use crate::protocol::EXEC_METHOD; +use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; +use crate::protocol::EXEC_TERMINATE_METHOD; +use crate::protocol::EXEC_WRITE_METHOD; +use crate::protocol::ExecExitedNotification; +use crate::protocol::ExecOutputDeltaNotification; +use crate::protocol::ExecOutputStream; +use crate::protocol::ExecParams; +use crate::protocol::ExecResponse; +use crate::protocol::INITIALIZE_METHOD; +use crate::protocol::INITIALIZED_METHOD; +use crate::protocol::InitializeResponse; +use crate::protocol::PROTOCOL_VERSION; +use crate::protocol::TerminateParams; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteParams; +use crate::protocol::WriteResponse; + +struct RunningProcess { + session: ExecCommandSession, + tty: bool, + stdout_buffer: Arc>, + stderr_buffer: Arc>, +} + +#[derive(Debug)] +struct BoundedBytesBuffer { + max_bytes: usize, + bytes: VecDeque, +} + +impl BoundedBytesBuffer { + fn new(max_bytes: usize) -> Self { + Self { + max_bytes, + bytes: VecDeque::with_capacity(max_bytes.min(8192)), + } + } + + fn push_chunk(&mut self, chunk: &[u8]) { + if self.max_bytes == 0 { + return; + } + for byte in chunk { + self.bytes.push_back(*byte); + if self.bytes.len() > self.max_bytes { + self.bytes.pop_front(); + } + } + } + + fn snapshot(&self) -> Vec { + self.bytes.iter().copied().collect() + } +} + +pub async fn run_main() -> Result<(), Box> { + let writer = Arc::new(Mutex::new(BufWriter::new(tokio::io::stdout()))); + let processes = Arc::new(Mutex::new(HashMap::::new())); + let mut lines = BufReader::new(tokio::io::stdin()).lines(); + + while let Some(line) = lines.next_line().await? { + if line.trim().is_empty() { + continue; + } + + let message = serde_json::from_str::(&line)?; + if let JSONRPCMessage::Request(request) = message { + handle_request(request, &writer, &processes).await; + continue; + } + + if let JSONRPCMessage::Notification(notification) = message { + if notification.method != INITIALIZED_METHOD { + send_error( + &writer, + RequestId::Integer(-1), + invalid_request(format!( + "unexpected notification method: {}", + notification.method + )), + ) + .await; + } + continue; + } + } + + let remaining = { + let mut processes = processes.lock().await; + processes + .drain() + .map(|(_, process)| process) + .collect::>() + }; + for process in remaining { + process.session.terminate(); + } + + Ok(()) +} + +async fn handle_request( + request: JSONRPCRequest, + writer: &Arc>>, + processes: &Arc>>, +) { + let response = match request.method.as_str() { + INITIALIZE_METHOD => serde_json::to_value(InitializeResponse { + protocol_version: PROTOCOL_VERSION.to_string(), + }) + .map_err(|err| internal_error(err.to_string())), + EXEC_METHOD => handle_exec_request(request.params, writer, processes).await, + EXEC_WRITE_METHOD => handle_write_request(request.params, processes).await, + EXEC_TERMINATE_METHOD => handle_terminate_request(request.params, processes).await, + other => Err(invalid_request(format!("unknown method: {other}"))), + }; + + match response { + Ok(result) => { + send_response( + writer, + JSONRPCResponse { + id: request.id, + result, + }, + ) + .await; + } + Err(err) => { + send_error(writer, request.id, err).await; + } + } +} + +async fn handle_exec_request( + params: Option, + writer: &Arc>>, + processes: &Arc>>, +) -> Result { + let params: ExecParams = serde_json::from_value(params.unwrap_or(serde_json::Value::Null)) + .map_err(|err| invalid_params(err.to_string()))?; + + let (program, args) = params + .argv + .split_first() + .ok_or_else(|| invalid_params("argv must not be empty".to_string()))?; + + let spawned = 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 + } + .map_err(|err| internal_error(err.to_string()))?; + + let stdout_buffer = Arc::new(StdMutex::new(BoundedBytesBuffer::new( + params.output_bytes_cap, + ))); + let stderr_buffer = Arc::new(StdMutex::new(BoundedBytesBuffer::new( + params.output_bytes_cap, + ))); + + let process_id = params.process_id.clone(); + { + let mut process_map = processes.lock().await; + if process_map.contains_key(&process_id) { + spawned.session.terminate(); + return Err(invalid_request(format!( + "process {} already exists", + params.process_id + ))); + } + process_map.insert( + process_id.clone(), + RunningProcess { + session: spawned.session, + tty: params.tty, + stdout_buffer: Arc::clone(&stdout_buffer), + stderr_buffer: Arc::clone(&stderr_buffer), + }, + ); + } + + tokio::spawn(stream_output( + process_id.clone(), + ExecOutputStream::Stdout, + spawned.stdout_rx, + Arc::clone(writer), + Arc::clone(&stdout_buffer), + )); + tokio::spawn(stream_output( + process_id.clone(), + ExecOutputStream::Stderr, + spawned.stderr_rx, + Arc::clone(writer), + Arc::clone(&stderr_buffer), + )); + tokio::spawn(watch_exit( + process_id.clone(), + spawned.exit_rx, + Arc::clone(writer), + Arc::clone(processes), + )); + + serde_json::to_value(ExecResponse { + process_id, + running: true, + exit_code: None, + stdout: None, + stderr: None, + }) + .map_err(|err| internal_error(err.to_string())) +} + +async fn handle_write_request( + params: Option, + processes: &Arc>>, +) -> Result { + let params: WriteParams = serde_json::from_value(params.unwrap_or(serde_json::Value::Null)) + .map_err(|err| invalid_params(err.to_string()))?; + + let writer_tx = { + let process_map = processes.lock().await; + let process = process_map + .get(¶ms.process_id) + .ok_or_else(|| invalid_request(format!("unknown process id {}", 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()))?; + + serde_json::to_value(WriteResponse { accepted: true }) + .map_err(|err| internal_error(err.to_string())) +} + +async fn handle_terminate_request( + params: Option, + processes: &Arc>>, +) -> Result { + let params: TerminateParams = serde_json::from_value(params.unwrap_or(serde_json::Value::Null)) + .map_err(|err| invalid_params(err.to_string()))?; + + let process = { + let mut process_map = processes.lock().await; + process_map.remove(¶ms.process_id) + }; + + if let Some(process) = process { + process.session.terminate(); + serde_json::to_value(TerminateResponse { running: true }) + .map_err(|err| internal_error(err.to_string())) + } else { + serde_json::to_value(TerminateResponse { running: false }) + .map_err(|err| internal_error(err.to_string())) + } +} + +async fn stream_output( + process_id: String, + stream: ExecOutputStream, + mut receiver: tokio::sync::mpsc::Receiver>, + writer: Arc>>, + buffer: Arc>, +) { + while let Some(chunk) = receiver.recv().await { + if let Ok(mut guard) = buffer.lock() { + guard.push_chunk(&chunk); + } + let notification = ExecOutputDeltaNotification { + process_id: process_id.clone(), + stream, + chunk: chunk.into(), + }; + if send_notification(&writer, EXEC_OUTPUT_DELTA_METHOD, ¬ification) + .await + .is_err() + { + break; + } + } +} + +async fn watch_exit( + process_id: String, + exit_rx: tokio::sync::oneshot::Receiver, + writer: Arc>>, + processes: Arc>>, +) { + let exit_code = exit_rx.await.unwrap_or(-1); + let removed = { + let mut processes = processes.lock().await; + processes.remove(&process_id) + }; + if let Some(process) = removed { + let _ = process.stdout_buffer.lock().map(|buffer| buffer.snapshot()); + let _ = process.stderr_buffer.lock().map(|buffer| buffer.snapshot()); + } + let _ = send_notification( + &writer, + EXEC_EXITED_METHOD, + &ExecExitedNotification { + process_id, + exit_code, + }, + ) + .await; +} + +async fn send_response( + writer: &Arc>>, + response: JSONRPCResponse, +) { + let _ = send_message(writer, JSONRPCMessage::Response(response)).await; +} + +async fn send_error( + writer: &Arc>>, + id: RequestId, + error: JSONRPCErrorError, +) { + let _ = send_message(writer, JSONRPCMessage::Error(JSONRPCError { error, id })).await; +} + +async fn send_notification( + writer: &Arc>>, + method: &str, + params: &T, +) -> Result<(), serde_json::Error> { + send_message( + writer, + JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(serde_json::to_value(params)?), + }), + ) + .await + .map_err(serde_json::Error::io) +} + +async fn send_message( + writer: &Arc>>, + message: JSONRPCMessage, +) -> std::io::Result<()> { + let encoded = + serde_json::to_vec(&message).map_err(|err| std::io::Error::other(err.to_string()))?; + let mut writer = writer.lock().await; + writer.write_all(&encoded).await?; + writer.write_all(b"\n").await?; + writer.flush().await +} + +fn invalid_request(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32600, + data: None, + message, + } +} + +fn invalid_params(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32602, + data: None, + message, + } +} + +fn internal_error(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32603, + data: None, + message, + } +} diff --git a/codex-rs/exec-server/tests/stdio_smoke.rs b/codex-rs/exec-server/tests/stdio_smoke.rs new file mode 100644 index 0000000000..79f264327c --- /dev/null +++ b/codex-rs/exec-server/tests/stdio_smoke.rs @@ -0,0 +1,141 @@ +#![cfg(unix)] + +use std::process::Stdio; +use std::time::Duration; + +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::ExecParams; +use codex_exec_server::ExecServerClient; +use codex_exec_server::ExecServerLaunchCommand; +use codex_exec_server::InitializeParams; +use codex_exec_server::InitializeResponse; +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)] +async fn exec_server_accepts_initialize_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()?; + + 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 { + client_name: "exec-server-test".to_string(), + })?), + trace: None, + }); + stdin + .write_all(format!("{}\n", serde_json::to_string(&initialize)?).as_bytes()) + .await?; + + let response_line = timeout(Duration::from_secs(5), stdout.next_line()).await??; + let response_line = response_line.expect("response line"); + let response: JSONRPCMessage = serde_json::from_str(&response_line)?; + let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else { + panic!("expected initialize response"); + }; + assert_eq!(id, RequestId::Integer(1)); + let initialize_response: InitializeResponse = serde_json::from_value(result)?; + assert_eq!(initialize_response.protocol_version, "exec-server.v0"); + + let initialized = JSONRPCMessage::Notification(JSONRPCNotification { + method: "initialized".to_string(), + params: Some(serde_json::json!({})), + }); + stdin + .write_all(format!("{}\n", serde_json::to_string(&initialized)?).as_bytes()) + .await?; + + child.start_kill()?; + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +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 client = ExecServerClient::spawn(ExecServerLaunchCommand { + program: cargo_bin("codex-exec-server")?, + args: Vec::new(), + }) + .await?; + + let process = client + .start_process(ExecParams { + process_id: "2001".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, + output_bytes_cap: 4096, + arg0: None, + }) + .await?; + + let mut output = process.output_receiver(); + assert!( + recv_until_contains(&mut output, "ready") + .await? + .contains("ready"), + "expected initial ready output" + ); + + process + .writer_sender() + .send(b"hello\n".to_vec()) + .await + .expect("write should succeed"); + + assert!( + recv_until_contains(&mut output, "echo:hello") + .await? + .contains("echo:hello"), + "expected echoed output" + ); + + process.terminate(); + Ok(()) +} + +async fn recv_until_contains( + output: &mut broadcast::Receiver>, + needle: &str, +) -> anyhow::Result { + 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 chunk = timeout(remaining, output.recv()).await??; + collected.push_str(&String::from_utf8_lossy(&chunk)); + if collected.contains(needle) { + return Ok(collected); + } + } +}