diff --git a/crates/tools/src/lib.rs b/crates/tools/src/lib.rs index 790dabcf3..8df330c17 100644 --- a/crates/tools/src/lib.rs +++ b/crates/tools/src/lib.rs @@ -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; diff --git a/crates/tools/src/view_image.rs b/crates/tools/src/view_image.rs index 255695b9a..d6448410e 100644 --- a/crates/tools/src/view_image.rs +++ b/crates/tools/src/view_image.rs @@ -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, 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) -> 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>, + } + + #[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 { + 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 { + self.local.stat(r).await + } + async fn read(&self, r: ReadRequest) -> Result { + self.local.read(r).await + } + async fn write(&self, r: WriteRequest) -> Result { + self.local.write(r).await + } + async fn edit(&self, r: EditRequest) -> Result { + self.local.edit(r).await + } + async fn list(&self, r: ListRequest) -> Result { + self.local.list(r).await + } + async fn glob(&self, r: GlobRequest) -> Result { + self.local.glob(r).await + } + async fn grep(&self, r: GrepRequest) -> Result { + self.local.grep(r).await + } + async fn start_command(&self, r: CommandRequest) -> Result { + self.local.start_command(r).await + } + async fn command_status(&self, r: CommandHandle) -> Result { + self.local.command_status(r).await + } + async fn command_output( + &self, + r: CommandOutputRequest, + ) -> Result { + 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) -> 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() { diff --git a/crates/worker-runtime/src/worker_source.rs b/crates/worker-runtime/src/worker_source.rs index 22f81aa79..43bda3fb0 100644 --- a/crates/worker-runtime/src/worker_source.rs +++ b/crates/worker-runtime/src/worker_source.rs @@ -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, + ) -> Result { + 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 { + self.execute_binary_with_optional_timeout(request, self.request_timeout) + } + + fn execute_binary_with_timeout( + &self, + request: WorkspaceBinaryRequest, + timeout: Duration, + ) -> Result { + 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) -> (String, std::thread::JoinHandle<(String, Vec)>) { + 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::().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()); diff --git a/crates/worker-runtime/src/workspace_request.rs b/crates/worker-runtime/src/workspace_request.rs index 9d71442a9..1eaa95490 100644 --- a/crates/worker-runtime/src/workspace_request.rs +++ b/crates/worker-runtime/src/workspace_request.rs @@ -255,7 +255,7 @@ impl RuntimeWorkspaceRequestClient { request: RuntimeWorkspaceRequest, ) -> Result { 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 { + 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 { 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 diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 069938d1f..a1cfa7c16 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -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, }; diff --git a/crates/worker/src/wip.rs b/crates/worker/src/wip.rs index 3a3687f90..e193957e1 100644 --- a/crates/worker/src/wip.rs +++ b/crates/worker/src/wip.rs @@ -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, @@ -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) -> 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), +} + 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>, execution: ToolExecutionContext, - ) -> Result< - ( - Response>, - Option>, - WipAuditOutcome, - ), - ToolError, - > { + ) -> Result<(Response>, Option, 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, + _context: WipCallContext, + ) -> Result { + 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::(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>>>, + requests: Arc>>, + } + + #[async_trait] + impl LlmClient for NativeImageClient { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + async fn stream(&self, request: LlmRequest) -> Result { + 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::(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(); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index a4bae5f46..ca5fde341 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -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>, + pub max_response_bytes: usize, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkspaceBinaryResponse { + pub status: u16, + pub body: Vec, +} + +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 { + Err(WorkspaceClientError::Unavailable( + "Workspace binary transport is unavailable".to_string(), + )) + } + + fn execute_binary_with_timeout( + &self, + request: WorkspaceBinaryRequest, + _timeout: Duration, + ) -> Result { + 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 { + 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 { + 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)>>, + } + + 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 { + panic!("binary request must not use JSON fallback") + } + fn execute_binary( + &self, + request: WorkspaceBinaryRequest, + ) -> Result { + 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 { + 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());