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:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user