Simplify exec-server client backends

Extract local and JSON-RPC backend helpers, collapse the remaining local-or-remote dispatch into one generic request helper, and keep the process test helper in its own module.

Co-authored-by: Codex <noreply@openai.com>

codex-rs/exec-server/src/client/jsonrpc_backend.rs

@@ -0,0 +1,51 @@
+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 serde::Serialize;
+use tokio::sync::mpsc;
+
+use super::ExecServerError;
+
+pub(super) struct JsonRpcBackend {
+    write_tx: mpsc::Sender<JSONRPCMessage>,
+}
+
+impl JsonRpcBackend {
+    pub(super) fn new(write_tx: mpsc::Sender<JSONRPCMessage>) -> Self {
+        Self { write_tx }
+    }
+
+    pub(super) async fn notify<P: Serialize>(
+        &self,
+        method: &str,
+        params: &P,
+    ) -> Result<(), ExecServerError> {
+        let params = serde_json::to_value(params)?;
+        self.write_tx
+            .send(JSONRPCMessage::Notification(JSONRPCNotification {
+                method: method.to_string(),
+                params: Some(params),
+            }))
+            .await
+            .map_err(|_| ExecServerError::Closed)
+    }
+
+    pub(super) async fn send_request<P: Serialize>(
+        &self,
+        request_id: RequestId,
+        method: &str,
+        params: &P,
+    ) -> Result<(), ExecServerError> {
+        let params = serde_json::to_value(params)?;
+        self.write_tx
+            .send(JSONRPCMessage::Request(JSONRPCRequest {
+                id: request_id,
+                method: method.to_string(),
+                params: Some(params),
+                trace: None,
+            }))
+            .await
+            .map_err(|_| ExecServerError::Closed)
+    }
+}

codex-rs/exec-server/src/client/local_backend.rs

@@ -0,0 +1,141 @@
+use std::sync::Arc;
+
+use codex_app_server_protocol::FsCopyParams;
+use codex_app_server_protocol::FsCopyResponse;
+use codex_app_server_protocol::FsCreateDirectoryParams;
+use codex_app_server_protocol::FsCreateDirectoryResponse;
+use codex_app_server_protocol::FsGetMetadataParams;
+use codex_app_server_protocol::FsGetMetadataResponse;
+use codex_app_server_protocol::FsReadDirectoryParams;
+use codex_app_server_protocol::FsReadDirectoryResponse;
+use codex_app_server_protocol::FsReadFileParams;
+use codex_app_server_protocol::FsReadFileResponse;
+use codex_app_server_protocol::FsRemoveParams;
+use codex_app_server_protocol::FsRemoveResponse;
+use codex_app_server_protocol::FsWriteFileParams;
+use codex_app_server_protocol::FsWriteFileResponse;
+use tokio::sync::Mutex;
+
+use crate::protocol::ExecParams;
+use crate::protocol::ExecResponse;
+use crate::protocol::INITIALIZED_METHOD;
+use crate::protocol::InitializeResponse;
+use crate::protocol::ReadParams;
+use crate::protocol::ReadResponse;
+use crate::protocol::TerminateParams;
+use crate::protocol::TerminateResponse;
+use crate::protocol::WriteParams;
+use crate::protocol::WriteResponse;
+use crate::server::ExecServerHandler;
+
+use super::ExecServerError;
+use super::server_result_to_client;
+
+#[derive(Clone)]
+pub(super) struct LocalBackend {
+    handler: Arc<Mutex<ExecServerHandler>>,
+}
+
+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<InitializeResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.initialize())
+    }
+
+    pub(super) async fn notify(&self, method: &str) -> Result<(), ExecServerError> {
+        match method {
+            INITIALIZED_METHOD => self
+                .handler
+                .lock()
+                .await
+                .initialized()
+                .map_err(ExecServerError::Protocol),
+            other => Err(ExecServerError::Protocol(format!(
+                "unsupported in-process notification method `{other}`"
+            ))),
+        }
+    }
+
+    pub(super) async fn exec(&self, params: ExecParams) -> Result<ExecResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.exec(params).await)
+    }
+
+    pub(super) async fn exec_read(
+        &self,
+        params: ReadParams,
+    ) -> Result<ReadResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.exec_read(params).await)
+    }
+
+    pub(super) async fn exec_write(
+        &self,
+        params: WriteParams,
+    ) -> Result<WriteResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.exec_write(params).await)
+    }
+
+    pub(super) async fn terminate(
+        &self,
+        params: TerminateParams,
+    ) -> Result<TerminateResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.terminate(params).await)
+    }
+
+    pub(super) async fn fs_read_file(
+        &self,
+        params: FsReadFileParams,
+    ) -> Result<FsReadFileResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.fs_read_file(params).await)
+    }
+
+    pub(super) async fn fs_write_file(
+        &self,
+        params: FsWriteFileParams,
+    ) -> Result<FsWriteFileResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.fs_write_file(params).await)
+    }
+
+    pub(super) async fn fs_create_directory(
+        &self,
+        params: FsCreateDirectoryParams,
+    ) -> Result<FsCreateDirectoryResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.fs_create_directory(params).await)
+    }
+
+    pub(super) async fn fs_get_metadata(
+        &self,
+        params: FsGetMetadataParams,
+    ) -> Result<FsGetMetadataResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.fs_get_metadata(params).await)
+    }
+
+    pub(super) async fn fs_read_directory(
+        &self,
+        params: FsReadDirectoryParams,
+    ) -> Result<FsReadDirectoryResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.fs_read_directory(params).await)
+    }
+
+    pub(super) async fn fs_remove(
+        &self,
+        params: FsRemoveParams,
+    ) -> Result<FsRemoveResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.fs_remove(params).await)
+    }
+
+    pub(super) async fn fs_copy(
+        &self,
+        params: FsCopyParams,
+    ) -> Result<FsCopyResponse, ExecServerError> {
+        server_result_to_client(self.handler.lock().await.fs_copy(params).await)
+    }
+}

codex-rs/exec-server/src/client/process.rs

@@ -0,0 +1,72 @@
+use std::sync::Arc;
+use std::sync::Mutex as StdMutex;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+
+use tokio::sync::broadcast;
+
+use super::ExecServerClient;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(super) struct ExecServerOutput {
+    pub(super) stream: crate::protocol::ExecOutputStream,
+    pub(super) chunk: Vec<u8>,
+}
+
+pub(super) struct ExecServerProcess {
+    pub(super) process_id: String,
+    pub(super) output_rx: broadcast::Receiver<ExecServerOutput>,
+    pub(super) status: Arc<RemoteProcessStatus>,
+    pub(super) client: ExecServerClient,
+}
+
+impl ExecServerProcess {
+    pub(super) fn output_receiver(&self) -> broadcast::Receiver<ExecServerOutput> {
+        self.output_rx.resubscribe()
+    }
+
+    pub(super) fn has_exited(&self) -> bool {
+        self.status.has_exited()
+    }
+
+    pub(super) fn exit_code(&self) -> Option<i32> {
+        self.status.exit_code()
+    }
+
+    pub(super) fn terminate(&self) {
+        let client = self.client.clone();
+        let process_id = self.process_id.clone();
+        tokio::spawn(async move {
+            let _ = client.terminate(&process_id).await;
+        });
+    }
+}
+
+pub(super) struct RemoteProcessStatus {
+    exited: AtomicBool,
+    exit_code: StdMutex<Option<i32>>,
+}
+
+impl RemoteProcessStatus {
+    pub(super) fn new() -> Self {
+        Self {
+            exited: AtomicBool::new(false),
+            exit_code: StdMutex::new(None),
+        }
+    }
+
+    pub(super) fn has_exited(&self) -> bool {
+        self.exited.load(Ordering::SeqCst)
+    }
+
+    pub(super) fn exit_code(&self) -> Option<i32> {
+        self.exit_code.lock().ok().and_then(|guard| *guard)
+    }
+
+    pub(super) fn mark_exited(&self, exit_code: Option<i32>) {
+        self.exited.store(true, Ordering::SeqCst);
+        if let Ok(mut guard) = self.exit_code.lock() {
+            *guard = exit_code;
+        }
+    }
+}

codex-rs/exec-server/src/client/tests.rs

@@ -250,12 +250,7 @@ async fn connect_in_process_rejects_writes_to_unknown_processes() {
         Err(err) => panic!("failed to connect in-process client: {err}"),
     };
 
-    let result = client
-        .write_process(crate::protocol::WriteParams {
-            process_id: "missing".to_string(),
-            chunk: b"input".to_vec().into(),
-        })
-        .await;
+    let result = client.write("missing", b"input".to_vec()).await;
 
     match result {
         Err(ExecServerError::Server { code, message }) => {
@@ -290,7 +285,7 @@ async fn connect_in_process_terminate_marks_process_exited() {
         Err(err) => panic!("failed to start in-process child: {err}"),
     };
 
-    if let Err(err) = client.terminate_session(&process.process_id).await {
+    if let Err(err) = client.terminate(&process.process_id).await {
         panic!("failed to terminate in-process child: {err}");
     }