From 833738bc85fb6688d01c34a94668710b7a947acb Mon Sep 17 00:00:00 2001 From: Nick Pismenkov Date: Fri, 6 Mar 2026 18:41:06 -0800 Subject: [PATCH] review fixes --- src/agent/dispatcher.rs | 30 ++++++++++++--------- src/app.rs | 42 ++++++++++++++---------------- src/channels/web/mod.rs | 18 +++++++++---- src/channels/web/server.rs | 31 +++++++++++++--------- src/channels/web/sse.rs | 4 +++ src/channels/web/static/app.js | 42 ++++++++++++++++++++++++++---- src/llm/rig_adapter.rs | 2 +- src/tools/builtin/image_analyze.rs | 15 ++++++----- 8 files changed, 119 insertions(+), 65 deletions(-) diff --git a/src/agent/dispatcher.rs b/src/agent/dispatcher.rs index 6c2b518b..9c0b5e07 100644 --- a/src/agent/dispatcher.rs +++ b/src/agent/dispatcher.rs @@ -17,6 +17,16 @@ use crate::error::Error; use crate::llm::{ChatMessage, Reasoning, ReasoningContext, RespondResult}; use crate::tools::redact_params; +/// Represents image generation sentinel data in tool output. +#[derive(serde::Deserialize)] +struct ImageGeneratedSentinel<'a> { + #[serde(rename = "type")] + ty: &'a str, + data: &'a str, + media_type: &'a str, + path: &'a str, +} + /// Result of the agentic loop execution. pub(super) enum AgenticLoopResult { /// Completed with a response. @@ -642,25 +652,21 @@ impl Agent { .await; // Check for image_generated sentinel and emit SSE event - if let Ok(result_json) = - serde_json::from_str::(output) - && let Some("image_generated") = - result_json.get("type").and_then(|v| v.as_str()) - && let (Some(data), Some(media_type), Some(path)) = ( - result_json.get("data").and_then(|v| v.as_str()), - result_json.get("media_type").and_then(|v| v.as_str()), - result_json.get("path").and_then(|v| v.as_str()), - ) + if let Ok(sentinel) = + serde_json::from_str::(output) + && sentinel.ty == "image_generated" { - let data_url = - format!("data:{};base64,{}", media_type, data); + let data_url = format!( + "data:{};base64,{}", + sentinel.media_type, sentinel.data + ); let _ = self .channels .send_status( &message.channel, StatusUpdate::ImageGenerated { data_url, - path: path.to_string(), + path: sentinel.path.to_string(), }, &message.metadata, ) diff --git a/src/app.rs b/src/app.rs index efe2c7bd..2b0ee46c 100644 --- a/src/app.rs +++ b/src/app.rs @@ -380,19 +380,17 @@ impl AppBuilder { // Register image tools if image generation models are available match llm.list_models().await { Ok(models) => { - if crate::llm::image_models::has_image_generation_model(&models) { - if let Some(image_model) = - crate::llm::image_models::suggest_image_model(&models) - { - tools.register_image_tools( - self.config.llm.nearai.clone(), - Arc::clone(&ws), - ); - tracing::info!( - "Image generation tools registered (model: {})", - image_model - ); - } + if let Some(image_model) = + crate::llm::image_models::suggest_image_model(&models) + { + tools.register_image_tools( + self.config.llm.nearai.clone(), + Arc::clone(&ws), + ); + tracing::info!( + "Image generation tools registered (model: {})", + image_model + ); } else { tracing::debug!( "No image generation models detected in available models: {:?}", @@ -401,16 +399,14 @@ impl AppBuilder { } // Register vision analysis tool if vision models are available - if crate::llm::vision_models::has_vision_model(&models) { - if let Some(vision_model) = - crate::llm::vision_models::suggest_vision_model(&models) - { - tools.register_vision_tools(Arc::clone(&ws)); - tracing::info!( - "Image analysis tool registered (vision model: {})", - vision_model - ); - } + if let Some(vision_model) = + crate::llm::vision_models::suggest_vision_model(&models) + { + tools.register_vision_tools(Arc::clone(&ws)); + tracing::info!( + "Image analysis tool registered (vision model: {})", + vision_model + ); } else { tracing::debug!("No vision-capable models detected in available models"); } diff --git a/src/channels/web/mod.rs b/src/channels/web/mod.rs index aa292461..68710b4a 100644 --- a/src/channels/web/mod.rs +++ b/src/channels/web/mod.rs @@ -369,11 +369,19 @@ impl Channel for GatewayChannel { success, message, }, - StatusUpdate::ImageGenerated { data_url, path } => SseEvent::ImageGenerated { - data_url, - path, - thread_id, - }, + StatusUpdate::ImageGenerated { data_url, path } => { + tracing::debug!( + path = %path, + data_url_len = data_url.len(), + thread_id = ?thread_id, + "Broadcasting ImageGenerated SSE event" + ); + 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 44edb7cf..72c4feb5 100644 --- a/src/channels/web/server.rs +++ b/src/channels/web/server.rs @@ -963,18 +963,25 @@ async fn chat_history_handler( tool_calls: t .tool_calls .iter() - .map(|tc| ToolCallInfo { - name: tc.name.clone(), - has_result: tc.result.is_some(), - has_error: tc.error.is_some(), - result_preview: tc.result.as_ref().map(|r| { - let s = match r { - serde_json::Value::String(s) => s.clone(), - other => other.to_string(), - }; - truncate_preview(&s, 500) - }), - error: tc.error.clone(), + .map(|tc| { + // Image tools need full results (large base64 data), don't truncate + let limit = match tc.name.as_str() { + "image_generate" | "image_edit" | "image_analyze" => usize::MAX, + _ => 500, + }; + ToolCallInfo { + name: tc.name.clone(), + has_result: tc.result.is_some(), + has_error: tc.error.is_some(), + result_preview: tc.result.as_ref().map(|r| { + let s = match r { + serde_json::Value::String(s) => s.clone(), + other => other.to_string(), + }; + truncate_preview(&s, limit) + }), + error: tc.error.clone(), + } }) .collect(), }) diff --git a/src/channels/web/sse.rs b/src/channels/web/sse.rs index 809da048..55991fb2 100644 --- a/src/channels/web/sse.rs +++ b/src/channels/web/sse.rs @@ -55,6 +55,10 @@ impl SseManager { /// Broadcast an event to all connected clients. pub fn broadcast(&self, event: SseEvent) { + // Log image events for debugging + if matches!(&event, SseEvent::ImageGenerated { .. }) { + tracing::debug!("Broadcasting image_generated SSE event to all connected clients"); + } // Ignore send errors (no receivers is fine) let _ = self.tx.send(event); } diff --git a/src/channels/web/static/app.js b/src/channels/web/static/app.js index 1d1ff4ae..4b2e2a70 100644 --- a/src/channels/web/static/app.js +++ b/src/channels/web/static/app.js @@ -392,7 +392,12 @@ function connectSSE() { eventSource.addEventListener('image_generated', (e) => { const data = JSON.parse(e.data); - if (!isCurrentThread(data.thread_id)) return; + console.log('Received image_generated event:', { thread_id: data.thread_id, path: data.path, data_url_len: data.data_url ? data.data_url.length : 0 }); + if (!isCurrentThread(data.thread_id)) { + console.log('Image event ignored: not current thread', { currentThreadId, eventThreadId: data.thread_id }); + return; + } + console.log('Adding generated image to chat'); addGeneratedImage(data.data_url, data.path); }); @@ -950,6 +955,7 @@ function finalizeActivityGroup() { function addGeneratedImage(dataUrl, path) { const container = document.getElementById('chat-messages'); + console.log('addGeneratedImage called', { dataUrl_len: dataUrl ? dataUrl.length : 0, path }); const card = document.createElement('div'); card.className = 'generated-image-card'; @@ -957,6 +963,8 @@ function addGeneratedImage(dataUrl, path) { img.src = dataUrl; img.alt = 'Generated image'; img.className = 'generated-image'; + img.onerror = () => console.error('Failed to load image from data URL:', dataUrl.substring(0, 100)); + img.onload = () => console.log('Image loaded successfully from data URL'); const pathLabel = document.createElement('div'); pathLabel.className = 'generated-image-path'; @@ -965,6 +973,7 @@ function addGeneratedImage(dataUrl, path) { card.appendChild(img); card.appendChild(pathLabel); container.appendChild(card); + console.log('Image card appended to DOM'); container.scrollTop = container.scrollHeight; } @@ -1333,10 +1342,33 @@ function createToolCallsSummaryElement(toolCalls) { item.appendChild(nameSpan); if (tc.result_preview) { - const preview = document.createElement('div'); - preview.className = 'tool-call-preview'; - preview.textContent = tc.result_preview; - item.appendChild(preview); + // Check if this is an image result + try { + const parsed = JSON.parse(tc.result_preview); + if (parsed.type === 'image_generated' && parsed.data && parsed.media_type) { + const dataUrl = `data:${parsed.media_type};base64,${parsed.data}`; + const imgDiv = document.createElement('div'); + imgDiv.className = 'generated-image-card'; + const img = document.createElement('img'); + img.src = dataUrl; + img.alt = 'Generated image'; + img.className = 'generated-image'; + imgDiv.appendChild(img); + item.appendChild(imgDiv); + } else { + // Regular text result + const preview = document.createElement('div'); + preview.className = 'tool-call-preview'; + preview.textContent = tc.result_preview; + item.appendChild(preview); + } + } catch { + // Not JSON, display as text + const preview = document.createElement('div'); + preview.className = 'tool-call-preview'; + preview.textContent = tc.result_preview; + item.appendChild(preview); + } } if (tc.error) { const errDiv = document.createElement('div'); diff --git a/src/llm/rig_adapter.rs b/src/llm/rig_adapter.rs index f9b75efe..85a468cc 100644 --- a/src/llm/rig_adapter.rs +++ b/src/llm/rig_adapter.rs @@ -236,7 +236,7 @@ fn convert_messages(messages: &[ChatMessage]) -> (Option, Vec = vec![UserContent::text(&msg.content)]; for img in &msg.images { - let media_type = match img.media_type.as_str() { + let media_type = match img.media_type.to_lowercase().as_str() { "image/jpeg" => ImageMediaType::JPEG, "image/png" => ImageMediaType::PNG, "image/gif" => ImageMediaType::GIF, diff --git a/src/tools/builtin/image_analyze.rs b/src/tools/builtin/image_analyze.rs index 8e18bfe3..4c91aa36 100644 --- a/src/tools/builtin/image_analyze.rs +++ b/src/tools/builtin/image_analyze.rs @@ -25,13 +25,14 @@ impl ImageAnalyzeTool { /// Infer media type from file extension. fn infer_media_type(path: &str) -> &'static str { - if path.ends_with(".png") || path.ends_with(".b64") { + let lower_path = path.to_lowercase(); + if lower_path.ends_with(".png") || lower_path.ends_with(".b64") { "image/png" - } else if path.ends_with(".jpg") || path.ends_with(".jpeg") { + } else if lower_path.ends_with(".jpg") || lower_path.ends_with(".jpeg") { "image/jpeg" - } else if path.ends_with(".gif") { + } else if lower_path.ends_with(".gif") { "image/gif" - } else if path.ends_with(".webp") { + } else if lower_path.ends_with(".webp") { "image/webp" } else { "image/png" // Default to PNG @@ -202,14 +203,14 @@ mod tests { #[test] fn test_infer_media_type_uppercase_extension_defaults() { - // Uppercase extensions don't match the lowercase checks, so defaults to PNG + // Uppercase extensions are now case-insensitively matched assert_eq!( ImageAnalyzeTool::infer_media_type("images/test.PNG"), - "image/png" // Defaults to PNG for unknown extension + "image/png" ); assert_eq!( ImageAnalyzeTool::infer_media_type("images/test.JPG"), - "image/png" // Defaults to PNG for unknown extension + "image/jpeg" ); }