feat: add bounded Workspace binary transport and native WIP attachments
Keep JSON and raw-byte authority paths distinct. Runtime signs exact method/query/worker/body; binary requests are bounded, reject target normalization and redirects, and sanitize transport errors. Native attachments accompany the validated WIP value and exact validator through existing ToolOutput capture/replay. Reuse scoped Workdir binary reads with complete SHA256 verification. T-723 validation: binary transport/reviewer denial, native Invoke image capture and history replay, and bounded Workdir image helpers pass. Complete changed-crate suites and root cargo check passed on the combined Ticket implementation.
This commit is contained in:
@@ -21,7 +21,7 @@ mod grep;
|
||||
mod read;
|
||||
mod routing;
|
||||
mod search_target;
|
||||
mod view_image;
|
||||
pub mod view_image;
|
||||
mod web;
|
||||
mod write;
|
||||
|
||||
@@ -33,7 +33,9 @@ pub use glob::glob_tool;
|
||||
pub use grep::grep_tool;
|
||||
pub use read::read_tool;
|
||||
pub use tracker::{ChangeStat, Tracker};
|
||||
pub use view_image::{routed_view_image_tool, view_image_tool};
|
||||
pub use view_image::{
|
||||
detect_image_mime, read_workdir_bytes, routed_view_image_tool, view_image_tool,
|
||||
};
|
||||
pub use web::{web_fetch_tool, web_search_tool};
|
||||
pub use write::write_tool;
|
||||
|
||||
|
||||
@@ -7,7 +7,10 @@ use agen::tool::{
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
use workdir::{ReadBytesRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter};
|
||||
use sha2::{Digest, Sha256};
|
||||
use workdir::{
|
||||
ContentHash, ReadBytesRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter,
|
||||
};
|
||||
|
||||
use crate::error::ToolsError;
|
||||
|
||||
@@ -49,7 +52,7 @@ impl Tool for ViewImageTool {
|
||||
workdir::WorkdirSessionCapability::Read,
|
||||
)?;
|
||||
let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?;
|
||||
let bytes = read_image_bytes(selected.session.as_ref(), path.clone())
|
||||
let bytes = read_workdir_bytes(selected.session.as_ref(), path.clone(), MAX_IMAGE_BYTES)
|
||||
.await
|
||||
.map_err(ToolsError::from)?;
|
||||
let mime_type = detect_image_mime(&bytes).ok_or_else(|| {
|
||||
@@ -70,9 +73,13 @@ impl Tool for ViewImageTool {
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_image_bytes(
|
||||
/// Read arbitrary bytes through the scoped provider, bounded by `max_total_bytes`.
|
||||
/// Checks every chunk's identity, range, size and hash, then verifies the assembled
|
||||
/// SHA256 against the provider's pinned [`ContentHash`]. No MIME restriction applies.
|
||||
pub async fn read_workdir_bytes(
|
||||
session: &dyn workdir::WorkdirSession,
|
||||
path: WorkdirPath,
|
||||
max_total_bytes: usize,
|
||||
) -> Result<Vec<u8>, workdir::WorkdirError> {
|
||||
let mut bytes = Vec::new();
|
||||
let mut offset = 0_u64;
|
||||
@@ -116,16 +123,16 @@ async fn read_image_bytes(
|
||||
"bounded binary read returned a truncated or invalid byte range".to_string(),
|
||||
));
|
||||
}
|
||||
if result.total_bytes > MAX_IMAGE_BYTES as u64 {
|
||||
if result.total_bytes > max_total_bytes as u64 {
|
||||
return Err(workdir::WorkdirError::InvalidArgument(format!(
|
||||
"image exceeds the {MAX_IMAGE_BYTES}-byte limit"
|
||||
"file exceeds the {max_total_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(),
|
||||
"file changed while it was being read; retry the read".to_string(),
|
||||
));
|
||||
}
|
||||
if expected_total.is_none() {
|
||||
@@ -146,7 +153,13 @@ async fn read_image_bytes(
|
||||
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(),
|
||||
"bounded binary read returned partial file bytes".to_string(),
|
||||
));
|
||||
}
|
||||
let assembled_hash: ContentHash = Sha256::digest(&bytes).into();
|
||||
if assembled_hash != result.content_hash {
|
||||
return Err(workdir::WorkdirError::Transport(
|
||||
"bounded binary read assembled bytes do not match the content hash".to_string(),
|
||||
));
|
||||
}
|
||||
return Ok(bytes);
|
||||
@@ -172,7 +185,8 @@ pub fn routed_view_image_tool(router: Arc<WorkdirSessionRouter>) -> ToolDefiniti
|
||||
})
|
||||
}
|
||||
|
||||
fn detect_image_mime(bytes: &[u8]) -> Option<&'static str> {
|
||||
/// Detect PNG, JPEG, GIF or WebP signatures without trusting a file extension.
|
||||
pub fn detect_image_mime(bytes: &[u8]) -> Option<&'static str> {
|
||||
if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
|
||||
Some("image/png")
|
||||
} else if bytes.starts_with(&[0xff, 0xd8, 0xff]) {
|
||||
@@ -189,6 +203,189 @@ fn detect_image_mime(bytes: &[u8]) -> Option<&'static str> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use manifest::Scope;
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
use tempfile::TempDir;
|
||||
use workdir::*;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct BinarySession {
|
||||
local: LocalWorkdirSession,
|
||||
responses: Mutex<VecDeque<ReadBytesResult>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WorkdirSession for BinarySession {
|
||||
fn workdir(&self) -> &Workdir {
|
||||
self.local.workdir()
|
||||
}
|
||||
fn capabilities(&self) -> WorkdirSessionCapabilities {
|
||||
self.local.capabilities()
|
||||
}
|
||||
async fn read_bytes(
|
||||
&self,
|
||||
request: ReadBytesRequest,
|
||||
) -> Result<ReadBytesResult, WorkdirError> {
|
||||
let response = self.responses.lock().unwrap().pop_front().unwrap();
|
||||
if request.offset > 0 {
|
||||
assert!(request.expected_hash.is_some());
|
||||
}
|
||||
assert!(request.max_bytes <= IMAGE_READ_CHUNK_BYTES);
|
||||
Ok(response)
|
||||
}
|
||||
async fn stat(&self, r: StatRequest) -> Result<StatResult, WorkdirError> {
|
||||
self.local.stat(r).await
|
||||
}
|
||||
async fn read(&self, r: ReadRequest) -> Result<ReadResult, WorkdirError> {
|
||||
self.local.read(r).await
|
||||
}
|
||||
async fn write(&self, r: WriteRequest) -> Result<WriteResult, WorkdirError> {
|
||||
self.local.write(r).await
|
||||
}
|
||||
async fn edit(&self, r: EditRequest) -> Result<EditResult, WorkdirError> {
|
||||
self.local.edit(r).await
|
||||
}
|
||||
async fn list(&self, r: ListRequest) -> Result<ListResult, WorkdirError> {
|
||||
self.local.list(r).await
|
||||
}
|
||||
async fn glob(&self, r: GlobRequest) -> Result<GlobResult, WorkdirError> {
|
||||
self.local.glob(r).await
|
||||
}
|
||||
async fn grep(&self, r: GrepRequest) -> Result<GrepResult, WorkdirError> {
|
||||
self.local.grep(r).await
|
||||
}
|
||||
async fn start_command(&self, r: CommandRequest) -> Result<CommandHandle, WorkdirError> {
|
||||
self.local.start_command(r).await
|
||||
}
|
||||
async fn command_status(&self, r: CommandHandle) -> Result<CommandStatus, WorkdirError> {
|
||||
self.local.command_status(r).await
|
||||
}
|
||||
async fn command_output(
|
||||
&self,
|
||||
r: CommandOutputRequest,
|
||||
) -> Result<CommandOutput, WorkdirError> {
|
||||
self.local.command_output(r).await
|
||||
}
|
||||
async fn cancel_command(&self, r: CommandHandle) -> Result<(), WorkdirError> {
|
||||
self.local.cancel_command(r).await
|
||||
}
|
||||
async fn close(&self) -> Result<(), WorkdirError> {
|
||||
self.local.close().await
|
||||
}
|
||||
}
|
||||
|
||||
fn binary_session(dir: &TempDir, responses: Vec<ReadBytesResult>) -> BinarySession {
|
||||
BinarySession {
|
||||
local: LocalWorkdirSession::new(
|
||||
Scope::writable(dir.path()).unwrap(),
|
||||
dir.path().into(),
|
||||
),
|
||||
responses: Mutex::new(responses.into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn chunk(bytes: &[u8], offset: u64, total: u64, hash: ContentHash) -> ReadBytesResult {
|
||||
ReadBytesResult {
|
||||
path: WorkdirPath::new("artifact.bin").unwrap(),
|
||||
bytes: bytes.to_vec(),
|
||||
offset,
|
||||
total_bytes: total,
|
||||
content_hash: hash,
|
||||
eof: offset + bytes.len() as u64 == total,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_binary_read_rejects_bytes_with_forged_consistent_hash() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let hash = Sha256::digest(b"abcd").into();
|
||||
let session = binary_session(
|
||||
&dir,
|
||||
vec![chunk(b"ab", 0, 4, hash), chunk(b"cx", 2, 4, hash)],
|
||||
);
|
||||
let error = read_workdir_bytes(&session, WorkdirPath::new("artifact.bin").unwrap(), 4)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
matches!(error, WorkdirError::Transport(message) if message.contains("content hash"))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_binary_read_rejects_invalid_provider_chunk_metadata() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let hash = Sha256::digest(b"abcd").into();
|
||||
let valid = chunk(b"abcd", 0, 4, hash);
|
||||
let mut cases = Vec::new();
|
||||
let mut wrong_path = valid.clone();
|
||||
wrong_path.path = WorkdirPath::new("other.bin").unwrap();
|
||||
cases.push(wrong_path);
|
||||
let mut wrong_offset = valid.clone();
|
||||
wrong_offset.offset = 1;
|
||||
cases.push(wrong_offset);
|
||||
let mut wrong_eof = valid.clone();
|
||||
wrong_eof.eof = false;
|
||||
cases.push(wrong_eof);
|
||||
let mut wrong_total = valid.clone();
|
||||
wrong_total.total_bytes = 3;
|
||||
cases.push(wrong_total);
|
||||
cases.push(chunk(b"", 0, 4, hash));
|
||||
cases.push(chunk(
|
||||
&vec![0; IMAGE_READ_CHUNK_BYTES + 1],
|
||||
0,
|
||||
(IMAGE_READ_CHUNK_BYTES + 1) as u64,
|
||||
hash,
|
||||
));
|
||||
for response in cases {
|
||||
let session = binary_session(&dir, vec![response]);
|
||||
let error = read_workdir_bytes(&session, valid.path.clone(), 16 * 1024 * 1024)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(error, WorkdirError::Transport(_)), "{error}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_binary_read_rejects_hash_or_size_changes_between_chunks() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let hash = Sha256::digest(b"abcd").into();
|
||||
for second in [chunk(b"cd", 2, 4, [0; 32]), chunk(b"cd", 2, 5, hash)] {
|
||||
let session = binary_session(&dir, vec![chunk(b"ab", 0, 4, hash), second]);
|
||||
let error = read_workdir_bytes(&session, WorkdirPath::new("artifact.bin").unwrap(), 5)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(error, WorkdirError::Conflict(_)), "{error}");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bounded_binary_read_accepts_arbitrary_bytes_at_caller_limit_and_rejects_larger_files()
|
||||
{
|
||||
let dir = TempDir::new().unwrap();
|
||||
let bytes = vec![0xff; 16 * 1024 * 1024];
|
||||
std::fs::write(dir.path().join("artifact.bin"), &bytes).unwrap();
|
||||
let session =
|
||||
LocalWorkdirSession::new(Scope::writable(dir.path()).unwrap(), dir.path().into());
|
||||
let path = WorkdirPath::new("artifact.bin").unwrap();
|
||||
assert_eq!(
|
||||
read_workdir_bytes(&session, path.clone(), bytes.len())
|
||||
.await
|
||||
.unwrap(),
|
||||
bytes
|
||||
);
|
||||
assert!(matches!(
|
||||
read_workdir_bytes(&session, path.clone(), bytes.len() - 1).await,
|
||||
Err(WorkdirError::InvalidArgument(_))
|
||||
));
|
||||
std::fs::write(dir.path().join("artifact.bin"), []).unwrap();
|
||||
assert!(
|
||||
read_workdir_bytes(&session, path, 0)
|
||||
.await
|
||||
.unwrap()
|
||||
.is_empty()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_supported_image_signatures_without_trusting_extensions() {
|
||||
|
||||
@@ -1,11 +1,13 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use crate::workspace_request::{RuntimeWorkspaceRequest, RuntimeWorkspaceRequestClient};
|
||||
use crate::workspace_request::{
|
||||
RuntimeWorkspaceRequest, RuntimeWorkspaceRequestClient, RuntimeWorkspaceRequestError,
|
||||
};
|
||||
use worker::{
|
||||
WorkspaceClient, WorkspaceClientError, WorkspacePromptCatalogResolution,
|
||||
WorkspacePromptProjection, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse,
|
||||
WorkspaceServerOperation,
|
||||
WorkspaceBinaryRequest, WorkspaceBinaryResponse, WorkspaceClient, WorkspaceClientError,
|
||||
WorkspacePromptCatalogResolution, WorkspacePromptProjection, WorkspaceRequest,
|
||||
WorkspaceRequestMethod, WorkspaceResponse, WorkspaceServerOperation,
|
||||
};
|
||||
|
||||
use crate::auth::{
|
||||
@@ -381,6 +383,64 @@ impl RuntimeOwnedWorkspaceClient {
|
||||
self
|
||||
}
|
||||
|
||||
fn execute_binary_with_optional_timeout(
|
||||
&self,
|
||||
request: WorkspaceBinaryRequest,
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
let method = match request.method {
|
||||
WorkspaceRequestMethod::Get => reqwest::Method::GET,
|
||||
WorkspaceRequestMethod::Post => reqwest::Method::POST,
|
||||
WorkspaceRequestMethod::Put => reqwest::Method::PUT,
|
||||
WorkspaceRequestMethod::Patch => reqwest::Method::PATCH,
|
||||
WorkspaceRequestMethod::Delete => reqwest::Method::DELETE,
|
||||
};
|
||||
let mut headers = reqwest::header::HeaderMap::new();
|
||||
if request.body.is_some() {
|
||||
headers.insert(
|
||||
reqwest::header::CONTENT_TYPE,
|
||||
reqwest::header::HeaderValue::from_static("application/octet-stream"),
|
||||
);
|
||||
}
|
||||
let response = self
|
||||
.request_client
|
||||
.execute_binary_blocking(RuntimeWorkspaceRequest {
|
||||
method,
|
||||
path_and_query: request.path,
|
||||
body: request.body.unwrap_or_default(),
|
||||
headers,
|
||||
permission: WORKSPACE_REQUEST_PERMISSION.to_string(),
|
||||
worker_id: Some(self.worker_id.clone()),
|
||||
timeout,
|
||||
max_response_bytes: request.max_response_bytes,
|
||||
})
|
||||
.map_err(|error| {
|
||||
// Never expose request URLs, signed proofs, or binary bodies in transport errors.
|
||||
let message = match error {
|
||||
RuntimeWorkspaceRequestError::InvalidRequest(_) => {
|
||||
"invalid Workspace binary request"
|
||||
}
|
||||
RuntimeWorkspaceRequestError::Sign(_) => {
|
||||
"Workspace binary request authorization failed"
|
||||
}
|
||||
RuntimeWorkspaceRequestError::ResponseTooLarge { .. } => {
|
||||
"Workspace binary response exceeded limit"
|
||||
}
|
||||
RuntimeWorkspaceRequestError::Transport { timeout: true, .. } => {
|
||||
"Workspace binary request timed out"
|
||||
}
|
||||
RuntimeWorkspaceRequestError::Transport { .. } => {
|
||||
"Workspace binary request failed"
|
||||
}
|
||||
};
|
||||
WorkspaceClientError::Request(message.to_string())
|
||||
})?;
|
||||
Ok(WorkspaceBinaryResponse {
|
||||
status: response.status.as_u16(),
|
||||
body: response.body,
|
||||
})
|
||||
}
|
||||
|
||||
fn execute_with_optional_timeout(
|
||||
&self,
|
||||
request: WorkspaceRequest,
|
||||
@@ -481,6 +541,21 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient {
|
||||
self.execute_with_optional_timeout(request, self.request_timeout)
|
||||
}
|
||||
|
||||
fn execute_binary(
|
||||
&self,
|
||||
request: WorkspaceBinaryRequest,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
self.execute_binary_with_optional_timeout(request, self.request_timeout)
|
||||
}
|
||||
|
||||
fn execute_binary_with_timeout(
|
||||
&self,
|
||||
request: WorkspaceBinaryRequest,
|
||||
timeout: Duration,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
self.execute_binary_with_optional_timeout(request, Some(timeout))
|
||||
}
|
||||
|
||||
fn execute_with_timeout(
|
||||
&self,
|
||||
request: WorkspaceRequest,
|
||||
@@ -789,6 +864,216 @@ mod tests {
|
||||
request_body_digest, verify_worker_mutation_source_proof,
|
||||
};
|
||||
|
||||
fn binary_loopback(response: Vec<u8>) -> (String, std::thread::JoinHandle<(String, Vec<u8>)>) {
|
||||
use std::io::{Read, Write};
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let base_url = format!("http://{}", listener.local_addr().unwrap());
|
||||
let server = std::thread::spawn(move || {
|
||||
let (mut stream, _) = listener.accept().unwrap();
|
||||
stream
|
||||
.set_read_timeout(Some(Duration::from_secs(5)))
|
||||
.unwrap();
|
||||
stream
|
||||
.set_write_timeout(Some(Duration::from_secs(5)))
|
||||
.unwrap();
|
||||
let mut head = Vec::new();
|
||||
while !head.ends_with(b"\r\n\r\n") {
|
||||
let mut byte = [0];
|
||||
stream.read_exact(&mut byte).unwrap();
|
||||
head.push(byte[0]);
|
||||
assert!(head.len() < 16 * 1024, "request headers too large");
|
||||
}
|
||||
let head = String::from_utf8(head).unwrap();
|
||||
let length = head
|
||||
.lines()
|
||||
.find_map(|line| {
|
||||
let (name, value) = line.split_once(':')?;
|
||||
name.eq_ignore_ascii_case("content-length")
|
||||
.then(|| value.trim().parse::<usize>().unwrap())
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let mut body = vec![0; length];
|
||||
stream.read_exact(&mut body).unwrap();
|
||||
// Oversize-response tests may close the connection before all bytes are sent.
|
||||
let _ = stream.write_all(&response);
|
||||
(head, body)
|
||||
});
|
||||
(base_url, server)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_roundtrip_signs_exact_query_method_worker_and_raw_body() {
|
||||
use crate::auth::{RuntimeRequestSourceExpectation, verify_runtime_request_source};
|
||||
let identity = RuntimeIdentityMaterial::generate("runtime-a").unwrap();
|
||||
for (method, wire_method, body, reply) in [
|
||||
(
|
||||
WorkspaceRequestMethod::Get,
|
||||
"GET",
|
||||
None,
|
||||
vec![0, 255, 128, 13],
|
||||
),
|
||||
(
|
||||
WorkspaceRequestMethod::Post,
|
||||
"POST",
|
||||
Some(vec![0, 255, 128, 13, 10]),
|
||||
vec![255, 0],
|
||||
),
|
||||
(
|
||||
WorkspaceRequestMethod::Put,
|
||||
"PUT",
|
||||
Some(vec![255; 16 * 1024 * 1024]),
|
||||
vec![0],
|
||||
),
|
||||
] {
|
||||
let path = "/api/w/workspace-a/drive/entries/a/content?offset=0&limit=65536&name=a%2Fb";
|
||||
let mut response = format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
|
||||
reply.len()
|
||||
)
|
||||
.into_bytes();
|
||||
response.extend_from_slice(&reply);
|
||||
let (base_url, server) = binary_loopback(response);
|
||||
let client =
|
||||
RuntimeOwnedWorkspaceClient::new("workspace-a", base_url, "runtime-a", "worker-a")
|
||||
.with_runtime_request_source(&identity, "server-a");
|
||||
let result = client
|
||||
.execute_binary_with_timeout(
|
||||
WorkspaceBinaryRequest {
|
||||
method,
|
||||
path: path.into(),
|
||||
body: body.clone(),
|
||||
max_response_bytes: reply.len(),
|
||||
},
|
||||
Duration::from_secs(5),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(result.status, 200);
|
||||
assert_eq!(result.body, reply);
|
||||
let (head, received) = server.join().unwrap();
|
||||
assert_eq!(received, body.unwrap_or_default());
|
||||
assert!(head.starts_with(&format!("{wire_method} {path} HTTP/1.1\r\n")));
|
||||
let header = |name: &str| {
|
||||
head.lines().find_map(|line| {
|
||||
let (key, value) = line.split_once(':')?;
|
||||
key.eq_ignore_ascii_case(name).then(|| value.trim())
|
||||
})
|
||||
};
|
||||
assert_eq!(header("x-yoi-runtime-id"), Some("runtime-a"));
|
||||
assert_eq!(header("x-yoi-worker-id"), Some("worker-a"));
|
||||
assert_eq!(header("authorization"), None);
|
||||
if method != WorkspaceRequestMethod::Get {
|
||||
assert_eq!(header("content-type"), Some("application/octet-stream"));
|
||||
}
|
||||
let digest = request_body_digest(&received);
|
||||
let expected = RuntimeRequestSourceExpectation {
|
||||
identity_id: "runtime-a",
|
||||
audience: "server-a",
|
||||
workspace_id: "workspace-a",
|
||||
worker_id: Some("worker-a"),
|
||||
permission: WORKSPACE_REQUEST_PERMISSION,
|
||||
method: wire_method,
|
||||
path,
|
||||
body_digest: &digest,
|
||||
now_unix: unix_now_seconds() as i64,
|
||||
};
|
||||
let token = header(RUNTIME_REQUEST_SOURCE_PROOF_HEADER).unwrap();
|
||||
verify_runtime_request_source(token, &identity.public_key, &expected).unwrap();
|
||||
let altered = RuntimeRequestSourceExpectation {
|
||||
body_digest: &request_body_digest(b"altered"),
|
||||
..expected.clone()
|
||||
};
|
||||
assert!(verify_runtime_request_source(token, &identity.public_key, &altered).is_err());
|
||||
let altered = RuntimeRequestSourceExpectation {
|
||||
path: "/api/w/workspace-a/drive/entries/a/content?offset=1",
|
||||
..expected
|
||||
};
|
||||
assert!(verify_runtime_request_source(token, &identity.public_key, &altered).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_responses_enforce_declared_and_streamed_limits_with_fixed_errors() {
|
||||
for response in [
|
||||
b"HTTP/1.1 200 OK\r\nContent-Length: 5\r\nConnection: close\r\n\r\nabcde".to_vec(),
|
||||
b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n2\r\nab\r\n3\r\ncde\r\n0\r\n\r\n".to_vec(),
|
||||
b"HTTP/1.1 500 Error\r\nConnection: close\r\n\r\nabcde".to_vec(),
|
||||
] {
|
||||
let (base_url, server) = binary_loopback(response);
|
||||
let client = RuntimeOwnedWorkspaceClient::new("workspace-a", base_url, "runtime-a", "worker-a");
|
||||
let error = client.execute_binary(WorkspaceBinaryRequest {
|
||||
method: WorkspaceRequestMethod::Get,
|
||||
path: "/api/w/workspace-a/drive/entries/a/content?private=secret".into(),
|
||||
body: None, max_response_bytes: 4,
|
||||
}).unwrap_err();
|
||||
server.join().unwrap();
|
||||
assert_eq!(error.to_string(), "workspace request failed: Workspace binary response exceeded limit");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_redirect_returns_status_without_forwarding_authentication() {
|
||||
let (base_url, server) = binary_loopback(b"HTTP/1.1 307 Temporary Redirect\r\nLocation: http://127.0.0.1:1/private\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_vec());
|
||||
let client =
|
||||
RuntimeOwnedWorkspaceClient::new("workspace-a", base_url, "runtime-a", "worker-a");
|
||||
let result = client
|
||||
.execute_binary(WorkspaceBinaryRequest {
|
||||
method: WorkspaceRequestMethod::Post,
|
||||
path: "/api/w/workspace-a/drive/upload".into(),
|
||||
body: Some(vec![0, 255]),
|
||||
max_response_bytes: 0,
|
||||
})
|
||||
.unwrap();
|
||||
server.join().unwrap();
|
||||
assert_eq!(result.status, 307);
|
||||
assert!(result.body.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_invalid_or_normalized_targets_fail_with_fixed_errors_before_http() {
|
||||
let client = RuntimeOwnedWorkspaceClient::new(
|
||||
"workspace-a",
|
||||
"http://127.0.0.1:1",
|
||||
"runtime-a",
|
||||
"worker-a",
|
||||
);
|
||||
for path in [
|
||||
"https://example.com/secret",
|
||||
"//example.com/secret",
|
||||
"/a/../secret",
|
||||
"/a#secret",
|
||||
"/a\\secret",
|
||||
"/a b",
|
||||
] {
|
||||
let error = client
|
||||
.execute_binary(WorkspaceBinaryRequest {
|
||||
method: WorkspaceRequestMethod::Get,
|
||||
path: path.into(),
|
||||
body: None,
|
||||
max_response_bytes: 4,
|
||||
})
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
error.to_string(),
|
||||
"workspace request failed: invalid Workspace binary request"
|
||||
);
|
||||
}
|
||||
let error = client
|
||||
.execute_binary_with_timeout(
|
||||
WorkspaceBinaryRequest {
|
||||
method: WorkspaceRequestMethod::Post,
|
||||
path: "/secret?private=value".into(),
|
||||
body: Some(vec![255]),
|
||||
max_response_bytes: 4,
|
||||
},
|
||||
Duration::from_secs(1),
|
||||
)
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
error.to_string(),
|
||||
"workspace request failed: Workspace binary request failed"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn current_prompt_projection_uses_the_shared_runtime_cache_without_http() {
|
||||
let cache = Arc::new(WorkspacePromptProjectionCache::default());
|
||||
|
||||
@@ -255,7 +255,7 @@ impl RuntimeWorkspaceRequestClient {
|
||||
request: RuntimeWorkspaceRequest,
|
||||
) -> Result<RuntimeWorkspaceResponse, RuntimeWorkspaceRequestError> {
|
||||
let client = self.clone();
|
||||
std::thread::spawn(move || client.execute_blocking_inner(request))
|
||||
std::thread::spawn(move || client.execute_blocking_inner(request, true))
|
||||
.join()
|
||||
.map_err(|_| RuntimeWorkspaceRequestError::Transport {
|
||||
message: "Workspace request thread panicked".to_string(),
|
||||
@@ -263,15 +263,45 @@ impl RuntimeWorkspaceRequestClient {
|
||||
})?
|
||||
}
|
||||
|
||||
fn execute_blocking_inner(
|
||||
/// Binary requests must retain their exact target and never follow redirects.
|
||||
pub(crate) fn execute_binary_blocking(
|
||||
&self,
|
||||
request: RuntimeWorkspaceRequest,
|
||||
) -> Result<RuntimeWorkspaceResponse, RuntimeWorkspaceRequestError> {
|
||||
let client = self.clone();
|
||||
std::thread::spawn(move || client.execute_blocking_inner(request, false))
|
||||
.join()
|
||||
.map_err(|_| RuntimeWorkspaceRequestError::Transport {
|
||||
message: "Workspace binary request thread panicked".to_string(),
|
||||
timeout: false,
|
||||
})?
|
||||
}
|
||||
|
||||
fn execute_blocking_inner(
|
||||
&self,
|
||||
request: RuntimeWorkspaceRequest,
|
||||
follow_redirects: bool,
|
||||
) -> Result<RuntimeWorkspaceResponse, RuntimeWorkspaceRequestError> {
|
||||
let prepared = self.prepare(&request)?;
|
||||
if !follow_redirects {
|
||||
let mut target = prepared.url.path().to_string();
|
||||
if let Some(query) = prepared.url.query() {
|
||||
target.push('?');
|
||||
target.push_str(query);
|
||||
}
|
||||
if prepared.url.fragment().is_some() || target != request.path_and_query {
|
||||
return Err(RuntimeWorkspaceRequestError::InvalidRequest(
|
||||
"binary request target must be exact".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
// The blocking client defaults to 30 seconds. Apply None explicitly so
|
||||
// provider-owned command waits are not cut short by an HTTP deadline.
|
||||
let client = reqwest::blocking::Client::builder()
|
||||
.timeout(request.timeout)
|
||||
let mut client_builder = reqwest::blocking::Client::builder().timeout(request.timeout);
|
||||
if !follow_redirects {
|
||||
client_builder = client_builder.redirect(reqwest::redirect::Policy::none());
|
||||
}
|
||||
let client = client_builder
|
||||
.build()
|
||||
.map_err(RuntimeWorkspaceRequestError::transport)?;
|
||||
let mut builder = client
|
||||
|
||||
@@ -65,9 +65,9 @@ pub use session_history::{
|
||||
pub use shared_state::WorkerSharedState;
|
||||
pub use worker::{
|
||||
LocalWorkingDirectory, SubjektivSessionAttributionLifecycle, Worker, WorkerError,
|
||||
WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient,
|
||||
WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspacePromptCatalogResolution,
|
||||
WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, WorkspaceServerOperation,
|
||||
WorkspaceWorkerDiscoveryRequest, apply_worker_manifest, marker_workspace_client,
|
||||
unavailable_workspace_client,
|
||||
WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceBinaryRequest,
|
||||
WorkspaceBinaryResponse, WorkspaceClient, WorkspaceClientError, WorkspaceId, WorkspaceIdError,
|
||||
WorkspacePromptCatalogResolution, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse,
|
||||
WorkspaceServerOperation, WorkspaceWorkerDiscoveryRequest, apply_worker_manifest,
|
||||
marker_workspace_client, unavailable_workspace_client,
|
||||
};
|
||||
|
||||
+334
-25
@@ -14,7 +14,9 @@ use std::sync::{Arc, Mutex};
|
||||
use agen::Engine;
|
||||
use agen::llm_client::client::LlmClient;
|
||||
use agen::state::Mutable;
|
||||
use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput};
|
||||
use agen::tool::{
|
||||
Attachment, Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use base64::Engine as _;
|
||||
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
|
||||
@@ -99,8 +101,8 @@ pub struct WipCallContext {
|
||||
pub security_context: String,
|
||||
}
|
||||
|
||||
/// Successful operation value. Compatibility projections can retain an ordinary
|
||||
/// ToolOutput so images and pruning semantics remain on the existing Engine path.
|
||||
/// Successful operation value. Native attachments and compatibility ToolOutputs
|
||||
/// use the existing Engine side channel, separate from the validated WIP value.
|
||||
pub struct WipOperationOutput {
|
||||
pub value: Value,
|
||||
tool_output: Option<ToolOutput>,
|
||||
@@ -125,6 +127,20 @@ impl WipOperationOutput {
|
||||
}
|
||||
}
|
||||
|
||||
/// Attach model-visible binary content without replacing the native result
|
||||
/// or its exact post-operation validator. Invoke returns these attachments
|
||||
/// through ToolOutput so Engine capture/history can persist their bytes.
|
||||
pub fn with_attachments(mut self, attachments: Vec<Attachment>) -> Self {
|
||||
self.tool_output
|
||||
.get_or_insert_with(|| ToolOutput {
|
||||
summary: String::new(),
|
||||
content: None,
|
||||
attachments: Vec::new(),
|
||||
})
|
||||
.attachments = attachments;
|
||||
self
|
||||
}
|
||||
|
||||
fn compatibility(value: Value, output: ToolOutput) -> Self {
|
||||
Self {
|
||||
value,
|
||||
@@ -134,6 +150,13 @@ impl WipOperationOutput {
|
||||
}
|
||||
}
|
||||
|
||||
enum WipToolOutput {
|
||||
/// Only binary attachments accompany the validated native completion.
|
||||
Native(ToolOutput),
|
||||
/// Existing compatibility/error output retains its summary and pruning policy.
|
||||
Passthrough(Result<ToolOutput, ToolError>),
|
||||
}
|
||||
|
||||
pub enum WipOperationError {
|
||||
Protocol(ProtocolError),
|
||||
Cancelled(ToolOutput),
|
||||
@@ -2227,15 +2250,22 @@ impl WipRuntime {
|
||||
guard.disarm();
|
||||
let completion = completion.map_err(client_tool_error)?;
|
||||
self.push_audit(prepared.id, path, operation, audit_outcome);
|
||||
if let Some(result) = passthrough {
|
||||
if result.is_ok() {
|
||||
self.metrics
|
||||
.operation_successes
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
let mut attachments = Vec::new();
|
||||
if let Some(passthrough) = passthrough {
|
||||
match passthrough {
|
||||
WipToolOutput::Passthrough(result) => {
|
||||
if result.is_ok() {
|
||||
self.metrics
|
||||
.operation_successes
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
WipToolOutput::Native(output) => attachments = output.attachments,
|
||||
}
|
||||
return result;
|
||||
}
|
||||
render_call_completion(completion).map(|output| {
|
||||
render_call_completion(completion).map(|mut output| {
|
||||
output.attachments = attachments;
|
||||
self.metrics
|
||||
.operation_successes
|
||||
.fetch_add(1, Ordering::Relaxed);
|
||||
@@ -2338,14 +2368,7 @@ impl WipRuntime {
|
||||
&self,
|
||||
request: &Request<Vec<u8>>,
|
||||
execution: ToolExecutionContext,
|
||||
) -> Result<
|
||||
(
|
||||
Response<Vec<u8>>,
|
||||
Option<Result<ToolOutput, ToolError>>,
|
||||
WipAuditOutcome,
|
||||
),
|
||||
ToolError,
|
||||
> {
|
||||
) -> Result<(Response<Vec<u8>>, Option<WipToolOutput>, WipAuditOutcome), ToolError> {
|
||||
let rejected = |error: ProtocolError| -> Result<_, ToolError> {
|
||||
Ok((
|
||||
encode_protocol_error_response(
|
||||
@@ -2385,6 +2408,7 @@ impl WipRuntime {
|
||||
return rejected(error);
|
||||
}
|
||||
let projection = projection.projection;
|
||||
let kind = projection.kind;
|
||||
let descriptor = projection.descriptor.clone();
|
||||
let target_validator = projection.object.validator.clone();
|
||||
let request_value =
|
||||
@@ -2443,8 +2467,10 @@ impl WipRuntime {
|
||||
.map_err(|error| ToolError::Internal(error.to_string()))?;
|
||||
return Ok((
|
||||
response,
|
||||
Some(Err(ToolError::ExecutionFailed(format!(
|
||||
"WIP operation outcome unknown; do not retry automatically: {message}"
|
||||
Some(WipToolOutput::Passthrough(Err(ToolError::ExecutionFailed(
|
||||
format!(
|
||||
"WIP operation outcome unknown; do not retry automatically: {message}"
|
||||
),
|
||||
)))),
|
||||
WipAuditOutcome::OutcomeUnknown,
|
||||
));
|
||||
@@ -2452,7 +2478,10 @@ impl WipRuntime {
|
||||
};
|
||||
Ok((
|
||||
response,
|
||||
output.tool_output.map(Ok),
|
||||
output.tool_output.map(|output| match kind {
|
||||
WipProjectionKind::Native => WipToolOutput::Native(output),
|
||||
WipProjectionKind::Compatibility => WipToolOutput::Passthrough(Ok(output)),
|
||||
}),
|
||||
WipAuditOutcome::Success,
|
||||
))
|
||||
}
|
||||
@@ -2477,7 +2506,9 @@ impl WipRuntime {
|
||||
.map_err(|encode| ToolError::Internal(encode.to_string()))?;
|
||||
Ok((
|
||||
response,
|
||||
Some(Err(ToolError::Cancelled(output))),
|
||||
Some(WipToolOutput::Passthrough(Err(ToolError::Cancelled(
|
||||
output,
|
||||
)))),
|
||||
WipAuditOutcome::Cancelled,
|
||||
))
|
||||
}
|
||||
@@ -2493,7 +2524,9 @@ impl WipRuntime {
|
||||
.map_err(|encode| ToolError::Internal(encode.to_string()))?;
|
||||
Ok((
|
||||
response,
|
||||
Some(Err(ToolError::Interrupted(output))),
|
||||
Some(WipToolOutput::Passthrough(Err(ToolError::Interrupted(
|
||||
output,
|
||||
)))),
|
||||
WipAuditOutcome::Interrupted,
|
||||
))
|
||||
}
|
||||
@@ -2506,8 +2539,10 @@ impl WipRuntime {
|
||||
.map_err(|encode| ToolError::Internal(encode.to_string()))?;
|
||||
Ok((
|
||||
response,
|
||||
Some(Err(ToolError::ExecutionFailed(format!(
|
||||
"WIP operation outcome unknown; do not retry automatically: {message}"
|
||||
Some(WipToolOutput::Passthrough(Err(ToolError::ExecutionFailed(
|
||||
format!(
|
||||
"WIP operation outcome unknown; do not retry automatically: {message}"
|
||||
),
|
||||
)))),
|
||||
WipAuditOutcome::OutcomeUnknown,
|
||||
))
|
||||
@@ -5072,6 +5107,280 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
struct NativeImage {
|
||||
invalid_result: bool,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl WipOperationHandler for NativeImage {
|
||||
async fn call(
|
||||
&self,
|
||||
_operation: &str,
|
||||
_arguments: &BTreeMap<String, Value>,
|
||||
_context: WipCallContext,
|
||||
) -> Result<WipOperationOutput, WipOperationError> {
|
||||
let value = if self.invalid_result {
|
||||
Value::Integer(1)
|
||||
} else {
|
||||
Value::String("Attached native image".into())
|
||||
};
|
||||
Ok(
|
||||
WipOperationOutput::native_with_validator(value, vec![7]).with_attachments(vec![
|
||||
Attachment::Image(agen::tool::ImageAttachment::new(
|
||||
"image/png",
|
||||
include_bytes!("../../tools/tests/fixtures/view-image/valid.png").to_vec(),
|
||||
)),
|
||||
]),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn native_image_registry(invalid_result: bool) -> WipMountRegistry {
|
||||
let mut registry = WipMountRegistry::new();
|
||||
registry.allocate_namespace("native", "native").unwrap();
|
||||
registry
|
||||
.mount(WipProjection {
|
||||
route: "/native/image".into(),
|
||||
capability: "native:image".into(),
|
||||
kind: WipProjectionKind::Native,
|
||||
object: Object {
|
||||
name: "image".into(),
|
||||
description: None,
|
||||
interfaces: vec![root_reference("yoi.native/image/v1")],
|
||||
r#ref: Some("native:image".into()),
|
||||
validator: Some(vec![1]),
|
||||
},
|
||||
interface: root_reference("yoi.native/image/v1"),
|
||||
descriptor: InterfaceDescriptor {
|
||||
format: INTERFACE_FORMAT_V1.into(),
|
||||
documentation: None,
|
||||
types: Vec::new(),
|
||||
operations: vec![OperationDeclaration {
|
||||
name: "view".into(),
|
||||
documentation: None,
|
||||
parameters: Vec::new(),
|
||||
returns: ReturnDeclaration {
|
||||
documentation: None,
|
||||
r#type: TypeExpr::String,
|
||||
},
|
||||
}],
|
||||
},
|
||||
interface_validator: Some(vec![1]),
|
||||
handler: Arc::new(NativeImage { invalid_result }),
|
||||
})
|
||||
.unwrap();
|
||||
registry
|
||||
}
|
||||
|
||||
fn native_image_invoke() -> String {
|
||||
json!({
|
||||
"path":"/native/image",
|
||||
"interface":{"scope":"/","name":"yoi.native/image/v1"},
|
||||
"operation":"view",
|
||||
"arguments":{}
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_invoke_returns_image_bytes_and_preserves_result_and_validator() {
|
||||
let mut engine = Engine::<_, Mutable, ()>::new_annotated(DummyClient);
|
||||
let runtime = install_wip_mode_with_mounts(
|
||||
&mut engine,
|
||||
None,
|
||||
"worker-a".into(),
|
||||
native_image_registry(false),
|
||||
)
|
||||
.unwrap();
|
||||
let tools = engine.tool_server_handle();
|
||||
tools.flush_pending();
|
||||
tools
|
||||
.call_tool(
|
||||
"Inspect",
|
||||
r#"{"path":"/native/image"}"#,
|
||||
ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let output = tools
|
||||
.call_tool(
|
||||
"Invoke",
|
||||
&native_image_invoke(),
|
||||
ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(output.summary, "WIP operation completed");
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Json>(output.content.as_deref().unwrap()).unwrap(),
|
||||
json!("Attached native image")
|
||||
);
|
||||
let [Attachment::Image(image)] = output.attachments.as_slice() else {
|
||||
panic!("Invoke must return an image, not a URL")
|
||||
};
|
||||
assert_eq!(image.mime_type(), "image/png");
|
||||
assert_eq!(
|
||||
image.data(),
|
||||
include_bytes!("../../tools/tests/fixtures/view-image/valid.png")
|
||||
);
|
||||
let state = runtime.state.lock().unwrap();
|
||||
let calls = state.client.call_history(&state.session).unwrap();
|
||||
assert!(
|
||||
matches!(&calls.back().unwrap().outcome, CallOutcome::Success(response) if response.result == Value::String("Attached native image".into()) && response.validator == Some(vec![7]))
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_attachments_do_not_bypass_return_value_validation() {
|
||||
let mut engine = Engine::<_, Mutable, ()>::new_annotated(DummyClient);
|
||||
install_wip_mode_with_mounts(
|
||||
&mut engine,
|
||||
None,
|
||||
"worker-a".into(),
|
||||
native_image_registry(true),
|
||||
)
|
||||
.unwrap();
|
||||
let tools = engine.tool_server_handle();
|
||||
tools.flush_pending();
|
||||
tools
|
||||
.call_tool(
|
||||
"Inspect",
|
||||
r#"{"path":"/native/image"}"#,
|
||||
ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let error = tools
|
||||
.call_tool(
|
||||
"Invoke",
|
||||
&native_image_invoke(),
|
||||
ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.to_string().contains("outcome unknown"), "{error}");
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct NativeImageClient {
|
||||
events: Arc<Mutex<VecDeque<Vec<agen::llm_client::event::Event>>>>,
|
||||
requests: Arc<Mutex<Vec<LlmRequest>>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LlmClient for NativeImageClient {
|
||||
fn clone_boxed(&self) -> Box<dyn LlmClient> {
|
||||
Box::new(self.clone())
|
||||
}
|
||||
async fn stream(&self, request: LlmRequest) -> Result<ResponseStream, ClientError> {
|
||||
self.requests.lock().unwrap().push(request);
|
||||
let events = self
|
||||
.events
|
||||
.lock()
|
||||
.unwrap()
|
||||
.pop_front()
|
||||
.expect("unexpected model request");
|
||||
Ok(Box::pin(stream::iter(events.into_iter().map(Ok))))
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_invoke_image_survives_engine_capture_and_durable_history_replay() {
|
||||
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
|
||||
let completed = || {
|
||||
Event::Status(StatusEvent {
|
||||
status: ResponseStatus::Completed,
|
||||
})
|
||||
};
|
||||
let client = NativeImageClient {
|
||||
events: Arc::new(Mutex::new(VecDeque::from([
|
||||
vec![
|
||||
Event::tool_use_start(0, "native-image-call", "Invoke"),
|
||||
Event::tool_input_delta(0, &native_image_invoke()),
|
||||
Event::tool_use_stop(0),
|
||||
completed(),
|
||||
],
|
||||
vec![completed()],
|
||||
]))),
|
||||
requests: Arc::new(Mutex::new(Vec::new())),
|
||||
};
|
||||
let mut engine = Engine::new(client.clone());
|
||||
install_wip_mode_with_mounts(
|
||||
&mut engine,
|
||||
None,
|
||||
"worker-a".into(),
|
||||
native_image_registry(false),
|
||||
)
|
||||
.unwrap();
|
||||
let tools = engine.tool_server_handle();
|
||||
tools.flush_pending();
|
||||
tools
|
||||
.call_tool(
|
||||
"Inspect",
|
||||
r#"{"path":"/native/image"}"#,
|
||||
ToolExecutionContext::direct(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let captured = Arc::new(Mutex::new(Vec::new()));
|
||||
let sink = Arc::clone(&captured);
|
||||
engine.on_history_append(move |item| {
|
||||
sink.lock().unwrap().push(item.clone());
|
||||
Ok(())
|
||||
});
|
||||
let mut history = agen::History::new();
|
||||
assert!(matches!(
|
||||
engine
|
||||
.run(&mut history, "view the native image")
|
||||
.await
|
||||
.result,
|
||||
agen::EngineRunExit::Finished
|
||||
));
|
||||
let saved = history.items().find(|item| matches!(item, agen::Item::ToolResult { call_id, .. } if call_id == "native-image-call")).unwrap();
|
||||
let agen::Item::ToolResult {
|
||||
attachments,
|
||||
is_error,
|
||||
content,
|
||||
..
|
||||
} = saved
|
||||
else {
|
||||
unreachable!()
|
||||
};
|
||||
assert!(!is_error);
|
||||
let [Attachment::Image(image)] = attachments.as_slice() else {
|
||||
panic!("durable Invoke result lost image")
|
||||
};
|
||||
assert_eq!(
|
||||
image.data(),
|
||||
include_bytes!("../../tools/tests/fixtures/view-image/valid.png")
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Json>(content.as_deref().unwrap()).unwrap(),
|
||||
json!("Attached native image")
|
||||
);
|
||||
assert!(captured.lock().unwrap().iter().any(|item| item == saved));
|
||||
assert!(
|
||||
client.requests.lock().unwrap()[1]
|
||||
.items
|
||||
.iter()
|
||||
.any(|item| item == saved)
|
||||
);
|
||||
let log = session_store::LogEntry::AnnotatedToolResult {
|
||||
ts: 1,
|
||||
entry: session_store::LoggedHistoryEntry {
|
||||
item: session_store::LoggedItem::from(saved),
|
||||
metadata: session_store::LoggedSessionHistoryMetadata {
|
||||
entry_id: session_store::LoggedSessionHistoryEntryId::new(),
|
||||
origin: session_store::LoggedSessionHistoryOrigin::LegacyUnknown,
|
||||
derivation: None,
|
||||
},
|
||||
},
|
||||
};
|
||||
let serialized = serde_json::to_string(&log).unwrap();
|
||||
let restored: session_store::LogEntry = serde_json::from_str(&serialized).unwrap();
|
||||
assert_eq!(&session_store::collect_state(&[restored]).history[0], saved);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn native_projection_preserves_descriptor_and_decodes_typed_arguments_end_to_end() {
|
||||
let mut registry = WipMountRegistry::new();
|
||||
|
||||
@@ -639,6 +639,27 @@ impl WorkspaceRequest {
|
||||
}
|
||||
}
|
||||
|
||||
/// Raw-byte request through a host-bound Workspace client, separate from JSON requests.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct WorkspaceBinaryRequest {
|
||||
pub method: WorkspaceRequestMethod,
|
||||
pub path: String,
|
||||
pub body: Option<Vec<u8>>,
|
||||
pub max_response_bytes: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct WorkspaceBinaryResponse {
|
||||
pub status: u16,
|
||||
pub body: Vec<u8>,
|
||||
}
|
||||
|
||||
impl WorkspaceBinaryResponse {
|
||||
pub fn is_success(&self) -> bool {
|
||||
(200..300).contains(&self.status)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum WorkspaceRequestMethod {
|
||||
Get,
|
||||
@@ -859,6 +880,24 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync {
|
||||
self.execute(request)
|
||||
}
|
||||
|
||||
/// Binary transport is unavailable unless a host-bound implementation opts in.
|
||||
fn execute_binary(
|
||||
&self,
|
||||
_request: WorkspaceBinaryRequest,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
Err(WorkspaceClientError::Unavailable(
|
||||
"Workspace binary transport is unavailable".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
fn execute_binary_with_timeout(
|
||||
&self,
|
||||
request: WorkspaceBinaryRequest,
|
||||
_timeout: Duration,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
self.execute_binary(request)
|
||||
}
|
||||
|
||||
fn execute_server_operation(
|
||||
&self,
|
||||
operation: WorkspaceServerOperation,
|
||||
@@ -961,6 +1000,31 @@ impl WorkspaceClient for ReviewerChildWorkspaceClient {
|
||||
self.inner.current_prompt_projection(minimum_revision)
|
||||
}
|
||||
|
||||
fn execute_binary(
|
||||
&self,
|
||||
request: WorkspaceBinaryRequest,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
if request.method != WorkspaceRequestMethod::Get {
|
||||
return Err(WorkspaceClientError::Unavailable(
|
||||
"Reviewer child binary Workspace authority is read-only".to_string(),
|
||||
));
|
||||
}
|
||||
self.inner.execute_binary(request)
|
||||
}
|
||||
|
||||
fn execute_binary_with_timeout(
|
||||
&self,
|
||||
request: WorkspaceBinaryRequest,
|
||||
timeout: Duration,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
if request.method != WorkspaceRequestMethod::Get {
|
||||
return Err(WorkspaceClientError::Unavailable(
|
||||
"Reviewer child binary Workspace authority is read-only".to_string(),
|
||||
));
|
||||
}
|
||||
self.inner.execute_binary_with_timeout(request, timeout)
|
||||
}
|
||||
|
||||
fn execute(
|
||||
&self,
|
||||
mut request: WorkspaceRequest,
|
||||
@@ -1157,6 +1221,123 @@ mod reviewer_client_tests {
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn binary_transport_defaults_to_denial_without_json_fallback() {
|
||||
let client = RecordingWorkspaceClient::default();
|
||||
let request = WorkspaceBinaryRequest {
|
||||
method: WorkspaceRequestMethod::Get,
|
||||
path: "/api/w/ws/drive/entries/a/content".into(),
|
||||
body: None,
|
||||
max_response_bytes: 64 * 1024,
|
||||
};
|
||||
assert!(matches!(
|
||||
client.execute_binary(request.clone()),
|
||||
Err(WorkspaceClientError::Unavailable(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
client.execute_binary_with_timeout(request, Duration::from_secs(1)),
|
||||
Err(WorkspaceClientError::Unavailable(_))
|
||||
));
|
||||
assert!(client.requests.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct BinaryRecordingClient {
|
||||
requests: Mutex<Vec<(WorkspaceBinaryRequest, Option<Duration>)>>,
|
||||
}
|
||||
|
||||
impl WorkspaceClient for BinaryRecordingClient {
|
||||
fn workspace_id(&self) -> Option<&str> {
|
||||
Some("ws")
|
||||
}
|
||||
fn kind(&self) -> &str {
|
||||
"binary-recording"
|
||||
}
|
||||
fn is_available(&self) -> bool {
|
||||
true
|
||||
}
|
||||
fn execute(&self, _: WorkspaceRequest) -> Result<WorkspaceResponse, WorkspaceClientError> {
|
||||
panic!("binary request must not use JSON fallback")
|
||||
}
|
||||
fn execute_binary(
|
||||
&self,
|
||||
request: WorkspaceBinaryRequest,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
self.requests.lock().unwrap().push((request, None));
|
||||
Ok(WorkspaceBinaryResponse {
|
||||
status: 200,
|
||||
body: vec![0, 255],
|
||||
})
|
||||
}
|
||||
fn execute_binary_with_timeout(
|
||||
&self,
|
||||
request: WorkspaceBinaryRequest,
|
||||
timeout: Duration,
|
||||
) -> Result<WorkspaceBinaryResponse, WorkspaceClientError> {
|
||||
self.requests.lock().unwrap().push((request, Some(timeout)));
|
||||
Ok(WorkspaceBinaryResponse {
|
||||
status: 200,
|
||||
body: vec![0, 255],
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reviewer_binary_get_preserves_request_limit_and_timeout() {
|
||||
let inner = Arc::new(BinaryRecordingClient::default());
|
||||
let client = reviewer_client(inner.clone());
|
||||
let request = WorkspaceBinaryRequest {
|
||||
method: WorkspaceRequestMethod::Get,
|
||||
path: "/api/w/ws/drive/entries/a/content?offset=2&limit=4".into(),
|
||||
body: None,
|
||||
max_response_bytes: 4,
|
||||
};
|
||||
assert_eq!(
|
||||
client.execute_binary(request.clone()).unwrap().body,
|
||||
vec![0, 255]
|
||||
);
|
||||
let timeout = Duration::from_secs(2);
|
||||
assert_eq!(
|
||||
client
|
||||
.execute_binary_with_timeout(request.clone(), timeout)
|
||||
.unwrap()
|
||||
.body,
|
||||
vec![0, 255]
|
||||
);
|
||||
assert_eq!(
|
||||
*inner.requests.lock().unwrap(),
|
||||
vec![(request.clone(), None), (request, Some(timeout))]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reviewer_binary_mutations_are_denied_even_for_attested_review_path() {
|
||||
let inner = Arc::new(BinaryRecordingClient::default());
|
||||
let client = reviewer_client(inner.clone());
|
||||
for method in [
|
||||
WorkspaceRequestMethod::Post,
|
||||
WorkspaceRequestMethod::Put,
|
||||
WorkspaceRequestMethod::Patch,
|
||||
WorkspaceRequestMethod::Delete,
|
||||
] {
|
||||
let request = WorkspaceBinaryRequest {
|
||||
method,
|
||||
path: "/api/w/ws/merge-requests/MR1/reviews".into(),
|
||||
body: Some(vec![0, 255]),
|
||||
max_response_bytes: 4,
|
||||
};
|
||||
assert!(matches!(
|
||||
client.execute_binary(request.clone()),
|
||||
Err(WorkspaceClientError::Unavailable(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
client.execute_binary_with_timeout(request, Duration::from_secs(1)),
|
||||
Err(WorkspaceClientError::Unavailable(_))
|
||||
));
|
||||
}
|
||||
assert!(inner.requests.lock().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reviewer_child_client_allows_typed_ticket_reads() {
|
||||
let inner = Arc::new(RecordingWorkspaceClient::default());
|
||||
|
||||
Reference in New Issue
Block a user