bamboo_tools/tools/
view_image.rs1use 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 path: String,
18}
19
20pub 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}