mirror of
https://github.com/openai/codex.git
synced 2026-11-05 20:10:33 +00:00
feat(exec-server): add batched text-prefix reads
This commit is contained in:
@@ -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!({
|
||||
|
||||
@@ -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<FsReadTextPrefixesBatchResponse, ExecServerError> {
|
||||
self.call(FS_READ_TEXT_PREFIXES_BATCH_METHOD, ¶ms).await
|
||||
}
|
||||
|
||||
pub async fn fs_open(&self, params: FsOpenParams) -> Result<FsOpenResponse, ExecServerError> {
|
||||
self.call(FS_OPEN_METHOD, ¶ms).await
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}],
|
||||
}
|
||||
}
|
||||
@@ -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)]
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user