From 84bb212aef25da549fb3abce082e9c880762c1a1 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 18 Mar 2026 23:05:11 +0000 Subject: [PATCH] exec-server: simplify transport and dispatch shape Co-authored-by: Codex --- codex-rs/exec-server/src/client.rs | 83 +- codex-rs/exec-server/src/server.rs | 26 +- codex-rs/exec-server/src/server/filesystem.rs | 4 +- codex-rs/exec-server/src/server/handler.rs | 198 +--- .../exec-server/src/server/handler/tests.rs | 852 ++++-------------- codex-rs/exec-server/src/server/jsonrpc.rs | 33 + codex-rs/exec-server/src/server/processor.rs | 493 +++++++--- codex-rs/exec-server/src/server/routing.rs | 585 ------------ codex-rs/exec-server/src/server/transport.rs | 55 +- 9 files changed, 664 insertions(+), 1665 deletions(-) create mode 100644 codex-rs/exec-server/src/server/jsonrpc.rs delete mode 100644 codex-rs/exec-server/src/server/routing.rs diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index bd0ec047cd..05f2de9f8d 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -31,8 +31,6 @@ use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; use serde::Serialize; use serde_json::Value; -use tokio::io::AsyncRead; -use tokio::io::AsyncWrite; use tokio::sync::Mutex; use tokio::sync::broadcast; use tokio::sync::mpsc; @@ -76,13 +74,13 @@ use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; use crate::protocol::WriteResponse; use crate::server::ExecServerHandler; -use crate::server::ExecServerOutboundMessage; use crate::server::ExecServerServerNotification; impl Default for ExecServerClientConnectOptions { fn default() -> Self { Self { client_name: "codex-core".to_string(), + auth_token: None, initialize_timeout: INITIALIZE_TIMEOUT, } } @@ -92,6 +90,7 @@ impl From for ExecServerClientConnectOptions { fn from(value: RemoteExecServerConnectArgs) -> Self { Self { client_name: value.client_name, + auth_token: value.auth_token, initialize_timeout: value.initialize_timeout, } } @@ -105,6 +104,7 @@ impl RemoteExecServerConnectArgs { Self { websocket_url, client_name, + auth_token: None, connect_timeout: CONNECT_TIMEOUT, initialize_timeout: INITIALIZE_TIMEOUT, } @@ -123,7 +123,6 @@ struct ExecServerProcess { process_id: String, output_rx: broadcast::Receiver, status: Arc, - client: ExecServerClient, } #[cfg(test)] @@ -135,18 +134,6 @@ impl ExecServerProcess { fn has_exited(&self) -> bool { self.status.has_exited() } - - fn exit_code(&self) -> Option { - self.status.exit_code() - } - - fn terminate(&self) { - let client = self.client.clone(); - let process_id = self.process_id.clone(); - tokio::spawn(async move { - let _ = client.terminate_session(&process_id).await; - }); - } } #[cfg(test)] @@ -168,10 +155,6 @@ impl RemoteProcessStatus { self.exited.load(Ordering::SeqCst) } - fn exit_code(&self) -> Option { - self.exit_code.lock().ok().and_then(|guard| *guard) - } - fn mark_exited(&self, exit_code: Option) { self.exited.store(true, Ordering::SeqCst); if let Ok(mut guard) = self.exit_code.lock() { @@ -348,21 +331,16 @@ impl ExecServerClient { pub async fn connect_in_process( options: ExecServerClientConnectOptions, ) -> Result { - let (outbound_tx, mut outgoing_rx) = mpsc::channel::(256); - let handler = Arc::new(Mutex::new(ExecServerHandler::new(outbound_tx))); + let (notification_tx, mut notification_rx) = + mpsc::channel::(256); + let handler = Arc::new(Mutex::new(ExecServerHandler::new(notification_tx, None))); let inner = Arc::new_cyclic(|weak| { let weak = weak.clone(); let reader_task = tokio::spawn(async move { - while let Some(message) = outgoing_rx.recv().await { - if let Some(inner) = weak.upgrade() - && let Err(err) = handle_in_process_outbound_message(&inner, message).await - { - warn!( - "in-process exec-server client closing after unexpected response: {err}" - ); - handle_transport_shutdown(&inner).await; - return; + while let Some(notification) = notification_rx.recv().await { + if let Some(inner) = weak.upgrade() { + handle_in_process_notification(&inner, notification).await; } } @@ -386,22 +364,6 @@ impl ExecServerClient { Ok(client) } - pub async fn connect_stdio( - stdin: W, - stdout: R, - options: ExecServerClientConnectOptions, - ) -> Result - where - R: AsyncRead + Unpin + Send + 'static, - W: AsyncWrite + Unpin + Send + 'static, - { - Self::connect( - JsonRpcConnection::from_stdio(stdout, stdin, "exec-server stdio".to_string()), - options, - ) - .await - } - pub async fn connect_websocket( args: RemoteExecServerConnectArgs, ) -> Result { @@ -521,7 +483,6 @@ impl ExecServerClient { process_id, output_rx, status, - client: self.clone(), }) } @@ -648,11 +609,15 @@ impl ExecServerClient { ) -> Result<(), ExecServerError> { let ExecServerClientConnectOptions { client_name, + auth_token, initialize_timeout, } = options; timeout(initialize_timeout, async { let _: InitializeResponse = self - .request_initialize(InitializeParams { client_name }) + .request_initialize(InitializeParams { + client_name, + auth_token, + }) .await?; self.notify(INITIALIZED_METHOD, &serde_json::json!({})) .await @@ -735,7 +700,7 @@ impl ExecServerClient { params: InitializeParams, ) -> Result { if let ClientBackend::InProcess { handler } = &self.inner.backend { - return server_result_to_client(handler.lock().await.initialize()); + return server_result_to_client(handler.lock().await.initialize(params)); } self.send_pending_request(INITIALIZE_METHOD, ¶ms, PendingRequest::Initialize) @@ -825,24 +790,6 @@ async fn send_jsonrpc_request( .map_err(|_| ExecServerError::Closed) } -async fn handle_in_process_outbound_message( - inner: &Arc, - message: ExecServerOutboundMessage, -) -> Result<(), ExecServerError> { - match message { - ExecServerOutboundMessage::Response { .. } | ExecServerOutboundMessage::Error { .. } => { - return Err(ExecServerError::Protocol( - "unexpected in-process RPC response".to_string(), - )); - } - ExecServerOutboundMessage::Notification(notification) => { - handle_in_process_notification(inner, notification).await; - } - } - - Ok(()) -} - async fn handle_in_process_notification( inner: &Arc, notification: ExecServerServerNotification, diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index 0faf9be40f..de1791ea27 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -1,21 +1,37 @@ mod filesystem; mod handler; +mod jsonrpc; mod processor; -mod routing; mod transport; pub(crate) use handler::ExecServerHandler; -pub(crate) use routing::ExecServerOutboundMessage; -pub(crate) use routing::ExecServerServerNotification; +pub(crate) use handler::ExecServerServerNotification; +pub(crate) use jsonrpc::internal_error; +pub(crate) use jsonrpc::invalid_params; +pub(crate) use jsonrpc::invalid_request; +pub(crate) use jsonrpc::unauthorized; pub use transport::ExecServerTransport; pub use transport::ExecServerTransportParseError; +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct ExecServerConfig { + pub auth_token: Option, +} + pub async fn run_main() -> Result<(), Box> { - run_main_with_transport(ExecServerTransport::Stdio).await + run_main_with_transport_and_config(ExecServerTransport::default(), ExecServerConfig::default()) + .await } pub async fn run_main_with_transport( transport: ExecServerTransport, ) -> Result<(), Box> { - transport::run_transport(transport).await + run_main_with_transport_and_config(transport, ExecServerConfig::default()).await +} + +pub async fn run_main_with_transport_and_config( + transport: ExecServerTransport, + config: ExecServerConfig, +) -> Result<(), Box> { + transport::run_transport(transport, config).await } diff --git a/codex-rs/exec-server/src/server/filesystem.rs b/codex-rs/exec-server/src/server/filesystem.rs index 2f0a71edf6..b9941a61a7 100644 --- a/codex-rs/exec-server/src/server/filesystem.rs +++ b/codex-rs/exec-server/src/server/filesystem.rs @@ -25,8 +25,8 @@ use codex_environment::Environment; use codex_environment::ExecutorFileSystem; use codex_environment::RemoveOptions; -use crate::server::routing::internal_error; -use crate::server::routing::invalid_request; +use crate::server::internal_error; +use crate::server::invalid_request; #[derive(Clone)] pub(crate) struct ExecServerFileSystem { diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index 35282f5ffe..58ba0e64af 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -36,14 +36,19 @@ use crate::protocol::ReadResponse; use crate::protocol::TerminateResponse; use crate::protocol::WriteResponse; use crate::server::filesystem::ExecServerFileSystem; -use crate::server::routing::ExecServerOutboundMessage; -use crate::server::routing::ExecServerServerNotification; -use crate::server::routing::internal_error; -use crate::server::routing::invalid_params; -use crate::server::routing::invalid_request; +use crate::server::internal_error; +use crate::server::invalid_params; +use crate::server::invalid_request; +use crate::server::unauthorized; const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024; +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum ExecServerServerNotification { + OutputDelta(ExecOutputDeltaNotification), + Exited(ExecExitedNotification), +} + #[derive(Clone)] struct RetainedOutputChunk { seq: u64, @@ -62,8 +67,9 @@ struct RunningProcess { } pub(crate) struct ExecServerHandler { - outbound_tx: mpsc::Sender, + notification_tx: mpsc::Sender, file_system: ExecServerFileSystem, + required_auth_token: Option, // Keyed by client-chosen logical `processId` scoped to this connection. // This is a protocol handle, not an OS pid. processes: Arc>>, @@ -72,10 +78,14 @@ pub(crate) struct ExecServerHandler { } impl ExecServerHandler { - pub(crate) fn new(outbound_tx: mpsc::Sender) -> Self { + pub(crate) fn new( + notification_tx: mpsc::Sender, + required_auth_token: Option, + ) -> Self { Self { - outbound_tx, + notification_tx, file_system: ExecServerFileSystem::default(), + required_auth_token, processes: Arc::new(Mutex::new(HashMap::new())), initialize_requested: false, initialized: false, @@ -105,12 +115,18 @@ impl ExecServerHandler { pub(crate) fn initialize( &mut self, + params: crate::protocol::InitializeParams, ) -> Result { if self.initialize_requested { return Err(invalid_request( "initialize may only be sent once per connection".to_string(), )); } + if let Some(required_auth_token) = &self.required_auth_token + && params.auth_token.as_deref() != Some(required_auth_token.as_str()) + { + return Err(unauthorized("invalid exec-server auth token".to_string())); + } self.initialize_requested = true; Ok(InitializeResponse { protocol_version: PROTOCOL_VERSION.to_string(), @@ -210,7 +226,7 @@ impl ExecServerHandler { ExecOutputStream::Stdout }, spawned.stdout_rx, - self.outbound_tx.clone(), + self.notification_tx.clone(), Arc::clone(&self.processes), Arc::clone(&output_notify), )); @@ -222,14 +238,14 @@ impl ExecServerHandler { ExecOutputStream::Stderr }, spawned.stderr_rx, - self.outbound_tx.clone(), + self.notification_tx.clone(), Arc::clone(&self.processes), Arc::clone(&output_notify), )); tokio::spawn(watch_exit( process_id.clone(), spawned.exit_rx, - self.outbound_tx.clone(), + self.notification_tx.clone(), Arc::clone(&self.processes), output_notify, )); @@ -402,153 +418,11 @@ impl ExecServerHandler { } } -#[cfg(test)] -impl ExecServerHandler { - async fn handle_message( - &mut self, - message: crate::server::routing::ExecServerInboundMessage, - ) -> Result<(), String> { - match message { - crate::server::routing::ExecServerInboundMessage::Request(request) => { - self.handle_request(request).await - } - crate::server::routing::ExecServerInboundMessage::Notification( - crate::server::routing::ExecServerClientNotification::Initialized, - ) => self.initialized(), - } - } - - async fn handle_request( - &mut self, - request: crate::server::routing::ExecServerRequest, - ) -> Result<(), String> { - let outbound = match request { - crate::server::routing::ExecServerRequest::Initialize { request_id, .. } => { - Self::request_outbound( - request_id, - self.initialize() - .map(crate::server::routing::ExecServerResponseMessage::Initialize), - ) - } - crate::server::routing::ExecServerRequest::Exec { request_id, params } => { - Self::request_outbound( - request_id, - self.exec(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::Exec), - ) - } - crate::server::routing::ExecServerRequest::Read { request_id, params } => { - Self::request_outbound( - request_id, - self.read(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::Read), - ) - } - crate::server::routing::ExecServerRequest::Write { request_id, params } => { - Self::request_outbound( - request_id, - self.write(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::Write), - ) - } - crate::server::routing::ExecServerRequest::Terminate { request_id, params } => { - Self::request_outbound( - request_id, - self.terminate(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::Terminate), - ) - } - crate::server::routing::ExecServerRequest::FsReadFile { request_id, params } => { - Self::request_outbound( - request_id, - self.fs_read_file(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::FsReadFile), - ) - } - crate::server::routing::ExecServerRequest::FsWriteFile { request_id, params } => { - Self::request_outbound( - request_id, - self.fs_write_file(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::FsWriteFile), - ) - } - crate::server::routing::ExecServerRequest::FsCreateDirectory { request_id, params } => { - Self::request_outbound( - request_id, - self.fs_create_directory(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::FsCreateDirectory), - ) - } - crate::server::routing::ExecServerRequest::FsGetMetadata { request_id, params } => { - Self::request_outbound( - request_id, - self.fs_get_metadata(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::FsGetMetadata), - ) - } - crate::server::routing::ExecServerRequest::FsReadDirectory { request_id, params } => { - Self::request_outbound( - request_id, - self.fs_read_directory(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::FsReadDirectory), - ) - } - crate::server::routing::ExecServerRequest::FsRemove { request_id, params } => { - Self::request_outbound( - request_id, - self.fs_remove(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::FsRemove), - ) - } - crate::server::routing::ExecServerRequest::FsCopy { request_id, params } => { - Self::request_outbound( - request_id, - self.fs_copy(params) - .await - .map(crate::server::routing::ExecServerResponseMessage::FsCopy), - ) - } - }; - self.outbound_tx - .send(outbound) - .await - .map_err(|_| "outbound channel closed".to_string()) - } - - fn request_outbound( - request_id: codex_app_server_protocol::RequestId, - result: Result< - crate::server::routing::ExecServerResponseMessage, - codex_app_server_protocol::JSONRPCErrorError, - >, - ) -> crate::server::routing::ExecServerOutboundMessage { - match result { - Ok(response) => crate::server::routing::ExecServerOutboundMessage::Response { - request_id, - response, - }, - Err(error) => { - crate::server::routing::ExecServerOutboundMessage::Error { request_id, error } - } - } - } -} - async fn stream_output( process_id: String, stream: ExecOutputStream, mut receiver: tokio::sync::mpsc::Receiver>, - outbound_tx: mpsc::Sender, + notification_tx: mpsc::Sender, processes: Arc>>, output_notify: Arc, ) { @@ -583,10 +457,8 @@ async fn stream_output( }; output_notify.notify_waiters(); - if outbound_tx - .send(ExecServerOutboundMessage::Notification( - ExecServerServerNotification::OutputDelta(notification), - )) + if notification_tx + .send(ExecServerServerNotification::OutputDelta(notification)) .await .is_err() { @@ -598,7 +470,7 @@ async fn stream_output( async fn watch_exit( process_id: String, exit_rx: tokio::sync::oneshot::Receiver, - outbound_tx: mpsc::Sender, + notification_tx: mpsc::Sender, processes: Arc>>, output_notify: Arc, ) { @@ -610,12 +482,12 @@ async fn watch_exit( } } output_notify.notify_waiters(); - let _ = outbound_tx - .send(ExecServerOutboundMessage::Notification( - ExecServerServerNotification::Exited(ExecExitedNotification { + let _ = notification_tx + .send(ExecServerServerNotification::Exited( + ExecExitedNotification { process_id, exit_code, - }), + }, )) .await; } diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs index c9682c16da..81eedac979 100644 --- a/codex-rs/exec-server/src/server/handler/tests.rs +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -1,322 +1,19 @@ 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::ExecParams; 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::TerminateParams; 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(), +fn exec_params(process_id: &str) -> ExecParams { + ExecParams { + process_id: process_id.to_string(), argv: vec![ "bash".to_string(), "-lc".to_string(), @@ -327,398 +24,203 @@ async fn exec_echoes_client_process_ids() { 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; +async fn initialized_handler() -> ExecServerHandler { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(8); + let mut handler = ExecServerHandler::new(notification_tx, None); + let response = handler + .initialize(InitializeParams { + client_name: "test".to_string(), + auth_token: None, + }) + .expect("initialize should succeed"); + assert_eq!(response.protocol_version, PROTOCOL_VERSION); + handler + .initialized() + .expect("initialized notification should succeed"); + handler } #[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); +async fn initialize_reports_protocol_version() { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(notification_tx, None); - 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 response = handler + .initialize(InitializeParams { + client_name: "test".to_string(), + auth_token: None, + }) + .expect("initialize should succeed"); - 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"); - }; + assert_eq!(response.protocol_version, PROTOCOL_VERSION); +} - 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}"); - } +#[tokio::test] +async fn exec_methods_require_initialize() { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let handler = ExecServerHandler::new(notification_tx, None); + + let error = handler + .exec(exec_params("proc-1")) + .await + .expect_err("exec should fail before initialize"); - 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"); + assert_eq!( + error.message, + "client must call initialize before using exec methods" + ); +} + +#[tokio::test] +async fn exec_methods_require_initialized_notification_after_initialize() { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(notification_tx, None); + let _ = handler + .initialize(InitializeParams { + client_name: "test".to_string(), + auth_token: None, + }) + .expect("initialize should succeed"); + + let error = handler + .exec(exec_params("proc-1")) + .await + .expect_err("exec should fail before initialized notification"); + + 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 (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(notification_tx, None); + + let error = handler + .initialized() + .expect_err("expected protocol error for early initialized notification"); + + assert_eq!( + error, + "received `initialized` notification before `initialize`" + ); +} + +#[tokio::test] +async fn initialize_may_only_be_sent_once_per_connection() { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(notification_tx, None); + let _ = handler + .initialize(InitializeParams { + client_name: "test".to_string(), + auth_token: None, + }) + .expect("first initialize should succeed"); + + let error = handler + .initialize(InitializeParams { + client_name: "test".to_string(), + auth_token: None, + }) + .expect_err("duplicate initialize should fail"); + + assert_eq!(error.code, -32600); + assert_eq!( + error.message, + "initialize may only be sent once per connection" + ); +} + +#[tokio::test] +async fn initialize_rejects_invalid_auth_token() { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(notification_tx, Some("secret-token".to_string())); + + let error = handler + .initialize(InitializeParams { + client_name: "test".to_string(), + auth_token: Some("wrong-token".to_string()), + }) + .expect_err("invalid auth token should fail"); + + assert_eq!(error.code, -32001); + assert_eq!(error.message, "invalid exec-server auth token"); +} + +#[tokio::test] +async fn exec_rejects_host_default_sandbox_mode() { + let handler = initialized_handler().await; + + let error = handler + .exec(ExecParams { + sandbox: Some(ExecSandboxConfig { + mode: ExecSandboxMode::HostDefault, + }), + ..exec_params("proc-1") + }) + .await + .expect_err("hostDefault sandbox should be rejected"); + + assert_eq!(error.code, -32600); + assert_eq!( + error.message, + "sandbox mode `hostDefault` is not supported by exec-server yet" + ); +} + +#[tokio::test] +async fn exec_rejects_duplicate_process_ids() { + let handler = initialized_handler().await; + let first = handler + .exec(exec_params("proc-1")) + .await + .expect("first exec should succeed"); + assert_eq!(first.process_id, "proc-1"); + + let error = handler + .exec(exec_params("proc-1")) + .await + .expect_err("duplicate process id should fail"); + + assert_eq!(error.code, -32600); + assert_eq!(error.message, "process proc-1 already exists"); 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); +async fn write_rejects_unknown_process_ids() { + let handler = initialized_handler().await; - if let Err(err) = handler - .handle_message(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test".to_string(), - }, - }, - )) + let error = handler + .write(WriteParams { + process_id: "missing".to_string(), + chunk: b"input".to_vec().into(), + }) .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}"); - } + .expect_err("writing to an unknown process should fail"); - 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); +async fn terminate_reports_missing_processes_as_not_running() { + let handler = initialized_handler().await; - 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), + let response = handler + .terminate(TerminateParams { + process_id: "missing".to_string(), }) .await - .expect("first read should succeed"); + .expect("terminate 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; + assert_eq!(response.running, false); } diff --git a/codex-rs/exec-server/src/server/jsonrpc.rs b/codex-rs/exec-server/src/server/jsonrpc.rs new file mode 100644 index 0000000000..6e6ac7cfe0 --- /dev/null +++ b/codex-rs/exec-server/src/server/jsonrpc.rs @@ -0,0 +1,33 @@ +use codex_app_server_protocol::JSONRPCErrorError; + +pub(crate) fn invalid_request(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32600, + data: None, + message, + } +} + +pub(crate) fn invalid_params(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32602, + data: None, + message, + } +} + +pub(crate) fn internal_error(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32603, + data: None, + message, + } +} + +pub(crate) fn unauthorized(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32001, + data: None, + message, + } +} diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 9dd9388b30..e0a002e938 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,3 +1,12 @@ +use codex_app_server_protocol::JSONRPCError; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::JSONRPCNotification; +use codex_app_server_protocol::JSONRPCRequest; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use serde::Serialize; +use serde::de::DeserializeOwned; use tokio::sync::mpsc; use tracing::debug; use tracing::warn; @@ -5,32 +14,49 @@ use tracing::warn; use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::server::handler::ExecServerHandler; -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 crate::server::routing::RoutedExecServerMessage; -use crate::server::routing::encode_outbound_message; -use crate::server::routing::route_jsonrpc_message; +use crate::protocol::EXEC_EXITED_METHOD; +use crate::protocol::EXEC_METHOD; +use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; +use crate::protocol::EXEC_READ_METHOD; +use crate::protocol::EXEC_TERMINATE_METHOD; +use crate::protocol::EXEC_WRITE_METHOD; +use crate::protocol::FS_COPY_METHOD; +use crate::protocol::FS_CREATE_DIRECTORY_METHOD; +use crate::protocol::FS_GET_METADATA_METHOD; +use crate::protocol::FS_READ_DIRECTORY_METHOD; +use crate::protocol::FS_READ_FILE_METHOD; +use crate::protocol::FS_REMOVE_METHOD; +use crate::protocol::FS_WRITE_FILE_METHOD; +use crate::protocol::INITIALIZE_METHOD; +use crate::protocol::INITIALIZED_METHOD; +use crate::server::ExecServerConfig; +use crate::server::ExecServerHandler; +use crate::server::ExecServerServerNotification; +use crate::server::internal_error; +use crate::server::invalid_params; +use crate::server::invalid_request; -pub(crate) async fn run_connection(connection: JsonRpcConnection) { +pub(crate) async fn run_connection(connection: JsonRpcConnection, config: ExecServerConfig) { 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()); + let json_outgoing_tx_for_notifications = json_outgoing_tx.clone(); + let (notification_tx, mut notification_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let mut handler = ExecServerHandler::new(notification_tx, config.auth_token); let outbound_task = tokio::spawn(async move { - while let Some(message) = outgoing_rx.recv().await { - let json_message = match encode_outbound_message(message) { + while let Some(notification) = notification_rx.recv().await { + let json_message = match notification_message(notification) { Ok(json_message) => json_message, Err(err) => { - warn!("failed to serialize exec-server outbound message: {err}"); + warn!("failed to serialize exec-server notification: {err}"); break; } }; - if json_outgoing_tx.send(json_message).await.is_err() { + if json_outgoing_tx_for_notifications + .send(json_message) + .await + .is_err() + { break; } } @@ -38,24 +64,20 @@ pub(crate) async fn run_connection(connection: JsonRpcConnection) { while let Some(event) = incoming_rx.recv().await { match event { - JsonRpcConnectionEvent::Message(message) => match route_jsonrpc_message(message) { - Ok(RoutedExecServerMessage::Inbound(message)) => { - if let Err(err) = dispatch_to_handler(&mut handler, message, &outgoing_tx).await - { + JsonRpcConnectionEvent::Message(message) => { + let maybe_response = match handle_connection_message(&mut handler, message).await { + Ok(maybe_response) => maybe_response, + Err(err) => { warn!("closing exec-server connection after protocol error: {err}"); break; } - } - Ok(RoutedExecServerMessage::ImmediateOutbound(message)) => { - if outgoing_tx.send(message).await.is_err() { - break; - } - } - Err(err) => { - warn!("closing exec-server connection after protocol error: {err}"); + }; + if let Some(response) = maybe_response + && json_outgoing_tx.send(response).await.is_err() + { break; } - }, + } JsonRpcConnectionEvent::Disconnected { reason } => { if let Some(reason) = reason { debug!("exec-server connection disconnected: {reason}"); @@ -67,122 +89,319 @@ pub(crate) async fn run_connection(connection: JsonRpcConnection) { handler.shutdown().await; drop(handler); - drop(outgoing_tx); let _ = outbound_task.await; } -async fn dispatch_to_handler( +async fn handle_connection_message( handler: &mut ExecServerHandler, - message: ExecServerInboundMessage, - outgoing_tx: &mpsc::Sender, -) -> Result<(), String> { + message: JSONRPCMessage, +) -> Result, String> { match message { - ExecServerInboundMessage::Request(request) => { - let outbound = match request { - ExecServerRequest::Initialize { request_id, .. } => request_outbound( - request_id, - handler - .initialize() - .map(ExecServerResponseMessage::Initialize), - ), - ExecServerRequest::Exec { request_id, params } => request_outbound( - request_id, - handler - .exec(params) - .await - .map(ExecServerResponseMessage::Exec), - ), - ExecServerRequest::Read { request_id, params } => request_outbound( - request_id, - handler - .read(params) - .await - .map(ExecServerResponseMessage::Read), - ), - ExecServerRequest::Write { request_id, params } => request_outbound( - request_id, - handler - .write(params) - .await - .map(ExecServerResponseMessage::Write), - ), - ExecServerRequest::Terminate { request_id, params } => request_outbound( - request_id, - handler - .terminate(params) - .await - .map(ExecServerResponseMessage::Terminate), - ), - ExecServerRequest::FsReadFile { request_id, params } => request_outbound( - request_id, - handler - .fs_read_file(params) - .await - .map(ExecServerResponseMessage::FsReadFile), - ), - ExecServerRequest::FsWriteFile { request_id, params } => request_outbound( - request_id, - handler - .fs_write_file(params) - .await - .map(ExecServerResponseMessage::FsWriteFile), - ), - ExecServerRequest::FsCreateDirectory { request_id, params } => request_outbound( - request_id, - handler - .fs_create_directory(params) - .await - .map(ExecServerResponseMessage::FsCreateDirectory), - ), - ExecServerRequest::FsGetMetadata { request_id, params } => request_outbound( - request_id, - handler - .fs_get_metadata(params) - .await - .map(ExecServerResponseMessage::FsGetMetadata), - ), - ExecServerRequest::FsReadDirectory { request_id, params } => request_outbound( - request_id, - handler - .fs_read_directory(params) - .await - .map(ExecServerResponseMessage::FsReadDirectory), - ), - ExecServerRequest::FsRemove { request_id, params } => request_outbound( - request_id, - handler - .fs_remove(params) - .await - .map(ExecServerResponseMessage::FsRemove), - ), - ExecServerRequest::FsCopy { request_id, params } => request_outbound( - request_id, - handler - .fs_copy(params) - .await - .map(ExecServerResponseMessage::FsCopy), - ), - }; - outgoing_tx - .send(outbound) - .await - .map_err(|_| "outbound channel closed".to_string()) + JSONRPCMessage::Request(request) => Ok(Some(dispatch_request(handler, request).await)), + JSONRPCMessage::Notification(notification) => { + handle_notification(handler, notification)?; + Ok(None) } - ExecServerInboundMessage::Notification(ExecServerClientNotification::Initialized) => { - handler.initialized() + JSONRPCMessage::Response(response) => Err(format!( + "unexpected client response for request id {:?}", + response.id + )), + JSONRPCMessage::Error(error) => Err(format!( + "unexpected client error for request id {:?}", + error.id + )), + } +} + +async fn dispatch_request( + handler: &mut ExecServerHandler, + request: JSONRPCRequest, +) -> JSONRPCMessage { + let JSONRPCRequest { + id, + method, + params, + trace: _, + } = request; + let params = params.unwrap_or(serde_json::Value::Null); + + match method.as_str() { + INITIALIZE_METHOD => request_response( + id, + parse_params(params).and_then(|params| handler.initialize(params)), + ), + EXEC_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.exec(params)).await, + ), + EXEC_READ_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.read(params)).await, + ), + EXEC_WRITE_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.write(params)).await, + ), + EXEC_TERMINATE_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.terminate(params)).await, + ), + FS_READ_FILE_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.fs_read_file(params)).await, + ), + FS_WRITE_FILE_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.fs_write_file(params)).await, + ), + FS_CREATE_DIRECTORY_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.fs_create_directory(params)).await, + ), + FS_GET_METADATA_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.fs_get_metadata(params)).await, + ), + FS_READ_DIRECTORY_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.fs_read_directory(params)).await, + ), + FS_REMOVE_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.fs_remove(params)).await, + ), + FS_COPY_METHOD => request_response( + id, + dispatch_async_request(params, |params| handler.fs_copy(params)).await, + ), + other => jsonrpc_error_response(id, invalid_request(format!("unknown method: {other}"))), + } +} + +fn handle_notification( + handler: &mut ExecServerHandler, + notification: JSONRPCNotification, +) -> Result<(), String> { + match notification.method.as_str() { + INITIALIZED_METHOD => handler.initialized(), + other => Err(format!("unexpected notification method: {other}")), + } +} + +fn parse_params

(params: serde_json::Value) -> Result +where + P: DeserializeOwned, +{ + serde_json::from_value(params).map_err(|err| invalid_params(err.to_string())) +} + +async fn dispatch_async_request( + params: serde_json::Value, + f: F, +) -> Result +where + P: DeserializeOwned, + F: FnOnce(P) -> Fut, + Fut: std::future::Future>, +{ + match parse_params(params) { + Ok(params) => f(params).await, + Err(err) => Err(err), + } +} + +fn request_response( + request_id: RequestId, + result: Result, +) -> JSONRPCMessage +where + T: Serialize, +{ + match result.and_then(serialize_response) { + Ok(result) => JSONRPCMessage::Response(JSONRPCResponse { + id: request_id, + result, + }), + Err(error) => JSONRPCMessage::Error(JSONRPCError { + id: request_id, + error, + }), + } +} + +fn serialize_response(response: T) -> Result +where + T: Serialize, +{ + serde_json::to_value(response).map_err(|err| internal_error(err.to_string())) +} + +fn jsonrpc_error_response(request_id: RequestId, error: JSONRPCErrorError) -> JSONRPCMessage { + JSONRPCMessage::Error(JSONRPCError { + id: request_id, + error, + }) +} + +fn notification_message( + notification: ExecServerServerNotification, +) -> Result { + match notification { + ExecServerServerNotification::OutputDelta(params) => { + typed_notification(EXEC_OUTPUT_DELTA_METHOD, params) + } + ExecServerServerNotification::Exited(params) => { + typed_notification(EXEC_EXITED_METHOD, params) } } } -fn request_outbound( - request_id: codex_app_server_protocol::RequestId, - result: Result, -) -> ExecServerOutboundMessage { - match result { - Ok(response) => ExecServerOutboundMessage::Response { - request_id, - response, - }, - Err(error) => ExecServerOutboundMessage::Error { request_id, error }, +fn typed_notification(method: &str, params: T) -> Result +where + T: Serialize, +{ + Ok(JSONRPCMessage::Notification(JSONRPCNotification { + method: method.to_string(), + params: Some(serde_json::to_value(params)?), + })) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use pretty_assertions::assert_eq; + + use super::dispatch_request; + use super::handle_connection_message; + use super::notification_message; + use crate::protocol::EXEC_METHOD; + use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; + use crate::protocol::ExecOutputDeltaNotification; + use crate::protocol::ExecOutputStream; + use crate::protocol::INITIALIZE_METHOD; + use crate::protocol::INITIALIZED_METHOD; + use crate::protocol::InitializeParams; + use crate::protocol::PROTOCOL_VERSION; + use crate::server::ExecServerHandler; + use crate::server::ExecServerServerNotification; + use codex_app_server_protocol::JSONRPCError; + use codex_app_server_protocol::JSONRPCMessage; + use codex_app_server_protocol::JSONRPCNotification; + use codex_app_server_protocol::JSONRPCRequest; + use codex_app_server_protocol::JSONRPCResponse; + use codex_app_server_protocol::RequestId; + + #[tokio::test] + async fn dispatch_initialize_returns_initialize_response() { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(notification_tx, None); + + let message = dispatch_request( + &mut handler, + JSONRPCRequest { + id: RequestId::Integer(1), + method: INITIALIZE_METHOD.to_string(), + params: Some( + serde_json::to_value(InitializeParams { + client_name: "test".to_string(), + auth_token: None, + }) + .expect("serialize initialize params"), + ), + trace: None, + }, + ) + .await; + + assert_eq!( + message, + JSONRPCMessage::Response(JSONRPCResponse { + id: RequestId::Integer(1), + result: serde_json::json!({ + "protocolVersion": PROTOCOL_VERSION, + }), + }) + ); + } + + #[tokio::test] + async fn dispatch_exec_returns_invalid_request_before_initialize() { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(notification_tx, None); + + let message = dispatch_request( + &mut handler, + JSONRPCRequest { + id: RequestId::Integer(7), + method: EXEC_METHOD.to_string(), + params: Some(serde_json::json!({ + "processId": "proc-1", + "argv": ["bash", "-lc", "true"], + "cwd": std::env::current_dir().expect("cwd"), + "env": HashMap::::new(), + "tty": true, + "arg0": null, + "sandbox": null, + })), + trace: None, + }, + ) + .await; + + let JSONRPCMessage::Error(JSONRPCError { id, error }) = message else { + panic!("expected invalid-request error"); + }; + assert_eq!(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 initialized_notification_before_initialize_is_protocol_error() { + let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1); + let mut handler = ExecServerHandler::new(notification_tx, None); + + let err = handle_connection_message( + &mut handler, + JSONRPCMessage::Notification(JSONRPCNotification { + method: INITIALIZED_METHOD.to_string(), + params: Some(serde_json::json!({})), + }), + ) + .await + .expect_err("expected early initialized to fail"); + + assert_eq!( + err, + "received `initialized` notification before `initialize`" + ); + } + + #[test] + fn notification_message_serializes_process_output() { + let message = notification_message(ExecServerServerNotification::OutputDelta( + ExecOutputDeltaNotification { + process_id: "proc-1".to_string(), + stream: ExecOutputStream::Stdout, + chunk: b"hello".to_vec().into(), + }, + )) + .expect("serialize notification"); + + assert_eq!( + message, + JSONRPCMessage::Notification(JSONRPCNotification { + method: EXEC_OUTPUT_DELTA_METHOD.to_string(), + params: Some(serde_json::json!({ + "processId": "proc-1", + "stream": "stdout", + "chunk": "aGVsbG8=", + })), + }) + ); } } diff --git a/codex-rs/exec-server/src/server/routing.rs b/codex-rs/exec-server/src/server/routing.rs deleted file mode 100644 index 7d523b49c9..0000000000 --- a/codex-rs/exec-server/src/server/routing.rs +++ /dev/null @@ -1,585 +0,0 @@ -use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCErrorError; -use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCRequest; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; -use serde::de::DeserializeOwned; - -use crate::protocol::EXEC_EXITED_METHOD; -use crate::protocol::EXEC_METHOD; -use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; -use crate::protocol::EXEC_READ_METHOD; -use crate::protocol::EXEC_TERMINATE_METHOD; -use crate::protocol::EXEC_WRITE_METHOD; -use crate::protocol::ExecExitedNotification; -use crate::protocol::ExecOutputDeltaNotification; -use crate::protocol::ExecParams; -use crate::protocol::ExecResponse; -use crate::protocol::FS_COPY_METHOD; -use crate::protocol::FS_CREATE_DIRECTORY_METHOD; -use crate::protocol::FS_GET_METADATA_METHOD; -use crate::protocol::FS_READ_DIRECTORY_METHOD; -use crate::protocol::FS_READ_FILE_METHOD; -use crate::protocol::FS_REMOVE_METHOD; -use crate::protocol::FS_WRITE_FILE_METHOD; -use crate::protocol::INITIALIZE_METHOD; -use crate::protocol::INITIALIZED_METHOD; -use crate::protocol::InitializeParams; -use crate::protocol::InitializeResponse; -use crate::protocol::ReadParams; -use crate::protocol::ReadResponse; -use crate::protocol::TerminateParams; -use crate::protocol::TerminateResponse; -use crate::protocol::WriteParams; -use crate::protocol::WriteResponse; -use codex_app_server_protocol::FsCopyParams; -use codex_app_server_protocol::FsCopyResponse; -use codex_app_server_protocol::FsCreateDirectoryParams; -use codex_app_server_protocol::FsCreateDirectoryResponse; -use codex_app_server_protocol::FsGetMetadataParams; -use codex_app_server_protocol::FsGetMetadataResponse; -use codex_app_server_protocol::FsReadDirectoryParams; -use codex_app_server_protocol::FsReadDirectoryResponse; -use codex_app_server_protocol::FsReadFileParams; -use codex_app_server_protocol::FsReadFileResponse; -use codex_app_server_protocol::FsRemoveParams; -use codex_app_server_protocol::FsRemoveResponse; -use codex_app_server_protocol::FsWriteFileParams; -use codex_app_server_protocol::FsWriteFileResponse; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum ExecServerInboundMessage { - Request(ExecServerRequest), - Notification(ExecServerClientNotification), -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum ExecServerRequest { - Initialize { - request_id: RequestId, - params: InitializeParams, - }, - Exec { - request_id: RequestId, - params: ExecParams, - }, - Read { - request_id: RequestId, - params: ReadParams, - }, - Write { - request_id: RequestId, - params: WriteParams, - }, - Terminate { - request_id: RequestId, - params: TerminateParams, - }, - FsReadFile { - request_id: RequestId, - params: FsReadFileParams, - }, - FsWriteFile { - request_id: RequestId, - params: FsWriteFileParams, - }, - FsCreateDirectory { - request_id: RequestId, - params: FsCreateDirectoryParams, - }, - FsGetMetadata { - request_id: RequestId, - params: FsGetMetadataParams, - }, - FsReadDirectory { - request_id: RequestId, - params: FsReadDirectoryParams, - }, - FsRemove { - request_id: RequestId, - params: FsRemoveParams, - }, - FsCopy { - request_id: RequestId, - params: FsCopyParams, - }, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum ExecServerClientNotification { - Initialized, -} - -#[derive(Debug, Clone, PartialEq)] -pub(crate) enum ExecServerOutboundMessage { - Response { - request_id: RequestId, - response: ExecServerResponseMessage, - }, - Error { - request_id: RequestId, - error: JSONRPCErrorError, - }, - Notification(ExecServerServerNotification), -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum ExecServerResponseMessage { - Initialize(InitializeResponse), - Exec(ExecResponse), - Read(ReadResponse), - Write(WriteResponse), - Terminate(TerminateResponse), - FsReadFile(FsReadFileResponse), - FsWriteFile(FsWriteFileResponse), - FsCreateDirectory(FsCreateDirectoryResponse), - FsGetMetadata(FsGetMetadataResponse), - FsReadDirectory(FsReadDirectoryResponse), - FsRemove(FsRemoveResponse), - FsCopy(FsCopyResponse), -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum ExecServerServerNotification { - OutputDelta(ExecOutputDeltaNotification), - Exited(ExecExitedNotification), -} - -#[derive(Debug, Clone, PartialEq)] -pub(crate) enum RoutedExecServerMessage { - Inbound(ExecServerInboundMessage), - ImmediateOutbound(ExecServerOutboundMessage), -} - -pub(crate) fn route_jsonrpc_message( - message: JSONRPCMessage, -) -> Result { - match message { - JSONRPCMessage::Request(request) => route_request(request), - JSONRPCMessage::Notification(notification) => route_notification(notification), - JSONRPCMessage::Response(response) => Err(format!( - "unexpected client response for request id {:?}", - response.id - )), - JSONRPCMessage::Error(error) => Err(format!( - "unexpected client error for request id {:?}", - error.id - )), - } -} - -pub(crate) fn encode_outbound_message( - message: ExecServerOutboundMessage, -) -> Result { - match message { - ExecServerOutboundMessage::Response { - request_id, - response, - } => Ok(JSONRPCMessage::Response(JSONRPCResponse { - id: request_id, - result: serialize_response(response)?, - })), - ExecServerOutboundMessage::Error { request_id, error } => { - Ok(JSONRPCMessage::Error(JSONRPCError { - id: request_id, - error, - })) - } - ExecServerOutboundMessage::Notification(notification) => Ok(JSONRPCMessage::Notification( - serialize_notification(notification)?, - )), - } -} - -pub(crate) fn invalid_request(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32600, - data: None, - message, - } -} - -pub(crate) fn invalid_params(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32602, - data: None, - message, - } -} - -pub(crate) fn internal_error(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32603, - data: None, - message, - } -} - -fn route_request(request: JSONRPCRequest) -> Result { - match request.method.as_str() { - INITIALIZE_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::Initialize { request_id, params } - })), - EXEC_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::Exec { request_id, params } - })), - EXEC_READ_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::Read { request_id, params } - })), - EXEC_WRITE_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::Write { request_id, params } - })), - EXEC_TERMINATE_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::Terminate { request_id, params } - })), - FS_READ_FILE_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::FsReadFile { request_id, params } - })), - FS_WRITE_FILE_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::FsWriteFile { request_id, params } - })), - FS_CREATE_DIRECTORY_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::FsCreateDirectory { request_id, params } - })), - FS_GET_METADATA_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::FsGetMetadata { request_id, params } - })), - FS_READ_DIRECTORY_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::FsReadDirectory { request_id, params } - })), - FS_REMOVE_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::FsRemove { request_id, params } - })), - FS_COPY_METHOD => Ok(parse_request_params(request, |request_id, params| { - ExecServerRequest::FsCopy { request_id, params } - })), - other => Ok(RoutedExecServerMessage::ImmediateOutbound( - ExecServerOutboundMessage::Error { - request_id: request.id, - error: invalid_request(format!("unknown method: {other}")), - }, - )), - } -} - -fn route_notification( - notification: JSONRPCNotification, -) -> Result { - match notification.method.as_str() { - INITIALIZED_METHOD => Ok(RoutedExecServerMessage::Inbound( - ExecServerInboundMessage::Notification(ExecServerClientNotification::Initialized), - )), - other => Err(format!("unexpected notification method: {other}")), - } -} - -fn parse_request_params(request: JSONRPCRequest, build: F) -> RoutedExecServerMessage -where - P: DeserializeOwned, - F: FnOnce(RequestId, P) -> ExecServerRequest, -{ - let request_id = request.id; - match serde_json::from_value::

(request.params.unwrap_or(serde_json::Value::Null)) { - Ok(params) => RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Request(build( - request_id, params, - ))), - Err(err) => RoutedExecServerMessage::ImmediateOutbound(ExecServerOutboundMessage::Error { - request_id, - error: invalid_params(err.to_string()), - }), - } -} - -fn serialize_response( - response: ExecServerResponseMessage, -) -> Result { - match response { - ExecServerResponseMessage::Initialize(response) => serde_json::to_value(response), - ExecServerResponseMessage::Exec(response) => serde_json::to_value(response), - ExecServerResponseMessage::Read(response) => serde_json::to_value(response), - ExecServerResponseMessage::Write(response) => serde_json::to_value(response), - ExecServerResponseMessage::Terminate(response) => serde_json::to_value(response), - ExecServerResponseMessage::FsReadFile(response) => serde_json::to_value(response), - ExecServerResponseMessage::FsWriteFile(response) => serde_json::to_value(response), - ExecServerResponseMessage::FsCreateDirectory(response) => serde_json::to_value(response), - ExecServerResponseMessage::FsGetMetadata(response) => serde_json::to_value(response), - ExecServerResponseMessage::FsReadDirectory(response) => serde_json::to_value(response), - ExecServerResponseMessage::FsRemove(response) => serde_json::to_value(response), - ExecServerResponseMessage::FsCopy(response) => serde_json::to_value(response), - } -} - -fn serialize_notification( - notification: ExecServerServerNotification, -) -> Result { - match notification { - ExecServerServerNotification::OutputDelta(params) => Ok(JSONRPCNotification { - method: EXEC_OUTPUT_DELTA_METHOD.to_string(), - params: Some(serde_json::to_value(params)?), - }), - ExecServerServerNotification::Exited(params) => Ok(JSONRPCNotification { - method: EXEC_EXITED_METHOD.to_string(), - params: Some(serde_json::to_value(params)?), - }), - } -} - -#[cfg(test)] -mod tests { - use pretty_assertions::assert_eq; - use serde_json::json; - - use super::ExecServerClientNotification; - use super::ExecServerInboundMessage; - use super::ExecServerOutboundMessage; - use super::ExecServerRequest; - use super::ExecServerResponseMessage; - use super::ExecServerServerNotification; - use super::RoutedExecServerMessage; - use super::encode_outbound_message; - use super::route_jsonrpc_message; - use crate::protocol::EXEC_EXITED_METHOD; - use crate::protocol::EXEC_METHOD; - 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; - use codex_app_server_protocol::JSONRPCMessage; - use codex_app_server_protocol::JSONRPCNotification; - use codex_app_server_protocol::JSONRPCRequest; - use codex_app_server_protocol::JSONRPCResponse; - use codex_app_server_protocol::RequestId; - - #[test] - fn routes_initialize_requests_to_typed_variants() { - let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(1), - method: INITIALIZE_METHOD.to_string(), - params: Some(json!({ "clientName": "test-client" })), - trace: None, - })) - .expect("initialize request should route"); - - assert_eq!( - routed, - RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Request( - ExecServerRequest::Initialize { - request_id: RequestId::Integer(1), - params: InitializeParams { - client_name: "test-client".to_string(), - }, - }, - )) - ); - } - - #[test] - fn malformed_exec_params_return_immediate_error_outbound() { - let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(2), - method: EXEC_METHOD.to_string(), - params: Some(json!({ "processId": "proc-1" })), - trace: None, - })) - .expect("exec request should route"); - - let RoutedExecServerMessage::ImmediateOutbound(ExecServerOutboundMessage::Error { - request_id, - error, - }) = routed - else { - panic!("expected invalid-params error outbound"); - }; - assert_eq!(request_id, RequestId::Integer(2)); - assert_eq!(error.code, -32602); - } - - #[test] - fn routes_initialized_notifications_to_typed_variants() { - let routed = route_jsonrpc_message(JSONRPCMessage::Notification(JSONRPCNotification { - method: INITIALIZED_METHOD.to_string(), - params: Some(json!({})), - })) - .expect("initialized notification should route"); - - assert_eq!( - routed, - RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Notification( - ExecServerClientNotification::Initialized, - )) - ); - } - - #[test] - fn serializes_typed_notifications_back_to_jsonrpc() { - let message = encode_outbound_message(ExecServerOutboundMessage::Notification( - ExecServerServerNotification::Exited(ExecExitedNotification { - process_id: "proc-1".to_string(), - exit_code: 0, - }), - )) - .expect("notification should serialize"); - - assert_eq!( - message, - JSONRPCMessage::Notification(JSONRPCNotification { - method: EXEC_EXITED_METHOD.to_string(), - params: Some(json!({ - "processId": "proc-1", - "exitCode": 0, - })), - }) - ); - } - - #[test] - fn serializes_typed_responses_back_to_jsonrpc() { - let message = encode_outbound_message(ExecServerOutboundMessage::Response { - request_id: RequestId::Integer(3), - response: ExecServerResponseMessage::Exec(ExecResponse { - process_id: "proc-1".to_string(), - }), - }) - .expect("response should serialize"); - - assert_eq!( - message, - JSONRPCMessage::Response(codex_app_server_protocol::JSONRPCResponse { - id: RequestId::Integer(3), - result: json!({ - "processId": "proc-1", - }), - }) - ); - } - - #[test] - fn routes_exec_requests_with_typed_params() { - 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, - })), - trace: None, - })) - .expect("exec request 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: 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, - }), - } - ); - } - - #[test] - fn unknown_request_methods_return_immediate_invalid_request_errors() { - let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(5), - method: "process/unknown".to_string(), - params: Some(json!({})), - trace: None, - })) - .expect("unknown request should still route"); - - assert_eq!( - routed, - RoutedExecServerMessage::ImmediateOutbound(ExecServerOutboundMessage::Error { - request_id: RequestId::Integer(5), - error: super::invalid_request("unknown method: process/unknown".to_string()), - }) - ); - } - - #[test] - fn unexpected_client_notifications_are_rejected() { - let err = route_jsonrpc_message(JSONRPCMessage::Notification(JSONRPCNotification { - method: "process/output".to_string(), - params: Some(json!({})), - })) - .expect_err("unexpected client notification should fail"); - - assert_eq!(err, "unexpected notification method: process/output"); - } - - #[test] - fn unexpected_client_responses_are_rejected() { - let err = route_jsonrpc_message(JSONRPCMessage::Response(JSONRPCResponse { - id: RequestId::Integer(6), - result: json!({}), - })) - .expect_err("unexpected client response should fail"); - - assert_eq!(err, "unexpected client response for request id Integer(6)"); - } -} diff --git a/codex-rs/exec-server/src/server/transport.rs b/codex-rs/exec-server/src/server/transport.rs index b653c0b79b..c98aafa17a 100644 --- a/codex-rs/exec-server/src/server/transport.rs +++ b/codex-rs/exec-server/src/server/transport.rs @@ -6,11 +6,11 @@ use tokio_tungstenite::accept_async; use tracing::warn; use crate::connection::JsonRpcConnection; +use crate::server::ExecServerConfig; use crate::server::processor::run_connection; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum ExecServerTransport { - Stdio, WebSocket { bind_address: SocketAddr }, } @@ -25,7 +25,7 @@ impl std::fmt::Display for ExecServerTransportParseError { match self { ExecServerTransportParseError::UnsupportedListenUrl(listen_url) => write!( f, - "unsupported --listen URL `{listen_url}`; expected `stdio://` or `ws://IP:PORT`" + "unsupported --listen URL `{listen_url}`; expected `ws://IP:PORT`" ), ExecServerTransportParseError::InvalidWebSocketListenUrl(listen_url) => write!( f, @@ -38,13 +38,9 @@ impl std::fmt::Display for ExecServerTransportParseError { impl std::error::Error for ExecServerTransportParseError {} impl ExecServerTransport { - pub const DEFAULT_LISTEN_URL: &str = "stdio://"; + pub const DEFAULT_LISTEN_URL: &str = "ws://127.0.0.1:0"; pub fn from_listen_url(listen_url: &str) -> Result { - if listen_url == Self::DEFAULT_LISTEN_URL { - return Ok(Self::Stdio); - } - if let Some(socket_addr) = listen_url.strip_prefix("ws://") { let bind_address = socket_addr.parse::().map_err(|_| { ExecServerTransportParseError::InvalidWebSocketListenUrl(listen_url.to_string()) @@ -58,6 +54,16 @@ impl ExecServerTransport { } } +impl Default for ExecServerTransport { + fn default() -> Self { + Self::WebSocket { + bind_address: "127.0.0.1:0".parse().unwrap_or_else(|err| { + panic!("default exec-server websocket bind address should parse: {err}") + }), + } + } +} + impl FromStr for ExecServerTransport { type Err = ExecServerTransportParseError; @@ -68,25 +74,18 @@ impl FromStr for ExecServerTransport { pub(crate) async fn run_transport( transport: ExecServerTransport, + config: ExecServerConfig, ) -> Result<(), Box> { match transport { - ExecServerTransport::Stdio => { - run_connection(JsonRpcConnection::from_stdio( - tokio::io::stdin(), - tokio::io::stdout(), - "exec-server stdio".to_string(), - )) - .await; - Ok(()) - } ExecServerTransport::WebSocket { bind_address } => { - run_websocket_listener(bind_address).await + run_websocket_listener(bind_address, config).await } } } async fn run_websocket_listener( bind_address: SocketAddr, + config: ExecServerConfig, ) -> Result<(), Box> { let listener = TcpListener::bind(bind_address).await?; let local_addr = listener.local_addr()?; @@ -94,13 +93,17 @@ async fn run_websocket_listener( loop { let (stream, peer_addr) = listener.accept().await?; + let config = config.clone(); tokio::spawn(async move { match accept_async(stream).await { Ok(websocket) => { - run_connection(JsonRpcConnection::from_websocket( - websocket, - format!("exec-server websocket {peer_addr}"), - )) + run_connection( + JsonRpcConnection::from_websocket( + websocket, + format!("exec-server websocket {peer_addr}"), + ), + config, + ) .await; } Err(err) => { @@ -124,14 +127,6 @@ mod tests { use super::ExecServerTransport; - #[test] - fn exec_server_transport_parses_stdio_listen_url() { - let transport = - ExecServerTransport::from_listen_url(ExecServerTransport::DEFAULT_LISTEN_URL) - .expect("stdio listen URL should parse"); - assert_eq!(transport, ExecServerTransport::Stdio); - } - #[test] fn exec_server_transport_parses_websocket_listen_url() { let transport = ExecServerTransport::from_listen_url("ws://127.0.0.1:1234") @@ -160,7 +155,7 @@ mod tests { .expect_err("unsupported scheme should fail"); assert_eq!( err.to_string(), - "unsupported --listen URL `http://127.0.0.1:1234`; expected `stdio://` or `ws://IP:PORT`" + "unsupported --listen URL `http://127.0.0.1:1234`; expected `ws://IP:PORT`" ); } }