From 68b5a3835390f18be91b6c398624b89998c32c20 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Thu, 9 Jul 2026 12:20:27 +0000 Subject: [PATCH] feat(exec-server): batch upward filesystem searches (cherry picked from commit aacdc03d678eba1adfd09627848a9bd5b701623d) --- codex-rs/exec-server-protocol/src/protocol.rs | 58 ++++++ codex-rs/exec-server/src/client.rs | 10 + codex-rs/exec-server/src/lib.rs | 5 + .../exec-server/src/remote_file_system.rs | 76 +++++++ .../src/remote_file_system_path_uri_tests.rs | 191 ++++++++++++++++++ .../src/server/file_system_handler.rs | 62 ++++++ codex-rs/exec-server/src/server/handler.rs | 10 + codex-rs/exec-server/src/server/registry.rs | 8 + .../exec-server/tests/file_system/shared.rs | 110 ++++++++++ codex-rs/file-system/src/find_up.rs | 152 +++++++++++++- codex-rs/file-system/src/lib.rs | 16 ++ 11 files changed, 697 insertions(+), 1 deletion(-) diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index fbf01c322d..f5e667e03c 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -6,6 +6,7 @@ pub use codex_file_system::FindUpErrorPolicy; pub use codex_file_system::FindUpMatchKind; pub use codex_file_system::FindUpOptions; pub use codex_file_system::FindUpOutcome; +pub use codex_file_system::FindUpRequest; pub use codex_file_system::WalkOptions; pub use codex_file_system::WalkOutcome; use codex_network_proxy::ManagedNetworkSandboxContext; @@ -38,6 +39,7 @@ pub const FS_CREATE_DIRECTORY_METHOD: &str = "fs/createDirectory"; pub const FS_GET_METADATA_METHOD: &str = "fs/getMetadata"; pub const FS_GET_METADATA_BATCH_METHOD: &str = "fs/getMetadataBatch"; pub const FS_FIND_UP_METHOD: &str = "fs/findUp"; +pub const FS_FIND_UP_BATCH_METHOD: &str = "fs/findUpBatch"; pub const FS_CANONICALIZE_METHOD: &str = "fs/canonicalize"; pub const FS_READ_DIRECTORY_METHOD: &str = "fs/readDirectory"; pub const FS_WALK_METHOD: &str = "fs/walk"; @@ -386,6 +388,32 @@ pub struct FsFindUpParams { pub type FsFindUpResponse = FindUpOutcome; +/// Maximum number of independent searches accepted by one `fs/findUpBatch` request. +pub const FS_FIND_UP_BATCH_MAX_REQUESTS: usize = 64; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FsFindUpBatchParams { + /// Searches are evaluated and returned in this order. + pub requests: Vec, + /// One sandbox context shared by every search in the batch. + pub sandbox: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum FsFindUpBatchResult { + Outcome { outcome: FindUpOutcome }, + Error { error: JSONRPCErrorError }, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FsFindUpBatchResponse { + /// One result for each requested search, in request order. + pub results: Vec, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct FsCanonicalizeParams { @@ -617,6 +645,8 @@ mod tests { use super::ExecParams; use super::FindUpErrorPolicy; use super::FindUpMatchKind; + use super::FsFindUpBatchParams; + use super::FsFindUpBatchResponse; use super::FsFindUpParams; use super::FsGetMetadataBatchResponse; use super::FsReadFileParams; @@ -730,6 +760,18 @@ mod tests { "sandbox": null, })) .expect_err("native absolute find-up start should not deserialize as a URI"); + serde_json::from_value::(serde_json::json!({ + "requests": [{ + "start": native_path.to_string_lossy(), + "options": { + "candidateRelativePaths": ["marker"], + "matchKind": FindUpMatchKind::Any, + "nonNotFoundErrorPolicy": FindUpErrorPolicy::Propagate, + }, + }], + "sandbox": null, + })) + .expect_err("native absolute batched find-up start should not deserialize as a URI"); let sandbox = FileSystemSandboxContext::from_permission_profile_with_cwd( PermissionProfile::default(), @@ -764,6 +806,22 @@ mod tests { ); } + #[test] + fn find_up_batch_error_result_round_trips() { + let serialized = serde_json::json!({ + "results": [{ + "type": "error", + "error": { "code": -32000, "message": "failed" }, + }], + }); + let response = serde_json::from_value::(serialized.clone()) + .expect("deserialize find-up batch"); + assert_eq!( + serde_json::to_value(response).expect("serialize find-up batch"), + serialized + ); + } + #[test] fn http_request_timeout_treats_omitted_and_null_as_no_timeout() { let omitted: HttpRequestParams = serde_json::from_value(serde_json::json!({ diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 4c6a62301c..3afbd4e811 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -53,6 +53,7 @@ use crate::protocol::FS_CANONICALIZE_METHOD; use crate::protocol::FS_CLOSE_METHOD; use crate::protocol::FS_COPY_METHOD; use crate::protocol::FS_CREATE_DIRECTORY_METHOD; +use crate::protocol::FS_FIND_UP_BATCH_METHOD; use crate::protocol::FS_FIND_UP_METHOD; use crate::protocol::FS_GET_METADATA_BATCH_METHOD; use crate::protocol::FS_GET_METADATA_METHOD; @@ -71,6 +72,8 @@ use crate::protocol::FsCopyParams; use crate::protocol::FsCopyResponse; use crate::protocol::FsCreateDirectoryParams; use crate::protocol::FsCreateDirectoryResponse; +use crate::protocol::FsFindUpBatchParams; +use crate::protocol::FsFindUpBatchResponse; use crate::protocol::FsFindUpParams; use crate::protocol::FsFindUpResponse; use crate::protocol::FsGetMetadataBatchParams; @@ -726,6 +729,13 @@ impl ExecServerClient { self.call(FS_FIND_UP_METHOD, ¶ms).await } + pub async fn fs_find_up_batch( + &self, + params: FsFindUpBatchParams, + ) -> Result { + self.call(FS_FIND_UP_BATCH_METHOD, ¶ms).await + } + pub async fn fs_canonicalize( &self, params: FsCanonicalizeParams, diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 2b5f70aa25..74c9e3700c 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -58,6 +58,7 @@ pub use codex_file_system::FindUpMatch; pub use codex_file_system::FindUpMatchKind; pub use codex_file_system::FindUpOptions; pub use codex_file_system::FindUpOutcome; +pub use codex_file_system::FindUpRequest; pub use codex_file_system::MAX_FIND_UP_CANDIDATE_BYTES; pub use codex_file_system::MAX_FIND_UP_CANDIDATES; pub use codex_file_system::MAX_FIND_UP_TOTAL_CANDIDATE_BYTES; @@ -109,6 +110,7 @@ pub use protocol::ExecOutputDeltaNotification; pub use protocol::ExecOutputStream; pub use protocol::ExecParams; pub use protocol::ExecResponse; +pub use protocol::FS_FIND_UP_BATCH_MAX_REQUESTS; pub use protocol::FS_GET_METADATA_BATCH_MAX_PATHS; pub use protocol::FsCanonicalizeParams; pub use protocol::FsCanonicalizeResponse; @@ -118,6 +120,9 @@ pub use protocol::FsCopyParams; pub use protocol::FsCopyResponse; pub use protocol::FsCreateDirectoryParams; pub use protocol::FsCreateDirectoryResponse; +pub use protocol::FsFindUpBatchParams; +pub use protocol::FsFindUpBatchResponse; +pub use protocol::FsFindUpBatchResult; pub use protocol::FsFindUpParams; pub use protocol::FsFindUpResponse; pub use protocol::FsGetMetadataBatchParams; diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index 3d136cac24..f8b0e0cbe1 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -1,6 +1,7 @@ use base64::Engine as _; use base64::engine::general_purpose::STANDARD; use codex_utils_path_uri::PathUri; +use futures::StreamExt; use tokio::io; use tracing::trace; @@ -15,6 +16,7 @@ use crate::FileSystemResult; use crate::FileSystemSandboxContext; use crate::FindUpOptions; use crate::FindUpOutcome; +use crate::FindUpRequest; use crate::ReadDirectoryEntry; use crate::RemoveOptions; use crate::WalkOptions; @@ -23,6 +25,8 @@ use crate::client::LazyRemoteExecServerClient; use crate::protocol::FsCanonicalizeParams; use crate::protocol::FsCopyParams; use crate::protocol::FsCreateDirectoryParams; +use crate::protocol::FsFindUpBatchParams; +use crate::protocol::FsFindUpBatchResult; use crate::protocol::FsFindUpParams; use crate::protocol::FsGetMetadataBatchParams; use crate::protocol::FsGetMetadataBatchResult; @@ -36,6 +40,7 @@ use crate::protocol::FsWriteFileParams; const INVALID_REQUEST_ERROR_CODE: i64 = -32600; const METHOD_NOT_FOUND_ERROR_CODE: i64 = -32601; const NOT_FOUND_ERROR_CODE: i64 = -32004; +const MAX_CONCURRENT_FIND_UP_FALLBACKS: usize = 8; #[path = "remote_file_stream.rs"] mod file_stream; @@ -242,6 +247,53 @@ impl RemoteFileSystem { } } + async fn find_up_batch( + &self, + requests: &[FindUpRequest], + sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult>> { + let client = self.client.get().await.map_err(map_remote_error)?; + let response = match client + .fs_find_up_batch(FsFindUpBatchParams { + requests: requests.to_vec(), + sandbox: remote_sandbox_context(sandbox), + }) + .await + { + Ok(response) => response, + Err(ExecServerError::Server { + code: METHOD_NOT_FOUND_ERROR_CODE, + .. + }) => { + return find_up_batch_via_individual(self, requests, sandbox).await; + } + Err(error) => return Err(map_remote_error(error)), + }; + if response.results.len() != requests.len() { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "remote fs/findUpBatch returned {} results for {} requests", + response.results.len(), + requests.len() + ), + )); + } + Ok(response + .results + .into_iter() + .map(|result| match result { + FsFindUpBatchResult::Outcome { outcome } => Ok(outcome), + FsFindUpBatchResult::Error { error } => { + Err(map_remote_error(ExecServerError::Server { + code: error.code, + message: error.message, + })) + } + }) + .collect()) + } + async fn read_directory( &self, path: &PathUri, @@ -410,6 +462,14 @@ impl ExecutorFileSystem for RemoteFileSystem { Box::pin(RemoteFileSystem::find_up(self, start, options, sandbox)) } + fn find_up_batch<'a>( + &'a self, + requests: &'a [FindUpRequest], + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec>> { + Box::pin(RemoteFileSystem::find_up_batch(self, requests, sandbox)) + } + fn read_directory<'a>( &'a self, path: &'a PathUri, @@ -453,6 +513,22 @@ impl ExecutorFileSystem for RemoteFileSystem { } } +async fn find_up_batch_via_individual( + file_system: &RemoteFileSystem, + requests: &[FindUpRequest], + sandbox: Option<&FileSystemSandboxContext>, +) -> FileSystemResult>> { + Ok(futures::stream::iter(requests.iter().cloned()) + .map(|request| async move { + file_system + .find_up(&request.start, &request.options, sandbox) + .await + }) + .buffered(MAX_CONCURRENT_FIND_UP_FALLBACKS) + .collect() + .await) +} + fn remote_sandbox_context( sandbox: Option<&FileSystemSandboxContext>, ) -> Option { diff --git a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs index f3cf0ef3bc..a9b1fb5273 100644 --- a/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs +++ b/codex-rs/exec-server/src/remote_file_system_path_uri_tests.rs @@ -1,7 +1,11 @@ #![allow(clippy::expect_used)] +use codex_exec_server_protocol::JSONRPCError; +use codex_exec_server_protocol::JSONRPCErrorError; use codex_exec_server_protocol::JSONRPCMessage; use codex_exec_server_protocol::JSONRPCResponse; +use codex_file_system::FindUpErrorPolicy; +use codex_file_system::FindUpMatchKind; use codex_protocol::models::PermissionProfile; use codex_protocol::permissions::FileSystemAccessMode; use codex_protocol::permissions::FileSystemPath; @@ -25,7 +29,12 @@ use tokio_tungstenite::tungstenite::Message; use super::*; use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT; use crate::client_api::ExecServerTransportParams; +use crate::protocol::FS_FIND_UP_BATCH_METHOD; +use crate::protocol::FS_FIND_UP_METHOD; use crate::protocol::FS_READ_FILE_METHOD; +use crate::protocol::FsFindUpBatchParams; +use crate::protocol::FsFindUpBatchResponse; +use crate::protocol::FsFindUpParams; use crate::protocol::FsReadFileParams; use crate::protocol::FsReadFileResponse; use crate::protocol::INITIALIZE_METHOD; @@ -82,6 +91,188 @@ async fn remote_file_system_sends_path_and_sandbox_cwd_uris_without_native_conve server.await.expect("recording server should succeed"); } +#[tokio::test] +async fn find_up_batch_method_not_found_falls_back_concurrently_and_preserves_order() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!("ws://{}", listener.local_addr().expect("listener address")); + let starts = vec![ + PathUri::parse("file:///C:/Users/Alice/src").expect("valid drive URI"), + PathUri::parse("file://server/share/src").expect("valid UNC URI"), + ]; + let expected_starts = starts.clone(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("listener should accept"); + let mut websocket = accept_async(stream) + .await + .expect("websocket handshake should succeed"); + complete_websocket_initialize(&mut websocket).await; + + let batch_request = match read_jsonrpc_websocket(&mut websocket).await { + JSONRPCMessage::Request(request) if request.method == FS_FIND_UP_BATCH_METHOD => { + request + } + other => panic!("expected fs/findUpBatch request, got {other:?}"), + }; + let batch_params: FsFindUpBatchParams = serde_json::from_value( + batch_request + .params + .expect("fs/findUpBatch params should exist"), + ) + .expect("fs/findUpBatch params should deserialize"); + assert_eq!( + batch_params + .requests + .iter() + .map(|request| request.start.clone()) + .collect::>(), + expected_starts + ); + write_jsonrpc_websocket( + &mut websocket, + JSONRPCMessage::Error(JSONRPCError { + error: JSONRPCErrorError { + code: -32601, + data: None, + message: "method not found".to_string(), + }, + id: batch_request.id, + }), + ) + .await; + + // Read both fallbacks before responding. A serial fallback would time out here. + let mut singles = Vec::new(); + for _ in 0..2 { + let request = match read_jsonrpc_websocket(&mut websocket).await { + JSONRPCMessage::Request(request) if request.method == FS_FIND_UP_METHOD => request, + other => panic!("expected fs/findUp request, got {other:?}"), + }; + let params: FsFindUpParams = serde_json::from_value( + request + .params + .clone() + .expect("fs/findUp params should exist"), + ) + .expect("fs/findUp params should deserialize"); + singles.push((request, params)); + } + for (request, params) in singles.into_iter().rev() { + let visited_ancestor_count = if params.start == expected_starts[0] { + 1 + } else { + 2 + }; + write_jsonrpc_websocket( + &mut websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(FindUpOutcome { + matched: None, + visited_ancestor_count, + metadata_probe_count: visited_ancestor_count, + ignored_error_count: 0, + ignored_errors: Vec::new(), + ignored_errors_truncated: false, + }) + .expect("find-up response should serialize"), + }), + ) + .await; + } + }); + let file_system = RemoteFileSystem::new(LazyRemoteExecServerClient::new( + ExecServerTransportParams::websocket_url( + websocket_url, + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ), + )); + let options = FindUpOptions { + candidate_relative_paths: vec!["marker".to_string()], + match_kind: FindUpMatchKind::Any, + non_not_found_error_policy: FindUpErrorPolicy::Propagate, + }; + let requests = starts + .into_iter() + .map(|start| FindUpRequest { + start, + options: options.clone(), + }) + .collect::>(); + + let outcomes = file_system + .find_up_batch(&requests, /*sandbox*/ None) + .await + .expect("method-not-found should use single-search fallback") + .into_iter() + .collect::, _>>() + .expect("fallback searches should succeed"); + + assert_eq!( + outcomes + .iter() + .map(|outcome| outcome.visited_ancestor_count) + .collect::>(), + vec![1, 2] + ); + server.await.expect("recording server should succeed"); +} + +#[tokio::test] +async fn find_up_batch_rejects_malformed_response_cardinality_without_legacy_fallback() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!("ws://{}", listener.local_addr().expect("listener address")); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("listener should accept"); + let mut websocket = accept_async(stream) + .await + .expect("websocket handshake should succeed"); + complete_websocket_initialize(&mut websocket).await; + let request = match read_jsonrpc_websocket(&mut websocket).await { + JSONRPCMessage::Request(request) if request.method == FS_FIND_UP_BATCH_METHOD => { + request + } + other => panic!("expected fs/findUpBatch request, got {other:?}"), + }; + write_jsonrpc_websocket( + &mut websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(FsFindUpBatchResponse { + results: Vec::new(), + }) + .expect("find-up batch response should serialize"), + }), + ) + .await; + }); + let file_system = RemoteFileSystem::new(LazyRemoteExecServerClient::new( + ExecServerTransportParams::websocket_url( + websocket_url, + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ), + )); + let requests = [FindUpRequest { + start: PathUri::parse("file:///C:/Users/Alice/src").expect("valid drive URI"), + options: FindUpOptions { + candidate_relative_paths: vec!["marker".to_string()], + match_kind: FindUpMatchKind::Any, + non_not_found_error_policy: FindUpErrorPolicy::Propagate, + }, + }]; + + let error = file_system + .find_up_batch(&requests, /*sandbox*/ None) + .await + .expect_err("malformed cardinality should be surfaced to the consumer"); + + assert_eq!(error.kind(), std::io::ErrorKind::InvalidData); + server.await.expect("recording server should succeed"); +} + async fn record_read_file_params( expected_requests: usize, ) -> ( diff --git a/codex-rs/exec-server/src/server/file_system_handler.rs b/codex-rs/exec-server/src/server/file_system_handler.rs index 3be81fb64f..9523ac5915 100644 --- a/codex-rs/exec-server/src/server/file_system_handler.rs +++ b/codex-rs/exec-server/src/server/file_system_handler.rs @@ -11,6 +11,7 @@ use crate::ExecutorFileSystem; use crate::RemoveOptions; use crate::file_read::FileReadHandleManager; use crate::local_file_system::LocalFileSystem; +use crate::protocol::FS_FIND_UP_BATCH_MAX_REQUESTS; use crate::protocol::FS_GET_METADATA_BATCH_MAX_PATHS; use crate::protocol::FS_WRITE_FILE_METHOD; use crate::protocol::FsCanonicalizeParams; @@ -21,6 +22,9 @@ use crate::protocol::FsCopyParams; use crate::protocol::FsCopyResponse; use crate::protocol::FsCreateDirectoryParams; use crate::protocol::FsCreateDirectoryResponse; +use crate::protocol::FsFindUpBatchParams; +use crate::protocol::FsFindUpBatchResponse; +use crate::protocol::FsFindUpBatchResult; use crate::protocol::FsFindUpParams; use crate::protocol::FsFindUpResponse; use crate::protocol::FsGetMetadataBatchParams; @@ -217,6 +221,30 @@ impl FileSystemHandler { .map_err(map_fs_error) } + pub(crate) async fn find_up_batch( + &self, + params: FsFindUpBatchParams, + ) -> Result { + if params.requests.len() > FS_FIND_UP_BATCH_MAX_REQUESTS { + return Err(invalid_request(format!( + "find-up batch must not exceed {FS_FIND_UP_BATCH_MAX_REQUESTS} requests" + ))); + } + let mut results = Vec::with_capacity(params.requests.len()); + for request in params.requests { + match self + .file_system + .find_up(&request.start, &request.options, params.sandbox.as_ref()) + .await + .map_err(map_fs_error) + { + Ok(outcome) => results.push(FsFindUpBatchResult::Outcome { outcome }), + Err(error) => results.push(FsFindUpBatchResult::Error { error }), + } + } + Ok(FsFindUpBatchResponse { results }) + } + pub(crate) async fn canonicalize( &self, params: FsCanonicalizeParams, @@ -419,4 +447,38 @@ mod tests { ); } + #[tokio::test] + async fn find_up_batch_rejects_too_many_requests_before_dispatch() { + let handler = FileSystemHandler::new( + ExecServerRuntimePaths::new( + std::env::current_exe().expect("current exe"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths"), + ); + let request = crate::FindUpRequest { + start: PathUri::from_host_native_path(std::env::current_dir().expect("current dir")) + .expect("current dir URI"), + options: crate::FindUpOptions { + candidate_relative_paths: vec!["marker".to_string()], + match_kind: crate::FindUpMatchKind::Any, + non_not_found_error_policy: crate::FindUpErrorPolicy::Propagate, + }, + }; + + let error = handler + .find_up_batch(FsFindUpBatchParams { + requests: vec![request; FS_FIND_UP_BATCH_MAX_REQUESTS + 1], + sandbox: None, + }) + .await + .expect_err("oversized find-up batch should fail"); + + assert_eq!(error.code, -32600); + assert_eq!( + error.message, + format!("find-up batch must not exceed {FS_FIND_UP_BATCH_MAX_REQUESTS} requests") + ); + } + } diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index 6ef67a4362..df10453641 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -25,6 +25,8 @@ use crate::protocol::FsCopyParams; use crate::protocol::FsCopyResponse; use crate::protocol::FsCreateDirectoryParams; use crate::protocol::FsCreateDirectoryResponse; +use crate::protocol::FsFindUpBatchParams; +use crate::protocol::FsFindUpBatchResponse; use crate::protocol::FsFindUpParams; use crate::protocol::FsFindUpResponse; use crate::protocol::FsGetMetadataBatchParams; @@ -320,6 +322,14 @@ impl ExecServerHandler { self.file_system.find_up(params).await } + pub(crate) async fn fs_find_up_batch( + &self, + params: FsFindUpBatchParams, + ) -> Result { + self.require_initialized_for("filesystem")?; + self.file_system.find_up_batch(params).await + } + pub(crate) async fn fs_canonicalize( &self, params: FsCanonicalizeParams, diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index 272b32f673..c4bc9eea61 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -11,6 +11,7 @@ use crate::protocol::FS_CANONICALIZE_METHOD; use crate::protocol::FS_CLOSE_METHOD; use crate::protocol::FS_COPY_METHOD; use crate::protocol::FS_CREATE_DIRECTORY_METHOD; +use crate::protocol::FS_FIND_UP_BATCH_METHOD; use crate::protocol::FS_FIND_UP_METHOD; use crate::protocol::FS_GET_METADATA_BATCH_METHOD; use crate::protocol::FS_GET_METADATA_METHOD; @@ -25,6 +26,7 @@ use crate::protocol::FsCanonicalizeParams; use crate::protocol::FsCloseParams; use crate::protocol::FsCopyParams; use crate::protocol::FsCreateDirectoryParams; +use crate::protocol::FsFindUpBatchParams; use crate::protocol::FsFindUpParams; use crate::protocol::FsGetMetadataBatchParams; use crate::protocol::FsGetMetadataParams; @@ -153,6 +155,12 @@ pub(crate) fn build_router() -> RpcRouter { handler.fs_find_up(params).await }, ); + router.request( + FS_FIND_UP_BATCH_METHOD, + |handler: Arc, params: FsFindUpBatchParams| async move { + handler.fs_find_up_batch(params).await + }, + ); router.request( FS_CANONICALIZE_METHOD, |handler: Arc, params: FsCanonicalizeParams| async move { diff --git a/codex-rs/exec-server/tests/file_system/shared.rs b/codex-rs/exec-server/tests/file_system/shared.rs index 17c540fd2e..3fb1a2f828 100644 --- a/codex-rs/exec-server/tests/file_system/shared.rs +++ b/codex-rs/exec-server/tests/file_system/shared.rs @@ -3,12 +3,14 @@ use anyhow::Result; use codex_exec_server::CopyOptions; use codex_exec_server::CreateDirectoryOptions; use codex_exec_server::FILE_READ_CHUNK_SIZE; +use codex_exec_server::FS_FIND_UP_BATCH_MAX_REQUESTS; use codex_exec_server::FileMetadata; use codex_exec_server::FindUpErrorPolicy; use codex_exec_server::FindUpMatch; use codex_exec_server::FindUpMatchKind; use codex_exec_server::FindUpOptions; use codex_exec_server::FindUpOutcome; +use codex_exec_server::FindUpRequest; use codex_exec_server::MAX_FIND_UP_CANDIDATE_BYTES; use codex_exec_server::MAX_FIND_UP_CANDIDATES; use codex_exec_server::MAX_FIND_UP_TOTAL_CANDIDATE_BYTES; @@ -200,6 +202,8 @@ async fn file_system_find_up_preserves_ancestor_and_candidate_order( visited_ancestor_count: 1, metadata_probe_count: 2, ignored_error_count: 0, + ignored_errors: Vec::new(), + ignored_errors_truncated: false, } ); @@ -217,6 +221,8 @@ async fn file_system_find_up_preserves_ancestor_and_candidate_order( visited_ancestor_count: 1, metadata_probe_count: 1, ignored_error_count: 0, + ignored_errors: Vec::new(), + ignored_errors_truncated: false, } ); Ok(()) @@ -275,6 +281,110 @@ async fn file_system_find_up_enforces_bounds_and_sandbox( Ok(()) } +#[test_case(FileSystemImplementation::Local ; "local")] +#[test_case(FileSystemImplementation::Remote ; "remote")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn file_system_find_up_batch_preserves_order_errors_and_sandbox( + implementation: FileSystemImplementation, +) -> Result<()> { + let context = create_file_system_context(implementation).await?; + let file_system = context.file_system; + let tmp = TempDir::new()?; + let parent = tmp.path().join("parent"); + let start = parent.join("child"); + std::fs::create_dir_all(&start)?; + std::fs::write(start.join("child.marker"), "child")?; + std::fs::write(parent.join("parent.marker"), "parent")?; + let start_uri = PathUri::from_host_native_path(&start)?; + let parent_uri = PathUri::from_host_native_path(&parent)?; + let sandbox = read_only_sandbox(tmp.path().to_path_buf()); + let options = |candidate_relative_paths| FindUpOptions { + candidate_relative_paths, + match_kind: FindUpMatchKind::File, + non_not_found_error_policy: FindUpErrorPolicy::Propagate, + }; + let requests = vec![ + FindUpRequest { + start: start_uri.clone(), + options: options(vec!["child.marker".to_string()]), + }, + FindUpRequest { + start: start_uri.clone(), + options: options(vec!["../invalid.marker".to_string()]), + }, + FindUpRequest { + start: start_uri.clone(), + options: options(vec!["parent.marker".to_string()]), + }, + ]; + + let results = file_system + .find_up_batch(&requests, Some(&sandbox)) + .await + .with_context(|| format!("mode={implementation}"))?; + let actual = results + .into_iter() + .map(|result| result.map_err(|error| error.kind())) + .collect::>(); + + assert_eq!( + actual, + vec![ + Ok(FindUpOutcome { + matched: Some(FindUpMatch { + ancestor: start_uri.clone(), + path: PathUri::from_host_native_path(start.join("child.marker"))?, + }), + visited_ancestor_count: 1, + metadata_probe_count: 1, + ignored_error_count: 0, + ignored_errors: Vec::new(), + ignored_errors_truncated: false, + }), + Err(std::io::ErrorKind::InvalidInput), + Ok(FindUpOutcome { + matched: Some(FindUpMatch { + ancestor: parent_uri, + path: PathUri::from_host_native_path(parent.join("parent.marker"))?, + }), + visited_ancestor_count: 2, + metadata_probe_count: 2, + ignored_error_count: 0, + ignored_errors: Vec::new(), + ignored_errors_truncated: false, + }), + ] + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_file_system_find_up_batch_enforces_request_cap() -> Result<()> { + let context = create_file_system_context(FileSystemImplementation::Remote).await?; + let file_system = context.file_system; + let tmp = TempDir::new()?; + let start = PathUri::from_host_native_path(tmp.path())?; + let request = FindUpRequest { + start, + options: FindUpOptions { + candidate_relative_paths: vec!["marker".to_string()], + match_kind: FindUpMatchKind::Any, + non_not_found_error_policy: FindUpErrorPolicy::Propagate, + }, + }; + + let error = file_system + .find_up_batch( + &vec![request; FS_FIND_UP_BATCH_MAX_REQUESTS + 1], + /*sandbox*/ None, + ) + .await + .expect_err("oversized find-up batch should fail before work"); + + assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput); + Ok(()) +} + #[test_case(FileSystemImplementation::Local ; "local")] #[test_case(FileSystemImplementation::Remote ; "remote")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/codex-rs/file-system/src/find_up.rs b/codex-rs/file-system/src/find_up.rs index cd225df6bb..4fc9d3fc9b 100644 --- a/codex-rs/file-system/src/find_up.rs +++ b/codex-rs/file-system/src/find_up.rs @@ -7,9 +7,14 @@ use futures::StreamExt; use std::io; const MAX_CONCURRENT_PROBES: usize = 8; +const MAX_CONCURRENT_FIND_UP_REQUESTS: usize = 8; pub const MAX_FIND_UP_CANDIDATES: usize = 16; pub const MAX_FIND_UP_CANDIDATE_BYTES: usize = 256; pub const MAX_FIND_UP_TOTAL_CANDIDATE_BYTES: usize = 2 * 1024; +pub const MAX_FIND_UP_IGNORED_ERRORS: usize = 64; +pub const MAX_FIND_UP_IGNORED_ERROR_MESSAGE_BYTES: usize = 1024; +pub const MAX_FIND_UP_IGNORED_ERROR_DETAILS_BYTES: usize = 64 * 1024; +const FIND_UP_IGNORED_ERROR_ITEM_OVERHEAD_BYTES: usize = 64; /// Controls how an upward marker search handles metadata errors other than `NotFound`. #[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] @@ -50,6 +55,14 @@ pub struct FindUpOptions { pub non_not_found_error_policy: FindUpErrorPolicy, } +/// One independent upward filesystem search in a batch. +#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct FindUpRequest { + pub start: PathUri, + pub options: FindUpOptions, +} + /// First qualifying path found by an upward filesystem search. #[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] #[serde(rename_all = "camelCase")] @@ -58,6 +71,14 @@ pub struct FindUpMatch { pub path: PathUri, } +/// A recoverable metadata error encountered during an upward filesystem search. +#[derive(Clone, Debug, Eq, PartialEq, serde::Deserialize, serde::Serialize)] +#[serde(rename_all = "camelCase")] +pub struct FindUpError { + pub path: PathUri, + pub message: String, +} + /// Result and generic work counters from an upward filesystem search. #[derive(Clone, Debug, Default, Eq, PartialEq, serde::Deserialize, serde::Serialize)] #[serde(rename_all = "camelCase")] @@ -65,7 +86,14 @@ pub struct FindUpOutcome { pub matched: Option, pub visited_ancestor_count: usize, pub metadata_probe_count: usize, + /// Total number of ignored errors, including any omitted from `ignored_errors`. pub ignored_error_count: usize, + /// Ordered prefix of ignored errors, bounded by count, message size, and aggregate bytes. + #[serde(default)] + pub ignored_errors: Vec, + /// Whether one or more ignored errors could not be returned within the bounds. + #[serde(default)] + pub ignored_errors_truncated: bool, } pub(super) async fn find_up_via_metadata( @@ -77,6 +105,7 @@ pub(super) async fn find_up_via_metadata( validate_find_up_options(start, options)?; let mut outcome = FindUpOutcome::default(); + let mut ignored_error_detail_bytes = 0usize; for ancestor in start.ancestors() { outcome.visited_ancestor_count += 1; for candidate in &options.candidate_relative_paths { @@ -91,7 +120,15 @@ pub(super) async fn find_up_via_metadata( Err(err) if err.kind() == io::ErrorKind::NotFound => {} Err(err) => match options.non_not_found_error_policy { FindUpErrorPolicy::Propagate => return Err(err), - FindUpErrorPolicy::Ignore => outcome.ignored_error_count += 1, + FindUpErrorPolicy::Ignore => { + outcome.ignored_error_count += 1; + record_ignored_error( + &mut outcome, + &mut ignored_error_detail_bytes, + path, + &err, + ); + } }, } } @@ -99,6 +136,119 @@ pub(super) async fn find_up_via_metadata( Ok(outcome) } +fn record_ignored_error( + outcome: &mut FindUpOutcome, + detail_bytes: &mut usize, + path: PathUri, + err: &io::Error, +) { + if outcome.ignored_errors_truncated { + return; + } + let Some(message) = bounded_error_message(err) else { + outcome.ignored_errors_truncated = true; + return; + }; + if outcome.ignored_errors.len() == MAX_FIND_UP_IGNORED_ERRORS { + outcome.ignored_errors_truncated = true; + return; + } + let Some(fixed_item_bytes) = message + .len() + .checked_add(FIND_UP_IGNORED_ERROR_ITEM_OVERHEAD_BYTES) + else { + outcome.ignored_errors_truncated = true; + return; + }; + let Some(remaining_path_bytes) = MAX_FIND_UP_IGNORED_ERROR_DETAILS_BYTES + .checked_sub(*detail_bytes) + .and_then(|remaining| remaining.checked_sub(fixed_item_bytes)) + else { + outcome.ignored_errors_truncated = true; + return; + }; + let Some(path_bytes) = formatted_len_at_most(&path, remaining_path_bytes) else { + outcome.ignored_errors_truncated = true; + return; + }; + *detail_bytes += fixed_item_bytes + path_bytes; + outcome.ignored_errors.push(FindUpError { path, message }); +} + +fn formatted_len_at_most(value: &impl std::fmt::Display, max_bytes: usize) -> Option { + struct BoundedLength { + bytes: usize, + max_bytes: usize, + exceeded: bool, + } + + impl std::fmt::Write for BoundedLength { + fn write_str(&mut self, value: &str) -> std::fmt::Result { + let Some(bytes) = self.bytes.checked_add(value.len()) else { + self.exceeded = true; + return Ok(()); + }; + if bytes > self.max_bytes { + self.exceeded = true; + } else if !self.exceeded { + self.bytes = bytes; + } + Ok(()) + } + } + + let mut length = BoundedLength { + bytes: 0, + max_bytes, + exceeded: false, + }; + std::fmt::write(&mut length, format_args!("{value}")).ok()?; + (!length.exceeded).then_some(length.bytes) +} + +fn bounded_error_message(err: &io::Error) -> Option { + struct BoundedMessage { + value: String, + truncated: bool, + } + + impl std::fmt::Write for BoundedMessage { + fn write_str(&mut self, value: &str) -> std::fmt::Result { + if self.value.len().saturating_add(value.len()) + > MAX_FIND_UP_IGNORED_ERROR_MESSAGE_BYTES + { + self.truncated = true; + } else if !self.truncated { + self.value.push_str(value); + } + Ok(()) + } + } + + let mut message = BoundedMessage { + value: String::new(), + truncated: false, + }; + std::fmt::write(&mut message, format_args!("{err:#}")).ok()?; + (!message.truncated).then_some(message.value) +} + +pub(super) async fn find_up_batch_via_individual( + file_system: &(impl ExecutorFileSystem + ?Sized), + requests: &[FindUpRequest], + sandbox: Option<&FileSystemSandboxContext>, +) -> FileSystemResult>> { + Ok(futures::stream::iter(requests.iter().cloned()) + .map(|request| async move { + file_system + .find_up(&request.start, &request.options, sandbox) + .await + }) + .buffered(MAX_CONCURRENT_FIND_UP_REQUESTS) + .collect() + .await) +} + fn validate_find_up_options(start: &PathUri, options: &FindUpOptions) -> FileSystemResult<()> { let candidate_count = options.candidate_relative_paths.len(); if candidate_count == 0 || candidate_count > MAX_FIND_UP_CANDIDATES { diff --git a/codex-rs/file-system/src/lib.rs b/codex-rs/file-system/src/lib.rs index 5d8891474d..7207c65441 100644 --- a/codex-rs/file-system/src/lib.rs +++ b/codex-rs/file-system/src/lib.rs @@ -5,6 +5,7 @@ pub use find_up::FindUpMatch; pub use find_up::FindUpMatchKind; pub use find_up::FindUpOptions; pub use find_up::FindUpOutcome; +pub use find_up::FindUpRequest; pub use find_up::MAX_FIND_UP_CANDIDATE_BYTES; pub use find_up::MAX_FIND_UP_CANDIDATES; pub use find_up::MAX_FIND_UP_TOTAL_CANDIDATE_BYTES; @@ -341,6 +342,21 @@ pub trait ExecutorFileSystem: Send + Sync { self.find_up_via_metadata(start, options, sandbox) } + /// Runs multiple independent upward searches, preserving request order and per-search errors. + /// + /// Implementations with an optimized batch transport can override this method. The default + /// preserves compatibility by running the existing single-search operation with bounded + /// concurrency. + fn find_up_batch<'a>( + &'a self, + requests: &'a [FindUpRequest], + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec>> { + Box::pin(find_up::find_up_batch_via_individual( + self, requests, sandbox, + )) + } + /// Performs an ordered upward search using primitive metadata operations. /// /// Implementations with an optimized transport can use this as a compatibility fallback.