diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 46961f5b42..2a72d1f311 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -109,7 +109,6 @@ Request params: "PATH": "/usr/bin:/bin" }, "tty": true, - "outputBytesCap": 16384, "arg0": null } ``` @@ -122,19 +121,13 @@ Field definitions: - `env`: environment variables passed to the child process. - `tty`: when `true`, spawn a PTY-backed interactive process; when `false`, spawn a pipe-backed process with closed stdin. -- `outputBytesCap`: maximum retained stdout/stderr bytes per stream for the - in-memory buffer. Defaults to `codex_utils_pty::DEFAULT_OUTPUT_BYTES_CAP`. - `arg0`: optional argv0 override forwarded to `codex-utils-pty`. Response: ```json { - "processId": "proc-1", - "running": true, - "exitCode": null, - "stdout": null, - "stderr": null + "processId": "proc-1" } ``` @@ -262,6 +255,7 @@ The crate exports: - `ExecServerClientConnectOptions` - `RemoteExecServerConnectArgs` - `ExecServerLaunchCommand` +- `ExecServerOutput` - `ExecServerProcess` - `SpawnedExecServer` - `ExecServerError` @@ -293,6 +287,18 @@ Connect the client to an existing server transport: - `ExecServerClient::connect_stdio(...)` - `ExecServerClient::connect_websocket(...)` +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: + +- `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 + Spawning a local child process is deliberately separate: - `spawn_local_exec_server(...)` @@ -310,8 +316,8 @@ 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,"outputBytesCap":4096,"arg0":null}} -{"id":2,"result":{"processId":"proc-1","running":true,"exitCode":null,"stdout":null,"stderr":null}} +{"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,"result":{"processId":"proc-1"}} {"method":"command/exec/outputDelta","params":{"processId":"proc-1","stream":"stdout","chunk":"cmVhZHkK"}} ``` diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 8aa780dd34..ea08a6cb96 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -4,6 +4,7 @@ use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; +use std::time::Duration; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCErrorError; @@ -22,6 +23,7 @@ use tokio::sync::broadcast; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::task::JoinHandle; +use tokio::time::timeout; use tokio_tungstenite::connect_async; use tracing::debug; use tracing::warn; @@ -49,12 +51,14 @@ use crate::protocol::WriteResponse; #[derive(Debug, Clone, PartialEq, Eq)] pub struct ExecServerClientConnectOptions { pub client_name: String, + pub initialize_timeout: Duration, } impl Default for ExecServerClientConnectOptions { fn default() -> Self { Self { client_name: "codex-core".to_string(), + initialize_timeout: INITIALIZE_TIMEOUT, } } } @@ -63,19 +67,42 @@ impl Default for ExecServerClientConnectOptions { pub struct RemoteExecServerConnectArgs { pub websocket_url: String, pub client_name: String, + pub connect_timeout: Duration, + pub initialize_timeout: Duration, } impl From for ExecServerClientConnectOptions { fn from(value: RemoteExecServerConnectArgs) -> Self { Self { client_name: value.client_name, + initialize_timeout: value.initialize_timeout, } } } +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); + +impl RemoteExecServerConnectArgs { + pub fn new(websocket_url: String, client_name: String) -> Self { + Self { + websocket_url, + client_name, + connect_timeout: CONNECT_TIMEOUT, + initialize_timeout: INITIALIZE_TIMEOUT, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecServerOutput { + pub stream: crate::protocol::ExecOutputStream, + pub chunk: Vec, +} + pub struct ExecServerProcess { process_id: String, - output_rx: broadcast::Receiver>, + output_rx: broadcast::Receiver, writer_tx: mpsc::Sender>, status: Arc, client: ExecServerClient, @@ -96,7 +123,7 @@ impl ExecServerProcess { self.writer_tx.clone() } - pub fn output_receiver(&self) -> broadcast::Receiver> { + pub fn output_receiver(&self) -> broadcast::Receiver { self.output_rx.resubscribe() } @@ -109,7 +136,6 @@ impl ExecServerProcess { } 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 { @@ -157,7 +183,7 @@ impl RemoteProcessStatus { } struct RegisteredProcess { - output_tx: broadcast::Sender>, + output_tx: broadcast::Sender, status: Arc, } @@ -184,12 +210,16 @@ pub struct ExecServerClient { pub enum ExecServerError { #[error("failed to spawn exec-server: {0}")] Spawn(#[source] std::io::Error), + #[error("timed out connecting to exec-server websocket `{url}` after {timeout:?}")] + WebSocketConnectTimeout { url: String, timeout: Duration }, #[error("failed to connect to exec-server websocket `{url}`: {source}")] WebSocketConnect { url: String, #[source] source: tokio_tungstenite::tungstenite::Error, }, + #[error("timed out waiting for exec-server initialize handshake after {timeout:?}")] + InitializeTimedOut { timeout: Duration }, #[error("exec-server transport closed")] Closed, #[error("failed to serialize or deserialize exec-server JSON: {0}")] @@ -221,8 +251,13 @@ impl ExecServerClient { args: RemoteExecServerConnectArgs, ) -> Result { let websocket_url = args.websocket_url.clone(); - let (stream, _) = connect_async(websocket_url.as_str()) + let connect_timeout = args.connect_timeout; + let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str())) .await + .map_err(|_| ExecServerError::WebSocketConnectTimeout { + url: websocket_url.clone(), + timeout: connect_timeout, + })? .map_err(|source| ExecServerError::WebSocketConnect { url: websocket_url.clone(), source, @@ -325,32 +360,12 @@ impl ExecServerClient { return Err(err); } }; - - if !response.running { - status.mark_exited(response.exit_code); - } - - if let Some(stdout) = response.stdout { - let _ = self - .inner - .processes - .lock() - .await - .get(&process_id) - .map(|process| process.output_tx.send(stdout.into_inner())); - } - if let Some(stderr) = response.stderr { - let _ = self - .inner - .processes - .lock() - .await - .get(&process_id) - .map(|process| process.output_tx.send(stderr.into_inner())); - } - - if let Some(exit_code) = response.exit_code { - status.mark_exited(Some(exit_code)); + 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 { @@ -366,16 +381,21 @@ impl ExecServerClient { &self, options: ExecServerClientConnectOptions, ) -> Result<(), ExecServerError> { - let _: InitializeResponse = self - .request( - INITIALIZE_METHOD, - &InitializeParams { - client_name: options.client_name, - }, - ) - .await?; - self.notify(INITIALIZED_METHOD, &serde_json::json!({})) - .await + let ExecServerClientConnectOptions { + client_name, + initialize_timeout, + } = options; + timeout(initialize_timeout, async { + let _: InitializeResponse = self + .request(INITIALIZE_METHOD, &InitializeParams { client_name }) + .await?; + self.notify(INITIALIZED_METHOD, &serde_json::json!({})) + .await + }) + .await + .map_err(|_| ExecServerError::InitializeTimedOut { + timeout: initialize_timeout, + })? } async fn write_process(&self, params: WriteParams) -> Result { @@ -412,6 +432,7 @@ impl ExecServerClient { 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)); let (response_tx, response_rx) = oneshot::channel(); @@ -421,7 +442,6 @@ impl ExecServerClient { .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(), @@ -482,10 +502,13 @@ async fn handle_server_notification( EXEC_OUTPUT_DELTA_METHOD => { let params: ExecOutputDeltaNotification = serde_json::from_value(notification.params.unwrap_or(Value::Null))?; - let chunk = params.chunk.into_inner(); + 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(chunk); + let _ = process.output_tx.send(output); } } EXEC_EXITED_METHOD => { @@ -543,6 +566,9 @@ mod tests { use super::ExecServerClientConnectOptions; use super::ExecServerError; use crate::protocol::EXEC_METHOD; + use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; + use crate::protocol::EXEC_TERMINATE_METHOD; + use crate::protocol::ExecOutputStream; use crate::protocol::ExecParams; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; @@ -554,6 +580,13 @@ mod tests { use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; + fn test_options() -> ExecServerClientConnectOptions { + ExecServerClientConnectOptions { + client_name: "test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + } + } + async fn read_jsonrpc_line(lines: &mut tokio::io::Lines>) -> JSONRPCMessage where R: tokio::io::AsyncRead + Unpin, @@ -625,14 +658,8 @@ mod tests { assert_eq!(params, Some(serde_json::json!({}))); }); - let client = ExecServerClient::connect_stdio( - client_stdin, - client_stdout, - ExecServerClientConnectOptions { - client_name: "test-client".to_string(), - }, - ) - .await; + let client = + ExecServerClient::connect_stdio(client_stdin, client_stdout, test_options()).await; if let Err(err) = client { panic!("failed to connect test client: {err}"); } @@ -668,14 +695,8 @@ mod tests { .await; }); - let result = ExecServerClient::connect_stdio( - client_stdin, - client_stdout, - ExecServerClientConnectOptions { - client_name: "test-client".to_string(), - }, - ) - .await; + let result = + ExecServerClient::connect_stdio(client_stdin, client_stdout, test_options()).await; match result { Err(ExecServerError::Server { code, message }) => { @@ -736,9 +757,7 @@ mod tests { let client = match ExecServerClient::connect_stdio( client_stdin, client_stdout, - ExecServerClientConnectOptions { - client_name: "test-client".to_string(), - }, + test_options(), ) .await { @@ -753,7 +772,6 @@ mod tests { cwd: std::env::current_dir().unwrap_or_else(|err| panic!("missing cwd: {err}")), env: HashMap::new(), tty: true, - output_bytes_cap: 4096, arg0: None, }) .await; @@ -772,4 +790,215 @@ mod tests { "failed requests should not leave registered process state behind" ); } + + #[tokio::test] + async fn connect_stdio_times_out_during_initialize_handshake() { + let (client_stdin, server_reader) = tokio::io::duplex(4096); + let (_server_writer, client_stdout) = tokio::io::duplex(4096); + + tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let _ = read_jsonrpc_line(&mut lines).await; + tokio::time::sleep(Duration::from_millis(200)).await; + }); + + let result = ExecServerClient::connect_stdio( + client_stdin, + client_stdout, + ExecServerClientConnectOptions { + client_name: "test-client".to_string(), + initialize_timeout: Duration::from_millis(25), + }, + ) + .await; + + match result { + Err(ExecServerError::InitializeTimedOut { timeout }) => { + assert_eq!(timeout, Duration::from_millis(25)); + } + Err(err) => panic!("unexpected initialize timeout failure: {err}"), + Ok(_) => panic!("expected initialize timeout"), + } + } + + #[tokio::test] + async fn start_process_preserves_output_stream_metadata() { + let (client_stdin, server_reader) = tokio::io::duplex(4096); + let (mut server_writer, client_stdout) = tokio::io::duplex(4096); + + tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + + let initialize = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Request(initialize_request) = initialize else { + panic!("expected initialize request"); + }; + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: initialize_request.id, + result: serde_json::json!({ "protocolVersion": PROTOCOL_VERSION }), + }), + ) + .await; + + let initialized = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Notification(notification) = initialized else { + panic!("expected initialized notification"); + }; + assert_eq!(notification.method, INITIALIZED_METHOD); + + let exec_request = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Request(JSONRPCRequest { id, method, .. }) = exec_request else { + panic!("expected exec request"); + }; + assert_eq!(method, EXEC_METHOD); + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id, + result: serde_json::json!({ "processId": "proc-1" }), + }), + ) + .await; + tokio::time::sleep(Duration::from_millis(25)).await; + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Notification(JSONRPCNotification { + method: EXEC_OUTPUT_DELTA_METHOD.to_string(), + params: Some(serde_json::json!({ + "processId": "proc-1", + "stream": "stderr", + "chunk": "ZXJyb3IK" + })), + }), + ) + .await; + tokio::time::sleep(Duration::from_millis(100)).await; + }); + + let client = match ExecServerClient::connect_stdio( + client_stdin, + client_stdout, + test_options(), + ) + .await + { + Ok(client) => client, + Err(err) => panic!("failed to connect test client: {err}"), + }; + + let process = match 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 + { + Ok(process) => process, + Err(err) => panic!("failed to start process: {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.stream, ExecOutputStream::Stderr); + assert_eq!(output.chunk, b"error\n".to_vec()); + } + + #[tokio::test] + async fn terminate_does_not_mark_process_exited_before_exit_notification() { + let (client_stdin, server_reader) = tokio::io::duplex(4096); + let (mut server_writer, client_stdout) = tokio::io::duplex(4096); + + tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + + let initialize = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Request(initialize_request) = initialize else { + panic!("expected initialize request"); + }; + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: initialize_request.id, + result: serde_json::json!({ "protocolVersion": PROTOCOL_VERSION }), + }), + ) + .await; + + let initialized = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Notification(notification) = initialized else { + panic!("expected initialized notification"); + }; + assert_eq!(notification.method, INITIALIZED_METHOD); + + let exec_request = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Request(JSONRPCRequest { id, method, .. }) = exec_request else { + panic!("expected exec request"); + }; + assert_eq!(method, EXEC_METHOD); + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id, + result: serde_json::json!({ "processId": "proc-1" }), + }), + ) + .await; + + let terminate_request = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Request(JSONRPCRequest { id, method, .. }) = terminate_request + else { + panic!("expected terminate request"); + }; + assert_eq!(method, EXEC_TERMINATE_METHOD); + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id, + result: serde_json::json!({ "running": true }), + }), + ) + .await; + tokio::time::sleep(Duration::from_millis(100)).await; + }); + + let client = match ExecServerClient::connect_stdio( + client_stdin, + client_stdout, + test_options(), + ) + .await + { + Ok(client) => client, + Err(err) => panic!("failed to connect test client: {err}"), + }; + + let process = match 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 + { + Ok(process) => process, + Err(err) => panic!("failed to start process: {err}"), + }; + + process.terminate(); + tokio::time::sleep(Duration::from_millis(25)).await; + assert!(!process.has_exited(), "terminate should not imply exit"); + assert_eq!(process.exit_code(), None); + } } diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 4c975aa5e2..714e044e56 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -7,6 +7,7 @@ mod server; pub use client::ExecServerClient; pub use client::ExecServerClientConnectOptions; pub use client::ExecServerError; +pub use client::ExecServerOutput; pub use client::ExecServerProcess; pub use client::RemoteExecServerConnectArgs; pub use local::ExecServerLaunchCommand; diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs index dcec9a6d6a..452c6f2894 100644 --- a/codex-rs/exec-server/src/protocol.rs +++ b/codex-rs/exec-server/src/protocol.rs @@ -2,7 +2,6 @@ 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; @@ -51,8 +50,6 @@ pub struct ExecParams { pub cwd: PathBuf, pub env: HashMap, pub tty: bool, - #[serde(default = "default_output_bytes_cap")] - pub output_bytes_cap: usize, pub arg0: Option, } @@ -60,10 +57,6 @@ pub struct ExecParams { #[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)] @@ -113,10 +106,6 @@ pub struct ExecExitedNotification { 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 _; diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index a7d1dcbcb3..fa0081f5c9 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,7 +1,5 @@ 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; @@ -41,39 +39,6 @@ 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(crate) async fn run_connection(connection: JsonRpcConnection) { @@ -252,13 +217,6 @@ impl ExecServerConnectionProcessor { } .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 = self.processes.lock().await; @@ -273,8 +231,6 @@ impl ExecServerConnectionProcessor { RunningProcess { session: spawned.session, tty: params.tty, - stdout_buffer: Arc::clone(&stdout_buffer), - stderr_buffer: Arc::clone(&stderr_buffer), }, ); } @@ -284,14 +240,12 @@ impl ExecServerConnectionProcessor { ExecOutputStream::Stdout, spawned.stdout_rx, self.outgoing_tx.clone(), - Arc::clone(&stdout_buffer), )); tokio::spawn(stream_output( process_id.clone(), ExecOutputStream::Stderr, spawned.stderr_rx, self.outgoing_tx.clone(), - Arc::clone(&stderr_buffer), )); tokio::spawn(watch_exit( process_id.clone(), @@ -300,13 +254,7 @@ impl ExecServerConnectionProcessor { Arc::clone(&self.processes), )); - json_value(ExecResponse { - process_id, - running: true, - exit_code: None, - stdout: None, - stderr: None, - }) + json_value(ExecResponse { process_id }) } async fn handle_write_request( @@ -379,12 +327,8 @@ async fn stream_output( stream: ExecOutputStream, mut receiver: tokio::sync::mpsc::Receiver>, outgoing_tx: mpsc::Sender, - 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, @@ -406,13 +350,9 @@ async fn watch_exit( 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()); + processes.remove(&process_id); } let _ = send_notification( &outgoing_tx, diff --git a/codex-rs/exec-server/tests/stdio_smoke.rs b/codex-rs/exec-server/tests/stdio_smoke.rs index 0fd464ace0..da4b8f007c 100644 --- a/codex-rs/exec-server/tests/stdio_smoke.rs +++ b/codex-rs/exec-server/tests/stdio_smoke.rs @@ -9,6 +9,7 @@ use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; +use codex_exec_server::ExecOutputStream; use codex_exec_server::ExecParams; use codex_exec_server::ExecServerClient; use codex_exec_server::ExecServerClientConnectOptions; @@ -87,6 +88,7 @@ async fn exec_server_client_streams_output_and_accepts_writes() -> anyhow::Resul }, ExecServerClientConnectOptions { client_name: "exec-server-test".to_string(), + initialize_timeout: Duration::from_secs(5), }, ) .await?; @@ -104,16 +106,15 @@ async fn exec_server_client_streams_output_and_accepts_writes() -> anyhow::Resul cwd: std::env::current_dir()?, env, tty: true, - output_bytes_cap: 4096, arg0: None, }) .await?; let mut output = process.output_receiver(); + let (stream, ready_output) = recv_until_contains(&mut output, "ready").await?; + assert_eq!(stream, ExecOutputStream::Stdout); assert!( - recv_until_contains(&mut output, "ready") - .await? - .contains("ready"), + ready_output.contains("ready"), "expected initial ready output" ); @@ -123,10 +124,10 @@ async fn exec_server_client_streams_output_and_accepts_writes() -> anyhow::Resul .await .expect("write should succeed"); + let (stream, echoed_output) = recv_until_contains(&mut output, "echo:hello").await?; + assert_eq!(stream, ExecOutputStream::Stdout); assert!( - recv_until_contains(&mut output, "echo:hello") - .await? - .contains("echo:hello"), + echoed_output.contains("echo:hello"), "expected echoed output" ); @@ -155,6 +156,8 @@ async fn exec_server_client_connects_over_websocket() -> anyhow::Result<()> { let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { websocket_url, client_name: "exec-server-test".to_string(), + connect_timeout: Duration::from_secs(5), + initialize_timeout: Duration::from_secs(5), }) .await?; @@ -170,16 +173,15 @@ async fn exec_server_client_connects_over_websocket() -> anyhow::Result<()> { cwd: std::env::current_dir()?, env, tty: true, - output_bytes_cap: 4096, arg0: None, }) .await?; let mut output = process.output_receiver(); + let (stream, ready_output) = recv_until_contains(&mut output, "ready").await?; + assert_eq!(stream, ExecOutputStream::Stdout); assert!( - recv_until_contains(&mut output, "ready") - .await? - .contains("ready"), + ready_output.contains("ready"), "expected initial ready output" ); @@ -189,10 +191,10 @@ async fn exec_server_client_connects_over_websocket() -> anyhow::Result<()> { .await .expect("write should succeed"); + let (stream, echoed_output) = recv_until_contains(&mut output, "echo:hello").await?; + assert_eq!(stream, ExecOutputStream::Stdout); assert!( - recv_until_contains(&mut output, "echo:hello") - .await? - .contains("echo:hello"), + echoed_output.contains("echo:hello"), "expected echoed output" ); @@ -215,17 +217,17 @@ where } async fn recv_until_contains( - output: &mut broadcast::Receiver>, + output: &mut broadcast::Receiver, needle: &str, -) -> anyhow::Result { +) -> 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 chunk = timeout(remaining, output.recv()).await??; - collected.push_str(&String::from_utf8_lossy(&chunk)); + let output_event = timeout(remaining, output.recv()).await??; + collected.push_str(&String::from_utf8_lossy(&output_event.chunk)); if collected.contains(needle) { - return Ok(collected); + return Ok((output_event.stream, collected)); } } }