Codex Handbook
exec-server/src/remote_file_stream.rs 121 lines
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;            });        }    }}