| 1 | //! `image_analyze` tool — analyze images using a dedicated vision model. |
| 2 | |
| 3 | use std::path::{Component, Path, PathBuf}; |
| 4 | use std::time::Duration; |
| 5 | |
| 6 | use async_trait::async_trait; |
| 7 | use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; |
| 8 | use serde_json::{Value, json}; |
| 9 | |
| 10 | use crate::config::VisionModelConfig; |
| 11 | use crate::llm_client::{LlmError, RetryConfig, sanitize_http_error_body, with_retry}; |
| 12 | use crate::tools::spec::{ |
| 13 | ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, required_str, |
| 14 | }; |
| 15 | |
| 16 | const DEFAULT_VISION_MAX_OUTPUT_TOKENS: u32 = 4096; |
| 17 | |
| 18 | pub struct ImageAnalyzeTool { |
| 19 | config: VisionModelConfig, |
| 20 | client: reqwest::Client, |
| 21 | } |
| 22 | |
| 23 | impl ImageAnalyzeTool { |
| 24 | #[must_use] |
| 25 | pub fn new(config: VisionModelConfig) -> Self { |
| 26 | let client = crate::tls::reqwest_client_builder() |
| 27 | .timeout(Duration::from_secs(120)) |
| 28 | .build() |
| 29 | .expect("Failed to build HTTP client"); |
| 30 | Self { config, client } |
| 31 | } |
| 32 | |
| 33 | async fn read_image_file(path: &Path) -> Result<(String, String), ToolError> { |
| 34 | let bytes = tokio::fs::read(path) |
| 35 | .await |
| 36 | .map_err(|e| ToolError::execution_failed(format!("Failed to read image file: {e}")))?; |
| 37 | |
| 38 | let mime_type = Self::detect_mime_type(path)?; |
| 39 | let base64_data = BASE64.encode(&bytes); |
| 40 | Ok((base64_data, mime_type)) |
| 41 | } |
| 42 | |
| 43 | fn resolve_image_path(workspace: &Path, image_path: &str) -> Result<PathBuf, ToolError> { |
| 44 | let image_path_buf = Path::new(image_path); |
| 45 | if image_path_buf.components().any(|c| { |
| 46 | matches!( |
| 47 | c, |
| 48 | Component::Prefix(_) | Component::RootDir | Component::ParentDir |
| 49 | ) |
| 50 | }) { |
| 51 | return Err(ToolError::execution_failed( |
| 52 | "image_path must be a relative path within the workspace and cannot escape it.", |
| 53 | )); |
| 54 | } |
| 55 | |
| 56 | let workspace = workspace.canonicalize().map_err(|e| { |
| 57 | ToolError::execution_failed(format!("Failed to resolve workspace path: {e}")) |
| 58 | })?; |
| 59 | let candidate = workspace.join(image_path_buf); |
| 60 | let resolved = candidate.canonicalize().map_err(|e| { |
| 61 | ToolError::execution_failed(format!("Failed to resolve image file: {e}")) |
| 62 | })?; |
| 63 | if !resolved.starts_with(&workspace) { |
| 64 | return Err(ToolError::execution_failed( |
| 65 | "image_path must resolve within the workspace and cannot escape it.", |
| 66 | )); |
| 67 | } |
| 68 | Ok(resolved) |
| 69 | } |
| 70 | |
| 71 | fn detect_mime_type(path: &Path) -> Result<String, ToolError> { |
| 72 | let extension = path |
| 73 | .extension() |
| 74 | .and_then(|e| e.to_str()) |
| 75 | .unwrap_or("") |
| 76 | .to_lowercase(); |
| 77 | |
| 78 | match extension.as_str() { |
| 79 | "png" => Ok("image/png".to_string()), |
| 80 | "jpg" | "jpeg" => Ok("image/jpeg".to_string()), |
| 81 | "gif" => Ok("image/gif".to_string()), |
| 82 | "webp" => Ok("image/webp".to_string()), |
| 83 | "bmp" => Ok("image/bmp".to_string()), |
| 84 | _ => Err(ToolError::execution_failed(format!( |
| 85 | "Unsupported image format: {extension}" |
| 86 | ))), |
| 87 | } |
| 88 | } |
| 89 | |
| 90 | fn base_url(&self) -> String { |
| 91 | self.config |
| 92 | .base_url |
| 93 | .clone() |
| 94 | .unwrap_or_else(|| "https://api.openai.com/v1".to_string()) |
| 95 | } |
| 96 | |
| 97 | fn api_key(&self) -> String { |
| 98 | self.config.api_key.clone().unwrap_or_default() |
| 99 | } |
| 100 | |
| 101 | fn is_xiaomi_mimo_model(model: &str) -> bool { |
| 102 | let normalized = model.trim().to_ascii_lowercase(); |
| 103 | let normalized = normalized.strip_prefix("xiaomi/").unwrap_or(&normalized); |
| 104 | normalized.starts_with("mimo-") |
| 105 | } |
| 106 | |
| 107 | fn uses_max_completion_tokens(config: &VisionModelConfig) -> bool { |
| 108 | if Self::is_xiaomi_mimo_model(&config.model) { |
| 109 | return true; |
| 110 | } |
| 111 | |
| 112 | let base_url = config.base_url.as_deref().unwrap_or_default(); |
| 113 | let Ok(url) = reqwest::Url::parse(base_url) else { |
| 114 | return false; |
| 115 | }; |
| 116 | let Some(domain) = url.domain() else { |
| 117 | return false; |
| 118 | }; |
| 119 | |
| 120 | domain.eq_ignore_ascii_case("xiaomimimo.com") |
| 121 | || domain.to_ascii_lowercase().ends_with(".xiaomimimo.com") |
| 122 | } |
| 123 | |
| 124 | fn request_payload(&self, prompt: &str, image_data: &str, mime_type: &str) -> Value { |
| 125 | let mut payload = json!({ |
| 126 | "model": self.config.model, |
| 127 | "messages": [ |
| 128 | { |
| 129 | "role": "user", |
| 130 | "content": [ |
| 131 | {"type": "text", "text": prompt}, |
| 132 | { |
| 133 | "type": "image_url", |
| 134 | "image_url": { |
| 135 | "url": format!("data:{};base64,{}", mime_type, image_data) |
| 136 | } |
| 137 | } |
| 138 | ] |
| 139 | } |
| 140 | ], |
| 141 | "temperature": 0.7 |
| 142 | }); |
| 143 | |
| 144 | let token_limit_field = if Self::uses_max_completion_tokens(&self.config) { |
| 145 | "max_completion_tokens" |
| 146 | } else { |
| 147 | "max_tokens" |
| 148 | }; |
| 149 | payload[token_limit_field] = json!(DEFAULT_VISION_MAX_OUTPUT_TOKENS); |
| 150 | |
| 151 | payload |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | #[async_trait] |
| 156 | impl ToolSpec for ImageAnalyzeTool { |
| 157 | fn name(&self) -> &str { |
| 158 | "image_analyze" |
| 159 | } |
| 160 | |
| 161 | fn description(&self) -> &str { |
| 162 | "Analyze an image using the configured vision model. \ |
| 163 | Supports PNG, JPEG, GIF, WebP, and BMP formats." |
| 164 | } |
| 165 | |
| 166 | fn input_schema(&self) -> Value { |
| 167 | json!({ |
| 168 | "type": "object", |
| 169 | "properties": { |
| 170 | "image_path": { |
| 171 | "type": "string", |
| 172 | "description": "Path to the image file to analyze" |
| 173 | }, |
| 174 | "prompt": { |
| 175 | "type": "string", |
| 176 | "description": "Optional prompt to guide the analysis." |
| 177 | } |
| 178 | }, |
| 179 | "required": ["image_path"] |
| 180 | }) |
| 181 | } |
| 182 | |
| 183 | fn capabilities(&self) -> Vec<ToolCapability> { |
| 184 | vec![ToolCapability::ReadOnly] |
| 185 | } |
| 186 | |
| 187 | async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolResult, ToolError> { |
| 188 | let image_path = required_str(&input, "image_path")?; |
| 189 | let prompt = input |
| 190 | .get("prompt") |
| 191 | .and_then(|v| v.as_str()) |
| 192 | .unwrap_or("Describe this image in detail."); |
| 193 | |
| 194 | let resolved_path = Self::resolve_image_path(&context.workspace, image_path)?; |
| 195 | let (image_data, mime_type) = Self::read_image_file(&resolved_path).await?; |
| 196 | |
| 197 | let payload = self.request_payload(prompt, &image_data, &mime_type); |
| 198 | |
| 199 | let url = format!("{}/chat/completions", self.base_url()); |
| 200 | let api_key = self.api_key(); |
| 201 | |
| 202 | let retry_config = RetryConfig { |
| 203 | max_retries: 3, |
| 204 | initial_delay: 1.0, |
| 205 | max_delay: 30.0, |
| 206 | enabled: true, |
| 207 | ..Default::default() |
| 208 | }; |
| 209 | |
| 210 | let response = with_retry( |
| 211 | &retry_config, |
| 212 | || { |
| 213 | let client = self.client.clone(); |
| 214 | let url = url.clone(); |
| 215 | let api_key = api_key.clone(); |
| 216 | let payload = payload.clone(); |
| 217 | async move { |
| 218 | let response = client |
| 219 | .post(&url) |
| 220 | .header("Content-Type", "application/json") |
| 221 | .header("Authorization", format!("Bearer {api_key}")) |
| 222 | .json(&payload) |
| 223 | .send() |
| 224 | .await |
| 225 | .map_err(|e| LlmError::from_reqwest(&e))?; |
| 226 | |
| 227 | let status = response.status(); |
| 228 | if !status.is_success() { |
| 229 | let error_text = response |
| 230 | .text() |
| 231 | .await |
| 232 | .unwrap_or_else(|_| "Unknown error".to_string()); |
| 233 | let error_text = sanitize_http_error_body( |
| 234 | Some("Vision provider"), |
| 235 | status.as_u16(), |
| 236 | &error_text, |
| 237 | ); |
| 238 | return Err(LlmError::from_http_response(status.as_u16(), &error_text)); |
| 239 | } |
| 240 | Ok(response) |
| 241 | } |
| 242 | }, |
| 243 | None, |
| 244 | ) |
| 245 | .await |
| 246 | .map_err(|e| ToolError::execution_failed(format!("Vision API request failed: {e}")))?; |
| 247 | |
| 248 | let json: Value = response |
| 249 | .json() |
| 250 | .await |
| 251 | .map_err(|e| ToolError::execution_failed(format!("Failed to parse response: {e}")))?; |
| 252 | |
| 253 | let content = json |
| 254 | .get("choices") |
| 255 | .and_then(|c| c.get(0)) |
| 256 | .and_then(|c| c.get("message")) |
| 257 | .and_then(|m| m.get("content")) |
| 258 | .and_then(|c| c.as_str()) |
| 259 | .unwrap_or("") |
| 260 | .to_string(); |
| 261 | |
| 262 | let model = json |
| 263 | .get("model") |
| 264 | .and_then(|m| m.as_str()) |
| 265 | .unwrap_or(&self.config.model) |
| 266 | .to_string(); |
| 267 | |
| 268 | let result = json!({ |
| 269 | "analysis": content, |
| 270 | "model": model, |
| 271 | }); |
| 272 | |
| 273 | ToolResult::json(&result) |
| 274 | .map_err(|e| ToolError::execution_failed(format!("Failed to serialize result: {e}"))) |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | #[cfg(test)] |
| 279 | mod tests { |
| 280 | use super::*; |
| 281 | use tempfile::tempdir; |
| 282 | |
| 283 | #[cfg(unix)] |
| 284 | fn create_file_symlink( |
| 285 | target: &std::path::Path, |
| 286 | link: &std::path::Path, |
| 287 | ) -> std::io::Result<()> { |
| 288 | std::os::unix::fs::symlink(target, link) |
| 289 | } |
| 290 | |
| 291 | #[cfg(windows)] |
| 292 | fn create_file_symlink( |
| 293 | target: &std::path::Path, |
| 294 | link: &std::path::Path, |
| 295 | ) -> std::io::Result<()> { |
| 296 | std::os::windows::fs::symlink_file(target, link) |
| 297 | } |
| 298 | |
| 299 | fn fake_config() -> VisionModelConfig { |
| 300 | VisionModelConfig { |
| 301 | model: "test-vision-model".to_string(), |
| 302 | api_key: Some("test-key".to_string()), |
| 303 | base_url: Some("https://example.invalid/v1".to_string()), |
| 304 | } |
| 305 | } |
| 306 | |
| 307 | #[test] |
| 308 | fn tool_metadata_is_read_only_and_named_image_analyze() { |
| 309 | let tool = ImageAnalyzeTool::new(fake_config()); |
| 310 | assert_eq!(tool.name(), "image_analyze"); |
| 311 | assert!(tool.capabilities().contains(&ToolCapability::ReadOnly)); |
| 312 | } |
| 313 | |
| 314 | #[test] |
| 315 | fn mime_type_detection_covers_common_formats() { |
| 316 | for (ext, expected) in [ |
| 317 | ("png", "image/png"), |
| 318 | ("PNG", "image/png"), |
| 319 | ("jpg", "image/jpeg"), |
| 320 | ("jpeg", "image/jpeg"), |
| 321 | ("gif", "image/gif"), |
| 322 | ("webp", "image/webp"), |
| 323 | ("bmp", "image/bmp"), |
| 324 | ] { |
| 325 | let path = std::path::PathBuf::from(format!("test.{ext}")); |
| 326 | let mime = ImageAnalyzeTool::detect_mime_type(&path) |
| 327 | .unwrap_or_else(|_| panic!("must detect {ext}")); |
| 328 | assert_eq!(mime, expected); |
| 329 | } |
| 330 | } |
| 331 | |
| 332 | #[test] |
| 333 | fn mime_type_detection_rejects_unsupported_extension() { |
| 334 | let path = std::path::PathBuf::from("test.svg"); |
| 335 | let err = ImageAnalyzeTool::detect_mime_type(&path) |
| 336 | .expect_err("svg is intentionally out of scope for vision tool"); |
| 337 | assert!(err.to_string().contains("Unsupported image format")); |
| 338 | } |
| 339 | |
| 340 | #[test] |
| 341 | fn generic_vision_payload_uses_max_tokens() { |
| 342 | let tool = ImageAnalyzeTool::new(fake_config()); |
| 343 | |
| 344 | let payload = tool.request_payload("describe", "abc123", "image/png"); |
| 345 | |
| 346 | assert_eq!( |
| 347 | payload.get("max_tokens").and_then(Value::as_u64), |
| 348 | Some(u64::from(DEFAULT_VISION_MAX_OUTPUT_TOKENS)) |
| 349 | ); |
| 350 | assert!(payload.get("max_completion_tokens").is_none()); |
| 351 | } |
| 352 | |
| 353 | #[test] |
| 354 | fn xiaomi_mimo_vision_payload_uses_max_completion_tokens() { |
| 355 | let mut config = fake_config(); |
| 356 | config.model = "mimo-v2.5".to_string(); |
| 357 | config.base_url = Some("https://api.xiaomimimo.com/v1".to_string()); |
| 358 | let tool = ImageAnalyzeTool::new(config); |
| 359 | |
| 360 | let payload = tool.request_payload("describe", "abc123", "image/png"); |
| 361 | |
| 362 | assert_eq!( |
| 363 | payload.get("max_completion_tokens").and_then(Value::as_u64), |
| 364 | Some(u64::from(DEFAULT_VISION_MAX_OUTPUT_TOKENS)) |
| 365 | ); |
| 366 | assert!(payload.get("max_tokens").is_none()); |
| 367 | } |
| 368 | |
| 369 | #[test] |
| 370 | fn xiaomi_mimo_vision_payload_uses_max_completion_tokens_with_custom_proxy() { |
| 371 | let mut config = fake_config(); |
| 372 | config.model = "mimo-v2.5".to_string(); |
| 373 | config.base_url = Some("https://vision-proxy.example.invalid/v1".to_string()); |
| 374 | let tool = ImageAnalyzeTool::new(config); |
| 375 | |
| 376 | let payload = tool.request_payload("describe", "abc123", "image/png"); |
| 377 | |
| 378 | assert_eq!( |
| 379 | payload.get("max_completion_tokens").and_then(Value::as_u64), |
| 380 | Some(u64::from(DEFAULT_VISION_MAX_OUTPUT_TOKENS)) |
| 381 | ); |
| 382 | assert!(payload.get("max_tokens").is_none()); |
| 383 | } |
| 384 | |
| 385 | #[tokio::test] |
| 386 | async fn execute_rejects_absolute_path() { |
| 387 | // Trust-boundary pin: image_path must stay inside the workspace |
| 388 | // — an absolute path or a `..`-traversing path must reject |
| 389 | // before any base64 / API call. |
| 390 | let tmp = tempdir().expect("tempdir"); |
| 391 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 392 | let tool = ImageAnalyzeTool::new(fake_config()); |
| 393 | let outside_workspace = if cfg!(windows) { |
| 394 | r"C:\Windows\System32\drivers\etc\hosts" |
| 395 | } else { |
| 396 | "/etc/hosts" |
| 397 | }; |
| 398 | let err = tool |
| 399 | .execute(json!({"image_path": outside_workspace}), &ctx) |
| 400 | .await |
| 401 | .expect_err("absolute path must reject"); |
| 402 | assert!( |
| 403 | err.to_string() |
| 404 | .contains("relative path within the workspace"), |
| 405 | "error must call out the workspace boundary; got {err}" |
| 406 | ); |
| 407 | } |
| 408 | |
| 409 | #[tokio::test] |
| 410 | async fn execute_rejects_parent_dir_traversal() { |
| 411 | let tmp = tempdir().expect("tempdir"); |
| 412 | let ctx = ToolContext::new(tmp.path().to_path_buf()); |
| 413 | let tool = ImageAnalyzeTool::new(fake_config()); |
| 414 | let err = tool |
| 415 | .execute(json!({"image_path": "../escape.png"}), &ctx) |
| 416 | .await |
| 417 | .expect_err("`..`-traversal must reject"); |
| 418 | assert!( |
| 419 | err.to_string() |
| 420 | .contains("relative path within the workspace"), |
| 421 | "error must call out the workspace boundary; got {err}" |
| 422 | ); |
| 423 | } |
| 424 | |
| 425 | #[tokio::test] |
| 426 | async fn execute_rejects_symlink_that_resolves_outside_workspace() { |
| 427 | let workspace = tempdir().expect("workspace tempdir"); |
| 428 | let outside = tempdir().expect("outside tempdir"); |
| 429 | let outside_image = outside.path().join("outside.png"); |
| 430 | std::fs::write(&outside_image, b"not a real png").expect("write outside image"); |
| 431 | let link = workspace.path().join("linked.png"); |
| 432 | if let Err(err) = create_file_symlink(&outside_image, &link) { |
| 433 | eprintln!("skipping symlink assertion: {err}"); |
| 434 | return; |
| 435 | } |
| 436 | |
| 437 | let ctx = ToolContext::new(workspace.path().to_path_buf()); |
| 438 | let tool = ImageAnalyzeTool::new(fake_config()); |
| 439 | let err = tool |
| 440 | .execute(json!({"image_path": "linked.png"}), &ctx) |
| 441 | .await |
| 442 | .expect_err("symlink target outside workspace must reject before reading"); |
| 443 | assert!( |
| 444 | err.to_string().contains("resolve within the workspace"), |
| 445 | "error must call out the canonical workspace boundary; got {err}" |
| 446 | ); |
| 447 | } |
| 448 | } |
| 449 |