feat(exec-server): add batched text-prefix reads

This commit is contained in:
Adam Perry
2026-07-10 06:05:32 +00:00
parent 68b5a38353
commit cd00adbbd3
13 changed files with 542 additions and 0 deletions
@@ -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<PathUri>,
/// 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<FileSystemSandboxContext>,
}
#[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<FsReadTextPrefixesBatchResult>,
}
#[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::<FsReadTextPrefixesBatchResponse>(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!({
+10
View File
@@ -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, &params).await
}
pub async fn fs_read_text_prefixes_batch(
&self,
params: FsReadTextPrefixesBatchParams,
) -> Result<FsReadTextPrefixesBatchResponse, ExecServerError> {
self.call(FS_READ_TEXT_PREFIXES_BATCH_METHOD, &params).await
}
pub async fn fs_open(&self, params: FsOpenParams) -> Result<FsOpenResponse, ExecServerError> {
self.call(FS_OPEN_METHOD, &params).await
}
+7
View File
@@ -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;
@@ -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<FileSystemResult<TextFilePrefix>>> {
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")
@@ -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::<Result<Vec<_>, _>>()
.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")
@@ -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<Vec<FileSystemResult<TextFilePrefix>>> {
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<Vec<FileSystemResult<TextFilePrefix>>> {
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;
@@ -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,
}],
}
}
+8
View File
@@ -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,
@@ -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<FsReadTextPrefixesBatchResponse, JSONRPCErrorError> {
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;
@@ -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<FsReadTextPrefixesBatchResponse, JSONRPCErrorError> {
self.require_initialized_for("filesystem")?;
self.file_system.read_text_prefixes_batch(params).await
}
pub(crate) async fn fs_create_directory(
&self,
params: FsCreateDirectoryParams,
@@ -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<ExecServerHandler> {
handler.fs_read_file(params).await
},
);
router.request(
FS_READ_TEXT_PREFIXES_BATCH_METHOD,
|handler: Arc<ExecServerHandler>, params: FsReadTextPrefixesBatchParams| async move {
handler.fs_read_text_prefixes_batch(params).await
},
);
router.request(
FS_OPEN_METHOD,
|handler: Arc<ExecServerHandler>, params: FsOpenParams| async move {
@@ -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::<std::io::Result<Vec<_>>>()?;
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)]
+62
View File
@@ -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<FileSystemResult<TextFilePrefix>>> {
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<FileSystemResult<TextFilePrefix>>> {
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,