Refactor ExecServer filesystem split between local and remote (#15232)

For each feature we have:
1. Trait exposed on environment
2. **Local Implementation** of the trait
3. Remote implementation that uses the client to proxy via network
4. Handler implementation that handles PRC requests and calls into
**Local Implementation**

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

@@ -1,8 +1,10 @@
 use crate::ExecServerClient;
 use crate::ExecServerError;
 use crate::RemoteExecServerConnectArgs;
-use crate::fs;
-use crate::fs::ExecutorFileSystem;
+use crate::file_system::ExecutorFileSystem;
+use crate::local_file_system::LocalFileSystem;
+use crate::remote_file_system::RemoteFileSystem;
+use std::sync::Arc;
 
 #[derive(Clone, Default)]
 pub struct Environment {
@@ -56,8 +58,12 @@ impl Environment {
         self.remote_exec_server_client.as_ref()
     }
 
-    pub fn get_filesystem(&self) -> impl ExecutorFileSystem + use<> {
-        fs::LocalFileSystem
+    pub fn get_filesystem(&self) -> Arc<dyn ExecutorFileSystem> {
+        if let Some(client) = self.remote_exec_server_client.clone() {
+            Arc::new(RemoteFileSystem::new(client))
+        } else {
+            Arc::new(LocalFileSystem)
+        }
     }
 }
 

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

@@ -0,0 +1,65 @@
+use async_trait::async_trait;
+use codex_utils_absolute_path::AbsolutePathBuf;
+use tokio::io;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct CreateDirectoryOptions {
+    pub recursive: bool,
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct RemoveOptions {
+    pub recursive: bool,
+    pub force: bool,
+}
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub struct CopyOptions {
+    pub recursive: bool,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct FileMetadata {
+    pub is_directory: bool,
+    pub is_file: bool,
+    pub created_at_ms: i64,
+    pub modified_at_ms: i64,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct ReadDirectoryEntry {
+    pub file_name: String,
+    pub is_directory: bool,
+    pub is_file: bool,
+}
+
+pub type FileSystemResult<T> = io::Result<T>;
+
+#[async_trait]
+pub trait ExecutorFileSystem: Send + Sync {
+    async fn read_file(&self, path: &AbsolutePathBuf) -> FileSystemResult<Vec<u8>>;
+
+    async fn write_file(&self, path: &AbsolutePathBuf, contents: Vec<u8>) -> FileSystemResult<()>;
+
+    async fn create_directory(
+        &self,
+        path: &AbsolutePathBuf,
+        options: CreateDirectoryOptions,
+    ) -> FileSystemResult<()>;
+
+    async fn get_metadata(&self, path: &AbsolutePathBuf) -> FileSystemResult<FileMetadata>;
+
+    async fn read_directory(
+        &self,
+        path: &AbsolutePathBuf,
+    ) -> FileSystemResult<Vec<ReadDirectoryEntry>>;
+
+    async fn remove(&self, path: &AbsolutePathBuf, options: RemoveOptions) -> FileSystemResult<()>;
+
+    async fn copy(
+        &self,
+        source_path: &AbsolutePathBuf,
+        destination_path: &AbsolutePathBuf,
+        options: CopyOptions,
+    ) -> FileSystemResult<()>;
+}

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

@@ -2,8 +2,10 @@ mod client;
 mod client_api;
 mod connection;
 mod environment;
-mod fs;
+mod file_system;
+mod local_file_system;
 mod protocol;
+mod remote_file_system;
 mod rpc;
 mod server;
 
@@ -28,13 +30,13 @@ pub use codex_app_server_protocol::FsRemoveResponse;
 pub use codex_app_server_protocol::FsWriteFileParams;
 pub use codex_app_server_protocol::FsWriteFileResponse;
 pub use environment::Environment;
-pub use fs::CopyOptions;
-pub use fs::CreateDirectoryOptions;
-pub use fs::ExecutorFileSystem;
-pub use fs::FileMetadata;
-pub use fs::FileSystemResult;
-pub use fs::ReadDirectoryEntry;
-pub use fs::RemoveOptions;
+pub use file_system::CopyOptions;
+pub use file_system::CreateDirectoryOptions;
+pub use file_system::ExecutorFileSystem;
+pub use file_system::FileMetadata;
+pub use file_system::FileSystemResult;
+pub use file_system::ReadDirectoryEntry;
+pub use file_system::RemoveOptions;
 pub use protocol::ExecExitedNotification;
 pub use protocol::ExecOutputDeltaNotification;
 pub use protocol::ExecOutputStream;

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

@@ -7,70 +7,16 @@ use std::time::SystemTime;
 use std::time::UNIX_EPOCH;
 use tokio::io;
 
+use crate::CopyOptions;
+use crate::CreateDirectoryOptions;
+use crate::ExecutorFileSystem;
+use crate::FileMetadata;
+use crate::FileSystemResult;
+use crate::ReadDirectoryEntry;
+use crate::RemoveOptions;
+
 const MAX_READ_FILE_BYTES: u64 = 512 * 1024 * 1024;
 
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub struct CreateDirectoryOptions {
-    pub recursive: bool,
-}
-
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub struct RemoveOptions {
-    pub recursive: bool,
-    pub force: bool,
-}
-
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub struct CopyOptions {
-    pub recursive: bool,
-}
-
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub struct FileMetadata {
-    pub is_directory: bool,
-    pub is_file: bool,
-    pub created_at_ms: i64,
-    pub modified_at_ms: i64,
-}
-
-#[derive(Clone, Debug, Eq, PartialEq)]
-pub struct ReadDirectoryEntry {
-    pub file_name: String,
-    pub is_directory: bool,
-    pub is_file: bool,
-}
-
-pub type FileSystemResult<T> = io::Result<T>;
-
-#[async_trait]
-pub trait ExecutorFileSystem: Send + Sync {
-    async fn read_file(&self, path: &AbsolutePathBuf) -> FileSystemResult<Vec<u8>>;
-
-    async fn write_file(&self, path: &AbsolutePathBuf, contents: Vec<u8>) -> FileSystemResult<()>;
-
-    async fn create_directory(
-        &self,
-        path: &AbsolutePathBuf,
-        options: CreateDirectoryOptions,
-    ) -> FileSystemResult<()>;
-
-    async fn get_metadata(&self, path: &AbsolutePathBuf) -> FileSystemResult<FileMetadata>;
-
-    async fn read_directory(
-        &self,
-        path: &AbsolutePathBuf,
-    ) -> FileSystemResult<Vec<ReadDirectoryEntry>>;
-
-    async fn remove(&self, path: &AbsolutePathBuf, options: RemoveOptions) -> FileSystemResult<()>;
-
-    async fn copy(
-        &self,
-        source_path: &AbsolutePathBuf,
-        destination_path: &AbsolutePathBuf,
-        options: CopyOptions,
-    ) -> FileSystemResult<()>;
-}
-
 #[derive(Clone, Default)]
 pub(crate) struct LocalFileSystem;
 

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

@@ -0,0 +1,154 @@
+use async_trait::async_trait;
+use base64::Engine as _;
+use base64::engine::general_purpose::STANDARD;
+use codex_app_server_protocol::FsCopyParams;
+use codex_app_server_protocol::FsCreateDirectoryParams;
+use codex_app_server_protocol::FsGetMetadataParams;
+use codex_app_server_protocol::FsReadDirectoryParams;
+use codex_app_server_protocol::FsReadFileParams;
+use codex_app_server_protocol::FsRemoveParams;
+use codex_app_server_protocol::FsWriteFileParams;
+use codex_utils_absolute_path::AbsolutePathBuf;
+use tokio::io;
+
+use crate::CopyOptions;
+use crate::CreateDirectoryOptions;
+use crate::ExecServerClient;
+use crate::ExecServerError;
+use crate::ExecutorFileSystem;
+use crate::FileMetadata;
+use crate::FileSystemResult;
+use crate::ReadDirectoryEntry;
+use crate::RemoveOptions;
+
+const INVALID_REQUEST_ERROR_CODE: i64 = -32600;
+
+#[derive(Clone)]
+pub(crate) struct RemoteFileSystem {
+    client: ExecServerClient,
+}
+
+impl RemoteFileSystem {
+    pub(crate) fn new(client: ExecServerClient) -> Self {
+        Self { client }
+    }
+}
+
+#[async_trait]
+impl ExecutorFileSystem for RemoteFileSystem {
+    async fn read_file(&self, path: &AbsolutePathBuf) -> FileSystemResult<Vec<u8>> {
+        let response = self
+            .client
+            .fs_read_file(FsReadFileParams { path: path.clone() })
+            .await
+            .map_err(map_remote_error)?;
+        STANDARD.decode(response.data_base64).map_err(|err| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("remote fs/readFile returned invalid base64 dataBase64: {err}"),
+            )
+        })
+    }
+
+    async fn write_file(&self, path: &AbsolutePathBuf, contents: Vec<u8>) -> FileSystemResult<()> {
+        self.client
+            .fs_write_file(FsWriteFileParams {
+                path: path.clone(),
+                data_base64: STANDARD.encode(contents),
+            })
+            .await
+            .map_err(map_remote_error)?;
+        Ok(())
+    }
+
+    async fn create_directory(
+        &self,
+        path: &AbsolutePathBuf,
+        options: CreateDirectoryOptions,
+    ) -> FileSystemResult<()> {
+        self.client
+            .fs_create_directory(FsCreateDirectoryParams {
+                path: path.clone(),
+                recursive: Some(options.recursive),
+            })
+            .await
+            .map_err(map_remote_error)?;
+        Ok(())
+    }
+
+    async fn get_metadata(&self, path: &AbsolutePathBuf) -> FileSystemResult<FileMetadata> {
+        let response = self
+            .client
+            .fs_get_metadata(FsGetMetadataParams { path: path.clone() })
+            .await
+            .map_err(map_remote_error)?;
+        Ok(FileMetadata {
+            is_directory: response.is_directory,
+            is_file: response.is_file,
+            created_at_ms: response.created_at_ms,
+            modified_at_ms: response.modified_at_ms,
+        })
+    }
+
+    async fn read_directory(
+        &self,
+        path: &AbsolutePathBuf,
+    ) -> FileSystemResult<Vec<ReadDirectoryEntry>> {
+        let response = self
+            .client
+            .fs_read_directory(FsReadDirectoryParams { path: path.clone() })
+            .await
+            .map_err(map_remote_error)?;
+        Ok(response
+            .entries
+            .into_iter()
+            .map(|entry| ReadDirectoryEntry {
+                file_name: entry.file_name,
+                is_directory: entry.is_directory,
+                is_file: entry.is_file,
+            })
+            .collect())
+    }
+
+    async fn remove(&self, path: &AbsolutePathBuf, options: RemoveOptions) -> FileSystemResult<()> {
+        self.client
+            .fs_remove(FsRemoveParams {
+                path: path.clone(),
+                recursive: Some(options.recursive),
+                force: Some(options.force),
+            })
+            .await
+            .map_err(map_remote_error)?;
+        Ok(())
+    }
+
+    async fn copy(
+        &self,
+        source_path: &AbsolutePathBuf,
+        destination_path: &AbsolutePathBuf,
+        options: CopyOptions,
+    ) -> FileSystemResult<()> {
+        self.client
+            .fs_copy(FsCopyParams {
+                source_path: source_path.clone(),
+                destination_path: destination_path.clone(),
+                recursive: options.recursive,
+            })
+            .await
+            .map_err(map_remote_error)?;
+        Ok(())
+    }
+}
+
+fn map_remote_error(error: ExecServerError) -> io::Error {
+    match error {
+        ExecServerError::Server { code, message } if code == INVALID_REQUEST_ERROR_CODE => {
+            io::Error::new(io::ErrorKind::InvalidInput, message)
+        }
+        ExecServerError::Server { message, .. } => io::Error::other(message),
+        ExecServerError::Closed => {
+            io::Error::new(io::ErrorKind::BrokenPipe, "exec-server transport closed")
+        }
+        _ => io::Error::other(error.to_string()),
+    }
+}

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

@@ -1,4 +1,4 @@
-mod filesystem;
+mod file_system_handler;
 mod handler;
 mod processor;
 mod registry;

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

@@ -1,5 +1,4 @@
 use std::io;
-use std::sync::Arc;
 
 use base64::Engine as _;
 use base64::engine::general_purpose::STANDARD;
@@ -22,26 +21,18 @@ use codex_app_server_protocol::JSONRPCErrorError;
 
 use crate::CopyOptions;
 use crate::CreateDirectoryOptions;
-use crate::Environment;
 use crate::ExecutorFileSystem;
 use crate::RemoveOptions;
+use crate::local_file_system::LocalFileSystem;
 use crate::rpc::internal_error;
 use crate::rpc::invalid_request;
 
-#[derive(Clone)]
-pub(crate) struct ExecServerFileSystem {
-    file_system: Arc<dyn ExecutorFileSystem>,
+#[derive(Clone, Default)]
+pub(crate) struct FileSystemHandler {
+    file_system: LocalFileSystem,
 }
 
-impl Default for ExecServerFileSystem {
-    fn default() -> Self {
-        Self {
-            file_system: Arc::new(Environment::default().get_filesystem()),
-        }
-    }
-}
-
-impl ExecServerFileSystem {
+impl FileSystemHandler {
     pub(crate) async fn read_file(
         &self,
         params: FsReadFileParams,

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

@@ -43,7 +43,7 @@ use crate::rpc::RpcNotificationSender;
 use crate::rpc::internal_error;
 use crate::rpc::invalid_params;
 use crate::rpc::invalid_request;
-use crate::server::filesystem::ExecServerFileSystem;
+use crate::server::file_system_handler::FileSystemHandler;
 
 const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024;
 #[cfg(test)]
@@ -75,7 +75,7 @@ enum ProcessEntry {
 
 pub(crate) struct ExecServerHandler {
     notifications: RpcNotificationSender,
-    file_system: ExecServerFileSystem,
+    file_system: FileSystemHandler,
     processes: Arc<Mutex<HashMap<String, ProcessEntry>>>,
     initialize_requested: AtomicBool,
     initialized: AtomicBool,
@@ -85,7 +85,7 @@ impl ExecServerHandler {
     pub(crate) fn new(notifications: RpcNotificationSender) -> Self {
         Self {
             notifications,
-            file_system: ExecServerFileSystem::default(),
+            file_system: FileSystemHandler::default(),
             processes: Arc::new(Mutex::new(HashMap::new())),
             initialize_requested: AtomicBool::new(false),
             initialized: AtomicBool::new(false),

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

@@ -59,6 +59,7 @@ async fn run_websocket_listener(
     let listener = TcpListener::bind(bind_address).await?;
     let local_addr = listener.local_addr()?;
     tracing::info!("codex-exec-server listening on ws://{local_addr}");
+    println!("ws://{local_addr}");
 
     loop {
         let (stream, peer_addr) = listener.accept().await?;