From 05d6da45fbf7fd8e1300e9f296fed4a413d20a4e Mon Sep 17 00:00:00 2001 From: Hare Date: Wed, 30 Sep 2026 07:44:59 +0900 Subject: [PATCH 1/3] fix: add bounded binary reads for ViewImage --- crates/fs-operation/src/lib.rs | 100 +++++++- crates/fs-operation/src/local.rs | 136 ++++++++++- crates/fs-operation/src/operation.rs | 30 +++ crates/tools/src/view_image.rs | 113 +++++++-- crates/tools/tests/integration.rs | 218 ++++++++++++++++++ crates/workdir/src/external.rs | 66 +++++- crates/workdir/src/http.rs | 25 +- crates/workdir/src/lib.rs | 11 +- crates/workdir/src/local.rs | 23 +- crates/workdir/src/router.rs | 12 +- crates/workdir/src/scope.rs | 36 ++- .../src/feature/builtin/manage_workdir.rs | 11 +- crates/workspace-server/src/server.rs | 78 ++++++- crates/yoi/src/workdir_share.rs | 2 +- 14 files changed, 816 insertions(+), 45 deletions(-) diff --git a/crates/fs-operation/src/lib.rs b/crates/fs-operation/src/lib.rs index 465e3b6b9..7fa897df8 100644 --- a/crates/fs-operation/src/lib.rs +++ b/crates/fs-operation/src/lib.rs @@ -15,7 +15,8 @@ use thiserror::Error; pub use glob::run_glob; 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 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] fn cancellation_stops_a_read_already_in_progress() { use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; diff --git a/crates/fs-operation/src/local.rs b/crates/fs-operation/src/local.rs index aaff61910..34f60b92d 100644 --- a/crates/fs-operation/src/local.rs +++ b/crates/fs-operation/src/local.rs @@ -7,8 +7,9 @@ use sha2::{Digest, Sha256}; use crate::{ AtomicWriteMode, BoundedReadLimits, ContentHash, EditRequest, EditResult, EntryKind, - FsAccessPolicy, FsError, FsPath, ListEntry, ListRequest, ListResult, ReadRequest, ReadResult, - StatRequest, StatResult, WriteRequest, WriteResult, direct_symlink, + FsAccessPolicy, FsError, FsPath, ListEntry, ListRequest, ListResult, ReadBytesRequest, + ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, WriteRequest, WriteResult, + direct_symlink, }; 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 { + 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 { + 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, +) -> Result { + 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( root: &Path, request: WriteRequest, diff --git a/crates/fs-operation/src/operation.rs b/crates/fs-operation/src/operation.rs index f5765b6f6..ddb309a8d 100644 --- a/crates/fs-operation/src/operation.rs +++ b/crates/fs-operation/src/operation.rs @@ -207,6 +207,36 @@ pub struct ReadResult { 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, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct ReadBytesResult { + pub path: FsPath, + pub bytes: Vec, + #[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)] #[serde(deny_unknown_fields)] pub struct WriteRequest { diff --git a/crates/tools/src/view_image.rs b/crates/tools/src/view_image.rs index 1fc86ca44..255695b9a 100644 --- a/crates/tools/src/view_image.rs +++ b/crates/tools/src/view_image.rs @@ -7,12 +7,14 @@ use agen::tool::{ }; use async_trait::async_trait; use serde::Deserialize; -use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; +use workdir::{ReadBytesRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter}; use crate::error::ToolsError; /// Maximum image body accepted for one model request. 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. \ 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, )?; let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?; - let result = selected - .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, - }) + let bytes = read_image_bytes(selected.session.as_ref(), path.clone()) .await .map_err(ToolsError::from)?; - - 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(|| { + let mime_type = detect_image_mime(&bytes).ok_or_else(|| { ToolError::InvalidArgument( "unsupported image; expected PNG, JPEG, GIF, or WebP bytes".to_string(), ) })?; - let bytes = result.bytes.len(); + let byte_count = bytes.len(); Ok(ToolOutput { - summary: format!("Attached image {path} ({mime_type}, {bytes} bytes)"), + summary: format!("Attached image {path} ({mime_type}, {byte_count} bytes)"), content: None, attachments: vec![Attachment::Image(ImageAttachment::new( mime_type, - Arc::<[u8]>::from(result.bytes), + Arc::<[u8]>::from(bytes), ))], }) } } +async fn read_image_bytes( + session: &dyn workdir::WorkdirSession, + path: WorkdirPath, +) -> Result, 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 { routed_view_image_tool(crate::routing::singleton_router(session)) } diff --git a/crates/tools/tests/integration.rs b/crates/tools/tests/integration.rs index f1d346e60..eb4ca3c7f 100644 --- a/crates/tools/tests/integration.rs +++ b/crates/tools/tests/integration.rs @@ -399,6 +399,224 @@ async fn detach_and_reattach_same_alias_requires_a_fresh_read() { router.close_all().await.unwrap(); } +#[derive(Debug)] +struct ExternalValidatingSession { + inner: LocalWorkdirSession, + truncate_first_binary_response: bool, + mutate_after_first_binary_response: Option, +} + +#[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 { + self.inner.stat(request).await + } + + async fn read( + &self, + request: workdir::ReadRequest, + ) -> Result { + 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 { + 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 { + self.inner.write(request).await + } + + async fn edit( + &self, + request: workdir::EditRequest, + ) -> Result { + self.inner.edit(request).await + } + + async fn list( + &self, + request: workdir::ListRequest, + ) -> Result { + self.inner.list(request).await + } + + async fn glob( + &self, + request: workdir::GlobRequest, + ) -> Result { + self.inner.glob(request).await + } + + async fn grep( + &self, + request: workdir::GrepRequest, + ) -> Result { + self.inner.grep(request).await + } + + async fn start_command( + &self, + request: workdir::CommandRequest, + ) -> Result { + self.inner.start_command(request).await + } + + async fn command_status( + &self, + handle: workdir::CommandHandle, + ) -> Result { + self.inner.command_status(handle).await + } + + async fn command_output( + &self, + request: workdir::CommandOutputRequest, + ) -> Result { + 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] async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { let dir = TempDir::new().unwrap(); diff --git a/crates/workdir/src/external.rs b/crates/workdir/src/external.rs index 5557231ff..c94582867 100644 --- a/crates/workdir/src/external.rs +++ b/crates/workdir/src/external.rs @@ -12,7 +12,7 @@ use serde::{Deserialize, Serialize}; use crate::http::{WorkdirSessionOperation, WorkdirSessionOperationResult, WorkdirTransportError}; 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. pub const MAX_EXTERNAL_SERVER_FRAME_BYTES: usize = 1024 * 1024; /// 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.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) => { 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) && 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) => { result.entries.len() <= MAX_EXTERNAL_RESULT_ITEMS && result @@ -482,8 +497,8 @@ mod tests { use super::*; use crate::{ - CommandRequest, EditRequest, GrepResult, ListRequest, ReadRequest, Workdir, WorkdirPath, - WriteRequest, + CommandRequest, EditRequest, GrepResult, ListRequest, ReadBytesRequest, ReadBytesResult, + ReadRequest, Workdir, WorkdirPath, WriteRequest, }; fn assert_no_provider_authority(value: &Value) { @@ -614,7 +629,7 @@ mod tests { #[test] fn unsupported_protocol_versions_and_invalid_read_limits_fail_closed() { let future = r#"{ - "version": 2, + "version": 3, "message": {"kind": "heartbeat", "generation": 1, "sequence": 1} }"#; assert!(serde_json::from_str::(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] fn protocol_rejects_unknown_nested_fields_and_unbounded_requests() { let unknown = r#"{ - "version": 1, + "version": 2, "message": { "kind": "heartbeat", "generation": 1, diff --git a/crates/workdir/src/http.rs b/crates/workdir/src/http.rs index 0a3f7e8aa..4017056a0 100644 --- a/crates/workdir/src/http.rs +++ b/crates/workdir/src/http.rs @@ -11,9 +11,9 @@ use serde::{Deserialize, Serialize}; use crate::{ CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, - ReadRequest, ReadResult, StatRequest, StatResult, WorkdirError, WorkdirId, - WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, - WorkdirSessionCapabilities, WriteRequest, WriteResult, + ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, + WorkdirError, WorkdirId, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, + WorkdirSession, WorkdirSessionCapabilities, WriteRequest, WriteResult, }; /// Opaque Runtime-owned identifier for one ephemeral Workdir session. @@ -66,6 +66,7 @@ pub enum WorkdirSessionOperation { ScopeRulesOverlap(WorkdirScopeOverlapRequest), Stat(StatRequest), Read(ReadRequest), + ReadBytes(ReadBytesRequest), Write(WriteRequest), Edit(EditRequest), List(ListRequest), @@ -97,6 +98,7 @@ pub enum WorkdirSessionOperationResult { ScopeRulesOverlap { overlaps: bool }, Stat(StatResult), Read(ReadResult), + ReadBytes(ReadBytesResult), Write(WriteResult), Edit(EditResult), List(ListResult), @@ -135,6 +137,10 @@ pub async fn dispatch_workdir_session_operation( .read(request) .await .map(WorkdirSessionOperationResult::Read), + WorkdirSessionOperation::ReadBytes(request) => session + .read_bytes(request) + .await + .map(WorkdirSessionOperationResult::ReadBytes), WorkdirSessionOperation::Write(request) => session .write(request) .await @@ -572,6 +578,19 @@ mod client { } } + async fn read_bytes( + &self, + request: ReadBytesRequest, + ) -> Result { + match self + .operate(WorkdirSessionOperation::ReadBytes(request)) + .await? + { + WorkdirSessionOperationResult::ReadBytes(result) => Ok(result), + _ => Err(Self::mismatch("read_bytes")), + } + } + async fn write(&self, request: WriteRequest) -> Result { match self .operate(WorkdirSessionOperation::Write(request)) diff --git a/crates/workdir/src/lib.rs b/crates/workdir/src/lib.rs index 806a1431a..a390b1084 100644 --- a/crates/workdir/src/lib.rs +++ b/crates/workdir/src/lib.rs @@ -27,7 +27,8 @@ use tokio::sync::broadcast; pub use fs_operation::{ BoundedReadLimits, ContentHash, EditRequest, EditResult, EntryKind, FsPath as WorkdirPath, 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 local::{ @@ -220,6 +221,14 @@ pub trait WorkdirSession: std::fmt::Debug + Send + Sync { async fn stat(&self, request: StatRequest) -> Result; async fn read(&self, request: ReadRequest) -> Result; + async fn read_bytes( + &self, + _request: ReadBytesRequest, + ) -> Result { + Err(WorkdirError::UnsupportedOperation( + "bounded binary reads are not supported by this provider".to_string(), + )) + } async fn write(&self, request: WriteRequest) -> Result; async fn edit(&self, request: EditRequest) -> Result; async fn list(&self, request: ListRequest) -> Result; diff --git a/crates/workdir/src/local.rs b/crates/workdir/src/local.rs index 4a4a3b61e..33850729b 100644 --- a/crates/workdir/src/local.rs +++ b/crates/workdir/src/local.rs @@ -30,10 +30,10 @@ use crate::{ BoundedReadLimits, CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandSnapshot, CommandStatus, CommandStream, CommandStreamSlice, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, - ReadRequest, ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirPath, - WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, - WorkdirSessionCapabilities, WorkdirSessionCapability, WorkdirToolScopePermission, WriteRequest, - WriteResult, + ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir, + WorkdirError, WorkdirPath, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, + WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionCapability, + WorkdirToolScopePermission, WriteRequest, WriteResult, }; #[cfg(test)] use crate::{EntryKind, WriteOutcome}; @@ -1261,6 +1261,21 @@ impl WorkdirSession for LocalWorkdirSession { .map_err(|error| sanitize_error(error, &logical)) } + async fn read_bytes(&self, request: ReadBytesRequest) -> Result { + 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 { self.ensure_capability(WorkdirSessionCapability::Write)?; let logical = request.path.clone(); diff --git a/crates/workdir/src/router.rs b/crates/workdir/src/router.rs index eed9a9e50..270f0ec90 100644 --- a/crates/workdir/src/router.rs +++ b/crates/workdir/src/router.rs @@ -16,9 +16,10 @@ use tokio::sync::broadcast; use crate::{ CommandEvent, CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandSnapshot, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, - GrepResult, ListRequest, ListResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir, - WorkdirError, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, - WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult, + GrepResult, ListRequest, ListResult, ReadBytesRequest, ReadBytesResult, ReadRequest, + ReadResult, StatRequest, StatResult, Workdir, WorkdirError, WorkdirScopeAuthorizationRequest, + WorkdirScopeOverlapRequest, WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle, + WriteRequest, WriteResult, }; /// Stable Worker-local routing key for one attached Workdir. @@ -257,6 +258,11 @@ impl WorkdirSession for RoutedWorkdirSession { self.state.session.read(request).await } + async fn read_bytes(&self, request: ReadBytesRequest) -> Result { + let _active = self.state.enter()?; + self.state.session.read_bytes(request).await + } + async fn write(&self, request: WriteRequest) -> Result { let _active = self.state.enter()?; self.state.session.write(request).await diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index e76f59125..efe773d9f 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -6,7 +6,8 @@ use std::sync::{Arc, Mutex, Weak}; use async_trait::async_trait; use fs_operation::{ 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 schemars::JsonSchema; @@ -1011,6 +1012,18 @@ impl WorkdirSession for ScopedWorkdirSession { self.source.read(request).await } + async fn read_bytes( + &self, + mut request: ReadBytesRequest, + ) -> Result { + 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 { let _scope_guard = self.scope_lock.lock().await; let path = self @@ -1276,6 +1289,10 @@ impl WorkdirSession for ReadOnlyWorkdirSession { self.inner.read(request).await } + async fn read_bytes(&self, request: ReadBytesRequest) -> Result { + self.inner.read_bytes(request).await + } + async fn write(&self, _request: WriteRequest) -> Result { 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 { WriteRequest { path: fs_path(path), @@ -1887,6 +1913,14 @@ mod tests { child.read(read("readme.md")).await.unwrap().bytes, b"visible" ); + assert_eq!( + child + .read_bytes(read_bytes("readme.md")) + .await + .unwrap() + .bytes, + b"visible" + ); assert!(matches!( child.write(write("new.md", "no")).await, Err(WorkdirError::Denied(_)) diff --git a/crates/worker/src/feature/builtin/manage_workdir.rs b/crates/worker/src/feature/builtin/manage_workdir.rs index 517f7c236..54514dd75 100644 --- a/crates/worker/src/feature/builtin/manage_workdir.rs +++ b/crates/worker/src/feature/builtin/manage_workdir.rs @@ -18,8 +18,8 @@ use workdir::workspace::WorkspaceWorkdirSessionOperationRequest; use workdir::{ CommandHandle, CommandOutput, CommandOutputRequest, CommandRequest, CommandStatus, EditRequest, EditResult, GlobRequest, GlobResult, GrepRequest, GrepResult, ListRequest, ListResult, - ReadRequest, ReadResult, StatRequest, StatResult, Workdir, WorkdirError, - WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, + ReadBytesRequest, ReadBytesResult, ReadRequest, ReadResult, StatRequest, StatResult, Workdir, + WorkdirError, WorkdirScopeAuthorizationRequest, WorkdirScopeOverlapRequest, WorkdirSession, WorkdirSessionCapabilities, WorkdirSessionHandle, WriteRequest, WriteResult, }; @@ -357,6 +357,13 @@ impl WorkdirSession for WorkspaceAttachedWorkdirSession { } } + async fn read_bytes(&self, request: ReadBytesRequest) -> Result { + match self.operate(WorkdirSessionOperation::ReadBytes(request))? { + WorkdirSessionOperationResult::ReadBytes(result) => Ok(result), + _ => Err(Self::mismatch("read_bytes")), + } + } + async fn write(&self, request: WriteRequest) -> Result { match self.operate(WorkdirSessionOperation::Write(request))? { WorkdirSessionOperationResult::Write(result) => Ok(result), diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 31649fcb3..35f37069b 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -1074,6 +1074,7 @@ fn external_workdir_operation_kind(operation: &WorkdirSessionOperation) -> &'sta WorkdirSessionOperation::ScopeRulesOverlap(_) => "scope_rules_overlap", WorkdirSessionOperation::Stat(_) => "stat", WorkdirSessionOperation::Read(_) => "read", + WorkdirSessionOperation::ReadBytes(_) => "read_bytes", WorkdirSessionOperation::Write(_) => "write", WorkdirSessionOperation::Edit(_) => "edit", WorkdirSessionOperation::List(_) => "list", @@ -1141,6 +1142,19 @@ impl workdir::WorkdirSession for ExternalProviderWorkdirSession { _ => Err(Self::mismatch("read")), } } + async fn read_bytes( + &self, + request: workdir::ReadBytesRequest, + ) -> std::result::Result { + 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( &self, request: workdir::WriteRequest, @@ -14824,6 +14838,7 @@ async fn scoped_execute_current_worker_workdir_operation( | WorkdirSessionOperation::ScopeRulesOverlap(_) | WorkdirSessionOperation::Stat(_) | WorkdirSessionOperation::Read(_) + | WorkdirSessionOperation::ReadBytes(_) | WorkdirSessionOperation::Write(_) | WorkdirSessionOperation::Edit(_) | WorkdirSessionOperation::List(_) @@ -14890,6 +14905,10 @@ async fn execute_workdir_session_operation( .read(request) .await .map(WorkdirSessionOperationResult::Read), + WorkdirSessionOperation::ReadBytes(request) => session + .read_bytes(request) + .await + .map(WorkdirSessionOperationResult::ReadBytes), WorkdirSessionOperation::Write(request) => session .write(request) .await @@ -17642,7 +17661,7 @@ async fn serve_external_workdir_provider( Ok(operation) => operation, Err(message) => { let _ = response.send(Err(WorkdirTransportError { - code: workdir::http::WorkdirTransportErrorCode::Unsupported, + code: workdir::http::WorkdirTransportErrorCode::InvalidRequest, message, })); continue; @@ -28109,6 +28128,8 @@ mod tests { ); let provider_root = tempfile::tempdir().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 local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned( workdir::Workdir::new(&grant.working_directory_id), @@ -28178,6 +28199,61 @@ mod tests { let read_result = read.await.unwrap().unwrap(); 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::(&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 broker_session = broker_session.clone(); let expected_hash = read_result.content_hash; diff --git a/crates/yoi/src/workdir_share.rs b/crates/yoi/src/workdir_share.rs index 2330f6419..a2d3bc6b2 100644 --- a/crates/yoi/src/workdir_share.rs +++ b/crates/yoi/src/workdir_share.rs @@ -903,7 +903,7 @@ async fn execute_operation( Ok(result) => ExternalWorkdirOperationOutcome::Completed { result }, Err(message) => ExternalWorkdirOperationOutcome::Failed { error: ExternalWorkdirOperationError::from_transport_error(WorkdirTransportError { - code: workdir::http::WorkdirTransportErrorCode::Unsupported, + code: workdir::http::WorkdirTransportErrorCode::InvalidRequest, message, }), }, From 4e8a7fc1180e61fe9be92952221e4766016906e9 Mon Sep 17 00:00:00 2001 From: Hare Date: Wed, 30 Sep 2026 08:16:06 +0900 Subject: [PATCH 2/3] fix: address ViewImage external review findings --- Cargo.lock | 2 + .../tools/tests/fixtures/view-image/valid.gif | Bin 0 -> 828 bytes .../tools/tests/fixtures/view-image/valid.jpg | Bin 0 -> 222 bytes .../tools/tests/fixtures/view-image/valid.png | Bin 0 -> 94 bytes .../tests/fixtures/view-image/valid.webp | Bin 0 -> 64 bytes crates/tools/tests/integration.rs | 137 ++++++++-- crates/workdir/src/scope.rs | 16 +- crates/workspace-server/Cargo.toml | 2 + crates/workspace-server/src/server.rs | 254 +++++++++++++++--- crates/yoi/src/workdir_share.rs | 63 ++++- 10 files changed, 407 insertions(+), 67 deletions(-) create mode 100644 crates/tools/tests/fixtures/view-image/valid.gif create mode 100644 crates/tools/tests/fixtures/view-image/valid.jpg create mode 100644 crates/tools/tests/fixtures/view-image/valid.png create mode 100644 crates/tools/tests/fixtures/view-image/valid.webp diff --git a/Cargo.lock b/Cargo.lock index 8053c9987..d5f5b2402 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6383,6 +6383,7 @@ dependencies = [ name = "yoi-workspace-server" version = "0.1.0" dependencies = [ + "agen", "async-trait", "axum", "chrono", @@ -6410,6 +6411,7 @@ dependencies = [ "tokio", "tokio-tungstenite 0.29.0", "toml", + "tools", "tower", "tracing", "tracing-subscriber", diff --git a/crates/tools/tests/fixtures/view-image/valid.gif b/crates/tools/tests/fixtures/view-image/valid.gif new file mode 100644 index 0000000000000000000000000000000000000000..f1eda11f885ee0f7a2a975585bf48aedf82e71e2 GIT binary patch literal 828 zcmW-gF{&m6420*seTZOVDOfBEDMi#u*-B(JHX^KZ0Q(#!=a@Q&$xRj^Fuwh}5N0x8 zG5O=q-+%t)-)B7I)7u{b1du=h4GgfrJ+=|ZKm|H5feqY48i`C)q7#$Y#6A2`$U+sm zFoiAL<7hOpQH^d)V;lEKia`!)(1RK5;2!a@$VDxBF^gT?gDDPqs6!v-u!s8$$^a8A zu)zTrJP^SMX0U=CoZtp8h-4%)S;-3wkpZ3brs4i*q#1p_u@ zz|PLj!NCC}+1S7=4o*%^pa78K;^N}r;pXP!;}aAV`G17Ln}HE%JPKgL6cIpI!yt$u q^8WyXAO`~z0~0f&5(ASU1G6B*|62?^%#1(|kZsS<%@Cma|0V!D7Z Scope { @@ -94,6 +96,33 @@ fn setup_routed() -> ( (left, right, spill, router, Registry::new(definitions)) } +fn valid_gif_with_len(target_len: usize) -> Vec { + let fixture = include_bytes!("fixtures/view-image/valid.gif"); + assert_eq!( + fixture.last(), + Some(&0x3b), + "GIF fixture must end in trailer" + ); + let mut bytes = fixture[..fixture.len() - 1].to_vec(); + bytes.extend_from_slice(&[0x21, 0xfe]); + let mut remaining = target_len + .checked_sub(bytes.len() + 2) + .expect("target must leave room for comment terminator and GIF trailer"); + while remaining > 0 { + let mut block = remaining.min(256); + if remaining - block == 1 { + block -= 1; + } + assert!(block >= 2); + bytes.push((block - 1) as u8); + bytes.extend(std::iter::repeat_n(b'x', block - 1)); + remaining -= block; + } + bytes.extend_from_slice(&[0, 0x3b]); + assert_eq!(bytes.len(), target_len); + bytes +} + async fn call(tool: &Arc, input: serde_json::Value) -> agen::tool::ToolOutput { tool.execute(&input.to_string(), Default::default()) .await @@ -546,7 +575,9 @@ async fn external_view_image_uses_a_protocol_bounded_request() { truncate_first_binary_response: false, mutate_after_first_binary_response: None, }); - let png = b"\x89PNG\r\n\x1a\nexternal\0image\nbody\xff"; + let png = include_bytes!("fixtures/view-image/valid.png"); + assert!(png.contains(&0)); + assert!(png.iter().any(|byte| *byte >= 0x80)); std::fs::write(dir.path().join("image.png"), png).unwrap(); let (_meta, tool) = view_image_tool(session)(); @@ -562,7 +593,7 @@ async fn external_view_image_rejects_truncated_and_cross_version_chunks() { 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(); + std::fs::write(&image_path, include_bytes!("fixtures/view-image/valid.png")).unwrap(); let truncated: WorkdirSessionHandle = Arc::new(ExternalValidatingSession { inner: LocalWorkdirSession::new(scope.clone(), dir.path().to_path_buf()), truncate_first_binary_response: true, @@ -577,8 +608,7 @@ async fn external_view_image_rejects_truncated_and_cross_version_chunks() { "{error}" ); - let mut large = vec![b'x'; 1024 * 1024 + 32]; - large[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); + let large = valid_gif_with_len(1024 * 1024 + 32); std::fs::write(&image_path, large).unwrap(); let changing: WorkdirSessionHandle = Arc::new(ExternalValidatingSession { inner: LocalWorkdirSession::new(scope, dir.path().to_path_buf()), @@ -601,22 +631,98 @@ async fn view_image_enforces_the_ten_mib_boundary_without_partial_attachments() 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 mut image = valid_gif_with_len(LIMIT); + std::fs::write(dir.path().join("boundary.gif"), &image).unwrap(); - let output = call(&tool, json!({ "path": "boundary.png" })).await; - let agen::tool::Attachment::Image(image) = &output.attachments[0]; - assert_eq!(image.data(), png); + let output = call(&tool, json!({ "path": "boundary.gif" })).await; + let agen::tool::Attachment::Image(attachment) = &output.attachments[0]; + assert_eq!(attachment.mime_type(), "image/gif"); + assert_eq!(attachment.data(), image); - png.push(0); - std::fs::write(dir.path().join("boundary.png"), png).unwrap(); - let error = call_err(&tool, json!({ "path": "boundary.png" })) + image.push(0); + std::fs::write(dir.path().join("boundary.gif"), image).unwrap(); + let error = call_err(&tool, json!({ "path": "boundary.gif" })) .await .to_string(); assert!(error.contains("10485760-byte limit"), "{error}"); } +#[tokio::test] +async fn view_image_attaches_valid_supported_image_fixtures() { + 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 fixtures: [(&str, &[u8], &str); 4] = [ + ( + "valid.png", + include_bytes!("fixtures/view-image/valid.png"), + "image/png", + ), + ( + "valid.jpg", + include_bytes!("fixtures/view-image/valid.jpg"), + "image/jpeg", + ), + ( + "valid.gif", + include_bytes!("fixtures/view-image/valid.gif"), + "image/gif", + ), + ( + "valid.webp", + include_bytes!("fixtures/view-image/valid.webp"), + "image/webp", + ), + ]; + + for (name, fixture, mime) in fixtures { + std::fs::write(dir.path().join(name), fixture).unwrap(); + let output = call(&tool, json!({ "path": name })).await; + let agen::tool::Attachment::Image(image) = &output.attachments[0]; + assert_eq!(image.mime_type(), mime); + assert_eq!(image.data(), fixture); + } +} + +#[tokio::test] +async fn delegated_view_image_preserves_caller_relative_path_identity() { + let dir = TempDir::new().unwrap(); + std::fs::create_dir(dir.path().join("docs")).unwrap(); + let fixture = include_bytes!("fixtures/view-image/valid.png"); + std::fs::write(dir.path().join("docs/image.png"), fixture).unwrap(); + let source: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::new( + Scope::writable(dir.path()).unwrap(), + dir.path().to_path_buf(), + )); + let broker = WorkdirToolBroker::new(source); + let lease = broker + .scope(WorkdirToolScope { + rules: vec![WorkdirToolScopeRule { + target: WorkdirPath::new("docs").unwrap(), + permission: WorkdirToolScopePermission::Read, + recursive: true, + symlink_policy: Default::default(), + }], + cwd: WorkdirPath::new("docs").unwrap(), + command: false, + }) + .await + .unwrap(); + let (_meta, tool) = view_image_tool(lease.tool_session())(); + + let output = call(&tool, json!({ "path": "image.png" })).await; + assert!(output.summary.contains("image.png")); + assert!(!output.summary.contains("docs/image.png")); + let agen::tool::Attachment::Image(image) = &output.attachments[0]; + assert_eq!(image.data(), fixture); + + lease.close().await.unwrap(); + broker.close().await.unwrap(); +} + #[tokio::test] async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { let dir = TempDir::new().unwrap(); @@ -624,7 +730,7 @@ async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { let scope = scope_with_spill(dir.path(), spill.path()); let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf())); - let png = b"\x89PNG\r\n\x1a\nprivate-image-body"; + let png = include_bytes!("fixtures/view-image/valid.png"); std::fs::write(dir.path().join("image.png"), png).unwrap(); let definition = view_image_tool(session); let (_meta, tool) = definition(); @@ -635,7 +741,6 @@ async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { assert_eq!(image.mime_type(), "image/png"); assert_eq!(image.data(), png); let serialized = serde_json::to_string(&output).unwrap(); - assert!(!serialized.contains("private-image-body")); assert!(serialized.contains("attachments")); let restored: agen::tool::ToolOutput = serde_json::from_str(&serialized).unwrap(); let agen::tool::Attachment::Image(restored_image) = &restored.attachments[0]; diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index efe773d9f..a58e0b01d 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -1016,12 +1016,15 @@ impl WorkdirSession for ScopedWorkdirSession { &self, mut request: ReadBytesRequest, ) -> Result { + let caller_path = request.path.clone(); 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 + let mut result = self.source.read_bytes(request).await?; + result.path = caller_path; + Ok(result) } async fn write(&self, mut request: WriteRequest) -> Result { @@ -1913,14 +1916,9 @@ mod tests { child.read(read("readme.md")).await.unwrap().bytes, b"visible" ); - assert_eq!( - child - .read_bytes(read_bytes("readme.md")) - .await - .unwrap() - .bytes, - b"visible" - ); + let binary = child.read_bytes(read_bytes("readme.md")).await.unwrap(); + assert_eq!(binary.path, fs_path("readme.md")); + assert_eq!(binary.bytes, b"visible"); assert!(matches!( child.write(write("new.md", "no")).await, Err(WorkdirError::Denied(_)) diff --git a/crates/workspace-server/Cargo.toml b/crates/workspace-server/Cargo.toml index 2438e67c8..1c50732f0 100644 --- a/crates/workspace-server/Cargo.toml +++ b/crates/workspace-server/Cargo.toml @@ -50,7 +50,9 @@ zeroize.workspace = true webauthn-rs = { workspace = true } [dev-dependencies] +agen.workspace = true tempfile.workspace = true +tools.workspace = true tower = { workspace = true, features = ["util"] } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } webauthn-authenticator-rs = { version = "0.5.5", default-features = false, features = ["softpasskey"] } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 35f37069b..5aad53234 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -5,7 +5,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock, Weak}; use async_trait::async_trait; -use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}; +use axum::extract::ws::{CloseFrame, Message as WsMessage, WebSocket, WebSocketUpgrade}; use axum::extract::{DefaultBodyLimit, Extension, Path as AxumPath, Query, Request, State}; #[cfg(test)] use axum::http::header::{CACHE_CONTROL, ETAG, IF_NONE_MATCH}; @@ -17472,41 +17472,116 @@ async fn send_external_workdir_server_frame( ) } +async fn reject_external_workdir_provider(socket: &mut WebSocket, code: u16, mut reason: String) { + const MAX_CLOSE_REASON_BYTES: usize = 123; + if reason.len() > MAX_CLOSE_REASON_BYTES { + reason.truncate(MAX_CLOSE_REASON_BYTES); + } + let _ = socket + .send(WsMessage::Close(Some(CloseFrame { + code, + reason: reason.into(), + }))) + .await; +} + +fn external_workdir_registration_parse_error(text: &str) -> String { + let actual_version = serde_json::from_str::(text) + .ok() + .and_then(|value| value.get("version")?.as_u64()); + match actual_version { + Some(actual) + if actual != workdir::external::EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION as u64 => + { + format!( + "unsupported External Workdir protocol version {actual}; Backend requires {}", + workdir::external::EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION + ) + } + _ => "invalid External Workdir provider registration".to_string(), + } +} + async fn serve_external_workdir_provider( api: WorkspaceApi, initial_grant: ExternalWorkdirGrantRecord, mut socket: WebSocket, ) { - let registration = + let registration_text = match tokio::time::timeout(std::time::Duration::from_secs(10), socket.next()).await { - Ok(Some(Ok(WsMessage::Text(text)))) if text.len() <= 64 * 1024 => { - serde_json::from_str::(&text).ok() + Ok(Some(Ok(WsMessage::Text(text)))) if text.len() <= 64 * 1024 => text, + _ => { + reject_external_workdir_provider( + &mut socket, + 1008, + "External Workdir provider registration was missing or oversized".to_string(), + ) + .await; + return; } - _ => None, }; - let Some(ExternalWorkdirProviderFrame { + let registration_frame = + match serde_json::from_str::(®istration_text) { + Ok(frame) => frame, + Err(_) => { + let reason = external_workdir_registration_parse_error(®istration_text); + reject_external_workdir_provider(&mut socket, 1002, reason).await; + return; + } + }; + let ExternalWorkdirProviderFrame { message: ExternalWorkdirProviderMessage::Register { registration }, .. - }) = registration + } = registration_frame else { - let _ = socket.close().await; + reject_external_workdir_provider( + &mut socket, + 1008, + "expected External Workdir provider registration".to_string(), + ) + .await; return; }; let expected_capabilities = match external_workdir_capabilities(&initial_grant.permissions) { Ok(capabilities) => capabilities, Err(_) => { - let _ = socket.close().await; + reject_external_workdir_provider( + &mut socket, + 1011, + "Backend could not validate External Workdir capabilities".to_string(), + ) + .await; return; } }; + if registration.read_limits != workdir::BoundedReadLimits::EXTERNAL_DEFAULT { + let required = workdir::BoundedReadLimits::EXTERNAL_DEFAULT; + reject_external_workdir_provider( + &mut socket, + 1008, + format!( + "provider read limits source={} response={}; Backend requires source={} response={}", + registration.read_limits.max_source_bytes, + registration.read_limits.max_response_bytes, + required.max_source_bytes, + required.max_response_bytes, + ), + ) + .await; + return; + } if registration.grant_id.as_str() != initial_grant.grant_id || registration.workdir_id.as_str() != initial_grant.workdir_id || registration.provider_instance_id.as_str() != initial_grant.provider_instance_id || registration.capabilities != expected_capabilities - || registration.read_limits != workdir::BoundedReadLimits::EXTERNAL_DEFAULT || registration.generation == 0 { - let _ = socket.close().await; + reject_external_workdir_provider( + &mut socket, + 1008, + "External Workdir provider registration does not match the grant".to_string(), + ) + .await; return; } let now = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true); @@ -26473,6 +26548,26 @@ mod tests { WorkspaceRuntimeBinding, }; + fn valid_gif_with_len(target_len: usize) -> Vec { + let fixture = include_bytes!("../../tools/tests/fixtures/view-image/valid.gif"); + assert_eq!(fixture.last(), Some(&0x3b)); + let mut bytes = fixture[..fixture.len() - 1].to_vec(); + bytes.extend_from_slice(&[0x21, 0xfe]); + let mut remaining = target_len.checked_sub(bytes.len() + 2).unwrap(); + while remaining > 0 { + let mut block = remaining.min(256); + if remaining - block == 1 { + block -= 1; + } + bytes.push((block - 1) as u8); + bytes.extend(std::iter::repeat_n(b'x', block - 1)); + remaining -= block; + } + bytes.extend_from_slice(&[0, 0x3b]); + assert_eq!(bytes.len(), target_len); + bytes + } + #[test] fn external_permission_levels_map_to_existing_bits_and_reject_legacy_partial_grants() { let bits = |capabilities| { @@ -28018,6 +28113,50 @@ mod tests { }) }; + let (mut old_provider, _) = connect_async(&url).await.unwrap(); + let old_registration = serde_json::to_string(®ister(1)).unwrap().replacen( + r#""version":2"#, + r#""version":1"#, + 1, + ); + old_provider + .send(Message::Text(old_registration.into())) + .await + .unwrap(); + let rejected = old_provider.next().await.unwrap().unwrap(); + let Message::Close(Some(frame)) = rejected else { + panic!("old provider must receive a bounded compatibility close reason"); + }; + assert!( + frame + .reason + .contains("unsupported External Workdir protocol version 1") + ); + assert!(frame.reason.contains("requires 2")); + + let (mut smaller_provider, _) = connect_async(&url).await.unwrap(); + let mut smaller_registration = serde_json::to_value(register(1)).unwrap(); + smaller_registration["message"]["registration"]["read_limits"] = serde_json::json!({ + "max_source_bytes": 32 * 1024 * 1024, + "max_response_bytes": 512 * 1024, + }); + smaller_provider + .send(Message::Text( + serde_json::to_string(&smaller_registration).unwrap().into(), + )) + .await + .unwrap(); + let rejected = smaller_provider.next().await.unwrap().unwrap(); + let Message::Close(Some(frame)) = rejected else { + panic!("smaller provider limits must receive the required effective limits"); + }; + assert!(frame.reason.contains("response=524288"), "{}", frame.reason); + assert!( + frame.reason.contains("response=1048576"), + "{}", + frame.reason + ); + let (mut first, _) = connect_async(&url).await.unwrap(); first .send(Message::Text( @@ -28128,8 +28267,8 @@ mod tests { ); let provider_root = tempfile::tempdir().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 image_bytes = valid_gif_with_len(1024 * 1024 + 257); + fs::write(provider_root.path().join("image.gif"), &image_bytes).unwrap(); let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap(); let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned( workdir::Workdir::new(&grant.working_directory_id), @@ -28199,17 +28338,72 @@ mod tests { let read_result = read.await.unwrap().unwrap(); assert_eq!(read_result.bytes, b"one\ntwo\n"); - let read_bytes = { + let view_image = { 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 (_meta, tool) = tools::view_image_tool(broker_session)(); + tool.execute( + r#"{"path":"image.gif"}"#, + agen::tool::ToolExecutionContext::default(), + ) + .await + }) + }; + for expected_offset in [0_u64, 1024 * 1024] { + let (operation_id, operation) = loop { + let text = second.next().await.unwrap().unwrap().into_text().unwrap(); + let frame = serde_json::from_str::(&text).unwrap(); + if let ExternalWorkdirServerMessage::Operation { + generation: 2, + operation_id, + operation, + } = frame.message + { + break (operation_id, operation); + } + }; + let WorkdirSessionOperation::ReadBytes(request) = operation.as_inner() else { + panic!("ViewImage must use the bounded binary operation"); + }; + assert_eq!(request.offset, expected_offset); + assert!(request.max_bytes <= 1024 * 1024); + 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 output = view_image.await.unwrap().unwrap(); + let agen::tool::Attachment::Image(image) = &output.attachments[0]; + assert_eq!(image.mime_type(), "image/gif"); + assert_eq!(image.data(), image_bytes); + + let missing_image = { + let broker_session = broker_session.clone(); + tokio::spawn(async move { + let (_meta, tool) = tools::view_image_tool(broker_session)(); + tool.execute( + r#"{"path":"missing.gif"}"#, + agen::tool::ToolExecutionContext::default(), + ) + .await }) }; let (operation_id, operation) = loop { @@ -28228,21 +28422,22 @@ mod tests { operation.as_inner(), WorkdirSessionOperation::ReadBytes(_) )); - let result = workdir::http::dispatch_workdir_session_operation( + let workdir_error = workdir::http::dispatch_workdir_session_operation( &local_session, operation.into_inner(), ) .await - .unwrap() - .try_into() - .unwrap(); + .unwrap_err(); + let error = workdir::external::ExternalWorkdirOperationError::from_transport_error( + workdir::http::WorkdirTransportError::from_workdir_error(&workdir_error), + ); second .send(Message::Text( serde_json::to_string(&ExternalWorkdirProviderFrame::current( ExternalWorkdirProviderMessage::OperationResult { generation: 2, operation_id, - outcome: ExternalWorkdirOperationOutcome::Completed { result }, + outcome: ExternalWorkdirOperationOutcome::Failed { error }, }, )) .unwrap() @@ -28250,9 +28445,8 @@ mod tests { )) .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 error = missing_image.await.unwrap().unwrap_err().to_string(); + assert!(error.contains("not found"), "{error}"); let edit = { let broker_session = broker_session.clone(); diff --git a/crates/yoi/src/workdir_share.rs b/crates/yoi/src/workdir_share.rs index a2d3bc6b2..d6dd9d187 100644 --- a/crates/yoi/src/workdir_share.rs +++ b/crates/yoi/src/workdir_share.rs @@ -37,6 +37,39 @@ pub(crate) struct WorkdirShareOptions { const PROVIDER_OPERATION_TIMEOUT: Duration = Duration::from_secs(30); +fn validate_provider_registration(message: Message, generation: u64) -> Result<(), String> { + let registered = match message { + Message::Text(text) => text, + Message::Close(Some(frame)) => { + let reason = frame.reason.trim(); + return Err(if reason.is_empty() { + "Backend rejected the External Workdir provider registration".to_string() + } else { + format!("Backend rejected the External Workdir provider registration: {reason}") + }); + } + Message::Close(None) => { + return Err("Backend rejected the External Workdir provider registration".to_string()); + } + _ => { + return Err( + "Backend returned an invalid External Workdir registration frame".to_string(), + ); + } + }; + let frame = serde_json::from_str::(®istered) + .map_err(|error| format!("Backend returned an invalid External Workdir frame: {error}"))?; + if matches!( + frame.message, + ExternalWorkdirServerMessage::Registered { generation: registered_generation, .. } + if registered_generation == generation + ) { + Ok(()) + } else { + Err("Backend rejected the External Workdir provider registration".to_string()) + } +} + fn provider_operation_timeout(operation: &workdir::http::WorkdirSessionOperation) -> Duration { if matches!( operation, @@ -672,18 +705,7 @@ async fn serve_provider_connection( .map_err(|error| format!("External Workdir provider registration failed: {error}"))? } }; - let Message::Text(registered) = registered else { - return Err("Backend returned an invalid External Workdir registration frame".to_string()); - }; - let frame = serde_json::from_str::(®istered) - .map_err(|error| format!("Backend returned an invalid External Workdir frame: {error}"))?; - if !matches!( - frame.message, - ExternalWorkdirServerMessage::Registered { generation: registered_generation, .. } - if registered_generation == generation - ) { - return Err("Backend rejected the External Workdir provider registration".to_string()); - } + validate_provider_registration(registered, generation)?; println!("Connection: online (generation {generation})"); let (completion_sender, mut completions) = @@ -1047,6 +1069,23 @@ mod tests { } } + #[test] + fn provider_registration_surfaces_backend_close_reason() { + use tokio_tungstenite::tungstenite::protocol::CloseFrame; + use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; + + let error = validate_provider_registration( + Message::Close(Some(CloseFrame { + code: CloseCode::Policy, + reason: "provider read limits source=1024 response=512; Backend requires source=2048 response=1024".into(), + })), + 1, + ) + .unwrap_err(); + assert!(error.contains("provider read limits"), "{error}"); + assert!(error.contains("Backend requires"), "{error}"); + } + #[test] fn interactive_share_defaults_to_first_workspace_read_only_and_unlimited() { let workspaces = [workspace("workspace-a", "Alpha")]; From 717d28839dc9b2f13eb3b37208137b312a49c6e1 Mon Sep 17 00:00:00 2001 From: Hare Date: Wed, 30 Sep 2026 08:32:43 +0900 Subject: [PATCH 3/3] fix: exercise production CLI image provider --- Cargo.lock | 1 + crates/workspace-server/Cargo.toml | 1 + crates/workspace-server/src/server.rs | 584 +++++--------------------- crates/yoi/src/lib.rs | 1 + crates/yoi/src/main.rs | 3 +- crates/yoi/src/workdir_share.rs | 50 ++- 6 files changed, 139 insertions(+), 501 deletions(-) create mode 100644 crates/yoi/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index d5f5b2402..41fb46abc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6423,6 +6423,7 @@ dependencies = [ "workdir", "worker", "worker-runtime", + "yoi", "zeroize", ] diff --git a/crates/workspace-server/Cargo.toml b/crates/workspace-server/Cargo.toml index 1c50732f0..5c78de89c 100644 --- a/crates/workspace-server/Cargo.toml +++ b/crates/workspace-server/Cargo.toml @@ -53,6 +53,7 @@ webauthn-rs = { workspace = true } agen.workspace = true tempfile.workspace = true tools.workspace = true +yoi = { path = "../yoi" } tower = { workspace = true, features = ["util"] } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } webauthn-authenticator-rs = { version = "0.5.5", default-features = false, features = ["softpasskey"] } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 5aad53234..e2669a778 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -28044,7 +28044,7 @@ mod tests { } #[tokio::test] - async fn external_provider_websocket_routes_files_and_commands_then_reconnects_and_revokes() { + async fn external_provider_websocket_runs_production_cli_pump_for_images_and_revoke() { let dir = tempfile::tempdir().unwrap(); init_clean_git_workspace(dir.path()); let api = test_api(dir.path()).await; @@ -28184,20 +28184,49 @@ mod tests { tokio::time::sleep(std::time::Duration::from_millis(10)).await; } - let (mut second, _) = connect_async(&url).await.unwrap(); - second - .send(Message::Text( - serde_json::to_string(®ister(2)).unwrap().into(), - )) + let provider_root = tempfile::tempdir().unwrap(); + fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap(); + let image_bytes = valid_gif_with_len(1024 * 1024 + 257); + fs::write(provider_root.path().join("image.gif"), &image_bytes).unwrap(); + let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap(); + let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned( + workdir::Workdir::new(&grant.working_directory_id), + pinned_provider_root, + workdir::BoundedReadLimits::EXTERNAL_DEFAULT, + workdir::WorkdirSessionCapabilities::ALL, + ) + .unwrap(); + let (second, _) = connect_async(&url).await.unwrap(); + let registration = register(2); + let provider = tokio::spawn(async move { + yoi::workdir_share::serve_external_workdir_provider_socket( + second, + registration, + &local_session, + 2, + std::future::pending::>(), + ) .await - .unwrap(); - let registered = second.next().await.unwrap().unwrap().into_text().unwrap(); - assert!(matches!( - serde_json::from_str::(®istered) + }); + for _ in 0..50 { + if api + .store + .get_external_workdir_grant(TEST_WORKSPACE_ID, &grant.grant_id) .unwrap() - .message, - ExternalWorkdirServerMessage::Registered { generation: 2, .. } - )); + .is_some_and(|stored| stored.status == "online") + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert_eq!( + api.store + .get_external_workdir_grant(TEST_WORKSPACE_ID, &grant.grant_id) + .unwrap() + .unwrap() + .status, + "online" + ); api.store .upsert_workdir_registry(&WorkdirRegistryRecord { @@ -28265,18 +28294,6 @@ mod tests { external_capabilities, workdir::WorkdirSessionCapabilities::ALL ); - let provider_root = tempfile::tempdir().unwrap(); - fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap(); - let image_bytes = valid_gif_with_len(1024 * 1024 + 257); - fs::write(provider_root.path().join("image.gif"), &image_bytes).unwrap(); - let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap(); - let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned( - workdir::Workdir::new(&grant.working_directory_id), - pinned_provider_root, - workdir::BoundedReadLimits::EXTERNAL_DEFAULT, - workdir::WorkdirSessionCapabilities::ALL, - ) - .unwrap(); let session_lock = current_worker_session_lock(&api, &worker); let session_guard = session_lock.lock().await; let broker_session = @@ -28288,272 +28305,60 @@ mod tests { workdir::WorkdirSessionCapabilities::ALL ); drop(session_guard); - let read = { - let broker_session = broker_session.clone(); - tokio::spawn(async move { - broker_session - .read(workdir::ReadRequest { - path: workdir::WorkdirPath::new("events.jsonl").unwrap(), - offset: 0, - limit: 10, - max_bytes: 1024, - }) - .await + let read_result = broker_session + .read(workdir::ReadRequest { + path: workdir::WorkdirPath::new("events.jsonl").unwrap(), + offset: 0, + limit: 10, + max_bytes: 1024, }) - }; - let (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if let ExternalWorkdirServerMessage::Operation { - generation: 2, - operation_id, - operation, - } = frame.message - { - break (operation_id, operation); - } - }; - 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_result = read.await.unwrap().unwrap(); assert_eq!(read_result.bytes, b"one\ntwo\n"); - let view_image = { - let broker_session = broker_session.clone(); - tokio::spawn(async move { - let (_meta, tool) = tools::view_image_tool(broker_session)(); - tool.execute( - r#"{"path":"image.gif"}"#, - agen::tool::ToolExecutionContext::default(), - ) - .await - }) - }; - for expected_offset in [0_u64, 1024 * 1024] { - let (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if let ExternalWorkdirServerMessage::Operation { - generation: 2, - operation_id, - operation, - } = frame.message - { - break (operation_id, operation); - } - }; - let WorkdirSessionOperation::ReadBytes(request) = operation.as_inner() else { - panic!("ViewImage must use the bounded binary operation"); - }; - assert_eq!(request.offset, expected_offset); - assert!(request.max_bytes <= 1024 * 1024); - let result = workdir::http::dispatch_workdir_session_operation( - &local_session, - operation.into_inner(), + let (_meta, view_image) = tools::view_image_tool(broker_session.clone())(); + let output = view_image + .execute( + r#"{"path":"image.gif"}"#, + agen::tool::ToolExecutionContext::default(), ) .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 output = view_image.await.unwrap().unwrap(); let agen::tool::Attachment::Image(image) = &output.attachments[0]; assert_eq!(image.mime_type(), "image/gif"); assert_eq!(image.data(), image_bytes); - let missing_image = { - let broker_session = broker_session.clone(); - tokio::spawn(async move { - let (_meta, tool) = tools::view_image_tool(broker_session)(); - tool.execute( - r#"{"path":"missing.gif"}"#, - agen::tool::ToolExecutionContext::default(), - ) - .await - }) - }; - let (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&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 workdir_error = workdir::http::dispatch_workdir_session_operation( - &local_session, - operation.into_inner(), - ) - .await - .unwrap_err(); - let error = workdir::external::ExternalWorkdirOperationError::from_transport_error( - workdir::http::WorkdirTransportError::from_workdir_error(&workdir_error), - ); - second - .send(Message::Text( - serde_json::to_string(&ExternalWorkdirProviderFrame::current( - ExternalWorkdirProviderMessage::OperationResult { - generation: 2, - operation_id, - outcome: ExternalWorkdirOperationOutcome::Failed { error }, - }, - )) - .unwrap() - .into(), - )) + let error = view_image + .execute( + r#"{"path":"missing.gif"}"#, + agen::tool::ToolExecutionContext::default(), + ) .await - .unwrap(); - let error = missing_image.await.unwrap().unwrap_err().to_string(); + .unwrap_err() + .to_string(); assert!(error.contains("not found"), "{error}"); - let edit = { - let broker_session = broker_session.clone(); - let expected_hash = read_result.content_hash; - tokio::spawn(async move { - broker_session - .edit(workdir::EditRequest { - path: workdir::WorkdirPath::new("events.jsonl").unwrap(), - old_string: "two".to_string(), - new_string: "changed".to_string(), - replace_all: false, - expected_hash, - }) - .await + let edit = broker_session + .edit(workdir::EditRequest { + path: workdir::WorkdirPath::new("events.jsonl").unwrap(), + old_string: "two".to_string(), + new_string: "changed".to_string(), + replace_all: false, + expected_hash: read_result.content_hash, }) - }; - let (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if let ExternalWorkdirServerMessage::Operation { - generation: 2, - operation_id, - operation, - } = frame.message - { - break (operation_id, operation); - } - }; - assert!(matches!( - operation.as_inner(), - WorkdirSessionOperation::Edit(_) - )); - 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(); - assert_eq!(edit.await.unwrap().unwrap().replacements, 1); + assert_eq!(edit.replacements, 1); - let write = { - let broker_session = broker_session.clone(); - tokio::spawn(async move { - broker_session - .write(workdir::WriteRequest { - path: workdir::WorkdirPath::new("created.txt").unwrap(), - content: b"created".to_vec(), - expected_hash: None, - }) - .await + let write = broker_session + .write(workdir::WriteRequest { + path: workdir::WorkdirPath::new("created.txt").unwrap(), + content: b"created".to_vec(), + expected_hash: None, }) - }; - let (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if let ExternalWorkdirServerMessage::Operation { - generation: 2, - operation_id, - operation, - } = frame.message - { - break (operation_id, operation); - } - }; - assert!(matches!( - operation.as_inner(), - WorkdirSessionOperation::Write(_) - )); - 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(); - assert!(write.await.unwrap().unwrap().created); + assert!(write.created); assert_eq!( fs::read_to_string(provider_root.path().join("events.jsonl")).unwrap(), "one\nchanged\n" @@ -28562,64 +28367,17 @@ mod tests { fs::read_to_string(provider_root.path().join("created.txt")).unwrap(), "created" ); - let command_start = { - let broker_session = broker_session.clone(); - tokio::spawn(async move { - broker_session - .start_command(workdir::CommandRequest { - command: "sleep 30 & child=$!; echo $child > child.pid; wait $child" - .to_string(), - timeout_secs: 30, - output_limit: 1024, - cwd: workdir::WorkdirPath::root(), - spill_dir: Some("/worker/private/spill".into()), - tool_call_id: Some("call-e2e".to_string()), - }) - .await + let command_handle = broker_session + .start_command(workdir::CommandRequest { + command: "sleep 30 & child=$!; echo $child > child.pid; wait $child".to_string(), + timeout_secs: 30, + output_limit: 1024, + cwd: workdir::WorkdirPath::root(), + spill_dir: Some("/worker/private/spill".into()), + tool_call_id: Some("call-e2e".to_string()), }) - }; - let (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if let ExternalWorkdirServerMessage::Operation { - generation: 2, - operation_id, - operation, - } = frame.message - { - break (operation_id, operation); - } - }; - assert!(matches!( - operation.as_inner(), - WorkdirSessionOperation::CommandStart(workdir::CommandRequest { - spill_dir: None, - .. - }) - )); - 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 command_handle = command_start.await.unwrap().unwrap(); for _ in 0..100 { if provider_root.path().join("child.pid").is_file() { @@ -28629,162 +28387,25 @@ mod tests { } assert!(provider_root.path().join("child.pid").is_file()); - let command_output = { - let broker_session = broker_session.clone(); - let handle = command_handle.clone(); - tokio::spawn(async move { - broker_session - .command_output(workdir::CommandOutputRequest { - handle, - cursor: 0, - limit: 1024, - wait: false, - }) - .await - }) - }; - let (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if let ExternalWorkdirServerMessage::Operation { - generation: 2, - operation_id, - operation, - } = frame.message - { - break (operation_id, operation); - } - }; - assert!(matches!( - operation.as_inner(), - WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest { + let command_output = broker_session + .command_output(workdir::CommandOutputRequest { + handle: command_handle, + cursor: 0, + limit: 1024, wait: false, - .. }) - )); - 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 command_output = command_output.await.unwrap().unwrap(); assert_eq!(command_output.status, workdir::CommandStatus::Running); assert_eq!(command_output.output_path, None); - let revoker = { - let api = api.clone(); - let grant_id = grant.grant_id.clone(); - tokio::spawn(async move { - scoped_revoke_external_workdir_grant( - State(api), - Extension(test_browser_request_actor()), - AxumPath((TEST_WORKSPACE_ID.to_string(), grant_id)), - ) - .await - }) - }; - let (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if let ExternalWorkdirServerMessage::Operation { - generation: 2, - operation_id, - operation, - } = frame.message - { - break (operation_id, operation); - } - }; - assert!(matches!( - operation.as_inner(), - WorkdirSessionOperation::CommandCancel(_) - )); - 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 (operation_id, operation) = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if let ExternalWorkdirServerMessage::Operation { - generation: 2, - operation_id, - operation, - } = frame.message - { - break (operation_id, operation); - } - }; - assert!(matches!( - operation.as_inner(), - WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest { - wait: true, - .. - }) - )); - let provider_result = workdir::http::dispatch_workdir_session_operation( - &local_session, - operation.into_inner(), + let Json(revoked) = scoped_revoke_external_workdir_grant( + State(api.clone()), + Extension(test_browser_request_actor()), + AxumPath((TEST_WORKSPACE_ID.to_string(), grant.grant_id.clone())), ) .await .unwrap(); - assert!(matches!( - provider_result, - WorkdirSessionOperationResult::CommandOutput(ref output) - if output.status != workdir::CommandStatus::Running - )); - let result = provider_result.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 Json(revoked) = revoker.await.unwrap().unwrap(); assert_eq!(revoked.status, "revoked"); let active_links = api .store @@ -28796,27 +28417,12 @@ mod tests { assert_eq!(active_links.len(), 1); assert_eq!(active_links[0].alias, "repository"); assert_eq!(active_links[0].workdir_id, "analysis-repository"); - let revoke = loop { - let text = second.next().await.unwrap().unwrap().into_text().unwrap(); - let frame = serde_json::from_str::(&text).unwrap(); - if matches!(frame.message, ExternalWorkdirServerMessage::Revoke { .. }) { - break frame; - } - }; - assert!(matches!( - revoke.message, - ExternalWorkdirServerMessage::Revoke { generation: 2, .. } - )); - second - .send(Message::Text( - serde_json::to_string(&ExternalWorkdirProviderFrame::current( - ExternalWorkdirProviderMessage::RevokeAcknowledged { generation: 2 }, - )) - .unwrap() - .into(), - )) + let provider_end = tokio::time::timeout(std::time::Duration::from_secs(2), provider) .await + .expect("production CLI provider pump did not acknowledge revocation") + .unwrap() .unwrap(); + assert_eq!(provider_end, yoi::workdir_share::ProviderEnd::Revoked); server.abort(); } diff --git a/crates/yoi/src/lib.rs b/crates/yoi/src/lib.rs new file mode 100644 index 000000000..32e2a7501 --- /dev/null +++ b/crates/yoi/src/lib.rs @@ -0,0 +1 @@ +pub mod workdir_share; diff --git a/crates/yoi/src/main.rs b/crates/yoi/src/main.rs index 924f5b849..6c36ba56e 100644 --- a/crates/yoi/src/main.rs +++ b/crates/yoi/src/main.rs @@ -4,10 +4,11 @@ mod objective_cli; mod plugin_cli; mod session_cli; mod ticket_cli; -mod workdir_share; mod worker_cleanup_cli; mod workspace_bootstrap; +use yoi::workdir_share; + use std::collections::BTreeMap; use std::fmt; use std::fs; diff --git a/crates/yoi/src/workdir_share.rs b/crates/yoi/src/workdir_share.rs index d6dd9d187..b97bf96ab 100644 --- a/crates/yoi/src/workdir_share.rs +++ b/crates/yoi/src/workdir_share.rs @@ -25,7 +25,7 @@ use workdir::{ }; #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct WorkdirShareOptions { +pub struct WorkdirShareOptions { pub path: PathBuf, pub workspace_id: Option, pub backend_url: String, @@ -82,7 +82,7 @@ fn provider_operation_timeout(operation: &workdir::http::WorkdirSessionOperation } #[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ProviderEnd { +pub enum ProviderEnd { Disconnected, Revoked, Interrupted, @@ -221,7 +221,7 @@ struct ResolvedWorkdirShareOptions { permissions: server_api::ExternalWorkdirPermissions, } -pub(crate) fn parse_ttl(value: &str) -> Result { +pub fn parse_ttl(value: &str) -> Result { let (number, multiplier) = if let Some(value) = value.strip_suffix('h') { (value, 60 * 60) } else if let Some(value) = value.strip_suffix('m') { @@ -457,7 +457,7 @@ async fn resolve_share_options( configure_interactively(options, &workspaces, &mut input, &mut output) } -pub(crate) async fn run(options: WorkdirShareOptions) -> Result<(), String> { +pub async fn run(options: WorkdirShareOptions) -> Result<(), String> { let Some(options) = resolve_share_options(options).await? else { eprintln!("External Workdir share cancelled; no grant was created."); return Ok(()); @@ -669,7 +669,7 @@ async fn serve_provider_connection( .max_message_size(Some(workdir::external::MAX_EXTERNAL_SERVER_FRAME_BYTES)) .max_frame_size(Some(workdir::external::MAX_EXTERNAL_SERVER_FRAME_BYTES)); let connection = connect_async_with_config(request, Some(websocket_config), false); - let (mut socket, _) = tokio::select! { + let (socket, _) = tokio::select! { signal = tokio::signal::ctrl_c() => { signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?; revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?; @@ -691,11 +691,40 @@ async fn serve_provider_connection( read_limits: BoundedReadLimits::EXTERNAL_DEFAULT, }, }); + let end = + serve_external_workdir_provider_socket(socket, registration, session, generation, async { + tokio::signal::ctrl_c() + .await + .map_err(|error| format!("failed to wait for Ctrl-C: {error}")) + }) + .await?; + if end == ProviderEnd::Interrupted { + revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?; + } + Ok(end) +} + +/// Run the production External Workdir provider protocol over an established WebSocket. +/// +/// The CLI uses this pump after opening its authenticated Backend connection. Keeping the +/// protocol pump separate lets composed tests exercise the same registration, operation, +/// cancellation, and result-forwarding implementation without spawning a second CLI process. +pub async fn serve_external_workdir_provider_socket( + mut socket: tokio_tungstenite::WebSocketStream, + registration: ExternalWorkdirProviderFrame, + session: &LocalWorkdirSession, + generation: u64, + interrupt: I, +) -> Result +where + S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, + I: std::future::Future>, +{ + tokio::pin!(interrupt); send_provider_frame(&mut socket, ®istration).await?; let registered = tokio::select! { - signal = tokio::signal::ctrl_c() => { - signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?; - revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?; + signal = &mut interrupt => { + signal?; return Ok(ProviderEnd::Interrupted); } registered = tokio::time::timeout(Duration::from_secs(10), socket.next()) => { @@ -713,10 +742,9 @@ async fn serve_provider_connection( let mut operations = HashMap::::new(); loop { tokio::select! { - signal = tokio::signal::ctrl_c() => { - signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?; + signal = &mut interrupt => { + signal?; stop_provider_operations(session, &mut operations, &mut completions).await?; - revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?; return Ok(ProviderEnd::Interrupted); } completion = completions.recv() => {