diff --git a/codex-rs/exec-server/src/bin/codex-exec-server.rs b/codex-rs/exec-server/src/bin/codex-exec-server.rs index 16df84d9b6..2638ecd1b7 100644 --- a/codex-rs/exec-server/src/bin/codex-exec-server.rs +++ b/codex-rs/exec-server/src/bin/codex-exec-server.rs @@ -1,5 +1,5 @@ use clap::Parser; -use codex_exec_server::ExecServerTransport; +use codex_exec_server::DEFAULT_LISTEN_URL; #[derive(Debug, Parser)] struct ExecServerArgs { @@ -8,15 +8,15 @@ struct ExecServerArgs { #[arg( long = "listen", value_name = "URL", - default_value = ExecServerTransport::DEFAULT_LISTEN_URL + default_value = DEFAULT_LISTEN_URL )] - listen: ExecServerTransport, + listen: String, } #[tokio::main] async fn main() { let args = ExecServerArgs::parse(); - if let Err(err) = codex_exec_server::run_main_with_transport(args.listen).await { + if let Err(err) = codex_exec_server::run_main_with_transport(&args.listen).await { eprintln!("{err}"); std::process::exit(1); } diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 12bf0e17f9..2ffcd30196 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -42,7 +42,8 @@ pub use protocol::TerminateParams; pub use protocol::TerminateResponse; pub use protocol::WriteParams; pub use protocol::WriteResponse; -pub use server::ExecServerTransport; +pub use server::DEFAULT_LISTEN_URL; pub use server::ExecServerTransportParseError; pub use server::run_main; +pub use server::run_main_with_listen_url; pub use server::run_main_with_transport; diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index a57e164919..8d79883c57 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -4,7 +4,6 @@ use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; -use std::pin::Pin; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCErrorError; @@ -197,6 +196,10 @@ impl RpcClient { break; } } + JsonRpcConnectionEvent::MalformedMessage { reason } => { + warn!("JSON-RPC client closing after malformed message: {reason}"); + break; + } JsonRpcConnectionEvent::Disconnected { reason } => { let _ = event_tx.send(RpcClientEvent::Disconnected { reason }).await; drain_pending(&pending_for_reader).await; @@ -442,217 +445,6 @@ async fn drain_pending(pending: &Mutex>) { } } -#[derive(Debug)] -pub(crate) enum RpcServerOutboundMessage { - Response { - request_id: RequestId, - result: Value, - }, - Error { - request_id: RequestId, - error: JSONRPCErrorError, - }, - Notification(JSONRPCNotification), -} - -impl RpcServerOutboundMessage { - fn response(request_id: RequestId, result: Result) -> Self { - match result { - Ok(result) => Self::Response { - request_id, - result, - }, - Err(error) => Self::Error { - request_id, - error, - }, - } - } -} - -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 method_not_found(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32601, - data: None, - message, - } -} - -pub(crate) fn internal_error(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32603, - data: None, - message, - } -} - -pub(crate) fn encode_server_message( - message: RpcServerOutboundMessage, -) -> Result { - Ok(match message { - RpcServerOutboundMessage::Response { request_id, result } => { - JSONRPCMessage::Response(JSONRPCResponse { id: request_id, result }) - } - RpcServerOutboundMessage::Error { request_id, error } => { - JSONRPCMessage::Error(JSONRPCError { id: request_id, error }) - } - RpcServerOutboundMessage::Notification(notification) => { - JSONRPCMessage::Notification(notification) - } - }) -} - -#[derive(Clone)] -pub(crate) struct RpcNotificationSender { - tx: mpsc::Sender, -} - -impl RpcNotificationSender { - pub(crate) fn new(tx: mpsc::Sender) -> Self { - Self { tx } - } - - pub(crate) async fn notify( - &self, - method: &str, - params: &P, - ) -> Result<(), serde_json::Error> { - let params = serde_json::to_value(params)?; - self.tx - .send(RpcServerOutboundMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(params), - })) - .await - .map_err(|_| { - serde_json::Error::io(std::io::Error::new( - std::io::ErrorKind::BrokenPipe, - "JSON-RPC transport closed", - )) - }) - } -} - -type RpcRequestRoute = dyn Fn( - Arc, - codex_app_server_protocol::JSONRPCRequest, - ) -> Pin + Send>> - + Send - + Sync; - -type RpcNotificationRoute = dyn Fn( - Arc, - codex_app_server_protocol::JSONRPCNotification, - ) -> Pin> + Send>> - + Send - + Sync; - -pub(crate) struct RpcRouter { - request_routes: HashMap>>, - notification_routes: HashMap>>, -} - -impl RpcRouter { - pub(crate) fn new() -> Self { - Self { - request_routes: HashMap::new(), - notification_routes: HashMap::new(), - } - } - - pub(crate) fn request(&mut self, method: &str, handler: F) - where - P: DeserializeOwned + Send + 'static, - R: Serialize + Send + 'static, - F: Fn(Arc, P) -> Fut + Send + Sync + 'static, - Fut: std::future::Future> + Send + 'static, - { - let method = method.to_string(); - let handler = std::sync::Arc::new(handler); - self.request_routes.insert( - method, - Box::new( - move |server_handler: Arc, request: codex_app_server_protocol::JSONRPCRequest| { - let handler = std::sync::Arc::clone(&handler); - let params = serde_json::from_value::

(request.params.unwrap_or(Value::Null)) - .map_err(|error| invalid_params(error.to_string())); - let request_id = request.id; - Box::pin(async move { - let result = match params { - Ok(params) => handler(server_handler.clone(), params) - .await - .and_then(|value| { - serde_json::to_value(value) - .map_err(|error| invalid_params(error.to_string())) - }), - Err(error) => Err(error), - }; - RpcServerOutboundMessage::response(request_id, result) - }) - }, - ), - ); - } - - pub(crate) fn notification(&mut self, method: &str, handler: F) - where - P: DeserializeOwned + Send + 'static, - F: Fn(Arc, P) -> Fut + Send + Sync + 'static, - Fut: std::future::Future> + Send + 'static, - { - let method = method.to_string(); - let handler = std::sync::Arc::new(handler); - self.notification_routes.insert( - method, - Box::new( - move | - server_handler: Arc, - notification: codex_app_server_protocol::JSONRPCNotification| { - let handler = std::sync::Arc::clone(&handler); - let params = serde_json::from_value::

(notification.params.unwrap_or(Value::Null)) - .map_err(|err| err.to_string()); - Box::pin(async move { - match params { - Ok(params) => handler(server_handler.clone(), params).await, - Err(error) => Err(error), - } - }) - }, - ), - ); - } - - pub(crate) fn request_route( - &self, - method: &str, - ) -> Option<&Box>> { - self.request_routes.get(method) - } - - pub(crate) fn notification_route( - &self, - method: &str, - ) -> Option<&Box>> { - self.notification_routes.get(method) - } -} - #[cfg(test)] mod tests { use std::time::Duration; diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index 10e6b19c66..24c8297241 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -1,16 +1,15 @@ mod filesystem; mod handler; -mod jsonrpc; mod registry; mod processor; mod transport; pub(crate) use handler::ExecServerHandler; pub use transport::DEFAULT_LISTEN_URL; -pub use transport::ExecServerListenUrlParseError; +pub use transport::ExecServerListenUrlParseError as ExecServerTransportParseError; pub async fn run_main() -> Result<(), Box> { - run_main_with_listen_url(DEFAULT_LISTEN_URL).await + run_main_with_transport(DEFAULT_LISTEN_URL).await } pub async fn run_main_with_listen_url( @@ -18,3 +17,9 @@ pub async fn run_main_with_listen_url( ) -> Result<(), Box> { transport::run_transport(listen_url).await } + +pub async fn run_main_with_transport( + listen_url: &str, +) -> Result<(), Box> { + transport::run_transport(listen_url).await +} diff --git a/codex-rs/exec-server/tests/common.rs b/codex-rs/exec-server/tests/common.rs new file mode 100644 index 0000000000..aa076a6bc8 --- /dev/null +++ b/codex-rs/exec-server/tests/common.rs @@ -0,0 +1,182 @@ +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::JSONRPCNotification; +use codex_app_server_protocol::JSONRPCRequest; +use codex_app_server_protocol::RequestId; +use codex_utils_cargo_bin::cargo_bin; +use futures::{SinkExt, StreamExt}; +use tokio::io::AsyncBufReadExt; +use tokio::io::BufReader; +use tokio::process::Command; +use std::process::Stdio; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tokio::time::timeout; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::connect_async; + +enum OutgoingMessage { + Json(JSONRPCMessage), + RawText(String), +} + +pub struct ExecServer { + child: tokio::process::Child, + next_request_id: AtomicI64, + incoming_rx: mpsc::Receiver, + outgoing_tx: mpsc::Sender, + reader_task: JoinHandle<()>, + writer_task: JoinHandle<()>, +} + +impl ExecServer { + pub async fn send_request( + &mut self, + method: &str, + params: serde_json::Value, + ) -> anyhow::Result { + let request_id = RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst)); + let request = JSONRPCRequest { + id: request_id.clone(), + method: method.to_string(), + params: Some(params), + trace: None, + }; + self.outgoing_tx + .send(OutgoingMessage::Json(JSONRPCMessage::Request(request))) + .await?; + Ok(request_id) + } + + pub async fn send_notification( + &mut self, + method: &str, + params: serde_json::Value, + ) -> anyhow::Result<()> { + let notification = JSONRPCNotification { + method: method.to_string(), + params: Some(params), + }; + self.outgoing_tx + .send(OutgoingMessage::Json(JSONRPCMessage::Notification( + notification, + ))) + .await?; + Ok(()) + } + + pub async fn send_raw_text(&mut self, text: &str) -> anyhow::Result<()> { + self.outgoing_tx + .send(OutgoingMessage::RawText(text.to_string())) + .await?; + Ok(()) + } + + pub async fn next_event(&mut self) -> anyhow::Result { + self.incoming_rx + .recv() + .await + .ok_or_else(|| anyhow::anyhow!("exec-server closed before next event")) + } + + pub async fn wait_for_event( + &mut self, + predicate: impl Fn(&JSONRPCMessage) -> bool, + ) -> anyhow::Result { + loop { + let event = self.next_event().await?; + if predicate(&event) { + return Ok(event); + } + } + } + + pub async fn shutdown(&mut self) -> anyhow::Result<()> { + self.reader_task.abort(); + self.writer_task.abort(); + self.child.start_kill()?; + Ok(()) + } +} + +pub mod exec_server { + use super::*; + + pub async fn exec_server() -> anyhow::Result { + let binary = cargo_bin("codex-exec-server")?; + let mut child = Command::new(binary); + child.args(["--listen", "ws://127.0.0.1:0"]); + child.stdin(Stdio::null()); + child.stdout(Stdio::null()); + child.stderr(Stdio::piped()); + let mut child = child.spawn()?; + + let stderr = child.stderr.take().expect("stderr should be piped"); + let mut stderr_lines = BufReader::new(stderr).lines(); + let websocket_url = read_websocket_url(&mut stderr_lines).await?; + + let (websocket, _) = connect_async(websocket_url).await?; + let (mut outgoing_ws, mut incoming_ws) = websocket.split(); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(128); + let (incoming_tx, incoming_rx) = mpsc::channel::(128); + + let reader_task = tokio::spawn(async move { + while let Some(message) = incoming_ws.next().await { + let Ok(message) = message else { + break; + }; + let outgoing = match message { + Message::Text(text) => serde_json::from_str::(&text), + Message::Binary(bytes) => serde_json::from_slice::(&bytes), + _ => continue, + }; + if let Ok(message) = outgoing && let Err(_err) = incoming_tx.send(message).await { + break; + } + } + }); + + let writer_task = tokio::spawn(async move { + while let Some(message) = outgoing_rx.recv().await { + let outgoing = match message { + OutgoingMessage::Json(message) => { + match serde_json::to_string(&message) { + Ok(json) => Message::Text(json.into()), + Err(_) => continue, + } + } + OutgoingMessage::RawText(message) => Message::Text(message.into()), + }; + if outgoing_ws.send(outgoing).await.is_err() { + break; + } + } + }); + + Ok(ExecServer { + child, + next_request_id: AtomicI64::new(1), + incoming_rx, + outgoing_tx, + reader_task, + writer_task, + }) + } + + async fn read_websocket_url(lines: &mut tokio::io::Lines>) -> anyhow::Result + where + R: tokio::io::AsyncRead + Unpin, + { + let line = timeout(std::time::Duration::from_secs(5), lines.next_line()) + .await?? + .ok_or_else(|| anyhow::anyhow!("missing websocket startup banner"))?; + + let websocket_url = line + .split_whitespace() + .find(|part| part.starts_with("ws://")) + .ok_or_else(|| anyhow::anyhow!("missing websocket URL in startup banner: {line}"))?; + Ok(websocket_url.to_string()) + } +} diff --git a/codex-rs/exec-server/tests/process.rs b/codex-rs/exec-server/tests/process.rs index 95583c61aa..aac2181bb5 100644 --- a/codex-rs/exec-server/tests/process.rs +++ b/codex-rs/exec-server/tests/process.rs @@ -24,7 +24,7 @@ async fn exec_server_stubs_process_start_over_websocket() -> anyhow::Result<()> .wait_for_event(|event| { matches!( event, - JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id + JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if *id == initialize_id ) }) .await?; @@ -53,7 +53,7 @@ async fn exec_server_stubs_process_start_over_websocket() -> anyhow::Result<()> .wait_for_event(|event| { matches!( event, - JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id + JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if *id == process_start_id ) }) .await?; diff --git a/codex-rs/exec-server/tests/websocket.rs b/codex-rs/exec-server/tests/websocket.rs index d653da6e2b..c0c0543a82 100644 --- a/codex-rs/exec-server/tests/websocket.rs +++ b/codex-rs/exec-server/tests/websocket.rs @@ -44,7 +44,7 @@ async fn exec_server_reports_malformed_websocket_json_and_keeps_running() -> any .wait_for_event(|event| { matches!( event, - JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id + JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if *id == initialize_id ) }) .await?;