diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index f04dc3d5f8..ff58c2a7c5 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -121,7 +121,8 @@ Request params: "PATH": "/usr/bin:/bin" }, "tty": true, - "arg0": null + "arg0": null, + "sandbox": null } ``` @@ -133,6 +134,9 @@ Field definitions: - `tty`: when `true`, spawn a PTY-backed interactive process; when `false`, spawn a pipe-backed process with closed stdin. - `arg0`: optional argv0 override forwarded to `codex-utils-pty`. +- `sandbox`: optional sandbox config. Omit it for the current direct-spawn + behavior. Explicit `{"mode":"none"}` is accepted; `{"mode":"hostDefault"}` + is currently rejected until host-local sandbox materialization is wired up. Response: diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 8586461ae2..b9005feeb5 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -224,11 +224,23 @@ struct Inner { pending: Mutex>, events_tx: broadcast::Sender, next_request_id: AtomicI64, + transport_tasks: Vec>, reader_task: JoinHandle<()>, } impl Drop for Inner { fn drop(&mut self) { + if let ClientBackend::InProcess { handler } = &self.backend + && let Ok(handle) = tokio::runtime::Handle::try_current() + { + let handler = Arc::clone(handler); + handle.spawn(async move { + handler.lock().await.shutdown().await; + }); + } + for task in &self.transport_tasks { + task.abort(); + } self.reader_task.abort(); } } @@ -294,6 +306,7 @@ impl ExecServerClient { pending: Mutex::new(HashMap::new()), events_tx: broadcast::channel(256).0, next_request_id: AtomicI64::new(1), + transport_tasks: Vec::new(), reader_task, } }); @@ -349,7 +362,7 @@ impl ExecServerClient { connection: JsonRpcConnection, options: ExecServerClientConnectOptions, ) -> Result { - let (write_tx, mut incoming_rx) = connection.into_parts(); + let (write_tx, mut incoming_rx, transport_tasks) = connection.into_parts(); let inner = Arc::new_cyclic(|weak| { let weak = weak.clone(); let reader_task = tokio::spawn(async move { @@ -386,6 +399,7 @@ impl ExecServerClient { pending: Mutex::new(HashMap::new()), events_tx: broadcast::channel(256).0, next_request_id: AtomicI64::new(1), + transport_tasks, reader_task, } }); @@ -786,860 +800,4 @@ async fn handle_transport_shutdown(inner: &Arc) { } #[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::time::Duration; - - use pretty_assertions::assert_eq; - use tokio::io::AsyncBufReadExt; - use tokio::io::AsyncWriteExt; - use tokio::io::BufReader; - use tokio::time::timeout; - - 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; - use crate::protocol::ExecOutputStream; - use crate::protocol::ExecParams; - 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; - use codex_app_server_protocol::JSONRPCNotification; - 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, - { - let next_line = timeout(Duration::from_secs(1), lines.next_line()).await; - let line_result = match next_line { - Ok(line_result) => line_result, - Err(err) => panic!("timed out waiting for JSON-RPC line: {err}"), - }; - let maybe_line = match line_result { - Ok(maybe_line) => maybe_line, - Err(err) => panic!("failed to read JSON-RPC line: {err}"), - }; - let line = match maybe_line { - Some(line) => line, - None => panic!("server connection closed before JSON-RPC line arrived"), - }; - match serde_json::from_str::(&line) { - Ok(message) => message, - Err(err) => panic!("failed to parse JSON-RPC line: {err}"), - } - } - - async fn write_jsonrpc_line(writer: &mut W, message: JSONRPCMessage) - where - W: tokio::io::AsyncWrite + Unpin, - { - let encoded = match serde_json::to_string(&message) { - Ok(encoded) => encoded, - Err(err) => panic!("failed to encode JSON-RPC message: {err}"), - }; - if let Err(err) = writer.write_all(format!("{encoded}\n").as_bytes()).await { - panic!("failed to write JSON-RPC line: {err}"); - } - } - - #[tokio::test] - async fn connect_stdio_performs_initialize_handshake() { - let (client_stdin, server_reader) = tokio::io::duplex(4096); - let (mut server_writer, client_stdout) = tokio::io::duplex(4096); - - let server = tokio::spawn(async move { - let mut lines = BufReader::new(server_reader).lines(); - - let initialize = read_jsonrpc_line(&mut lines).await; - let JSONRPCMessage::Request(request) = initialize else { - panic!("expected initialize request"); - }; - assert_eq!(request.method, INITIALIZE_METHOD); - assert_eq!( - request.params, - Some(serde_json::json!({ "clientName": "test-client" })) - ); - write_jsonrpc_line( - &mut server_writer, - JSONRPCMessage::Response(JSONRPCResponse { - id: request.id, - result: serde_json::json!({ "protocolVersion": PROTOCOL_VERSION }), - }), - ) - .await; - - let initialized = read_jsonrpc_line(&mut lines).await; - let JSONRPCMessage::Notification(JSONRPCNotification { method, params }) = initialized - else { - panic!("expected initialized notification"); - }; - assert_eq!(method, INITIALIZED_METHOD); - assert_eq!(params, Some(serde_json::json!({}))); - }); - - let client = - ExecServerClient::connect_stdio(client_stdin, client_stdout, test_options()).await; - if let Err(err) = client { - panic!("failed to connect test client: {err}"); - } - - if let Err(err) = server.await { - panic!("server task failed: {err}"); - } - } - - #[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); - 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(request) = initialize else { - panic!("expected initialize request"); - }; - write_jsonrpc_line( - &mut server_writer, - JSONRPCMessage::Error(JSONRPCError { - id: request.id, - error: JSONRPCErrorError { - code: -32600, - message: "rejected".to_string(), - data: None, - }, - }), - ) - .await; - }); - - let result = - ExecServerClient::connect_stdio(client_stdin, client_stdout, test_options()).await; - - match result { - Err(ExecServerError::Server { code, message }) => { - assert_eq!(code, -32600); - assert_eq!(message, "rejected"); - } - Err(err) => panic!("unexpected initialize failure: {err}"), - Ok(_) => panic!("expected initialize failure"), - } - } - - #[tokio::test] - async fn start_process_cleans_up_registered_process_after_request_error() { - 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::Error(JSONRPCError { - id, - error: JSONRPCErrorError { - code: -32600, - message: "duplicate process".to_string(), - data: None, - }, - }), - ) - .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 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 result { - Err(ExecServerError::Server { code, message }) => { - assert_eq!(code, -32600); - assert_eq!(message, "duplicate process"); - } - Err(err) => panic!("unexpected start_process failure: {err}"), - Ok(_) => panic!("expected start_process failure"), - } - - assert!( - client.inner.pending.lock().await.is_empty(), - "failed requests should not leave pending request 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); - } - - #[tokio::test] - 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); - - 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": "other-proc" }), - }), - ) - .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}"), - }; - - assert_eq!(process.process_id, "other-proc"); - } - - #[tokio::test] - 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); - - 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": "stdout", - "chunk": "YWxpdmUK" - })), - }), - ) - .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 first_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 first process: {err}"), - }; - - let mut output = first_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::Stdout); - assert_eq!(output.chunk, b"alive\n".to_vec()); - } - - #[tokio::test] - async fn transport_shutdown_marks_processes_exited_without_exit_codes() { - 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; - drop(server_writer); - }); - - 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 _ = process; - } -} +mod tests; diff --git a/codex-rs/exec-server/src/client/tests.rs b/codex-rs/exec-server/src/client/tests.rs new file mode 100644 index 0000000000..8e12a6a525 --- /dev/null +++ b/codex-rs/exec-server/src/client/tests.rs @@ -0,0 +1,876 @@ +use std::collections::HashMap; +use std::time::Duration; + +use pretty_assertions::assert_eq; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::time::timeout; + +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; +use crate::protocol::ExecOutputStream; +use crate::protocol::ExecParams; +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; +use codex_app_server_protocol::JSONRPCNotification; +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, +{ + let next_line = timeout(Duration::from_secs(1), lines.next_line()).await; + let line_result = match next_line { + Ok(line_result) => line_result, + Err(err) => panic!("timed out waiting for JSON-RPC line: {err}"), + }; + let maybe_line = match line_result { + Ok(maybe_line) => maybe_line, + Err(err) => panic!("failed to read JSON-RPC line: {err}"), + }; + let line = match maybe_line { + Some(line) => line, + None => panic!("server connection closed before JSON-RPC line arrived"), + }; + match serde_json::from_str::(&line) { + Ok(message) => message, + Err(err) => panic!("failed to parse JSON-RPC line: {err}"), + } +} + +async fn write_jsonrpc_line(writer: &mut W, message: JSONRPCMessage) +where + W: tokio::io::AsyncWrite + Unpin, +{ + let encoded = match serde_json::to_string(&message) { + Ok(encoded) => encoded, + Err(err) => panic!("failed to encode JSON-RPC message: {err}"), + }; + if let Err(err) = writer.write_all(format!("{encoded}\n").as_bytes()).await { + panic!("failed to write JSON-RPC line: {err}"); + } +} + +#[tokio::test] +async fn connect_stdio_performs_initialize_handshake() { + let (client_stdin, server_reader) = tokio::io::duplex(4096); + let (mut server_writer, client_stdout) = tokio::io::duplex(4096); + + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + + let initialize = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Request(request) = initialize else { + panic!("expected initialize request"); + }; + assert_eq!(request.method, INITIALIZE_METHOD); + assert_eq!( + request.params, + Some(serde_json::json!({ "clientName": "test-client" })) + ); + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::json!({ "protocolVersion": PROTOCOL_VERSION }), + }), + ) + .await; + + let initialized = read_jsonrpc_line(&mut lines).await; + let JSONRPCMessage::Notification(JSONRPCNotification { method, params }) = initialized + else { + panic!("expected initialized notification"); + }; + assert_eq!(method, INITIALIZED_METHOD); + assert_eq!(params, Some(serde_json::json!({}))); + }); + + let client = ExecServerClient::connect_stdio(client_stdin, client_stdout, test_options()).await; + if let Err(err) = client { + panic!("failed to connect test client: {err}"); + } + + if let Err(err) = server.await { + panic!("server task failed: {err}"); + } +} + +#[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, + sandbox: 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, + sandbox: 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, + sandbox: 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, + sandbox: 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 dropping_in_process_client_terminates_running_processes() { + let marker_path = std::env::temp_dir().join(format!( + "codex-exec-server-inprocess-drop-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system time") + .as_nanos() + )); + let _ = std::fs::remove_file(&marker_path); + + { + let client = match ExecServerClient::connect_in_process(test_options()).await { + Ok(client) => client, + Err(err) => panic!("failed to connect in-process client: {err}"), + }; + + let _ = client + .exec(ExecParams { + process_id: "proc-1".to_string(), + argv: vec![ + "/bin/sh".to_string(), + "-c".to_string(), + format!("sleep 2; printf dropped > {}", marker_path.display()), + ], + cwd: std::env::current_dir().expect("cwd"), + env: HashMap::new(), + tty: false, + arg0: None, + sandbox: None, + }) + .await + .unwrap_or_else(|err| panic!("failed to start in-process child: {err}")); + } + + tokio::time::sleep(Duration::from_secs(3)).await; + assert!( + !marker_path.exists(), + "dropping the in-process client should terminate managed children" + ); + let _ = std::fs::remove_file(&marker_path); +} + +#[tokio::test] +async fn connect_stdio_returns_initialize_errors() { + 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(request) = initialize else { + panic!("expected initialize request"); + }; + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Error(JSONRPCError { + id: request.id, + error: JSONRPCErrorError { + code: -32600, + message: "rejected".to_string(), + data: None, + }, + }), + ) + .await; + }); + + let result = ExecServerClient::connect_stdio(client_stdin, client_stdout, test_options()).await; + + match result { + Err(ExecServerError::Server { code, message }) => { + assert_eq!(code, -32600); + assert_eq!(message, "rejected"); + } + Err(err) => panic!("unexpected initialize failure: {err}"), + Ok(_) => panic!("expected initialize failure"), + } +} + +#[tokio::test] +async fn start_process_cleans_up_registered_process_after_request_error() { + 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::Error(JSONRPCError { + id, + error: JSONRPCErrorError { + code: -32600, + message: "duplicate process".to_string(), + data: None, + }, + }), + ) + .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 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, + sandbox: None, + }) + .await; + + match result { + Err(ExecServerError::Server { code, message }) => { + assert_eq!(code, -32600); + assert_eq!(message, "duplicate process"); + } + Err(err) => panic!("unexpected start_process failure: {err}"), + Ok(_) => panic!("expected start_process failure"), + } + + assert!( + client.inner.pending.lock().await.is_empty(), + "failed requests should not leave pending request 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, + sandbox: 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, + sandbox: 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); +} + +#[tokio::test] +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); + + 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": "other-proc" }), + }), + ) + .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, + sandbox: None, + }) + .await + { + Ok(process) => process, + Err(err) => panic!("failed to start process: {err}"), + }; + + assert_eq!(process.process_id, "other-proc"); +} + +#[tokio::test] +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); + + 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": "stdout", + "chunk": "YWxpdmUK" + })), + }), + ) + .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 first_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, + sandbox: None, + }) + .await + { + Ok(process) => process, + Err(err) => panic!("failed to start first process: {err}"), + }; + + let mut output = first_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::Stdout); + assert_eq!(output.chunk, b"alive\n".to_vec()); +} + +#[tokio::test] +async fn transport_shutdown_marks_processes_exited_without_exit_codes() { + 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; + drop(server_writer); + }); + + 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, + sandbox: None, + }) + .await + { + Ok(process) => process, + Err(err) => panic!("failed to start process: {err}"), + }; + + let _ = process; +} diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 691a3dce72..f8cce422bf 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -22,6 +22,7 @@ pub(crate) enum JsonRpcConnectionEvent { pub(crate) struct JsonRpcConnection { outgoing_tx: mpsc::Sender, incoming_rx: mpsc::Receiver, + task_handles: Vec>, } impl JsonRpcConnection { @@ -35,7 +36,7 @@ impl JsonRpcConnection { let reader_label = connection_label.clone(); let incoming_tx_for_reader = incoming_tx.clone(); - tokio::spawn(async move { + let reader_task = tokio::spawn(async move { let mut lines = BufReader::new(reader).lines(); loop { match lines.next_line().await { @@ -66,7 +67,7 @@ impl JsonRpcConnection { } } Ok(None) => { - send_disconnected(&incoming_tx_for_reader, None).await; + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; break; } Err(err) => { @@ -83,7 +84,7 @@ impl JsonRpcConnection { } }); - tokio::spawn(async move { + let writer_task = tokio::spawn(async move { let mut writer = BufWriter::new(writer); while let Some(message) = outgoing_rx.recv().await { if let Err(err) = write_jsonrpc_line_message(&mut writer, &message).await { @@ -102,6 +103,7 @@ impl JsonRpcConnection { Self { outgoing_tx, incoming_rx, + task_handles: vec![reader_task, writer_task], } } @@ -115,7 +117,7 @@ impl JsonRpcConnection { let reader_label = connection_label.clone(); let incoming_tx_for_reader = incoming_tx.clone(); - tokio::spawn(async move { + let reader_task = tokio::spawn(async move { loop { match websocket_reader.next().await { Some(Ok(Message::Text(text))) => { @@ -165,7 +167,7 @@ impl JsonRpcConnection { } } Some(Ok(Message::Close(_))) => { - send_disconnected(&incoming_tx_for_reader, None).await; + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; break; } Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {} @@ -181,14 +183,14 @@ impl JsonRpcConnection { break; } None => { - send_disconnected(&incoming_tx_for_reader, None).await; + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; break; } } } }); - tokio::spawn(async move { + let writer_task = tokio::spawn(async move { while let Some(message) = outgoing_rx.recv().await { match serialize_jsonrpc_message(&message) { Ok(encoded) => { @@ -221,6 +223,7 @@ impl JsonRpcConnection { Self { outgoing_tx, incoming_rx, + task_handles: vec![reader_task, writer_task], } } @@ -229,8 +232,9 @@ impl JsonRpcConnection { ) -> ( mpsc::Sender, mpsc::Receiver, + Vec>, ) { - (self.outgoing_tx, self.incoming_rx) + (self.outgoing_tx, self.incoming_rx, self.task_handles) } } @@ -323,7 +327,7 @@ mod tests { let (connection_writer, reader_from_connection) = tokio::io::duplex(1024); let connection = JsonRpcConnection::from_stdio(connection_reader, connection_writer, "test".to_string()); - let (outgoing_tx, mut incoming_rx) = connection.into_parts(); + let (outgoing_tx, mut incoming_rx, _task_handles) = connection.into_parts(); let incoming_message = JSONRPCMessage::Request(JSONRPCRequest { id: RequestId::Integer(7), @@ -371,7 +375,7 @@ mod tests { let (connection_writer, _reader_from_connection) = tokio::io::duplex(1024); let connection = JsonRpcConnection::from_stdio(connection_reader, connection_writer, "test".to_string()); - let (_outgoing_tx, mut incoming_rx) = connection.into_parts(); + let (_outgoing_tx, mut incoming_rx, _task_handles) = connection.into_parts(); if let Err(err) = writer_to_connection.write_all(b"not-json\n").await { panic!("failed to write invalid JSON: {err}"); @@ -401,7 +405,7 @@ mod tests { let (connection_writer, _reader_from_connection) = tokio::io::duplex(1024); let connection = JsonRpcConnection::from_stdio(connection_reader, connection_writer, "test".to_string()); - let (_outgoing_tx, mut incoming_rx) = connection.into_parts(); + let (_outgoing_tx, mut incoming_rx, _task_handles) = connection.into_parts(); drop(writer_to_connection); let event = recv_event(&mut incoming_rx).await; diff --git a/codex-rs/exec-server/src/protocol.rs b/codex-rs/exec-server/src/protocol.rs index 7ed8e20ae4..20b5593450 100644 --- a/codex-rs/exec-server/src/protocol.rs +++ b/codex-rs/exec-server/src/protocol.rs @@ -54,6 +54,21 @@ pub struct ExecParams { pub env: HashMap, pub tty: bool, pub arg0: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sandbox: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ExecSandboxConfig { + pub mode: ExecSandboxMode, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum ExecSandboxMode { + None, + HostDefault, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index b3171a9c8f..b45bd4538e 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -6,6 +6,7 @@ use std::time::Duration; use codex_utils_pty::ExecCommandSession; use codex_utils_pty::TerminalSize; use tokio::sync::Mutex; +use tokio::sync::Notify; use tokio::sync::mpsc; use tracing::warn; @@ -13,6 +14,7 @@ use crate::protocol::ExecExitedNotification; use crate::protocol::ExecOutputDeltaNotification; use crate::protocol::ExecOutputStream; use crate::protocol::ExecResponse; +use crate::protocol::ExecSandboxMode; use crate::protocol::InitializeResponse; use crate::protocol::PROTOCOL_VERSION; use crate::protocol::ProcessOutputChunk; @@ -41,6 +43,7 @@ struct RunningProcess { retained_bytes: usize, next_seq: u64, exit_code: Option, + output_notify: Arc, } pub(crate) struct ExecServerHandler { @@ -129,6 +132,15 @@ impl ExecServerHandler { } } + if matches!( + params.sandbox.as_ref().map(|sandbox| sandbox.mode), + Some(ExecSandboxMode::HostDefault) + ) { + return Err(invalid_request( + "sandbox mode `hostDefault` is not supported by exec-server yet".to_string(), + )); + } + let (program, args) = params .argv .split_first() @@ -156,6 +168,7 @@ impl ExecServerHandler { } .map_err(|err| internal_error(err.to_string()))?; + let output_notify = Arc::new(Notify::new()); { let mut process_map = self.processes.lock().await; process_map.insert( @@ -167,6 +180,7 @@ impl ExecServerHandler { retained_bytes: 0, next_seq: 1, exit_code: None, + output_notify: Arc::clone(&output_notify), }, ); } @@ -181,6 +195,7 @@ impl ExecServerHandler { spawned.stdout_rx, self.outbound_tx.clone(), Arc::clone(&self.processes), + Arc::clone(&output_notify), )); tokio::spawn(stream_output( process_id.clone(), @@ -192,12 +207,14 @@ impl ExecServerHandler { spawned.stderr_rx, self.outbound_tx.clone(), Arc::clone(&self.processes), + Arc::clone(&output_notify), )); tokio::spawn(watch_exit( process_id.clone(), spawned.exit_rx, self.outbound_tx.clone(), Arc::clone(&self.processes), + output_notify, )); Ok(ExecResponse { process_id }) @@ -214,7 +231,7 @@ impl ExecServerHandler { let deadline = tokio::time::Instant::now() + wait; loop { - let response = { + let (response, output_notify) = { let process_map = self.processes.lock().await; let process = process_map.get(¶ms.process_id).ok_or_else(|| { invalid_request(format!("unknown process id {}", params.process_id)) @@ -240,12 +257,15 @@ impl ExecServerHandler { } } - ReadResponse { - chunks, - next_seq, - exited: process.exit_code.is_some(), - exit_code: process.exit_code, - } + ( + ReadResponse { + chunks, + next_seq, + exited: process.exit_code.is_some(), + exit_code: process.exit_code, + }, + Arc::clone(&process.output_notify), + ) }; if !response.chunks.is_empty() @@ -255,7 +275,11 @@ impl ExecServerHandler { return Ok(response); } - tokio::time::sleep(Duration::from_millis(10)).await; + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + return Ok(response); + } + let _ = tokio::time::timeout(remaining, output_notify.notified()).await; } } @@ -397,6 +421,7 @@ async fn stream_output( mut receiver: tokio::sync::mpsc::Receiver>, outbound_tx: mpsc::Sender, processes: Arc>>, + output_notify: Arc, ) { while let Some(chunk) = receiver.recv().await { let notification = { @@ -427,6 +452,7 @@ async fn stream_output( chunk: chunk.into(), } }; + output_notify.notify_waiters(); if outbound_tx .send(ExecServerOutboundMessage::Notification( @@ -445,6 +471,7 @@ async fn watch_exit( exit_rx: tokio::sync::oneshot::Receiver, outbound_tx: mpsc::Sender, processes: Arc>>, + output_notify: Arc, ) { let exit_code = exit_rx.await.unwrap_or(-1); { @@ -453,6 +480,7 @@ async fn watch_exit( process.exit_code = Some(exit_code); } } + output_notify.notify_waiters(); let _ = outbound_tx .send(ExecServerOutboundMessage::Notification( ExecServerServerNotification::Exited(ExecExitedNotification { @@ -464,660 +492,4 @@ async fn watch_exit( } #[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::collections::VecDeque; - use std::time::Duration; - - use pretty_assertions::assert_eq; - use tokio::time::timeout; - - use super::ExecServerHandler; - use super::RetainedOutputChunk; - use super::RunningProcess; - use crate::protocol::ExecOutputStream; - use crate::protocol::InitializeParams; - use crate::protocol::InitializeResponse; - use crate::protocol::PROTOCOL_VERSION; - use crate::protocol::ReadParams; - use crate::protocol::TerminateResponse; - use crate::protocol::WriteParams; - use crate::server::routing::ExecServerClientNotification; - use crate::server::routing::ExecServerInboundMessage; - use crate::server::routing::ExecServerOutboundMessage; - use crate::server::routing::ExecServerRequest; - use crate::server::routing::ExecServerResponseMessage; - use codex_app_server_protocol::RequestId; - - async fn recv_outbound( - outgoing_rx: &mut tokio::sync::mpsc::Receiver, - ) -> ExecServerOutboundMessage { - let recv_result = timeout(Duration::from_secs(1), outgoing_rx.recv()).await; - let maybe_message = match recv_result { - Ok(maybe_message) => maybe_message, - Err(err) => panic!("timed out waiting for handler output: {err}"), - }; - match maybe_message { - Some(message) => message, - None => panic!("handler output channel closed unexpectedly"), - } - } - - #[tokio::test] - async fn initialize_response_reports_protocol_version() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(1); - let mut handler = ExecServerHandler::new(outgoing_tx); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("initialize should succeed: {err}"); - } - - assert_eq!( - recv_outbound(&mut outgoing_rx).await, - ExecServerOutboundMessage::Response { - request_id: RequestId::Integer(1), - response: ExecServerResponseMessage::Initialize(InitializeResponse { - protocol_version: PROTOCOL_VERSION.to_string(), - }), - } - ); - } - - #[tokio::test] - async fn exec_methods_require_initialize() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(1); - let mut handler = ExecServerHandler::new(outgoing_tx); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { - request_id: RequestId::Integer(7), - params: crate::protocol::ExecParams { - process_id: "proc-1".to_string(), - argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], - cwd: std::env::current_dir().expect("cwd"), - env: HashMap::new(), - tty: true, - arg0: None, - }, - })) - .await - { - panic!("request handling should not fail the handler: {err}"); - } - - let ExecServerOutboundMessage::Error { request_id, error } = - recv_outbound(&mut outgoing_rx).await - else { - panic!("expected invalid-request error"); - }; - assert_eq!(request_id, RequestId::Integer(7)); - assert_eq!(error.code, -32600); - assert_eq!( - error.message, - "client must call initialize before using exec methods" - ); - } - - #[tokio::test] - async fn exec_methods_require_initialized_notification_after_initialize() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); - let mut handler = ExecServerHandler::new(outgoing_tx); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outbound(&mut outgoing_rx).await; - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { - request_id: RequestId::Integer(2), - params: crate::protocol::ExecParams { - process_id: "proc-1".to_string(), - argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], - cwd: std::env::current_dir().expect("cwd"), - env: HashMap::new(), - tty: true, - arg0: None, - }, - })) - .await - { - panic!("request handling should not fail the handler: {err}"); - } - - let ExecServerOutboundMessage::Error { request_id, error } = - recv_outbound(&mut outgoing_rx).await - else { - panic!("expected invalid-request error"); - }; - assert_eq!(request_id, RequestId::Integer(2)); - assert_eq!(error.code, -32600); - assert_eq!( - error.message, - "client must send initialized before using exec methods" - ); - } - - #[tokio::test] - async fn initialized_before_initialize_is_a_protocol_error() { - let (outgoing_tx, _outgoing_rx) = tokio::sync::mpsc::channel(1); - let mut handler = ExecServerHandler::new(outgoing_tx); - - let result = handler - .handle_message(ExecServerInboundMessage::Notification( - ExecServerClientNotification::Initialized, - )) - .await; - - match result { - Err(err) => { - assert_eq!( - err, - "received `initialized` notification before `initialize`" - ); - } - Ok(()) => panic!("expected protocol error for early initialized notification"), - } - } - - #[tokio::test] - async fn initialize_may_only_be_sent_once_per_connection() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); - let mut handler = ExecServerHandler::new(outgoing_tx); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outbound(&mut outgoing_rx).await; - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(2), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("duplicate initialize should not fail the handler: {err}"); - } - - let ExecServerOutboundMessage::Error { request_id, error } = - recv_outbound(&mut outgoing_rx).await - else { - panic!("expected invalid-request error"); - }; - assert_eq!(request_id, RequestId::Integer(2)); - assert_eq!(error.code, -32600); - assert_eq!( - error.message, - "initialize may only be sent once per connection" - ); - } - - #[tokio::test] - 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); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outbound(&mut outgoing_rx).await; - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Notification( - ExecServerClientNotification::Initialized, - )) - .await - { - panic!("initialized should succeed: {err}"); - } - - let params = crate::protocol::ExecParams { - process_id: "proc-1".to_string(), - argv: vec![ - "bash".to_string(), - "-lc".to_string(), - "sleep 30".to_string(), - ], - cwd: std::env::current_dir().expect("cwd"), - env: HashMap::new(), - tty: false, - arg0: None, - }; - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { - request_id: RequestId::Integer(2), - params: params.clone(), - })) - .await - { - panic!("first exec should succeed: {err}"); - } - 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 { - process_id: "proc-2".to_string(), - argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], - ..params - }, - })) - .await - { - panic!("second exec should succeed: {err}"); - } - - let ExecServerOutboundMessage::Response { - request_id, - response: ExecServerResponseMessage::Exec(second_exec), - } = recv_outbound(&mut outgoing_rx).await - else { - panic!("expected second exec response"); - }; - assert_eq!(request_id, RequestId::Integer(3)); - assert_eq!(second_exec.process_id, "proc-2"); - - handler.shutdown().await; - } - - #[tokio::test] - async fn writes_to_pipe_backed_processes_are_rejected() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(4); - let mut handler = ExecServerHandler::new(outgoing_tx); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outbound(&mut outgoing_rx).await; - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Notification( - ExecServerClientNotification::Initialized, - )) - .await - { - panic!("initialized should succeed: {err}"); - } - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { - request_id: RequestId::Integer(2), - params: crate::protocol::ExecParams { - process_id: "proc-1".to_string(), - argv: vec![ - "bash".to_string(), - "-lc".to_string(), - "sleep 30".to_string(), - ], - cwd: std::env::current_dir().expect("cwd"), - env: HashMap::new(), - tty: false, - arg0: None, - }, - })) - .await - { - panic!("exec should succeed: {err}"); - } - 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: exec_response.process_id, - chunk: b"hello\n".to_vec().into(), - }, - }, - )) - .await - { - panic!("write should not fail the handler: {err}"); - } - - let ExecServerOutboundMessage::Error { request_id, error } = - recv_outbound(&mut outgoing_rx).await - else { - panic!("expected stdin-closed error"); - }; - assert_eq!(request_id, RequestId::Integer(3)); - assert_eq!(error.code, -32600); - assert_eq!(error.message, "stdin is closed for process proc-1"); - - handler.shutdown().await; - } - - #[tokio::test] - async fn writes_to_unknown_processes_are_rejected() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); - let mut handler = ExecServerHandler::new(outgoing_tx); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outbound(&mut outgoing_rx).await; - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Notification( - ExecServerClientNotification::Initialized, - )) - .await - { - panic!("initialized should succeed: {err}"); - } - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Write { - request_id: RequestId::Integer(2), - params: WriteParams { - process_id: "missing".to_string(), - chunk: b"hello\n".to_vec().into(), - }, - }, - )) - .await - { - panic!("write should not fail the handler: {err}"); - } - - let ExecServerOutboundMessage::Error { request_id, error } = - recv_outbound(&mut outgoing_rx).await - else { - panic!("expected unknown-process error"); - }; - assert_eq!(request_id, RequestId::Integer(2)); - assert_eq!(error.code, -32600); - assert_eq!(error.message, "unknown process id missing"); - } - - #[tokio::test] - async fn terminate_unknown_processes_report_running_false() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); - let mut handler = ExecServerHandler::new(outgoing_tx); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outbound(&mut outgoing_rx).await; - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Notification( - ExecServerClientNotification::Initialized, - )) - .await - { - panic!("initialized should succeed: {err}"); - } - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Terminate { - request_id: RequestId::Integer(2), - params: crate::protocol::TerminateParams { - process_id: "missing".to_string(), - }, - }, - )) - .await - { - panic!("terminate should not fail the handler: {err}"); - } - - assert_eq!( - recv_outbound(&mut outgoing_rx).await, - ExecServerOutboundMessage::Response { - request_id: RequestId::Integer(2), - response: ExecServerResponseMessage::Terminate(TerminateResponse { - running: false, - }), - } - ); - } - - #[tokio::test] - 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); - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outbound(&mut outgoing_rx).await; - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Notification( - ExecServerClientNotification::Initialized, - )) - .await - { - panic!("initialized should succeed: {err}"); - } - - let spawned = codex_utils_pty::spawn_pipe_process_no_stdin( - "bash", - &["-lc".to_string(), "sleep 30".to_string()], - std::env::current_dir().expect("cwd").as_path(), - &HashMap::new(), - &None, - ) - .await - .expect("spawn test process"); - { - let mut process_map = handler.processes.lock().await; - process_map.insert( - "proc-1".to_string(), - super::RunningProcess { - session: spawned.session, - tty: false, - output: std::collections::VecDeque::new(), - retained_bytes: 0, - next_seq: 1, - exit_code: None, - }, - ); - } - - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Terminate { - request_id: RequestId::Integer(2), - params: crate::protocol::TerminateParams { - process_id: "proc-1".to_string(), - }, - }, - )) - .await - { - panic!("terminate should not fail the handler: {err}"); - } - - assert_eq!( - recv_outbound(&mut outgoing_rx).await, - ExecServerOutboundMessage::Response { - request_id: RequestId::Integer(2), - response: ExecServerResponseMessage::Terminate(TerminateResponse { running: true }), - } - ); - - assert!( - handler.processes.lock().await.contains_key("proc-1"), - "terminated ids should stay reserved until exit cleanup removes them" - ); - - handler.shutdown().await; - } - - #[tokio::test] - async fn read_paginates_retained_output_without_skipping_omitted_chunks() { - let (outgoing_tx, _outgoing_rx) = tokio::sync::mpsc::channel(1); - let mut handler = ExecServerHandler::new(outgoing_tx); - let _ = handler.initialize().expect("initialize should succeed"); - handler.initialized().expect("initialized should succeed"); - - let spawned = codex_utils_pty::spawn_pipe_process_no_stdin( - "bash", - &["-lc".to_string(), "true".to_string()], - std::env::current_dir().expect("cwd").as_path(), - &HashMap::new(), - &None, - ) - .await - .expect("spawn test process"); - { - let mut process_map = handler.processes.lock().await; - process_map.insert( - "proc-1".to_string(), - RunningProcess { - session: spawned.session, - tty: false, - output: VecDeque::from([ - RetainedOutputChunk { - seq: 1, - stream: ExecOutputStream::Stdout, - chunk: b"abc".to_vec(), - }, - RetainedOutputChunk { - seq: 2, - stream: ExecOutputStream::Stderr, - chunk: b"def".to_vec(), - }, - ]), - retained_bytes: 6, - next_seq: 3, - exit_code: None, - }, - ); - } - - let first = handler - .read(ReadParams { - process_id: "proc-1".to_string(), - after_seq: Some(0), - max_bytes: Some(3), - wait_ms: Some(0), - }) - .await - .expect("first read should succeed"); - - assert_eq!(first.chunks.len(), 1); - assert_eq!(first.chunks[0].seq, 1); - assert_eq!(first.chunks[0].stream, ExecOutputStream::Stdout); - assert_eq!(first.chunks[0].chunk.clone().into_inner(), b"abc".to_vec()); - assert_eq!(first.next_seq, 2); - - let second = handler - .read(ReadParams { - process_id: "proc-1".to_string(), - after_seq: Some(first.next_seq - 1), - max_bytes: Some(3), - wait_ms: Some(0), - }) - .await - .expect("second read should succeed"); - - assert_eq!(second.chunks.len(), 1); - assert_eq!(second.chunks[0].seq, 2); - assert_eq!(second.chunks[0].stream, ExecOutputStream::Stderr); - assert_eq!(second.chunks[0].chunk.clone().into_inner(), b"def".to_vec()); - assert_eq!(second.next_seq, 3); - - handler.shutdown().await; - } -} +mod tests; diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs new file mode 100644 index 0000000000..c9682c16da --- /dev/null +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -0,0 +1,724 @@ +use std::collections::HashMap; +use std::collections::VecDeque; +use std::sync::Arc; +use std::time::Duration; + +use pretty_assertions::assert_eq; +use tokio::sync::Notify; +use tokio::time::timeout; + +use super::ExecServerHandler; +use super::RetainedOutputChunk; +use super::RunningProcess; +use crate::protocol::ExecOutputStream; +use crate::protocol::ExecSandboxConfig; +use crate::protocol::ExecSandboxMode; +use crate::protocol::InitializeParams; +use crate::protocol::InitializeResponse; +use crate::protocol::PROTOCOL_VERSION; +use crate::protocol::ReadParams; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteParams; +use crate::server::routing::ExecServerClientNotification; +use crate::server::routing::ExecServerInboundMessage; +use crate::server::routing::ExecServerOutboundMessage; +use crate::server::routing::ExecServerRequest; +use crate::server::routing::ExecServerResponseMessage; +use codex_app_server_protocol::RequestId; + +async fn recv_outbound( + outgoing_rx: &mut tokio::sync::mpsc::Receiver, +) -> ExecServerOutboundMessage { + let recv_result = timeout(Duration::from_secs(1), outgoing_rx.recv()).await; + let maybe_message = match recv_result { + Ok(maybe_message) => maybe_message, + Err(err) => panic!("timed out waiting for handler output: {err}"), + }; + match maybe_message { + Some(message) => message, + None => panic!("handler output channel closed unexpectedly"), + } +} + +#[tokio::test] +async fn initialize_response_reports_protocol_version() { + let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(outgoing_tx); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + + assert_eq!( + recv_outbound(&mut outgoing_rx).await, + ExecServerOutboundMessage::Response { + request_id: RequestId::Integer(1), + response: ExecServerResponseMessage::Initialize(InitializeResponse { + protocol_version: PROTOCOL_VERSION.to_string(), + }), + } + ); +} + +#[tokio::test] +async fn exec_methods_require_initialize() { + let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(outgoing_tx); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { + request_id: RequestId::Integer(7), + params: crate::protocol::ExecParams { + process_id: "proc-1".to_string(), + argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], + cwd: std::env::current_dir().expect("cwd"), + env: HashMap::new(), + tty: true, + arg0: None, + sandbox: None, + }, + })) + .await + { + panic!("request handling should not fail the handler: {err}"); + } + + let ExecServerOutboundMessage::Error { request_id, error } = + recv_outbound(&mut outgoing_rx).await + else { + panic!("expected invalid-request error"); + }; + assert_eq!(request_id, RequestId::Integer(7)); + assert_eq!(error.code, -32600); + assert_eq!( + error.message, + "client must call initialize before using exec methods" + ); +} + +#[tokio::test] +async fn exec_methods_require_initialized_notification_after_initialize() { + let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); + let mut handler = ExecServerHandler::new(outgoing_tx); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + let _ = recv_outbound(&mut outgoing_rx).await; + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { + request_id: RequestId::Integer(2), + params: crate::protocol::ExecParams { + process_id: "proc-1".to_string(), + argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], + cwd: std::env::current_dir().expect("cwd"), + env: HashMap::new(), + tty: true, + arg0: None, + sandbox: None, + }, + })) + .await + { + panic!("request handling should not fail the handler: {err}"); + } + + let ExecServerOutboundMessage::Error { request_id, error } = + recv_outbound(&mut outgoing_rx).await + else { + panic!("expected invalid-request error"); + }; + assert_eq!(request_id, RequestId::Integer(2)); + assert_eq!(error.code, -32600); + assert_eq!( + error.message, + "client must send initialized before using exec methods" + ); +} + +#[tokio::test] +async fn initialized_before_initialize_is_a_protocol_error() { + let (outgoing_tx, _outgoing_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(outgoing_tx); + + let result = handler + .handle_message(ExecServerInboundMessage::Notification( + ExecServerClientNotification::Initialized, + )) + .await; + + match result { + Err(err) => { + assert_eq!( + err, + "received `initialized` notification before `initialize`" + ); + } + Ok(()) => panic!("expected protocol error for early initialized notification"), + } +} + +#[tokio::test] +async fn initialize_may_only_be_sent_once_per_connection() { + let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); + let mut handler = ExecServerHandler::new(outgoing_tx); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + let _ = recv_outbound(&mut outgoing_rx).await; + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(2), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("duplicate initialize should not fail the handler: {err}"); + } + + let ExecServerOutboundMessage::Error { request_id, error } = + recv_outbound(&mut outgoing_rx).await + else { + panic!("expected invalid-request error"); + }; + assert_eq!(request_id, RequestId::Integer(2)); + assert_eq!(error.code, -32600); + assert_eq!( + error.message, + "initialize may only be sent once per connection" + ); +} + +#[tokio::test] +async fn host_default_sandbox_requests_are_rejected_until_supported() { + let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(3); + let mut handler = ExecServerHandler::new(outgoing_tx); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + let _ = recv_outbound(&mut outgoing_rx).await; + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Notification( + ExecServerClientNotification::Initialized, + )) + .await + { + panic!("initialized should succeed: {err}"); + } + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { + request_id: RequestId::Integer(2), + params: crate::protocol::ExecParams { + process_id: "proc-1".to_string(), + argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], + cwd: std::env::current_dir().expect("cwd"), + env: HashMap::new(), + tty: false, + arg0: None, + sandbox: Some(ExecSandboxConfig { + mode: ExecSandboxMode::HostDefault, + }), + }, + })) + .await + { + panic!("request handling should not fail the handler: {err}"); + } + + let ExecServerOutboundMessage::Error { request_id, error } = + recv_outbound(&mut outgoing_rx).await + else { + panic!("expected unsupported sandbox error"); + }; + assert_eq!(request_id, RequestId::Integer(2)); + assert_eq!(error.code, -32600); + assert_eq!( + error.message, + "sandbox mode `hostDefault` is not supported by exec-server yet" + ); +} + +#[tokio::test] +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); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + let _ = recv_outbound(&mut outgoing_rx).await; + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Notification( + ExecServerClientNotification::Initialized, + )) + .await + { + panic!("initialized should succeed: {err}"); + } + + let params = crate::protocol::ExecParams { + process_id: "proc-1".to_string(), + argv: vec![ + "bash".to_string(), + "-lc".to_string(), + "sleep 30".to_string(), + ], + cwd: std::env::current_dir().expect("cwd"), + env: HashMap::new(), + tty: false, + arg0: None, + sandbox: None, + }; + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { + request_id: RequestId::Integer(2), + params: params.clone(), + })) + .await + { + panic!("first exec should succeed: {err}"); + } + 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 { + process_id: "proc-2".to_string(), + argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], + ..params + }, + })) + .await + { + panic!("second exec should succeed: {err}"); + } + + let ExecServerOutboundMessage::Response { + request_id, + response: ExecServerResponseMessage::Exec(second_exec), + } = recv_outbound(&mut outgoing_rx).await + else { + panic!("expected second exec response"); + }; + assert_eq!(request_id, RequestId::Integer(3)); + assert_eq!(second_exec.process_id, "proc-2"); + + handler.shutdown().await; +} + +#[tokio::test] +async fn writes_to_pipe_backed_processes_are_rejected() { + let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(4); + let mut handler = ExecServerHandler::new(outgoing_tx); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + let _ = recv_outbound(&mut outgoing_rx).await; + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Notification( + ExecServerClientNotification::Initialized, + )) + .await + { + panic!("initialized should succeed: {err}"); + } + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec { + request_id: RequestId::Integer(2), + params: crate::protocol::ExecParams { + process_id: "proc-1".to_string(), + argv: vec![ + "bash".to_string(), + "-lc".to_string(), + "sleep 30".to_string(), + ], + cwd: std::env::current_dir().expect("cwd"), + env: HashMap::new(), + tty: false, + arg0: None, + sandbox: None, + }, + })) + .await + { + panic!("exec should succeed: {err}"); + } + 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: exec_response.process_id, + chunk: b"hello\n".to_vec().into(), + }, + }, + )) + .await + { + panic!("write should not fail the handler: {err}"); + } + + let ExecServerOutboundMessage::Error { request_id, error } = + recv_outbound(&mut outgoing_rx).await + else { + panic!("expected stdin-closed error"); + }; + assert_eq!(request_id, RequestId::Integer(3)); + assert_eq!(error.code, -32600); + assert_eq!(error.message, "stdin is closed for process proc-1"); + + handler.shutdown().await; +} + +#[tokio::test] +async fn writes_to_unknown_processes_are_rejected() { + let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); + let mut handler = ExecServerHandler::new(outgoing_tx); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + let _ = recv_outbound(&mut outgoing_rx).await; + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Notification( + ExecServerClientNotification::Initialized, + )) + .await + { + panic!("initialized should succeed: {err}"); + } + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Write { + request_id: RequestId::Integer(2), + params: WriteParams { + process_id: "missing".to_string(), + chunk: b"hello\n".to_vec().into(), + }, + }, + )) + .await + { + panic!("write should not fail the handler: {err}"); + } + + let ExecServerOutboundMessage::Error { request_id, error } = + recv_outbound(&mut outgoing_rx).await + else { + panic!("expected unknown-process error"); + }; + assert_eq!(request_id, RequestId::Integer(2)); + assert_eq!(error.code, -32600); + assert_eq!(error.message, "unknown process id missing"); +} + +#[tokio::test] +async fn terminate_unknown_processes_report_running_false() { + let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2); + let mut handler = ExecServerHandler::new(outgoing_tx); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + let _ = recv_outbound(&mut outgoing_rx).await; + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Notification( + ExecServerClientNotification::Initialized, + )) + .await + { + panic!("initialized should succeed: {err}"); + } + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Terminate { + request_id: RequestId::Integer(2), + params: crate::protocol::TerminateParams { + process_id: "missing".to_string(), + }, + }, + )) + .await + { + panic!("terminate should not fail the handler: {err}"); + } + + assert_eq!( + recv_outbound(&mut outgoing_rx).await, + ExecServerOutboundMessage::Response { + request_id: RequestId::Integer(2), + response: ExecServerResponseMessage::Terminate(TerminateResponse { running: false }), + } + ); +} + +#[tokio::test] +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); + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Initialize { + request_id: RequestId::Integer(1), + params: InitializeParams { + client_name: "test".to_string(), + }, + }, + )) + .await + { + panic!("initialize should succeed: {err}"); + } + let _ = recv_outbound(&mut outgoing_rx).await; + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Notification( + ExecServerClientNotification::Initialized, + )) + .await + { + panic!("initialized should succeed: {err}"); + } + + let spawned = codex_utils_pty::spawn_pipe_process_no_stdin( + "bash", + &["-lc".to_string(), "sleep 30".to_string()], + std::env::current_dir().expect("cwd").as_path(), + &HashMap::new(), + &None, + ) + .await + .expect("spawn test process"); + { + let mut process_map = handler.processes.lock().await; + process_map.insert( + "proc-1".to_string(), + super::RunningProcess { + session: spawned.session, + tty: false, + output: std::collections::VecDeque::new(), + retained_bytes: 0, + next_seq: 1, + exit_code: None, + output_notify: Arc::new(Notify::new()), + }, + ); + } + + if let Err(err) = handler + .handle_message(ExecServerInboundMessage::Request( + ExecServerRequest::Terminate { + request_id: RequestId::Integer(2), + params: crate::protocol::TerminateParams { + process_id: "proc-1".to_string(), + }, + }, + )) + .await + { + panic!("terminate should not fail the handler: {err}"); + } + + assert_eq!( + recv_outbound(&mut outgoing_rx).await, + ExecServerOutboundMessage::Response { + request_id: RequestId::Integer(2), + response: ExecServerResponseMessage::Terminate(TerminateResponse { running: true }), + } + ); + + assert!( + handler.processes.lock().await.contains_key("proc-1"), + "terminated ids should stay reserved until exit cleanup removes them" + ); + + handler.shutdown().await; +} + +#[tokio::test] +async fn read_paginates_retained_output_without_skipping_omitted_chunks() { + let (outgoing_tx, _outgoing_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(outgoing_tx); + let _ = handler.initialize().expect("initialize should succeed"); + handler.initialized().expect("initialized should succeed"); + + let spawned = codex_utils_pty::spawn_pipe_process_no_stdin( + "bash", + &["-lc".to_string(), "true".to_string()], + std::env::current_dir().expect("cwd").as_path(), + &HashMap::new(), + &None, + ) + .await + .expect("spawn test process"); + { + let mut process_map = handler.processes.lock().await; + process_map.insert( + "proc-1".to_string(), + RunningProcess { + session: spawned.session, + tty: false, + output: VecDeque::from([ + RetainedOutputChunk { + seq: 1, + stream: ExecOutputStream::Stdout, + chunk: b"abc".to_vec(), + }, + RetainedOutputChunk { + seq: 2, + stream: ExecOutputStream::Stderr, + chunk: b"def".to_vec(), + }, + ]), + retained_bytes: 6, + next_seq: 3, + exit_code: None, + output_notify: Arc::new(Notify::new()), + }, + ); + } + + let first = handler + .read(ReadParams { + process_id: "proc-1".to_string(), + after_seq: Some(0), + max_bytes: Some(3), + wait_ms: Some(0), + }) + .await + .expect("first read should succeed"); + + assert_eq!(first.chunks.len(), 1); + assert_eq!(first.chunks[0].seq, 1); + assert_eq!(first.chunks[0].stream, ExecOutputStream::Stdout); + assert_eq!(first.chunks[0].chunk.clone().into_inner(), b"abc".to_vec()); + assert_eq!(first.next_seq, 2); + + let second = handler + .read(ReadParams { + process_id: "proc-1".to_string(), + after_seq: Some(first.next_seq - 1), + max_bytes: Some(3), + wait_ms: Some(0), + }) + .await + .expect("second read should succeed"); + + assert_eq!(second.chunks.len(), 1); + assert_eq!(second.chunks[0].seq, 2); + assert_eq!(second.chunks[0].stream, ExecOutputStream::Stderr); + assert_eq!(second.chunks[0].chunk.clone().into_inner(), b"def".to_vec()); + assert_eq!(second.next_seq, 3); + + handler.shutdown().await; +} diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 812fcb67f3..60b7d0c7dd 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -16,7 +16,7 @@ use crate::server::routing::encode_outbound_message; use crate::server::routing::route_jsonrpc_message; pub(crate) async fn run_connection(connection: JsonRpcConnection) { - let (json_outgoing_tx, mut incoming_rx) = connection.into_parts(); + let (json_outgoing_tx, mut incoming_rx, _connection_tasks) = connection.into_parts(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); let mut handler = ExecServerHandler::new(outgoing_tx.clone()); diff --git a/codex-rs/exec-server/src/server/routing.rs b/codex-rs/exec-server/src/server/routing.rs index e70b371980..f900d7cf69 100644 --- a/codex-rs/exec-server/src/server/routing.rs +++ b/codex-rs/exec-server/src/server/routing.rs @@ -261,6 +261,8 @@ mod tests { use crate::protocol::ExecExitedNotification; use crate::protocol::ExecParams; use crate::protocol::ExecResponse; + use crate::protocol::ExecSandboxConfig; + use crate::protocol::ExecSandboxMode; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; @@ -407,6 +409,51 @@ mod tests { env: std::collections::HashMap::new(), tty: true, arg0: None, + sandbox: None, + } + ); + } + + #[test] + fn routes_exec_requests_with_optional_sandbox_config() { + let cwd = std::env::current_dir().expect("cwd"); + let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(4), + method: EXEC_METHOD.to_string(), + params: Some(json!({ + "processId": "proc-1", + "argv": ["bash", "-lc", "true"], + "cwd": cwd, + "env": {}, + "tty": true, + "arg0": null, + "sandbox": { + "mode": "none", + }, + })), + trace: None, + })) + .expect("exec request with sandbox should route"); + + let RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Request( + ExecServerRequest::Exec { request_id, params }, + )) = routed + else { + panic!("expected typed exec request"); + }; + assert_eq!(request_id, RequestId::Integer(4)); + assert_eq!( + params, + ExecParams { + process_id: "proc-1".to_string(), + argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()], + cwd: std::env::current_dir().expect("cwd"), + env: std::collections::HashMap::new(), + tty: true, + arg0: None, + sandbox: Some(ExecSandboxConfig { + mode: ExecSandboxMode::None, + }), } ); } diff --git a/codex-rs/exec-server/tests/stdio_smoke.rs b/codex-rs/exec-server/tests/stdio_smoke.rs index 041c7a4b75..f10650b799 100644 --- a/codex-rs/exec-server/tests/stdio_smoke.rs +++ b/codex-rs/exec-server/tests/stdio_smoke.rs @@ -109,6 +109,7 @@ async fn exec_server_client_streams_output_and_accepts_writes() -> anyhow::Resul env, tty: true, arg0: None, + sandbox: None, }) .await?; let process_id = response.process_id; @@ -174,6 +175,7 @@ async fn exec_server_client_connects_over_websocket() -> anyhow::Result<()> { env, tty: true, arg0: None, + sandbox: None, }) .await?; let process_id = response.process_id; @@ -248,6 +250,7 @@ async fn websocket_disconnect_terminates_processes_for_that_connection() -> anyh env, tty: false, arg0: None, + sandbox: None, }) .await?; }