fix: add bounded binary reads for ViewImage

This commit is contained in:
2026-09-30 07:44:59 +09:00
parent cafbb73a22
commit 05d6da45fb
14 changed files with 816 additions and 45 deletions
+99 -1
View File
@@ -15,7 +15,8 @@ use thiserror::Error;
pub use glob::run_glob; pub use glob::run_glob;
pub use local::{ pub use local::{
resolve_access_path, run_edit, run_list, run_read, run_read_bounded, run_stat, run_write, resolve_access_path, run_edit, run_list, run_read, run_read_bounded, run_read_bytes,
run_read_bytes_bounded, run_stat, run_write,
}; };
pub use operation::*; pub use operation::*;
pub use search::run_grep; pub use search::run_grep;
@@ -569,6 +570,103 @@ mod tests {
} }
} }
#[test]
fn bounded_byte_reads_preserve_binary_ranges_and_fence_versions() {
let directory = tempfile::tempdir().unwrap();
let root = directory.path().canonicalize().unwrap();
let original = b"\x89PNG\r\n\x1a\n\0line\n\xfftail";
std::fs::write(root.join("image.bin"), original).unwrap();
let access = RootAccess(root.clone());
let limits = BoundedReadLimits::new(1024, 8).unwrap();
let first = run_read_bytes_bounded(
&root,
ReadBytesRequest {
path: FsPath::new("image.bin").unwrap(),
offset: 0,
max_bytes: 8,
expected_hash: None,
},
&access,
limits,
)
.unwrap();
assert_eq!(first.bytes, original[..8]);
assert_eq!(first.total_bytes, original.len() as u64);
assert!(!first.eof);
let second = run_read_bytes_bounded(
&root,
ReadBytesRequest {
path: FsPath::new("image.bin").unwrap(),
offset: 8,
max_bytes: 8,
expected_hash: Some(first.content_hash),
},
&access,
limits,
)
.unwrap();
assert_eq!(second.bytes, original[8..16]);
assert_eq!(second.content_hash, first.content_hash);
std::fs::write(root.join("image.bin"), b"different bytes long!").unwrap();
let error = run_read_bytes_bounded(
&root,
ReadBytesRequest {
path: FsPath::new("image.bin").unwrap(),
offset: 16,
max_bytes: 8,
expected_hash: Some(first.content_hash),
},
&access,
limits,
)
.expect_err("a later chunk must not cross file versions");
assert!(matches!(error, FsError::Conflict(_)));
}
#[test]
fn bounded_byte_reads_reject_request_and_source_limit_overflow() {
let directory = tempfile::tempdir().unwrap();
let root = directory.path().canonicalize().unwrap();
std::fs::write(root.join("image.bin"), [0_u8; 17]).unwrap();
let access = RootAccess(root.clone());
let limits = BoundedReadLimits::new(16, 8).unwrap();
let response_error = run_read_bytes_bounded(
&root,
ReadBytesRequest {
path: FsPath::new("image.bin").unwrap(),
offset: 0,
max_bytes: 9,
expected_hash: None,
},
&access,
limits,
)
.expect_err("requests above the response ceiling must be rejected");
assert!(
matches!(response_error, FsError::InvalidArgument(message) if message.contains("8 bytes"))
);
let source_error = run_read_bytes_bounded(
&root,
ReadBytesRequest {
path: FsPath::new("image.bin").unwrap(),
offset: 0,
max_bytes: 8,
expected_hash: None,
},
&access,
limits,
)
.expect_err("sources above the provider ceiling must be rejected");
assert!(
matches!(source_error, FsError::InvalidArgument(message) if message.contains("provider limit 16"))
);
}
#[test] #[test]
fn cancellation_stops_a_read_already_in_progress() { fn cancellation_stops_a_read_already_in_progress() {
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+134 -2
View File
@@ -7,8 +7,9 @@ use sha2::{Digest, Sha256};
use crate::{ use crate::{
AtomicWriteMode, BoundedReadLimits, ContentHash, EditRequest, EditResult, EntryKind, AtomicWriteMode, BoundedReadLimits, ContentHash, EditRequest, EditResult, EntryKind,
FsAccessPolicy, FsError, FsPath, ListEntry, ListRequest, ListResult, ReadRequest, ReadResult, FsAccessPolicy, FsError, FsPath, ListEntry, ListRequest, ListResult, ReadBytesRequest,
StatRequest, StatResult, WriteRequest, WriteResult, direct_symlink, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult,
direct_symlink,
}; };
const READ_BUFFER_BYTES: usize = 16 * 1024; const READ_BUFFER_BYTES: usize = 16 * 1024;
@@ -195,6 +196,137 @@ fn run_read_with_limits(
}) })
} }
pub fn run_read_bytes(
root: &Path,
request: ReadBytesRequest,
access: &dyn FsAccessPolicy,
) -> Result<ReadBytesResult, FsError> {
run_read_bytes_with_limits(root, request, access, None)
}
pub fn run_read_bytes_bounded(
root: &Path,
request: ReadBytesRequest,
access: &dyn FsAccessPolicy,
limits: BoundedReadLimits,
) -> Result<ReadBytesResult, FsError> {
run_read_bytes_with_limits(root, request, access, Some(limits))
}
fn run_read_bytes_with_limits(
root: &Path,
request: ReadBytesRequest,
access: &dyn FsAccessPolicy,
limits: Option<BoundedReadLimits>,
) -> Result<ReadBytesResult, FsError> {
let logical = request.path;
if request.max_bytes == 0 {
return Err(FsError::InvalidArgument(
"byte read max_bytes must be greater than zero".to_string(),
));
}
if let Some(limits) = limits
&& request.max_bytes > limits.max_response_bytes
{
return Err(FsError::InvalidArgument(format!(
"byte read response limit is {} bytes",
limits.max_response_bytes
)));
}
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
let path = resolve(root, &logical)?;
let target = require_access(&path, &logical, access, false, false)?;
let file = access
.open_read_file(&path, &target)
.map_err(|error| map_io(&logical, error))?;
let metadata = file.metadata().map_err(|error| map_io(&logical, error))?;
if metadata.is_dir() {
return Err(FsError::IsDirectory(PathBuf::from(logical.as_str())));
}
if !metadata.is_file() {
return Err(FsError::InvalidArgument(
"byte read source must be a regular file".to_string(),
));
}
if let Some(limits) = limits
&& metadata.len() > limits.max_source_bytes
{
return Err(FsError::InvalidArgument(format!(
"byte read source exceeds provider limit {}",
limits.max_source_bytes
)));
}
let mut bytes = Vec::with_capacity(request.max_bytes.min(READ_BUFFER_BYTES));
let mut source_bytes = 0_u64;
let mut hasher = Sha256::new();
let mut reader = BufReader::with_capacity(READ_BUFFER_BYTES, file);
let mut buffer = [0_u8; READ_BUFFER_BYTES];
loop {
access
.check_cancelled()
.map_err(|error| map_io(&logical, error))?;
let read_limit = limits
.map(|limits| {
limits
.max_source_bytes
.saturating_sub(source_bytes)
.saturating_add(1)
.min(READ_BUFFER_BYTES as u64) as usize
})
.unwrap_or(READ_BUFFER_BYTES);
let read = reader
.read(&mut buffer[..read_limit])
.map_err(|error| map_io(&logical, error))?;
if read == 0 {
break;
}
let chunk_start = source_bytes;
source_bytes = source_bytes.saturating_add(read as u64);
if let Some(limits) = limits
&& source_bytes > limits.max_source_bytes
{
return Err(FsError::InvalidArgument(format!(
"byte read source exceeds provider limit {}",
limits.max_source_bytes
)));
}
hasher.update(&buffer[..read]);
if bytes.len() < request.max_bytes && source_bytes > request.offset {
let start = request.offset.saturating_sub(chunk_start).min(read as u64) as usize;
let retained = (request.max_bytes - bytes.len()).min(read - start);
bytes.extend_from_slice(&buffer[start..start + retained]);
}
}
let content_hash: ContentHash = hasher.finalize().into();
if request
.expected_hash
.is_some_and(|expected| expected != content_hash)
{
return Err(FsError::Conflict(logical.as_str().to_string()));
}
if request.offset > source_bytes {
return Err(FsError::InvalidArgument(format!(
"byte offset {} exceeds file size {source_bytes}",
request.offset
)));
}
let eof = request.offset.saturating_add(bytes.len() as u64) == source_bytes;
Ok(ReadBytesResult {
path: logical,
bytes,
offset: request.offset,
total_bytes: source_bytes,
content_hash,
eof,
})
}
pub fn run_write( pub fn run_write(
root: &Path, root: &Path,
request: WriteRequest, request: WriteRequest,
+30
View File
@@ -207,6 +207,36 @@ pub struct ReadResult {
pub truncated: bool, pub truncated: bool,
} }
/// Binary-safe byte range read with an optional whole-file version fence.
///
/// Unlike [`ReadRequest`], offsets and limits are bytes rather than lines. The
/// provider still hashes the complete bounded source so callers can reject
/// chunks from different file versions.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct ReadBytesRequest {
pub path: FsPath,
#[schemars(range(min = 0, max = 9_007_199_254_740_991_u64))]
pub offset: u64,
#[schemars(range(min = 1, max = 9_007_199_254_740_991_usize))]
pub max_bytes: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub expected_hash: Option<ContentHash>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct ReadBytesResult {
pub path: FsPath,
pub bytes: Vec<u8>,
#[schemars(range(min = 0, max = 9_007_199_254_740_991_u64))]
pub offset: u64,
#[schemars(range(min = 0, max = 9_007_199_254_740_991_u64))]
pub total_bytes: u64,
pub content_hash: ContentHash,
pub eof: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)] #[serde(deny_unknown_fields)]
pub struct WriteRequest { pub struct WriteRequest {
+92 -21
View File
@@ -7,12 +7,14 @@ use agen::tool::{
}; };
use async_trait::async_trait; use async_trait::async_trait;
use serde::Deserialize; use serde::Deserialize;
use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; use workdir::{ReadBytesRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter};
use crate::error::ToolsError; use crate::error::ToolsError;
/// Maximum image body accepted for one model request. /// Maximum image body accepted for one model request.
pub const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024; pub const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024;
/// One binary response remains within the External Workdir transport bound.
const IMAGE_READ_CHUNK_BYTES: usize = 1024 * 1024;
const DESCRIPTION: &str = "Attach an image from the selected Workdir attachment to the next model request. \ const DESCRIPTION: &str = "Attach an image from the selected Workdir attachment to the next model request. \
The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP. \ The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP. \
@@ -47,42 +49,111 @@ impl Tool for ViewImageTool {
workdir::WorkdirSessionCapability::Read, workdir::WorkdirSessionCapability::Read,
)?; )?;
let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?; let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?;
let result = selected let bytes = read_image_bytes(selected.session.as_ref(), path.clone())
.session
.read(ReadRequest {
path: path.clone(),
offset: 0,
limit: usize::MAX,
// The scoped provider enforces this cap while reading, rather
// than allocating an unbounded binary body first.
max_bytes: MAX_IMAGE_BYTES + 1,
})
.await .await
.map_err(ToolsError::from)?; .map_err(ToolsError::from)?;
let mime_type = detect_image_mime(&bytes).ok_or_else(|| {
if result.truncated || result.bytes.len() > MAX_IMAGE_BYTES {
return Err(ToolError::InvalidArgument(format!(
"image exceeds the {MAX_IMAGE_BYTES}-byte limit"
)));
}
let mime_type = detect_image_mime(&result.bytes).ok_or_else(|| {
ToolError::InvalidArgument( ToolError::InvalidArgument(
"unsupported image; expected PNG, JPEG, GIF, or WebP bytes".to_string(), "unsupported image; expected PNG, JPEG, GIF, or WebP bytes".to_string(),
) )
})?; })?;
let bytes = result.bytes.len(); let byte_count = bytes.len();
Ok(ToolOutput { Ok(ToolOutput {
summary: format!("Attached image {path} ({mime_type}, {bytes} bytes)"), summary: format!("Attached image {path} ({mime_type}, {byte_count} bytes)"),
content: None, content: None,
attachments: vec![Attachment::Image(ImageAttachment::new( attachments: vec![Attachment::Image(ImageAttachment::new(
mime_type, mime_type,
Arc::<[u8]>::from(result.bytes), Arc::<[u8]>::from(bytes),
))], ))],
}) })
} }
} }
async fn read_image_bytes(
session: &dyn workdir::WorkdirSession,
path: WorkdirPath,
) -> Result<Vec<u8>, workdir::WorkdirError> {
let mut bytes = Vec::new();
let mut offset = 0_u64;
let mut expected_hash = None;
let mut expected_total = None;
loop {
let remaining = expected_total
.map(|total: u64| total.saturating_sub(offset) as usize)
.unwrap_or(IMAGE_READ_CHUNK_BYTES);
let request_bytes = remaining.min(IMAGE_READ_CHUNK_BYTES).max(1);
let result = session
.read_bytes(ReadBytesRequest {
path: path.clone(),
offset,
max_bytes: request_bytes,
expected_hash,
})
.await?;
if result.path != path || result.offset != offset {
return Err(workdir::WorkdirError::Transport(
"bounded binary read returned a mismatched path or offset".to_string(),
));
}
if result.bytes.len() > request_bytes {
return Err(workdir::WorkdirError::Transport(
"bounded binary read exceeded the requested response size".to_string(),
));
}
let end = result
.offset
.checked_add(result.bytes.len() as u64)
.ok_or_else(|| {
workdir::WorkdirError::Transport(
"bounded binary read returned an invalid byte range".to_string(),
)
})?;
if end > result.total_bytes || result.eof != (end == result.total_bytes) {
return Err(workdir::WorkdirError::Transport(
"bounded binary read returned a truncated or invalid byte range".to_string(),
));
}
if result.total_bytes > MAX_IMAGE_BYTES as u64 {
return Err(workdir::WorkdirError::InvalidArgument(format!(
"image exceeds the {MAX_IMAGE_BYTES}-byte limit"
)));
}
if expected_total.is_some_and(|total| total != result.total_bytes)
|| expected_hash.is_some_and(|hash| hash != result.content_hash)
{
return Err(workdir::WorkdirError::Conflict(
"image changed while it was being read; retry ViewImage".to_string(),
));
}
if expected_total.is_none() {
bytes
.try_reserve_exact(result.total_bytes as usize)
.map_err(|_| workdir::WorkdirError::OperationFailed)?;
expected_total = Some(result.total_bytes);
expected_hash = Some(result.content_hash);
}
if result.bytes.is_empty() && !result.eof {
return Err(workdir::WorkdirError::Transport(
"bounded binary read made no progress".to_string(),
));
}
bytes.extend_from_slice(&result.bytes);
offset = end;
if result.eof {
if bytes.len() as u64 != result.total_bytes {
return Err(workdir::WorkdirError::Transport(
"bounded binary read returned a partial image".to_string(),
));
}
return Ok(bytes);
}
}
}
pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition { pub fn view_image_tool(session: WorkdirSessionHandle) -> ToolDefinition {
routed_view_image_tool(crate::routing::singleton_router(session)) routed_view_image_tool(crate::routing::singleton_router(session))
} }
+218
View File
@@ -399,6 +399,224 @@ async fn detach_and_reattach_same_alias_requires_a_fresh_read() {
router.close_all().await.unwrap(); router.close_all().await.unwrap();
} }
#[derive(Debug)]
struct ExternalValidatingSession {
inner: LocalWorkdirSession,
truncate_first_binary_response: bool,
mutate_after_first_binary_response: Option<std::path::PathBuf>,
}
#[async_trait::async_trait]
impl workdir::WorkdirSession for ExternalValidatingSession {
fn workdir(&self) -> &workdir::Workdir {
self.inner.workdir()
}
fn capabilities(&self) -> workdir::WorkdirSessionCapabilities {
self.inner.capabilities()
}
async fn stat(
&self,
request: workdir::StatRequest,
) -> Result<workdir::StatResult, workdir::WorkdirError> {
self.inner.stat(request).await
}
async fn read(
&self,
request: workdir::ReadRequest,
) -> Result<workdir::ReadResult, workdir::WorkdirError> {
let operation = workdir::http::WorkdirSessionOperation::Read(request);
let operation = workdir::external::ExternalWorkdirOperation::try_from(operation)
.map_err(workdir::WorkdirError::UnsupportedOperation)?;
let workdir::http::WorkdirSessionOperation::Read(request) = operation.into_inner() else {
unreachable!("validated Read changed operation kind")
};
self.inner.read(request).await
}
async fn read_bytes(
&self,
request: workdir::ReadBytesRequest,
) -> Result<workdir::ReadBytesResult, workdir::WorkdirError> {
let operation = workdir::http::WorkdirSessionOperation::ReadBytes(request);
let operation = workdir::external::ExternalWorkdirOperation::try_from(operation)
.map_err(workdir::WorkdirError::UnsupportedOperation)?;
let workdir::http::WorkdirSessionOperation::ReadBytes(request) = operation.into_inner()
else {
unreachable!("validated ReadBytes changed operation kind")
};
let result = self.inner.read_bytes(request).await?;
let result = workdir::external::ExternalWorkdirOperationResult::try_from(
workdir::http::WorkdirSessionOperationResult::ReadBytes(result),
)
.map_err(workdir::WorkdirError::UnsupportedOperation)?;
let workdir::http::WorkdirSessionOperationResult::ReadBytes(mut result) =
result.into_inner()
else {
unreachable!("validated ReadBytes result changed operation kind")
};
if result.offset == 0 {
if let Some(path) = &self.mutate_after_first_binary_response {
std::fs::write(path, vec![b'z'; result.total_bytes as usize]).unwrap();
}
if self.truncate_first_binary_response {
result.bytes.pop();
}
}
Ok(result)
}
async fn write(
&self,
request: workdir::WriteRequest,
) -> Result<workdir::WriteResult, workdir::WorkdirError> {
self.inner.write(request).await
}
async fn edit(
&self,
request: workdir::EditRequest,
) -> Result<workdir::EditResult, workdir::WorkdirError> {
self.inner.edit(request).await
}
async fn list(
&self,
request: workdir::ListRequest,
) -> Result<workdir::ListResult, workdir::WorkdirError> {
self.inner.list(request).await
}
async fn glob(
&self,
request: workdir::GlobRequest,
) -> Result<workdir::GlobResult, workdir::WorkdirError> {
self.inner.glob(request).await
}
async fn grep(
&self,
request: workdir::GrepRequest,
) -> Result<workdir::GrepResult, workdir::WorkdirError> {
self.inner.grep(request).await
}
async fn start_command(
&self,
request: workdir::CommandRequest,
) -> Result<workdir::CommandHandle, workdir::WorkdirError> {
self.inner.start_command(request).await
}
async fn command_status(
&self,
handle: workdir::CommandHandle,
) -> Result<workdir::CommandStatus, workdir::WorkdirError> {
self.inner.command_status(handle).await
}
async fn command_output(
&self,
request: workdir::CommandOutputRequest,
) -> Result<workdir::CommandOutput, workdir::WorkdirError> {
self.inner.command_output(request).await
}
async fn cancel_command(
&self,
handle: workdir::CommandHandle,
) -> Result<(), workdir::WorkdirError> {
self.inner.cancel_command(handle).await
}
async fn close(&self) -> Result<(), workdir::WorkdirError> {
self.inner.close().await
}
}
#[tokio::test]
async fn external_view_image_uses_a_protocol_bounded_request() {
let dir = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let scope = scope_with_spill(dir.path(), spill.path());
let session: WorkdirSessionHandle = Arc::new(ExternalValidatingSession {
inner: LocalWorkdirSession::new(scope, dir.path().to_path_buf()),
truncate_first_binary_response: false,
mutate_after_first_binary_response: None,
});
let png = b"\x89PNG\r\n\x1a\nexternal\0image\nbody\xff";
std::fs::write(dir.path().join("image.png"), png).unwrap();
let (_meta, tool) = view_image_tool(session)();
let output = call(&tool, json!({ "path": "image.png" })).await;
let agen::tool::Attachment::Image(image) = &output.attachments[0];
assert_eq!(image.mime_type(), "image/png");
assert_eq!(image.data(), png);
}
#[tokio::test]
async fn external_view_image_rejects_truncated_and_cross_version_chunks() {
let dir = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let image_path = dir.path().join("image.png");
let scope = scope_with_spill(dir.path(), spill.path());
std::fs::write(&image_path, b"\x89PNG\r\n\x1a\npartial").unwrap();
let truncated: WorkdirSessionHandle = Arc::new(ExternalValidatingSession {
inner: LocalWorkdirSession::new(scope.clone(), dir.path().to_path_buf()),
truncate_first_binary_response: true,
mutate_after_first_binary_response: None,
});
let (_meta, tool) = view_image_tool(truncated)();
let error = call_err(&tool, json!({ "path": "image.png" }))
.await
.to_string();
assert!(
error.contains("truncated") || error.contains("partial"),
"{error}"
);
let mut large = vec![b'x'; 1024 * 1024 + 32];
large[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n");
std::fs::write(&image_path, large).unwrap();
let changing: WorkdirSessionHandle = Arc::new(ExternalValidatingSession {
inner: LocalWorkdirSession::new(scope, dir.path().to_path_buf()),
truncate_first_binary_response: false,
mutate_after_first_binary_response: Some(image_path),
});
let (_meta, tool) = view_image_tool(changing)();
let error = call_err(&tool, json!({ "path": "image.png" }))
.await
.to_string();
assert!(error.contains("changed"), "{error}");
}
#[tokio::test]
async fn view_image_enforces_the_ten_mib_boundary_without_partial_attachments() {
const LIMIT: usize = 10 * 1024 * 1024;
let dir = TempDir::new().unwrap();
let spill = TempDir::new().unwrap();
let scope = scope_with_spill(dir.path(), spill.path());
let session: WorkdirSessionHandle =
Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf()));
let (_meta, tool) = view_image_tool(session)();
let mut png = vec![0xff; LIMIT];
png[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n");
std::fs::write(dir.path().join("boundary.png"), &png).unwrap();
let output = call(&tool, json!({ "path": "boundary.png" })).await;
let agen::tool::Attachment::Image(image) = &output.attachments[0];
assert_eq!(image.data(), png);
png.push(0);
std::fs::write(dir.path().join("boundary.png"), png).unwrap();
let error = call_err(&tool, json!({ "path": "boundary.png" }))
.await
.to_string();
assert!(error.contains("10485760-byte limit"), "{error}");
}
#[tokio::test] #[tokio::test]
async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { async fn view_image_reads_scoped_bytes_into_durable_tool_detail() {
let dir = TempDir::new().unwrap(); let dir = TempDir::new().unwrap();
+61 -5
View File
@@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize};
use crate::http::{WorkdirSessionOperation, WorkdirSessionOperationResult, WorkdirTransportError}; use crate::http::{WorkdirSessionOperation, WorkdirSessionOperationResult, WorkdirTransportError};
use crate::{BoundedReadLimits, WorkdirId, WorkdirPath, WorkdirSessionCapabilities}; use crate::{BoundedReadLimits, WorkdirId, WorkdirPath, WorkdirSessionCapabilities};
pub const EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION: u16 = 1; pub const EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION: u16 = 2;
/// Maximum serialized Backend-to-provider command frame size. /// Maximum serialized Backend-to-provider command frame size.
pub const MAX_EXTERNAL_SERVER_FRAME_BYTES: usize = 1024 * 1024; pub const MAX_EXTERNAL_SERVER_FRAME_BYTES: usize = 1024 * 1024;
/// Maximum serialized provider-to-Backend result frame size. Read bytes are /// Maximum serialized provider-to-Backend result frame size. Read bytes are
@@ -155,6 +155,12 @@ fn validate_external_operation(operation: &WorkdirSessionOperation) -> Result<()
&& request.limit <= 1_000_000 && request.limit <= 1_000_000
&& request.max_bytes <= BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes && request.max_bytes <= BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes
} }
WorkdirSessionOperation::ReadBytes(request) => {
is_root_relative(&request.path)
&& request.offset <= BoundedReadLimits::EXTERNAL_DEFAULT.max_source_bytes
&& (1..=BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes)
.contains(&request.max_bytes)
}
WorkdirSessionOperation::List(request) => { WorkdirSessionOperation::List(request) => {
is_root_relative(&request.path) && request.limit <= MAX_EXTERNAL_RESULT_ITEMS is_root_relative(&request.path) && request.limit <= MAX_EXTERNAL_RESULT_ITEMS
} }
@@ -275,6 +281,15 @@ fn validate_external_result(result: &WorkdirSessionOperationResult) -> Result<()
is_root_relative(&result.path) is_root_relative(&result.path)
&& result.bytes.len() <= BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes && result.bytes.len() <= BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes
} }
WorkdirSessionOperationResult::ReadBytes(result) => {
let returned_end = result.offset.checked_add(result.bytes.len() as u64);
is_root_relative(&result.path)
&& result.bytes.len() <= BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes
&& result.total_bytes <= BoundedReadLimits::EXTERNAL_DEFAULT.max_source_bytes
&& returned_end.is_some_and(|end| {
end <= result.total_bytes && result.eof == (end == result.total_bytes)
})
}
WorkdirSessionOperationResult::List(result) => { WorkdirSessionOperationResult::List(result) => {
result.entries.len() <= MAX_EXTERNAL_RESULT_ITEMS result.entries.len() <= MAX_EXTERNAL_RESULT_ITEMS
&& result && result
@@ -482,8 +497,8 @@ mod tests {
use super::*; use super::*;
use crate::{ use crate::{
CommandRequest, EditRequest, GrepResult, ListRequest, ReadRequest, Workdir, WorkdirPath, CommandRequest, EditRequest, GrepResult, ListRequest, ReadBytesRequest, ReadBytesResult,
WriteRequest, ReadRequest, Workdir, WorkdirPath, WriteRequest,
}; };
fn assert_no_provider_authority(value: &Value) { fn assert_no_provider_authority(value: &Value) {
@@ -614,7 +629,7 @@ mod tests {
#[test] #[test]
fn unsupported_protocol_versions_and_invalid_read_limits_fail_closed() { fn unsupported_protocol_versions_and_invalid_read_limits_fail_closed() {
let future = r#"{ let future = r#"{
"version": 2, "version": 3,
"message": {"kind": "heartbeat", "generation": 1, "sequence": 1} "message": {"kind": "heartbeat", "generation": 1, "sequence": 1}
}"#; }"#;
assert!(serde_json::from_str::<ExternalWorkdirServerFrame>(future).is_err()); assert!(serde_json::from_str::<ExternalWorkdirServerFrame>(future).is_err());
@@ -726,10 +741,51 @@ mod tests {
); );
} }
#[test]
fn protocol_bounds_binary_read_requests_and_rejects_partial_results() {
let path = WorkdirPath::new("image.png").unwrap();
let bounded = WorkdirSessionOperation::ReadBytes(ReadBytesRequest {
path: path.clone(),
offset: 0,
max_bytes: BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes,
expected_hash: None,
});
assert!(ExternalWorkdirOperation::try_from(bounded).is_ok());
let oversized = WorkdirSessionOperation::ReadBytes(ReadBytesRequest {
path: path.clone(),
offset: 0,
max_bytes: BoundedReadLimits::EXTERNAL_DEFAULT.max_response_bytes + 1,
expected_hash: None,
});
assert!(ExternalWorkdirOperation::try_from(oversized).is_err());
let hash = [7; 32];
let complete = WorkdirSessionOperationResult::ReadBytes(ReadBytesResult {
path: path.clone(),
bytes: vec![1, 2, 3],
offset: 0,
total_bytes: 3,
content_hash: hash,
eof: true,
});
assert!(ExternalWorkdirOperationResult::try_from(complete).is_ok());
let partial_marked_complete = WorkdirSessionOperationResult::ReadBytes(ReadBytesResult {
path,
bytes: vec![1, 2],
offset: 0,
total_bytes: 3,
content_hash: hash,
eof: true,
});
assert!(ExternalWorkdirOperationResult::try_from(partial_marked_complete).is_err());
}
#[test] #[test]
fn protocol_rejects_unknown_nested_fields_and_unbounded_requests() { fn protocol_rejects_unknown_nested_fields_and_unbounded_requests() {
let unknown = r#"{ let unknown = r#"{
"version": 1, "version": 2,
"message": { "message": {
"kind": "heartbeat", "kind": "heartbeat",
"generation": 1, "generation": 1,
+22 -3
View File
@@ -11,9 +11,9 @@ use serde::{Deserialize, Serialize};
use crate::{ use crate::{
CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest,
EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult,
ReadRequest, ReadResult, StatRequest, StatResult, WorkdirError, WorkdirId, ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult,
WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, WorkdirError, WorkdirId, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest,
WorkdirSessionCapabilities, WriteRequest, WriteResult, WorkdirSession, WorkdirSessionCapabilities, WriteRequest, WriteResult,
}; };
/// Opaque Runtime-owned identifier for one ephemeral Workdir session. /// Opaque Runtime-owned identifier for one ephemeral Workdir session.
@@ -66,6 +66,7 @@ pub enum WorkdirSessionOperation {
ScopeRulesOverlap(WorkdirScopeOverlapRequest), ScopeRulesOverlap(WorkdirScopeOverlapRequest),
Stat(StatRequest), Stat(StatRequest),
Read(ReadRequest), Read(ReadRequest),
ReadBytes(ReadBytesRequest),
Write(WriteRequest), Write(WriteRequest),
Edit(EditRequest), Edit(EditRequest),
List(ListRequest), List(ListRequest),
@@ -97,6 +98,7 @@ pub enum WorkdirSessionOperationResult {
ScopeRulesOverlap { overlaps: bool }, ScopeRulesOverlap { overlaps: bool },
Stat(StatResult), Stat(StatResult),
Read(ReadResult), Read(ReadResult),
ReadBytes(ReadBytesResult),
Write(WriteResult), Write(WriteResult),
Edit(EditResult), Edit(EditResult),
List(ListResult), List(ListResult),
@@ -135,6 +137,10 @@ pub async fn dispatch_workdir_session_operation(
.read(request) .read(request)
.await .await
.map(WorkdirSessionOperationResult::Read), .map(WorkdirSessionOperationResult::Read),
WorkdirSessionOperation::ReadBytes(request) => session
.read_bytes(request)
.await
.map(WorkdirSessionOperationResult::ReadBytes),
WorkdirSessionOperation::Write(request) => session WorkdirSessionOperation::Write(request) => session
.write(request) .write(request)
.await .await
@@ -572,6 +578,19 @@ mod client {
} }
} }
async fn read_bytes(
&self,
request: ReadBytesRequest,
) -> Result<ReadBytesResult, WorkdirError> {
match self
.operate(WorkdirSessionOperation::ReadBytes(request))
.await?
{
WorkdirSessionOperationResult::ReadBytes(result) => Ok(result),
_ => Err(Self::mismatch("read_bytes")),
}
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> { async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
match self match self
.operate(WorkdirSessionOperation::Write(request)) .operate(WorkdirSessionOperation::Write(request))
+10 -1
View File
@@ -27,7 +27,8 @@ use tokio::sync::broadcast;
pub use fs_operation::{ pub use fs_operation::{
BoundedReadLimits, ContentHash, EditRequest, EditResult, EntryKind, FsPath as WorkdirPath, BoundedReadLimits, ContentHash, EditRequest, EditResult, EntryKind, FsPath as WorkdirPath,
GlobRequest, GlobResult, GrepOutputMode, GrepRequest, GrepResult, ListEntry, ListRequest, GlobRequest, GlobResult, GrepOutputMode, GrepRequest, GrepResult, ListEntry, ListRequest,
ListResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult, ListResult, ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest,
StatResult, WriteRequest, WriteResult,
}; };
pub use http::dispatch_workdir_session_operation; pub use http::dispatch_workdir_session_operation;
pub use local::{ pub use local::{
@@ -220,6 +221,14 @@ pub trait WorkdirSession: std::fmt::Debug + Send + Sync {
async fn stat(&self, request: StatRequest) -> Result<StatResult, WorkdirError>; async fn stat(&self, request: StatRequest) -> Result<StatResult, WorkdirError>;
async fn read(&self, request: ReadRequest) -> Result<ReadResult, WorkdirError>; async fn read(&self, request: ReadRequest) -> Result<ReadResult, WorkdirError>;
async fn read_bytes(
&self,
_request: ReadBytesRequest,
) -> Result<ReadBytesResult, WorkdirError> {
Err(WorkdirError::UnsupportedOperation(
"bounded binary reads are not supported by this provider".to_string(),
))
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError>; async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError>;
async fn edit(&self, request: EditRequest) -> Result<EditResult, WorkdirError>; async fn edit(&self, request: EditRequest) -> Result<EditResult, WorkdirError>;
async fn list(&self, request: ListRequest) -> Result<ListResult, WorkdirError>; async fn list(&self, request: ListRequest) -> Result<ListResult, WorkdirError>;
+19 -4
View File
@@ -30,10 +30,10 @@ use crate::{
BoundedReadLimits, CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, BoundedReadLimits, CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest,
CommandRequest, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, EditRequest, CommandRequest, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, EditRequest,
EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult,
ReadRequest, ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirPath, ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir,
WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, WorkdirError, WorkdirPath, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest,
WorkdirSessionCapabilities, WorkdirSessionCapability, WorkdirToolScopePermission, WriteRequest, WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionCapability,
WriteResult, WorkdirToolScopePermission, WriteRequest, WriteResult,
}; };
#[cfg(test)] #[cfg(test)]
use crate::{EntryKind, WriteOutcome}; use crate::{EntryKind, WriteOutcome};
@@ -1261,6 +1261,21 @@ impl WorkdirSession for LocalWorkdirSession {
.map_err(|error| sanitize_error(error, &logical)) .map_err(|error| sanitize_error(error, &logical))
} }
async fn read_bytes(&self, request: ReadBytesRequest) -> Result<ReadBytesResult, WorkdirError> {
self.ensure_capability(WorkdirSessionCapability::Read)?;
let logical = request.path.clone();
self.validate_operation_path(&request.path)?;
let access = self.scope_access();
match self.inner.read_limits {
Some(limits) => {
fs_operation::run_read_bytes_bounded(&self.inner.root, request, &access, limits)
}
None => fs_operation::run_read_bytes(&self.inner.root, request, &access),
}
.map_err(WorkdirError::from)
.map_err(|error| sanitize_error(error, &logical))
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> { async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
self.ensure_capability(WorkdirSessionCapability::Write)?; self.ensure_capability(WorkdirSessionCapability::Write)?;
let logical = request.path.clone(); let logical = request.path.clone();
+9 -3
View File
@@ -16,9 +16,10 @@ use tokio::sync::broadcast;
use crate::{ use crate::{
CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest,
CommandSnapshot, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, CommandSnapshot, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest,
GrepResult, ListRequest, ListResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir, GrepResult, ListRequest, ListResult, ReadBytesRequest, ReadBytesResult, ReadRequest,
WorkdirError, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirScopeAuthorizationRequest,
WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult, WorkdirScopeOverlapRequest, WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle,
WriteRequest, WriteResult,
}; };
/// Stable Worker-local routing key for one attached Workdir. /// Stable Worker-local routing key for one attached Workdir.
@@ -257,6 +258,11 @@ impl WorkdirSession for RoutedWorkdirSession {
self.state.session.read(request).await self.state.session.read(request).await
} }
async fn read_bytes(&self, request: ReadBytesRequest) -> Result<ReadBytesResult, WorkdirError> {
let _active = self.state.enter()?;
self.state.session.read_bytes(request).await
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> { async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
let _active = self.state.enter()?; let _active = self.state.enter()?;
self.state.session.write(request).await self.state.session.write(request).await
+35 -1
View File
@@ -6,7 +6,8 @@ use std::sync::{Arc, Mutex, Weak};
use async_trait::async_trait; use async_trait::async_trait;
use fs_operation::{ use fs_operation::{
EditRequest, EditResult, FsPath, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, EditRequest, EditResult, FsPath, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest,
ListResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult, ListResult, ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest,
StatResult, WriteRequest, WriteResult,
}; };
use manifest::SymlinkPolicy; use manifest::SymlinkPolicy;
use schemars::JsonSchema; use schemars::JsonSchema;
@@ -1011,6 +1012,18 @@ impl WorkdirSession for ScopedWorkdirSession {
self.source.read(request).await self.source.read(request).await
} }
async fn read_bytes(
&self,
mut request: ReadBytesRequest,
) -> Result<ReadBytesResult, WorkdirError> {
let path = self
.resolve_operation_path(&request.path, WorkdirToolScopePermission::Read)
.await?;
self.ensure_read(&path, WorkdirSessionCapability::Read)?;
request.path = path;
self.source.read_bytes(request).await
}
async fn write(&self, mut request: WriteRequest) -> Result<WriteResult, WorkdirError> { async fn write(&self, mut request: WriteRequest) -> Result<WriteResult, WorkdirError> {
let _scope_guard = self.scope_lock.lock().await; let _scope_guard = self.scope_lock.lock().await;
let path = self let path = self
@@ -1276,6 +1289,10 @@ impl WorkdirSession for ReadOnlyWorkdirSession {
self.inner.read(request).await self.inner.read(request).await
} }
async fn read_bytes(&self, request: ReadBytesRequest) -> Result<ReadBytesResult, WorkdirError> {
self.inner.read_bytes(request).await
}
async fn write(&self, _request: WriteRequest) -> Result<WriteResult, WorkdirError> { async fn write(&self, _request: WriteRequest) -> Result<WriteResult, WorkdirError> {
Err(WorkdirError::Denied("read-only workdir session".into())) Err(WorkdirError::Denied("read-only workdir session".into()))
} }
@@ -1656,6 +1673,15 @@ mod tests {
} }
} }
fn read_bytes(path: &str) -> ReadBytesRequest {
ReadBytesRequest {
path: fs_path(path),
offset: 0,
max_bytes: 1024,
expected_hash: None,
}
}
fn write(path: &str, content: &str) -> WriteRequest { fn write(path: &str, content: &str) -> WriteRequest {
WriteRequest { WriteRequest {
path: fs_path(path), path: fs_path(path),
@@ -1887,6 +1913,14 @@ mod tests {
child.read(read("readme.md")).await.unwrap().bytes, child.read(read("readme.md")).await.unwrap().bytes,
b"visible" b"visible"
); );
assert_eq!(
child
.read_bytes(read_bytes("readme.md"))
.await
.unwrap()
.bytes,
b"visible"
);
assert!(matches!( assert!(matches!(
child.write(write("new.md", "no")).await, child.write(write("new.md", "no")).await,
Err(WorkdirError::Denied(_)) Err(WorkdirError::Denied(_))
@@ -18,8 +18,8 @@ use workdir::workspace::WorkspaceWorkdirSessionOperationRequest;
use workdir::{ use workdir::{
CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest,
EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult,
ReadRequest, ReadResult, StatRequest, StatResult, Workdir, WorkdirError, ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir,
WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, WorkdirError, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession,
WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult, WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult,
}; };
@@ -357,6 +357,13 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession {
} }
} }
async fn read_bytes(&self, request: ReadBytesRequest) -> Result<ReadBytesResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::ReadBytes(request))? {
WorkdirSessionOperationResult::ReadBytes(result) => Ok(result),
_ => Err(Self::mismatch("read_bytes")),
}
}
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> { async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
match self.operate(WorkdirSessionOperation::Write(request))? { match self.operate(WorkdirSessionOperation::Write(request))? {
WorkdirSessionOperationResult::Write(result) => Ok(result), WorkdirSessionOperationResult::Write(result) => Ok(result),
+77 -1
View File
@@ -1074,6 +1074,7 @@ fn external_workdir_operation_kind(operation: &WorkdirSessionOperation) -> &'sta
WorkdirSessionOperation::ScopeRulesOverlap(_) => "scope_rules_overlap", WorkdirSessionOperation::ScopeRulesOverlap(_) => "scope_rules_overlap",
WorkdirSessionOperation::Stat(_) => "stat", WorkdirSessionOperation::Stat(_) => "stat",
WorkdirSessionOperation::Read(_) => "read", WorkdirSessionOperation::Read(_) => "read",
WorkdirSessionOperation::ReadBytes(_) => "read_bytes",
WorkdirSessionOperation::Write(_) => "write", WorkdirSessionOperation::Write(_) => "write",
WorkdirSessionOperation::Edit(_) => "edit", WorkdirSessionOperation::Edit(_) => "edit",
WorkdirSessionOperation::List(_) => "list", WorkdirSessionOperation::List(_) => "list",
@@ -1141,6 +1142,19 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession {
_ => Err(Self::mismatch("read")), _ => Err(Self::mismatch("read")),
} }
} }
async fn read_bytes(
&self,
request: workdir::ReadBytesRequest,
) -> std::result::Result<workdir::ReadBytesResult, workdir::WorkdirError> {
self.ensure_capability(workdir::WorkdirSessionCapability::Read)?;
match self
.operate(WorkdirSessionOperation::ReadBytes(request))
.await?
{
WorkdirSessionOperationResult::ReadBytes(value) => Ok(value),
_ => Err(Self::mismatch("read_bytes")),
}
}
async fn write( async fn write(
&self, &self,
request: workdir::WriteRequest, request: workdir::WriteRequest,
@@ -14824,6 +14838,7 @@ async fn scoped_execute_current_worker_workdir_operation(
| WorkdirSessionOperation::ScopeRulesOverlap(_) | WorkdirSessionOperation::ScopeRulesOverlap(_)
| WorkdirSessionOperation::Stat(_) | WorkdirSessionOperation::Stat(_)
| WorkdirSessionOperation::Read(_) | WorkdirSessionOperation::Read(_)
| WorkdirSessionOperation::ReadBytes(_)
| WorkdirSessionOperation::Write(_) | WorkdirSessionOperation::Write(_)
| WorkdirSessionOperation::Edit(_) | WorkdirSessionOperation::Edit(_)
| WorkdirSessionOperation::List(_) | WorkdirSessionOperation::List(_)
@@ -14890,6 +14905,10 @@ async fn execute_workdir_session_operation(
.read(request) .read(request)
.await .await
.map(WorkdirSessionOperationResult::Read), .map(WorkdirSessionOperationResult::Read),
WorkdirSessionOperation::ReadBytes(request) => session
.read_bytes(request)
.await
.map(WorkdirSessionOperationResult::ReadBytes),
WorkdirSessionOperation::Write(request) => session WorkdirSessionOperation::Write(request) => session
.write(request) .write(request)
.await .await
@@ -17642,7 +17661,7 @@ async fn serve_external_workdir_provider(
Ok(operation) => operation, Ok(operation) => operation,
Err(message) => { Err(message) => {
let _ = response.send(Err(WorkdirTransportError { let _ = response.send(Err(WorkdirTransportError {
code: workdir::http::WorkdirTransportErrorCode::Unsupported, code: workdir::http::WorkdirTransportErrorCode::InvalidRequest,
message, message,
})); }));
continue; continue;
@@ -28109,6 +28128,8 @@ mod tests {
); );
let provider_root = tempfile::tempdir().unwrap(); let provider_root = tempfile::tempdir().unwrap();
fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap(); fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap();
let image_bytes = b"\x89PNG\r\n\x1a\nexternal\0image\n\xff";
fs::write(provider_root.path().join("image.png"), image_bytes).unwrap();
let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap(); let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap();
let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned( let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned(
workdir::Workdir::new(&grant.working_directory_id), workdir::Workdir::new(&grant.working_directory_id),
@@ -28178,6 +28199,61 @@ mod tests {
let read_result = read.await.unwrap().unwrap(); let read_result = read.await.unwrap().unwrap();
assert_eq!(read_result.bytes, b"one\ntwo\n"); assert_eq!(read_result.bytes, b"one\ntwo\n");
let read_bytes = {
let broker_session = broker_session.clone();
tokio::spawn(async move {
broker_session
.read_bytes(workdir::ReadBytesRequest {
path: workdir::WorkdirPath::new("image.png").unwrap(),
offset: 0,
max_bytes: 1024,
expected_hash: None,
})
.await
})
};
let (operation_id, operation) = loop {
let text = second.next().await.unwrap().unwrap().into_text().unwrap();
let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
if let ExternalWorkdirServerMessage::Operation {
generation: 2,
operation_id,
operation,
} = frame.message
{
break (operation_id, operation);
}
};
assert!(matches!(
operation.as_inner(),
WorkdirSessionOperation::ReadBytes(_)
));
let result = workdir::http::dispatch_workdir_session_operation(
&local_session,
operation.into_inner(),
)
.await
.unwrap()
.try_into()
.unwrap();
second
.send(Message::Text(
serde_json::to_string(&ExternalWorkdirProviderFrame::current(
ExternalWorkdirProviderMessage::OperationResult {
generation: 2,
operation_id,
outcome: ExternalWorkdirOperationOutcome::Completed { result },
},
))
.unwrap()
.into(),
))
.await
.unwrap();
let read_bytes_result = read_bytes.await.unwrap().unwrap();
assert_eq!(read_bytes_result.bytes, image_bytes);
assert!(read_bytes_result.eof);
let edit = { let edit = {
let broker_session = broker_session.clone(); let broker_session = broker_session.clone();
let expected_hash = read_result.content_hash; let expected_hash = read_result.content_hash;
+1 -1
View File
@@ -903,7 +903,7 @@ async fn execute_operation(
Ok(result) => ExternalWorkdirOperationOutcome::Completed { result }, Ok(result) => ExternalWorkdirOperationOutcome::Completed { result },
Err(message) => ExternalWorkdirOperationOutcome::Failed { Err(message) => ExternalWorkdirOperationOutcome::Failed {
error: ExternalWorkdirOperationError::from_transport_error(WorkdirTransportError { error: ExternalWorkdirOperationError::from_transport_error(WorkdirTransportError {
code: workdir::http::WorkdirTransportErrorCode::Unsupported, code: workdir::http::WorkdirTransportErrorCode::InvalidRequest,
message, message,
}), }),
}, },