Merge remote-tracking branch 'origin/develop' into work/T-661-memory-document

This commit is contained in:
2026-09-30 08:39:27 +09:00
22 changed files with 1220 additions and 471 deletions
Generated
+3
View File
@@ -6383,6 +6383,7 @@ dependencies = [
name = "yoi-workspace-server" name = "yoi-workspace-server"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"agen",
"async-trait", "async-trait",
"axum", "axum",
"chrono", "chrono",
@@ -6410,6 +6411,7 @@ dependencies = [
"tokio", "tokio",
"tokio-tungstenite 0.29.0", "tokio-tungstenite 0.29.0",
"toml", "toml",
"tools",
"tower", "tower",
"tracing", "tracing",
"tracing-subscriber", "tracing-subscriber",
@@ -6421,6 +6423,7 @@ dependencies = [
"workdir", "workdir",
"worker", "worker",
"worker-runtime", "worker-runtime",
"yoi",
"zeroize", "zeroize",
] ]
+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))
} }
Binary file not shown.

After

Width:  |  Height:  |  Size: 828 B

Binary file not shown.

After

Width:  |  Height:  |  Size: 222 B

Binary file not shown.

After

Width:  |  Height:  |  Size: 94 B

Binary file not shown.

After

Width:  |  Height:  |  Size: 64 B

+326 -3
View File
@@ -18,7 +18,9 @@ use tools::{
Tracker, core_builtin_tools, routed_builtin_tools, routed_view_image_tool, view_image_tool, Tracker, core_builtin_tools, routed_builtin_tools, routed_view_image_tool, view_image_tool,
}; };
use workdir::{ use workdir::{
LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionHandle, WorkdirSessionRouter, LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirPath, WorkdirSessionHandle,
WorkdirSessionRouter, WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission,
WorkdirToolScopeRule,
}; };
fn scope_with_spill(workspace: &Path, spill: &Path) -> Scope { fn scope_with_spill(workspace: &Path, spill: &Path) -> Scope {
@@ -94,6 +96,33 @@ fn setup_routed() -> (
(left, right, spill, router, Registry::new(definitions)) (left, right, spill, router, Registry::new(definitions))
} }
fn valid_gif_with_len(target_len: usize) -> Vec<u8> {
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<dyn Tool>, input: serde_json::Value) -> agen::tool::ToolOutput { async fn call(tool: &Arc<dyn Tool>, input: serde_json::Value) -> agen::tool::ToolOutput {
tool.execute(&input.to_string(), Default::default()) tool.execute(&input.to_string(), Default::default())
.await .await
@@ -399,6 +428,301 @@ 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 = 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)();
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, 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,
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 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()),
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 image = valid_gif_with_len(LIMIT);
std::fs::write(dir.path().join("boundary.gif"), &image).unwrap();
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);
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] #[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();
@@ -406,7 +730,7 @@ async fn view_image_reads_scoped_bytes_into_durable_tool_detail() {
let scope = scope_with_spill(dir.path(), spill.path()); let scope = scope_with_spill(dir.path(), spill.path());
let session: WorkdirSessionHandle = let session: WorkdirSessionHandle =
Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf())); 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(); std::fs::write(dir.path().join("image.png"), png).unwrap();
let definition = view_image_tool(session); let definition = view_image_tool(session);
let (_meta, tool) = definition(); let (_meta, tool) = definition();
@@ -417,7 +741,6 @@ async fn view_image_reads_scoped_bytes_into_durable_tool_detail() {
assert_eq!(image.mime_type(), "image/png"); assert_eq!(image.mime_type(), "image/png");
assert_eq!(image.data(), png); assert_eq!(image.data(), png);
let serialized = serde_json::to_string(&output).unwrap(); let serialized = serde_json::to_string(&output).unwrap();
assert!(!serialized.contains("private-image-body"));
assert!(serialized.contains("attachments")); assert!(serialized.contains("attachments"));
let restored: agen::tool::ToolOutput = serde_json::from_str(&serialized).unwrap(); let restored: agen::tool::ToolOutput = serde_json::from_str(&serialized).unwrap();
let agen::tool::Attachment::Image(restored_image) = &restored.attachments[0]; let agen::tool::Attachment::Image(restored_image) = &restored.attachments[0];
+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
+33 -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,21 @@ 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 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;
let mut result = self.source.read_bytes(request).await?;
result.path = caller_path;
Ok(result)
}
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 +1292,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 +1676,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 +1916,9 @@ mod tests {
child.read(read("readme.md")).await.unwrap().bytes, child.read(read("readme.md")).await.unwrap().bytes,
b"visible" 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!( 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),
+3
View File
@@ -50,7 +50,10 @@ zeroize.workspace = true
webauthn-rs = { workspace = true } webauthn-rs = { workspace = true }
[dev-dependencies] [dev-dependencies]
agen.workspace = true
tempfile.workspace = true tempfile.workspace = true
tools.workspace = true
yoi = { path = "../yoi" }
tower = { workspace = true, features = ["util"] } tower = { workspace = true, features = ["util"] }
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
webauthn-authenticator-rs = { version = "0.5.5", default-features = false, features = ["softpasskey"] } webauthn-authenticator-rs = { version = "0.5.5", default-features = false, features = ["softpasskey"] }
+251 -375
View File
@@ -5,7 +5,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock, Weak}; use std::sync::{Arc, Mutex, RwLock, Weak};
use async_trait::async_trait; 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}; use axum::extract::{DefaultBodyLimit, Extension, Path as AxumPath, Query, Request, State};
#[cfg(test)] #[cfg(test)]
use axum::http::header::{CACHE_CONTROL, ETAG, IF_NONE_MATCH}; use axum::http::header::{CACHE_CONTROL, ETAG, IF_NONE_MATCH};
@@ -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
@@ -17453,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::<serde_json::Value>(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( async fn serve_external_workdir_provider(
api: WorkspaceApi, api: WorkspaceApi,
initial_grant: ExternalWorkdirGrantRecord, initial_grant: ExternalWorkdirGrantRecord,
mut socket: WebSocket, mut socket: WebSocket,
) { ) {
let registration = let registration_text =
match tokio::time::timeout(std::time::Duration::from_secs(10), socket.next()).await { match tokio::time::timeout(std::time::Duration::from_secs(10), socket.next()).await {
Ok(Some(Ok(WsMessage::Text(text)))) if text.len() <= 64 * 1024 => { Ok(Some(Ok(WsMessage::Text(text)))) if text.len() <= 64 * 1024 => text,
serde_json::from_str::<ExternalWorkdirProviderFrame>(&text).ok() _ => {
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::<ExternalWorkdirProviderFrame>(&registration_text) {
Ok(frame) => frame,
Err(_) => {
let reason = external_workdir_registration_parse_error(&registration_text);
reject_external_workdir_provider(&mut socket, 1002, reason).await;
return;
}
};
let ExternalWorkdirProviderFrame {
message: ExternalWorkdirProviderMessage::Register { registration }, message: ExternalWorkdirProviderMessage::Register { registration },
.. ..
}) = registration } = registration_frame
else { else {
let _ = socket.close().await; reject_external_workdir_provider(
&mut socket,
1008,
"expected External Workdir provider registration".to_string(),
)
.await;
return; return;
}; };
let expected_capabilities = match external_workdir_capabilities(&initial_grant.permissions) { let expected_capabilities = match external_workdir_capabilities(&initial_grant.permissions) {
Ok(capabilities) => capabilities, Ok(capabilities) => capabilities,
Err(_) => { Err(_) => {
let _ = socket.close().await; reject_external_workdir_provider(
&mut socket,
1011,
"Backend could not validate External Workdir capabilities".to_string(),
)
.await;
return; 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 if registration.grant_id.as_str() != initial_grant.grant_id
|| registration.workdir_id.as_str() != initial_grant.workdir_id || registration.workdir_id.as_str() != initial_grant.workdir_id
|| registration.provider_instance_id.as_str() != initial_grant.provider_instance_id || registration.provider_instance_id.as_str() != initial_grant.provider_instance_id
|| registration.capabilities != expected_capabilities || registration.capabilities != expected_capabilities
|| registration.read_limits != workdir::BoundedReadLimits::EXTERNAL_DEFAULT
|| registration.generation == 0 || 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; return;
} }
let now = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true); let now = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true);
@@ -17642,7 +17736,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;
@@ -26454,6 +26548,26 @@ mod tests {
WorkspaceRuntimeBinding, WorkspaceRuntimeBinding,
}; };
fn valid_gif_with_len(target_len: usize) -> Vec<u8> {
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] #[test]
fn external_permission_levels_map_to_existing_bits_and_reject_legacy_partial_grants() { fn external_permission_levels_map_to_existing_bits_and_reject_legacy_partial_grants() {
let bits = |capabilities| { let bits = |capabilities| {
@@ -27930,7 +28044,7 @@ mod tests {
} }
#[tokio::test] #[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(); let dir = tempfile::tempdir().unwrap();
init_clean_git_workspace(dir.path()); init_clean_git_workspace(dir.path());
let api = test_api(dir.path()).await; let api = test_api(dir.path()).await;
@@ -27999,6 +28113,50 @@ mod tests {
}) })
}; };
let (mut old_provider, _) = connect_async(&url).await.unwrap();
let old_registration = serde_json::to_string(&register(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(); let (mut first, _) = connect_async(&url).await.unwrap();
first first
.send(Message::Text( .send(Message::Text(
@@ -28026,20 +28184,49 @@ mod tests {
tokio::time::sleep(std::time::Duration::from_millis(10)).await; tokio::time::sleep(std::time::Duration::from_millis(10)).await;
} }
let (mut second, _) = connect_async(&url).await.unwrap(); let provider_root = tempfile::tempdir().unwrap();
second fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap();
.send(Message::Text( let image_bytes = valid_gif_with_len(1024 * 1024 + 257);
serde_json::to_string(&register(2)).unwrap().into(), fs::write(provider_root.path().join("image.gif"), &image_bytes).unwrap();
)) let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap();
.await 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(); .unwrap();
let registered = second.next().await.unwrap().unwrap().into_text().unwrap(); let (second, _) = connect_async(&url).await.unwrap();
assert!(matches!( let registration = register(2);
serde_json::from_str::<ExternalWorkdirServerFrame>(&registered) let provider = tokio::spawn(async move {
yoi::workdir_share::serve_external_workdir_provider_socket(
second,
registration,
&local_session,
2,
std::future::pending::<std::result::Result<(), String>>(),
)
.await
});
for _ in 0..50 {
if api
.store
.get_external_workdir_grant(TEST_WORKSPACE_ID, &grant.grant_id)
.unwrap() .unwrap()
.message, .is_some_and(|stored| stored.status == "online")
ExternalWorkdirServerMessage::Registered { generation: 2, .. } {
)); 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 api.store
.upsert_workdir_registry(&WorkdirRegistryRecord { .upsert_workdir_registry(&WorkdirRegistryRecord {
@@ -28107,16 +28294,6 @@ mod tests {
external_capabilities, external_capabilities,
workdir::WorkdirSessionCapabilities::ALL workdir::WorkdirSessionCapabilities::ALL
); );
let provider_root = tempfile::tempdir().unwrap();
fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").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_lock = current_worker_session_lock(&api, &worker);
let session_guard = session_lock.lock().await; let session_guard = session_lock.lock().await;
let broker_session = let broker_session =
@@ -28128,10 +28305,7 @@ mod tests {
workdir::WorkdirSessionCapabilities::ALL workdir::WorkdirSessionCapabilities::ALL
); );
drop(session_guard); drop(session_guard);
let read = { let read_result = broker_session
let broker_session = broker_session.clone();
tokio::spawn(async move {
broker_session
.read(workdir::ReadRequest { .read(workdir::ReadRequest {
path: workdir::WorkdirPath::new("events.jsonl").unwrap(), path: workdir::WorkdirPath::new("events.jsonl").unwrap(),
offset: 0, offset: 0,
@@ -28139,151 +28313,52 @@ mod tests {
max_bytes: 1024, max_bytes: 1024,
}) })
.await .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);
}
};
let result = workdir::http::dispatch_workdir_session_operation(
&local_session,
operation.into_inner(),
)
.await
.unwrap()
.try_into()
.unwrap(); .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"); assert_eq!(read_result.bytes, b"one\ntwo\n");
let edit = { let (_meta, view_image) = tools::view_image_tool(broker_session.clone())();
let broker_session = broker_session.clone(); let output = view_image
let expected_hash = read_result.content_hash; .execute(
tokio::spawn(async move { r#"{"path":"image.gif"}"#,
broker_session agen::tool::ToolExecutionContext::default(),
)
.await
.unwrap();
let agen::tool::Attachment::Image(image) = &output.attachments[0];
assert_eq!(image.mime_type(), "image/gif");
assert_eq!(image.data(), image_bytes);
let error = view_image
.execute(
r#"{"path":"missing.gif"}"#,
agen::tool::ToolExecutionContext::default(),
)
.await
.unwrap_err()
.to_string();
assert!(error.contains("not found"), "{error}");
let edit = broker_session
.edit(workdir::EditRequest { .edit(workdir::EditRequest {
path: workdir::WorkdirPath::new("events.jsonl").unwrap(), path: workdir::WorkdirPath::new("events.jsonl").unwrap(),
old_string: "two".to_string(), old_string: "two".to_string(),
new_string: "changed".to_string(), new_string: "changed".to_string(),
replace_all: false, replace_all: false,
expected_hash, expected_hash: read_result.content_hash,
}) })
.await .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::Edit(_)
));
let result = workdir::http::dispatch_workdir_session_operation(
&local_session,
operation.into_inner(),
)
.await
.unwrap()
.try_into()
.unwrap(); .unwrap();
second assert_eq!(edit.replacements, 1);
.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);
let write = { let write = broker_session
let broker_session = broker_session.clone();
tokio::spawn(async move {
broker_session
.write(workdir::WriteRequest { .write(workdir::WriteRequest {
path: workdir::WorkdirPath::new("created.txt").unwrap(), path: workdir::WorkdirPath::new("created.txt").unwrap(),
content: b"created".to_vec(), content: b"created".to_vec(),
expected_hash: None, expected_hash: None,
}) })
.await .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::Write(_)
));
let result = workdir::http::dispatch_workdir_session_operation(
&local_session,
operation.into_inner(),
)
.await
.unwrap()
.try_into()
.unwrap(); .unwrap();
second assert!(write.created);
.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_eq!( assert_eq!(
fs::read_to_string(provider_root.path().join("events.jsonl")).unwrap(), fs::read_to_string(provider_root.path().join("events.jsonl")).unwrap(),
"one\nchanged\n" "one\nchanged\n"
@@ -28292,13 +28367,9 @@ mod tests {
fs::read_to_string(provider_root.path().join("created.txt")).unwrap(), fs::read_to_string(provider_root.path().join("created.txt")).unwrap(),
"created" "created"
); );
let command_start = { let command_handle = broker_session
let broker_session = broker_session.clone();
tokio::spawn(async move {
broker_session
.start_command(workdir::CommandRequest { .start_command(workdir::CommandRequest {
command: "sleep 30 & child=$!; echo $child > child.pid; wait $child" command: "sleep 30 & child=$!; echo $child > child.pid; wait $child".to_string(),
.to_string(),
timeout_secs: 30, timeout_secs: 30,
output_limit: 1024, output_limit: 1024,
cwd: workdir::WorkdirPath::root(), cwd: workdir::WorkdirPath::root(),
@@ -28306,50 +28377,7 @@ mod tests {
tool_call_id: Some("call-e2e".to_string()), tool_call_id: Some("call-e2e".to_string()),
}) })
.await .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::CommandStart(workdir::CommandRequest {
spill_dir: None,
..
})
));
let result = workdir::http::dispatch_workdir_session_operation(
&local_session,
operation.into_inner(),
)
.await
.unwrap()
.try_into()
.unwrap(); .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 { for _ in 0..100 {
if provider_root.path().join("child.pid").is_file() { if provider_root.path().join("child.pid").is_file() {
@@ -28359,162 +28387,25 @@ mod tests {
} }
assert!(provider_root.path().join("child.pid").is_file()); assert!(provider_root.path().join("child.pid").is_file());
let command_output = { let command_output = broker_session
let broker_session = broker_session.clone();
let handle = command_handle.clone();
tokio::spawn(async move {
broker_session
.command_output(workdir::CommandOutputRequest { .command_output(workdir::CommandOutputRequest {
handle, handle: command_handle,
cursor: 0, cursor: 0,
limit: 1024, limit: 1024,
wait: false, wait: false,
}) })
.await .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::CommandOutput(workdir::CommandOutputRequest {
wait: false,
..
})
));
let result = workdir::http::dispatch_workdir_session_operation(
&local_session,
operation.into_inner(),
)
.await
.unwrap()
.try_into()
.unwrap(); .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.status, workdir::CommandStatus::Running);
assert_eq!(command_output.output_path, None); assert_eq!(command_output.output_path, None);
let revoker = { let Json(revoked) = scoped_revoke_external_workdir_grant(
let api = api.clone(); State(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()), Extension(test_browser_request_actor()),
AxumPath((TEST_WORKSPACE_ID.to_string(), grant_id)), AxumPath((TEST_WORKSPACE_ID.to_string(), grant.grant_id.clone())),
)
.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::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::<ExternalWorkdirServerFrame>(&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(),
) )
.await .await
.unwrap(); .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"); assert_eq!(revoked.status, "revoked");
let active_links = api let active_links = api
.store .store
@@ -28526,27 +28417,12 @@ mod tests {
assert_eq!(active_links.len(), 1); assert_eq!(active_links.len(), 1);
assert_eq!(active_links[0].alias, "repository"); assert_eq!(active_links[0].alias, "repository");
assert_eq!(active_links[0].workdir_id, "analysis-repository"); assert_eq!(active_links[0].workdir_id, "analysis-repository");
let revoke = loop { let provider_end = tokio::time::timeout(std::time::Duration::from_secs(2), provider)
let text = second.next().await.unwrap().unwrap().into_text().unwrap();
let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&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(),
))
.await .await
.expect("production CLI provider pump did not acknowledge revocation")
.unwrap()
.unwrap(); .unwrap();
assert_eq!(provider_end, yoi::workdir_share::ProviderEnd::Revoked);
server.abort(); server.abort();
} }
+1
View File
@@ -0,0 +1 @@
pub mod workdir_share;
+2 -1
View File
@@ -4,10 +4,11 @@ mod objective_cli;
mod plugin_cli; mod plugin_cli;
mod session_cli; mod session_cli;
mod ticket_cli; mod ticket_cli;
mod workdir_share;
mod worker_cleanup_cli; mod worker_cleanup_cli;
mod workspace_bootstrap; mod workspace_bootstrap;
use yoi::workdir_share;
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::fmt; use std::fmt;
use std::fs; use std::fs;
+91 -24
View File
@@ -25,7 +25,7 @@ use workdir::{
}; };
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct WorkdirShareOptions { pub struct WorkdirShareOptions {
pub path: PathBuf, pub path: PathBuf,
pub workspace_id: Option<String>, pub workspace_id: Option<String>,
pub backend_url: String, pub backend_url: String,
@@ -37,6 +37,39 @@ pub(crate) struct WorkdirShareOptions {
const PROVIDER_OPERATION_TIMEOUT: Duration = Duration::from_secs(30); 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::<ExternalWorkdirServerFrame>(&registered)
.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 { fn provider_operation_timeout(operation: &workdir::http::WorkdirSessionOperation) -> Duration {
if matches!( if matches!(
operation, operation,
@@ -49,7 +82,7 @@ fn provider_operation_timeout(operation: &workdir::http::WorkdirSessionOperation
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq)] #[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum ProviderEnd { pub enum ProviderEnd {
Disconnected, Disconnected,
Revoked, Revoked,
Interrupted, Interrupted,
@@ -188,7 +221,7 @@ struct ResolvedWorkdirShareOptions {
permissions: server_api::ExternalWorkdirPermissions, permissions: server_api::ExternalWorkdirPermissions,
} }
pub(crate) fn parse_ttl(value: &str) -> Result<Duration, String> { pub fn parse_ttl(value: &str) -> Result<Duration, String> {
let (number, multiplier) = if let Some(value) = value.strip_suffix('h') { let (number, multiplier) = if let Some(value) = value.strip_suffix('h') {
(value, 60 * 60) (value, 60 * 60)
} else if let Some(value) = value.strip_suffix('m') { } else if let Some(value) = value.strip_suffix('m') {
@@ -424,7 +457,7 @@ async fn resolve_share_options(
configure_interactively(options, &workspaces, &mut input, &mut output) 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 { let Some(options) = resolve_share_options(options).await? else {
eprintln!("External Workdir share cancelled; no grant was created."); eprintln!("External Workdir share cancelled; no grant was created.");
return Ok(()); return Ok(());
@@ -636,7 +669,7 @@ async fn serve_provider_connection(
.max_message_size(Some(workdir::external::MAX_EXTERNAL_SERVER_FRAME_BYTES)) .max_message_size(Some(workdir::external::MAX_EXTERNAL_SERVER_FRAME_BYTES))
.max_frame_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 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 = tokio::signal::ctrl_c() => {
signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?; signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?;
revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?; revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?;
@@ -658,11 +691,40 @@ async fn serve_provider_connection(
read_limits: BoundedReadLimits::EXTERNAL_DEFAULT, 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<S, I>(
mut socket: tokio_tungstenite::WebSocketStream<S>,
registration: ExternalWorkdirProviderFrame,
session: &LocalWorkdirSession,
generation: u64,
interrupt: I,
) -> Result<ProviderEnd, String>
where
S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
I: std::future::Future<Output = Result<(), String>>,
{
tokio::pin!(interrupt);
send_provider_frame(&mut socket, &registration).await?; send_provider_frame(&mut socket, &registration).await?;
let registered = tokio::select! { let registered = tokio::select! {
signal = tokio::signal::ctrl_c() => { signal = &mut interrupt => {
signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?; signal?;
revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?;
return Ok(ProviderEnd::Interrupted); return Ok(ProviderEnd::Interrupted);
} }
registered = tokio::time::timeout(Duration::from_secs(10), socket.next()) => { registered = tokio::time::timeout(Duration::from_secs(10), socket.next()) => {
@@ -672,18 +734,7 @@ async fn serve_provider_connection(
.map_err(|error| format!("External Workdir provider registration failed: {error}"))? .map_err(|error| format!("External Workdir provider registration failed: {error}"))?
} }
}; };
let Message::Text(registered) = registered else { validate_provider_registration(registered, generation)?;
return Err("Backend returned an invalid External Workdir registration frame".to_string());
};
let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&registered)
.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());
}
println!("Connection: online (generation {generation})"); println!("Connection: online (generation {generation})");
let (completion_sender, mut completions) = let (completion_sender, mut completions) =
@@ -691,10 +742,9 @@ async fn serve_provider_connection(
let mut operations = HashMap::<String, OperationTask>::new(); let mut operations = HashMap::<String, OperationTask>::new();
loop { loop {
tokio::select! { tokio::select! {
signal = tokio::signal::ctrl_c() => { signal = &mut interrupt => {
signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?; signal?;
stop_provider_operations(session, &mut operations, &mut completions).await?; stop_provider_operations(session, &mut operations, &mut completions).await?;
revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?;
return Ok(ProviderEnd::Interrupted); return Ok(ProviderEnd::Interrupted);
} }
completion = completions.recv() => { completion = completions.recv() => {
@@ -903,7 +953,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,
}), }),
}, },
@@ -1047,6 +1097,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] #[test]
fn interactive_share_defaults_to_first_workspace_read_only_and_unlimited() { fn interactive_share_defaults_to_first_workspace_read_only_and_unlimited() {
let workspaces = [workspace("workspace-a", "Alpha")]; let workspaces = [workspace("workspace-a", "Alpha")];