diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index 2a72d1f311..5c487c584b 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -17,11 +17,13 @@ unified-exec in this PR; it is only the standalone transport layer. The internal shape is intentionally closer to `app-server` than the first cut: - transport adapters are separate from the per-connection request processor +- JSON-RPC route matching is separate from the stateful exec handler - the client only speaks the protocol; it does not spawn a server subprocess - local child-process launch is handled by a separate helper/factory layer That split is meant to leave reusable seams if exec-server and app-server later -share transport or JSON-RPC connection utilities. +share transport or JSON-RPC connection utilities. It also keeps the core +handler testable without the RPC server implementation itself. ## Transport diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index 2a7ee08d7d..f979ef8a2b 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -1,4 +1,6 @@ +mod handler; mod processor; +mod routing; mod transport; pub use transport::ExecServerTransport; diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs new file mode 100644 index 0000000000..baccc9a281 --- /dev/null +++ b/codex-rs/exec-server/src/server/handler.rs @@ -0,0 +1,562 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use codex_utils_pty::ExecCommandSession; +use codex_utils_pty::TerminalSize; +use tokio::sync::Mutex; +use tokio::sync::mpsc; + +use crate::protocol::ExecExitedNotification; +use crate::protocol::ExecOutputDeltaNotification; +use crate::protocol::ExecOutputStream; +use crate::protocol::ExecResponse; +use crate::protocol::InitializeResponse; +use crate::protocol::PROTOCOL_VERSION; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteResponse; +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::ExecServerServerNotification; +use crate::server::routing::internal_error; +use crate::server::routing::invalid_params; +use crate::server::routing::invalid_request; + +struct RunningProcess { + session: ExecCommandSession, + tty: bool, +} + +pub(crate) struct ExecServerHandler { + outbound_tx: mpsc::Sender, + processes: Arc>>, + initialize_requested: bool, + initialized: bool, +} + +impl ExecServerHandler { + pub(crate) fn new(outbound_tx: mpsc::Sender) -> Self { + Self { + outbound_tx, + processes: Arc::new(Mutex::new(HashMap::new())), + initialize_requested: false, + initialized: false, + } + } + + pub(crate) async fn shutdown(&self) { + let remaining = { + let mut processes = self.processes.lock().await; + processes + .drain() + .map(|(_, process)| process) + .collect::>() + }; + for process in remaining { + process.session.terminate(); + } + } + + pub(crate) async fn handle_message( + &mut self, + message: ExecServerInboundMessage, + ) -> Result<(), String> { + match message { + ExecServerInboundMessage::Request(request) => self.handle_request(request).await, + ExecServerInboundMessage::Notification(notification) => { + self.handle_notification(notification) + } + } + } + + async fn handle_request(&mut self, request: ExecServerRequest) -> Result<(), String> { + match request { + ExecServerRequest::Initialize { request_id, .. } => { + let result = self + .handle_initialize_request() + .map(ExecServerResponseMessage::Initialize); + self.send_request_result(request_id, result).await; + } + ExecServerRequest::Exec { request_id, params } => { + self.send_request_result( + request_id, + match self.require_initialized() { + Ok(()) => self + .handle_exec_request(params) + .await + .map(ExecServerResponseMessage::Exec), + Err(err) => Err(err), + }, + ) + .await; + } + ExecServerRequest::Write { request_id, params } => { + self.send_request_result( + request_id, + match self.require_initialized() { + Ok(()) => self + .handle_write_request(params) + .await + .map(ExecServerResponseMessage::Write), + Err(err) => Err(err), + }, + ) + .await; + } + ExecServerRequest::Terminate { request_id, params } => { + self.send_request_result( + request_id, + match self.require_initialized() { + Ok(()) => self + .handle_terminate_request(params) + .await + .map(ExecServerResponseMessage::Terminate), + Err(err) => Err(err), + }, + ) + .await; + } + } + Ok(()) + } + + fn handle_notification( + &mut self, + notification: ExecServerClientNotification, + ) -> Result<(), String> { + match notification { + ExecServerClientNotification::Initialized => { + if !self.initialize_requested { + return Err("received `initialized` notification before `initialize`".into()); + } + self.initialized = true; + Ok(()) + } + } + } + + fn handle_initialize_request( + &mut self, + ) -> Result { + if self.initialize_requested { + return Err(invalid_request( + "initialize may only be sent once per connection".to_string(), + )); + } + self.initialize_requested = true; + Ok(InitializeResponse { + protocol_version: PROTOCOL_VERSION.to_string(), + }) + } + + fn require_initialized(&self) -> Result<(), codex_app_server_protocol::JSONRPCErrorError> { + if !self.initialize_requested { + return Err(invalid_request( + "client must call initialize before using exec methods".to_string(), + )); + } + if !self.initialized { + return Err(invalid_request( + "client must send initialized before using exec methods".to_string(), + )); + } + Ok(()) + } + + async fn handle_exec_request( + &self, + params: crate::protocol::ExecParams, + ) -> Result { + let (program, args) = params + .argv + .split_first() + .ok_or_else(|| invalid_params("argv must not be empty".to_string()))?; + + let spawned = if params.tty { + codex_utils_pty::spawn_pty_process( + program, + args, + params.cwd.as_path(), + ¶ms.env, + ¶ms.arg0, + TerminalSize::default(), + ) + .await + } else { + codex_utils_pty::spawn_pipe_process_no_stdin( + program, + args, + params.cwd.as_path(), + ¶ms.env, + ¶ms.arg0, + ) + .await + } + .map_err(|err| internal_error(err.to_string()))?; + + let process_id = params.process_id.clone(); + { + let mut process_map = self.processes.lock().await; + if process_map.contains_key(&process_id) { + spawned.session.terminate(); + return Err(invalid_request(format!( + "process {process_id} already exists" + ))); + } + process_map.insert( + process_id.clone(), + RunningProcess { + session: spawned.session, + tty: params.tty, + }, + ); + } + + tokio::spawn(stream_output( + process_id.clone(), + ExecOutputStream::Stdout, + spawned.stdout_rx, + self.outbound_tx.clone(), + )); + tokio::spawn(stream_output( + process_id.clone(), + ExecOutputStream::Stderr, + spawned.stderr_rx, + self.outbound_tx.clone(), + )); + tokio::spawn(watch_exit( + process_id.clone(), + spawned.exit_rx, + self.outbound_tx.clone(), + Arc::clone(&self.processes), + )); + + Ok(ExecResponse { process_id }) + } + + async fn handle_write_request( + &self, + params: crate::protocol::WriteParams, + ) -> Result { + let writer_tx = { + 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)) + })?; + if !process.tty { + return Err(invalid_request(format!( + "stdin is closed for process {}", + params.process_id + ))); + } + process.session.writer_sender() + }; + + writer_tx + .send(params.chunk.into_inner()) + .await + .map_err(|_| internal_error("failed to write to process stdin".to_string()))?; + + Ok(WriteResponse { accepted: true }) + } + + async fn handle_terminate_request( + &self, + params: crate::protocol::TerminateParams, + ) -> Result { + let process = { + let mut process_map = self.processes.lock().await; + process_map.remove(¶ms.process_id) + }; + + Ok(if let Some(process) = process { + process.session.terminate(); + TerminateResponse { running: true } + } else { + TerminateResponse { running: false } + }) + } + + async fn send_request_result( + &self, + request_id: codex_app_server_protocol::RequestId, + result: Result, + ) { + let outbound = match result { + Ok(response) => ExecServerOutboundMessage::Response { + request_id, + response, + }, + Err(error) => ExecServerOutboundMessage::Error { request_id, error }, + }; + self.send_outbound(outbound).await; + } + + async fn send_outbound(&self, outbound: ExecServerOutboundMessage) { + let _ = self.outbound_tx.send(outbound).await; + } +} + +async fn stream_output( + process_id: String, + stream: ExecOutputStream, + mut receiver: tokio::sync::mpsc::Receiver>, + outbound_tx: mpsc::Sender, +) { + while let Some(chunk) = receiver.recv().await { + if outbound_tx + .send(ExecServerOutboundMessage::Notification( + ExecServerServerNotification::OutputDelta(ExecOutputDeltaNotification { + process_id: process_id.clone(), + stream, + chunk: chunk.into(), + }), + )) + .await + .is_err() + { + break; + } + } +} + +async fn watch_exit( + process_id: String, + exit_rx: tokio::sync::oneshot::Receiver, + outbound_tx: mpsc::Sender, + processes: Arc>>, +) { + let exit_code = exit_rx.await.unwrap_or(-1); + { + let mut processes = processes.lock().await; + processes.remove(&process_id); + } + let _ = outbound_tx + .send(ExecServerOutboundMessage::Notification( + ExecServerServerNotification::Exited(ExecExitedNotification { + process_id, + exit_code, + }), + )) + .await; +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::time::Duration; + + use pretty_assertions::assert_eq; + use tokio::time::timeout; + + use super::ExecServerHandler; + use crate::protocol::InitializeParams; + use crate::protocol::InitializeResponse; + use crate::protocol::PROTOCOL_VERSION; + 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" + ); + } +} diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index fa0081f5c9..e84a14eb8a 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,58 +1,56 @@ -use std::collections::HashMap; -use std::sync::Arc; - -use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCErrorError; -use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCNotification; -use codex_app_server_protocol::JSONRPCRequest; -use codex_app_server_protocol::JSONRPCResponse; -use codex_app_server_protocol::RequestId; -use codex_utils_pty::ExecCommandSession; -use codex_utils_pty::TerminalSize; -use tokio::sync::Mutex; use tokio::sync::mpsc; use tracing::debug; use tracing::warn; +use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; -use crate::protocol::EXEC_EXITED_METHOD; -use crate::protocol::EXEC_METHOD; -use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; -use crate::protocol::EXEC_TERMINATE_METHOD; -use crate::protocol::EXEC_WRITE_METHOD; -use crate::protocol::ExecExitedNotification; -use crate::protocol::ExecOutputDeltaNotification; -use crate::protocol::ExecOutputStream; -use crate::protocol::ExecParams; -use crate::protocol::ExecResponse; -use crate::protocol::INITIALIZE_METHOD; -use crate::protocol::INITIALIZED_METHOD; -use crate::protocol::InitializeResponse; -use crate::protocol::PROTOCOL_VERSION; -use crate::protocol::TerminateParams; -use crate::protocol::TerminateResponse; -use crate::protocol::WriteParams; -use crate::protocol::WriteResponse; - -struct RunningProcess { - session: ExecCommandSession, - tty: bool, -} +use crate::server::handler::ExecServerHandler; +use crate::server::routing::ExecServerOutboundMessage; +use crate::server::routing::RoutedExecServerMessage; +use crate::server::routing::encode_outbound_message; +use crate::server::routing::route_jsonrpc_message; pub(crate) async fn run_connection(connection: JsonRpcConnection) { - let (outgoing_tx, mut incoming_rx) = connection.into_parts(); - let mut processor = ExecServerConnectionProcessor::new(outgoing_tx); + let (json_outgoing_tx, mut incoming_rx) = connection.into_parts(); + let (outgoing_tx, mut outgoing_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let mut handler = ExecServerHandler::new(outgoing_tx.clone()); + + let outbound_task = tokio::spawn(async move { + while let Some(message) = outgoing_rx.recv().await { + let json_message = match encode_outbound_message(message) { + Ok(json_message) => json_message, + Err(err) => { + warn!("failed to serialize exec-server outbound message: {err}"); + break; + } + }; + if json_outgoing_tx.send(json_message).await.is_err() { + break; + } + } + }); while let Some(event) = incoming_rx.recv().await { match event { - JsonRpcConnectionEvent::Message(message) => { - if let Err(err) = processor.handle_message(message).await { + JsonRpcConnectionEvent::Message(message) => match route_jsonrpc_message(message) { + Ok(RoutedExecServerMessage::Inbound(message)) => { + if let Err(err) = handler.handle_message(message).await { + 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}"); break; } - } + }, JsonRpcConnectionEvent::Disconnected { reason } => { if let Some(reason) = reason { debug!("exec-server connection disconnected: {reason}"); @@ -62,545 +60,8 @@ pub(crate) async fn run_connection(connection: JsonRpcConnection) { } } - processor.shutdown().await; -} - -struct ExecServerConnectionProcessor { - outgoing_tx: mpsc::Sender, - processes: Arc>>, - initialize_requested: bool, - initialized: bool, -} - -impl ExecServerConnectionProcessor { - fn new(outgoing_tx: mpsc::Sender) -> Self { - Self { - outgoing_tx, - processes: Arc::new(Mutex::new(HashMap::new())), - initialize_requested: false, - initialized: false, - } - } - - async fn shutdown(&self) { - let remaining = { - let mut processes = self.processes.lock().await; - processes - .drain() - .map(|(_, process)| process) - .collect::>() - }; - for process in remaining { - process.session.terminate(); - } - } - - async fn handle_message(&mut self, message: JSONRPCMessage) -> Result<(), String> { - match message { - JSONRPCMessage::Request(request) => self.handle_request(request).await, - JSONRPCMessage::Notification(notification) => self.handle_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 - )), - } - } - - async fn handle_request(&mut self, request: JSONRPCRequest) -> Result<(), String> { - let response = match request.method.as_str() { - INITIALIZE_METHOD => self.handle_initialize_request(), - EXEC_METHOD => match self.require_initialized() { - Ok(()) => self.handle_exec_request(request.params).await, - Err(err) => Err(err), - }, - EXEC_WRITE_METHOD => match self.require_initialized() { - Ok(()) => self.handle_write_request(request.params).await, - Err(err) => Err(err), - }, - EXEC_TERMINATE_METHOD => match self.require_initialized() { - Ok(()) => self.handle_terminate_request(request.params).await, - Err(err) => Err(err), - }, - other => Err(invalid_request(format!("unknown method: {other}"))), - }; - - match response { - Ok(result) => { - self.send_response(JSONRPCResponse { - id: request.id, - result, - }) - .await; - } - Err(error) => { - self.send_error(request.id, error).await; - } - } - - Ok(()) - } - - fn handle_notification(&mut self, notification: JSONRPCNotification) -> Result<(), String> { - match notification.method.as_str() { - INITIALIZED_METHOD => { - if !self.initialize_requested { - return Err("received `initialized` notification before `initialize`".into()); - } - self.initialized = true; - Ok(()) - } - other => Err(format!("unexpected notification method: {other}")), - } - } - - fn handle_initialize_request(&mut self) -> Result { - if self.initialize_requested { - return Err(invalid_request( - "initialize may only be sent once per connection".to_string(), - )); - } - self.initialize_requested = true; - json_value(InitializeResponse { - protocol_version: PROTOCOL_VERSION.to_string(), - }) - } - - fn require_initialized(&self) -> Result<(), JSONRPCErrorError> { - if !self.initialize_requested { - return Err(invalid_request( - "client must call initialize before using exec methods".to_string(), - )); - } - if !self.initialized { - return Err(invalid_request( - "client must send initialized before using exec methods".to_string(), - )); - } - Ok(()) - } - - async fn handle_exec_request( - &self, - params: Option, - ) -> Result { - let params: ExecParams = serde_json::from_value(params.unwrap_or(serde_json::Value::Null)) - .map_err(|err| invalid_params(err.to_string()))?; - - let (program, args) = params - .argv - .split_first() - .ok_or_else(|| invalid_params("argv must not be empty".to_string()))?; - - let spawned = if params.tty { - codex_utils_pty::spawn_pty_process( - program, - args, - params.cwd.as_path(), - ¶ms.env, - ¶ms.arg0, - TerminalSize::default(), - ) - .await - } else { - codex_utils_pty::spawn_pipe_process_no_stdin( - program, - args, - params.cwd.as_path(), - ¶ms.env, - ¶ms.arg0, - ) - .await - } - .map_err(|err| internal_error(err.to_string()))?; - - let process_id = params.process_id.clone(); - { - let mut process_map = self.processes.lock().await; - if process_map.contains_key(&process_id) { - spawned.session.terminate(); - return Err(invalid_request(format!( - "process {process_id} already exists" - ))); - } - process_map.insert( - process_id.clone(), - RunningProcess { - session: spawned.session, - tty: params.tty, - }, - ); - } - - tokio::spawn(stream_output( - process_id.clone(), - ExecOutputStream::Stdout, - spawned.stdout_rx, - self.outgoing_tx.clone(), - )); - tokio::spawn(stream_output( - process_id.clone(), - ExecOutputStream::Stderr, - spawned.stderr_rx, - self.outgoing_tx.clone(), - )); - tokio::spawn(watch_exit( - process_id.clone(), - spawned.exit_rx, - self.outgoing_tx.clone(), - Arc::clone(&self.processes), - )); - - json_value(ExecResponse { process_id }) - } - - async fn handle_write_request( - &self, - params: Option, - ) -> Result { - let params: WriteParams = serde_json::from_value(params.unwrap_or(serde_json::Value::Null)) - .map_err(|err| invalid_params(err.to_string()))?; - - let writer_tx = { - let process_map = self.processes.lock().await; - let process = process_map.get(¶ms.process_id).ok_or_else(|| { - invalid_request(format!("unknown process id {}", params.process_id)) - })?; - if !process.tty { - return Err(invalid_request(format!( - "stdin is closed for process {}", - params.process_id - ))); - } - process.session.writer_sender() - }; - - writer_tx - .send(params.chunk.into_inner()) - .await - .map_err(|_| internal_error("failed to write to process stdin".to_string()))?; - - json_value(WriteResponse { accepted: true }) - } - - async fn handle_terminate_request( - &self, - params: Option, - ) -> Result { - let params: TerminateParams = - serde_json::from_value(params.unwrap_or(serde_json::Value::Null)) - .map_err(|err| invalid_params(err.to_string()))?; - - let process = { - let mut process_map = self.processes.lock().await; - process_map.remove(¶ms.process_id) - }; - - if let Some(process) = process { - process.session.terminate(); - json_value(TerminateResponse { running: true }) - } else { - json_value(TerminateResponse { running: false }) - } - } - - async fn send_response(&self, response: JSONRPCResponse) { - let _ = self - .outgoing_tx - .send(JSONRPCMessage::Response(response)) - .await; - } - - async fn send_error(&self, id: RequestId, error: JSONRPCErrorError) { - let _ = self - .outgoing_tx - .send(JSONRPCMessage::Error(JSONRPCError { error, id })) - .await; - } -} - -async fn stream_output( - process_id: String, - stream: ExecOutputStream, - mut receiver: tokio::sync::mpsc::Receiver>, - outgoing_tx: mpsc::Sender, -) { - while let Some(chunk) = receiver.recv().await { - let notification = ExecOutputDeltaNotification { - process_id: process_id.clone(), - stream, - chunk: chunk.into(), - }; - if send_notification(&outgoing_tx, EXEC_OUTPUT_DELTA_METHOD, ¬ification) - .await - .is_err() - { - break; - } - } -} - -async fn watch_exit( - process_id: String, - exit_rx: tokio::sync::oneshot::Receiver, - outgoing_tx: mpsc::Sender, - processes: Arc>>, -) { - let exit_code = exit_rx.await.unwrap_or(-1); - { - let mut processes = processes.lock().await; - processes.remove(&process_id); - } - let _ = send_notification( - &outgoing_tx, - EXEC_EXITED_METHOD, - &ExecExitedNotification { - process_id, - exit_code, - }, - ) - .await; -} - -async fn send_notification( - outgoing_tx: &mpsc::Sender, - method: &str, - params: &T, -) -> Result<(), serde_json::Error> { - outgoing_tx - .send(JSONRPCMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(serde_json::to_value(params)?), - })) - .await - .map_err(|_| serde_json::Error::io(std::io::Error::other("connection closed"))) -} - -fn json_value(value: T) -> Result { - serde_json::to_value(value).map_err(|err| internal_error(err.to_string())) -} - -fn invalid_request(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32600, - data: None, - message, - } -} - -fn invalid_params(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32602, - data: None, - message, - } -} - -fn internal_error(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32603, - data: None, - message, - } -} - -#[cfg(test)] -mod tests { - use std::time::Duration; - - use pretty_assertions::assert_eq; - use serde_json::json; - use tokio::time::timeout; - - use super::ExecServerConnectionProcessor; - use crate::protocol::EXEC_METHOD; - use crate::protocol::INITIALIZE_METHOD; - use crate::protocol::INITIALIZED_METHOD; - use crate::protocol::PROTOCOL_VERSION; - use codex_app_server_protocol::JSONRPCMessage; - use codex_app_server_protocol::JSONRPCNotification; - use codex_app_server_protocol::JSONRPCRequest; - use codex_app_server_protocol::RequestId; - - fn request(id: i64, method: &str, params: serde_json::Value) -> JSONRPCMessage { - JSONRPCMessage::Request(JSONRPCRequest { - id: RequestId::Integer(id), - method: method.to_string(), - params: Some(params), - trace: None, - }) - } - - async fn recv_outgoing_json( - outgoing_rx: &mut tokio::sync::mpsc::Receiver, - ) -> serde_json::Value { - 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 processor output: {err}"), - }; - let message = match maybe_message { - Some(message) => message, - None => panic!("processor output channel closed unexpectedly"), - }; - serde_json::to_value(message) - .unwrap_or_else(|err| panic!("failed to serialize processor output: {err}")) - } - - #[tokio::test] - async fn initialize_response_reports_protocol_version() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(1); - let mut processor = ExecServerConnectionProcessor::new(outgoing_tx); - - if let Err(err) = processor - .handle_message(request( - 1, - INITIALIZE_METHOD, - json!({ "clientName": "test" }), - )) - .await - { - panic!("initialize should succeed: {err}"); - } - - let outgoing = recv_outgoing_json(&mut outgoing_rx).await; - assert_eq!( - outgoing, - json!({ - "id": 1, - "result": { - "protocolVersion": PROTOCOL_VERSION - } - }) - ); - } - - #[tokio::test] - async fn exec_methods_require_initialize() { - let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(1); - let mut processor = ExecServerConnectionProcessor::new(outgoing_tx); - - if let Err(err) = processor - .handle_message(request(7, EXEC_METHOD, json!({ "processId": "proc-1" }))) - .await - { - panic!("request handling should not fail the connection: {err}"); - } - - let outgoing = recv_outgoing_json(&mut outgoing_rx).await; - assert_eq!( - outgoing, - json!({ - "id": 7, - "error": { - "code": -32600, - "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 processor = ExecServerConnectionProcessor::new(outgoing_tx); - - if let Err(err) = processor - .handle_message(request( - 1, - INITIALIZE_METHOD, - json!({ "clientName": "test" }), - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outgoing_json(&mut outgoing_rx).await; - - if let Err(err) = processor - .handle_message(request(2, EXEC_METHOD, json!({ "processId": "proc-1" }))) - .await - { - panic!("request handling should not fail the connection: {err}"); - } - - let outgoing = recv_outgoing_json(&mut outgoing_rx).await; - assert_eq!( - outgoing, - json!({ - "id": 2, - "error": { - "code": -32600, - "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 processor = ExecServerConnectionProcessor::new(outgoing_tx); - - let result = processor - .handle_message(JSONRPCMessage::Notification(JSONRPCNotification { - method: INITIALIZED_METHOD.to_string(), - params: Some(json!({})), - })) - .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 processor = ExecServerConnectionProcessor::new(outgoing_tx); - - if let Err(err) = processor - .handle_message(request( - 1, - INITIALIZE_METHOD, - json!({ "clientName": "test" }), - )) - .await - { - panic!("initialize should succeed: {err}"); - } - let _ = recv_outgoing_json(&mut outgoing_rx).await; - - if let Err(err) = processor - .handle_message(request( - 2, - INITIALIZE_METHOD, - json!({ "clientName": "test" }), - )) - .await - { - panic!("duplicate initialize should not fail the connection: {err}"); - } - - let outgoing = recv_outgoing_json(&mut outgoing_rx).await; - assert_eq!( - outgoing, - json!({ - "id": 2, - "error": { - "code": -32600, - "message": "initialize may only be sent once per connection" - } - }) - ); - } + handler.shutdown().await; + drop(handler); + drop(outgoing_tx); + let _ = outbound_task.await; } diff --git a/codex-rs/exec-server/src/server/routing.rs b/codex-rs/exec-server/src/server/routing.rs new file mode 100644 index 0000000000..aba9c61c93 --- /dev/null +++ b/codex-rs/exec-server/src/server/routing.rs @@ -0,0 +1,400 @@ +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_TERMINATE_METHOD; +use crate::protocol::EXEC_WRITE_METHOD; +use crate::protocol::ExecExitedNotification; +use crate::protocol::ExecOutputDeltaNotification; +use crate::protocol::ExecParams; +use crate::protocol::ExecResponse; +use crate::protocol::INITIALIZE_METHOD; +use crate::protocol::INITIALIZED_METHOD; +use crate::protocol::InitializeParams; +use crate::protocol::InitializeResponse; +use crate::protocol::TerminateParams; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteParams; +use crate::protocol::WriteResponse; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(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, + }, + Write { + request_id: RequestId, + params: WriteParams, + }, + Terminate { + request_id: RequestId, + params: TerminateParams, + }, +} + +#[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), + Write(WriteResponse), + Terminate(TerminateResponse), +} + +#[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_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 } + })), + 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::Write(response) => serde_json::to_value(response), + ExecServerResponseMessage::Terminate(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::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::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, + } + ); + } +}