Merge remote-tracking branch 'origin/develop' into work/T-661-memory-document
This commit is contained in:
Generated
+3
@@ -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",
|
||||
@@ -6421,6 +6423,7 @@ dependencies = [
|
||||
"workdir",
|
||||
"worker",
|
||||
"worker-runtime",
|
||||
"yoi",
|
||||
"zeroize",
|
||||
]
|
||||
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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<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(
|
||||
root: &Path,
|
||||
request: WriteRequest,
|
||||
|
||||
@@ -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<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)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct WriteRequest {
|
||||
|
||||
@@ -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<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 {
|
||||
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 |
@@ -18,7 +18,9 @@ use tools::{
|
||||
Tracker, core_builtin_tools, routed_builtin_tools, routed_view_image_tool, view_image_tool,
|
||||
};
|
||||
use workdir::{
|
||||
LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionHandle, WorkdirSessionRouter,
|
||||
LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirPath, WorkdirSessionHandle,
|
||||
WorkdirSessionRouter, WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission,
|
||||
WorkdirToolScopeRule,
|
||||
};
|
||||
|
||||
fn scope_with_spill(workspace: &Path, spill: &Path) -> Scope {
|
||||
@@ -94,6 +96,33 @@ fn setup_routed() -> (
|
||||
(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 {
|
||||
tool.execute(&input.to_string(), Default::default())
|
||||
.await
|
||||
@@ -399,6 +428,301 @@ 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<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]
|
||||
async fn view_image_reads_scoped_bytes_into_durable_tool_detail() {
|
||||
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 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();
|
||||
@@ -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.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];
|
||||
|
||||
@@ -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::<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]
|
||||
fn protocol_rejects_unknown_nested_fields_and_unbounded_requests() {
|
||||
let unknown = r#"{
|
||||
"version": 1,
|
||||
"version": 2,
|
||||
"message": {
|
||||
"kind": "heartbeat",
|
||||
"generation": 1,
|
||||
|
||||
@@ -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<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> {
|
||||
match self
|
||||
.operate(WorkdirSessionOperation::Write(request))
|
||||
|
||||
@@ -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<StatResult, 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 edit(&self, request: EditRequest) -> Result<EditResult, WorkdirError>;
|
||||
async fn list(&self, request: ListRequest) -> Result<ListResult, WorkdirError>;
|
||||
|
||||
@@ -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<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> {
|
||||
self.ensure_capability(WorkdirSessionCapability::Write)?;
|
||||
let logical = request.path.clone();
|
||||
|
||||
@@ -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<ReadBytesResult, WorkdirError> {
|
||||
let _active = self.state.enter()?;
|
||||
self.state.session.read_bytes(request).await
|
||||
}
|
||||
|
||||
async fn write(&self, request: WriteRequest) -> Result<WriteResult, WorkdirError> {
|
||||
let _active = self.state.enter()?;
|
||||
self.state.session.write(request).await
|
||||
|
||||
@@ -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,21 @@ impl WorkdirSession for ScopedWorkdirSession {
|
||||
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> {
|
||||
let _scope_guard = self.scope_lock.lock().await;
|
||||
let path = self
|
||||
@@ -1276,6 +1292,10 @@ impl WorkdirSession for ReadOnlyWorkdirSession {
|
||||
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> {
|
||||
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 {
|
||||
WriteRequest {
|
||||
path: fs_path(path),
|
||||
@@ -1887,6 +1916,9 @@ mod tests {
|
||||
child.read(read("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(_))
|
||||
|
||||
@@ -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<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> {
|
||||
match self.operate(WorkdirSessionOperation::Write(request))? {
|
||||
WorkdirSessionOperationResult::Write(result) => Ok(result),
|
||||
|
||||
@@ -50,7 +50,10 @@ zeroize.workspace = true
|
||||
webauthn-rs = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
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"] }
|
||||
|
||||
@@ -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};
|
||||
@@ -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<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(
|
||||
&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
|
||||
@@ -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(
|
||||
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::<ExternalWorkdirProviderFrame>(&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::<ExternalWorkdirProviderFrame>(®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);
|
||||
@@ -17642,7 +17736,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;
|
||||
@@ -26454,6 +26548,26 @@ mod tests {
|
||||
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]
|
||||
fn external_permission_levels_map_to_existing_bits_and_reject_legacy_partial_grants() {
|
||||
let bits = |capabilities| {
|
||||
@@ -27930,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;
|
||||
@@ -27999,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(
|
||||
@@ -28026,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::<std::result::Result<(), String>>(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let registered = second.next().await.unwrap().unwrap().into_text().unwrap();
|
||||
assert!(matches!(
|
||||
serde_json::from_str::<ExternalWorkdirServerFrame>(®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 {
|
||||
@@ -28107,16 +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 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 =
|
||||
@@ -28128,162 +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::<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();
|
||||
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 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 (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();
|
||||
second
|
||||
.send(Message::Text(
|
||||
serde_json::to_string(&ExternalWorkdirProviderFrame::current(
|
||||
ExternalWorkdirProviderMessage::OperationResult {
|
||||
generation: 2,
|
||||
operation_id,
|
||||
outcome: ExternalWorkdirOperationOutcome::Completed { result },
|
||||
},
|
||||
))
|
||||
.unwrap()
|
||||
.into(),
|
||||
))
|
||||
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();
|
||||
assert_eq!(edit.await.unwrap().unwrap().replacements, 1);
|
||||
let agen::tool::Attachment::Image(image) = &output.attachments[0];
|
||||
assert_eq!(image.mime_type(), "image/gif");
|
||||
assert_eq!(image.data(), image_bytes);
|
||||
|
||||
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 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 {
|
||||
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::<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();
|
||||
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_eq!(edit.replacements, 1);
|
||||
|
||||
let write = broker_session
|
||||
.write(workdir::WriteRequest {
|
||||
path: workdir::WorkdirPath::new("created.txt").unwrap(),
|
||||
content: b"created".to_vec(),
|
||||
expected_hash: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(write.created);
|
||||
assert_eq!(
|
||||
fs::read_to_string(provider_root.path().join("events.jsonl")).unwrap(),
|
||||
"one\nchanged\n"
|
||||
@@ -28292,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::<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();
|
||||
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() {
|
||||
@@ -28359,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::<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 {
|
||||
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::<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(),
|
||||
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
|
||||
@@ -28526,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::<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(),
|
||||
))
|
||||
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();
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
pub mod workdir_share;
|
||||
@@ -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;
|
||||
|
||||
@@ -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<String>,
|
||||
pub backend_url: String,
|
||||
@@ -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::<ExternalWorkdirServerFrame>(®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,
|
||||
@@ -49,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,
|
||||
@@ -188,7 +221,7 @@ struct ResolvedWorkdirShareOptions {
|
||||
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') {
|
||||
(value, 60 * 60)
|
||||
} 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)
|
||||
}
|
||||
|
||||
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(());
|
||||
@@ -636,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?;
|
||||
@@ -658,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<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, ®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()) => {
|
||||
@@ -672,18 +734,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::<ExternalWorkdirServerFrame>(®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) =
|
||||
@@ -691,10 +742,9 @@ async fn serve_provider_connection(
|
||||
let mut operations = HashMap::<String, OperationTask>::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() => {
|
||||
@@ -903,7 +953,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,
|
||||
}),
|
||||
},
|
||||
@@ -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]
|
||||
fn interactive_share_defaults_to_first_workspace_read_only_and_unlimited() {
|
||||
let workspaces = [workspace("workspace-a", "Alpha")];
|
||||
|
||||
Reference in New Issue
Block a user