From 553c306c52170a1340e426815e459e94b8e14f4f Mon Sep 17 00:00:00 2001 From: Illia Polosukhin Date: Mon, 9 Mar 2026 03:41:27 +0000 Subject: [PATCH] feat: full image support across all channels (#725) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: full image support across all channels End-to-end image handling: upload, generation, analysis, editing, and rendering across web gateway, HTTP webhook, WASM (Telegram/Slack), and REPL channels. Builds on the attachment infrastructure from #596 and draws inspiration from PR #641's image pipeline approach — credit to that PR's author for the sentinel JSON pattern and base64-in-JSON upload design. Key changes: - Image upload in web UI (file picker, paste, preview strip) - Image generation tool (FLUX/DALL-E via /v1/images/generations) - Image edit tool (multipart /v1/images/edits with fallback) - Image analysis tool (vision model for workspace images) - Model detection utilities (image_models.rs, vision_models.rs) - Sentinel JSON detection in dispatcher for generated image rendering - StatusUpdate::ImageGenerated → SSE/WS/REPL/WASM broadcast - HTTP webhook attachment support (base64, 5MB/file, 10MB total) - WASM channel image download (Telegram via file API, Slack via host HTTP) - Tool registration wiring in app.rs [skip-regression-check] Co-Authored-By: Claude Opus 4.6 * fix: address PR #725 review comments (16 issues) - SecretString for API keys in all image tools (image_gen, image_edit, image_analyze) - Binary image read via tokio::fs::read instead of DB-backed workspace.read() - Replace Arc with Option base_dir (workspace has no filesystem API) - ApprovalRequirement::UnlessAutoApproved for cost-sensitive image tools - Scope sentinel detection to image_generate/image_edit tool names only - Skip ToolResult preview broadcast for image sentinels (avoids multi-MB base64 in SSE) - Extract shared media_type_from_path() to builtin/mod.rs - Rename fallback_chat_edit → fallback_generate with tracing::warn - Increase gateway body limit from 1MB to 10MB for image uploads - Increase webhook body limit to 15MB (base64 overhead) - Log warning on invalid base64 in images_to_attachments - Client-side image size limits (5MB/file, 5 images max) in app.js - aria-label on attach button for accessibility - Update body_too_large test for new 10MB limit [skip-regression-check] Co-Authored-By: Claude Opus 4.6 * fix: add Slack file size check before download (PR review item #15) Skip downloading files larger than 20 MB in the Slack WASM channel to avoid excessive memory use and slow downloads in the WASM runtime. Logs a warning when a file is skipped. Also bumps channel versions for Slack and Telegram (prior branch changes). [skip-regression-check] Co-Authored-By: Claude Opus 4.6 * style: cargo fmt Co-Authored-By: Claude Opus 4.6 * fix(security): add path validation and approval requirement to image tools Add sandbox path validation via validate_path() to both ImageAnalyzeTool and ImageEditTool to prevent path traversal attacks that could exfiltrate arbitrary files through external vision/edit APIs. Also fix ImageAnalyzeTool::requires_approval to return UnlessAutoApproved, consistent with ImageEditTool and ImageGenerateTool. Co-Authored-By: Claude Opus 4.6 * fix: post-download size guards and empty data_url sentinel check - Slack: add post-download size check on actual bytes when metadata size_bytes is absent, preventing bypass of the 20MB limit - Telegram: add 20MB download size limit (matching Slack) enforced in download_telegram_file() after receiving response bytes - Dispatcher: skip broadcasting ImageGenerated SSE event when data_url is empty from unwrap_or_default(), log warning instead Closes correctness issues #3, #4, #5 from PR #725 review. Co-Authored-By: Claude Opus 4.6 * fix: use mime_guess for media type detection, add alt attrs and media_type validation - Replace hardcoded media type mapping with mime_guess crate (already in deps) - Add alt attributes to img elements in web UI for accessibility - Validate media_type starts with "image/" in images_to_attachments() - Update bmp test assertion to match mime_guess behavior Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: Zaki --- .gitignore | 1 + channels-src/slack/Cargo.lock | 2 +- channels-src/slack/Cargo.toml | 2 +- channels-src/slack/src/lib.rs | 104 ++++++++++ channels-src/telegram/Cargo.lock | 2 +- channels-src/telegram/Cargo.toml | 2 +- channels-src/telegram/src/lib.rs | 62 +++++- registry/channels/slack.json | 2 +- registry/channels/telegram.json | 2 +- src/agent/dispatcher.rs | 92 ++++++++- src/app.rs | 43 ++++ src/channels/channel.rs | 7 + src/channels/http.rs | 132 +++++++++++- src/channels/repl.rs | 7 + src/channels/wasm/wrapper.rs | 8 + src/channels/web/mod.rs | 5 + src/channels/web/server.rs | 63 +++++- src/channels/web/sse.rs | 1 + src/channels/web/static/app.js | 118 ++++++++++- src/channels/web/static/index.html | 3 + src/channels/web/static/style.css | 91 ++++++++ src/channels/web/types.rs | 26 +++ src/channels/web/ws.rs | 9 + src/llm/image_models.rs | 95 +++++++++ src/llm/mod.rs | 3 + src/llm/vision_models.rs | 104 ++++++++++ src/tools/builtin/image_analyze.rs | 250 ++++++++++++++++++++++ src/tools/builtin/image_edit.rs | 322 +++++++++++++++++++++++++++++ src/tools/builtin/image_gen.rs | 251 ++++++++++++++++++++++ src/tools/builtin/mod.rs | 16 ++ src/tools/registry.rs | 49 +++++ tests/openai_compat_integration.rs | 4 +- 32 files changed, 1851 insertions(+), 27 deletions(-) create mode 100644 src/llm/image_models.rs create mode 100644 src/llm/vision_models.rs create mode 100644 src/tools/builtin/image_analyze.rs create mode 100644 src/tools/builtin/image_edit.rs create mode 100644 src/tools/builtin/image_gen.rs diff --git a/.gitignore b/.gitignore index 17bdb86d..f03e691c 100644 --- a/.gitignore +++ b/.gitignore @@ -23,4 +23,5 @@ bench-results/ # WASM build artifacts (loaded from disk, not bundled) *.wasm +# Traces trace_*.json diff --git a/channels-src/slack/Cargo.lock b/channels-src/slack/Cargo.lock index 4e646b06..08b69e1c 100644 --- a/channels-src/slack/Cargo.lock +++ b/channels-src/slack/Cargo.lock @@ -267,7 +267,7 @@ dependencies = [ [[package]] name = "slack-channel" -version = "0.1.0" +version = "0.2.1" dependencies = [ "hex", "hmac", diff --git a/channels-src/slack/Cargo.toml b/channels-src/slack/Cargo.toml index bc8c7434..e5445abb 100644 --- a/channels-src/slack/Cargo.toml +++ b/channels-src/slack/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "slack-channel" -version = "0.2.0" +version = "0.2.1" edition = "2021" description = "Slack Events API channel for IronClaw" license = "MIT OR Apache-2.0" diff --git a/channels-src/slack/src/lib.rs b/channels-src/slack/src/lib.rs index 71f1e731..24f01df3 100644 --- a/channels-src/slack/src/lib.rs +++ b/channels-src/slack/src/lib.rs @@ -357,10 +357,108 @@ fn extract_slack_attachments(files: &Option>) -> Vec Result, String> { + let headers = serde_json::json!({}); + + let result = channel_host::http_request("GET", url, &headers.to_string(), None, None); + + let response = result.map_err(|e| format!("Slack file download failed: {}", e))?; + + if response.status != 200 { + let body_str = String::from_utf8_lossy(&response.body); + return Err(format!( + "Slack file download returned {}: {}", + response.status, body_str + )); + } + + Ok(response.body) +} + +/// Download file bytes and store them via the host for processing. +/// +/// Downloads all file types (images, documents, etc.) so the host-side +/// middleware can process them (vision pipeline for images, text extraction +/// for documents, transcription for audio, etc.). +/// Maximum file size to download (20 MB). Files larger than this are skipped +/// to avoid excessive memory use and slow downloads in the WASM runtime. +const MAX_DOWNLOAD_SIZE_BYTES: u64 = 20 * 1024 * 1024; + +fn download_and_store_slack_files(attachments: &[InboundAttachment]) { + for att in attachments { + let Some(ref url) = att.source_url else { + continue; + }; + + // Skip files that exceed the size limit + if let Some(size) = att.size_bytes { + if size > MAX_DOWNLOAD_SIZE_BYTES { + channel_host::log( + channel_host::LogLevel::Warn, + &format!( + "Skipping Slack file download: {} bytes exceeds {} MB limit (id={})", + size, + MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024), + att.id + ), + ); + continue; + } + } + + match download_slack_file(url) { + Ok(bytes) => { + // Post-download size guard: metadata size_bytes is optional, + // so a file with no size info could bypass the pre-download check. + if bytes.len() as u64 > MAX_DOWNLOAD_SIZE_BYTES { + channel_host::log( + channel_host::LogLevel::Warn, + &format!( + "Discarding Slack file after download: {} bytes exceeds {} MB limit (id={})", + bytes.len(), + MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024), + att.id + ), + ); + continue; + } + + channel_host::log( + channel_host::LogLevel::Info, + &format!( + "Downloaded Slack file: {} bytes, mime={}", + bytes.len(), + att.mime_type + ), + ); + if let Err(e) = channel_host::store_attachment_data(&att.id, &bytes) { + channel_host::log( + channel_host::LogLevel::Error, + &format!("Failed to store Slack file data: {}", e), + ); + } + } + Err(e) => { + channel_host::log( + channel_host::LogLevel::Error, + &format!("Failed to download Slack file: {}", e), + ); + } + } + } +} + /// Handle a Slack event and emit message if applicable. fn handle_slack_event(event: SlackEvent, team_id: Option, _event_id: Option) { let attachments = extract_slack_attachments(&event.files); + // Download and store file attachments for host-side processing + download_and_store_slack_files(&attachments); + match event.event_type.as_str() { // Direct mention of the bot (always in a channel, not a DM) "app_mention" => { @@ -722,4 +820,10 @@ mod tests { let event: SlackEvent = serde_json::from_str(json).unwrap(); assert!(event.files.is_none()); } + + #[test] + fn test_max_download_size_constant() { + // Verify the constant is 20 MB + assert_eq!(MAX_DOWNLOAD_SIZE_BYTES, 20 * 1024 * 1024); + } } diff --git a/channels-src/telegram/Cargo.lock b/channels-src/telegram/Cargo.lock index 67c27867..8d40f01e 100644 --- a/channels-src/telegram/Cargo.lock +++ b/channels-src/telegram/Cargo.lock @@ -212,7 +212,7 @@ dependencies = [ [[package]] name = "telegram-channel" -version = "0.2.0" +version = "0.2.1" dependencies = [ "serde", "serde_json", diff --git a/channels-src/telegram/Cargo.toml b/channels-src/telegram/Cargo.toml index 93a1eb57..182e5f5d 100644 --- a/channels-src/telegram/Cargo.toml +++ b/channels-src/telegram/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "telegram-channel" -version = "0.2.0" +version = "0.2.1" edition = "2021" description = "Telegram Bot API channel for IronClaw" license = "MIT OR Apache-2.0" diff --git a/channels-src/telegram/src/lib.rs b/channels-src/telegram/src/lib.rs index c3ab9050..d8718ebb 100644 --- a/channels-src/telegram/src/lib.rs +++ b/channels-src/telegram/src/lib.rs @@ -878,10 +878,6 @@ fn send_message( // Voice File Download // ============================================================================ -/// Download a voice file from Telegram by file_id. -/// -/// 1. Call getFile to get the file_path. -/// 2. Download the file bytes from /file/bot{TOKEN}/{file_path}. /// Percent-encode a string for safe use as a URL query parameter value. fn percent_encode(s: &str) -> String { let mut out = String::with_capacity(s.len()); @@ -898,6 +894,10 @@ fn percent_encode(s: &str) -> String { out } +/// Maximum file size to download (20 MB). Files larger than this are discarded +/// to avoid excessive memory use and slow downloads in the WASM runtime. +const MAX_DOWNLOAD_SIZE_BYTES: u64 = 20 * 1024 * 1024; + fn download_telegram_file(file_id: &str) -> Result, String> { // Reject file_id containing curly braces to prevent credential placeholder injection if file_id.contains('{') || file_id.contains('}') { @@ -965,6 +965,16 @@ fn download_telegram_file(file_id: &str) -> Result, String> { )); } + // Post-download size guard: Telegram metadata file_size is optional, + // so enforce the limit on actual downloaded bytes. + if response.body.len() as u64 > MAX_DOWNLOAD_SIZE_BYTES { + return Err(format!( + "Downloaded file exceeds {} MB limit ({} bytes)", + MAX_DOWNLOAD_SIZE_BYTES / (1024 * 1024), + response.body.len() + )); + } + Ok(response.body) } @@ -1535,6 +1545,39 @@ fn download_and_store_voice(attachments: &[InboundAttachment]) { } } +/// Download image file bytes and store them via the host for the vision pipeline. +/// +/// Separated from `extract_attachments` so that function stays pure (no host +/// calls) and remains testable in native unit tests. +fn download_and_store_images(attachments: &[InboundAttachment]) { + for att in attachments { + if !att.mime_type.starts_with("image/") { + continue; + } + + match download_telegram_file(&att.id) { + Ok(bytes) => { + channel_host::log( + channel_host::LogLevel::Info, + &format!("Downloaded image file: {} bytes", bytes.len()), + ); + if let Err(e) = channel_host::store_attachment_data(&att.id, &bytes) { + channel_host::log( + channel_host::LogLevel::Error, + &format!("Failed to store image data: {}", e), + ); + } + } + Err(e) => { + channel_host::log( + channel_host::LogLevel::Error, + &format!("Failed to download image file: {}", e), + ); + } + } + } +} + /// Returns true if the attachment should be downloaded for document text extraction. /// /// Excludes voice (handled by transcription), image (vision pipeline), @@ -1608,6 +1651,9 @@ fn handle_message(message: TelegramMessage) { // Download and store voice attachments for host-side transcription download_and_store_voice(&attachments); + // Download and store image attachments for host-side vision pipeline + download_and_store_images(&attachments); + // Download and store document attachments for host-side text extraction download_and_store_documents(&mut attachments); @@ -1681,7 +1727,7 @@ fn handle_message(message: TelegramMessage) { let username_opt = from.username.as_deref(); let is_allowed = allowed.contains(&"*".to_string()) || allowed.contains(&id_str) - || username_opt.map_or(false, |u| allowed.contains(&u.to_string())); + || username_opt.is_some_and(|u| allowed.contains(&u.to_string())); if !is_allowed { if is_private && dm_policy == "pairing" { @@ -2605,4 +2651,10 @@ mod tests { assert!(!is_downloadable_document(&make("audio/mpeg", Some("song.mp3")))); assert!(!is_downloadable_document(&make("video/mp4", Some("clip.mp4")))); } + + #[test] + fn test_max_download_size_constant() { + // Verify the constant is 20 MB, matching the Slack channel limit + assert_eq!(MAX_DOWNLOAD_SIZE_BYTES, 20 * 1024 * 1024); + } } diff --git a/registry/channels/slack.json b/registry/channels/slack.json index 58a6e10e..f123798f 100644 --- a/registry/channels/slack.json +++ b/registry/channels/slack.json @@ -2,7 +2,7 @@ "name": "slack", "display_name": "Slack Channel", "kind": "channel", - "version": "0.2.0", + "version": "0.2.1", "wit_version": "0.3.0", "description": "Talk to your agent in Slack", "keywords": [ diff --git a/registry/channels/telegram.json b/registry/channels/telegram.json index d28234f9..45bf5426 100644 --- a/registry/channels/telegram.json +++ b/registry/channels/telegram.json @@ -2,7 +2,7 @@ "name": "telegram", "display_name": "Telegram Channel", "kind": "channel", - "version": "0.2.0", + "version": "0.2.1", "wit_version": "0.3.0", "description": "Talk to your agent through a Telegram bot", "keywords": [ diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 2754f4d6..b59ff92f 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -681,8 +681,53 @@ impl Agent { .into()) }); - // Send ToolResult preview - if let Ok(ref output) = tool_result + // Detect image generation sentinel in tool output + // (only from image tools — avoids parsing all tool outputs) + let is_image_sentinel = if let Ok(ref output) = tool_result + && matches!(tc.name.as_str(), "image_generate" | "image_edit") + { + if let Ok(sentinel) = + serde_json::from_str::(output) + && sentinel.get("type").and_then(|v| v.as_str()) + == Some("image_generated") + { + let data_url = sentinel + .get("data") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + let path = sentinel + .get("path") + .and_then(|v| v.as_str()) + .map(String::from); + // Skip broadcasting if data_url is empty to avoid + // sending a broken ImageGenerated SSE event. + if data_url.is_empty() { + tracing::warn!( + "Image generation sentinel has empty data URL, skipping broadcast" + ); + } else { + let _ = self + .channels + .send_status( + &message.channel, + StatusUpdate::ImageGenerated { data_url, path }, + &message.metadata, + ) + .await; + } + true + } else { + false + } + } else { + false + }; + + // Send ToolResult preview (skip for image sentinels to avoid + // broadcasting multi-MB base64 data as a preview) + if !is_image_sentinel + && let Ok(ref output) = tool_result && !output.is_empty() { let _ = self @@ -2124,4 +2169,47 @@ mod tests { "Error should include the underlying reason, got: {formatted}" ); } + + #[test] + fn test_image_sentinel_empty_data_url_should_be_skipped() { + // Regression: unwrap_or_default() on missing "data" field produces an empty + // string. Broadcasting an empty data_url would send a broken SSE event. + let sentinel = serde_json::json!({ + "type": "image_generated", + "path": "/tmp/image.png" + // "data" field is missing + }); + + let data_url = sentinel + .get("data") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + + assert!( + data_url.is_empty(), + "Missing 'data' field should produce empty string" + ); + // The fix: empty data_url means we skip broadcasting + } + + #[test] + fn test_image_sentinel_present_data_url_is_valid() { + let sentinel = serde_json::json!({ + "type": "image_generated", + "data": "data:image/png;base64,abc123", + "path": "/tmp/image.png" + }); + + let data_url = sentinel + .get("data") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + + assert!( + !data_url.is_empty(), + "Present 'data' field should produce non-empty string" + ); + } } diff --git a/src/app.rs b/src/app.rs index 738d659c..42b4c569 100644 --- a/src/app.rs +++ b/src/app.rs @@ -400,6 +400,49 @@ impl AppBuilder { None }; + // Register image/vision tools if we have a workspace and LLM API credentials + if workspace.is_some() { + let (api_base, api_key_opt) = if let Some(ref provider) = self.config.llm.provider { + ( + provider.base_url.clone(), + provider.api_key.as_ref().map(|s| { + use secrecy::ExposeSecret; + s.expose_secret().to_string() + }), + ) + } else { + ( + self.config.llm.nearai.base_url.clone(), + self.config.llm.nearai.api_key.as_ref().map(|s| { + use secrecy::ExposeSecret; + s.expose_secret().to_string() + }), + ) + }; + + if let Some(api_key) = api_key_opt { + // Check for image generation models + let model_name = self + .config + .llm + .provider + .as_ref() + .map(|p| p.model.clone()) + .unwrap_or_else(|| self.config.llm.nearai.model.clone()); + let models = vec![model_name.clone()]; + let gen_model = crate::llm::image_models::suggest_image_model(&models) + .unwrap_or("flux-1.1-pro") + .to_string(); + tools.register_image_tools(api_base.clone(), api_key.clone(), gen_model, None); + + // Check for vision models + let vision_model = crate::llm::vision_models::suggest_vision_model(&models) + .unwrap_or(&model_name) + .to_string(); + tools.register_vision_tools(api_base, api_key, vision_model, None); + } + } + // Register builder tool if enabled if self.config.builder.enabled && (self.config.agent.allow_local_tools || !self.config.sandbox.enabled) diff --git a/src/channels/channel.rs b/src/channels/channel.rs index 3ab5c1f6..e126ca1f 100644 --- a/src/channels/channel.rs +++ b/src/channels/channel.rs @@ -231,6 +231,13 @@ pub enum StatusUpdate { success: bool, message: String, }, + /// An image was generated by a tool. + ImageGenerated { + /// Base64 data URL of the generated image. + data_url: String, + /// Optional workspace path where the image was saved. + path: Option, + }, } impl StatusUpdate { diff --git a/src/channels/http.rs b/src/channels/http.rs index 87cd2051..74799b04 100644 --- a/src/channels/http.rs +++ b/src/channels/http.rs @@ -17,7 +17,9 @@ use tokio::sync::{RwLock, mpsc, oneshot}; use tokio_stream::wrappers::ReceiverStream; use uuid::Uuid; -use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse}; +use crate::channels::{ + AttachmentKind, Channel, IncomingAttachment, IncomingMessage, MessageStream, OutgoingResponse, +}; use crate::config::HttpConfig; use crate::error::ChannelError; @@ -46,8 +48,9 @@ struct RateLimitState { request_count: u32, } -/// Maximum JSON body size for webhook requests (64 KB). -const MAX_BODY_BYTES: usize = 64 * 1024; +/// Maximum JSON body size for webhook requests (15 MB, to support base64 image attachments +/// with ~33% overhead from base64 encoding). +const MAX_BODY_BYTES: usize = 15 * 1024 * 1024; /// Maximum number of pending wait-for-response requests. const MAX_PENDING_RESPONSES: usize = 100; @@ -115,8 +118,34 @@ struct WebhookRequest { /// Whether to wait for a synchronous response. #[serde(default)] wait_for_response: bool, + /// Optional file attachments (base64-encoded). + #[serde(default)] + attachments: Vec, } +/// A file attachment in a webhook request. +#[derive(Debug, Deserialize)] +struct AttachmentData { + /// MIME type (e.g. "image/png", "application/pdf"). + mime_type: String, + /// Optional filename. + #[serde(default)] + filename: Option, + /// Base64-encoded file data. + #[serde(default)] + data_base64: Option, + /// URL to fetch the file from (not downloaded server-side for SSRF prevention). + #[serde(default)] + url: Option, +} + +/// Maximum size per attachment (5 MB decoded). +const MAX_ATTACHMENT_BYTES: usize = 5 * 1024 * 1024; +/// Maximum total attachment size (10 MB decoded). +const MAX_TOTAL_ATTACHMENT_BYTES: usize = 10 * 1024 * 1024; +/// Maximum number of attachments per request. +const MAX_ATTACHMENTS: usize = 5; + #[derive(Debug, Serialize)] struct WebhookResponse { /// Message ID assigned to this request. @@ -211,15 +240,106 @@ async fn webhook_handler( ); } - let msg = IncomingMessage::new("http", &state.user_id, &req.content).with_metadata( + // Validate and decode attachments + let attachments = if !req.attachments.is_empty() { + if req.attachments.len() > MAX_ATTACHMENTS { + return ( + StatusCode::BAD_REQUEST, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some(format!("Too many attachments (max {})", MAX_ATTACHMENTS)), + }), + ); + } + + let mut decoded_attachments = Vec::new(); + let mut total_bytes: usize = 0; + for att in &req.attachments { + if let Some(ref b64) = att.data_base64 { + use base64::Engine; + let data = match base64::engine::general_purpose::STANDARD.decode(b64) { + Ok(d) => d, + Err(_) => { + return ( + StatusCode::BAD_REQUEST, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Invalid base64 in attachment".to_string()), + }), + ); + } + }; + if data.len() > MAX_ATTACHMENT_BYTES { + return ( + StatusCode::PAYLOAD_TOO_LARGE, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some(format!( + "Attachment too large (max {} bytes)", + MAX_ATTACHMENT_BYTES + )), + }), + ); + } + total_bytes += data.len(); + if total_bytes > MAX_TOTAL_ATTACHMENT_BYTES { + return ( + StatusCode::PAYLOAD_TOO_LARGE, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Total attachment size exceeds limit".to_string()), + }), + ); + } + decoded_attachments.push(IncomingAttachment { + id: Uuid::new_v4().to_string(), + kind: AttachmentKind::from_mime_type(&att.mime_type), + mime_type: att.mime_type.clone(), + filename: att.filename.clone(), + size_bytes: Some(data.len() as u64), + source_url: None, + storage_key: None, + extracted_text: None, + data, + duration_secs: None, + }); + } else if let Some(ref url) = att.url { + // URL-only attachment: set source_url but don't download (SSRF prevention) + decoded_attachments.push(IncomingAttachment { + id: Uuid::new_v4().to_string(), + kind: AttachmentKind::from_mime_type(&att.mime_type), + mime_type: att.mime_type.clone(), + filename: att.filename.clone(), + size_bytes: None, + source_url: Some(url.clone()), + storage_key: None, + extracted_text: None, + data: Vec::new(), + duration_secs: None, + }); + } + } + decoded_attachments + } else { + Vec::new() + }; + + let mut msg = IncomingMessage::new("http", &state.user_id, &req.content).with_metadata( serde_json::json!({ "wait_for_response": req.wait_for_response, }), ); + if !attachments.is_empty() { + msg = msg.with_attachments(attachments); + } + if let Some(thread_id) = &req.thread_id { - let msg = msg.with_thread(thread_id); - return process_message(state, msg, req.wait_for_response).await; + msg = msg.with_thread(thread_id); } process_message(state, msg, req.wait_for_response).await diff --git a/src/channels/repl.rs b/src/channels/repl.rs index b1f06ec2..33adc23f 100644 --- a/src/channels/repl.rs +++ b/src/channels/repl.rs @@ -600,6 +600,13 @@ impl Channel for ReplChannel { eprintln!("\x1b[31m {extension_name}: {message}\x1b[0m"); } } + StatusUpdate::ImageGenerated { path, .. } => { + if let Some(ref p) = path { + eprintln!("\x1b[36m [image] {p}\x1b[0m"); + } else { + eprintln!("\x1b[36m [image generated]\x1b[0m"); + } + } } Ok(()) } diff --git a/src/channels/wasm/wrapper.rs b/src/channels/wasm/wrapper.rs index cac0cb1f..3b788e89 100644 --- a/src/channels/wasm/wrapper.rs +++ b/src/channels/wasm/wrapper.rs @@ -2809,6 +2809,14 @@ fn status_to_wit(status: &StatusUpdate, metadata: &serde_json::Value) -> wit_cha ), metadata_json, }, + StatusUpdate::ImageGenerated { path, .. } => wit_channel::StatusUpdate { + status: wit_channel::StatusType::Status, + message: match path { + Some(p) => format!("[image] {}", p), + None => "[image generated]".to_string(), + }, + metadata_json, + }, } } diff --git a/src/channels/web/mod.rs b/src/channels/web/mod.rs index 92e8ac5f..0fcf228e 100644 --- a/src/channels/web/mod.rs +++ b/src/channels/web/mod.rs @@ -386,6 +386,11 @@ impl Channel for GatewayChannel { success, message, }, + StatusUpdate::ImageGenerated { data_url, path } => SseEvent::ImageGenerated { + data_url, + path, + thread_id, + }, }; self.state.sse.broadcast(event); diff --git a/src/channels/web/server.rs b/src/channels/web/server.rs index 1c6e7f85..d6605eee 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -351,7 +351,7 @@ pub async fn start_server( .merge(statics) .merge(projects) .merge(protected) - .layer(DefaultBodyLimit::max(1024 * 1024)) // 1 MB max request body + .layer(DefaultBodyLimit::max(10 * 1024 * 1024)) // 10 MB max request body (image uploads) .layer(cors) .layer(SetResponseHeaderLayer::if_not_present( header::X_CONTENT_TYPE_OPTIONS, @@ -608,6 +608,56 @@ async fn oauth_callback_handler( // --- Chat handlers --- +/// Convert web gateway `ImageData` to `IncomingAttachment` objects. +pub(crate) fn images_to_attachments( + images: &[ImageData], +) -> Vec { + use base64::Engine; + images + .iter() + .enumerate() + .filter_map(|(i, img)| { + if !img.media_type.starts_with("image/") { + tracing::warn!( + "Skipping image {i}: invalid media type '{}' (must start with 'image/')", + img.media_type + ); + return None; + } + let data = match base64::engine::general_purpose::STANDARD.decode(&img.data) { + Ok(d) => d, + Err(e) => { + tracing::warn!("Skipping image {i}: invalid base64 data: {e}"); + return None; + } + }; + Some(crate::channels::IncomingAttachment { + id: format!("web-image-{i}"), + kind: crate::channels::AttachmentKind::Image, + mime_type: img.media_type.clone(), + filename: Some(format!("image-{i}.{}", mime_to_ext(&img.media_type))), + size_bytes: Some(data.len() as u64), + source_url: None, + storage_key: None, + extracted_text: None, + data, + duration_secs: None, + }) + }) + .collect() +} + +/// Map MIME type to file extension. +fn mime_to_ext(mime: &str) -> &str { + match mime { + "image/png" => "png", + "image/gif" => "gif", + "image/webp" => "webp", + "image/svg+xml" => "svg", + _ => "jpg", + } +} + async fn chat_send_handler( State(state): State>, headers: axum::http::HeaderMap, @@ -641,11 +691,18 @@ async fn chat_send_handler( msg = msg.with_metadata(serde_json::json!({"thread_id": thread_id})); } + // Convert uploaded images to IncomingAttachments + if !req.images.is_empty() { + let attachments = images_to_attachments(&req.images); + msg = msg.with_attachments(attachments); + } + let msg_id = msg.id; tracing::debug!( - "[chat_send_handler] Created message id={}, content={:?}", + "[chat_send_handler] Created message id={}, content={:?}, images={}", msg_id, - req.content + req.content, + req.images.len() ); let tx_guard = state.msg_tx.read().await; diff --git a/src/channels/web/sse.rs b/src/channels/web/sse.rs index e1e2b270..6d9c4142 100644 --- a/src/channels/web/sse.rs +++ b/src/channels/web/sse.rs @@ -142,6 +142,7 @@ impl SseManager { SseEvent::JobStatus { .. } => "job_status", SseEvent::JobResult { .. } => "job_result", SseEvent::Heartbeat => "heartbeat", + SseEvent::ImageGenerated { .. } => "image_generated", SseEvent::ExtensionStatus { .. } => "extension_status", }; Ok(Event::default().event(event_type).data(data)) diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index 87c83ede..573ce5f2 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -18,6 +18,7 @@ let unreadThreads = new Map(); // thread_id -> unread count let _loadThreadsTimer = null; const JOB_EVENTS_CAP = 500; const MEMORY_SEARCH_QUERY_MAX_LENGTH = 100; +let stagedImages = []; // --- Slash Commands --- @@ -389,6 +390,12 @@ function connectSSE() { if (currentTab === 'extensions') loadExtensions(); }); + eventSource.addEventListener('image_generated', (e) => { + const data = JSON.parse(e.data); + if (!isCurrentThread(data.thread_id)) return; + addGeneratedImage(data.data_url, data.path); + }); + eventSource.addEventListener('error', (e) => { if (e.data) { const data = JSON.parse(e.data); @@ -446,16 +453,23 @@ function sendMessage() { return; } const content = input.value.trim(); - if (!content) return; + if (!content && stagedImages.length === 0) return; - addMessage('user', content); + addMessage('user', content || '(images attached)'); input.value = ''; autoResizeTextarea(input); input.focus(); + const body = { content, thread_id: currentThreadId || undefined, timezone: Intl.DateTimeFormat().resolvedOptions().timeZone }; + if (stagedImages.length > 0) { + body.images = stagedImages.map(img => ({ media_type: img.media_type, data: img.data })); + stagedImages = []; + renderImagePreviews(); + } + apiFetch('/api/chat/send', { method: 'POST', - body: { content, thread_id: currentThreadId || undefined, timezone: Intl.DateTimeFormat().resolvedOptions().timeZone }, + body: body, }).catch((err) => { addMessage('system', 'Failed to send: ' + err.message); }); @@ -472,6 +486,104 @@ function enableChatInput() { if (btn) btn.disabled = false; } +// --- Image Upload --- + +function renderImagePreviews() { + const strip = document.getElementById('image-preview-strip'); + strip.innerHTML = ''; + stagedImages.forEach((img, idx) => { + const container = document.createElement('div'); + container.className = 'image-preview-container'; + + const preview = document.createElement('img'); + preview.className = 'image-preview'; + preview.src = img.dataUrl; + preview.alt = 'Attached image'; + + const removeBtn = document.createElement('button'); + removeBtn.className = 'image-preview-remove'; + removeBtn.textContent = '\u00d7'; + removeBtn.addEventListener('click', () => { + stagedImages.splice(idx, 1); + renderImagePreviews(); + }); + + container.appendChild(preview); + container.appendChild(removeBtn); + strip.appendChild(container); + }); +} + +const MAX_IMAGE_SIZE_BYTES = 5 * 1024 * 1024; // 5 MB per image +const MAX_STAGED_IMAGES = 5; + +function handleImageFiles(files) { + Array.from(files).forEach(file => { + if (!file.type.startsWith('image/')) return; + if (file.size > MAX_IMAGE_SIZE_BYTES) { + alert(`Image "${file.name}" exceeds 5 MB limit (${(file.size / 1024 / 1024).toFixed(1)} MB)`); + return; + } + if (stagedImages.length >= MAX_STAGED_IMAGES) { + alert(`Maximum ${MAX_STAGED_IMAGES} images allowed per message`); + return; + } + const reader = new FileReader(); + reader.onload = function(e) { + const dataUrl = e.target.result; + const commaIdx = dataUrl.indexOf(','); + const meta = dataUrl.substring(0, commaIdx); // e.g. "data:image/png;base64" + const base64 = dataUrl.substring(commaIdx + 1); + const mediaType = meta.replace('data:', '').replace(';base64', ''); + stagedImages.push({ media_type: mediaType, data: base64, dataUrl: dataUrl }); + renderImagePreviews(); + }; + reader.readAsDataURL(file); + }); +} + +document.getElementById('attach-btn').addEventListener('click', () => { + document.getElementById('image-file-input').click(); +}); + +document.getElementById('image-file-input').addEventListener('change', (e) => { + handleImageFiles(e.target.files); + e.target.value = ''; +}); + +document.getElementById('chat-input').addEventListener('paste', (e) => { + const items = (e.clipboardData || e.originalEvent.clipboardData).items; + for (let i = 0; i < items.length; i++) { + if (items[i].kind === 'file' && items[i].type.startsWith('image/')) { + const file = items[i].getAsFile(); + if (file) handleImageFiles([file]); + } + } +}); + +function addGeneratedImage(dataUrl, path) { + const container = document.getElementById('chat-messages'); + const card = document.createElement('div'); + card.className = 'generated-image-card'; + + const img = document.createElement('img'); + img.className = 'generated-image'; + img.src = dataUrl; + img.alt = 'Generated image'; + + card.appendChild(img); + + if (path) { + const pathLabel = document.createElement('div'); + pathLabel.className = 'generated-image-path'; + pathLabel.textContent = path; + card.appendChild(pathLabel); + } + + container.appendChild(card); + container.scrollTop = container.scrollHeight; +} + // --- Slash Autocomplete --- function showSlashAutocomplete(matches) { diff --git a/src/channels/web/static/index.html b/src/channels/web/static/index.html index be8a0c9e..385b0086 100644 --- a/src/channels/web/static/index.html +++ b/src/channels/web/static/index.html @@ -130,7 +130,10 @@
+
+ +
diff --git a/src/channels/web/static/style.css b/src/channels/web/static/style.css index a21775fb..192e63f5 100644 --- a/src/channels/web/static/style.css +++ b/src/channels/web/static/style.css @@ -1272,6 +1272,7 @@ body { /* Chat input */ .chat-input { display: flex; + flex-wrap: wrap; padding: 12px 16px max(12px, env(safe-area-inset-bottom)) 16px; gap: 8px; background: var(--bg-secondary); @@ -3761,3 +3762,93 @@ mark { text-overflow: ellipsis; white-space: nowrap; } + +/* Image Upload */ +.attach-btn { + background: none; + border: none; + cursor: pointer; + font-size: 1.2em; + padding: 8px; + align-self: flex-end; + color: var(--text-secondary); + transition: color 0.2s; + min-height: 40px; + display: flex; + align-items: center; + justify-content: center; +} + +.attach-btn:hover { + color: var(--text); +} + +.image-preview-strip { + display: flex; + flex-direction: row; + gap: 8px; + padding: 4px; + overflow-x: auto; + min-height: 0; + width: 100%; +} + +.image-preview-strip:empty { + display: none; +} + +.image-preview-container { + position: relative; + display: inline-block; + flex-shrink: 0; +} + +.image-preview { + width: 60px; + height: 60px; + border-radius: 6px; + object-fit: cover; + display: block; +} + +.image-preview-remove { + position: absolute; + top: -6px; + right: -6px; + width: 18px; + height: 18px; + border-radius: 50%; + background: var(--danger); + color: #fff; + border: none; + font-size: 12px; + line-height: 18px; + text-align: center; + cursor: pointer; + padding: 0; +} + +.image-preview-remove:hover { + background: #c33; +} + +/* Generated Image */ +.generated-image-card { + max-width: 512px; + margin: 8px 0; + border-radius: 8px; + overflow: hidden; + border: 1px solid var(--border); +} + +.generated-image { + max-width: 100%; + display: block; +} + +.generated-image-path { + font-size: 12px; + color: var(--text-secondary); + padding: 4px 8px; + background: var(--bg-secondary); +} diff --git a/src/channels/web/types.rs b/src/channels/web/types.rs index 7d65965d..4d85c671 100644 --- a/src/channels/web/types.rs +++ b/src/channels/web/types.rs @@ -5,11 +5,23 @@ use uuid::Uuid; // --- Chat --- +/// Base64-encoded image data sent from the web frontend. +#[derive(Debug, Clone, Deserialize)] +pub struct ImageData { + /// MIME type (e.g., "image/png", "image/jpeg"). + pub media_type: String, + /// Base64-encoded image data (without data: URL prefix). + pub data: String, +} + #[derive(Debug, Deserialize)] pub struct SendMessageRequest { pub content: String, pub thread_id: Option, pub timezone: Option, + /// Optional images attached to the message. + #[serde(default)] + pub images: Vec, } #[derive(Debug, Serialize)] @@ -220,6 +232,16 @@ pub enum SseEvent { session_id: Option, }, + /// An image was generated by a tool. + #[serde(rename = "image_generated")] + ImageGenerated { + data_url: String, + #[serde(skip_serializing_if = "Option::is_none")] + path: Option, + #[serde(skip_serializing_if = "Option::is_none")] + thread_id: Option, + }, + /// Extension activation status change (WASM channels). #[serde(rename = "extension_status")] ExtensionStatus { @@ -615,6 +637,9 @@ pub enum WsClientMessage { content: String, thread_id: Option, timezone: Option, + /// Optional images attached to the message. + #[serde(default)] + images: Vec, }, /// Approve or deny a pending tool execution. #[serde(rename = "approval")] @@ -681,6 +706,7 @@ impl WsServerMessage { SseEvent::JobToolResult { .. } => "job_tool_result", SseEvent::JobStatus { .. } => "job_status", SseEvent::JobResult { .. } => "job_result", + SseEvent::ImageGenerated { .. } => "image_generated", SseEvent::ExtensionStatus { .. } => "extension_status", }; let data = serde_json::to_value(event).unwrap_or(serde_json::Value::Null); diff --git a/src/channels/web/ws.rs b/src/channels/web/ws.rs index e9e3c8e6..1736ae7e 100644 --- a/src/channels/web/ws.rs +++ b/src/channels/web/ws.rs @@ -160,6 +160,7 @@ async fn handle_client_message( content, thread_id, timezone, + images, } => { let mut incoming = IncomingMessage::new("gateway", user_id, &content); if let Some(ref tz) = timezone { @@ -169,6 +170,12 @@ async fn handle_client_message( incoming = incoming.with_thread(tid); } + // Convert uploaded images to IncomingAttachments + if !images.is_empty() { + let attachments = crate::channels::web::server::images_to_attachments(&images); + incoming = incoming.with_attachments(attachments); + } + let tx_guard = state.msg_tx.read().await; if let Some(ref tx) = *tx_guard { if tx.send(incoming).await.is_err() { @@ -357,6 +364,7 @@ mod tests { content: "hello agent".to_string(), thread_id: Some("t1".to_string()), timezone: None, + images: Vec::new(), }, &state, "user1", @@ -382,6 +390,7 @@ mod tests { content: "hello".to_string(), thread_id: None, timezone: None, + images: Vec::new(), }, &state, "user1", diff --git a/src/llm/image_models.rs b/src/llm/image_models.rs new file mode 100644 index 00000000..651c6703 --- /dev/null +++ b/src/llm/image_models.rs @@ -0,0 +1,95 @@ +//! Image generation model detection utilities. + +/// Known image generation model families. +const IMAGE_GEN_PATTERNS: &[&str] = &[ + "flux", + "dall-e", + "dalle", + "stable-diffusion", + "sdxl", + "imagen", + "midjourney", + "ideogram", + "playground", +]; + +/// Check if a model name indicates an image generation model. +pub fn is_image_generation_model(model: &str) -> bool { + let lower = model.to_lowercase(); + IMAGE_GEN_PATTERNS.iter().any(|p| lower.contains(p)) +} + +/// Suggest the best image generation model from a list of available models. +/// +/// Priority: FLUX > DALL-E > Stable Diffusion > others. +pub fn suggest_image_model(models: &[String]) -> Option<&str> { + let priorities: &[&str] = &[ + "flux", + "dall-e", + "dalle", + "stable-diffusion", + "sdxl", + "imagen", + ]; + for priority in priorities { + if let Some(model) = models.iter().find(|m| m.to_lowercase().contains(priority)) { + return Some(model); + } + } + // Fall back to any image gen model + models.iter().find_map(|m| { + if is_image_generation_model(m) { + Some(m.as_str()) + } else { + None + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn detects_flux_models() { + assert!(is_image_generation_model( + "black-forest-labs/FLUX.1-schnell" + )); + assert!(is_image_generation_model("flux-pro")); + } + + #[test] + fn detects_dalle_models() { + assert!(is_image_generation_model("dall-e-3")); + assert!(is_image_generation_model("dalle-3")); + } + + #[test] + fn rejects_non_image_models() { + assert!(!is_image_generation_model("gpt-4o")); + assert!(!is_image_generation_model("claude-3-sonnet")); + assert!(!is_image_generation_model("llama-3.1-70b")); + } + + #[test] + fn suggests_flux_first() { + let models = vec![ + "gpt-4o".to_string(), + "dall-e-3".to_string(), + "flux-pro".to_string(), + ]; + assert_eq!(suggest_image_model(&models), Some("flux-pro")); + } + + #[test] + fn suggests_dalle_without_flux() { + let models = vec!["gpt-4o".to_string(), "dall-e-3".to_string()]; + assert_eq!(suggest_image_model(&models), Some("dall-e-3")); + } + + #[test] + fn returns_none_when_no_image_models() { + let models = vec!["gpt-4o".to_string(), "claude-3-sonnet".to_string()]; + assert_eq!(suggest_image_model(&models), None); + } +} diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 8945a887..388ad290 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -22,6 +22,9 @@ mod rig_adapter; pub mod session; pub mod smart_routing; +pub mod image_models; +pub mod vision_models; + pub use circuit_breaker::{CircuitBreakerConfig, CircuitBreakerProvider}; pub use failover::{CooldownConfig, FailoverProvider}; pub use nearai_chat::{ModelInfo, NearAiChatProvider}; diff --git a/src/llm/vision_models.rs b/src/llm/vision_models.rs new file mode 100644 index 00000000..27e1b1d9 --- /dev/null +++ b/src/llm/vision_models.rs @@ -0,0 +1,104 @@ +//! Vision model detection utilities. + +/// Known vision-capable model families. +const VISION_PATTERNS: &[&str] = &[ + "claude-3", + "claude-4", + "gpt-4o", + "gpt-4-turbo", + "gpt-4-vision", + "gemini-pro-vision", + "gemini-1.5", + "gemini-2", + "llava", + "cogvlm", + "internvl", + "qwen-vl", + "qwen2-vl", + "pixtral", +]; + +/// Check if a model name indicates vision capabilities. +pub fn is_vision_model(model: &str) -> bool { + let lower = model.to_lowercase(); + VISION_PATTERNS.iter().any(|p| lower.contains(p)) +} + +/// Suggest the best vision model from a list of available models. +/// +/// Priority: Claude > GPT-4 > Gemini > others. +pub fn suggest_vision_model(models: &[String]) -> Option<&str> { + let priorities: &[&str] = &[ + "claude-3", + "claude-4", + "gpt-4o", + "gpt-4-turbo", + "gpt-4-vision", + "gemini", + "llava", + "pixtral", + ]; + for priority in priorities { + if let Some(model) = models.iter().find(|m| m.to_lowercase().contains(priority)) { + return Some(model); + } + } + models.iter().find_map(|m| { + if is_vision_model(m) { + Some(m.as_str()) + } else { + None + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn detects_claude_vision() { + assert!(is_vision_model("claude-3-5-sonnet-20241022")); + assert!(is_vision_model("claude-3-opus")); + assert!(is_vision_model("claude-4-sonnet")); + } + + #[test] + fn detects_gpt4_vision() { + assert!(is_vision_model("gpt-4o")); + assert!(is_vision_model("gpt-4-turbo")); + assert!(is_vision_model("gpt-4-vision-preview")); + } + + #[test] + fn detects_other_vision_models() { + assert!(is_vision_model("gemini-1.5-pro")); + assert!(is_vision_model("llava-v1.6")); + assert!(is_vision_model("pixtral-12b")); + } + + #[test] + fn rejects_non_vision_models() { + assert!(!is_vision_model("gpt-3.5-turbo")); + assert!(!is_vision_model("llama-3.1-70b")); + assert!(!is_vision_model("mistral-7b")); + } + + #[test] + fn suggests_claude_first() { + let models = vec![ + "gpt-4o".to_string(), + "claude-3-5-sonnet-20241022".to_string(), + ]; + assert_eq!( + suggest_vision_model(&models), + Some("claude-3-5-sonnet-20241022") + ); + } + + #[test] + fn returns_none_when_no_vision_models() { + let models = vec!["gpt-3.5-turbo".to_string(), "llama-3.1-70b".to_string()]; + assert_eq!(suggest_vision_model(&models), None); + } +} diff --git a/src/tools/builtin/image_analyze.rs b/src/tools/builtin/image_analyze.rs new file mode 100644 index 00000000..b1f8a62f --- /dev/null +++ b/src/tools/builtin/image_analyze.rs @@ -0,0 +1,250 @@ +//! Image analysis tool using vision-capable LLM models. + +use std::path::PathBuf; + +use async_trait::async_trait; +use base64::Engine; +use secrecy::{ExposeSecret, SecretString}; + +use crate::context::JobContext; +use crate::tools::builtin::path_utils::validate_path; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput}; + +/// Tool for analyzing images using a vision-capable model. +pub struct ImageAnalyzeTool { + /// API base URL. + api_base_url: String, + /// Bearer token for API auth. + api_key: SecretString, + /// Vision-capable model name. + model: String, + /// HTTP client. + client: reqwest::Client, + /// Optional base directory for resolving relative image paths. + base_dir: Option, +} + +impl ImageAnalyzeTool { + /// Create a new image analysis tool. + pub fn new( + api_base_url: String, + api_key: String, + model: String, + base_dir: Option, + ) -> Self { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(120)) + .build() + .unwrap_or_default(); + Self { + api_base_url, + api_key: SecretString::from(api_key), + model, + client, + base_dir, + } + } + + /// Read binary image bytes from filesystem. + /// + /// Validates the path against the base directory sandbox to prevent + /// path traversal attacks, then reads the file bytes. + async fn read_image_bytes(&self, image_path: &str) -> Result, ToolError> { + let resolved = validate_path(image_path, self.base_dir.as_deref())?; + + tokio::fs::read(&resolved) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Failed to read image file: {e}"))) + } +} + +#[async_trait] +impl Tool for ImageAnalyzeTool { + fn name(&self) -> &str { + "image_analyze" + } + + fn description(&self) -> &str { + "Analyze an image using a vision-capable AI model. Provide a workspace path to the image and an optional analysis question." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "image_path": { + "type": "string", + "description": "Path to the image file in the workspace (e.g., 'images/photo.jpg')" + }, + "question": { + "type": "string", + "description": "Specific question to answer about the image. Defaults to general analysis.", + "default": "Describe this image in detail." + } + }, + "required": ["image_path"] + }) + } + + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { + ApprovalRequirement::UnlessAutoApproved + } + + fn requires_sanitization(&self) -> bool { + true + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let image_path = params + .get("image_path") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing required 'image_path' parameter".to_string()) + })?; + + let question = params + .get("question") + .and_then(|v| v.as_str()) + .unwrap_or("Describe this image in detail."); + + // Read binary image bytes directly from filesystem + let image_bytes = self.read_image_bytes(image_path).await?; + if image_bytes.is_empty() { + return Err(ToolError::ExecutionFailed( + "Image file is empty".to_string(), + )); + } + + let media_type = super::media_type_from_path(image_path); + let b64 = base64::engine::general_purpose::STANDARD.encode(&image_bytes); + let data_url = format!("data:{media_type};base64,{b64}"); + + // Call vision model via chat completions API + let url = format!( + "{}/v1/chat/completions", + self.api_base_url.trim_end_matches('/') + ); + + let request_body = serde_json::json!({ + "model": &self.model, + "messages": [{ + "role": "user", + "content": [ + { + "type": "text", + "text": question + }, + { + "type": "image_url", + "image_url": { + "url": data_url + } + } + ] + }], + "max_tokens": 2048 + }); + + let response = self + .client + .post(&url) + .bearer_auth(self.api_key.expose_secret()) + .json(&request_body) + .send() + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Vision API request failed: {e}")))?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(ToolError::ExecutionFailed(format!( + "Vision API returned {status}: {body}" + ))); + } + + let resp: serde_json::Value = response.json().await.map_err(|e| { + ToolError::ExecutionFailed(format!("Failed to parse vision API response: {e}")) + })?; + + let analysis = resp + .pointer("/choices/0/message/content") + .and_then(|v| v.as_str()) + .unwrap_or("No analysis available."); + + Ok(ToolOutput::text(analysis, start.elapsed())) + } +} + +#[cfg(test)] +mod tests { + use super::super::media_type_from_path; + use super::*; + use tempfile::TempDir; + + #[test] + fn test_media_type_detection() { + assert_eq!(media_type_from_path("photo.png"), "image/png"); + assert_eq!(media_type_from_path("photo.jpg"), "image/jpeg"); + assert_eq!(media_type_from_path("photo.jpeg"), "image/jpeg"); + assert_eq!(media_type_from_path("photo.gif"), "image/gif"); + assert_eq!(media_type_from_path("photo.webp"), "image/webp"); + assert_eq!(media_type_from_path("photo.bmp"), "image/bmp"); + assert_eq!(media_type_from_path("photo.svg"), "image/svg+xml"); + } + + #[test] + fn test_requires_approval_returns_unless_auto_approved() { + let tool = ImageAnalyzeTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "gpt-4o".to_string(), + None, + ); + assert_eq!( + tool.requires_approval(&serde_json::json!({})), + ApprovalRequirement::UnlessAutoApproved + ); + } + + #[tokio::test] + async fn test_read_image_bytes_rejects_path_traversal() { + let dir = TempDir::new().unwrap(); + let tool = ImageAnalyzeTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "gpt-4o".to_string(), + Some(dir.path().to_path_buf()), + ); + + let result = tool.read_image_bytes("../../etc/passwd").await; + assert!( + result.is_err(), + "Should reject path traversal, got: {:?}", + result + ); + } + + #[tokio::test] + async fn test_read_image_bytes_rejects_absolute_path_outside_sandbox() { + let dir = TempDir::new().unwrap(); + let tool = ImageAnalyzeTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "gpt-4o".to_string(), + Some(dir.path().to_path_buf()), + ); + + let result = tool.read_image_bytes("/etc/passwd").await; + assert!( + result.is_err(), + "Should reject absolute path outside sandbox, got: {:?}", + result + ); + } +} diff --git a/src/tools/builtin/image_edit.rs b/src/tools/builtin/image_edit.rs new file mode 100644 index 00000000..818454cc --- /dev/null +++ b/src/tools/builtin/image_edit.rs @@ -0,0 +1,322 @@ +//! Image editing tool using cloud API. + +use std::path::PathBuf; + +use async_trait::async_trait; +use secrecy::{ExposeSecret, SecretString}; + +use crate::context::JobContext; +use crate::tools::builtin::path_utils::validate_path; +use crate::tools::tool::{ApprovalRequirement, Tool, ToolError, ToolOutput}; + +/// Tool for editing images using an AI image editing API. +pub struct ImageEditTool { + /// API base URL. + api_base_url: String, + /// Bearer token for API auth. + api_key: SecretString, + /// Model to use. + model: String, + /// HTTP client. + client: reqwest::Client, + /// Optional base directory for resolving relative image paths. + base_dir: Option, +} + +impl ImageEditTool { + /// Create a new image edit tool. + pub fn new( + api_base_url: String, + api_key: String, + model: String, + base_dir: Option, + ) -> Self { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(180)) + .build() + .unwrap_or_default(); + Self { + api_base_url, + api_key: SecretString::from(api_key), + model, + client, + base_dir, + } + } + + /// Read binary image bytes from filesystem. + /// + /// Validates the path against the base directory sandbox to prevent + /// path traversal attacks, then reads the file bytes. + async fn read_image_bytes(&self, image_path: &str) -> Result, ToolError> { + let resolved = validate_path(image_path, self.base_dir.as_deref())?; + + tokio::fs::read(&resolved) + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Failed to read image file: {e}"))) + } +} + +#[async_trait] +impl Tool for ImageEditTool { + fn name(&self) -> &str { + "image_edit" + } + + fn description(&self) -> &str { + "Edit an existing image using an AI model. Provide the workspace path to the source image and a text prompt describing the desired edits." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Text description of the edits to apply to the image", + "maxLength": 4000 + }, + "image_path": { + "type": "string", + "description": "Path to the source image in the workspace (e.g., 'images/photo.jpg')" + } + }, + "required": ["prompt", "image_path"] + }) + } + + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { + ApprovalRequirement::UnlessAutoApproved + } + + fn requires_sanitization(&self) -> bool { + false + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let prompt = params + .get("prompt") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing required 'prompt' parameter".to_string()) + })?; + + let image_path = params + .get("image_path") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing required 'image_path' parameter".to_string()) + })?; + + if prompt.len() > 4000 { + return Err(ToolError::InvalidParameters( + "Prompt exceeds 4000 character limit".to_string(), + )); + } + + // Read binary image bytes directly from filesystem + let image_bytes = self.read_image_bytes(image_path).await?; + if image_bytes.is_empty() { + return Err(ToolError::ExecutionFailed( + "Source image file is empty".to_string(), + )); + } + + let media_type = super::media_type_from_path(image_path); + + // Use multipart form for image edit API + let url = format!( + "{}/v1/images/edits", + self.api_base_url.trim_end_matches('/') + ); + + let form = reqwest::multipart::Form::new() + .text("model", self.model.clone()) + .text("prompt", prompt.to_string()) + .text("response_format", "b64_json") + .part( + "image", + reqwest::multipart::Part::bytes(image_bytes) + .mime_str(&media_type) + .map_err(|e| ToolError::ExecutionFailed(format!("Invalid media type: {e}")))? + .file_name("image"), + ); + + let response = self + .client + .post(&url) + .bearer_auth(self.api_key.expose_secret()) + .multipart(form) + .send() + .await + .map_err(|e| ToolError::ExecutionFailed(format!("Image edit request failed: {e}")))?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + + // Fall back to generation if edits endpoint not available + if status.as_u16() == 404 { + tracing::warn!( + "Image edit endpoint returned 404, falling back to generation API. \ + Note: the source image will NOT be used — a new image will be generated from the prompt alone." + ); + return self.fallback_generate(prompt, start).await; + } + + return Err(ToolError::ExecutionFailed(format!( + "Image edit API returned {status}: {body}" + ))); + } + + let resp: serde_json::Value = response.json().await.map_err(|e| { + ToolError::ExecutionFailed(format!("Failed to parse image edit response: {e}")) + })?; + + let edited_data = resp + .pointer("/data/0/b64_json") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExecutionFailed("No image data in edit response".to_string()) + })?; + + let sentinel = serde_json::json!({ + "type": "image_generated", + "data": format!("data:image/png;base64,{}", edited_data), + "media_type": "image/png", + "prompt": prompt, + "source_path": image_path + }); + + Ok(ToolOutput::text(sentinel.to_string(), start.elapsed())) + } +} + +impl ImageEditTool { + /// Fallback: generate a new image from the prompt when the edit endpoint is unavailable. + /// + /// The source image is NOT used — this generates a completely new image. + /// The response includes a `note` field warning the user. + async fn fallback_generate( + &self, + prompt: &str, + start: std::time::Instant, + ) -> Result { + let url = format!( + "{}/v1/images/generations", + self.api_base_url.trim_end_matches('/') + ); + + let request_body = serde_json::json!({ + "model": &self.model, + "prompt": prompt, + "size": "1024x1024", + "response_format": "b64_json", + "n": 1 + }); + + let response = self + .client + .post(&url) + .bearer_auth(self.api_key.expose_secret()) + .json(&request_body) + .send() + .await + .map_err(|e| { + ToolError::ExecutionFailed(format!("Fallback image generation failed: {e}")) + })?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(ToolError::ExecutionFailed(format!( + "Fallback generation API returned {status}: {body}" + ))); + } + + let resp: serde_json::Value = response.json().await.map_err(|e| { + ToolError::ExecutionFailed(format!("Failed to parse fallback response: {e}")) + })?; + + let image_data = resp + .pointer("/data/0/b64_json") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExecutionFailed("No image data in fallback response".to_string()) + })?; + + let sentinel = serde_json::json!({ + "type": "image_generated", + "data": format!("data:image/png;base64,{}", image_data), + "media_type": "image/png", + "prompt": prompt, + "note": "Generated new image (edit endpoint unavailable — source image was NOT used)" + }); + + Ok(ToolOutput::text(sentinel.to_string(), start.elapsed())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use tempfile::TempDir; + + #[test] + fn test_tool_metadata() { + let tool = ImageEditTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "flux-1".to_string(), + None, + ); + assert_eq!(tool.name(), "image_edit"); + assert!(!tool.requires_sanitization()); + assert_eq!( + tool.requires_approval(&serde_json::json!({})), + ApprovalRequirement::UnlessAutoApproved + ); + } + + #[tokio::test] + async fn test_read_image_bytes_rejects_path_traversal() { + let dir = TempDir::new().unwrap(); + let tool = ImageEditTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "flux-1".to_string(), + Some(dir.path().to_path_buf()), + ); + + let result = tool.read_image_bytes("../../etc/passwd").await; + assert!( + result.is_err(), + "Should reject path traversal, got: {:?}", + result + ); + } + + #[tokio::test] + async fn test_read_image_bytes_rejects_absolute_path_outside_sandbox() { + let dir = TempDir::new().unwrap(); + let tool = ImageEditTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "flux-1".to_string(), + Some(dir.path().to_path_buf()), + ); + + let result = tool.read_image_bytes("/etc/passwd").await; + assert!( + result.is_err(), + "Should reject absolute path outside sandbox, got: {:?}", + result + ); + } +} diff --git a/src/tools/builtin/image_gen.rs b/src/tools/builtin/image_gen.rs new file mode 100644 index 00000000..c87b10d7 --- /dev/null +++ b/src/tools/builtin/image_gen.rs @@ -0,0 +1,251 @@ +//! Image generation tool using cloud API. + +use async_trait::async_trait; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; + +use crate::context::JobContext; +use crate::tools::tool::ApprovalRequirement; +use crate::tools::{Tool, ToolError, ToolOutput}; + +/// Tool for generating images using FLUX or compatible image generation APIs. +pub struct ImageGenerateTool { + /// API base URL (e.g., "https://cloud-api.near.ai"). + api_base_url: String, + /// Bearer token for API auth. + api_key: SecretString, + /// Model to use (e.g., "black-forest-labs/FLUX.1-schnell"). + model: String, + /// HTTP client. + client: reqwest::Client, +} + +#[derive(Debug, Serialize)] +struct ImageGenRequest { + model: String, + prompt: String, + size: String, + response_format: String, + n: u32, +} + +#[derive(Debug, Deserialize)] +struct ImageGenResponse { + data: Vec, +} + +#[derive(Debug, Deserialize)] +#[allow(dead_code)] +struct ImageGenData { + b64_json: Option, + url: Option, +} + +impl ImageGenerateTool { + /// Create a new image generation tool. + pub fn new(api_base_url: String, api_key: String, model: String) -> Self { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(180)) + .build() + .unwrap_or_default(); + Self { + api_base_url, + api_key: SecretString::from(api_key), + model, + client, + } + } +} + +#[async_trait] +impl Tool for ImageGenerateTool { + fn name(&self) -> &str { + "image_generate" + } + + fn description(&self) -> &str { + "Generate an image from a text prompt using an AI image generation model (e.g., FLUX). Returns the generated image data." + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "prompt": { + "type": "string", + "description": "Text description of the image to generate (max 4000 chars)", + "maxLength": 4000 + }, + "size": { + "type": "string", + "description": "Image dimensions", + "enum": ["1024x1024", "1792x1024", "1024x1792"], + "default": "1024x1024" + } + }, + "required": ["prompt"] + }) + } + + fn requires_approval(&self, _params: &serde_json::Value) -> ApprovalRequirement { + ApprovalRequirement::UnlessAutoApproved + } + + fn requires_sanitization(&self) -> bool { + false + } + + async fn execute( + &self, + params: serde_json::Value, + _ctx: &JobContext, + ) -> Result { + let start = std::time::Instant::now(); + + let prompt = params + .get("prompt") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::InvalidParameters("Missing required 'prompt' parameter".to_string()) + })?; + + if prompt.len() > 4000 { + return Err(ToolError::InvalidParameters( + "Prompt exceeds 4000 character limit".to_string(), + )); + } + + let size = params + .get("size") + .and_then(|v| v.as_str()) + .unwrap_or("1024x1024"); + + // Validate size + if !["1024x1024", "1792x1024", "1024x1792"].contains(&size) { + return Err(ToolError::InvalidParameters(format!( + "Invalid size '{}'. Must be 1024x1024, 1792x1024, or 1024x1792", + size + ))); + } + + let url = format!( + "{}/v1/images/generations", + self.api_base_url.trim_end_matches('/') + ); + + let request_body = ImageGenRequest { + model: self.model.clone(), + prompt: prompt.to_string(), + size: size.to_string(), + response_format: "b64_json".to_string(), + n: 1, + }; + + let response = self + .client + .post(&url) + .bearer_auth(self.api_key.expose_secret()) + .json(&request_body) + .send() + .await + .map_err(|e| { + ToolError::ExecutionFailed(format!("Image generation request failed: {e}")) + })?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(ToolError::ExecutionFailed(format!( + "Image generation API returned {status}: {body}" + ))); + } + + let gen_response: ImageGenResponse = response.json().await.map_err(|e| { + ToolError::ExecutionFailed(format!("Failed to parse image generation response: {e}")) + })?; + + let image_data = gen_response + .data + .first() + .and_then(|d| d.b64_json.as_deref()) + .ok_or_else(|| ToolError::ExecutionFailed("No image data in response".to_string()))?; + + // Return sentinel JSON for image display + let sentinel = serde_json::json!({ + "type": "image_generated", + "data": format!("data:image/png;base64,{}", image_data), + "media_type": "image/png", + "prompt": prompt, + "size": size + }); + + Ok(ToolOutput::text(sentinel.to_string(), start.elapsed())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_tool_metadata() { + let tool = ImageGenerateTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "flux-1".to_string(), + ); + assert_eq!(tool.name(), "image_generate"); + assert_eq!( + tool.requires_approval(&serde_json::json!({})), + ApprovalRequirement::UnlessAutoApproved + ); + + let schema = tool.parameters_schema(); + assert!(schema["properties"]["prompt"].is_object()); + assert!(schema["properties"]["size"].is_object()); + } + + #[tokio::test] + async fn test_missing_prompt() { + let tool = ImageGenerateTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "flux-1".to_string(), + ); + let ctx = JobContext::default(); + let result = tool.execute(serde_json::json!({}), &ctx).await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_invalid_size() { + let tool = ImageGenerateTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "flux-1".to_string(), + ); + let ctx = JobContext::default(); + let result = tool + .execute( + serde_json::json!({"prompt": "a cat", "size": "999x999"}), + &ctx, + ) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_prompt_too_long() { + let tool = ImageGenerateTool::new( + "https://api.example.com".to_string(), + "test-key".to_string(), + "flux-1".to_string(), + ); + let ctx = JobContext::default(); + let long_prompt = "x".repeat(4001); + let result = tool + .execute(serde_json::json!({"prompt": long_prompt}), &ctx) + .await; + assert!(result.is_err()); + } +} diff --git a/src/tools/builtin/mod.rs b/src/tools/builtin/mod.rs index bbbc7056..0b181986 100644 --- a/src/tools/builtin/mod.rs +++ b/src/tools/builtin/mod.rs @@ -40,5 +40,21 @@ pub use shell::ShellTool; pub use skill_tools::{SkillInstallTool, SkillListTool, SkillRemoveTool, SkillSearchTool}; pub use time::TimeTool; mod html_converter; +pub mod image_analyze; +pub mod image_edit; +pub mod image_gen; pub use html_converter::convert_html_to_markdown; +pub use image_analyze::ImageAnalyzeTool; +pub use image_edit::ImageEditTool; +pub use image_gen::ImageGenerateTool; + +/// Detect image media type from file extension via `mime_guess`. +/// Falls back to `image/jpeg` for unrecognized or non-image extensions. +pub(crate) fn media_type_from_path(path: &str) -> String { + mime_guess::from_path(path) + .first_raw() + .filter(|m| m.starts_with("image/")) + .unwrap_or("image/jpeg") + .to_string() +} diff --git a/src/tools/registry.rs b/src/tools/registry.rs index 44552541..7d78cc24 100644 --- a/src/tools/registry.rs +++ b/src/tools/registry.rs @@ -71,6 +71,9 @@ const PROTECTED_TOOL_NAMES: &[&str] = &[ "message", "web_fetch", "restart", + "image_generate", + "image_edit", + "image_analyze", ]; /// Registry of available tools. @@ -475,6 +478,52 @@ impl ToolRegistry { } } + /// Register image generation and editing tools. + /// + /// These tools allow the LLM to generate and edit images using cloud APIs. + /// Requires an API base URL, API key, and model name for the image generation backend. + pub fn register_image_tools( + &self, + api_base_url: String, + api_key: String, + gen_model: String, + base_dir: Option, + ) { + use crate::tools::builtin::{ImageEditTool, ImageGenerateTool}; + self.register_sync(Arc::new(ImageGenerateTool::new( + api_base_url.clone(), + api_key.clone(), + gen_model.clone(), + ))); + self.register_sync(Arc::new(ImageEditTool::new( + api_base_url, + api_key, + gen_model, + base_dir, + ))); + tracing::info!("Registered 2 image tools (generate, edit)"); + } + + /// Register vision/image analysis tools. + /// + /// These tools allow the LLM to analyze images using a vision-capable model. + pub fn register_vision_tools( + &self, + api_base_url: String, + api_key: String, + vision_model: String, + base_dir: Option, + ) { + use crate::tools::builtin::ImageAnalyzeTool; + self.register_sync(Arc::new(ImageAnalyzeTool::new( + api_base_url, + api_key, + vision_model, + base_dir, + ))); + tracing::info!("Registered 1 vision tool (analyze)"); + } + /// Register the software builder tool. /// /// The builder tool allows the agent to create new software including WASM tools, diff --git a/tests/openai_compat_integration.rs b/tests/openai_compat_integration.rs index 7f3f1e7f..501fa1aa 100644 --- a/tests/openai_compat_integration.rs +++ b/tests/openai_compat_integration.rs @@ -730,8 +730,8 @@ async fn test_chat_completions_body_too_large() { let (addr, _state, _mock_state) = start_test_server().await; let url = format!("http://{}/v1/chat/completions", addr); - // Build a payload over 1 MB (the gateway's DefaultBodyLimit) - let big_content = "x".repeat(2 * 1024 * 1024); + // Build a payload over 10 MB (the gateway's DefaultBodyLimit) + let big_content = "x".repeat(11 * 1024 * 1024); let resp = client() .post(&url) .bearer_auth(AUTH_TOKEN)