diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index a6d3df0fef..b90d23f7a4 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::collections::HashSet; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -116,6 +117,7 @@ impl RpcNotificationSender { pub(crate) struct RpcRouter { request_routes: HashMap<&'static str, RequestRoute>, + batchable_request_methods: HashSet<&'static str>, notification_routes: HashMap<&'static str, NotificationRoute>, } @@ -123,6 +125,7 @@ impl Default for RpcRouter { fn default() -> Self { Self { request_routes: HashMap::new(), + batchable_request_methods: HashSet::new(), notification_routes: HashMap::new(), } } @@ -143,6 +146,7 @@ where F: Fn(Arc, P) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { + self.batchable_request_methods.remove(method); self.request_routes.insert( method, Box::new(move |state, request| { @@ -172,12 +176,24 @@ where ); } + pub(crate) fn batchable_request(&mut self, method: &'static str, handler: F) + where + P: DeserializeOwned + Send + 'static, + R: Serialize + Send + 'static, + F: Fn(Arc, P) -> Fut + Send + Sync + 'static, + Fut: Future> + Send + 'static, + { + self.request(method, handler); + self.batchable_request_methods.insert(method); + } + pub(crate) fn request_with_id(&mut self, method: &'static str, handler: F) where P: DeserializeOwned + Send + 'static, F: Fn(Arc, RequestId, P) -> Fut + Send + Sync + 'static, Fut: Future> + Send + 'static, { + self.batchable_request_methods.remove(method); self.request_routes.insert( method, Box::new(move |state, request| { @@ -226,6 +242,10 @@ where self.request_routes.get(method) } + pub(crate) fn is_request_batchable(&self, method: &str) -> bool { + self.batchable_request_methods.contains(method) + } + pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { self.notification_routes.get(method) } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index d1ccd57463..ca0fa1a427 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -10,10 +10,6 @@ use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::MAX_RPC_BATCH_REQUESTS; -use crate::protocol::FS_CANONICALIZE_METHOD; -use crate::protocol::FS_GET_METADATA_METHOD; -use crate::protocol::FS_READ_DIRECTORY_METHOD; -use crate::protocol::FS_READ_FILE_METHOD; use crate::rpc::RpcNotificationSender; use crate::rpc::RpcRouter; use crate::rpc::RpcServerOutboundMessage; @@ -194,7 +190,7 @@ async fn run_connection( async move { match message { codex_app_server_protocol::JSONRPCMessage::Request(request) - if !is_batchable_request_method(request.method.as_str()) => + if !router.is_request_batchable(request.method.as_str()) => { Some(RpcServerOutboundMessage::Error { request_id: request.id, @@ -276,18 +272,6 @@ async fn run_connection( let _ = outbound_task.await; } -fn is_batchable_request_method(method: &str) -> bool { - // Batch handling is only for remote skill discovery lookups. Keep it read-only so concurrent - // execution cannot reorder mutations, process I/O, HTTP side effects, or file handle lifetimes. - matches!( - method, - FS_CANONICALIZE_METHOD - | FS_GET_METADATA_METHOD - | FS_READ_DIRECTORY_METHOD - | FS_READ_FILE_METHOD - ) -} - async fn dispatch_request( router: Arc>, handler: Arc, diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 8f48aeaf99..29dd7af6f4 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -93,7 +93,7 @@ pub(crate) fn build_router() -> RpcRouter { handler.terminate(params).await }, ); - router.request( + router.batchable_request( FS_READ_FILE_METHOD, |handler: Arc, params: FsReadFileParams| async move { handler.fs_read_file(params).await @@ -129,19 +129,19 @@ pub(crate) fn build_router() -> RpcRouter { handler.fs_create_directory(params).await }, ); - router.request( + router.batchable_request( FS_GET_METADATA_METHOD, |handler: Arc, params: FsGetMetadataParams| async move { handler.fs_get_metadata(params).await }, ); - router.request( + router.batchable_request( FS_CANONICALIZE_METHOD, |handler: Arc, params: FsCanonicalizeParams| async move { handler.fs_canonicalize(params).await }, ); - router.request( + router.batchable_request( FS_READ_DIRECTORY_METHOD, |handler: Arc, params: FsReadDirectoryParams| async move { handler.fs_read_directory(params).await diff --git a/codex-rs/exec-server/tests/file_system/shared.rs b/codex-rs/exec-server/tests/file_system/shared.rs index 08ba2efdd8..2d7cc1c668 100644 --- a/codex-rs/exec-server/tests/file_system/shared.rs +++ b/codex-rs/exec-server/tests/file_system/shared.rs @@ -4,6 +4,8 @@ use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::FILE_READ_CHUNK_SIZE; use codex_exec_server::FileMetadata; +use codex_exec_server::FileSystemOperation; +use codex_exec_server::FileSystemOperationOutput; use codex_exec_server::ReadDirectoryEntry; use codex_exec_server::RemoveOptions; use codex_protocol::models::AdditionalPermissionProfile; @@ -196,6 +198,55 @@ async fn file_system_read_file_returns_bytes( Ok(()) } +#[test_case(FileSystemImplementation::Local ; "local")] +#[test_case(FileSystemImplementation::Remote ; "remote")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn file_system_batch_preserves_operation_order( + implementation: FileSystemImplementation, +) -> Result<()> { + let context = create_file_system_context(implementation).await?; + let file_system = context.file_system; + + let tmp = TempDir::new()?; + let first_path = tmp.path().join("first.txt"); + let second_path = tmp.path().join("second.txt"); + std::fs::write(&first_path, "first")?; + std::fs::write(&second_path, "second")?; + let first_uri = PathUri::from_path(&first_path)?; + let second_uri = PathUri::from_path(&second_path)?; + + let outputs = file_system + .execute_batch( + vec![ + FileSystemOperation::ReadFile { + path: first_uri.clone(), + }, + FileSystemOperation::Canonicalize { + path: second_uri.clone(), + }, + FileSystemOperation::ReadFile { path: second_uri }, + ], + /*sandbox*/ None, + ) + .await + .with_context(|| format!("mode={implementation}"))? + .into_iter() + .collect::>>()?; + + assert_eq!( + outputs, + vec![ + FileSystemOperationOutput::ReadFile(b"first".to_vec()), + FileSystemOperationOutput::Canonicalize(PathUri::from_path(std::fs::canonicalize( + second_path + )?)?), + FileSystemOperationOutput::ReadFile(b"second".to_vec()), + ] + ); + + Ok(()) +} + #[test_case(FileSystemImplementation::Local ; "local")] #[test_case(FileSystemImplementation::Remote ; "remote")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)]