From cd00adbbd374f7b0a68508a19135fd7fab077e5f Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Thu, 9 Jul 2026 16:15:06 +0000 Subject: [PATCH] feat(exec-server): add batched text-prefix reads --- codex-rs/exec-server-protocol/src/protocol.rs | 63 +++++++++++ codex-rs/exec-server/src/client.rs | 10 ++ codex-rs/exec-server/src/lib.rs | 7 ++ .../exec-server/src/remote_file_system.rs | 18 ++++ .../src/remote_file_system_path_uri_tests.rs | 94 ++++++++++++++++ .../exec-server/src/remote_text_prefix.rs | 100 ++++++++++++++++++ .../src/remote_text_prefix_tests.rs | 26 +++++ codex-rs/exec-server/src/rpc.rs | 8 ++ .../src/server/file_system_handler.rs | 54 ++++++++++ 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 | 82 ++++++++++++++ codex-rs/file-system/src/lib.rs | 62 +++++++++++ 13 files changed, 542 insertions(+) create mode 100644 codex-rs/exec-server/src/remote_text_prefix.rs create mode 100644 codex-rs/exec-server/src/remote_text_prefix_tests.rs diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index f5e667e03c..1e2770c6de 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -31,6 +31,7 @@ pub const EXEC_EXITED_METHOD: &str = "process/exited"; pub const EXEC_CLOSED_METHOD: &str = "process/closed"; pub const ENVIRONMENT_INFO_METHOD: &str = "environment/info"; pub const FS_READ_FILE_METHOD: &str = "fs/readFile"; +pub const FS_READ_TEXT_PREFIXES_BATCH_METHOD: &str = "fs/readTextPrefixesBatch"; pub const FS_OPEN_METHOD: &str = "fs/open"; pub const FS_READ_BLOCK_METHOD: &str = "fs/readBlock"; pub const FS_CLOSE_METHOD: &str = "fs/close"; @@ -270,6 +271,46 @@ pub struct FsReadFileResponse { pub data_base64: String, } +/// Maximum number of paths accepted by one `fs/readTextPrefixesBatch` request. +pub const FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS: usize = 64; +/// Maximum caller-selected raw prefix size for one `fs/readTextPrefixesBatch` item. +pub const FS_READ_TEXT_PREFIXES_BATCH_MAX_PREFIX_BYTES: usize = 16 * 1024; +/// Maximum raw prefix bytes returned by one full batch. Per-item base64 encoding expands this to +/// at most 1,398,272 bytes before per-item JSON framing. +pub const FS_READ_TEXT_PREFIXES_BATCH_MAX_RAW_RESPONSE_BYTES: usize = + FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS * FS_READ_TEXT_PREFIXES_BATCH_MAX_PREFIX_BYTES; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FsReadTextPrefixesBatchParams { + /// Paths are evaluated and returned in this order. + pub paths: Vec, + /// Maximum raw UTF-8 bytes returned for each path. + pub prefix_byte_limit: usize, + /// One sandbox context shared by every path in the batch. + pub sandbox: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "camelCase")] +pub enum FsReadTextPrefixesBatchResult { + Data { + #[serde(rename = "dataBase64")] + data_base64: String, + complete: bool, + }, + Error { + error: JSONRPCErrorError, + }, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FsReadTextPrefixesBatchResponse { + /// One result for each requested path, in request order. + pub results: Vec, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct FsOpenParams { @@ -650,6 +691,7 @@ mod tests { use super::FsFindUpParams; use super::FsGetMetadataBatchResponse; use super::FsReadFileParams; + use super::FsReadTextPrefixesBatchResponse; use super::HttpRequestParams; use super::InitializeResponse; use super::ProcessId; @@ -790,6 +832,27 @@ mod tests { .expect_err("native absolute sandbox cwd should not deserialize as a URI"); } + #[test] + fn read_text_prefixes_batch_mixed_results_round_trip() { + let serialized = serde_json::json!({ + "results": [ + { "type": "data", "dataBase64": "aGVsbG8=", "complete": false }, + { + "type": "error", + "error": { "code": -32004, "message": "missing" }, + }, + ], + }); + let response = + serde_json::from_value::(serialized.clone()) + .expect("deserialize text-prefix batch"); + + assert_eq!( + serde_json::to_value(response).expect("serialize text-prefix batch"), + serialized + ); + } + #[test] fn metadata_batch_error_result_round_trips() { let serialized = serde_json::json!({ diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 3afbd4e811..492c9b598a 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -61,6 +61,7 @@ use crate::protocol::FS_OPEN_METHOD; use crate::protocol::FS_READ_BLOCK_METHOD; use crate::protocol::FS_READ_DIRECTORY_METHOD; use crate::protocol::FS_READ_FILE_METHOD; +use crate::protocol::FS_READ_TEXT_PREFIXES_BATCH_METHOD; use crate::protocol::FS_REMOVE_METHOD; use crate::protocol::FS_WALK_METHOD; use crate::protocol::FS_WRITE_FILE_METHOD; @@ -88,6 +89,8 @@ use crate::protocol::FsReadDirectoryParams; use crate::protocol::FsReadDirectoryResponse; use crate::protocol::FsReadFileParams; use crate::protocol::FsReadFileResponse; +use crate::protocol::FsReadTextPrefixesBatchParams; +use crate::protocol::FsReadTextPrefixesBatchResponse; use crate::protocol::FsRemoveParams; use crate::protocol::FsRemoveResponse; use crate::protocol::FsWalkParams; @@ -676,6 +679,13 @@ impl ExecServerClient { self.call(FS_READ_FILE_METHOD, ¶ms).await } + pub async fn fs_read_text_prefixes_batch( + &self, + params: FsReadTextPrefixesBatchParams, + ) -> Result { + self.call(FS_READ_TEXT_PREFIXES_BATCH_METHOD, ¶ms).await + } + pub async fn fs_open(&self, params: FsOpenParams) -> Result { self.call(FS_OPEN_METHOD, ¶ms).await } diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 74c9e3700c..ec21704f42 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -64,6 +64,7 @@ pub use codex_file_system::MAX_FIND_UP_CANDIDATES; pub use codex_file_system::MAX_FIND_UP_TOTAL_CANDIDATE_BYTES; pub use codex_file_system::ReadDirectoryEntry; pub use codex_file_system::RemoveOptions; +pub use codex_file_system::TextFilePrefix; pub use codex_file_system::WalkEntry; pub use codex_file_system::WalkEntryKind; pub use codex_file_system::WalkError; @@ -112,6 +113,9 @@ 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::FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS; +pub use protocol::FS_READ_TEXT_PREFIXES_BATCH_MAX_PREFIX_BYTES; +pub use protocol::FS_READ_TEXT_PREFIXES_BATCH_MAX_RAW_RESPONSE_BYTES; pub use protocol::FsCanonicalizeParams; pub use protocol::FsCanonicalizeResponse; pub use protocol::FsCloseParams; @@ -139,6 +143,9 @@ pub use protocol::FsReadDirectoryParams; pub use protocol::FsReadDirectoryResponse; pub use protocol::FsReadFileParams; pub use protocol::FsReadFileResponse; +pub use protocol::FsReadTextPrefixesBatchParams; +pub use protocol::FsReadTextPrefixesBatchResponse; +pub use protocol::FsReadTextPrefixesBatchResult; pub use protocol::FsRemoveParams; pub use protocol::FsRemoveResponse; pub use protocol::FsWalkParams; diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index f8b0e0cbe1..3df0109a11 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/codex-rs/exec-server/src/remote_file_system.rs @@ -19,6 +19,7 @@ use crate::FindUpOutcome; use crate::FindUpRequest; use crate::ReadDirectoryEntry; use crate::RemoveOptions; +use crate::TextFilePrefix; use crate::WalkOptions; use crate::WalkOutcome; use crate::client::LazyRemoteExecServerClient; @@ -38,12 +39,15 @@ use crate::protocol::FsWalkParams; use crate::protocol::FsWriteFileParams; const INVALID_REQUEST_ERROR_CODE: i64 = -32600; +const INVALID_DATA_ERROR_CODE: i64 = -32005; 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; +#[path = "remote_text_prefix.rs"] +mod text_prefix; pub(crate) struct RemoteFileSystem { client: LazyRemoteExecServerClient, @@ -409,6 +413,17 @@ impl ExecutorFileSystem for RemoteFileSystem { Box::pin(RemoteFileSystem::read_file(self, path, sandbox)) } + fn read_text_prefixes_batch<'a>( + &'a self, + paths: &'a [PathUri], + max_bytes: usize, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec>> { + Box::pin(RemoteFileSystem::read_text_prefixes_batch( + self, paths, max_bytes, sandbox, + )) + } + fn read_file_stream<'a>( &'a self, path: &'a PathUri, @@ -556,6 +571,9 @@ fn map_remote_error(error: ExecServerError) -> io::Error { ExecServerError::Server { code, message } if code == INVALID_REQUEST_ERROR_CODE => { io::Error::new(io::ErrorKind::InvalidInput, message) } + ExecServerError::Server { code, message } if code == INVALID_DATA_ERROR_CODE => { + io::Error::new(io::ErrorKind::InvalidData, message) + } ExecServerError::Server { message, .. } => io::Error::other(message), ExecServerError::Closed | ExecServerError::Disconnected(_) => { io::Error::new(io::ErrorKind::BrokenPipe, "exec-server transport closed") 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 a9b1fb5273..66ab831096 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 @@ -32,11 +32,13 @@ 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::FS_READ_TEXT_PREFIXES_BATCH_METHOD; use crate::protocol::FsFindUpBatchParams; use crate::protocol::FsFindUpBatchResponse; use crate::protocol::FsFindUpParams; use crate::protocol::FsReadFileParams; use crate::protocol::FsReadFileResponse; +use crate::protocol::FsReadTextPrefixesBatchParams; use crate::protocol::INITIALIZE_METHOD; use crate::protocol::INITIALIZED_METHOD; use crate::protocol::InitializeResponse; @@ -219,6 +221,98 @@ async fn find_up_batch_method_not_found_falls_back_concurrently_and_preserves_or server.await.expect("recording server should succeed"); } +#[tokio::test] +async fn text_prefix_batch_method_not_found_falls_back_to_complete_read() { + 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 path = PathUri::parse("file:///C:/Users/Alice/skill.txt").expect("valid drive URI"); + let expected_path = path.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_READ_TEXT_PREFIXES_BATCH_METHOD => + { + request + } + other => panic!("expected fs/readTextPrefixesBatch request, got {other:?}"), + }; + let params: FsReadTextPrefixesBatchParams = serde_json::from_value( + batch_request + .params + .expect("text-prefix batch params should exist"), + ) + .expect("text-prefix batch params should deserialize"); + assert_eq!(params.paths, vec![expected_path.clone()]); + assert_eq!(params.prefix_byte_limit, 3); + 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; + + let request = match read_jsonrpc_websocket(&mut websocket).await { + JSONRPCMessage::Request(request) if request.method == FS_READ_FILE_METHOD => request, + other => panic!("expected fs/readFile request, got {other:?}"), + }; + let params: FsReadFileParams = serde_json::from_value( + request + .params + .clone() + .expect("read-file params should exist"), + ) + .expect("read-file params should deserialize"); + assert_eq!(params.path, expected_path); + write_jsonrpc_websocket( + &mut websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(FsReadFileResponse { + data_base64: "YWJjZA==".to_string(), + }) + .expect("read response should serialize"), + }), + ) + .await; + }); + let file_system = RemoteFileSystem::new(LazyRemoteExecServerClient::new( + ExecServerTransportParams::websocket_url( + websocket_url, + DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, + ), + )); + + let actual = file_system + .read_text_prefixes_batch(&[path], 3, /*sandbox*/ None) + .await + .expect("legacy fallback should succeed") + .into_iter() + .collect::, _>>() + .expect("legacy reads should succeed"); + + assert_eq!( + actual, + vec![TextFilePrefix { + text: "abc".to_string(), + complete: false, + }] + ); + 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") diff --git a/codex-rs/exec-server/src/remote_text_prefix.rs b/codex-rs/exec-server/src/remote_text_prefix.rs new file mode 100644 index 0000000000..a90d483b85 --- /dev/null +++ b/codex-rs/exec-server/src/remote_text_prefix.rs @@ -0,0 +1,100 @@ +use std::io; + +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD; +use codex_utils_path_uri::PathUri; + +use super::METHOD_NOT_FOUND_ERROR_CODE; +use super::RemoteFileSystem; +use super::map_remote_error; +use super::remote_sandbox_context; +use crate::ExecServerError; +use crate::ExecutorFileSystem; +use crate::FileSystemResult; +use crate::FileSystemSandboxContext; +use crate::TextFilePrefix; +use crate::protocol::FsReadTextPrefixesBatchParams; +use crate::protocol::FsReadTextPrefixesBatchResponse; +use crate::protocol::FsReadTextPrefixesBatchResult; + +impl RemoteFileSystem { + pub(super) async fn read_text_prefixes_batch( + &self, + paths: &[PathUri], + max_bytes: usize, + sandbox: Option<&FileSystemSandboxContext>, + ) -> FileSystemResult>> { + let client = self.client.get().await.map_err(map_remote_error)?; + let response = match client + .fs_read_text_prefixes_batch(FsReadTextPrefixesBatchParams { + paths: paths.to_vec(), + prefix_byte_limit: max_bytes, + sandbox: remote_sandbox_context(sandbox), + }) + .await + { + Ok(response) => response, + Err(ExecServerError::Server { + code: METHOD_NOT_FOUND_ERROR_CODE, + .. + }) => { + return self + .read_text_prefixes_batch_via_reads(paths, max_bytes, sandbox) + .await; + } + Err(error) => return Err(map_remote_error(error)), + }; + decode_response(response, paths.len(), max_bytes) + } +} + +fn decode_response( + response: FsReadTextPrefixesBatchResponse, + expected_results: usize, + max_bytes: usize, +) -> FileSystemResult>> { + if response.results.len() != expected_results { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "remote fs/readTextPrefixesBatch returned {} results for {expected_results} paths", + response.results.len() + ), + )); + } + response + .results + .into_iter() + .map(|result| match result { + FsReadTextPrefixesBatchResult::Data { + data_base64, + complete, + } => { + let bytes = STANDARD.decode(data_base64).map_err(|error| { + io::Error::new(io::ErrorKind::InvalidData, format!( + "remote fs/readTextPrefixesBatch returned invalid base64 dataBase64: {error}" + )) + })?; + if bytes.len() > max_bytes { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "remote fs/readTextPrefixesBatch returned an oversized prefix", + )); + } + String::from_utf8(bytes) + .map(|text| Ok(TextFilePrefix { text, complete })) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) + } + FsReadTextPrefixesBatchResult::Error { error } => { + Ok(Err(map_remote_error(ExecServerError::Server { + code: error.code, + message: error.message, + }))) + } + }) + .collect() +} + +#[cfg(test)] +#[path = "remote_text_prefix_tests.rs"] +mod tests; diff --git a/codex-rs/exec-server/src/remote_text_prefix_tests.rs b/codex-rs/exec-server/src/remote_text_prefix_tests.rs new file mode 100644 index 0000000000..6e5715cf2b --- /dev/null +++ b/codex-rs/exec-server/src/remote_text_prefix_tests.rs @@ -0,0 +1,26 @@ +use pretty_assertions::assert_eq; + +use super::*; + +#[test] +fn rejects_malformed_cardinality_base64_and_utf8() { + for response in [ + FsReadTextPrefixesBatchResponse { + results: Vec::new(), + }, + data_response("not-base64"), + data_response("/w=="), + ] { + let error = decode_response(response, 1, 4).expect_err("malformed response"); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + } +} + +fn data_response(data_base64: &str) -> FsReadTextPrefixesBatchResponse { + FsReadTextPrefixesBatchResponse { + results: vec![FsReadTextPrefixesBatchResult::Data { + data_base64: data_base64.to_string(), + complete: false, + }], + } +} diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 9671f25d39..b5e79662a8 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -580,6 +580,14 @@ pub(crate) fn not_found(message: String) -> JSONRPCErrorError { } } +pub(crate) fn invalid_data(message: String) -> JSONRPCErrorError { + JSONRPCErrorError { + code: -32005, + data: None, + message, + } +} + pub(crate) fn internal_error(message: String) -> JSONRPCErrorError { JSONRPCErrorError { code: -32603, 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 9523ac5915..e870da1081 100644 --- a/codex-rs/exec-server/src/server/file_system_handler.rs +++ b/codex-rs/exec-server/src/server/file_system_handler.rs @@ -9,10 +9,13 @@ use crate::CreateDirectoryOptions; use crate::ExecServerRuntimePaths; use crate::ExecutorFileSystem; use crate::RemoveOptions; +use crate::TextFilePrefix; 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_READ_TEXT_PREFIXES_BATCH_MAX_PATHS; +use crate::protocol::FS_READ_TEXT_PREFIXES_BATCH_MAX_PREFIX_BYTES; use crate::protocol::FS_WRITE_FILE_METHOD; use crate::protocol::FsCanonicalizeParams; use crate::protocol::FsCanonicalizeResponse; @@ -41,6 +44,9 @@ use crate::protocol::FsReadDirectoryParams; use crate::protocol::FsReadDirectoryResponse; use crate::protocol::FsReadFileParams; use crate::protocol::FsReadFileResponse; +use crate::protocol::FsReadTextPrefixesBatchParams; +use crate::protocol::FsReadTextPrefixesBatchResponse; +use crate::protocol::FsReadTextPrefixesBatchResult; use crate::protocol::FsRemoveParams; use crate::protocol::FsRemoveResponse; use crate::protocol::FsWalkParams; @@ -48,6 +54,7 @@ use crate::protocol::FsWalkResponse; use crate::protocol::FsWriteFileParams; use crate::protocol::FsWriteFileResponse; use crate::rpc::internal_error; +use crate::rpc::invalid_data; use crate::rpc::invalid_request; use crate::rpc::not_found; @@ -128,6 +135,45 @@ impl FileSystemHandler { }) } + pub(crate) async fn read_text_prefixes_batch( + &self, + params: FsReadTextPrefixesBatchParams, + ) -> Result { + if params.paths.len() > FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS { + return Err(invalid_request(format!( + "text-prefix batch must not exceed {FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS} paths" + ))); + } + if params.prefix_byte_limit > FS_READ_TEXT_PREFIXES_BATCH_MAX_PREFIX_BYTES { + return Err(invalid_request(format!( + "text-prefix byte limit must not exceed {FS_READ_TEXT_PREFIXES_BATCH_MAX_PREFIX_BYTES}" + ))); + } + let mut results = Vec::with_capacity(params.paths.len()); + for path in params.paths { + let result = self + .file_system + .read_file(&path, params.sandbox.as_ref()) + .await + .and_then(|bytes| { + String::from_utf8(bytes) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)) + }) + .map(|text| TextFilePrefix::from_complete_text(text, params.prefix_byte_limit)) + .map_err(map_text_read_error); + match result { + Ok(prefix) => { + results.push(FsReadTextPrefixesBatchResult::Data { + data_base64: STANDARD.encode(prefix.text), + complete: prefix.complete, + }); + } + Err(error) => results.push(FsReadTextPrefixesBatchResult::Error { error }), + } + } + Ok(FsReadTextPrefixesBatchResponse { results }) + } + pub(crate) async fn write_file( &self, params: FsWriteFileParams, @@ -341,6 +387,14 @@ fn map_fs_error(err: io::Error) -> JSONRPCErrorError { } } +fn map_text_read_error(err: io::Error) -> JSONRPCErrorError { + if err.kind() == io::ErrorKind::InvalidData { + invalid_data(err.to_string()) + } else { + map_fs_error(err) + } +} + #[cfg(test)] mod tests { use codex_protocol::protocol::NetworkAccess; diff --git a/codex-rs/exec-server/src/server/handler.rs b/codex-rs/exec-server/src/server/handler.rs index df10453641..a8e56490cf 100644 --- a/codex-rs/exec-server/src/server/handler.rs +++ b/codex-rs/exec-server/src/server/handler.rs @@ -41,6 +41,8 @@ use crate::protocol::FsReadDirectoryParams; use crate::protocol::FsReadDirectoryResponse; use crate::protocol::FsReadFileParams; use crate::protocol::FsReadFileResponse; +use crate::protocol::FsReadTextPrefixesBatchParams; +use crate::protocol::FsReadTextPrefixesBatchResponse; use crate::protocol::FsRemoveParams; use crate::protocol::FsRemoveResponse; use crate::protocol::FsWalkParams; @@ -290,6 +292,14 @@ impl ExecServerHandler { self.file_system.write_file(params).await } + pub(crate) async fn fs_read_text_prefixes_batch( + &self, + params: FsReadTextPrefixesBatchParams, + ) -> Result { + self.require_initialized_for("filesystem")?; + self.file_system.read_text_prefixes_batch(params).await + } + pub(crate) async fn fs_create_directory( &self, params: FsCreateDirectoryParams, diff --git a/codex-rs/exec-server/src/server/registry.rs b/codex-rs/exec-server/src/server/registry.rs index c4bc9eea61..8db0b6eed6 100644 --- a/codex-rs/exec-server/src/server/registry.rs +++ b/codex-rs/exec-server/src/server/registry.rs @@ -19,6 +19,7 @@ use crate::protocol::FS_OPEN_METHOD; use crate::protocol::FS_READ_BLOCK_METHOD; use crate::protocol::FS_READ_DIRECTORY_METHOD; use crate::protocol::FS_READ_FILE_METHOD; +use crate::protocol::FS_READ_TEXT_PREFIXES_BATCH_METHOD; use crate::protocol::FS_REMOVE_METHOD; use crate::protocol::FS_WALK_METHOD; use crate::protocol::FS_WRITE_FILE_METHOD; @@ -34,6 +35,7 @@ use crate::protocol::FsOpenParams; use crate::protocol::FsReadBlockParams; use crate::protocol::FsReadDirectoryParams; use crate::protocol::FsReadFileParams; +use crate::protocol::FsReadTextPrefixesBatchParams; use crate::protocol::FsRemoveParams; use crate::protocol::FsWalkParams; use crate::protocol::FsWriteFileParams; @@ -107,6 +109,12 @@ pub(crate) fn build_router() -> RpcRouter { handler.fs_read_file(params).await }, ); + router.request( + FS_READ_TEXT_PREFIXES_BATCH_METHOD, + |handler: Arc, params: FsReadTextPrefixesBatchParams| async move { + handler.fs_read_text_prefixes_batch(params).await + }, + ); router.request( FS_OPEN_METHOD, |handler: Arc, params: FsOpenParams| 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 3fb1a2f828..7839731b73 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::FS_FIND_UP_BATCH_MAX_REQUESTS; +use codex_exec_server::FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS; +use codex_exec_server::FS_READ_TEXT_PREFIXES_BATCH_MAX_PREFIX_BYTES; use codex_exec_server::FileMetadata; use codex_exec_server::FindUpErrorPolicy; use codex_exec_server::FindUpMatch; @@ -16,6 +18,7 @@ use codex_exec_server::MAX_FIND_UP_CANDIDATES; use codex_exec_server::MAX_FIND_UP_TOTAL_CANDIDATE_BYTES; use codex_exec_server::ReadDirectoryEntry; use codex_exec_server::RemoveOptions; +use codex_exec_server::TextFilePrefix; use codex_exec_server::WalkEntry; use codex_exec_server::WalkEntryKind; use codex_exec_server::WalkOptions; @@ -549,6 +552,85 @@ async fn file_system_read_file_text_returns_string( Ok(()) } +#[test_case(FileSystemImplementation::Local ; "local")] +#[test_case(FileSystemImplementation::Remote ; "remote")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn file_system_text_prefix_batch_preserves_text_semantics_and_order( + implementation: FileSystemImplementation, +) -> Result<()> { + let context = create_file_system_context(implementation).await?; + let file_system = context.file_system; + let tmp = TempDir::new()?; + let multibyte = tmp.path().join("multibyte.txt"); + let missing = tmp.path().join("missing.txt"); + let exact = tmp.path().join("exact.txt"); + let invalid_tail = tmp.path().join("invalid-tail.txt"); + std::fs::write(&multibyte, "abéz")?; + std::fs::write(&exact, "hey")?; + std::fs::write(&invalid_tail, b"abc\xff")?; + let paths = [&multibyte, &missing, &exact, &invalid_tail] + .into_iter() + .map(PathUri::from_host_native_path) + .collect::>>()?; + + let results = file_system + .read_text_prefixes_batch(&paths, 3, /*sandbox*/ None) + .await + .with_context(|| format!("mode={implementation}"))?; + + assert_eq!( + results[0].as_ref().expect("multibyte prefix"), + &TextFilePrefix { + text: "ab".to_string(), + complete: false, + } + ); + assert_eq!( + results[1].as_ref().expect_err("missing file").kind(), + std::io::ErrorKind::NotFound + ); + assert_eq!( + results[2].as_ref().expect("exact prefix"), + &TextFilePrefix { + text: "hey".to_string(), + complete: true, + } + ); + assert_eq!( + results[3] + .as_ref() + .expect_err("invalid omitted tail") + .kind(), + std::io::ErrorKind::InvalidData + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_text_prefix_batch_enforces_request_caps() -> Result<()> { + let context = create_file_system_context(FileSystemImplementation::Remote).await?; + let file_system = context.file_system; + let path = PathUri::from_host_native_path(std::env::current_dir()?)?; + let paths = vec![path.clone(); FS_READ_TEXT_PREFIXES_BATCH_MAX_PATHS + 1]; + + let path_error = file_system + .read_text_prefixes_batch(&paths, 1, /*sandbox*/ None) + .await + .expect_err("path cap should be enforced"); + let prefix_error = file_system + .read_text_prefixes_batch( + &[path], + FS_READ_TEXT_PREFIXES_BATCH_MAX_PREFIX_BYTES + 1, + /*sandbox*/ None, + ) + .await + .expect_err("prefix cap should be enforced"); + + assert_eq!(path_error.kind(), std::io::ErrorKind::InvalidInput); + assert_eq!(prefix_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/lib.rs b/codex-rs/file-system/src/lib.rs index 7207c65441..27e31cd4f5 100644 --- a/codex-rs/file-system/src/lib.rs +++ b/codex-rs/file-system/src/lib.rs @@ -26,6 +26,7 @@ use codex_protocol::protocol::SandboxPolicy; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; use futures::Stream; +use futures::StreamExt; use std::collections::HashSet; use std::collections::VecDeque; use std::future::Future; @@ -37,6 +38,7 @@ use std::task::Poll; /// Maximum chunk size returned by [`ExecutorFileSystem::read_file_stream`]. pub const FILE_READ_CHUNK_SIZE: usize = 1024 * 1024; +const DEFAULT_BATCH_TEXT_PREFIX_READ_CONCURRENCY: usize = 64; const MAX_WALK_DEPTH: usize = 64; const MAX_WALK_DIRECTORIES: usize = 10_000; const MAX_WALK_ENTRIES: usize = 50_000; @@ -70,6 +72,29 @@ pub struct FileMetadata { pub modified_at_ms: i64, } +/// A UTF-8 file prefix returned by a bounded text read. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct TextFilePrefix { + pub text: String, + /// Whether `text` contains the complete file rather than a bounded prefix. + pub complete: bool, +} + +impl TextFilePrefix { + /// Produces the longest prefix no larger than `max_bytes` that ends on a UTF-8 boundary. + pub fn from_complete_text(mut text: String, max_bytes: usize) -> Self { + let complete = text.len() <= max_bytes; + if !complete { + let mut end = max_bytes; + while !text.is_char_boundary(end) { + end -= 1; + } + text.truncate(end); + } + Self { text, complete } + } +} + #[derive(Clone, Debug, Eq, PartialEq)] pub struct ReadDirectoryEntry { pub file_name: String, @@ -294,6 +319,43 @@ pub trait ExecutorFileSystem: Send + Sync { }) } + /// Reads bounded UTF-8 prefixes for multiple files, preserving path order and per-file errors. + /// + /// The complete file is decoded before its prefix is selected, so invalid UTF-8 outside the + /// returned prefix remains an error. Implementations with an optimized batch transport can + /// override this method. + fn read_text_prefixes_batch<'a>( + &'a self, + paths: &'a [PathUri], + max_bytes: usize, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec>> { + self.read_text_prefixes_batch_via_reads(paths, max_bytes, sandbox) + } + + /// Performs bounded text-prefix reads using the primitive complete-text operation. + /// + /// Optimized transports can use this as a compatibility fallback. + fn read_text_prefixes_batch_via_reads<'a>( + &'a self, + paths: &'a [PathUri], + max_bytes: usize, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec>> { + Box::pin(async move { + let results = futures::stream::iter(paths.iter().cloned()) + .map(|path| async move { + self.read_file_text(&path, sandbox) + .await + .map(|text| TextFilePrefix::from_complete_text(text, max_bytes)) + }) + .buffered(DEFAULT_BATCH_TEXT_PREFIX_READ_CONCURRENCY) + .collect() + .await; + Ok(results) + }) + } + fn write_file<'a>( &'a self, path: &'a PathUri,