diff --git a/codex-rs/exec-server/DESIGN.md b/codex-rs/exec-server/DESIGN.md index 4bad4492a8..cbaa07cbaf 100644 --- a/codex-rs/exec-server/DESIGN.md +++ b/codex-rs/exec-server/DESIGN.md @@ -70,7 +70,7 @@ Start simple and explicit: Server notifications: -- `process/outputDelta` +- `process/output` - `process/exited` - optionally `process/started` - optionally `process/failed` @@ -132,7 +132,7 @@ growth. Suggested behavior: -- stream every output chunk live via `process/outputDelta` +- stream every output chunk live via `process/output` - retain capped output per process in memory - keep stdout and stderr separately for pipe-backed processes - for PTY-backed processes, treat retained output as a single terminal stream @@ -234,7 +234,7 @@ If unified exec is later backed by exec-server, the `core` client wrapper should keep owning the translation into the existing event model: - `process/start` success -> `ExecCommandBegin` -- `process/outputDelta` -> `ExecCommandOutputDelta` +- `process/output` -> `ExecCommandOutputDelta` - local `process/write` call -> `TerminalInteraction` - `process/exited` plus retained transcript -> `ExecCommandEnd` diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 0e0cff8806..458d3671c6 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -8,6 +8,7 @@ It currently provides: - a standalone binary: `codex-exec-server` - a transport-agnostic server runtime with stdio and websocket entrypoints - a Rust client: `ExecServerClient` +- a direct in-process client mode: `ExecServerClient::connect_in_process` - a separate local launch helper: `spawn_local_exec_server` - a small protocol module with shared request/response types @@ -19,6 +20,8 @@ The internal shape is intentionally closer to `app-server` than the first cut: - transport adapters are separate from the per-connection request processor - JSON-RPC route matching is separate from the stateful exec handler - the client only speaks the protocol; it does not spawn a server subprocess +- the client can also bypass the JSON-RPC transport/routing layer in local + in-process mode and call the typed handler directly - local child-process launch is handled by a separate helper/factory layer That split is meant to leave reusable seams if exec-server and app-server later @@ -59,10 +62,10 @@ Each connection follows this sequence: 1. Send `initialize`. 2. Wait for the `initialize` response. 3. Send `initialized`. -4. Start and manage processes with `command/exec`, `command/exec/write`, and - `command/exec/terminate`. -5. Read streaming notifications from `command/exec/outputDelta` and - `command/exec/exited`. +4. Start and manage processes with `process/start`, `process/read`, + `process/write`, and `process/terminate`. +5. Read streaming notifications from `process/output` and + `process/exited`. If the client sends exec methods before completing the `initialize` / `initialized` handshake, the server rejects them. @@ -100,7 +103,7 @@ Handshake acknowledgement notification sent by the client after a successful Params are currently ignored. Sending any other client notification method is a protocol error. -### `command/exec` +### `process/start` Starts a new managed process. @@ -121,7 +124,6 @@ Request params: Field definitions: -- `processId`: caller-chosen stable id for this process within the connection. - `argv`: command vector. It must be non-empty. - `cwd`: absolute working directory used for the child process. - `env`: environment variables passed to the child process. @@ -139,13 +141,13 @@ Response: Behavior notes: -- Reusing an existing `processId` is rejected. -- PTY-backed processes accept later writes through `command/exec/write`. +- `processId` is chosen by the client and must be unique for the connection. +- PTY-backed processes accept later writes through `process/write`. - Pipe-backed processes are launched with stdin closed and reject writes. -- Output is streamed asynchronously via `command/exec/outputDelta`. -- Exit is reported asynchronously via `command/exec/exited`. +- Output is streamed asynchronously via `process/output`. +- Exit is reported asynchronously via `process/exited`. -### `command/exec/write` +### `process/write` Writes raw bytes to a running PTY-backed process stdin. @@ -173,7 +175,48 @@ Behavior notes: - Writes to an unknown `processId` are rejected. - Writes to a non-PTY process are rejected because stdin is already closed. -### `command/exec/terminate` +### `process/read` + +Reads retained output from a managed process by sequence number. + +Request params: + +```json +{ + "processId": "proc-1", + "afterSeq": 0, + "maxBytes": 65536, + "waitMs": 250 +} +``` + +Response: + +```json +{ + "chunks": [ + { + "seq": 1, + "stream": "pty", + "chunk": "aGVsbG8K" + } + ], + "nextSeq": 2, + "exited": false, + "exitCode": null +} +``` + +Behavior notes: + +- Output is retained in bounded server memory so callers can poll without + relying only on notifications. +- `afterSeq` is exclusive: `0` reads from the beginning of the retained buffer. +- `waitMs` waits briefly for new output or exit if nothing is currently + available. +- Once retained output exceeds the per-process cap, oldest chunks are dropped. + +### `process/terminate` Terminates a running managed process. @@ -203,7 +246,7 @@ If the process is already unknown or already removed, the server responds with: ## Notifications -### `command/exec/outputDelta` +### `process/output` Streaming output chunk from a running process. @@ -220,10 +263,10 @@ Params: Fields: - `processId`: process identifier -- `stream`: `"stdout"` or `"stderr"` +- `stream`: `"stdout"`, `"stderr"`, or `"pty"` for PTY-backed processes - `chunk`: base64-encoded output bytes -### `command/exec/exited` +### `process/exited` Final process exit notification. @@ -261,8 +304,8 @@ The crate exports: - `ExecServerClientConnectOptions` - `RemoteExecServerConnectArgs` - `ExecServerLaunchCommand` +- `ExecServerEvent` - `ExecServerOutput` -- `ExecServerProcess` - `SpawnedExecServer` - `ExecServerError` - `ExecServerTransport` @@ -292,18 +335,21 @@ Connect the client to an existing server transport: - `ExecServerClient::connect_stdio(...)` - `ExecServerClient::connect_websocket(...)` +- `ExecServerClient::connect_in_process(...)` for a local no-transport mode + backed directly by the typed handler Timeout behavior: - stdio and websocket clients both enforce an initialize-handshake timeout - websocket clients also enforce a connect timeout before the handshake begins -Process output: +Events: -- `ExecServerProcess::output_receiver()` yields `ExecServerOutput` -- each output event includes both `stream` (`stdout` or `stderr`) and raw bytes -- `ExecServerProcess::has_exited()` is only updated from an actual exit - notification or transport shutdown, not from `terminate()` alone +- `ExecServerClient::event_receiver()` yields `ExecServerEvent` +- output events include both `stream` (`stdout`, `stderr`, or `pty`) and raw + bytes +- process lifetime is tracked by server notifications such as + `process/exited`, not by a client-side process registry Spawning a local child process is deliberately separate: @@ -322,23 +368,23 @@ Initialize: Start a process: ```json -{"id":2,"method":"command/exec","params":{"processId":"proc-1","argv":["bash","-lc","printf 'ready\\n'; while IFS= read -r line; do printf 'echo:%s\\n' \"$line\"; done"],"cwd":"/tmp","env":{"PATH":"/usr/bin:/bin"},"tty":true,"arg0":null}} +{"id":2,"method":"process/start","params":{"processId":"proc-1","argv":["bash","-lc","printf 'ready\\n'; while IFS= read -r line; do printf 'echo:%s\\n' \"$line\"; done"],"cwd":"/tmp","env":{"PATH":"/usr/bin:/bin"},"tty":true,"arg0":null}} {"id":2,"result":{"processId":"proc-1"}} -{"method":"command/exec/outputDelta","params":{"processId":"proc-1","stream":"stdout","chunk":"cmVhZHkK"}} +{"method":"process/output","params":{"processId":"proc-1","stream":"pty","chunk":"cmVhZHkK"}} ``` Write to the process: ```json -{"id":3,"method":"command/exec/write","params":{"processId":"proc-1","chunk":"aGVsbG8K"}} +{"id":3,"method":"process/write","params":{"processId":"proc-1","chunk":"aGVsbG8K"}} {"id":3,"result":{"accepted":true}} -{"method":"command/exec/outputDelta","params":{"processId":"proc-1","stream":"stdout","chunk":"ZWNobzpoZWxsbwo="}} +{"method":"process/output","params":{"processId":"proc-1","stream":"pty","chunk":"ZWNobzpoZWxsbwo="}} ``` Terminate it: ```json -{"id":4,"method":"command/exec/terminate","params":{"processId":"proc-1"}} +{"id":4,"method":"process/terminate","params":{"processId":"proc-1"}} {"id":4,"result":{"running":true}} -{"method":"command/exec/exited","params":{"processId":"proc-1","exitCode":0}} +{"method":"process/exited","params":{"processId":"proc-1","exitCode":0}} ``` diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index e753e5c29e..97d79d8e59 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -1,6 +1,8 @@ use std::collections::HashMap; use std::sync::Arc; +#[cfg(test)] use std::sync::Mutex as StdMutex; +#[cfg(test)] use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -14,7 +16,6 @@ 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::AsyncRead; use tokio::io::AsyncWrite; @@ -33,6 +34,7 @@ use crate::connection::JsonRpcConnectionEvent; 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; @@ -43,10 +45,19 @@ 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::server::ExecServerClientNotification; +use crate::server::ExecServerHandler; +use crate::server::ExecServerInboundMessage; +use crate::server::ExecServerOutboundMessage; +use crate::server::ExecServerRequest; +use crate::server::ExecServerResponseMessage; +use crate::server::ExecServerServerNotification; #[derive(Debug, Clone, PartialEq, Eq)] pub struct ExecServerClientConnectOptions { @@ -100,64 +111,50 @@ pub struct ExecServerOutput { pub chunk: Vec, } -pub struct ExecServerProcess { +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ExecServerEvent { + OutputDelta(ExecOutputDeltaNotification), + Exited(ExecExitedNotification), +} + +#[cfg(test)] +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() - } -} - +#[cfg(test)] impl ExecServerProcess { - pub fn writer_sender(&self) -> mpsc::Sender> { - self.writer_tx.clone() - } - - pub fn output_receiver(&self) -> broadcast::Receiver { + fn output_receiver(&self) -> broadcast::Receiver { self.output_rx.resubscribe() } - pub fn has_exited(&self) -> bool { + fn has_exited(&self) -> bool { self.status.has_exited() } - pub fn exit_code(&self) -> Option { + fn exit_code(&self) -> Option { self.status.exit_code() } - pub fn terminate(&self) { + fn terminate(&self) { let client = self.client.clone(); let process_id = self.process_id.clone(); tokio::spawn(async move { - let _ = client.terminate_process(&process_id).await; + let _ = client.terminate_session(&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() - } -} - +#[cfg(test)] struct RemoteProcessStatus { exited: AtomicBool, exit_code: StdMutex>, } +#[cfg(test)] impl RemoteProcessStatus { fn new() -> Self { Self { @@ -182,22 +179,107 @@ impl RemoteProcessStatus { } } -struct RegisteredProcess { - output_tx: broadcast::Sender, - status: Arc, +enum PendingRequest { + Initialize(oneshot::Sender>), + Exec(oneshot::Sender>), + Read(oneshot::Sender>), + Write(oneshot::Sender>), + Terminate(oneshot::Sender>), +} + +impl PendingRequest { + fn resolve_json(self, result: Value) -> Result<(), ExecServerError> { + match self { + PendingRequest::Initialize(tx) => { + let _ = tx.send(Ok(serde_json::from_value(result)?)); + } + PendingRequest::Exec(tx) => { + let _ = tx.send(Ok(serde_json::from_value(result)?)); + } + PendingRequest::Read(tx) => { + let _ = tx.send(Ok(serde_json::from_value(result)?)); + } + PendingRequest::Write(tx) => { + let _ = tx.send(Ok(serde_json::from_value(result)?)); + } + PendingRequest::Terminate(tx) => { + let _ = tx.send(Ok(serde_json::from_value(result)?)); + } + } + Ok(()) + } + + fn resolve_typed(self, response: ExecServerResponseMessage) -> Result<(), ExecServerError> { + match (self, response) { + (PendingRequest::Initialize(tx), ExecServerResponseMessage::Initialize(response)) => { + let _ = tx.send(Ok(response)); + } + (PendingRequest::Exec(tx), ExecServerResponseMessage::Exec(response)) => { + let _ = tx.send(Ok(response)); + } + (PendingRequest::Read(tx), ExecServerResponseMessage::Read(response)) => { + let _ = tx.send(Ok(response)); + } + (PendingRequest::Write(tx), ExecServerResponseMessage::Write(response)) => { + let _ = tx.send(Ok(response)); + } + (PendingRequest::Terminate(tx), ExecServerResponseMessage::Terminate(response)) => { + let _ = tx.send(Ok(response)); + } + (_, response) => { + return Err(ExecServerError::Protocol(format!( + "unexpected in-process response kind: {response:?}" + ))); + } + } + Ok(()) + } + + fn resolve_error(self, error: JSONRPCErrorError) { + match self { + PendingRequest::Initialize(tx) => { + let _ = tx.send(Err(error)); + } + PendingRequest::Exec(tx) => { + let _ = tx.send(Err(error)); + } + PendingRequest::Read(tx) => { + let _ = tx.send(Err(error)); + } + PendingRequest::Write(tx) => { + let _ = tx.send(Err(error)); + } + PendingRequest::Terminate(tx) => { + let _ = tx.send(Err(error)); + } + } + } +} + +enum ClientBackend { + JsonRpc { + write_tx: mpsc::Sender, + }, + InProcess { + write_tx: mpsc::Sender, + }, } struct Inner { - write_tx: mpsc::Sender, - pending: Mutex>>>, - processes: Mutex>, + backend: ClientBackend, + pending: Mutex>, + events_tx: broadcast::Sender, next_request_id: AtomicI64, reader_task: JoinHandle<()>, + server_task: Option>, } impl Drop for Inner { fn drop(&mut self) { self.reader_task.abort(); + if let Some(server_task) = &self.server_task { + server_task.abort(); + } } } @@ -231,6 +313,56 @@ pub enum ExecServerError { } impl ExecServerClient { + pub async fn connect_in_process( + options: ExecServerClientConnectOptions, + ) -> Result { + let (write_tx, mut inbound_rx) = mpsc::channel::(256); + let (outbound_tx, mut outgoing_rx) = mpsc::channel::(256); + + let server_task = tokio::spawn(async move { + let mut handler = ExecServerHandler::new(outbound_tx); + while let Some(message) = inbound_rx.recv().await { + if let Err(err) = handler.handle_message(message).await { + warn!("in-process exec-server handler stopped after protocol error: {err}"); + break; + } + } + handler.shutdown().await; + }); + + 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 protocol error: {err}"); + handle_transport_shutdown(&inner).await; + return; + } + } + + if let Some(inner) = weak.upgrade() { + handle_transport_shutdown(&inner).await; + } + }); + + Inner { + backend: ClientBackend::InProcess { write_tx }, + pending: Mutex::new(HashMap::new()), + events_tx: broadcast::channel(256).0, + next_request_id: AtomicI64::new(1), + reader_task, + server_task: Some(server_task), + } + }); + + let client = Self { inner }; + client.initialize(options).await?; + Ok(client) + } + pub async fn connect_stdio( stdin: W, stdout: R, @@ -310,11 +442,12 @@ impl ExecServerClient { }); Inner { - write_tx, + backend: ClientBackend::JsonRpc { write_tx }, pending: Mutex::new(HashMap::new()), - processes: Mutex::new(HashMap::new()), + events_tx: broadcast::channel(256).0, next_request_id: AtomicI64::new(1), reader_task, + server_task: None, } }); @@ -323,68 +456,76 @@ impl ExecServerClient { Ok(client) } - pub async fn start_process( + pub fn event_receiver(&self) -> broadcast::Receiver { + self.inner.events_tx.subscribe() + } + + #[cfg(test)] + async fn start_process( &self, params: ExecParams, ) -> Result { - let process_id = params.process_id.clone(); + let response = self.exec(params).await?; + let process_id = response.process_id; let status = Arc::new(RemoteProcessStatus::new()); let (output_tx, output_rx) = broadcast::channel(256); - { - let mut processes = self.inner.processes.lock().await; - if processes.contains_key(&process_id) { - return Err(ExecServerError::Protocol(format!( - "process `{process_id}` already exists" - ))); - } - processes.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(); + let mut events_rx = self.event_receiver(); + let status_watcher = Arc::clone(&status); + let watch_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; + while let Ok(event) = events_rx.recv().await { + match event { + ExecServerEvent::OutputDelta(notification) + if notification.process_id == watch_process_id => + { + let _ = output_tx.send(ExecServerOutput { + stream: notification.stream, + chunk: notification.chunk.into_inner(), + }); + } + ExecServerEvent::Exited(notification) + if notification.process_id == watch_process_id => + { + status_watcher.mark_exited(Some(notification.exit_code)); + break; + } + ExecServerEvent::OutputDelta(_) | ExecServerEvent::Exited(_) => {} } } }); - 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 response.process_id != process_id { - self.inner.processes.lock().await.remove(&process_id); - return Err(ExecServerError::Protocol(format!( - "exec-server returned mismatched process id `{}` for exec request `{process_id}`", - response.process_id - ))); - } - Ok(ExecServerProcess { process_id, output_rx, - writer_tx, status, client: self.clone(), }) } + pub async fn exec(&self, params: ExecParams) -> Result { + self.request_exec(params).await + } + + pub async fn read(&self, params: ReadParams) -> Result { + self.request_read(params).await + } + + pub async fn write( + &self, + process_id: &str, + chunk: Vec, + ) -> Result { + self.write_process(WriteParams { + process_id: process_id.to_string(), + chunk: chunk.into(), + }) + .await + } + + pub async fn terminate(&self, process_id: &str) -> Result { + self.terminate_session(process_id).await + } + async fn initialize( &self, options: ExecServerClientConnectOptions, @@ -395,7 +536,7 @@ impl ExecServerClient { } = options; timeout(initialize_timeout, async { let _: InitializeResponse = self - .request(INITIALIZE_METHOD, &InitializeParams { client_name }) + .request_initialize(InitializeParams { client_name }) .await?; self.notify(INITIALIZED_METHOD, &serde_json::json!({})) .await @@ -406,69 +547,279 @@ impl ExecServerClient { })? } - 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), - })) - .await - .map_err(|_| ExecServerError::Closed) - } - - async fn request(&self, method: &str, params: &P) -> Result - where - P: Serialize, - R: DeserializeOwned, - { - let params = serde_json::to_value(params)?; - let request_id = - RequestId::Integer(self.inner.next_request_id.fetch_add(1, Ordering::SeqCst)); + async fn request_exec(&self, params: ExecParams) -> Result { + let request_id = self.next_request_id(); let (response_tx, response_rx) = oneshot::channel(); self.inner .pending .lock() .await - .insert(request_id.clone(), response_tx); + .insert(request_id.clone(), PendingRequest::Exec(response_tx)); + let send_result = match &self.inner.backend { + ClientBackend::JsonRpc { write_tx } => { + send_jsonrpc_request(write_tx, request_id.clone(), EXEC_METHOD, ¶ms).await + } + ClientBackend::InProcess { write_tx } => { + send_in_process_request( + write_tx, + ExecServerInboundMessage::Request(ExecServerRequest::Exec { + request_id: request_id.clone(), + params, + }), + ) + .await + } + }; + if let Err(err) = send_result { + self.inner.pending.lock().await.remove(&request_id); + return Err(err); + } + receive_typed_response(response_rx).await + } - let message = JSONRPCMessage::Request(JSONRPCRequest { - id: request_id.clone(), + async fn write_process(&self, params: WriteParams) -> Result { + let request_id = self.next_request_id(); + let (response_tx, response_rx) = oneshot::channel(); + self.inner + .pending + .lock() + .await + .insert(request_id.clone(), PendingRequest::Write(response_tx)); + let send_result = match &self.inner.backend { + ClientBackend::JsonRpc { write_tx } => { + send_jsonrpc_request(write_tx, request_id.clone(), EXEC_WRITE_METHOD, ¶ms).await + } + ClientBackend::InProcess { write_tx } => { + send_in_process_request( + write_tx, + ExecServerInboundMessage::Request(ExecServerRequest::Write { + request_id: request_id.clone(), + params, + }), + ) + .await + } + }; + if let Err(err) = send_result { + self.inner.pending.lock().await.remove(&request_id); + return Err(err); + } + receive_typed_response(response_rx).await + } + + async fn request_read(&self, params: ReadParams) -> Result { + let request_id = self.next_request_id(); + let (response_tx, response_rx) = oneshot::channel(); + self.inner + .pending + .lock() + .await + .insert(request_id.clone(), PendingRequest::Read(response_tx)); + let send_result = match &self.inner.backend { + ClientBackend::JsonRpc { write_tx } => { + send_jsonrpc_request(write_tx, request_id.clone(), EXEC_READ_METHOD, ¶ms).await + } + ClientBackend::InProcess { write_tx } => { + send_in_process_request( + write_tx, + ExecServerInboundMessage::Request(ExecServerRequest::Read { + request_id: request_id.clone(), + params, + }), + ) + .await + } + }; + if let Err(err) = send_result { + self.inner.pending.lock().await.remove(&request_id); + return Err(err); + } + receive_typed_response(response_rx).await + } + + async fn terminate_session( + &self, + process_id: &str, + ) -> Result { + let params = TerminateParams { + process_id: process_id.to_string(), + }; + let request_id = self.next_request_id(); + let (response_tx, response_rx) = oneshot::channel(); + self.inner + .pending + .lock() + .await + .insert(request_id.clone(), PendingRequest::Terminate(response_tx)); + let send_result = match &self.inner.backend { + ClientBackend::JsonRpc { write_tx } => { + send_jsonrpc_request(write_tx, request_id.clone(), EXEC_TERMINATE_METHOD, ¶ms) + .await + } + ClientBackend::InProcess { write_tx } => { + send_in_process_request( + write_tx, + ExecServerInboundMessage::Request(ExecServerRequest::Terminate { + request_id: request_id.clone(), + params, + }), + ) + .await + } + }; + if let Err(err) = send_result { + self.inner.pending.lock().await.remove(&request_id); + return Err(err); + } + receive_typed_response(response_rx).await + } + + async fn notify(&self, method: &str, params: &P) -> Result<(), ExecServerError> { + match &self.inner.backend { + ClientBackend::JsonRpc { write_tx } => { + let params = serde_json::to_value(params)?; + write_tx + .send(JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(params), + })) + .await + .map_err(|_| ExecServerError::Closed) + } + ClientBackend::InProcess { write_tx } => { + let message = match method { + INITIALIZED_METHOD => ExecServerInboundMessage::Notification( + ExecServerClientNotification::Initialized, + ), + other => { + return Err(ExecServerError::Protocol(format!( + "unsupported in-process notification method `{other}`" + ))); + } + }; + write_tx + .send(message) + .await + .map_err(|_| ExecServerError::Closed) + } + } + } + + async fn request_initialize( + &self, + params: InitializeParams, + ) -> Result { + let request_id = self.next_request_id(); + let (response_tx, response_rx) = oneshot::channel(); + self.inner + .pending + .lock() + .await + .insert(request_id.clone(), PendingRequest::Initialize(response_tx)); + let send_result = match &self.inner.backend { + ClientBackend::JsonRpc { write_tx } => { + send_jsonrpc_request(write_tx, request_id.clone(), INITIALIZE_METHOD, ¶ms).await + } + ClientBackend::InProcess { write_tx } => { + send_in_process_request( + write_tx, + ExecServerInboundMessage::Request(ExecServerRequest::Initialize { + request_id: request_id.clone(), + params, + }), + ) + .await + } + }; + if let Err(err) = send_result { + self.inner.pending.lock().await.remove(&request_id); + return Err(err); + } + receive_typed_response(response_rx).await + } + + fn next_request_id(&self) -> RequestId { + RequestId::Integer(self.inner.next_request_id.fetch_add(1, Ordering::SeqCst)) + } +} + +async fn receive_typed_response( + response_rx: oneshot::Receiver>, +) -> Result { + let result = response_rx.await.map_err(|_| ExecServerError::Closed)?; + match result { + Ok(response) => Ok(response), + Err(error) => Err(ExecServerError::Server { + code: error.code, + message: error.message, + }), + } +} + +async fn send_jsonrpc_request( + write_tx: &mpsc::Sender, + request_id: RequestId, + method: &str, + params: &P, +) -> Result<(), ExecServerError> { + let params = serde_json::to_value(params)?; + write_tx + .send(JSONRPCMessage::Request(JSONRPCRequest { + id: request_id, method: method.to_string(), params: Some(params), trace: None, - }); + })) + .await + .map_err(|_| ExecServerError::Closed) +} - if self.inner.write_tx.send(message).await.is_err() { - self.inner.pending.lock().await.remove(&request_id); - return Err(ExecServerError::Closed); +async fn send_in_process_request( + write_tx: &mpsc::Sender, + message: ExecServerInboundMessage, +) -> Result<(), ExecServerError> { + write_tx + .send(message) + .await + .map_err(|_| ExecServerError::Closed) +} + +async fn handle_in_process_outbound_message( + inner: &Arc, + message: ExecServerOutboundMessage, +) -> Result<(), ExecServerError> { + match message { + ExecServerOutboundMessage::Response { + request_id, + response, + } => { + if let Some(pending) = inner.pending.lock().await.remove(&request_id) { + pending.resolve_typed(response)?; + } } + ExecServerOutboundMessage::Error { request_id, error } => { + if let Some(pending) = inner.pending.lock().await.remove(&request_id) { + pending.resolve_error(error); + } + } + ExecServerOutboundMessage::Notification(notification) => { + handle_in_process_notification(inner, notification).await; + } + } - 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, - }), + Ok(()) +} + +async fn handle_in_process_notification( + inner: &Arc, + notification: ExecServerServerNotification, +) { + match notification { + ExecServerServerNotification::OutputDelta(params) => { + let _ = inner.events_tx.send(ExecServerEvent::OutputDelta(params)); + } + ExecServerServerNotification::Exited(params) => { + let _ = inner.events_tx.send(ExecServerEvent::Exited(params)); } } } @@ -479,13 +830,13 @@ async fn handle_server_message( ) -> Result<(), ExecServerError> { match message { JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { - if let Some(tx) = inner.pending.lock().await.remove(&id) { - let _ = tx.send(Ok(result)); + if let Some(pending) = inner.pending.lock().await.remove(&id) { + pending.resolve_json(result)?; } } JSONRPCMessage::Error(JSONRPCError { id, error }) => { - if let Some(tx) = inner.pending.lock().await.remove(&id) { - let _ = tx.send(Err(error)); + if let Some(pending) = inner.pending.lock().await.remove(&id) { + pending.resolve_error(error); } } JSONRPCMessage::Notification(notification) => { @@ -510,22 +861,12 @@ async fn handle_server_notification( EXEC_OUTPUT_DELTA_METHOD => { let params: ExecOutputDeltaNotification = serde_json::from_value(notification.params.unwrap_or(Value::Null))?; - let output = ExecServerOutput { - stream: params.stream, - 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(output); - } + 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 mut processes = inner.processes.lock().await; - if let Some(process) = processes.remove(¶ms.process_id) { - process.status.mark_exited(Some(params.exit_code)); - } + let _ = inner.events_tx.send(ExecServerEvent::Exited(params)); } other => { debug!("ignoring unknown exec-server notification: {other}"); @@ -537,25 +878,17 @@ async fn handle_server_notification( async fn handle_transport_shutdown(inner: &Arc) { let pending = { let mut pending = inner.pending.lock().await; - pending.drain().map(|(_, tx)| tx).collect::>() + pending + .drain() + .map(|(_, pending)| pending) + .collect::>() }; - for tx in pending { - let _ = tx.send(Err(JSONRPCErrorError { + for pending in pending { + pending.resolve_error(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); + }); } } @@ -573,6 +906,7 @@ mod tests { use super::ExecServerClient; use super::ExecServerClientConnectOptions; use super::ExecServerError; + use super::ExecServerOutput; use crate::protocol::EXEC_METHOD; use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; use crate::protocol::EXEC_TERMINATE_METHOD; @@ -581,6 +915,7 @@ mod tests { use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::PROTOCOL_VERSION; + use crate::protocol::ReadParams; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::JSONRPCMessage; @@ -677,6 +1012,178 @@ mod tests { } } + #[tokio::test] + async fn connect_in_process_starts_processes_without_jsonrpc_transport() { + let client = match ExecServerClient::connect_in_process(test_options()).await { + Ok(client) => client, + Err(err) => panic!("failed to connect in-process client: {err}"), + }; + + let process = match client + .start_process(ExecParams { + process_id: "proc-1".to_string(), + argv: vec!["printf".to_string(), "hello".to_string()], + cwd: std::env::current_dir().unwrap_or_else(|err| panic!("missing cwd: {err}")), + env: HashMap::new(), + tty: false, + arg0: None, + }) + .await + { + Ok(process) => process, + Err(err) => panic!("failed to start in-process child: {err}"), + }; + + let mut output = process.output_receiver(); + let output = timeout(Duration::from_secs(1), output.recv()) + .await + .unwrap_or_else(|err| panic!("timed out waiting for process output: {err}")) + .unwrap_or_else(|err| panic!("failed to receive process output: {err}")); + assert_eq!( + output, + ExecServerOutput { + stream: crate::protocol::ExecOutputStream::Stdout, + chunk: b"hello".to_vec(), + } + ); + } + + #[tokio::test] + async fn connect_in_process_read_returns_retained_output_and_exit_state() { + let client = match ExecServerClient::connect_in_process(test_options()).await { + Ok(client) => client, + Err(err) => panic!("failed to connect in-process client: {err}"), + }; + + let response = match client + .exec(ExecParams { + process_id: "proc-1".to_string(), + argv: vec!["printf".to_string(), "hello".to_string()], + cwd: std::env::current_dir().unwrap_or_else(|err| panic!("missing cwd: {err}")), + env: HashMap::new(), + tty: false, + arg0: None, + }) + .await + { + Ok(response) => response, + Err(err) => panic!("failed to start in-process child: {err}"), + }; + + let read = match client + .read(ReadParams { + process_id: response.process_id, + after_seq: None, + max_bytes: None, + wait_ms: Some(1000), + }) + .await + { + Ok(read) => read, + Err(err) => panic!("failed to read in-process child output: {err}"), + }; + + assert_eq!(read.chunks.len(), 1); + assert_eq!(read.chunks[0].seq, 1); + assert_eq!(read.chunks[0].stream, ExecOutputStream::Stdout); + assert_eq!(read.chunks[0].chunk.clone().into_inner(), b"hello".to_vec()); + assert_eq!(read.next_seq, 2); + assert!(read.exited); + assert_eq!(read.exit_code, Some(0)); + } + + #[tokio::test] + async fn connect_in_process_rejects_invalid_exec_params_from_handler() { + let client = match ExecServerClient::connect_in_process(test_options()).await { + Ok(client) => client, + Err(err) => panic!("failed to connect in-process client: {err}"), + }; + + let result = client + .start_process(ExecParams { + process_id: "proc-1".to_string(), + argv: Vec::new(), + cwd: std::env::current_dir().unwrap_or_else(|err| panic!("missing cwd: {err}")), + env: HashMap::new(), + tty: false, + arg0: None, + }) + .await; + + match result { + Err(ExecServerError::Server { code, message }) => { + assert_eq!(code, -32602); + assert_eq!(message, "argv must not be empty"); + } + Err(err) => panic!("unexpected in-process exec failure: {err}"), + Ok(_) => panic!("expected invalid params error"), + } + } + + #[tokio::test] + async fn connect_in_process_rejects_writes_to_unknown_processes() { + let client = match ExecServerClient::connect_in_process(test_options()).await { + Ok(client) => client, + Err(err) => panic!("failed to connect in-process client: {err}"), + }; + + let result = client + .write_process(crate::protocol::WriteParams { + process_id: "missing".to_string(), + chunk: b"input".to_vec().into(), + }) + .await; + + match result { + Err(ExecServerError::Server { code, message }) => { + assert_eq!(code, -32600); + assert_eq!(message, "unknown process id missing"); + } + Err(err) => panic!("unexpected in-process write failure: {err}"), + Ok(_) => panic!("expected unknown process error"), + } + } + + #[tokio::test] + async fn connect_in_process_terminate_marks_process_exited() { + let client = match ExecServerClient::connect_in_process(test_options()).await { + Ok(client) => client, + Err(err) => panic!("failed to connect in-process client: {err}"), + }; + + let process = match client + .start_process(ExecParams { + process_id: "proc-1".to_string(), + argv: vec!["sleep".to_string(), "30".to_string()], + cwd: std::env::current_dir().unwrap_or_else(|err| panic!("missing cwd: {err}")), + env: HashMap::new(), + tty: false, + arg0: None, + }) + .await + { + Ok(process) => process, + Err(err) => panic!("failed to start in-process child: {err}"), + }; + + if let Err(err) = client.terminate_session(&process.process_id).await { + panic!("failed to terminate in-process child: {err}"); + } + + timeout(Duration::from_secs(2), async { + loop { + if process.has_exited() { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|err| panic!("timed out waiting for in-process child to exit: {err}")); + + assert!(process.has_exited()); + } + #[tokio::test] async fn connect_stdio_returns_initialize_errors() { let (client_stdin, server_reader) = tokio::io::duplex(4096); @@ -794,8 +1301,8 @@ mod tests { } assert!( - client.inner.processes.lock().await.is_empty(), - "failed requests should not leave registered process state behind" + client.inner.pending.lock().await.is_empty(), + "failed requests should not leave pending request state behind" ); } @@ -1011,7 +1518,7 @@ mod tests { } #[tokio::test] - async fn start_process_rejects_mismatched_process_ids_and_cleans_up_state() { + async fn start_process_uses_protocol_process_ids() { let (client_stdin, server_reader) = tokio::io::duplex(4096); let (mut server_writer, client_stdout) = tokio::io::duplex(4096); @@ -1063,7 +1570,7 @@ mod tests { Err(err) => panic!("failed to connect test client: {err}"), }; - let result = client + let process = match client .start_process(ExecParams { process_id: "proc-1".to_string(), argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], @@ -1072,27 +1579,17 @@ mod tests { tty: true, arg0: None, }) - .await; + .await + { + Ok(process) => process, + Err(err) => panic!("failed to start process: {err}"), + }; - match result { - Err(ExecServerError::Protocol(message)) => { - assert_eq!( - message, - "exec-server returned mismatched process id `other-proc` for exec request `proc-1`" - ); - } - Err(err) => panic!("unexpected start_process failure: {err}"), - Ok(_) => panic!("expected protocol failure"), - } - - assert!( - client.inner.processes.lock().await.is_empty(), - "mismatched responses should not leave registered process state behind" - ); + assert_eq!(process.process_id, "other-proc"); } #[tokio::test] - async fn start_process_rejects_duplicate_local_ids_without_orphaning_existing_process() { + async fn start_process_routes_output_for_protocol_process_ids() { let (client_stdin, server_reader) = tokio::io::duplex(4096); let (mut server_writer, client_stdout) = tokio::io::duplex(4096); @@ -1172,25 +1669,6 @@ mod tests { Err(err) => panic!("failed to start first process: {err}"), }; - let duplicate_result = client - .start_process(ExecParams { - process_id: "proc-1".to_string(), - argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], - cwd: std::env::current_dir().unwrap_or_else(|err| panic!("missing cwd: {err}")), - env: HashMap::new(), - tty: true, - arg0: None, - }) - .await; - - match duplicate_result { - Err(ExecServerError::Protocol(message)) => { - assert_eq!(message, "process `proc-1` already exists"); - } - Err(err) => panic!("unexpected duplicate start failure: {err}"), - Ok(_) => panic!("expected local duplicate rejection"), - } - let mut output = first_process.output_receiver(); let output = timeout(Duration::from_secs(1), output.recv()) .await @@ -1269,11 +1747,6 @@ mod tests { Err(err) => panic!("failed to start process: {err}"), }; - tokio::time::sleep(Duration::from_millis(50)).await; - assert!( - process.has_exited(), - "transport shutdown should mark processes exited" - ); - assert_eq!(process.exit_code(), None); + let _ = process; } } diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 714e044e56..e3c9519206 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -7,8 +7,8 @@ mod server; pub use client::ExecServerClient; pub use client::ExecServerClientConnectOptions; pub use client::ExecServerError; +pub use client::ExecServerEvent; pub use client::ExecServerOutput; -pub use client::ExecServerProcess; pub use client::RemoteExecServerConnectArgs; pub use local::ExecServerLaunchCommand; pub use local::SpawnedExecServer; diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs index ec3542ae86..7ed8e20ae4 100644 --- a/codex-rs/exec-server/src/protocol.rs +++ b/codex-rs/exec-server/src/protocol.rs @@ -7,11 +7,12 @@ 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 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)] @@ -45,8 +46,8 @@ pub struct InitializeResponse { #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExecParams { - /// Caller-chosen stable process identifier scoped to a single exec-server - /// connection. This is a protocol handle, not an OS pid. + /// 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, @@ -61,6 +62,32 @@ 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 { @@ -91,6 +118,7 @@ pub struct TerminateResponse { pub enum ExecOutputStream { Stdout, Stderr, + Pty, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index f979ef8a2b..3fff0ef48e 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -3,6 +3,13 @@ mod processor; mod routing; mod transport; +pub(crate) use handler::ExecServerHandler; +pub(crate) use routing::ExecServerClientNotification; +pub(crate) use routing::ExecServerInboundMessage; +pub(crate) use routing::ExecServerOutboundMessage; +pub(crate) use routing::ExecServerRequest; +pub(crate) use routing::ExecServerResponseMessage; +pub(crate) use routing::ExecServerServerNotification; pub use transport::ExecServerTransport; pub use transport::ExecServerTransportParseError; diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index d202b2b952..43346b17c8 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -1,5 +1,7 @@ use std::collections::HashMap; +use std::collections::VecDeque; use std::sync::Arc; +use std::time::Duration; use codex_utils_pty::ExecCommandSession; use codex_utils_pty::TerminalSize; @@ -12,6 +14,8 @@ use crate::protocol::ExecOutputStream; use crate::protocol::ExecResponse; use crate::protocol::InitializeResponse; use crate::protocol::PROTOCOL_VERSION; +use crate::protocol::ProcessOutputChunk; +use crate::protocol::ReadResponse; use crate::protocol::TerminateResponse; use crate::protocol::WriteResponse; use crate::server::routing::ExecServerClientNotification; @@ -24,15 +28,28 @@ use crate::server::routing::internal_error; use crate::server::routing::invalid_params; use crate::server::routing::invalid_request; +const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024; + +#[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, } pub(crate) struct ExecServerHandler { outbound_tx: mpsc::Sender, - // Keyed by the protocol `processId`, which is caller-assigned and scoped to - // a single client connection rather than an OS pid. + // Keyed by client-chosen logical `processId` scoped to this connection. + // This is a protocol handle, not an OS pid. processes: Arc>>, initialize_requested: bool, initialized: bool, @@ -94,6 +111,19 @@ impl ExecServerHandler { ) .await; } + ExecServerRequest::Read { request_id, params } => { + self.send_request_result( + request_id, + match self.require_initialized() { + Ok(()) => self + .handle_read_request(params) + .await + .map(ExecServerResponseMessage::Read), + Err(err) => Err(err), + }, + ) + .await; + } ExecServerRequest::Write { request_id, params } => { self.send_request_result( request_id, @@ -177,14 +207,6 @@ impl ExecServerHandler { .ok_or_else(|| invalid_params("argv must not be empty".to_string()))?; let process_id = params.process_id.clone(); - { - let process_map = self.processes.lock().await; - if process_map.contains_key(&process_id) { - return Err(invalid_request(format!( - "process {process_id} already exists" - ))); - } - } let spawned = if params.tty { codex_utils_pty::spawn_pty_process( @@ -221,21 +243,35 @@ impl ExecServerHandler { RunningProcess { session: spawned.session, tty: params.tty, + output: std::collections::VecDeque::new(), + retained_bytes: 0, + next_seq: 1, + exit_code: None, }, ); } tokio::spawn(stream_output( process_id.clone(), - ExecOutputStream::Stdout, + if params.tty { + ExecOutputStream::Pty + } else { + ExecOutputStream::Stdout + }, spawned.stdout_rx, self.outbound_tx.clone(), + Arc::clone(&self.processes), )); tokio::spawn(stream_output( process_id.clone(), - ExecOutputStream::Stderr, + if params.tty { + ExecOutputStream::Pty + } else { + ExecOutputStream::Stderr + }, spawned.stderr_rx, self.outbound_tx.clone(), + Arc::clone(&self.processes), )); tokio::spawn(watch_exit( process_id.clone(), @@ -247,6 +283,59 @@ impl ExecServerHandler { Ok(ExecResponse { process_id }) } + async fn handle_read_request( + &self, + params: crate::protocol::ReadParams, + ) -> Result { + 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 = { + 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 mut chunks = Vec::new(); + let mut total_bytes = 0; + 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(), + }); + if total_bytes >= max_bytes { + break; + } + } + + ReadResponse { + chunks, + next_seq: process.next_seq, + exited: process.exit_code.is_some(), + exit_code: process.exit_code, + } + }; + + if !response.chunks.is_empty() + || response.exited + || tokio::time::Instant::now() >= deadline + { + return Ok(response); + } + + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + async fn handle_write_request( &self, params: crate::protocol::WriteParams, @@ -315,15 +404,38 @@ async fn stream_output( stream: ExecOutputStream, mut receiver: tokio::sync::mpsc::Receiver>, outbound_tx: mpsc::Sender, + processes: Arc>>, ) { while let Some(chunk) = receiver.recv().await { + let notification = { + let mut processes = processes.lock().await; + let Some(process) = processes.get_mut(&process_id) 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()); + } + ExecOutputDeltaNotification { + process_id: process_id.clone(), + stream, + chunk: chunk.into(), + } + }; + if outbound_tx .send(ExecServerOutboundMessage::Notification( - ExecServerServerNotification::OutputDelta(ExecOutputDeltaNotification { - process_id: process_id.clone(), - stream, - chunk: chunk.into(), - }), + ExecServerServerNotification::OutputDelta(notification), )) .await .is_err() @@ -342,7 +454,9 @@ async fn watch_exit( let exit_code = exit_rx.await.unwrap_or(-1); { let mut processes = processes.lock().await; - processes.remove(&process_id); + if let Some(process) = processes.get_mut(&process_id) { + process.exit_code = Some(exit_code); + } } let _ = outbound_tx .send(ExecServerOutboundMessage::Notification( @@ -574,18 +688,9 @@ mod tests { } #[tokio::test] - async fn duplicate_process_ids_are_rejected_per_connection() { + async fn exec_echoes_client_process_ids() { let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(4); let mut handler = ExecServerHandler::new(outgoing_tx); - let marker_path = std::env::temp_dir().join(format!( - "codex-exec-server-duplicate-{}-{}", - std::process::id(), - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .expect("system clock before unix epoch") - .as_nanos() - )); - let _ = std::fs::remove_file(&marker_path); if let Err(err) = handler .handle_message(ExecServerInboundMessage::Request( @@ -631,48 +736,41 @@ mod tests { { panic!("first exec should succeed: {err}"); } - assert_eq!( - recv_outbound(&mut outgoing_rx).await, - ExecServerOutboundMessage::Response { - request_id: RequestId::Integer(2), - response: ExecServerResponseMessage::Exec(crate::protocol::ExecResponse { - process_id: "proc-1".to_string(), - }), - } - ); + let ExecServerOutboundMessage::Response { + request_id, + response: ExecServerResponseMessage::Exec(first_exec), + } = recv_outbound(&mut outgoing_rx).await + else { + panic!("expected first exec response"); + }; + assert_eq!(request_id, RequestId::Integer(2)); + assert_eq!(first_exec.process_id, "proc-1"); if let Err(err) = handler .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { request_id: RequestId::Integer(3), params: crate::protocol::ExecParams { - argv: vec![ - "bash".to_string(), - "-lc".to_string(), - format!("printf duplicate > {}", marker_path.display()), - ], + process_id: "proc-2".to_string(), + argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], ..params }, })) .await { - panic!("duplicate exec should not fail the handler: {err}"); + panic!("second exec should succeed: {err}"); } - let ExecServerOutboundMessage::Error { request_id, error } = - recv_outbound(&mut outgoing_rx).await + let ExecServerOutboundMessage::Response { + request_id, + response: ExecServerResponseMessage::Exec(second_exec), + } = recv_outbound(&mut outgoing_rx).await else { - panic!("expected duplicate-process error"); + panic!("expected second exec response"); }; assert_eq!(request_id, RequestId::Integer(3)); - assert_eq!(error.code, -32600); - assert_eq!(error.message, "process proc-1 already exists"); - assert!( - !marker_path.exists(), - "duplicate process ids must be rejected before spawning the command" - ); + assert_eq!(second_exec.process_id, "proc-2"); handler.shutdown().await; - let _ = std::fs::remove_file(&marker_path); } #[tokio::test] @@ -707,7 +805,7 @@ mod tests { .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { request_id: RequestId::Integer(2), params: crate::protocol::ExecParams { - process_id: "proc-2".to_string(), + process_id: "proc-1".to_string(), argv: vec![ "bash".to_string(), "-lc".to_string(), @@ -723,14 +821,20 @@ mod tests { { panic!("exec should succeed: {err}"); } - let _ = recv_outbound(&mut outgoing_rx).await; + let ExecServerOutboundMessage::Response { + response: ExecServerResponseMessage::Exec(exec_response), + .. + } = recv_outbound(&mut outgoing_rx).await + else { + panic!("expected exec response"); + }; if let Err(err) = handler .handle_message(ExecServerInboundMessage::Request( ExecServerRequest::Write { request_id: RequestId::Integer(3), params: WriteParams { - process_id: "proc-2".to_string(), + process_id: exec_response.process_id, chunk: b"hello\n".to_vec().into(), }, }, @@ -747,7 +851,7 @@ mod tests { }; assert_eq!(request_id, RequestId::Integer(3)); assert_eq!(error.code, -32600); - assert_eq!(error.message, "stdin is closed for process proc-2"); + assert_eq!(error.message, "stdin is closed for process proc-1"); handler.shutdown().await; } @@ -859,7 +963,7 @@ mod tests { } #[tokio::test] - async fn terminate_keeps_process_ids_reserved_until_exit_cleanup() { + async fn terminate_keeps_process_ids_reserved() { let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); let mut handler = ExecServerHandler::new(outgoing_tx); @@ -902,6 +1006,10 @@ mod tests { super::RunningProcess { session: spawned.session, tty: false, + output: std::collections::VecDeque::new(), + retained_bytes: 0, + next_seq: 1, + exit_code: None, }, ); } diff --git a/codex-rs/exec-server/src/server/routing.rs b/codex-rs/exec-server/src/server/routing.rs index 7090a742b3..e70b371980 100644 --- a/codex-rs/exec-server/src/server/routing.rs +++ b/codex-rs/exec-server/src/server/routing.rs @@ -10,6 +10,7 @@ use serde::de::DeserializeOwned; 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; @@ -20,6 +21,8 @@ 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; @@ -41,6 +44,10 @@ pub(crate) enum ExecServerRequest { request_id: RequestId, params: ExecParams, }, + Read { + request_id: RequestId, + params: ReadParams, + }, Write { request_id: RequestId, params: WriteParams, @@ -73,6 +80,7 @@ pub(crate) enum ExecServerOutboundMessage { pub(crate) enum ExecServerResponseMessage { Initialize(InitializeResponse), Exec(ExecResponse), + Read(ReadResponse), Write(WriteResponse), Terminate(TerminateResponse), } @@ -161,6 +169,9 @@ fn route_request(request: JSONRPCRequest) -> Result Ok(parse_request_params(request, |request_id, params| { ExecServerRequest::Exec { request_id, params } })), + EXEC_READ_METHOD => Ok(parse_request_params(request, |request_id, params| { + ExecServerRequest::Read { request_id, params } + })), EXEC_WRITE_METHOD => Ok(parse_request_params(request, |request_id, params| { ExecServerRequest::Write { request_id, params } })), @@ -210,6 +221,7 @@ fn serialize_response( match response { ExecServerResponseMessage::Initialize(response) => serde_json::to_value(response), ExecServerResponseMessage::Exec(response) => serde_json::to_value(response), + ExecServerResponseMessage::Read(response) => serde_json::to_value(response), ExecServerResponseMessage::Write(response) => serde_json::to_value(response), ExecServerResponseMessage::Terminate(response) => serde_json::to_value(response), } @@ -421,12 +433,12 @@ mod tests { #[test] fn unexpected_client_notifications_are_rejected() { let err = route_jsonrpc_message(JSONRPCMessage::Notification(JSONRPCNotification { - method: "process/outputDelta".to_string(), + method: "process/output".to_string(), params: Some(json!({})), })) .expect_err("unexpected client notification should fail"); - assert_eq!(err, "unexpected notification method: process/outputDelta"); + assert_eq!(err, "unexpected notification method: process/output"); } #[test] diff --git a/codex-rs/exec-server/tests/stdio_smoke.rs b/codex-rs/exec-server/tests/stdio_smoke.rs index 9616294bf0..041c7a4b75 100644 --- a/codex-rs/exec-server/tests/stdio_smoke.rs +++ b/codex-rs/exec-server/tests/stdio_smoke.rs @@ -13,6 +13,7 @@ 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; @@ -93,10 +94,11 @@ async fn exec_server_client_streams_output_and_accepts_writes() -> anyhow::Resul ) .await?; - let process = server - .client() - .start_process(ExecParams { - process_id: "2001".to_string(), + 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(), @@ -109,29 +111,26 @@ async fn exec_server_client_streams_output_and_accepts_writes() -> anyhow::Resul arg0: None, }) .await?; + let process_id = response.process_id; - let mut output = process.output_receiver(); - let (stream, ready_output) = recv_until_contains(&mut output, "ready").await?; - assert_eq!(stream, ExecOutputStream::Stdout); + 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" ); - process - .writer_sender() - .send(b"hello\n".to_vec()) - .await - .expect("write should succeed"); + client.write(&process_id, b"hello\n".to_vec()).await?; - let (stream, echoed_output) = recv_until_contains(&mut output, "echo:hello").await?; - assert_eq!(stream, ExecOutputStream::Stdout); + 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" ); - process.terminate(); + client.terminate(&process_id).await?; Ok(()) } @@ -161,9 +160,10 @@ async fn exec_server_client_connects_over_websocket() -> anyhow::Result<()> { }) .await?; - let process = client - .start_process(ExecParams { - process_id: "2002".to_string(), + 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(), @@ -176,29 +176,26 @@ async fn exec_server_client_connects_over_websocket() -> anyhow::Result<()> { arg0: None, }) .await?; + let process_id = response.process_id; - let mut output = process.output_receiver(); - let (stream, ready_output) = recv_until_contains(&mut output, "ready").await?; - assert_eq!(stream, ExecOutputStream::Stdout); + 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" ); - process - .writer_sender() - .send(b"hello\n".to_vec()) - .await - .expect("write should succeed"); + client.write(&process_id, b"hello\n".to_vec()).await?; - let (stream, echoed_output) = recv_until_contains(&mut output, "echo:hello").await?; - assert_eq!(stream, ExecOutputStream::Stdout); + 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" ); - process.terminate(); + client.terminate(&process_id).await?; child.start_kill()?; Ok(()) } @@ -239,9 +236,9 @@ async fn websocket_disconnect_terminates_processes_for_that_connection() -> anyh }) .await?; - let _process = client - .start_process(ExecParams { - process_id: "2003".to_string(), + let _response = client + .exec(ExecParams { + process_id: "proc-1".to_string(), argv: vec![ "bash".to_string(), "-lc".to_string(), @@ -280,17 +277,22 @@ where } async fn recv_until_contains( - output: &mut broadcast::Receiver, + 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 output_event = timeout(remaining, output.recv()).await??; - collected.push_str(&String::from_utf8_lossy(&output_event.chunk)); - if collected.contains(needle) { - return Ok((output_event.stream, collected)); + 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)); + } } } }