fix: add bounded binary reads for ViewImage

crates/tools/src/view_image.rs

@@ -7,12 +7,14 @@ use agen::tool::{
 };
 use async_trait::async_trait;
 use serde::Deserialize;
-use workdir::{ReadRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter};
+use workdir::{ReadBytesRequest, WorkdirPath, WorkdirSessionHandle, WorkdirSessionRouter};
 
 use crate::error::ToolsError;
 
 /// Maximum image body accepted for one model request.
 pub const MAX_IMAGE_BYTES: usize = 10 * 1024 * 1024;
+/// One binary response remains within the External Workdir transport bound.
+const IMAGE_READ_CHUNK_BYTES: usize = 1024 * 1024;
 
 const DESCRIPTION: &str = "Attach an image from the selected Workdir attachment to the next model request. \
 The path must be logical and Workdir-relative. Supported formats: PNG, JPEG, GIF, and WebP. \
@@ -47,42 +49,111 @@ impl Tool for ViewImageTool {
             workdir::WorkdirSessionCapability::Read,
         )?;
         let path = WorkdirPath::new(&input.path).map_err(ToolsError::from)?;
-        let result = selected
-            .session
-            .read(ReadRequest {
-                path: path.clone(),
-                offset: 0,
-                limit: usize::MAX,
-                // The scoped provider enforces this cap while reading, rather
-                // than allocating an unbounded binary body first.
-                max_bytes: MAX_IMAGE_BYTES + 1,
-            })
+        let bytes = read_image_bytes(selected.session.as_ref(), path.clone())
             .await
             .map_err(ToolsError::from)?;
-
-        if result.truncated || result.bytes.len() > MAX_IMAGE_BYTES {
-            return Err(ToolError::InvalidArgument(format!(
-                "image exceeds the {MAX_IMAGE_BYTES}-byte limit"
-            )));
-        }
-        let mime_type = detect_image_mime(&result.bytes).ok_or_else(|| {
+        let mime_type = detect_image_mime(&bytes).ok_or_else(|| {
             ToolError::InvalidArgument(
                 "unsupported image; expected PNG, JPEG, GIF, or WebP bytes".to_string(),
             )
         })?;
-        let bytes = result.bytes.len();
+        let byte_count = bytes.len();
 
         Ok(ToolOutput {
-            summary: format!("Attached image {path} ({mime_type}, {bytes} bytes)"),
+            summary: format!("Attached image {path} ({mime_type}, {byte_count} bytes)"),
             content: None,
             attachments: vec![Attachment::Image(ImageAttachment::new(
                 mime_type,
-                Arc::<[u8]>::from(result.bytes),
+                Arc::<[u8]>::from(bytes),
             ))],
         })
     }
 }
 
+async fn read_image_bytes(
+    session: &dyn workdir::WorkdirSession,
+    path: WorkdirPath,
+) -> Result<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))
 }

crates/tools/tests/integration.rs

@@ -399,6 +399,224 @@ async fn detach_and_reattach_same_alias_requires_a_fresh_read() {
     router.close_all().await.unwrap();
 }
 
+#[derive(Debug)]
+struct ExternalValidatingSession {
+    inner: LocalWorkdirSession,
+    truncate_first_binary_response: bool,
+    mutate_after_first_binary_response: Option<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 = b"\x89PNG\r\n\x1a\nexternal\0image\nbody\xff";
+    std::fs::write(dir.path().join("image.png"), png).unwrap();
+    let (_meta, tool) = view_image_tool(session)();
+
+    let output = call(&tool, json!({ "path": "image.png" })).await;
+    let agen::tool::Attachment::Image(image) = &output.attachments[0];
+    assert_eq!(image.mime_type(), "image/png");
+    assert_eq!(image.data(), png);
+}
+
+#[tokio::test]
+async fn external_view_image_rejects_truncated_and_cross_version_chunks() {
+    let dir = TempDir::new().unwrap();
+    let spill = TempDir::new().unwrap();
+    let image_path = dir.path().join("image.png");
+    let scope = scope_with_spill(dir.path(), spill.path());
+    std::fs::write(&image_path, b"\x89PNG\r\n\x1a\npartial").unwrap();
+    let truncated: WorkdirSessionHandle = Arc::new(ExternalValidatingSession {
+        inner: LocalWorkdirSession::new(scope.clone(), dir.path().to_path_buf()),
+        truncate_first_binary_response: true,
+        mutate_after_first_binary_response: None,
+    });
+    let (_meta, tool) = view_image_tool(truncated)();
+    let error = call_err(&tool, json!({ "path": "image.png" }))
+        .await
+        .to_string();
+    assert!(
+        error.contains("truncated") || error.contains("partial"),
+        "{error}"
+    );
+
+    let mut large = vec![b'x'; 1024 * 1024 + 32];
+    large[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n");
+    std::fs::write(&image_path, large).unwrap();
+    let changing: WorkdirSessionHandle = Arc::new(ExternalValidatingSession {
+        inner: LocalWorkdirSession::new(scope, dir.path().to_path_buf()),
+        truncate_first_binary_response: false,
+        mutate_after_first_binary_response: Some(image_path),
+    });
+    let (_meta, tool) = view_image_tool(changing)();
+    let error = call_err(&tool, json!({ "path": "image.png" }))
+        .await
+        .to_string();
+    assert!(error.contains("changed"), "{error}");
+}
+
+#[tokio::test]
+async fn view_image_enforces_the_ten_mib_boundary_without_partial_attachments() {
+    const LIMIT: usize = 10 * 1024 * 1024;
+    let dir = TempDir::new().unwrap();
+    let spill = TempDir::new().unwrap();
+    let scope = scope_with_spill(dir.path(), spill.path());
+    let session: WorkdirSessionHandle =
+        Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf()));
+    let (_meta, tool) = view_image_tool(session)();
+    let mut png = vec![0xff; LIMIT];
+    png[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n");
+    std::fs::write(dir.path().join("boundary.png"), &png).unwrap();
+
+    let output = call(&tool, json!({ "path": "boundary.png" })).await;
+    let agen::tool::Attachment::Image(image) = &output.attachments[0];
+    assert_eq!(image.data(), png);
+
+    png.push(0);
+    std::fs::write(dir.path().join("boundary.png"), png).unwrap();
+    let error = call_err(&tool, json!({ "path": "boundary.png" }))
+        .await
+        .to_string();
+    assert!(error.contains("10485760-byte limit"), "{error}");
+}
+
 #[tokio::test]
 async fn view_image_reads_scoped_bytes_into_durable_tool_detail() {
     let dir = TempDir::new().unwrap();