From 0a846a2625a56d2a913e516c476abd5a1d71c0e1 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 18 Mar 2026 14:30:57 -0700 Subject: [PATCH] Add generic exec-server RPC foundation Co-authored-by: Codex --- codex-rs/Cargo.lock | 3 + codex-rs/exec-server/Cargo.toml | 4 + .../exec-server/src/bin/codex-exec-server.rs | 18 +- codex-rs/exec-server/src/client.rs | 431 +++++++++--------- .../exec-server/src/client/local_backend.rs | 44 ++ codex-rs/exec-server/src/client_api.rs | 17 + codex-rs/exec-server/src/connection.rs | 266 +++++++++++ codex-rs/exec-server/src/lib.rs | 13 +- codex-rs/exec-server/src/local.rs | 70 +++ codex-rs/exec-server/src/rpc.rs | 426 ++++++++++++++--- codex-rs/exec-server/src/server.rs | 132 +----- codex-rs/exec-server/src/server/handler.rs | 40 ++ codex-rs/exec-server/src/server/processor.rs | 107 +++++ codex-rs/exec-server/src/server/registry.rs | 26 ++ codex-rs/exec-server/src/server/transport.rs | 166 +++++++ codex-rs/exec-server/src/server_process.rs | 46 -- codex-rs/exec-server/tests/stdio_smoke.rs | 8 +- 17 files changed, 1368 insertions(+), 449 deletions(-) create mode 100644 codex-rs/exec-server/src/client/local_backend.rs create mode 100644 codex-rs/exec-server/src/client_api.rs create mode 100644 codex-rs/exec-server/src/connection.rs create mode 100644 codex-rs/exec-server/src/local.rs create mode 100644 codex-rs/exec-server/src/server/handler.rs create mode 100644 codex-rs/exec-server/src/server/processor.rs create mode 100644 codex-rs/exec-server/src/server/registry.rs create mode 100644 codex-rs/exec-server/src/server/transport.rs delete mode 100644 codex-rs/exec-server/src/server_process.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 8aa93c9687..11b3fa6e81 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2009,14 +2009,17 @@ version = "0.0.0" dependencies = [ "anyhow", "base64 0.22.1", + "clap", "codex-app-server-protocol", "codex-utils-cargo-bin", "codex-utils-pty", + "futures", "pretty_assertions", "serde", "serde_json", "thiserror 2.0.18", "tokio", + "tokio-tungstenite", "tracing", ] diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 7b1ec13cef..1b47760975 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/codex-rs/exec-server/Cargo.toml @@ -13,8 +13,10 @@ workspace = true [dependencies] base64 = { workspace = true } +clap = { workspace = true, features = ["derive"] } codex-app-server-protocol = { workspace = true } codex-utils-pty = { workspace = true } +futures = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } thiserror = { workspace = true } @@ -22,11 +24,13 @@ tokio = { workspace = true, features = [ "io-std", "io-util", "macros", + "net", "process", "rt-multi-thread", "sync", "time", ] } +tokio-tungstenite = { workspace = true } tracing = { workspace = true } [dev-dependencies] 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 399167c1a9..16df84d9b6 100644 --- a/codex-rs/exec-server/src/bin/codex-exec-server.rs +++ b/codex-rs/exec-server/src/bin/codex-exec-server.rs @@ -1,6 +1,22 @@ +use clap::Parser; +use codex_exec_server::ExecServerTransport; + +#[derive(Debug, Parser)] +struct ExecServerArgs { + /// Transport endpoint URL. Supported values: `stdio://` (default), + /// `ws://IP:PORT`. + #[arg( + long = "listen", + value_name = "URL", + default_value = ExecServerTransport::DEFAULT_LISTEN_URL + )] + listen: ExecServerTransport, +} + #[tokio::main] async fn main() { - if let Err(err) = codex_exec_server::run_main().await { + let args = ExecServerArgs::parse(); + 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/client.rs b/codex-rs/exec-server/src/client.rs index 7176d84343..9830771a00 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -1,54 +1,95 @@ -use std::collections::HashMap; use std::sync::Arc; -use std::sync::Mutex as StdMutex; -use std::sync::atomic::AtomicI64; -use std::sync::atomic::Ordering; +use std::time::Duration; -use codex_app_server_protocol::JSONRPCError; -use codex_app_server_protocol::JSONRPCErrorError; -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 serde_json::Value; -use tokio::io::AsyncBufReadExt; -use tokio::io::AsyncWriteExt; -use tokio::io::BufReader; -use tokio::process::Child; -use tokio::sync::Mutex; -use tokio::sync::mpsc; -use tokio::sync::oneshot; -use tokio::task::JoinHandle; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::time::timeout; +use tokio_tungstenite::connect_async; use tracing::warn; +use crate::client_api::ExecServerClientConnectOptions; +use crate::client_api::RemoteExecServerConnectArgs; +use crate::connection::JsonRpcConnection; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeParams; use crate::protocol::InitializeResponse; -use crate::server_process::ExecServerLaunchCommand; -use crate::server_process::spawn_stdio_exec_server; +use crate::rpc::RpcCallError; +use crate::rpc::RpcClient; +use crate::rpc::RpcClientEvent; + +mod local_backend; +use local_backend::LocalBackend; + +const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); + +impl Default for ExecServerClientConnectOptions { + fn default() -> Self { + Self { + client_name: "codex-core".to_string(), + initialize_timeout: INITIALIZE_TIMEOUT, + } + } +} + +impl From for ExecServerClientConnectOptions { + fn from(value: RemoteExecServerConnectArgs) -> Self { + Self { + client_name: value.client_name, + initialize_timeout: value.initialize_timeout, + } + } +} + +impl RemoteExecServerConnectArgs { + pub fn new(websocket_url: String, client_name: String) -> Self { + Self { + websocket_url, + client_name, + connect_timeout: CONNECT_TIMEOUT, + initialize_timeout: INITIALIZE_TIMEOUT, + } + } +} + +enum ClientBackend { + Remote(RpcClient), + InProcess(LocalBackend), +} + +impl ClientBackend { + fn as_local(&self) -> Option<&LocalBackend> { + match self { + ClientBackend::Remote(_) => None, + ClientBackend::InProcess(backend) => Some(backend), + } + } + + fn as_remote(&self) -> Option<&RpcClient> { + match self { + ClientBackend::Remote(client) => Some(client), + ClientBackend::InProcess(_) => None, + } + } +} struct Inner { - child: StdMutex>, - write_tx: mpsc::UnboundedSender, - pending: Mutex>>>, - next_request_id: AtomicI64, - reader_task: JoinHandle<()>, - writer_task: JoinHandle<()>, + backend: ClientBackend, + reader_task: tokio::task::JoinHandle<()>, } impl Drop for Inner { fn drop(&mut self) { - self.reader_task.abort(); - self.writer_task.abort(); - if let Ok(mut child_guard) = self.child.lock() - && let Some(child) = child_guard.as_mut() + if let Some(backend) = self.backend.as_local() + && let Ok(handle) = tokio::runtime::Handle::try_current() { - let _ = child.start_kill(); + let backend = backend.clone(); + handle.spawn(async move { + backend.shutdown().await; + }); } + self.reader_task.abort(); } } @@ -61,6 +102,16 @@ pub struct ExecServerClient { pub enum ExecServerError { #[error("failed to spawn exec-server: {0}")] Spawn(#[source] std::io::Error), + #[error("timed out connecting to exec-server websocket `{url}` after {timeout:?}")] + WebSocketConnectTimeout { url: String, timeout: Duration }, + #[error("failed to connect to exec-server websocket `{url}`: {source}")] + WebSocketConnect { + url: String, + #[source] + source: tokio_tungstenite::tungstenite::Error, + }, + #[error("timed out waiting for exec-server initialize handshake after {timeout:?}")] + InitializeTimedOut { timeout: Duration }, #[error("exec-server transport closed")] Closed, #[error("failed to serialize or deserialize exec-server JSON: {0}")] @@ -72,205 +123,145 @@ pub enum ExecServerError { } impl ExecServerClient { - pub async fn spawn(command: ExecServerLaunchCommand) -> Result { - let crate::server_process::SpawnedStdioExecServer { - child, - stdin, - stdout, - } = spawn_stdio_exec_server(command)?; - - let (write_tx, mut write_rx) = mpsc::unbounded_channel::(); - let writer_task = tokio::spawn(async move { - let mut stdin = stdin; - while let Some(message) = write_rx.recv().await { - let encoded = match serde_json::to_vec(&message) { - Ok(encoded) => encoded, - Err(err) => { - warn!("failed to encode exec-server message: {err}"); - break; - } - }; - if stdin.write_all(&encoded).await.is_err() { - break; - } - if stdin.write_all(b"\n").await.is_err() { - break; - } - if stdin.flush().await.is_err() { - break; - } - } + pub async fn connect_in_process( + options: ExecServerClientConnectOptions, + ) -> Result { + let backend = LocalBackend::new(crate::server::ExecServerHandler::new()); + let inner = Arc::new(Inner { + backend: ClientBackend::InProcess(backend), + reader_task: tokio::spawn(async {}), }); - - let pending = Mutex::new(HashMap::< - RequestId, - oneshot::Sender>, - >::new()); - let inner = Arc::new_cyclic(move |weak| { - let weak = weak.clone(); - let reader_task = tokio::spawn(async move { - let mut lines = BufReader::new(stdout).lines(); - loop { - let Some(inner) = weak.upgrade() else { - break; - }; - match lines.next_line().await { - Ok(Some(line)) => { - if line.trim().is_empty() { - continue; - } - - match serde_json::from_str::(&line) { - Ok(message) => { - if let Err(err) = handle_server_message(&inner, message).await { - warn!("failed to handle exec-server message: {err}"); - break; - } - } - Err(err) => { - warn!("failed to parse exec-server message: {err}"); - break; - } - } - } - Ok(None) => break, - Err(err) => { - warn!("failed to read exec-server stdout: {err}"); - break; - } - } - } - - if let Some(inner) = weak.upgrade() { - fail_pending_requests(&inner).await; - } - }); - - Inner { - child: StdMutex::new(Some(child)), - write_tx, - pending, - next_request_id: AtomicI64::new(1), - reader_task, - writer_task, - } - }); - let client = Self { inner }; - client - .initialize(InitializeParams { - client_name: "codex-core".to_string(), - }) - .await?; - client.send_notification(INITIALIZED_METHOD, serde_json::json!({}))?; + client.initialize(options).await?; 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 { + let websocket_url = args.websocket_url.clone(); + let connect_timeout = args.connect_timeout; + let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str())) + .await + .map_err(|_| ExecServerError::WebSocketConnectTimeout { + url: websocket_url.clone(), + timeout: connect_timeout, + })? + .map_err(|source| ExecServerError::WebSocketConnect { + url: websocket_url.clone(), + source, + })?; + + Self::connect( + JsonRpcConnection::from_websocket( + stream, + format!("exec-server websocket {websocket_url}"), + ), + args.into(), + ) + .await + } + pub async fn initialize( &self, - params: InitializeParams, + options: ExecServerClientConnectOptions, ) -> Result { - self.send_request(INITIALIZE_METHOD, params).await - } + let ExecServerClientConnectOptions { + client_name, + initialize_timeout, + } = options; - async fn send_request(&self, method: &str, params: P) -> Result - where - P: Serialize, - R: DeserializeOwned, - { - let id = RequestId::Integer(self.inner.next_request_id.fetch_add(1, Ordering::SeqCst)); - let params = serde_json::to_value(params)?; - let (tx, rx) = oneshot::channel(); - self.inner.pending.lock().await.insert(id.clone(), tx); - - if let Err(err) = self - .inner - .write_tx - .send(JSONRPCMessage::Request(JSONRPCRequest { - id: id.clone(), - method: method.to_string(), - params: Some(params), - trace: None, - })) - { - let _ = self.inner.pending.lock().await.remove(&id); - return Err(ExecServerError::Protocol(format!( - "failed to queue exec-server request: {err}" - ))); - } - - let result = rx.await.map_err(|_| ExecServerError::Closed)??; - Ok(serde_json::from_value(result)?) - } - - fn send_notification

(&self, method: &str, params: P) -> Result<(), ExecServerError> - where - P: Serialize, - { - let params = serde_json::to_value(params)?; - self.inner - .write_tx - .send(JSONRPCMessage::Notification(JSONRPCNotification { - method: method.to_string(), - params: Some(params), - })) - .map_err(|err| { - ExecServerError::Protocol(format!( - "failed to queue exec-server notification: {err}" - )) - }) - } -} - -impl From for ExecServerError { - fn from(error: JSONRPCErrorError) -> Self { - Self::Server { - code: error.code, - message: error.message, - } - } -} - -async fn handle_server_message( - inner: &Inner, - message: JSONRPCMessage, -) -> Result<(), ExecServerError> { - match message { - JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { - if let Some(tx) = inner.pending.lock().await.remove(&id) { - let _ = tx.send(Ok(result)); - } - Ok(()) - } - JSONRPCMessage::Error(JSONRPCError { id, error }) => { - if let Some(tx) = inner.pending.lock().await.remove(&id) { - let _ = tx.send(Err(error)); - Ok(()) + timeout(initialize_timeout, async { + let response = if let Some(backend) = self.inner.backend.as_local() { + backend.initialize().await? } else { - Err(ExecServerError::Server { - code: error.code, - message: error.message, - }) + let params = InitializeParams { client_name }; + let Some(remote) = self.inner.backend.as_remote() else { + return Err(ExecServerError::Protocol( + "remote backend missing during initialize".to_string(), + )); + }; + remote.call(INITIALIZE_METHOD, ¶ms).await? + }; + self.notify_initialized().await?; + Ok(response) + }) + .await + .map_err(|_| ExecServerError::InitializeTimedOut { + timeout: initialize_timeout, + })? + } + + async fn connect( + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + ) -> Result { + let (rpc_client, mut events_rx) = RpcClient::new(connection); + let reader_task = tokio::spawn(async move { + while let Some(event) = events_rx.recv().await { + match event { + RpcClientEvent::Notification(notification) => { + warn!( + "ignoring unexpected exec-server notification during stub phase: {}", + notification.method + ); + } + RpcClientEvent::Disconnected { reason } => { + if let Some(reason) = reason { + warn!("exec-server client transport disconnected: {reason}"); + } + return; + } + } } + }); + + let client = Self { + inner: Arc::new(Inner { + backend: ClientBackend::Remote(rpc_client), + reader_task, + }), + }; + client.initialize(options).await?; + Ok(client) + } + + async fn notify_initialized(&self) -> Result<(), ExecServerError> { + match &self.inner.backend { + ClientBackend::Remote(client) => client + .notify(INITIALIZED_METHOD, &serde_json::json!({})) + .await + .map_err(ExecServerError::Json), + ClientBackend::InProcess(backend) => backend.initialized().await, } - JSONRPCMessage::Notification(notification) => Err(ExecServerError::Protocol(format!( - "unexpected exec-server notification: {}", - notification.method - ))), - JSONRPCMessage::Request(request) => Err(ExecServerError::Protocol(format!( - "unexpected exec-server request: {}", - request.method - ))), } } -async fn fail_pending_requests(inner: &Inner) { - let mut pending = inner.pending.lock().await; - for (_, tx) in pending.drain() { - let _ = tx.send(Err(JSONRPCErrorError { - code: -32000, - message: "exec-server transport closed".to_string(), - data: None, - })); +impl From for ExecServerError { + fn from(value: RpcCallError) -> Self { + match value { + RpcCallError::Closed => Self::Closed, + RpcCallError::Json(err) => Self::Json(err), + RpcCallError::Server(error) => Self::Server { + code: error.code, + message: error.message, + }, + } } } diff --git a/codex-rs/exec-server/src/client/local_backend.rs b/codex-rs/exec-server/src/client/local_backend.rs new file mode 100644 index 0000000000..195847bf14 --- /dev/null +++ b/codex-rs/exec-server/src/client/local_backend.rs @@ -0,0 +1,44 @@ +use std::sync::Arc; + +use tokio::sync::Mutex; + +use crate::protocol::InitializeResponse; +use crate::server::ExecServerHandler; + +use super::ExecServerError; + +#[derive(Clone)] +pub(super) struct LocalBackend { + handler: Arc>, +} + +impl LocalBackend { + pub(super) fn new(handler: ExecServerHandler) -> Self { + Self { + handler: Arc::new(Mutex::new(handler)), + } + } + + pub(super) async fn shutdown(&self) { + self.handler.lock().await.shutdown().await; + } + + pub(super) async fn initialize(&self) -> Result { + self.handler + .lock() + .await + .initialize() + .map_err(|error| ExecServerError::Server { + code: error.code, + message: error.message, + }) + } + + pub(super) async fn initialized(&self) -> Result<(), ExecServerError> { + self.handler + .lock() + .await + .initialized() + .map_err(ExecServerError::Protocol) + } +} diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs new file mode 100644 index 0000000000..6e89763416 --- /dev/null +++ b/codex-rs/exec-server/src/client_api.rs @@ -0,0 +1,17 @@ +use std::time::Duration; + +/// Connection options for any exec-server client transport. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecServerClientConnectOptions { + pub client_name: String, + pub initialize_timeout: Duration, +} + +/// WebSocket connection arguments for a remote exec-server. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RemoteExecServerConnectArgs { + pub websocket_url: String, + pub client_name: String, + pub connect_timeout: Duration, + pub initialize_timeout: Duration, +} diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs new file mode 100644 index 0000000000..6c9f66d9a8 --- /dev/null +++ b/codex-rs/exec-server/src/connection.rs @@ -0,0 +1,266 @@ +use codex_app_server_protocol::JSONRPCMessage; +use futures::SinkExt; +use futures::StreamExt; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::io::BufWriter; +use tokio::sync::mpsc; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::tungstenite::Message; + +pub(crate) const CHANNEL_CAPACITY: usize = 128; + +#[derive(Debug)] +pub(crate) enum JsonRpcConnectionEvent { + Message(JSONRPCMessage), + Disconnected { reason: Option }, +} + +pub(crate) struct JsonRpcConnection { + outgoing_tx: mpsc::Sender, + incoming_rx: mpsc::Receiver, + task_handles: Vec>, +} + +impl JsonRpcConnection { + pub(crate) fn from_stdio(reader: R, writer: W, connection_label: String) -> Self + where + R: AsyncRead + Unpin + Send + 'static, + W: AsyncWrite + Unpin + Send + 'static, + { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); + + let reader_label = connection_label.clone(); + let incoming_tx_for_reader = incoming_tx.clone(); + let reader_task = tokio::spawn(async move { + let mut lines = BufReader::new(reader).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => { + if line.trim().is_empty() { + continue; + } + match serde_json::from_str::(&line) { + Ok(message) => { + if incoming_tx_for_reader + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + break; + } + } + Err(err) => { + send_disconnected( + &incoming_tx_for_reader, + Some(format!( + "failed to parse JSON-RPC message from {reader_label}: {err}" + )), + ) + .await; + break; + } + } + } + Ok(None) => { + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; + break; + } + Err(err) => { + send_disconnected( + &incoming_tx_for_reader, + Some(format!( + "failed to read JSON-RPC message from {reader_label}: {err}" + )), + ) + .await; + break; + } + } + } + }); + + let writer_task = tokio::spawn(async move { + let mut writer = BufWriter::new(writer); + while let Some(message) = outgoing_rx.recv().await { + if let Err(err) = write_jsonrpc_line_message(&mut writer, &message).await { + send_disconnected( + &incoming_tx, + Some(format!( + "failed to write JSON-RPC message to {connection_label}: {err}" + )), + ) + .await; + break; + } + } + }); + + Self { + outgoing_tx, + incoming_rx, + task_handles: vec![reader_task, writer_task], + } + } + + pub(crate) fn from_websocket(stream: WebSocketStream, connection_label: String) -> Self + where + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, + { + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (mut websocket_writer, mut websocket_reader) = stream.split(); + + let reader_label = connection_label.clone(); + let incoming_tx_for_reader = incoming_tx.clone(); + let reader_task = tokio::spawn(async move { + loop { + match websocket_reader.next().await { + Some(Ok(Message::Text(text))) => { + match serde_json::from_str::(text.as_ref()) { + Ok(message) => { + if incoming_tx_for_reader + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + break; + } + } + Err(err) => { + send_disconnected( + &incoming_tx_for_reader, + Some(format!( + "failed to parse websocket JSON-RPC message from {reader_label}: {err}" + )), + ) + .await; + break; + } + } + } + Some(Ok(Message::Binary(bytes))) => { + match serde_json::from_slice::(bytes.as_ref()) { + Ok(message) => { + if incoming_tx_for_reader + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + break; + } + } + Err(err) => { + send_disconnected( + &incoming_tx_for_reader, + Some(format!( + "failed to parse websocket JSON-RPC message from {reader_label}: {err}" + )), + ) + .await; + break; + } + } + } + Some(Ok(Message::Close(_))) => { + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; + break; + } + Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {} + Some(Ok(_)) => {} + Some(Err(err)) => { + send_disconnected( + &incoming_tx_for_reader, + Some(format!( + "failed to read websocket JSON-RPC message from {reader_label}: {err}" + )), + ) + .await; + break; + } + None => { + send_disconnected(&incoming_tx_for_reader, /*reason*/ None).await; + break; + } + } + } + }); + + let writer_task = tokio::spawn(async move { + while let Some(message) = outgoing_rx.recv().await { + match serialize_jsonrpc_message(&message) { + Ok(encoded) => { + if let Err(err) = websocket_writer.send(Message::Text(encoded.into())).await + { + send_disconnected( + &incoming_tx, + Some(format!( + "failed to write websocket JSON-RPC message to {connection_label}: {err}" + )), + ) + .await; + break; + } + } + Err(err) => { + send_disconnected( + &incoming_tx, + Some(format!( + "failed to serialize JSON-RPC message for {connection_label}: {err}" + )), + ) + .await; + break; + } + } + } + }); + + Self { + outgoing_tx, + incoming_rx, + task_handles: vec![reader_task, writer_task], + } + } + + pub(crate) fn into_parts( + self, + ) -> ( + mpsc::Sender, + mpsc::Receiver, + Vec>, + ) { + (self.outgoing_tx, self.incoming_rx, self.task_handles) + } +} + +async fn send_disconnected( + incoming_tx: &mpsc::Sender, + reason: Option, +) { + let _ = incoming_tx + .send(JsonRpcConnectionEvent::Disconnected { reason }) + .await; +} + +async fn write_jsonrpc_line_message( + writer: &mut BufWriter, + message: &JSONRPCMessage, +) -> std::io::Result<()> +where + W: AsyncWrite + Unpin, +{ + let encoded = + serialize_jsonrpc_message(message).map_err(|err| std::io::Error::other(err.to_string()))?; + writer.write_all(encoded.as_bytes()).await?; + writer.write_all(b"\n").await?; + writer.flush().await +} + +fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result { + serde_json::to_string(message) +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index e1f9c05c71..e204d6e084 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -1,12 +1,21 @@ mod client; +mod client_api; +mod connection; +mod local; mod protocol; mod rpc; mod server; -mod server_process; pub use client::ExecServerClient; pub use client::ExecServerError; +pub use client_api::ExecServerClientConnectOptions; +pub use client_api::RemoteExecServerConnectArgs; +pub use local::ExecServerLaunchCommand; +pub use local::SpawnedExecServer; +pub use local::spawn_local_exec_server; pub use protocol::InitializeParams; pub use protocol::InitializeResponse; +pub use server::ExecServerTransport; +pub use server::ExecServerTransportParseError; pub use server::run_main; -pub use server_process::ExecServerLaunchCommand; +pub use server::run_main_with_transport; diff --git a/codex-rs/exec-server/src/local.rs b/codex-rs/exec-server/src/local.rs new file mode 100644 index 0000000000..20b59a983c --- /dev/null +++ b/codex-rs/exec-server/src/local.rs @@ -0,0 +1,70 @@ +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::Mutex as StdMutex; + +use tokio::process::Child; +use tokio::process::Command; + +use crate::client::ExecServerClient; +use crate::client::ExecServerError; +use crate::client_api::ExecServerClientConnectOptions; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ExecServerLaunchCommand { + pub program: PathBuf, + pub args: Vec, +} + +pub struct SpawnedExecServer { + client: ExecServerClient, + child: StdMutex>, +} + +impl SpawnedExecServer { + pub fn client(&self) -> &ExecServerClient { + &self.client + } +} + +impl Drop for SpawnedExecServer { + fn drop(&mut self) { + if let Ok(mut child_guard) = self.child.lock() + && let Some(child) = child_guard.as_mut() + { + let _ = child.start_kill(); + } + } +} + +pub async fn spawn_local_exec_server( + command: ExecServerLaunchCommand, + options: ExecServerClientConnectOptions, +) -> Result { + let mut child = Command::new(&command.program); + child.args(&command.args); + child.stdin(Stdio::piped()); + child.stdout(Stdio::piped()); + child.stderr(Stdio::inherit()); + child.kill_on_drop(true); + + let mut child = child.spawn().map_err(ExecServerError::Spawn)?; + let stdin = child.stdin.take().ok_or_else(|| { + ExecServerError::Protocol("exec-server stdin was not captured".to_string()) + })?; + let stdout = child.stdout.take().ok_or_else(|| { + ExecServerError::Protocol("exec-server stdout was not captured".to_string()) + })?; + + let client = match ExecServerClient::connect_stdio(stdin, stdout, options).await { + Ok(client) => client, + Err(err) => { + let _ = child.start_kill(); + return Err(err); + } + }; + + Ok(SpawnedExecServer { + client, + child: StdMutex::new(Some(child)), + }) +} diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 522120e307..9ebd832ad3 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -1,4 +1,9 @@ use std::collections::HashMap; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; use codex_app_server_protocol::JSONRPCError; use codex_app_server_protocol::JSONRPCErrorError; @@ -7,12 +12,29 @@ 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 serde_json::Value; +use tokio::sync::Mutex; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; +use tracing::warn; -#[derive(Debug, Clone, PartialEq)] -pub(crate) enum RpcServerInboundMessage { - Request(JSONRPCRequest), +use crate::connection::JsonRpcConnection; +use crate::connection::JsonRpcConnectionEvent; + +type PendingRequest = oneshot::Sender>; +type BoxFuture = Pin + Send + 'static>>; +type RequestRoute = + Box, JSONRPCRequest) -> BoxFuture + Send + Sync>; +type NotificationRoute = + Box, JSONRPCNotification) -> BoxFuture> + Send + Sync>; + +#[derive(Debug)] +pub(crate) enum RpcClientEvent { Notification(JSONRPCNotification), + Disconnected { reason: Option }, } #[derive(Debug, Clone, PartialEq)] @@ -25,17 +47,46 @@ pub(crate) enum RpcServerOutboundMessage { request_id: RequestId, error: JSONRPCErrorError, }, + #[allow(dead_code)] + Notification(JSONRPCNotification), } -type RequestRoute = Box I + Send + Sync>; -type NotificationRoute = Box Result + Send + Sync>; - -pub(crate) struct RpcRouter { - request_routes: HashMap<&'static str, RequestRoute>, - notification_routes: HashMap<&'static str, NotificationRoute>, +#[allow(dead_code)] +#[derive(Clone)] +pub(crate) struct RpcNotificationSender { + outgoing_tx: mpsc::Sender, } -impl Default for RpcRouter { +impl RpcNotificationSender { + pub(crate) fn new(outgoing_tx: mpsc::Sender) -> Self { + Self { outgoing_tx } + } + + #[allow(dead_code)] + pub(crate) async fn notify( + &self, + method: &str, + params: &P, + ) -> Result<(), JSONRPCErrorError> { + let params = serde_json::to_value(params).map_err(|err| internal_error(err.to_string()))?; + self.outgoing_tx + .send(RpcServerOutboundMessage::Notification( + JSONRPCNotification { + method: method.to_string(), + params: Some(params), + }, + )) + .await + .map_err(|_| internal_error("RPC connection closed while sending notification".into())) + } +} + +pub(crate) struct RpcRouter { + request_routes: HashMap<&'static str, RequestRoute>, + notification_routes: HashMap<&'static str, NotificationRoute>, +} + +impl Default for RpcRouter { fn default() -> Self { Self { request_routes: HashMap::new(), @@ -44,68 +95,216 @@ impl Default for RpcRouter { } } -impl RpcRouter { +impl RpcRouter +where + S: Send + Sync + 'static, +{ pub(crate) fn new() -> Self { Self::default() } - pub(crate) fn raw_request(&mut self, method: &'static str, route: F) + pub(crate) fn request(&mut self, method: &'static str, handler: F) where - F: Fn(JSONRPCRequest) -> I + Send + Sync + 'static, + P: DeserializeOwned + Send + 'static, + R: Serialize + Send + 'static, + F: Fn(Arc, P) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { - self.request_routes.insert(method, Box::new(route)); + self.request_routes.insert( + method, + Box::new(move |state, request| { + let request_id = request.id; + let params = request.params; + let response = + decode_request_params::

(params).map(|params| handler(state, params)); + Box::pin(async move { + let response = match response { + Ok(response) => response.await, + Err(error) => { + return RpcServerOutboundMessage::Error { request_id, error }; + } + }; + match response { + Ok(result) => match serde_json::to_value(result) { + Ok(result) => RpcServerOutboundMessage::Response { request_id, result }, + Err(err) => RpcServerOutboundMessage::Error { + request_id, + error: internal_error(err.to_string()), + }, + }, + Err(error) => RpcServerOutboundMessage::Error { request_id, error }, + } + }) + }), + ); } - pub(crate) fn notification(&mut self, method: &'static str, route: F) + pub(crate) fn notification(&mut self, method: &'static str, handler: F) where - F: Fn(JSONRPCNotification) -> Result + Send + Sync + 'static, + P: DeserializeOwned + Send + 'static, + F: Fn(Arc, P) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, { - self.notification_routes.insert(method, Box::new(route)); + self.notification_routes.insert( + method, + Box::new(move |state, notification| { + let params = decode_notification_params::

(notification.params) + .map(|params| handler(state, params)); + Box::pin(async move { + let handler = match params { + Ok(handler) => handler, + Err(err) => return Err(err), + }; + handler.await + }) + }), + ); } - pub(crate) fn route_message( - &self, - message: JSONRPCMessage, - unknown_request: impl FnOnce(JSONRPCRequest) -> I, - ) -> Result { - match route_server_message(message)? { - RpcServerInboundMessage::Request(request) => { - if let Some(route) = self.request_routes.get(request.method.as_str()) { - Ok(route(request)) - } else { - Ok(unknown_request(request)) - } - } - RpcServerInboundMessage::Notification(notification) => { - let Some(route) = self.notification_routes.get(notification.method.as_str()) else { - return Err(format!( - "unexpected notification method: {}", - notification.method - )); - }; - route(notification) - } - } + pub(crate) fn request_route(&self, method: &str) -> Option<&RequestRoute> { + self.request_routes.get(method) + } + + pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { + self.notification_routes.get(method) } } -pub(crate) fn route_server_message( - message: JSONRPCMessage, -) -> Result { - match message { - JSONRPCMessage::Request(request) => Ok(RpcServerInboundMessage::Request(request)), - JSONRPCMessage::Notification(notification) => { - Ok(RpcServerInboundMessage::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) struct RpcClient { + write_tx: mpsc::Sender, + pending: Arc>>, + next_request_id: AtomicI64, + transport_tasks: Vec>, + reader_task: JoinHandle<()>, +} + +impl RpcClient { + pub(crate) fn new(connection: JsonRpcConnection) -> (Self, mpsc::Receiver) { + let (write_tx, mut incoming_rx, transport_tasks) = connection.into_parts(); + let pending = Arc::new(Mutex::new(HashMap::::new())); + let (event_tx, event_rx) = mpsc::channel(128); + + let pending_for_reader = Arc::clone(&pending); + let reader_task = tokio::spawn(async move { + while let Some(event) = incoming_rx.recv().await { + match event { + JsonRpcConnectionEvent::Message(message) => { + if let Err(err) = + handle_server_message(&pending_for_reader, &event_tx, message).await + { + warn!("JSON-RPC client closing after protocol error: {err}"); + break; + } + } + JsonRpcConnectionEvent::Disconnected { reason } => { + let _ = event_tx.send(RpcClientEvent::Disconnected { reason }).await; + drain_pending(&pending_for_reader).await; + return; + } + } + } + + let _ = event_tx + .send(RpcClientEvent::Disconnected { reason: None }) + .await; + drain_pending(&pending_for_reader).await; + }); + + ( + Self { + write_tx, + pending, + next_request_id: AtomicI64::new(1), + transport_tasks, + reader_task, + }, + event_rx, + ) } + + pub(crate) async fn notify( + &self, + method: &str, + params: &P, + ) -> Result<(), serde_json::Error> { + let params = serde_json::to_value(params)?; + self.write_tx + .send(JSONRPCMessage::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", + )) + }) + } + + pub(crate) async fn call(&self, method: &str, params: &P) -> Result + where + P: Serialize, + T: DeserializeOwned, + { + let request_id = RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst)); + let (response_tx, response_rx) = oneshot::channel(); + self.pending + .lock() + .await + .insert(request_id.clone(), response_tx); + + let params = match serde_json::to_value(params) { + Ok(params) => params, + Err(err) => { + self.pending.lock().await.remove(&request_id); + return Err(RpcCallError::Json(err)); + } + }; + if self + .write_tx + .send(JSONRPCMessage::Request(JSONRPCRequest { + id: request_id.clone(), + method: method.to_string(), + params: Some(params), + trace: None, + })) + .await + .is_err() + { + self.pending.lock().await.remove(&request_id); + return Err(RpcCallError::Closed); + } + + let result = response_rx.await.map_err(|_| RpcCallError::Closed)?; + let response = match result { + Ok(response) => response, + Err(error) => return Err(RpcCallError::Server(error)), + }; + serde_json::from_value(response).map_err(RpcCallError::Json) + } + + #[cfg(test)] + #[allow(dead_code)] + pub(crate) async fn pending_request_count(&self) -> usize { + self.pending.lock().await.len() + } +} + +impl Drop for RpcClient { + fn drop(&mut self) { + for task in &self.transport_tasks { + task.abort(); + } + self.reader_task.abort(); + } +} + +#[derive(Debug)] +pub(crate) enum RpcCallError { + Closed, + Json(serde_json::Error), + Server(JSONRPCErrorError), } pub(crate) fn encode_server_message( @@ -124,5 +323,120 @@ pub(crate) fn encode_server_message( error, })) } + RpcServerOutboundMessage::Notification(notification) => { + Ok(JSONRPCMessage::Notification(notification)) + } + } +} + +pub(crate) fn invalid_request(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32600, + data: None, + message, + } +} + +pub(crate) fn method_not_found(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32601, + 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 decode_request_params

(params: Option) -> Result +where + P: DeserializeOwned, +{ + decode_params(params).map_err(|err| invalid_params(err.to_string())) +} + +fn decode_notification_params

(params: Option) -> Result +where + P: DeserializeOwned, +{ + decode_params(params).map_err(|err| err.to_string()) +} + +fn decode_params

(params: Option) -> Result +where + P: DeserializeOwned, +{ + let params = params.unwrap_or(Value::Null); + match serde_json::from_value(params.clone()) { + Ok(params) => Ok(params), + Err(err) => { + if matches!(params, Value::Object(ref map) if map.is_empty()) { + serde_json::from_value(Value::Null).map_err(|_| err) + } else { + Err(err) + } + } + } +} + +async fn handle_server_message( + pending: &Mutex>, + event_tx: &mpsc::Sender, + message: JSONRPCMessage, +) -> Result<(), String> { + match message { + JSONRPCMessage::Response(JSONRPCResponse { id, result }) => { + if let Some(pending) = pending.lock().await.remove(&id) { + let _ = pending.send(Ok(result)); + } + } + JSONRPCMessage::Error(JSONRPCError { id, error }) => { + if let Some(pending) = pending.lock().await.remove(&id) { + let _ = pending.send(Err(error)); + } + } + JSONRPCMessage::Notification(notification) => { + let _ = event_tx + .send(RpcClientEvent::Notification(notification)) + .await; + } + JSONRPCMessage::Request(request) => { + return Err(format!( + "unexpected JSON-RPC request from remote server: {}", + request.method + )); + } + } + + Ok(()) +} + +async fn drain_pending(pending: &Mutex>) { + let pending = { + let mut pending = pending.lock().await; + pending + .drain() + .map(|(_, pending)| pending) + .collect::>() + }; + for pending in pending { + let _ = pending.send(Err(JSONRPCErrorError { + code: -32000, + data: None, + message: "JSON-RPC transport closed".to_string(), + })); } } diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index cc81b9f70c..ba074e617f 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -1,126 +1,18 @@ -use codex_app_server_protocol::JSONRPCErrorError; -use codex_app_server_protocol::JSONRPCMessage; -use codex_app_server_protocol::JSONRPCRequest; -use codex_app_server_protocol::RequestId; -use tokio::io::AsyncBufReadExt; -use tokio::io::AsyncWriteExt; -use tokio::io::BufReader; +mod handler; +mod processor; +mod registry; +mod transport; -use crate::protocol::INITIALIZE_METHOD; -use crate::protocol::INITIALIZED_METHOD; -use crate::protocol::InitializeResponse; -use crate::protocol::PROTOCOL_VERSION; -use crate::rpc::RpcRouter; -use crate::rpc::RpcServerOutboundMessage; -use crate::rpc::encode_server_message; +pub(crate) use handler::ExecServerHandler; +pub use transport::ExecServerTransport; +pub use transport::ExecServerTransportParseError; pub async fn run_main() -> Result<(), Box> { - let mut stdin = BufReader::new(tokio::io::stdin()).lines(); - let mut stdout = tokio::io::stdout(); - - while let Some(line) = stdin.next_line().await? { - if line.trim().is_empty() { - continue; - } - - let message = serde_json::from_str::(&line)?; - let mut router = RpcRouter::new(); - router.raw_request(INITIALIZE_METHOD, handle_request); - router.notification(INITIALIZED_METHOD, |_| Ok(None)); - - match router.route_message(message, unknown_request) { - Ok(Some(outbound)) => { - send_message(&mut stdout, outbound).await?; - } - Ok(None) => {} - Err(message) => { - send_error( - &mut stdout, - RequestId::Integer(-1), - invalid_request(message), - ) - .await?; - } - } - } - - Ok(()) + run_main_with_transport(ExecServerTransport::Stdio).await } -fn handle_request(request: JSONRPCRequest) -> Option { - let result = match serde_json::to_value(InitializeResponse { - protocol_version: PROTOCOL_VERSION.to_string(), - }) { - Ok(result) => result, - Err(err) => { - return Some(RpcServerOutboundMessage::Error { - request_id: request.id, - error: internal_error(err.to_string()), - }); - } - }; - Some(RpcServerOutboundMessage::Response { - request_id: request.id, - result, - }) -} - -fn unknown_request(request: JSONRPCRequest) -> Option { - Some(RpcServerOutboundMessage::Error { - request_id: request.id, - error: method_not_implemented(format!( - "exec-server stub does not implement `{}` yet", - request.method - )), - }) -} - -async fn send_error( - stdout: &mut tokio::io::Stdout, - id: RequestId, - error: JSONRPCErrorError, -) -> Result<(), std::io::Error> { - send_message( - stdout, - RpcServerOutboundMessage::Error { - request_id: id, - error, - }, - ) - .await -} - -async fn send_message( - stdout: &mut tokio::io::Stdout, - message: RpcServerOutboundMessage, -) -> Result<(), std::io::Error> { - let message = encode_server_message(message).map_err(std::io::Error::other)?; - let encoded = serde_json::to_vec(&message).map_err(std::io::Error::other)?; - stdout.write_all(&encoded).await?; - stdout.write_all(b"\n").await?; - stdout.flush().await -} - -fn invalid_request(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32600, - message, - data: None, - } -} - -fn method_not_implemented(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32601, - message, - data: None, - } -} - -fn internal_error(message: String) -> JSONRPCErrorError { - JSONRPCErrorError { - code: -32603, - message, - data: None, - } +pub async fn run_main_with_transport( + transport: ExecServerTransport, +) -> Result<(), Box> { + transport::run_transport(transport).await } 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..4d033030b8 --- /dev/null +++ b/codex-rs/exec-server/src/server/handler.rs @@ -0,0 +1,40 @@ +use codex_app_server_protocol::JSONRPCErrorError; + +use crate::protocol::InitializeResponse; +use crate::protocol::PROTOCOL_VERSION; + +pub(crate) struct ExecServerHandler { + initialize_requested: bool, + initialized: bool, +} + +impl ExecServerHandler { + pub(crate) fn new() -> Self { + Self { + initialize_requested: false, + initialized: false, + } + } + + pub(crate) async fn shutdown(&self) {} + + pub(crate) fn initialize(&mut self) -> Result { + if self.initialize_requested { + return Err(crate::rpc::invalid_request( + "initialize may only be sent once per connection".to_string(), + )); + } + self.initialize_requested = true; + Ok(InitializeResponse { + protocol_version: PROTOCOL_VERSION.to_string(), + }) + } + + pub(crate) fn initialized(&mut self) -> Result<(), String> { + if !self.initialize_requested { + return Err("received `initialized` notification before `initialize`".into()); + } + self.initialized = true; + Ok(()) + } +} diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs new file mode 100644 index 0000000000..aa4da92d0f --- /dev/null +++ b/codex-rs/exec-server/src/server/processor.rs @@ -0,0 +1,107 @@ +use std::sync::Arc; + +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::rpc::RpcNotificationSender; +use crate::rpc::RpcServerOutboundMessage; +use crate::rpc::encode_server_message; +use crate::rpc::method_not_found; +use crate::server::ExecServerHandler; +use crate::server::registry::build_router; + +pub(crate) async fn run_connection(connection: JsonRpcConnection) { + let router = Arc::new(build_router()); + let (json_outgoing_tx, mut incoming_rx, _connection_tasks) = connection.into_parts(); + let (outgoing_tx, mut outgoing_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let handler = Arc::new(Mutex::new(ExecServerHandler::new())); + let _notifications = RpcNotificationSender::new(outgoing_tx.clone()); + + let outbound_task = tokio::spawn(async move { + while let Some(message) = outgoing_rx.recv().await { + let json_message = match encode_server_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) => match message { + codex_app_server_protocol::JSONRPCMessage::Request(request) => { + if let Some(route) = router.request_route(request.method.as_str()) { + let route = route(handler.clone(), request); + let outgoing_tx = outgoing_tx.clone(); + tokio::spawn(async move { + let message = route.await; + let _ = outgoing_tx.send(message).await; + }); + } else if outgoing_tx + .send(RpcServerOutboundMessage::Error { + request_id: request.id, + error: method_not_found(format!( + "exec-server stub does not implement `{}` yet", + request.method + )), + }) + .await + .is_err() + { + break; + } + } + codex_app_server_protocol::JSONRPCMessage::Notification(notification) => { + let Some(route) = router.notification_route(notification.method.as_str()) + else { + warn!( + "closing exec-server connection after unexpected notification: {}", + notification.method + ); + break; + }; + if let Err(err) = route(handler.clone(), notification).await { + warn!("closing exec-server connection after protocol error: {err}"); + break; + } + } + codex_app_server_protocol::JSONRPCMessage::Response(response) => { + warn!( + "closing exec-server connection after unexpected client response: {:?}", + response.id + ); + break; + } + codex_app_server_protocol::JSONRPCMessage::Error(error) => { + warn!( + "closing exec-server connection after unexpected client error: {:?}", + error.id + ); + break; + } + }, + JsonRpcConnectionEvent::Disconnected { reason } => { + if let Some(reason) = reason { + debug!("exec-server connection disconnected: {reason}"); + } + break; + } + } + } + + handler.lock().await.shutdown().await; + drop(outgoing_tx); + let _ = outbound_task.await; +} diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs new file mode 100644 index 0000000000..34aad179c6 --- /dev/null +++ b/codex-rs/exec-server/src/server/registry.rs @@ -0,0 +1,26 @@ +use std::sync::Arc; + +use tokio::sync::Mutex; + +use crate::protocol::INITIALIZE_METHOD; +use crate::protocol::INITIALIZED_METHOD; +use crate::protocol::InitializeParams; +use crate::rpc::RpcRouter; +use crate::server::ExecServerHandler; + +pub(crate) fn build_router() -> RpcRouter> { + let mut router = RpcRouter::new(); + router.request( + INITIALIZE_METHOD, + |handler: Arc>, _params: InitializeParams| async move { + handler.lock().await.initialize() + }, + ); + router.notification( + INITIALIZED_METHOD, + |handler: Arc>, (): ()| async move { + handler.lock().await.initialized() + }, + ); + router +} diff --git a/codex-rs/exec-server/src/server/transport.rs b/codex-rs/exec-server/src/server/transport.rs new file mode 100644 index 0000000000..b653c0b79b --- /dev/null +++ b/codex-rs/exec-server/src/server/transport.rs @@ -0,0 +1,166 @@ +use std::net::SocketAddr; +use std::str::FromStr; + +use tokio::net::TcpListener; +use tokio_tungstenite::accept_async; +use tracing::warn; + +use crate::connection::JsonRpcConnection; +use crate::server::processor::run_connection; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ExecServerTransport { + Stdio, + WebSocket { bind_address: SocketAddr }, +} + +#[derive(Debug, Clone, Eq, PartialEq)] +pub enum ExecServerTransportParseError { + UnsupportedListenUrl(String), + InvalidWebSocketListenUrl(String), +} + +impl std::fmt::Display for ExecServerTransportParseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ExecServerTransportParseError::UnsupportedListenUrl(listen_url) => write!( + f, + "unsupported --listen URL `{listen_url}`; expected `stdio://` or `ws://IP:PORT`" + ), + ExecServerTransportParseError::InvalidWebSocketListenUrl(listen_url) => write!( + f, + "invalid websocket --listen URL `{listen_url}`; expected `ws://IP:PORT`" + ), + } + } +} + +impl std::error::Error for ExecServerTransportParseError {} + +impl ExecServerTransport { + pub const DEFAULT_LISTEN_URL: &str = "stdio://"; + + 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()) + })?; + return Ok(Self::WebSocket { bind_address }); + } + + Err(ExecServerTransportParseError::UnsupportedListenUrl( + listen_url.to_string(), + )) + } +} + +impl FromStr for ExecServerTransport { + type Err = ExecServerTransportParseError; + + fn from_str(s: &str) -> Result { + Self::from_listen_url(s) + } +} + +pub(crate) async fn run_transport( + transport: ExecServerTransport, +) -> 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 + } + } +} + +async fn run_websocket_listener( + bind_address: SocketAddr, +) -> Result<(), Box> { + let listener = TcpListener::bind(bind_address).await?; + let local_addr = listener.local_addr()?; + print_websocket_startup_banner(local_addr); + + loop { + let (stream, peer_addr) = listener.accept().await?; + tokio::spawn(async move { + match accept_async(stream).await { + Ok(websocket) => { + run_connection(JsonRpcConnection::from_websocket( + websocket, + format!("exec-server websocket {peer_addr}"), + )) + .await; + } + Err(err) => { + warn!( + "failed to accept exec-server websocket connection from {peer_addr}: {err}" + ); + } + } + }); + } +} + +#[allow(clippy::print_stderr)] +fn print_websocket_startup_banner(addr: SocketAddr) { + eprintln!("codex-exec-server listening on ws://{addr}"); +} + +#[cfg(test)] +mod tests { + use pretty_assertions::assert_eq; + + 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") + .expect("websocket listen URL should parse"); + assert_eq!( + transport, + ExecServerTransport::WebSocket { + bind_address: "127.0.0.1:1234".parse().expect("valid socket address"), + } + ); + } + + #[test] + fn exec_server_transport_rejects_invalid_websocket_listen_url() { + let err = ExecServerTransport::from_listen_url("ws://localhost:1234") + .expect_err("hostname bind address should be rejected"); + assert_eq!( + err.to_string(), + "invalid websocket --listen URL `ws://localhost:1234`; expected `ws://IP:PORT`" + ); + } + + #[test] + fn exec_server_transport_rejects_unsupported_listen_url() { + let err = ExecServerTransport::from_listen_url("http://127.0.0.1:1234") + .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`" + ); + } +} diff --git a/codex-rs/exec-server/src/server_process.rs b/codex-rs/exec-server/src/server_process.rs deleted file mode 100644 index 787f7d8c02..0000000000 --- a/codex-rs/exec-server/src/server_process.rs +++ /dev/null @@ -1,46 +0,0 @@ -use std::path::PathBuf; -use std::process::Stdio; - -use tokio::process::Child; -use tokio::process::ChildStdin; -use tokio::process::ChildStdout; -use tokio::process::Command; - -use crate::client::ExecServerError; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ExecServerLaunchCommand { - pub program: PathBuf, - pub args: Vec, -} - -pub(crate) struct SpawnedStdioExecServer { - pub(crate) child: Child, - pub(crate) stdin: ChildStdin, - pub(crate) stdout: ChildStdout, -} - -pub(crate) fn spawn_stdio_exec_server( - command: ExecServerLaunchCommand, -) -> Result { - let mut child = Command::new(&command.program); - child.args(&command.args); - child.stdin(Stdio::piped()); - child.stdout(Stdio::piped()); - child.stderr(Stdio::inherit()); - child.kill_on_drop(true); - - let mut child = child.spawn().map_err(ExecServerError::Spawn)?; - let stdin = child.stdin.take().ok_or_else(|| { - ExecServerError::Protocol("exec-server stdin was not captured".to_string()) - })?; - let stdout = child.stdout.take().ok_or_else(|| { - ExecServerError::Protocol("exec-server stdout was not captured".to_string()) - })?; - - Ok(SpawnedStdioExecServer { - child, - stdin, - stdout, - }) -} diff --git a/codex-rs/exec-server/tests/stdio_smoke.rs b/codex-rs/exec-server/tests/stdio_smoke.rs index f2095129a6..6a89318133 100644 --- a/codex-rs/exec-server/tests/stdio_smoke.rs +++ b/codex-rs/exec-server/tests/stdio_smoke.rs @@ -66,7 +66,7 @@ async fn exec_server_accepts_initialize_over_stdio() -> anyhow::Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn exec_server_stubs_command_exec_over_stdio() -> anyhow::Result<()> { +async fn exec_server_stubs_process_start_over_stdio() -> anyhow::Result<()> { let binary = cargo_bin("codex-exec-server")?; let mut child = Command::new(binary); child.stdin(Stdio::piped()); @@ -93,7 +93,7 @@ async fn exec_server_stubs_command_exec_over_stdio() -> anyhow::Result<()> { let exec = JSONRPCMessage::Request(JSONRPCRequest { id: RequestId::Integer(2), - method: "command/exec".to_string(), + method: "process/start".to_string(), params: Some(serde_json::json!({ "processId": "proc-1", "argv": ["true"], @@ -113,13 +113,13 @@ async fn exec_server_stubs_command_exec_over_stdio() -> anyhow::Result<()> { let response: JSONRPCMessage = serde_json::from_str(&response_line)?; let JSONRPCMessage::Error(codex_app_server_protocol::JSONRPCError { id, error }) = response else { - panic!("expected command/exec stub error"); + panic!("expected process/start stub error"); }; assert_eq!(id, RequestId::Integer(2)); assert_eq!(error.code, -32601); assert_eq!( error.message, - "exec-server stub does not implement `command/exec` yet" + "exec-server stub does not implement `process/start` yet" ); child.start_kill()?;