Add exec-server exec RPC implementation (#15090)

Stacked PR 2/3, based on the stub PR.

Adds the exec RPC implementation and process/event flow in exec-server
only.

---------

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

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

@@ -1,20 +1,65 @@
 use std::sync::Arc;
 use std::time::Duration;
 
+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 codex_app_server_protocol::JSONRPCNotification;
+use serde_json::Value;
+use tokio::sync::broadcast;
+use tokio::sync::mpsc;
 use tokio::time::timeout;
 use tokio_tungstenite::connect_async;
+use tracing::debug;
 use tracing::warn;
 
 use crate::client_api::ExecServerClientConnectOptions;
+use crate::client_api::ExecServerEvent;
 use crate::client_api::RemoteExecServerConnectArgs;
 use crate::connection::JsonRpcConnection;
+use crate::protocol::EXEC_EXITED_METHOD;
+use crate::protocol::EXEC_METHOD;
+use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
+use crate::protocol::EXEC_READ_METHOD;
+use crate::protocol::EXEC_TERMINATE_METHOD;
+use crate::protocol::EXEC_WRITE_METHOD;
+use crate::protocol::ExecExitedNotification;
+use crate::protocol::ExecOutputDeltaNotification;
+use crate::protocol::ExecParams;
+use crate::protocol::ExecResponse;
+use crate::protocol::FS_COPY_METHOD;
+use crate::protocol::FS_CREATE_DIRECTORY_METHOD;
+use crate::protocol::FS_GET_METADATA_METHOD;
+use crate::protocol::FS_READ_DIRECTORY_METHOD;
+use crate::protocol::FS_READ_FILE_METHOD;
+use crate::protocol::FS_REMOVE_METHOD;
+use crate::protocol::FS_WRITE_FILE_METHOD;
 use crate::protocol::INITIALIZE_METHOD;
 use crate::protocol::INITIALIZED_METHOD;
 use crate::protocol::InitializeParams;
 use crate::protocol::InitializeResponse;
+use crate::protocol::ReadParams;
+use crate::protocol::ReadResponse;
+use crate::protocol::TerminateParams;
+use crate::protocol::TerminateResponse;
+use crate::protocol::WriteParams;
+use crate::protocol::WriteResponse;
 use crate::rpc::RpcCallError;
 use crate::rpc::RpcClient;
 use crate::rpc::RpcClientEvent;
+use crate::rpc::RpcNotificationSender;
+use crate::rpc::RpcServerOutboundMessage;
 
 mod local_backend;
 use local_backend::LocalBackend;
@@ -74,6 +119,7 @@ impl ClientBackend {
 
 struct Inner {
     backend: ClientBackend,
+    events_tx: broadcast::Sender<ExecServerEvent>,
     reader_task: tokio::task::JoinHandle<()>,
 }
 
@@ -124,11 +170,32 @@ impl ExecServerClient {
     pub async fn connect_in_process(
         options: ExecServerClientConnectOptions,
     ) -> Result<Self, ExecServerError> {
-        let backend = LocalBackend::new(crate::server::ExecServerHandler::new());
-        let inner = Arc::new(Inner {
-            backend: ClientBackend::InProcess(backend),
-            reader_task: tokio::spawn(async {}),
+        let (outgoing_tx, mut outgoing_rx) = mpsc::channel::<RpcServerOutboundMessage>(256);
+        let backend = LocalBackend::new(crate::server::ExecServerHandler::new(
+            RpcNotificationSender::new(outgoing_tx),
+        ));
+        let inner = Arc::new_cyclic(|weak| {
+            let weak = weak.clone();
+            let reader_task = tokio::spawn(async move {
+                while let Some(message) = outgoing_rx.recv().await {
+                    if let Some(inner) = weak.upgrade()
+                        && let Err(err) = handle_in_process_outbound_message(&inner, message).await
+                    {
+                        warn!(
+                            "in-process exec-server client closing after unexpected response: {err}"
+                        );
+                        return;
+                    }
+                }
+            });
+
+            Inner {
+                backend: ClientBackend::InProcess(backend),
+                events_tx: broadcast::channel(256).0,
+                reader_task,
+            }
         });
+
         let client = Self { inner };
         client.initialize(options).await?;
         Ok(client)
@@ -160,6 +227,10 @@ impl ExecServerClient {
         .await
     }
 
+    pub fn event_receiver(&self) -> broadcast::Receiver<ExecServerEvent> {
+        self.inner.events_tx.subscribe()
+    }
+
     pub async fn initialize(
         &self,
         options: ExecServerClientConnectOptions,
@@ -190,36 +261,234 @@ impl ExecServerClient {
         })?
     }
 
+    pub async fn exec(&self, params: ExecParams) -> Result<ExecResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.exec(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during exec".to_string(),
+            ));
+        };
+        remote.call(EXEC_METHOD, &params).await.map_err(Into::into)
+    }
+
+    pub async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.exec_read(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during read".to_string(),
+            ));
+        };
+        remote
+            .call(EXEC_READ_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn write(
+        &self,
+        process_id: &str,
+        chunk: Vec<u8>,
+    ) -> Result<WriteResponse, ExecServerError> {
+        let params = WriteParams {
+            process_id: process_id.to_string(),
+            chunk: chunk.into(),
+        };
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.exec_write(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during write".to_string(),
+            ));
+        };
+        remote
+            .call(EXEC_WRITE_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError> {
+        let params = TerminateParams {
+            process_id: process_id.to_string(),
+        };
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.terminate(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during terminate".to_string(),
+            ));
+        };
+        remote
+            .call(EXEC_TERMINATE_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn fs_read_file(
+        &self,
+        params: FsReadFileParams,
+    ) -> Result<FsReadFileResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.fs_read_file(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during fs/readFile".to_string(),
+            ));
+        };
+        remote
+            .call(FS_READ_FILE_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn fs_write_file(
+        &self,
+        params: FsWriteFileParams,
+    ) -> Result<FsWriteFileResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.fs_write_file(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during fs/writeFile".to_string(),
+            ));
+        };
+        remote
+            .call(FS_WRITE_FILE_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn fs_create_directory(
+        &self,
+        params: FsCreateDirectoryParams,
+    ) -> Result<FsCreateDirectoryResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.fs_create_directory(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during fs/createDirectory".to_string(),
+            ));
+        };
+        remote
+            .call(FS_CREATE_DIRECTORY_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn fs_get_metadata(
+        &self,
+        params: FsGetMetadataParams,
+    ) -> Result<FsGetMetadataResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.fs_get_metadata(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during fs/getMetadata".to_string(),
+            ));
+        };
+        remote
+            .call(FS_GET_METADATA_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn fs_read_directory(
+        &self,
+        params: FsReadDirectoryParams,
+    ) -> Result<FsReadDirectoryResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.fs_read_directory(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during fs/readDirectory".to_string(),
+            ));
+        };
+        remote
+            .call(FS_READ_DIRECTORY_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn fs_remove(
+        &self,
+        params: FsRemoveParams,
+    ) -> Result<FsRemoveResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.fs_remove(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during fs/remove".to_string(),
+            ));
+        };
+        remote
+            .call(FS_REMOVE_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
+    pub async fn fs_copy(&self, params: FsCopyParams) -> Result<FsCopyResponse, ExecServerError> {
+        if let Some(backend) = self.inner.backend.as_local() {
+            return backend.fs_copy(params).await;
+        }
+        let Some(remote) = self.inner.backend.as_remote() else {
+            return Err(ExecServerError::Protocol(
+                "remote backend missing during fs/copy".to_string(),
+            ));
+        };
+        remote
+            .call(FS_COPY_METHOD, &params)
+            .await
+            .map_err(Into::into)
+    }
+
     async fn connect(
         connection: JsonRpcConnection,
         options: ExecServerClientConnectOptions,
     ) -> Result<Self, ExecServerError> {
         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}");
+        let inner = Arc::new_cyclic(|weak| {
+            let weak = weak.clone();
+            let reader_task = tokio::spawn(async move {
+                while let Some(event) = events_rx.recv().await {
+                    match event {
+                        RpcClientEvent::Notification(notification) => {
+                            if let Some(inner) = weak.upgrade()
+                                && let Err(err) =
+                                    handle_server_notification(&inner, notification).await
+                            {
+                                warn!("exec-server client closing after protocol error: {err}");
+                                return;
+                            }
+                        }
+                        RpcClientEvent::Disconnected { reason } => {
+                            if let Some(reason) = reason {
+                                warn!("exec-server client transport disconnected: {reason}");
+                            }
+                            return;
                         }
-                        return;
                     }
                 }
+            });
+
+            Inner {
+                backend: ClientBackend::Remote(rpc_client),
+                events_tx: broadcast::channel(256).0,
+                reader_task,
             }
         });
 
-        let client = Self {
-            inner: Arc::new(Inner {
-                backend: ClientBackend::Remote(rpc_client),
-                reader_task,
-            }),
-        };
+        let client = Self { inner };
         client.initialize(options).await?;
         Ok(client)
     }
@@ -247,3 +516,39 @@ impl From<RpcCallError> for ExecServerError {
         }
     }
 }
+
+async fn handle_in_process_outbound_message(
+    inner: &Arc<Inner>,
+    message: RpcServerOutboundMessage,
+) -> Result<(), ExecServerError> {
+    match message {
+        RpcServerOutboundMessage::Response { .. } | RpcServerOutboundMessage::Error { .. } => Err(
+            ExecServerError::Protocol("unexpected in-process RPC response".to_string()),
+        ),
+        RpcServerOutboundMessage::Notification(notification) => {
+            handle_server_notification(inner, notification).await
+        }
+    }
+}
+
+async fn handle_server_notification(
+    inner: &Arc<Inner>,
+    notification: JSONRPCNotification,
+) -> Result<(), ExecServerError> {
+    match notification.method.as_str() {
+        EXEC_OUTPUT_DELTA_METHOD => {
+            let params: ExecOutputDeltaNotification =
+                serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
+            let _ = inner.events_tx.send(ExecServerEvent::OutputDelta(params));
+        }
+        EXEC_EXITED_METHOD => {
+            let params: ExecExitedNotification =
+                serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
+            let _ = inner.events_tx.send(ExecServerEvent::Exited(params));
+        }
+        other => {
+            debug!("ignoring unknown exec-server notification: {other}");
+        }
+    }
+    Ok(())
+}