Skip to main content

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::path::Path;
9
10use super::parameter_schema;
11use super::read_tracker::{self, MAX_TRACKED_FILE_SIZE};
12
13#[derive(Debug, Deserialize, schemars::JsonSchema)]
14#[serde(deny_unknown_fields)]
15struct ViewImageArgs {
16    /// Absolute path to a local PNG, JPEG, GIF, or WebP image
17    path: String,
18}
19
20/// Read a local raster image and return it through Bamboo's multimodal tool-result channel.
21pub struct ViewImageTool;
22
23impl ViewImageTool {
24    pub fn new() -> Self {
25        Self
26    }
27
28    fn supported_mime(bytes: &[u8]) -> Option<&'static str> {
29        if bytes.starts_with(b"\x89PNG\r\n\x1a\n") {
30            Some("image/png")
31        } else if bytes.starts_with(b"\xff\xd8\xff") {
32            Some("image/jpeg")
33        } else if bytes.starts_with(b"GIF87a") || bytes.starts_with(b"GIF89a") {
34            Some("image/gif")
35        } else if bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WEBP" {
36            Some("image/webp")
37        } else {
38            None
39        }
40    }
41}
42
43impl Default for ViewImageTool {
44    fn default() -> Self {
45        Self::new()
46    }
47}
48
49#[async_trait]
50impl Tool for ViewImageTool {
51    fn name(&self) -> &str {
52        "ViewImage"
53    }
54
55    fn description(&self) -> &str {
56        "View a local PNG, JPEG, GIF, or WebP image. Vision-capable main models receive image content directly. For a main model without Vision, the configured Vision model returns a labeled, untrusted text description. This fallback applies only to this tool; direct conversation images still require Vision on the main model. A missing or unsupported Vision model returns a configuration error."
57    }
58
59    fn classify(&self, _args: &serde_json::Value) -> ToolClass {
60        ToolClass::READONLY_PARALLEL
61    }
62
63    fn parameters_schema(&self) -> serde_json::Value {
64        parameter_schema::for_arguments::<ViewImageArgs>()
65    }
66
67    async fn invoke(
68        &self,
69        args: serde_json::Value,
70        ctx: ToolCtx,
71    ) -> Result<ToolOutcome, ToolError> {
72        let parsed: ViewImageArgs = serde_json::from_value(args).map_err(|error| {
73            ToolError::InvalidArguments(format!("Invalid ViewImage args: {error}"))
74        })?;
75        let raw_path = parsed.path.trim();
76        if raw_path.is_empty() {
77            return Err(ToolError::InvalidArguments(
78                "path must be a non-empty absolute path".to_string(),
79            ));
80        }
81
82        let path = Path::new(raw_path);
83        if !path.is_absolute() {
84            return Err(ToolError::InvalidArguments(
85                "path must be an absolute path".to_string(),
86            ));
87        }
88
89        let metadata = tokio::fs::metadata(path)
90            .await
91            .map_err(|error| ToolError::Execution(format!("Failed to inspect image: {error}")))?;
92        if !metadata.is_file() {
93            return Err(ToolError::InvalidArguments(format!(
94                "Image path must point to a regular file: {}",
95                path.display()
96            )));
97        }
98        if metadata.len() > MAX_TRACKED_FILE_SIZE {
99            return Err(ToolError::Execution(format!(
100                "Image is {} bytes, which exceeds the maximum ViewImage size of {} bytes ({} MB)",
101                metadata.len(),
102                MAX_TRACKED_FILE_SIZE,
103                MAX_TRACKED_FILE_SIZE / 1024 / 1024
104            )));
105        }
106
107        let stable = read_tracker::stable_read(raw_path).await.map_err(|error| {
108            ToolError::Execution(format!("Failed to read stable image: {error}"))
109        })?;
110        let bytes = stable.bytes();
111        let mime_type = Self::supported_mime(bytes).ok_or_else(|| {
112            ToolError::InvalidArguments(
113                "Unsupported image format; expected PNG, JPEG, GIF, or WebP content".to_string(),
114            )
115        })?;
116
117        if let Some(session_id) = ctx.session_id() {
118            read_tracker::mark_stable_read(session_id, raw_path, &stable).await;
119        }
120
121        Ok(ToolOutcome::Completed(ToolResult {
122            success: true,
123            result: json!({
124                "path": raw_path,
125                "mime_type": mime_type,
126                "size_bytes": bytes.len(),
127                "delivery": "base64_image"
128            })
129            .to_string(),
130            display_preference: None,
131            images: vec![ToolResultImage {
132                mime_type: mime_type.to_string(),
133                data: STANDARD.encode(bytes),
134            }],
135        }))
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142
143    #[tokio::test]
144    async fn view_image_returns_detected_base64_image() {
145        let file = tempfile::Builder::new().suffix(".bin").tempfile().unwrap();
146        let bytes = b"\x89PNG\r\n\x1a\nview-image-test";
147        tokio::fs::write(file.path(), bytes).await.unwrap();
148
149        let output = ViewImageTool::new()
150            .invoke(
151                json!({"path": file.path().to_string_lossy()}),
152                ToolCtx::none("view-image"),
153            )
154            .await
155            .unwrap();
156        let ToolOutcome::Completed(result) = output else {
157            panic!("expected completed ViewImage result");
158        };
159
160        assert!(result.success);
161        assert_eq!(result.images.len(), 1);
162        assert_eq!(result.images[0].mime_type, "image/png");
163        assert_eq!(result.images[0].data, STANDARD.encode(bytes));
164        assert!(!result.result.contains(&result.images[0].data));
165    }
166
167    #[tokio::test]
168    async fn view_image_rejects_relative_and_non_image_paths() {
169        let relative = ViewImageTool::new()
170            .invoke(json!({"path": "image.png"}), ToolCtx::none("relative"))
171            .await;
172        assert!(matches!(relative, Err(ToolError::InvalidArguments(_))));
173
174        let file = tempfile::NamedTempFile::new().unwrap();
175        tokio::fs::write(file.path(), b"plain text").await.unwrap();
176        let non_image = ViewImageTool::new()
177            .invoke(
178                json!({"path": file.path().to_string_lossy()}),
179                ToolCtx::none("non-image"),
180            )
181            .await;
182        assert!(matches!(non_image, Err(ToolError::InvalidArguments(_))));
183    }
184
185    #[tokio::test]
186    async fn view_image_rejects_oversized_files_before_reading() {
187        let file = tempfile::NamedTempFile::new().unwrap();
188        file.as_file().set_len(MAX_TRACKED_FILE_SIZE + 1).unwrap();
189
190        let result = ViewImageTool::new()
191            .invoke(
192                json!({"path": file.path().to_string_lossy()}),
193                ToolCtx::none("oversized"),
194            )
195            .await;
196        assert!(matches!(result, Err(ToolError::Execution(_))));
197    }
198}