bamboo_tools/tools/
view_image.rs

1use async_trait::async_trait;
2use bamboo_agent_core::{
3    Tool, ToolClass, ToolCtx, ToolError, ToolOutcome, ToolResult, ToolResultImage,
4};
5use base64::{engine::general_purpose::STANDARD, Engine as _};
6use serde::Deserialize;
7use serde_json::json;
8use std::future::Future;
9use std::path::Path;
10use std::sync::Arc;
11
12use futures::future::BoxFuture;
13
14use super::parameter_schema;
15use super::read_tracker::{self, MAX_TRACKED_FILE_SIZE};
16
17#[derive(Debug, Deserialize, schemars::JsonSchema)]
18#[serde(deny_unknown_fields)]
19struct ViewImageArgs {
20    /// Absolute local image path or exact image reference supplied by this conversation
21    path: String,
22}
23
24/// Host-owned reader for references from the executing Session's original image parts.
25/// It is scoped to dispatch, never deserialized from model arguments or persisted.
26pub type ViewImageReferenceReader =
27    Arc<dyn Fn(String) -> BoxFuture<'static, Result<Vec<u8>, ToolError>> + Send + Sync>;
28
29pub const MAX_VIEW_IMAGE_BYTES: usize = MAX_TRACKED_FILE_SIZE as usize;
30
31struct ImageReferenceScope {
32    session_id: String,
33    reader: ViewImageReferenceReader,
34}
35
36tokio::task_local! {
37    static IMAGE_REFERENCES: ImageReferenceScope;
38}
39
40/// Keep reference access inside the same inline tool dispatch and cancellation scope.
41pub async fn with_image_reference_reader<F: Future>(
42    session_id: &str,
43    reader: ViewImageReferenceReader,
44    dispatch: F,
45) -> F::Output {
46    IMAGE_REFERENCES
47        .scope(
48            ImageReferenceScope {
49                session_id: session_id.to_string(),
50                reader,
51            },
52            dispatch,
53        )
54        .await
55}
56
57fn is_image_reference(path: &str) -> bool {
58    path.starts_with("bamboo-attachment://") || path.starts_with("bamboo-image://")
59}
60
61async fn read_image_reference(ctx: &ToolCtx, reference: &str) -> Result<Vec<u8>, ToolError> {
62    let reader = IMAGE_REFERENCES
63        .try_with(|scope| {
64            (ctx.session_id() == Some(scope.session_id.as_str())).then(|| scope.reader.clone())
65        })
66        .ok()
67        .flatten()
68        .ok_or_else(|| {
69            ToolError::Execution(
70                "Image reference is unavailable outside its executing conversation. Use the exact reference from the current conversation or an authorized local image path.".into(),
71            )
72        })?;
73    reader(reference.to_string()).await
74}
75
76/// Read a local raster image or a host-bound conversation image reference.
77pub struct ViewImageTool;
78
79impl ViewImageTool {
80    pub fn new() -> Self {
81        Self
82    }
83
84    fn supported_mime(bytes: &[u8]) -> Option<&'static str> {
85        if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
86            Some("image/png")
87        } else if bytes.starts_with(b"\xff\xd8\xff") {
88            Some("image/jpeg")
89        } else if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
90            Some("image/gif")
91        } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
92            Some("image/webp")
93        } else {
94            None
95        }
96    }
97}
98
99impl Default for ViewImageTool {
100    fn default() -> Self {
101        Self::new()
102    }
103}
104
105#[async_trait]
106impl Tool for ViewImageTool {
107    fn name(&self) -> &str {
108        "ViewImage"
109    }
110
111    fn description(&self) -> &str {
112        "View a local PNG, JPEG, GIF, or WebP image, or the exact image reference supplied in this conversation. Vision-capable main models receive image content directly. For a main without Vision, an explicitly configured Vision model returns a labeled, untrusted description. Uploaded images remain preserved; only this requested tool call sends them to the configured Vision model. Missing or unsupported Vision configuration returns actionable feedback. Never claim to have inspected an image before receiving a successful result."
113    }
114
115    fn classify(&self, _args: &serde_json::Value) -> ToolClass {
116        ToolClass::READONLY_PARALLEL
117    }
118
119    fn parameters_schema(&self) -> serde_json::Value {
120        parameter_schema::for_arguments::<ViewImageArgs>()
121    }
122
123    async fn invoke(
124        &self,
125        args: serde_json::Value,
126        ctx: ToolCtx,
127    ) -> Result<ToolOutcome, ToolError> {
128        let parsed: ViewImageArgs = serde_json::from_value(args).map_err(|error| {
129            ToolError::InvalidArguments(format!("Invalid ViewImage args: {error}"))
130        })?;
131        let raw_path = parsed.path.trim();
132        if raw_path.is_empty() {
133            return Err(ToolError::InvalidArguments(
134                "path must be a non-empty absolute path".to_string(),
135            ));
136        }
137
138        let (reference_bytes, stable) = if is_image_reference(raw_path) {
139            // Never pass an unavailable/forged reference to the ordinary file branch.
140            (Some(read_image_reference(&ctx, raw_path).await?), None)
141        } else {
142            let path = Path::new(raw_path);
143            if !path.is_absolute() {
144                return Err(ToolError::InvalidArguments(
145                    "path must be an absolute local path or a current conversation image reference"
146                        .to_string(),
147                ));
148            }
149
150            let metadata = tokio::fs::metadata(path).await.map_err(|error| {
151                ToolError::Execution(format!("Failed to inspect image: {error}"))
152            })?;
153            if !metadata.is_file() {
154                return Err(ToolError::InvalidArguments(format!(
155                    "Image path must point to a regular file: {}",
156                    path.display()
157                )));
158            }
159            if metadata.len() > MAX_TRACKED_FILE_SIZE {
160                return Err(ToolError::Execution(format!(
161                    "Image is {} bytes, which exceeds the maximum ViewImage size of {} bytes ({} MB)",
162                    metadata.len(),
163                    MAX_TRACKED_FILE_SIZE,
164                    MAX_TRACKED_FILE_SIZE / 1024 / 1024
165                )));
166            }
167
168            let stable = read_tracker::stable_read(raw_path).await.map_err(|error| {
169                ToolError::Execution(format!("Failed to read stable image: {error}"))
170            })?;
171            (None, Some(stable))
172        };
173        let bytes = match (&reference_bytes, &stable) {
174            (Some(bytes), _) => bytes.as_slice(),
175            (_, Some(stable)) => stable.bytes(),
176            _ => unreachable!("image source selected above"),
177        };
178        if bytes.len() > MAX_VIEW_IMAGE_BYTES {
179            return Err(ToolError::Execution(format!(
180                "Image exceeds the maximum ViewImage size of {MAX_VIEW_IMAGE_BYTES} bytes"
181            )));
182        }
183        let mime_type = Self::supported_mime(bytes).ok_or_else(|| {
184            ToolError::InvalidArguments(
185                "Unsupported image format; expected PNG, JPEG, GIF, or WebP content".to_string(),
186            )
187        })?;
188
189        if let (Some(session_id), Some(stable)) = (ctx.session_id(), stable.as_ref()) {
190            read_tracker::mark_stable_read(session_id, raw_path, stable).await;
191        }
192
193        Ok(ToolOutcome::Completed(ToolResult {
194            success: true,
195            result: json!({
196                "path": raw_path,
197                "mime_type": mime_type,
198                "size_bytes": bytes.len(),
199                "delivery": "base64_image"
200            })
201            .to_string(),
202            display_preference: None,
203            images: vec![ToolResultImage {
204                mime_type: mime_type.to_string(),
205                data: STANDARD.encode(bytes),
206            }],
207        }))
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[tokio::test]
216    async fn view_image_returns_detected_base64_image() {
217        let file = tempfile::Builder::new().suffix(".bin").tempfile().unwrap();
218        let bytes = b"\x89PNG\r\n\x1a\nview-image-test";
219        tokio::fs::write(file.path(), bytes).await.unwrap();
220
221        let output = ViewImageTool::new()
222            .invoke(
223                json!({"path": file.path().to_string_lossy()}),
224                ToolCtx::none("view-image"),
225            )
226            .await
227            .unwrap();
228        let ToolOutcome::Completed(result) = output else {
229            panic!("expected completed ViewImage result");
230        };
231
232        assert!(result.success);
233        assert_eq!(result.images.len(), 1);
234        assert_eq!(result.images[0].mime_type, "image/png");
235        assert_eq!(result.images[0].data, STANDARD.encode(bytes));
236        assert!(!result.result.contains(&result.images[0].data));
237    }
238
239    #[tokio::test]
240    async fn view_image_rejects_relative_and_non_image_paths() {
241        let relative = ViewImageTool::new()
242            .invoke(json!({"path": "image.png"}), ToolCtx::none("relative"))
243            .await;
244        assert!(matches!(relative, Err(ToolError::InvalidArguments(_))));
245
246        let file = tempfile::NamedTempFile::new().unwrap();
247        tokio::fs::write(file.path(), b"plain text").await.unwrap();
248        let non_image = ViewImageTool::new()
249            .invoke(
250                json!({"path": file.path().to_string_lossy()}),
251                ToolCtx::none("non-image"),
252            )
253            .await;
254        assert!(matches!(non_image, Err(ToolError::InvalidArguments(_))));
255    }
256
257    #[tokio::test]
258    async fn view_image_rejects_oversized_files_before_reading() {
259        let file = tempfile::NamedTempFile::new().unwrap();
260        file.as_file().set_len(MAX_TRACKED_FILE_SIZE + 1).unwrap();
261
262        let result = ViewImageTool::new()
263            .invoke(
264                json!({"path": file.path().to_string_lossy()}),
265                ToolCtx::none("oversized"),
266            )
267            .await;
268        assert!(matches!(result, Err(ToolError::Execution(_))));
269    }
270}
bamboo_tools/tools/view_image.rs

bamboo_tools/tools/
view_image.rs
