返回 CodeWhale
tools.rs
根目录 / crates / tui / src / vision / tools.rs
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::client::CodewhaleClient;
11 use crate::config::ApiProvider;
12 use crate::config::VisionModelConfig;
13 use crate::llm_client::{LlmError, RetryConfig, sanitize_http_error_body, with_retry};
14 use crate::tools::spec::{
15 ToolCapability, ToolContext, ToolError, ToolResult, ToolSpec, required_str,
16 };
17
18 pub struct ImageAnalyzeTool {
19 config: VisionModelConfig,
20 client: reqwest::Client,
21 route_client: Option<CodewhaleClient>,
22 }
23
24 impl ImageAnalyzeTool {
25 #[cfg(test)]
26 #[must_use]
27 pub fn new(config: VisionModelConfig) -> Self {
28 Self::new_with_route_client(config, None)
29 }
30
31 #[must_use]
32 pub fn new_with_route_client(
33 config: VisionModelConfig,
34 route_client: Option<CodewhaleClient>,
35 ) -> Self {
36 let client = crate::tls::reqwest_client_builder()
37 .timeout(Duration::from_secs(120))
38 .build()
39 .expect("Failed to build HTTP client");
40 Self {
41 config,
42 client,
43 route_client,
44 }
45 }
46
47 async fn read_image_file(path: &Path) -> Result<(String, String), ToolError> {
48 let bytes = tokio::fs::read(path)
49 .await
50 .map_err(|e| ToolError::execution_failed(format!("Failed to read image file: {e}")))?;
51
52 let mime_type = Self::detect_mime_type(path)?;
53 let base64_data = BASE64.encode(&bytes);
54 Ok((base64_data, mime_type))
55 }
56
57 fn resolve_image_path(workspace: &Path, image_path: &str) -> Result<PathBuf, ToolError> {
58 let image_path_buf = Path::new(image_path);
59 if image_path_buf.components().any(|c| {
60 matches!(
61 c,
62 Component::Prefix(_) | Component::RootDir | Component::ParentDir
63 )
64 }) {
65 return Err(ToolError::execution_failed(
66 "image_path must be a relative path within the workspace and cannot escape it.",
67 ));
68 }
69
70 let workspace = workspace.canonicalize().map_err(|e| {
71 ToolError::execution_failed(format!("Failed to resolve workspace path: {e}"))
72 })?;
73 let candidate = workspace.join(image_path_buf);
74 let resolved = candidate.canonicalize().map_err(|e| {
75 ToolError::execution_failed(format!("Failed to resolve image file: {e}"))
76 })?;
77 if !resolved.starts_with(&workspace) {
78 return Err(ToolError::execution_failed(
79 "image_path must resolve within the workspace and cannot escape it.",
80 ));
81 }
82 Ok(resolved)
83 }
84
85 fn detect_mime_type(path: &Path) -> Result<String, ToolError> {
86 let extension = path
87 .extension()
88 .and_then(|e| e.to_str())
89 .unwrap_or("")
90 .to_lowercase();
91
92 match extension.as_str() {
93 "png" => Ok("image/png".to_string()),
94 "jpg" | "jpeg" => Ok("image/jpeg".to_string()),
95 "gif" => Ok("image/gif".to_string()),
96 "webp" => Ok("image/webp".to_string()),
97 "bmp" => Ok("image/bmp".to_string()),
98 _ => Err(ToolError::execution_failed(format!(
99 "Unsupported image format: {extension}"
100 ))),
101 }
102 }
103
104 fn base_url(&self) -> String {
105 self.config
106 .base_url
107 .clone()
108 .unwrap_or_else(|| "https://api.openai.com/v1".to_string())
109 }
110
111 fn api_key(&self) -> String {
112 self.config.api_key.clone().unwrap_or_default()
113 }
114
115 fn is_xiaomi_mimo_model(model: &str) -> bool {
116 let normalized = model.trim().to_ascii_lowercase();
117 let normalized = normalized.strip_prefix("xiaomi/").unwrap_or(&normalized);
118 normalized.starts_with("mimo-")
119 }
120
121 fn uses_max_completion_tokens(config: &VisionModelConfig) -> bool {
122 if Self::is_xiaomi_mimo_model(&config.model) {
123 return true;
124 }
125
126 let base_url = config.base_url.as_deref().unwrap_or_default();
127 let Ok(url) = reqwest::Url::parse(base_url) else {
128 return false;
129 };
130 let Some(domain) = url.domain() else {
131 return false;
132 };
133
134 domain.eq_ignore_ascii_case("xiaomimimo.com")
135 || domain.to_ascii_lowercase().ends_with(".xiaomimimo.com")
136 }
137
138 fn request_payload(&self, prompt: &str, image_data: &str, mime_type: &str) -> Value {
139 let mut payload = json!({
140 "model": self.config.model,
141 "messages": [
142 {
143 "role": "user",
144 "content": [
145 {"type": "text", "text": prompt},
146 {
147 "type": "image_url",
148 "image_url": {
149 "url": format!("data:{};base64,{}", mime_type, image_data)
150 }
151 }
152 ]
153 }
154 ]
155 });
156
157 let token_limit_field = if Self::uses_max_completion_tokens(&self.config) {
158 "max_completion_tokens"
159 } else {
160 "max_tokens"
161 };
162 let configured_base = self.base_url();
163 let route_cap = self
164 .route_client
165 .as_ref()
166 .filter(|client| {
167 client.base_url().trim_end_matches('/') == configured_base.trim_end_matches('/')
168 })
169 .map_or_else(
170 || {
171 // A standalone `[vision_model]` route has no resolved
172 // max-model-len fact. Do not guess one or let a process
173 // override turn a capability maximum into an unbounded
174 // request; a matched active client above carries exact
175 // route limits when the vision route is shared.
176 crate::route_budget::effective_max_output_tokens_for_route(
177 ApiProvider::Custom,
178 &self.config.model,
179 None,
180 )
181 .min(65_536)
182 },
183 |client| client.effective_max_output_tokens(&self.config.model),
184 );
185 payload[token_limit_field] = json!(route_cap);
186 if let Some(client) = self.route_client.as_ref().filter(|client| {
187 client.base_url().trim_end_matches('/') == configured_base.trim_end_matches('/')
188 }) {
189 client.apply_provider_routing(&mut payload);
190 }
191
192 payload
193 }
194 }
195
196 #[async_trait]
197 impl ToolSpec for ImageAnalyzeTool {
198 fn name(&self) -> &str {
199 "image_analyze"
200 }
201
202 fn description(&self) -> &str {
203 "Analyze an image using the configured vision model. \
204 Supports PNG, JPEG, GIF, WebP, and BMP formats."
205 }
206
207 fn input_schema(&self) -> Value {
208 json!({
209 "type": "object",
210 "properties": {
211 "image_path": {
212 "type": "string",
213 "description": "Path to the image file to analyze"
214 },
215 "prompt": {
216 "type": "string",
217 "description": "Optional prompt to guide the analysis."
218 }
219 },
220 "required": ["image_path"]
221 })
222 }
223
224 fn capabilities(&self) -> Vec<ToolCapability> {
225 vec![ToolCapability::ReadOnly]
226 }
227
228 async fn execute(&self, input: Value, context: &ToolContext) -> Result<ToolResult, ToolError> {
229 let image_path = required_str(&input, "image_path")?;
230 let prompt = input
231 .get("prompt")
232 .and_then(|v| v.as_str())
233 .unwrap_or("Describe this image in detail.");
234
235 let resolved_path = Self::resolve_image_path(&context.workspace, image_path)?;
236 let (image_data, mime_type) = Self::read_image_file(&resolved_path).await?;
237
238 let payload = self.request_payload(prompt, &image_data, &mime_type);
239
240 let url = format!("{}/chat/completions", self.base_url());
241 let api_key = self.api_key();
242
243 let retry_config = RetryConfig {
244 max_retries: 3,
245 initial_delay: 1.0,
246 max_delay: 30.0,
247 enabled: true,
248 ..Default::default()
249 };
250 let _inference = match self.route_client.as_ref() {
251 Some(client) => client.acquire_remote_control_inference_permit().await,
252 None => Some(crate::client::acquire_remote_control_inference_participant().await),
253 };
254
255 let response = with_retry(
256 &retry_config,
257 || {
258 let client = self.client.clone();
259 let url = url.clone();
260 let api_key = api_key.clone();
261 let payload = payload.clone();
262 async move {
263 let response = client
264 .post(&url)
265 .header("Content-Type", "application/json")
266 .header("Authorization", format!("Bearer {api_key}"))
267 .json(&payload)
268 .send()
269 .await
270 .map_err(|e| LlmError::from_reqwest(&e))?;
271
272 let status = response.status();
273 if !status.is_success() {
274 let error_text = response
275 .text()
276 .await
277 .unwrap_or_else(|_| "Unknown error".to_string());
278 let error_text = sanitize_http_error_body(
279 Some("Vision provider"),
280 status.as_u16(),
281 &error_text,
282 );
283 return Err(LlmError::from_http_response(status.as_u16(), &error_text));
284 }
285 Ok(response)
286 }
287 },
288 None,
289 )
290 .await
291 .map_err(|e| ToolError::execution_failed(format!("Vision API request failed: {e}")))?;
292
293 let json: Value = response
294 .json()
295 .await
296 .map_err(|e| ToolError::execution_failed(format!("Failed to parse response: {e}")))?;
297
298 let content = json
299 .get("choices")
300 .and_then(|c| c.get(0))
301 .and_then(|c| c.get("message"))
302 .and_then(|m| m.get("content"))
303 .and_then(|c| c.as_str())
304 .unwrap_or("")
305 .to_string();
306
307 let model = json
308 .get("model")
309 .and_then(|m| m.as_str())
310 .unwrap_or(&self.config.model)
311 .to_string();
312
313 let result = json!({
314 "analysis": content,
315 "model": model,
316 });
317
318 ToolResult::json(&result)
319 .map_err(|e| ToolError::execution_failed(format!("Failed to serialize result: {e}")))
320 }
321 }
322
323 #[cfg(test)]
324 mod tests {
325 use super::*;
326 use tempfile::tempdir;
327
328 #[cfg(unix)]
329 fn create_file_symlink(
330 target: &std::path::Path,
331 link: &std::path::Path,
332 ) -> std::io::Result<()> {
333 std::os::unix::fs::symlink(target, link)
334 }
335
336 #[cfg(windows)]
337 fn create_file_symlink(
338 target: &std::path::Path,
339 link: &std::path::Path,
340 ) -> std::io::Result<()> {
341 std::os::windows::fs::symlink_file(target, link)
342 }
343
344 fn fake_config() -> VisionModelConfig {
345 VisionModelConfig {
346 model: "test-vision-model".to_string(),
347 api_key: Some("test-key".to_string()),
348 base_url: Some("https://example.invalid/v1".to_string()),
349 }
350 }
351
352 fn standalone_vision_cap(model: &str) -> u64 {
353 u64::from(
354 crate::route_budget::effective_max_output_tokens_for_route(
355 ApiProvider::Custom,
356 model,
357 None,
358 )
359 .min(65_536),
360 )
361 }
362
363 #[test]
364 fn tool_metadata_is_read_only_and_named_image_analyze() {
365 let tool = ImageAnalyzeTool::new(fake_config());
366 assert_eq!(tool.name(), "image_analyze");
367 assert!(tool.capabilities().contains(&ToolCapability::ReadOnly));
368 }
369
370 #[test]
371 fn mime_type_detection_covers_common_formats() {
372 for (ext, expected) in [
373 ("png", "image/png"),
374 ("PNG", "image/png"),
375 ("jpg", "image/jpeg"),
376 ("jpeg", "image/jpeg"),
377 ("gif", "image/gif"),
378 ("webp", "image/webp"),
379 ("bmp", "image/bmp"),
380 ] {
381 let path = std::path::PathBuf::from(format!("test.{ext}"));
382 let mime = ImageAnalyzeTool::detect_mime_type(&path)
383 .unwrap_or_else(|_| panic!("must detect {ext}"));
384 assert_eq!(mime, expected);
385 }
386 }
387
388 #[test]
389 fn mime_type_detection_rejects_unsupported_extension() {
390 let path = std::path::PathBuf::from("test.svg");
391 let err = ImageAnalyzeTool::detect_mime_type(&path)
392 .expect_err("svg is intentionally out of scope for vision tool");
393 assert!(err.to_string().contains("Unsupported image format"));
394 }
395
396 #[test]
397 fn generic_vision_payload_uses_max_tokens() {
398 let tool = ImageAnalyzeTool::new(fake_config());
399
400 let payload = tool.request_payload("describe", "abc123", "image/png");
401
402 assert_eq!(
403 payload.get("max_tokens").and_then(Value::as_u64),
404 Some(standalone_vision_cap(&tool.config.model))
405 );
406 assert!(payload.get("temperature").is_none());
407 assert!(payload.get("max_completion_tokens").is_none());
408 }
409
410 #[test]
411 fn xiaomi_mimo_vision_payload_uses_max_completion_tokens() {
412 let mut config = fake_config();
413 config.model = "mimo-v2.5".to_string();
414 config.base_url = Some("https://api.xiaomimimo.com/v1".to_string());
415 let tool = ImageAnalyzeTool::new(config);
416
417 let payload = tool.request_payload("describe", "abc123", "image/png");
418
419 assert_eq!(
420 payload.get("max_completion_tokens").and_then(Value::as_u64),
421 Some(standalone_vision_cap(&tool.config.model))
422 );
423 assert!(payload.get("temperature").is_none());
424 assert!(payload.get("max_tokens").is_none());
425 }
426
427 #[test]
428 fn xiaomi_mimo_vision_payload_uses_max_completion_tokens_with_custom_proxy() {
429 let mut config = fake_config();
430 config.model = "mimo-v2.5".to_string();
431 config.base_url = Some("https://vision-proxy.example.invalid/v1".to_string());
432 let tool = ImageAnalyzeTool::new(config);
433
434 let payload = tool.request_payload("describe", "abc123", "image/png");
435
436 assert_eq!(
437 payload.get("max_completion_tokens").and_then(Value::as_u64),
438 Some(standalone_vision_cap(&tool.config.model))
439 );
440 assert!(payload.get("max_tokens").is_none());
441 }
442
443 #[test]
444 fn vision_vendor_pin_requires_the_matching_bound_route() {
445 let _lock = crate::test_support::lock_test_env();
446 let base_url = "http://127.0.0.1:18080/v1";
447 let client = CodewhaleClient::new(&crate::config::Config {
448 provider: Some("openrouter".to_string()),
449 providers: Some(crate::config::ProvidersConfig {
450 openrouter: crate::config::ProviderConfig {
451 api_key: Some("fixture-openrouter-key".to_string()),
452 base_url: Some(base_url.to_string()),
453 model: Some("fixture/vision".to_string()),
454 vendor: Some("chutes/region-fixture".to_string()),
455 ..Default::default()
456 },
457 ..Default::default()
458 }),
459 ..Default::default()
460 })
461 .unwrap();
462 for (vision_base, matched_client, pinned) in [
463 (base_url, Some(client.clone()), true),
464 ("http://127.0.0.1:18081/v1", Some(client.clone()), false),
465 (base_url, None, false),
466 ] {
467 let tool = ImageAnalyzeTool::new_with_route_client(
468 VisionModelConfig {
469 model: "fixture/vision".to_string(),
470 api_key: Some("fixture-vision-key".to_string()),
471 base_url: Some(vision_base.to_string()),
472 },
473 matched_client,
474 );
475 let body = tool.request_payload("describe", "abc123", "image/png");
476 if pinned {
477 assert_eq!(
478 body["provider"],
479 json!({
480 "order": ["chutes/region-fixture"], "allow_fallbacks": false
481 })
482 );
483 } else {
484 assert!(body.get("provider").is_none());
485 }
486 }
487 }
488
489 #[test]
490 fn matched_vision_route_uses_bound_client_window_cap() {
491 let _lock = crate::test_support::lock_test_env();
492 let _canonical =
493 crate::test_support::EnvVarGuard::set("CODEWHALE_MAX_OUTPUT_TOKENS", "384000");
494 let base_url = "http://127.0.0.1:18080/v1".to_string();
495 let model = "DeepSeek-V4-Flash".to_string();
496 let client = CodewhaleClient::new(&crate::config::Config {
497 provider: Some("vllm".to_string()),
498 providers: Some(crate::config::ProvidersConfig {
499 vllm: crate::config::ProviderConfig {
500 base_url: Some(base_url.clone()),
501 model: Some(model.clone()),
502 context_window: Some(327_680),
503 ..crate::config::ProviderConfig::default()
504 },
505 ..crate::config::ProvidersConfig::default()
506 }),
507 ..crate::config::Config::default()
508 })
509 .expect("bound vLLM client");
510 let tool = ImageAnalyzeTool::new_with_route_client(
511 VisionModelConfig {
512 model,
513 api_key: None,
514 base_url: Some(base_url),
515 },
516 Some(client),
517 );
518
519 let payload = tool.request_payload("describe", "abc123", "image/png");
520 assert_eq!(payload["max_tokens"], 325_632);
521 }
522
523 #[tokio::test]
524 async fn execute_rejects_absolute_path() {
525 // Trust-boundary pin: image_path must stay inside the workspace
526 // — an absolute path or a `..`-traversing path must reject
527 // before any base64 / API call.
528 let tmp = tempdir().expect("tempdir");
529 let ctx = ToolContext::new(tmp.path().to_path_buf());
530 let tool = ImageAnalyzeTool::new(fake_config());
531 let outside_workspace = if cfg!(windows) {
532 r"C:\Windows\System32\drivers\etc\hosts"
533 } else {
534 "/etc/hosts"
535 };
536 let err = tool
537 .execute(json!({"image_path": outside_workspace}), &ctx)
538 .await
539 .expect_err("absolute path must reject");
540 assert!(
541 err.to_string()
542 .contains("relative path within the workspace"),
543 "error must call out the workspace boundary; got {err}"
544 );
545 }
546
547 #[tokio::test]
548 async fn execute_rejects_parent_dir_traversal() {
549 let tmp = tempdir().expect("tempdir");
550 let ctx = ToolContext::new(tmp.path().to_path_buf());
551 let tool = ImageAnalyzeTool::new(fake_config());
552 let err = tool
553 .execute(json!({"image_path": "../escape.png"}), &ctx)
554 .await
555 .expect_err("`..`-traversal must reject");
556 assert!(
557 err.to_string()
558 .contains("relative path within the workspace"),
559 "error must call out the workspace boundary; got {err}"
560 );
561 }
562
563 #[tokio::test]
564 async fn execute_rejects_symlink_that_resolves_outside_workspace() {
565 let workspace = tempdir().expect("workspace tempdir");
566 let outside = tempdir().expect("outside tempdir");
567 let outside_image = outside.path().join("outside.png");
568 std::fs::write(&outside_image, b"not a real png").expect("write outside image");
569 let link = workspace.path().join("linked.png");
570 if let Err(err) = create_file_symlink(&outside_image, &link) {
571 eprintln!("skipping symlink assertion: {err}");
572 return;
573 }
574
575 let ctx = ToolContext::new(workspace.path().to_path_buf());
576 let tool = ImageAnalyzeTool::new(fake_config());
577 let err = tool
578 .execute(json!({"image_path": "linked.png"}), &ctx)
579 .await
580 .expect_err("symlink target outside workspace must reject before reading");
581 assert!(
582 err.to_string().contains("resolve within the workspace"),
583 "error must call out the canonical workspace boundary; got {err}"
584 );
585 }
586 }
587
587 lines RUST