mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
[codex] exec-server: stream files in chunks (#28354)
## Why `fs/readFile` buffers the entire file in one response, which makes large remote reads expensive and prevents callers from applying backpressure. We need an opt-in streaming path with bounded block sizes while preserving the existing single-call API for small and sandboxed reads. ## What changed - Add `ExecServerClient::stream`, returning a named `FileReadStream` that implements `futures::Stream` and yields immutable 1 MiB byte blocks. - Add internal `fs/open`, `fs/readBlock`, and `fs/close` RPCs. `fs/readBlock` accepts an explicit offset and length. - Keep unsandboxed files open between block reads, cap open handles per connection, and clean them up on EOF, error, stream drop, explicit close, or connection shutdown. - Reject platform-sandboxed streaming opens instead of turning the one-shot sandbox helper into a persistent server. Existing `fs/readFile` behavior is unchanged. ## Testing - `just test -p codex-exec-server` - Integration coverage for 1 MiB chunking, exact block-boundary EOF, sandbox rejection, and continued reads from the opened file after path replacement. - Handle-manager coverage for non-sequential offsets, variable block lengths, the 128-handle limit, and capacity release after close.
This commit is contained in:
committed by
GitHub
Unverified
parent
76135cbe7e
commit
a4711b88dd
Generated
+4
@@ -2891,6 +2891,7 @@ dependencies = [
|
||||
"ctor 0.6.3",
|
||||
"futures",
|
||||
"http 1.4.0",
|
||||
"libc",
|
||||
"pretty_assertions",
|
||||
"prost 0.14.3",
|
||||
"reqwest 0.12.28",
|
||||
@@ -2906,6 +2907,7 @@ dependencies = [
|
||||
"toml 0.9.11+spec-1.1.0",
|
||||
"tracing",
|
||||
"uuid",
|
||||
"windows-sys 0.52.0",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
@@ -3040,9 +3042,11 @@ dependencies = [
|
||||
name = "codex-file-system"
|
||||
version = "0.0.0"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"codex-protocol",
|
||||
"codex-utils-absolute-path",
|
||||
"codex-utils-path-uri",
|
||||
"futures",
|
||||
"serde",
|
||||
]
|
||||
|
||||
|
||||
@@ -886,24 +886,14 @@ async fn external_agent_config_import_returns_before_background_session_import_f
|
||||
let session_path = session_dir.join("session.jsonl");
|
||||
std::fs::create_dir_all(&project_root)?;
|
||||
std::fs::create_dir_all(&session_dir)?;
|
||||
std::fs::write(
|
||||
&session_path,
|
||||
serde_json::json!({
|
||||
"type": "user",
|
||||
"cwd": &project_root,
|
||||
"timestamp": &recent_timestamp,
|
||||
"message": { "content": "first request" },
|
||||
})
|
||||
.to_string(),
|
||||
)?;
|
||||
|
||||
let project_config_dir = project_root.join(".codex");
|
||||
std::fs::create_dir_all(&project_config_dir)?;
|
||||
let project_config = project_config_dir.join("config.toml");
|
||||
let status = std::process::Command::new("mkfifo")
|
||||
.arg(&project_config)
|
||||
.status()?;
|
||||
assert!(status.success());
|
||||
let session_contents = serde_json::json!({
|
||||
"type": "user",
|
||||
"cwd": &project_root,
|
||||
"timestamp": &recent_timestamp,
|
||||
"message": { "content": "first request" },
|
||||
})
|
||||
.to_string();
|
||||
std::fs::write(&session_path, &session_contents)?;
|
||||
|
||||
let home_dir = codex_home.path().display().to_string();
|
||||
let mut mcp =
|
||||
@@ -926,6 +916,12 @@ async fn external_agent_config_import_returns_before_background_session_import_f
|
||||
assert_eq!(detected.items.len(), 1);
|
||||
let detected_items = detected.items;
|
||||
|
||||
std::fs::remove_file(&session_path)?;
|
||||
let status = std::process::Command::new("mkfifo")
|
||||
.arg(&session_path)
|
||||
.status()?;
|
||||
assert!(status.success());
|
||||
|
||||
let request_id = mcp
|
||||
.send_raw_request(
|
||||
"externalAgentConfig/import",
|
||||
@@ -964,17 +960,17 @@ async fn external_agent_config_import_returns_before_background_session_import_f
|
||||
let response: ExternalAgentConfigImportResponse = to_response(response)?;
|
||||
let duplicate_import_id = assert_import_response(response);
|
||||
|
||||
let writer = tokio::spawn(async move {
|
||||
let mut file = tokio::fs::OpenOptions::new()
|
||||
.write(true)
|
||||
.open(&project_config)
|
||||
.await?;
|
||||
file.write_all(b"\n").await
|
||||
});
|
||||
timeout(DEFAULT_TIMEOUT, writer).await???;
|
||||
|
||||
let mut completed_import_ids = Vec::new();
|
||||
for _ in 0..2 {
|
||||
timeout(DEFAULT_TIMEOUT, async {
|
||||
let mut file = tokio::fs::OpenOptions::new()
|
||||
.write(true)
|
||||
.open(&session_path)
|
||||
.await?;
|
||||
file.write_all(session_contents.as_bytes()).await
|
||||
})
|
||||
.await??;
|
||||
|
||||
let notification = timeout(
|
||||
DEFAULT_TIMEOUT,
|
||||
mcp.read_stream_until_notification_message("externalAgentConfig/import/completed"),
|
||||
|
||||
@@ -3,6 +3,7 @@ use codex_file_system::CopyOptions;
|
||||
use codex_file_system::CreateDirectoryOptions;
|
||||
use codex_file_system::ExecutorFileSystemFuture;
|
||||
use codex_file_system::FileMetadata;
|
||||
use codex_file_system::FileSystemReadStream;
|
||||
use codex_file_system::FileSystemSandboxContext;
|
||||
use codex_file_system::ReadDirectoryEntry;
|
||||
use codex_file_system::RemoveOptions;
|
||||
@@ -36,6 +37,19 @@ impl ExecutorFileSystem for TestFileSystem {
|
||||
})
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
_sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(async {
|
||||
Err(std::io::Error::new(
|
||||
std::io::ErrorKind::Unsupported,
|
||||
"test filesystem does not support streaming reads",
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
|
||||
@@ -8,6 +8,7 @@ use codex_exec_server::EnvironmentManager;
|
||||
use codex_exec_server::ExecutorFileSystem;
|
||||
use codex_exec_server::ExecutorFileSystemFuture;
|
||||
use codex_exec_server::FileMetadata;
|
||||
use codex_exec_server::FileSystemReadStream;
|
||||
use codex_exec_server::FileSystemResult;
|
||||
use codex_exec_server::FileSystemSandboxContext;
|
||||
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
|
||||
@@ -89,6 +90,14 @@ impl ExecutorFileSystem for SyntheticPluginFileSystem {
|
||||
})
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
_sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(async { Self::unsupported() })
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
|
||||
@@ -11,6 +11,7 @@ use codex_exec_server::CreateDirectoryOptions;
|
||||
use codex_exec_server::Environment;
|
||||
use codex_exec_server::ExecutorFileSystemFuture;
|
||||
use codex_exec_server::FileMetadata;
|
||||
use codex_exec_server::FileSystemReadStream;
|
||||
use codex_exec_server::FileSystemSandboxContext;
|
||||
use codex_exec_server::LOCAL_FS;
|
||||
use codex_exec_server::ReadDirectoryEntry;
|
||||
@@ -140,6 +141,19 @@ impl ExecutorFileSystem for FailingFileSystem {
|
||||
Box::pin(FailingFileSystem::read_file(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
_sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(async {
|
||||
Err(io::Error::new(
|
||||
io::ErrorKind::Unsupported,
|
||||
"failing filesystem does not support streaming reads",
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
|
||||
@@ -45,11 +45,20 @@ tokio = { workspace = true, features = [
|
||||
"sync",
|
||||
"time",
|
||||
] }
|
||||
tokio-util = { workspace = true, features = ["rt"] }
|
||||
tokio-util = { workspace = true, features = ["io", "rt"] }
|
||||
tokio-tungstenite = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
uuid = { workspace = true, features = ["v4"] }
|
||||
|
||||
[target.'cfg(unix)'.dependencies]
|
||||
libc = { workspace = true }
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
windows-sys = { version = "0.52", features = [
|
||||
"Win32_Foundation",
|
||||
"Win32_Storage_FileSystem",
|
||||
] }
|
||||
|
||||
[dev-dependencies]
|
||||
anyhow = { workspace = true }
|
||||
codex-test-binary-support = { workspace = true }
|
||||
|
||||
@@ -345,6 +345,8 @@ invalid or unavailable paths. For compatibility, requests also accept native
|
||||
absolute path strings and normalize them to `file:` URIs:
|
||||
|
||||
- `fs/readFile`
|
||||
- `fs/open`, `fs/readBlock`, and `fs/close` (internal transport for
|
||||
`ExecutorFileSystem::read_file_stream`)
|
||||
- `fs/writeFile`
|
||||
- `fs/createDirectory`
|
||||
- `fs/getMetadata`
|
||||
|
||||
@@ -45,21 +45,30 @@ use crate::protocol::ExecOutputDeltaNotification;
|
||||
use crate::protocol::ExecParams;
|
||||
use crate::protocol::ExecResponse;
|
||||
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_GET_METADATA_METHOD;
|
||||
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_REMOVE_METHOD;
|
||||
use crate::protocol::FS_WRITE_FILE_METHOD;
|
||||
use crate::protocol::FsCanonicalizeParams;
|
||||
use crate::protocol::FsCanonicalizeResponse;
|
||||
use crate::protocol::FsCloseParams;
|
||||
use crate::protocol::FsCloseResponse;
|
||||
use crate::protocol::FsCopyParams;
|
||||
use crate::protocol::FsCopyResponse;
|
||||
use crate::protocol::FsCreateDirectoryParams;
|
||||
use crate::protocol::FsCreateDirectoryResponse;
|
||||
use crate::protocol::FsGetMetadataParams;
|
||||
use crate::protocol::FsGetMetadataResponse;
|
||||
use crate::protocol::FsOpenParams;
|
||||
use crate::protocol::FsOpenResponse;
|
||||
use crate::protocol::FsReadBlockParams;
|
||||
use crate::protocol::FsReadBlockResponse;
|
||||
use crate::protocol::FsReadDirectoryParams;
|
||||
use crate::protocol::FsReadDirectoryResponse;
|
||||
use crate::protocol::FsReadFileParams;
|
||||
@@ -430,6 +439,24 @@ impl ExecServerClient {
|
||||
self.call(FS_READ_FILE_METHOD, ¶ms).await
|
||||
}
|
||||
|
||||
pub async fn fs_open(&self, params: FsOpenParams) -> Result<FsOpenResponse, ExecServerError> {
|
||||
self.call(FS_OPEN_METHOD, ¶ms).await
|
||||
}
|
||||
|
||||
pub async fn fs_read_block(
|
||||
&self,
|
||||
params: FsReadBlockParams,
|
||||
) -> Result<FsReadBlockResponse, ExecServerError> {
|
||||
self.call(FS_READ_BLOCK_METHOD, ¶ms).await
|
||||
}
|
||||
|
||||
pub async fn fs_close(
|
||||
&self,
|
||||
params: FsCloseParams,
|
||||
) -> Result<FsCloseResponse, ExecServerError> {
|
||||
self.call(FS_CLOSE_METHOD, ¶ms).await
|
||||
}
|
||||
|
||||
pub async fn fs_write_file(
|
||||
&self,
|
||||
params: FsWriteFileParams,
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
use std::collections::HashMap;
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::sync::Arc;
|
||||
|
||||
use codex_file_system::FILE_READ_CHUNK_SIZE;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
const MAX_OPEN_FILE_READS: usize = 128;
|
||||
|
||||
#[derive(Debug, Eq, PartialEq)]
|
||||
pub(crate) struct FileReadBlock {
|
||||
pub(crate) bytes: Vec<u8>,
|
||||
pub(crate) eof: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub(crate) struct FileReadHandleManager {
|
||||
handles: Arc<Mutex<HashMap<String, Arc<File>>>>,
|
||||
}
|
||||
|
||||
impl FileReadHandleManager {
|
||||
pub(crate) async fn open(
|
||||
&self,
|
||||
handle_id: String,
|
||||
file: tokio::fs::File,
|
||||
) -> io::Result<String> {
|
||||
let file = Arc::new(file.into_std().await);
|
||||
let mut handles = self.handles.lock().await;
|
||||
if handles.contains_key(&handle_id) {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("file read handle `{handle_id}` already exists"),
|
||||
));
|
||||
}
|
||||
if handles.len() >= MAX_OPEN_FILE_READS {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("at most {MAX_OPEN_FILE_READS} file reads may be open per connection"),
|
||||
));
|
||||
}
|
||||
handles.insert(handle_id.clone(), file);
|
||||
Ok(handle_id)
|
||||
}
|
||||
|
||||
pub(crate) async fn read_block(
|
||||
&self,
|
||||
handle_id: &str,
|
||||
offset: u64,
|
||||
len: usize,
|
||||
) -> io::Result<FileReadBlock> {
|
||||
validate_read_block_len(len)?;
|
||||
let file = {
|
||||
let handles = self.handles.lock().await;
|
||||
handles
|
||||
.get(handle_id)
|
||||
.cloned()
|
||||
.ok_or_else(|| unknown_handle_error(handle_id))?
|
||||
};
|
||||
let result =
|
||||
match tokio::task::spawn_blocking(move || read_block_at(&file, offset, len)).await {
|
||||
Ok(result) => result,
|
||||
Err(error) => Err(io::Error::other(format!(
|
||||
"file read task stopped unexpectedly: {error}"
|
||||
))),
|
||||
};
|
||||
if result.is_err() {
|
||||
self.close(handle_id).await;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
pub(crate) async fn close(&self, handle_id: &str) {
|
||||
self.handles.lock().await.remove(handle_id);
|
||||
}
|
||||
|
||||
pub(crate) async fn close_all(&self) {
|
||||
self.handles.lock().await.clear();
|
||||
}
|
||||
}
|
||||
|
||||
fn read_block_at(file: &File, offset: u64, len: usize) -> io::Result<FileReadBlock> {
|
||||
let mut bytes = vec![0; len];
|
||||
let mut bytes_read = 0;
|
||||
while bytes_read < len {
|
||||
let read_offset = offset.checked_add(bytes_read as u64).ok_or_else(|| {
|
||||
io::Error::new(io::ErrorKind::InvalidInput, "file read offset overflowed")
|
||||
})?;
|
||||
match read_file_at(file, &mut bytes[bytes_read..], read_offset) {
|
||||
Ok(0) => break,
|
||||
Ok(read) => bytes_read += read,
|
||||
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
bytes.truncate(bytes_read);
|
||||
Ok(FileReadBlock {
|
||||
eof: bytes_read < len,
|
||||
bytes,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn read_file_at(file: &File, bytes: &mut [u8], offset: u64) -> io::Result<usize> {
|
||||
std::os::unix::fs::FileExt::read_at(file, bytes, offset)
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn read_file_at(file: &File, bytes: &mut [u8], offset: u64) -> io::Result<usize> {
|
||||
std::os::windows::fs::FileExt::seek_read(file, bytes, offset)
|
||||
}
|
||||
|
||||
fn validate_read_block_len(len: usize) -> io::Result<()> {
|
||||
if !(1..=FILE_READ_CHUNK_SIZE).contains(&len) {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("file read block length must be between 1 and {FILE_READ_CHUNK_SIZE}"),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn unknown_handle_error(handle_id: &str) -> io::Error {
|
||||
io::Error::new(
|
||||
io::ErrorKind::NotFound,
|
||||
format!("unknown file read handle `{handle_id}`"),
|
||||
)
|
||||
}
|
||||
@@ -5,6 +5,7 @@ mod connection;
|
||||
mod environment;
|
||||
mod environment_provider;
|
||||
mod environment_toml;
|
||||
mod file_read;
|
||||
mod fs_helper;
|
||||
mod fs_helper_main;
|
||||
mod fs_sandbox;
|
||||
@@ -15,6 +16,7 @@ mod noise_relay;
|
||||
mod process;
|
||||
mod process_id;
|
||||
mod protocol;
|
||||
mod regular_file;
|
||||
mod relay;
|
||||
mod relay_proto;
|
||||
mod remote;
|
||||
@@ -39,7 +41,9 @@ pub use codex_file_system::CopyOptions;
|
||||
pub use codex_file_system::CreateDirectoryOptions;
|
||||
pub use codex_file_system::ExecutorFileSystem;
|
||||
pub use codex_file_system::ExecutorFileSystemFuture;
|
||||
pub use codex_file_system::FILE_READ_CHUNK_SIZE;
|
||||
pub use codex_file_system::FileMetadata;
|
||||
pub use codex_file_system::FileSystemReadStream;
|
||||
pub use codex_file_system::FileSystemResult;
|
||||
pub use codex_file_system::FileSystemSandboxContext;
|
||||
pub use codex_file_system::ReadDirectoryEntry;
|
||||
@@ -67,6 +71,7 @@ pub use process::ExecProcessEventReceiver;
|
||||
pub use process::ExecProcessFuture;
|
||||
pub use process::StartedExecProcess;
|
||||
pub use process_id::ProcessId;
|
||||
pub use protocol::ByteChunk;
|
||||
pub use protocol::EnvironmentInfo;
|
||||
pub use protocol::ExecClosedNotification;
|
||||
pub use protocol::ExecEnvPolicy;
|
||||
@@ -77,12 +82,18 @@ pub use protocol::ExecParams;
|
||||
pub use protocol::ExecResponse;
|
||||
pub use protocol::FsCanonicalizeParams;
|
||||
pub use protocol::FsCanonicalizeResponse;
|
||||
pub use protocol::FsCloseParams;
|
||||
pub use protocol::FsCloseResponse;
|
||||
pub use protocol::FsCopyParams;
|
||||
pub use protocol::FsCopyResponse;
|
||||
pub use protocol::FsCreateDirectoryParams;
|
||||
pub use protocol::FsCreateDirectoryResponse;
|
||||
pub use protocol::FsGetMetadataParams;
|
||||
pub use protocol::FsGetMetadataResponse;
|
||||
pub use protocol::FsOpenParams;
|
||||
pub use protocol::FsOpenResponse;
|
||||
pub use protocol::FsReadBlockParams;
|
||||
pub use protocol::FsReadBlockResponse;
|
||||
pub use protocol::FsReadDirectoryEntry;
|
||||
pub use protocol::FsReadDirectoryParams;
|
||||
pub use protocol::FsReadDirectoryResponse;
|
||||
|
||||
@@ -7,21 +7,33 @@ use std::sync::LazyLock;
|
||||
use std::time::SystemTime;
|
||||
use std::time::UNIX_EPOCH;
|
||||
use tokio::io;
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio_util::io::ReaderStream;
|
||||
|
||||
use crate::CopyOptions;
|
||||
use crate::CreateDirectoryOptions;
|
||||
use crate::ExecServerRuntimePaths;
|
||||
use crate::ExecutorFileSystem;
|
||||
use crate::ExecutorFileSystemFuture;
|
||||
use crate::FILE_READ_CHUNK_SIZE;
|
||||
use crate::FileMetadata;
|
||||
use crate::FileSystemReadStream;
|
||||
use crate::FileSystemResult;
|
||||
use crate::FileSystemSandboxContext;
|
||||
use crate::ReadDirectoryEntry;
|
||||
use crate::RemoveOptions;
|
||||
use crate::regular_file;
|
||||
use crate::sandboxed_file_system::SandboxedFileSystem;
|
||||
|
||||
const MAX_READ_FILE_BYTES: u64 = 512 * 1024 * 1024;
|
||||
|
||||
fn file_too_large_error() -> io::Error {
|
||||
io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("file is too large to read: limit is {MAX_READ_FILE_BYTES} bytes"),
|
||||
)
|
||||
}
|
||||
|
||||
pub static LOCAL_FS: LazyLock<Arc<dyn ExecutorFileSystem>> =
|
||||
LazyLock::new(|| -> Arc<dyn ExecutorFileSystem> { Arc::new(LocalFileSystem::unsandboxed()) });
|
||||
|
||||
@@ -79,6 +91,20 @@ impl LocalFileSystem {
|
||||
}
|
||||
|
||||
impl LocalFileSystem {
|
||||
pub(crate) async fn open_file_for_read(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
sandbox: Option<&FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<tokio::fs::File> {
|
||||
if sandbox.is_some_and(FileSystemSandboxContext::should_run_in_sandbox) {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
"streaming file reads do not support platform sandboxing",
|
||||
));
|
||||
}
|
||||
self.unsandboxed.open_file_for_read(path, sandbox).await
|
||||
}
|
||||
|
||||
async fn canonicalize(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
@@ -97,6 +123,15 @@ impl LocalFileSystem {
|
||||
file_system.read_file(path, sandbox).await
|
||||
}
|
||||
|
||||
async fn read_file_stream(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
sandbox: Option<&FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<FileSystemReadStream> {
|
||||
let (file_system, sandbox) = self.file_system_for(sandbox)?;
|
||||
file_system.read_file_stream(path, sandbox).await
|
||||
}
|
||||
|
||||
async fn write_file(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
@@ -176,6 +211,14 @@ impl ExecutorFileSystem for LocalFileSystem {
|
||||
Box::pin(LocalFileSystem::read_file(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(LocalFileSystem::read_file_stream(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
@@ -239,6 +282,17 @@ impl ExecutorFileSystem for LocalFileSystem {
|
||||
}
|
||||
|
||||
impl UnsandboxedFileSystem {
|
||||
async fn open_file_for_read(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
sandbox: Option<&FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<tokio::fs::File> {
|
||||
reject_platform_sandbox_context(sandbox)?;
|
||||
self.file_system
|
||||
.open_file_for_read(path, /*sandbox*/ None)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn canonicalize(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
@@ -257,6 +311,17 @@ impl UnsandboxedFileSystem {
|
||||
self.file_system.read_file(path, /*sandbox*/ None).await
|
||||
}
|
||||
|
||||
async fn read_file_stream(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
sandbox: Option<&FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<FileSystemReadStream> {
|
||||
reject_platform_sandbox_context(sandbox)?;
|
||||
self.file_system
|
||||
.read_file_stream(path, /*sandbox*/ None)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn write_file(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
@@ -349,6 +414,14 @@ impl ExecutorFileSystem for UnsandboxedFileSystem {
|
||||
Box::pin(UnsandboxedFileSystem::read_file(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(UnsandboxedFileSystem::read_file_stream(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
@@ -414,6 +487,16 @@ impl ExecutorFileSystem for UnsandboxedFileSystem {
|
||||
}
|
||||
|
||||
impl DirectFileSystem {
|
||||
async fn open_file_for_read(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
sandbox: Option<&FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<tokio::fs::File> {
|
||||
reject_sandbox_context(sandbox)?;
|
||||
let path = path.to_abs_path()?;
|
||||
regular_file::open(path.as_path()).await
|
||||
}
|
||||
|
||||
async fn canonicalize(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
@@ -431,16 +514,31 @@ impl DirectFileSystem {
|
||||
path: &PathUri,
|
||||
sandbox: Option<&FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<Vec<u8>> {
|
||||
reject_sandbox_context(sandbox)?;
|
||||
let path = path.to_abs_path()?;
|
||||
let metadata = tokio::fs::metadata(path.as_path()).await?;
|
||||
let file = self.open_file_for_read(path, sandbox).await?;
|
||||
let metadata = file.metadata().await?;
|
||||
if metadata.len() > MAX_READ_FILE_BYTES {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("file is too large to read: limit is {MAX_READ_FILE_BYTES} bytes"),
|
||||
));
|
||||
return Err(file_too_large_error());
|
||||
}
|
||||
tokio::fs::read(path.as_path()).await
|
||||
let mut bytes = Vec::with_capacity(metadata.len() as usize);
|
||||
file.take(MAX_READ_FILE_BYTES + 1)
|
||||
.read_to_end(&mut bytes)
|
||||
.await?;
|
||||
if bytes.len() as u64 > MAX_READ_FILE_BYTES {
|
||||
return Err(file_too_large_error());
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
async fn read_file_stream(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
sandbox: Option<&FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<FileSystemReadStream> {
|
||||
let file = self.open_file_for_read(path, sandbox).await?;
|
||||
Ok(FileSystemReadStream::new(ReaderStream::with_capacity(
|
||||
file,
|
||||
FILE_READ_CHUNK_SIZE,
|
||||
)))
|
||||
}
|
||||
|
||||
async fn write_file(
|
||||
@@ -609,6 +707,14 @@ impl ExecutorFileSystem for DirectFileSystem {
|
||||
Box::pin(DirectFileSystem::read_file(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(DirectFileSystem::read_file_stream(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
|
||||
@@ -21,6 +21,9 @@ 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(crate) const FS_OPEN_METHOD: &str = "fs/open";
|
||||
pub(crate) const FS_READ_BLOCK_METHOD: &str = "fs/readBlock";
|
||||
pub(crate) const FS_CLOSE_METHOD: &str = "fs/close";
|
||||
pub const FS_WRITE_FILE_METHOD: &str = "fs/writeFile";
|
||||
pub const FS_CREATE_DIRECTORY_METHOD: &str = "fs/createDirectory";
|
||||
pub const FS_GET_METADATA_METHOD: &str = "fs/getMetadata";
|
||||
@@ -210,6 +213,45 @@ pub struct FsReadFileResponse {
|
||||
pub data_base64: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsOpenParams {
|
||||
pub handle_id: String,
|
||||
pub path: PathUri,
|
||||
pub sandbox: Option<FileSystemSandboxContext>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsOpenResponse {
|
||||
pub handle_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsReadBlockParams {
|
||||
pub handle_id: String,
|
||||
pub offset: u64,
|
||||
pub len: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsReadBlockResponse {
|
||||
pub chunk: ByteChunk,
|
||||
pub eof: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsCloseParams {
|
||||
pub handle_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsCloseResponse {}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsWriteFileParams {
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
|
||||
pub(crate) async fn open(path: &Path) -> io::Result<tokio::fs::File> {
|
||||
let mut options = tokio::fs::OpenOptions::new();
|
||||
options.read(true);
|
||||
configure_open(&mut options);
|
||||
|
||||
let file = options.open(path).await?;
|
||||
if !is_disk_file(&file) || !file.metadata().await?.is_file() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidInput,
|
||||
format!("path `{}` is not a file", path.display()),
|
||||
));
|
||||
}
|
||||
Ok(file)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn configure_open(options: &mut tokio::fs::OpenOptions) {
|
||||
options.custom_flags(libc::O_NONBLOCK);
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn configure_open(options: &mut tokio::fs::OpenOptions) {
|
||||
use windows_sys::Win32::Storage::FileSystem::SECURITY_IDENTIFICATION;
|
||||
|
||||
options.security_qos_flags(SECURITY_IDENTIFICATION);
|
||||
}
|
||||
|
||||
#[cfg(not(any(unix, windows)))]
|
||||
fn configure_open(_options: &mut tokio::fs::OpenOptions) {}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn is_disk_file(file: &tokio::fs::File) -> bool {
|
||||
use std::os::windows::io::AsRawHandle;
|
||||
use windows_sys::Win32::Foundation::HANDLE;
|
||||
use windows_sys::Win32::Storage::FileSystem::FILE_TYPE_DISK;
|
||||
use windows_sys::Win32::Storage::FileSystem::GetFileType;
|
||||
|
||||
// SAFETY: `file` owns this handle for the duration of the call.
|
||||
unsafe { GetFileType(file.as_raw_handle() as HANDLE) == FILE_TYPE_DISK }
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
fn is_disk_file(_file: &tokio::fs::File) -> bool {
|
||||
true
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
use bytes::Bytes;
|
||||
use codex_utils_path_uri::PathUri;
|
||||
use tokio::io;
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::map_remote_error;
|
||||
use crate::ExecServerClient;
|
||||
use crate::FILE_READ_CHUNK_SIZE;
|
||||
use crate::FileSystemReadStream;
|
||||
use crate::FileSystemResult;
|
||||
use crate::FileSystemSandboxContext;
|
||||
use crate::protocol::FS_READ_BLOCK_METHOD;
|
||||
use crate::protocol::FsCloseParams;
|
||||
use crate::protocol::FsOpenParams;
|
||||
use crate::protocol::FsReadBlockParams;
|
||||
|
||||
struct FileReadRegistration {
|
||||
client: ExecServerClient,
|
||||
handle_id: String,
|
||||
runtime: Option<tokio::runtime::Handle>,
|
||||
active: bool,
|
||||
}
|
||||
|
||||
pub(super) async fn open(
|
||||
client: ExecServerClient,
|
||||
path: PathUri,
|
||||
sandbox: Option<FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<FileSystemReadStream> {
|
||||
let registration = FileReadRegistration {
|
||||
client,
|
||||
handle_id: Uuid::new_v4().simple().to_string(),
|
||||
runtime: tokio::runtime::Handle::try_current().ok(),
|
||||
active: true,
|
||||
};
|
||||
registration
|
||||
.client
|
||||
.fs_open(FsOpenParams {
|
||||
handle_id: registration.handle_id.clone(),
|
||||
path,
|
||||
sandbox,
|
||||
})
|
||||
.await
|
||||
.map_err(map_remote_error)?;
|
||||
Ok(FileSystemReadStream::new(futures::stream::try_unfold(
|
||||
Some((registration, 0_u64)),
|
||||
|state| async move {
|
||||
let Some((mut registration, offset)) = state else {
|
||||
return Ok(None);
|
||||
};
|
||||
let response = registration
|
||||
.client
|
||||
.fs_read_block(FsReadBlockParams {
|
||||
handle_id: registration.handle_id.clone(),
|
||||
offset,
|
||||
len: FILE_READ_CHUNK_SIZE,
|
||||
})
|
||||
.await
|
||||
.map_err(map_remote_error)?;
|
||||
let chunk = Bytes::from(response.chunk.into_inner());
|
||||
if chunk.len() > FILE_READ_CHUNK_SIZE {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
format!(
|
||||
"{FS_READ_BLOCK_METHOD} returned {} bytes, maximum is {}",
|
||||
chunk.len(),
|
||||
FILE_READ_CHUNK_SIZE
|
||||
),
|
||||
));
|
||||
}
|
||||
if response.eof {
|
||||
if registration
|
||||
.client
|
||||
.fs_close(FsCloseParams {
|
||||
handle_id: registration.handle_id.clone(),
|
||||
})
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
registration.active = false;
|
||||
}
|
||||
return if chunk.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some((chunk, None)))
|
||||
};
|
||||
}
|
||||
if chunk.is_empty() {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
format!("{FS_READ_BLOCK_METHOD} returned an empty non-terminal block"),
|
||||
));
|
||||
}
|
||||
let next_offset = offset.checked_add(chunk.len() as u64).ok_or_else(|| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
format!("{FS_READ_BLOCK_METHOD} offset overflowed after {offset} bytes"),
|
||||
)
|
||||
})?;
|
||||
Ok(Some((chunk, Some((registration, next_offset)))))
|
||||
},
|
||||
)))
|
||||
}
|
||||
|
||||
impl Drop for FileReadRegistration {
|
||||
fn drop(&mut self) {
|
||||
if !self.active {
|
||||
return;
|
||||
}
|
||||
let client = self.client.clone();
|
||||
let handle_id = self.handle_id.clone();
|
||||
let runtime = self
|
||||
.runtime
|
||||
.clone()
|
||||
.or_else(|| tokio::runtime::Handle::try_current().ok());
|
||||
if let Some(runtime) = runtime {
|
||||
runtime.spawn(async move {
|
||||
let _ = client.fs_close(FsCloseParams { handle_id }).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -10,6 +10,7 @@ use crate::ExecServerError;
|
||||
use crate::ExecutorFileSystem;
|
||||
use crate::ExecutorFileSystemFuture;
|
||||
use crate::FileMetadata;
|
||||
use crate::FileSystemReadStream;
|
||||
use crate::FileSystemResult;
|
||||
use crate::FileSystemSandboxContext;
|
||||
use crate::ReadDirectoryEntry;
|
||||
@@ -27,6 +28,9 @@ use crate::protocol::FsWriteFileParams;
|
||||
const INVALID_REQUEST_ERROR_CODE: i64 = -32600;
|
||||
const NOT_FOUND_ERROR_CODE: i64 = -32004;
|
||||
|
||||
#[path = "remote_file_stream.rs"]
|
||||
mod file_stream;
|
||||
|
||||
pub(crate) struct RemoteFileSystem {
|
||||
client: LazyRemoteExecServerClient,
|
||||
}
|
||||
@@ -76,6 +80,22 @@ impl RemoteFileSystem {
|
||||
})
|
||||
}
|
||||
|
||||
async fn read_file_stream(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
sandbox: Option<&FileSystemSandboxContext>,
|
||||
) -> FileSystemResult<FileSystemReadStream> {
|
||||
if sandbox.is_some_and(FileSystemSandboxContext::should_run_in_sandbox) {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::Unsupported,
|
||||
"streaming file reads do not support platform sandboxing",
|
||||
));
|
||||
}
|
||||
trace!("remote fs read_file_stream");
|
||||
let client = self.client.get().await.map_err(map_remote_error)?;
|
||||
file_stream::open(client, path.clone(), remote_sandbox_context(sandbox)).await
|
||||
}
|
||||
|
||||
async fn write_file(
|
||||
&self,
|
||||
path: &PathUri,
|
||||
@@ -222,6 +242,14 @@ impl ExecutorFileSystem for RemoteFileSystem {
|
||||
Box::pin(RemoteFileSystem::read_file(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(RemoteFileSystem::read_file_stream(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
|
||||
@@ -10,6 +10,7 @@ use crate::ExecServerRuntimePaths;
|
||||
use crate::ExecutorFileSystem;
|
||||
use crate::ExecutorFileSystemFuture;
|
||||
use crate::FileMetadata;
|
||||
use crate::FileSystemReadStream;
|
||||
use crate::FileSystemResult;
|
||||
use crate::FileSystemSandboxContext;
|
||||
use crate::ReadDirectoryEntry;
|
||||
@@ -265,6 +266,19 @@ impl ExecutorFileSystem for SandboxedFileSystem {
|
||||
Box::pin(SandboxedFileSystem::read_file(self, path, sandbox))
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
_sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(async {
|
||||
Err(io::Error::new(
|
||||
io::ErrorKind::Unsupported,
|
||||
"streaming file reads do not support platform sandboxing",
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
|
||||
@@ -9,16 +9,23 @@ use crate::CreateDirectoryOptions;
|
||||
use crate::ExecServerRuntimePaths;
|
||||
use crate::ExecutorFileSystem;
|
||||
use crate::RemoveOptions;
|
||||
use crate::file_read::FileReadHandleManager;
|
||||
use crate::local_file_system::LocalFileSystem;
|
||||
use crate::protocol::FS_WRITE_FILE_METHOD;
|
||||
use crate::protocol::FsCanonicalizeParams;
|
||||
use crate::protocol::FsCanonicalizeResponse;
|
||||
use crate::protocol::FsCloseParams;
|
||||
use crate::protocol::FsCloseResponse;
|
||||
use crate::protocol::FsCopyParams;
|
||||
use crate::protocol::FsCopyResponse;
|
||||
use crate::protocol::FsCreateDirectoryParams;
|
||||
use crate::protocol::FsCreateDirectoryResponse;
|
||||
use crate::protocol::FsGetMetadataParams;
|
||||
use crate::protocol::FsGetMetadataResponse;
|
||||
use crate::protocol::FsOpenParams;
|
||||
use crate::protocol::FsOpenResponse;
|
||||
use crate::protocol::FsReadBlockParams;
|
||||
use crate::protocol::FsReadBlockResponse;
|
||||
use crate::protocol::FsReadDirectoryEntry;
|
||||
use crate::protocol::FsReadDirectoryParams;
|
||||
use crate::protocol::FsReadDirectoryResponse;
|
||||
@@ -32,18 +39,69 @@ use crate::rpc::internal_error;
|
||||
use crate::rpc::invalid_request;
|
||||
use crate::rpc::not_found;
|
||||
|
||||
const MAX_FILE_READ_HANDLE_ID_BYTES: usize = 32;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct FileSystemHandler {
|
||||
file_system: LocalFileSystem,
|
||||
file_reads: FileReadHandleManager,
|
||||
}
|
||||
|
||||
impl FileSystemHandler {
|
||||
pub(crate) fn new(runtime_paths: ExecServerRuntimePaths) -> Self {
|
||||
Self {
|
||||
file_system: LocalFileSystem::with_runtime_paths(runtime_paths),
|
||||
file_reads: FileReadHandleManager::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn shutdown(&self) {
|
||||
self.file_reads.close_all().await;
|
||||
}
|
||||
|
||||
pub(crate) async fn open(
|
||||
&self,
|
||||
params: FsOpenParams,
|
||||
) -> Result<FsOpenResponse, JSONRPCErrorError> {
|
||||
validate_file_read_handle_id(¶ms.handle_id)?;
|
||||
let file = self
|
||||
.file_system
|
||||
.open_file_for_read(¶ms.path, params.sandbox.as_ref())
|
||||
.await
|
||||
.map_err(map_fs_error)?;
|
||||
let handle_id = self
|
||||
.file_reads
|
||||
.open(params.handle_id, file)
|
||||
.await
|
||||
.map_err(map_fs_error)?;
|
||||
Ok(FsOpenResponse { handle_id })
|
||||
}
|
||||
|
||||
pub(crate) async fn read_block(
|
||||
&self,
|
||||
params: FsReadBlockParams,
|
||||
) -> Result<FsReadBlockResponse, JSONRPCErrorError> {
|
||||
validate_file_read_handle_id(¶ms.handle_id)?;
|
||||
let block = self
|
||||
.file_reads
|
||||
.read_block(¶ms.handle_id, params.offset, params.len)
|
||||
.await
|
||||
.map_err(map_fs_error)?;
|
||||
Ok(FsReadBlockResponse {
|
||||
chunk: block.bytes.into(),
|
||||
eof: block.eof,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn close(
|
||||
&self,
|
||||
params: FsCloseParams,
|
||||
) -> Result<FsCloseResponse, JSONRPCErrorError> {
|
||||
validate_file_read_handle_id(¶ms.handle_id)?;
|
||||
self.file_reads.close(¶ms.handle_id).await;
|
||||
Ok(FsCloseResponse {})
|
||||
}
|
||||
|
||||
pub(crate) async fn read_file(
|
||||
&self,
|
||||
params: FsReadFileParams,
|
||||
@@ -176,6 +234,15 @@ impl FileSystemHandler {
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_file_read_handle_id(handle_id: &str) -> Result<(), JSONRPCErrorError> {
|
||||
if handle_id.len() > MAX_FILE_READ_HANDLE_ID_BYTES {
|
||||
return Err(invalid_request(format!(
|
||||
"file read handle ID must not exceed {MAX_FILE_READ_HANDLE_ID_BYTES} bytes"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn map_fs_error(err: io::Error) -> JSONRPCErrorError {
|
||||
match err.kind() {
|
||||
io::ErrorKind::NotFound => not_found(err.to_string()),
|
||||
|
||||
@@ -19,12 +19,18 @@ use crate::protocol::ExecParams;
|
||||
use crate::protocol::ExecResponse;
|
||||
use crate::protocol::FsCanonicalizeParams;
|
||||
use crate::protocol::FsCanonicalizeResponse;
|
||||
use crate::protocol::FsCloseParams;
|
||||
use crate::protocol::FsCloseResponse;
|
||||
use crate::protocol::FsCopyParams;
|
||||
use crate::protocol::FsCopyResponse;
|
||||
use crate::protocol::FsCreateDirectoryParams;
|
||||
use crate::protocol::FsCreateDirectoryResponse;
|
||||
use crate::protocol::FsGetMetadataParams;
|
||||
use crate::protocol::FsGetMetadataResponse;
|
||||
use crate::protocol::FsOpenParams;
|
||||
use crate::protocol::FsOpenResponse;
|
||||
use crate::protocol::FsReadBlockParams;
|
||||
use crate::protocol::FsReadBlockResponse;
|
||||
use crate::protocol::FsReadDirectoryParams;
|
||||
use crate::protocol::FsReadDirectoryResponse;
|
||||
use crate::protocol::FsReadFileParams;
|
||||
@@ -87,6 +93,7 @@ impl ExecServerHandler {
|
||||
self.background_task_shutdown.cancel();
|
||||
self.background_tasks.close();
|
||||
self.background_tasks.wait().await;
|
||||
self.file_system.shutdown().await;
|
||||
if let Some(session) = self.session() {
|
||||
session.detach().await;
|
||||
}
|
||||
@@ -234,6 +241,30 @@ impl ExecServerHandler {
|
||||
self.file_system.read_file(params).await
|
||||
}
|
||||
|
||||
pub(crate) async fn fs_open(
|
||||
&self,
|
||||
params: FsOpenParams,
|
||||
) -> Result<FsOpenResponse, JSONRPCErrorError> {
|
||||
self.require_initialized_for("filesystem")?;
|
||||
self.file_system.open(params).await
|
||||
}
|
||||
|
||||
pub(crate) async fn fs_read_block(
|
||||
&self,
|
||||
params: FsReadBlockParams,
|
||||
) -> Result<FsReadBlockResponse, JSONRPCErrorError> {
|
||||
self.require_initialized_for("filesystem")?;
|
||||
self.file_system.read_block(params).await
|
||||
}
|
||||
|
||||
pub(crate) async fn fs_close(
|
||||
&self,
|
||||
params: FsCloseParams,
|
||||
) -> Result<FsCloseResponse, JSONRPCErrorError> {
|
||||
self.require_initialized_for("filesystem")?;
|
||||
self.file_system.close(params).await
|
||||
}
|
||||
|
||||
pub(crate) async fn fs_write_file(
|
||||
&self,
|
||||
params: FsWriteFileParams,
|
||||
|
||||
@@ -8,17 +8,23 @@ use crate::protocol::EXEC_TERMINATE_METHOD;
|
||||
use crate::protocol::EXEC_WRITE_METHOD;
|
||||
use crate::protocol::ExecParams;
|
||||
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_GET_METADATA_METHOD;
|
||||
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_REMOVE_METHOD;
|
||||
use crate::protocol::FS_WRITE_FILE_METHOD;
|
||||
use crate::protocol::FsCanonicalizeParams;
|
||||
use crate::protocol::FsCloseParams;
|
||||
use crate::protocol::FsCopyParams;
|
||||
use crate::protocol::FsCreateDirectoryParams;
|
||||
use crate::protocol::FsGetMetadataParams;
|
||||
use crate::protocol::FsOpenParams;
|
||||
use crate::protocol::FsReadBlockParams;
|
||||
use crate::protocol::FsReadDirectoryParams;
|
||||
use crate::protocol::FsReadFileParams;
|
||||
use crate::protocol::FsRemoveParams;
|
||||
@@ -93,6 +99,24 @@ pub(crate) fn build_router() -> RpcRouter<ExecServerHandler> {
|
||||
handler.fs_read_file(params).await
|
||||
},
|
||||
);
|
||||
router.request(
|
||||
FS_OPEN_METHOD,
|
||||
|handler: Arc<ExecServerHandler>, params: FsOpenParams| async move {
|
||||
handler.fs_open(params).await
|
||||
},
|
||||
);
|
||||
router.request(
|
||||
FS_READ_BLOCK_METHOD,
|
||||
|handler: Arc<ExecServerHandler>, params: FsReadBlockParams| async move {
|
||||
handler.fs_read_block(params).await
|
||||
},
|
||||
);
|
||||
router.request(
|
||||
FS_CLOSE_METHOD,
|
||||
|handler: Arc<ExecServerHandler>, params: FsCloseParams| async move {
|
||||
handler.fs_close(params).await
|
||||
},
|
||||
);
|
||||
router.request(
|
||||
FS_WRITE_FILE_METHOD,
|
||||
|handler: Arc<ExecServerHandler>, params: FsWriteFileParams| async move {
|
||||
|
||||
@@ -0,0 +1,396 @@
|
||||
mod common;
|
||||
|
||||
use anyhow::Result;
|
||||
use codex_exec_server::Environment;
|
||||
use codex_exec_server::ExecServerClient;
|
||||
use codex_exec_server::ExecServerError;
|
||||
use codex_exec_server::ExecutorFileSystem;
|
||||
use codex_exec_server::FileSystemSandboxContext;
|
||||
use codex_exec_server::FsCloseParams;
|
||||
use codex_exec_server::FsOpenParams;
|
||||
use codex_exec_server::FsReadBlockParams;
|
||||
use codex_exec_server::FsReadBlockResponse;
|
||||
use codex_exec_server::RemoteExecServerConnectArgs;
|
||||
use codex_protocol::models::PermissionProfile;
|
||||
use codex_protocol::permissions::FileSystemAccessMode;
|
||||
use codex_protocol::permissions::FileSystemPath;
|
||||
use codex_protocol::permissions::FileSystemSandboxEntry;
|
||||
use codex_protocol::permissions::FileSystemSandboxPolicy;
|
||||
use codex_protocol::permissions::NetworkSandboxPolicy;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use codex_utils_path_uri::PathUri;
|
||||
use futures::TryStreamExt;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::sync::Arc;
|
||||
#[cfg(any(unix, windows))]
|
||||
use std::time::Duration;
|
||||
use tempfile::TempDir;
|
||||
#[cfg(windows)]
|
||||
use tokio::net::windows::named_pipe::ServerOptions;
|
||||
#[cfg(any(unix, windows))]
|
||||
use tokio::time::timeout;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::common::exec_server::exec_server;
|
||||
|
||||
const BLOCK_SIZE: usize = 1024 * 1024;
|
||||
const OPEN_FILE_LIMIT: usize = 128;
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_stops_after_an_exact_block_boundary() -> Result<()> {
|
||||
let server = exec_server().await?;
|
||||
let file_system = connect_file_system(server.websocket_url())?;
|
||||
let tmp = TempDir::new()?;
|
||||
let path = tmp.path().join("exact-blocks.bin");
|
||||
std::fs::write(&path, vec![b'x'; BLOCK_SIZE * 2])?;
|
||||
|
||||
let chunks = file_system
|
||||
.read_file_stream(&PathUri::from_path(path)?, /*sandbox*/ None)
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
chunks.iter().map(bytes::Bytes::len).collect::<Vec<_>>(),
|
||||
vec![BLOCK_SIZE, BLOCK_SIZE]
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn completed_streams_release_handle_capacity() -> Result<()> {
|
||||
let server = exec_server().await?;
|
||||
let file_system = connect_file_system(server.websocket_url())?;
|
||||
let tmp = TempDir::new()?;
|
||||
let path = tmp.path().join("repeated.txt");
|
||||
std::fs::write(&path, b"repeated")?;
|
||||
let path = PathUri::from_path(path)?;
|
||||
|
||||
for _ in 0..=OPEN_FILE_LIMIT {
|
||||
let chunks = file_system
|
||||
.read_file_stream(&path, /*sandbox*/ None)
|
||||
.await?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
assert_eq!(chunks, vec![bytes::Bytes::from_static(b"repeated")]);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_rejects_platform_sandbox() -> Result<()> {
|
||||
let server = exec_server().await?;
|
||||
let file_system = connect_file_system(server.websocket_url())?;
|
||||
let tmp = TempDir::new()?;
|
||||
let path = tmp.path().join("sandboxed.txt");
|
||||
std::fs::write(&path, "sandboxed hello")?;
|
||||
|
||||
let result = file_system
|
||||
.read_file_stream(
|
||||
&PathUri::from_path(&path)?,
|
||||
Some(&read_only_sandbox(tmp.path().to_path_buf())),
|
||||
)
|
||||
.await;
|
||||
|
||||
let Err(error) = result else {
|
||||
panic!("sandboxed stream should be rejected");
|
||||
};
|
||||
assert_eq!(error.kind(), std::io::ErrorKind::Unsupported);
|
||||
assert_eq!(
|
||||
error.to_string(),
|
||||
"streaming file reads do not support platform sandboxing"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn file_reads_reject_fifo_without_waiting_for_a_writer() -> Result<()> {
|
||||
let server = exec_server().await?;
|
||||
let file_system = connect_file_system(server.websocket_url())?;
|
||||
let tmp = TempDir::new()?;
|
||||
let path = tmp.path().join("named-pipe");
|
||||
let output = std::process::Command::new("mkfifo").arg(&path).output()?;
|
||||
if !output.status.success() {
|
||||
anyhow::bail!(
|
||||
"mkfifo failed: stdout={} stderr={}",
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
);
|
||||
}
|
||||
|
||||
let path_uri = PathUri::from_path(&path)?;
|
||||
let read_error = timeout(
|
||||
Duration::from_secs(1),
|
||||
file_system.read_file(&path_uri, /*sandbox*/ None),
|
||||
)
|
||||
.await
|
||||
.expect("reading a FIFO should not wait for a writer")
|
||||
.expect_err("reading a FIFO should be rejected");
|
||||
let stream_result = timeout(
|
||||
Duration::from_secs(1),
|
||||
file_system.read_file_stream(&path_uri, /*sandbox*/ None),
|
||||
)
|
||||
.await
|
||||
.expect("streaming a FIFO should not wait for a writer");
|
||||
let Err(stream_error) = stream_result else {
|
||||
panic!("streaming a FIFO should be rejected");
|
||||
};
|
||||
let expected = format!("path `{}` is not a file", path.display());
|
||||
assert_eq!(
|
||||
(read_error.to_string(), stream_error.to_string()),
|
||||
(expected.clone(), expected)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
#[tokio::test]
|
||||
async fn file_reads_reject_named_pipes() -> Result<()> {
|
||||
let server = exec_server().await?;
|
||||
let file_system = connect_file_system(server.websocket_url())?;
|
||||
|
||||
let read_path = format!(r"\\.\pipe\codex-fs-read-{}", Uuid::new_v4());
|
||||
let _read_pipe = ServerOptions::new()
|
||||
.first_pipe_instance(true)
|
||||
.create(&read_path)?;
|
||||
let read_error = timeout(
|
||||
Duration::from_secs(1),
|
||||
file_system.read_file(
|
||||
&PathUri::from_path(std::path::Path::new(&read_path))?,
|
||||
/*sandbox*/ None,
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("reading a named pipe should not hang")
|
||||
.expect_err("reading a named pipe should be rejected");
|
||||
|
||||
let stream_path = format!(r"\\.\pipe\codex-fs-stream-{}", Uuid::new_v4());
|
||||
let _stream_pipe = ServerOptions::new()
|
||||
.first_pipe_instance(true)
|
||||
.create(&stream_path)?;
|
||||
let stream_result = timeout(
|
||||
Duration::from_secs(1),
|
||||
file_system.read_file_stream(
|
||||
&PathUri::from_path(std::path::Path::new(&stream_path))?,
|
||||
/*sandbox*/ None,
|
||||
),
|
||||
)
|
||||
.await
|
||||
.expect("streaming a named pipe should not hang");
|
||||
let Err(stream_error) = stream_result else {
|
||||
panic!("streaming a named pipe should be rejected");
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
(read_error.kind(), stream_error.kind()),
|
||||
(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn stream_keeps_reading_the_open_file_after_path_replacement() -> Result<()> {
|
||||
let server = exec_server().await?;
|
||||
let file_system = connect_file_system(server.websocket_url())?;
|
||||
let tmp = TempDir::new()?;
|
||||
let path = tmp.path().join("replaceable.bin");
|
||||
std::fs::write(&path, vec![b'a'; BLOCK_SIZE + 1])?;
|
||||
let mut stream = file_system
|
||||
.read_file_stream(&PathUri::from_path(&path)?, /*sandbox*/ None)
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
stream.try_next().await?,
|
||||
Some(bytes::Bytes::from(vec![b'a'; BLOCK_SIZE]))
|
||||
);
|
||||
let replacement = tmp.path().join("replacement.bin");
|
||||
std::fs::write(&replacement, vec![b'b'; BLOCK_SIZE + 1])?;
|
||||
std::fs::remove_file(&path)?;
|
||||
std::fs::rename(replacement, &path)?;
|
||||
|
||||
assert_eq!(
|
||||
stream.try_next().await?,
|
||||
Some(bytes::Bytes::from_static(b"a"))
|
||||
);
|
||||
assert_eq!(stream.try_next().await?, None);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_block_supports_non_sequential_offsets_and_lengths() -> Result<()> {
|
||||
let mut server = exec_server().await?;
|
||||
let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
|
||||
server.websocket_url().to_string(),
|
||||
"file-stream-protocol-test".to_string(),
|
||||
))
|
||||
.await?;
|
||||
let tmp = TempDir::new()?;
|
||||
let path = tmp.path().join("non-sequential.bin");
|
||||
std::fs::write(&path, b"0123456789")?;
|
||||
let open = client
|
||||
.fs_open(FsOpenParams {
|
||||
handle_id: Uuid::new_v4().simple().to_string(),
|
||||
path: PathUri::from_path(path)?,
|
||||
sandbox: None,
|
||||
})
|
||||
.await?;
|
||||
|
||||
let mut blocks = Vec::new();
|
||||
for (offset, len) in [(6, 3), (1, 2), (8, 4), (0, 2)] {
|
||||
blocks.push(
|
||||
client
|
||||
.fs_read_block(FsReadBlockParams {
|
||||
handle_id: open.handle_id.clone(),
|
||||
offset,
|
||||
len,
|
||||
})
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
blocks,
|
||||
vec![
|
||||
FsReadBlockResponse {
|
||||
chunk: b"678".to_vec().into(),
|
||||
eof: false,
|
||||
},
|
||||
FsReadBlockResponse {
|
||||
chunk: b"12".to_vec().into(),
|
||||
eof: false,
|
||||
},
|
||||
FsReadBlockResponse {
|
||||
chunk: b"89".to_vec().into(),
|
||||
eof: true,
|
||||
},
|
||||
FsReadBlockResponse {
|
||||
chunk: b"01".to_vec().into(),
|
||||
eof: false,
|
||||
},
|
||||
]
|
||||
);
|
||||
client
|
||||
.fs_close(FsCloseParams {
|
||||
handle_id: open.handle_id,
|
||||
})
|
||||
.await?;
|
||||
drop(client);
|
||||
server.shutdown().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn open_enforces_the_per_connection_limit_and_close_releases_capacity() -> Result<()> {
|
||||
let mut server = exec_server().await?;
|
||||
let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
|
||||
server.websocket_url().to_string(),
|
||||
"file-stream-protocol-test".to_string(),
|
||||
))
|
||||
.await?;
|
||||
let tmp = TempDir::new()?;
|
||||
let path = tmp.path().join("limited.bin");
|
||||
std::fs::write(&path, b"limited")?;
|
||||
let path = PathUri::from_path(path)?;
|
||||
let mut handles = Vec::with_capacity(OPEN_FILE_LIMIT);
|
||||
for _ in 0..OPEN_FILE_LIMIT {
|
||||
let open = client
|
||||
.fs_open(FsOpenParams {
|
||||
handle_id: Uuid::new_v4().simple().to_string(),
|
||||
path: path.clone(),
|
||||
sandbox: None,
|
||||
})
|
||||
.await?;
|
||||
handles.push(open.handle_id);
|
||||
}
|
||||
|
||||
let error = client
|
||||
.fs_open(FsOpenParams {
|
||||
handle_id: Uuid::new_v4().simple().to_string(),
|
||||
path: path.clone(),
|
||||
sandbox: None,
|
||||
})
|
||||
.await
|
||||
.expect_err("opening beyond the limit should fail");
|
||||
let ExecServerError::Server { code, message } = error else {
|
||||
anyhow::bail!("expected server error, got {error:?}");
|
||||
};
|
||||
assert_eq!(
|
||||
(code, message),
|
||||
(
|
||||
-32600,
|
||||
format!("at most {OPEN_FILE_LIMIT} file reads may be open per connection"),
|
||||
)
|
||||
);
|
||||
|
||||
client
|
||||
.fs_close(FsCloseParams {
|
||||
handle_id: handles.remove(0),
|
||||
})
|
||||
.await?;
|
||||
client
|
||||
.fs_open(FsOpenParams {
|
||||
handle_id: Uuid::new_v4().simple().to_string(),
|
||||
path,
|
||||
sandbox: None,
|
||||
})
|
||||
.await?;
|
||||
drop(client);
|
||||
server.shutdown().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn open_rejects_handle_ids_longer_than_32_bytes() -> Result<()> {
|
||||
let server = exec_server().await?;
|
||||
let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs::new(
|
||||
server.websocket_url().to_string(),
|
||||
"file-stream-protocol-test".to_string(),
|
||||
))
|
||||
.await?;
|
||||
let tmp = TempDir::new()?;
|
||||
let path = tmp.path().join("handle-id-limit.bin");
|
||||
std::fs::write(&path, b"limited")?;
|
||||
|
||||
let error = client
|
||||
.fs_open(FsOpenParams {
|
||||
handle_id: "x".repeat(33),
|
||||
path: PathUri::from_path(path)?,
|
||||
sandbox: None,
|
||||
})
|
||||
.await
|
||||
.expect_err("oversized handle ID should fail");
|
||||
|
||||
let ExecServerError::Server { code, message } = error else {
|
||||
anyhow::bail!("expected server error, got {error:?}");
|
||||
};
|
||||
assert_eq!(
|
||||
(code, message),
|
||||
(
|
||||
-32600,
|
||||
"file read handle ID must not exceed 32 bytes".to_string(),
|
||||
)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn connect_file_system(websocket_url: &str) -> Result<Arc<dyn ExecutorFileSystem>> {
|
||||
let environment = Environment::create_for_tests(Some(websocket_url.to_string()))?;
|
||||
Ok(environment.get_filesystem())
|
||||
}
|
||||
|
||||
fn read_only_sandbox(path: std::path::PathBuf) -> FileSystemSandboxContext {
|
||||
let path = AbsolutePathBuf::from_absolute_path(&path)
|
||||
.unwrap_or_else(|err| panic!("sandbox path should be absolute: {err}"));
|
||||
FileSystemSandboxContext::from_permission_profile(PermissionProfile::from_runtime_permissions(
|
||||
&FileSystemSandboxPolicy::restricted(vec![FileSystemSandboxEntry {
|
||||
path: FileSystemPath::Path { path },
|
||||
access: FileSystemAccessMode::Read,
|
||||
}]),
|
||||
NetworkSandboxPolicy::Restricted,
|
||||
))
|
||||
}
|
||||
@@ -2,6 +2,7 @@ use anyhow::Context;
|
||||
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::FileMetadata;
|
||||
use codex_exec_server::ReadDirectoryEntry;
|
||||
use codex_exec_server::RemoveOptions;
|
||||
@@ -11,6 +12,7 @@ use codex_protocol::models::PermissionProfile;
|
||||
use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy;
|
||||
use codex_sandboxing::policy_transforms::effective_network_sandbox_policy;
|
||||
use codex_utils_path_uri::PathUri;
|
||||
use futures::TryStreamExt;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::path::Path;
|
||||
use tempfile::TempDir;
|
||||
@@ -194,6 +196,45 @@ async fn file_system_read_file_returns_bytes(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test_case(FileSystemImplementation::Local ; "local")]
|
||||
#[test_case(FileSystemImplementation::Remote ; "remote")]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn file_system_read_file_stream_returns_bounded_chunks(
|
||||
implementation: FileSystemImplementation,
|
||||
) -> Result<()> {
|
||||
let context = create_file_system_context(implementation).await?;
|
||||
let file_system = context.file_system;
|
||||
|
||||
let tmp = TempDir::new()?;
|
||||
let file_path = tmp.path().join("blocks.bin");
|
||||
let contents = (0..FILE_READ_CHUNK_SIZE * 2 + 17)
|
||||
.map(|index| (index % 251) as u8)
|
||||
.collect::<Vec<_>>();
|
||||
std::fs::write(&file_path, &contents)?;
|
||||
|
||||
let chunks = file_system
|
||||
.read_file_stream(&PathUri::from_path(file_path)?, /*sandbox*/ None)
|
||||
.await
|
||||
.with_context(|| format!("mode={implementation}"))?
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
|
||||
assert!(
|
||||
chunks
|
||||
.iter()
|
||||
.all(|chunk| !chunk.is_empty() && chunk.len() <= FILE_READ_CHUNK_SIZE)
|
||||
);
|
||||
assert_eq!(
|
||||
chunks
|
||||
.iter()
|
||||
.flat_map(|chunk| chunk.iter().copied())
|
||||
.collect::<Vec<_>>(),
|
||||
contents
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test_case(FileSystemImplementation::Local ; "local")]
|
||||
#[test_case(FileSystemImplementation::Remote ; "remote")]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
|
||||
@@ -8,6 +8,7 @@ use codex_exec_server::CreateDirectoryOptions;
|
||||
use codex_exec_server::ExecutorFileSystem;
|
||||
use codex_exec_server::ExecutorFileSystemFuture;
|
||||
use codex_exec_server::FileMetadata;
|
||||
use codex_exec_server::FileSystemReadStream;
|
||||
use codex_exec_server::FileSystemResult;
|
||||
use codex_exec_server::FileSystemSandboxContext;
|
||||
use codex_exec_server::ReadDirectoryEntry;
|
||||
@@ -73,6 +74,14 @@ impl ExecutorFileSystem for SyntheticExecutorFileSystem {
|
||||
})
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
_sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(async { Self::unsupported() })
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
|
||||
@@ -12,6 +12,7 @@ use codex_exec_server::EnvironmentManager;
|
||||
use codex_exec_server::ExecutorFileSystem;
|
||||
use codex_exec_server::ExecutorFileSystemFuture;
|
||||
use codex_exec_server::FileMetadata;
|
||||
use codex_exec_server::FileSystemReadStream;
|
||||
use codex_exec_server::FileSystemSandboxContext;
|
||||
use codex_exec_server::ReadDirectoryEntry;
|
||||
use codex_exec_server::RemoveOptions;
|
||||
@@ -109,6 +110,19 @@ impl ExecutorFileSystem for SyntheticFileSystem {
|
||||
Box::pin(SyntheticFileSystem::read_file(self, path))
|
||||
}
|
||||
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
_sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> {
|
||||
Box::pin(async {
|
||||
Err(io::Error::new(
|
||||
io::ErrorKind::Unsupported,
|
||||
"synthetic filesystem does not support streaming reads",
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
fn write_file<'a>(
|
||||
&'a self,
|
||||
_path: &'a PathUri,
|
||||
|
||||
@@ -8,9 +8,11 @@ license.workspace = true
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
bytes = { workspace = true }
|
||||
codex-protocol = { workspace = true }
|
||||
codex-utils-absolute-path = { workspace = true }
|
||||
codex-utils-path-uri = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
|
||||
[lib]
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
use bytes::Bytes;
|
||||
use codex_protocol::config_types::WindowsSandboxLevel;
|
||||
use codex_protocol::models::ManagedFileSystemPermissions;
|
||||
use codex_protocol::models::PermissionProfile;
|
||||
@@ -10,10 +11,16 @@ use codex_protocol::permissions::NetworkSandboxPolicy;
|
||||
use codex_protocol::protocol::SandboxPolicy;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use codex_utils_path_uri::PathUri;
|
||||
use futures::Stream;
|
||||
use std::future::Future;
|
||||
use std::io;
|
||||
use std::path::Path;
|
||||
use std::pin::Pin;
|
||||
use std::task::Context;
|
||||
use std::task::Poll;
|
||||
|
||||
/// Maximum chunk size returned by [`ExecutorFileSystem::read_file_stream`].
|
||||
pub const FILE_READ_CHUNK_SIZE: usize = 1024 * 1024;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub struct CreateDirectoryOptions {
|
||||
@@ -155,6 +162,28 @@ pub type FileSystemResult<T> = io::Result<T>;
|
||||
pub type ExecutorFileSystemFuture<'a, T> =
|
||||
Pin<Box<dyn Future<Output = FileSystemResult<T>> + Send + 'a>>;
|
||||
|
||||
/// Stream of immutable chunks read from an [`ExecutorFileSystem`].
|
||||
pub struct FileSystemReadStream {
|
||||
inner: Pin<Box<dyn Stream<Item = FileSystemResult<Bytes>> + Send + 'static>>,
|
||||
}
|
||||
|
||||
impl FileSystemReadStream {
|
||||
/// Wraps a filesystem byte stream.
|
||||
pub fn new(stream: impl Stream<Item = FileSystemResult<Bytes>> + Send + 'static) -> Self {
|
||||
Self {
|
||||
inner: Box::pin(stream),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for FileSystemReadStream {
|
||||
type Item = FileSystemResult<Bytes>;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
self.inner.as_mut().poll_next(cx)
|
||||
}
|
||||
}
|
||||
|
||||
/// Abstract filesystem access used by components that may operate locally or via
|
||||
/// a remote environment.
|
||||
pub trait ExecutorFileSystem: Send + Sync {
|
||||
@@ -171,6 +200,13 @@ pub trait ExecutorFileSystem: Send + Sync {
|
||||
sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, Vec<u8>>;
|
||||
|
||||
/// Reads a file as a stream of chunks no larger than [`FILE_READ_CHUNK_SIZE`].
|
||||
fn read_file_stream<'a>(
|
||||
&'a self,
|
||||
path: &'a PathUri,
|
||||
sandbox: Option<&'a FileSystemSandboxContext>,
|
||||
) -> ExecutorFileSystemFuture<'a, FileSystemReadStream>;
|
||||
|
||||
/// Reads a file and decodes it as UTF-8 text.
|
||||
fn read_file_text<'a>(
|
||||
&'a self,
|
||||
|
||||
Reference in New Issue
Block a user