@probelabs/probe 0.6.0-rc102 → 0.6.0-rc103

package/cjs/index.cjs

@@ -1543,7 +1543,7 @@ async function delegate({ task, timeout = 300, debug = false, currentIteration =
       console.error(`[DELEGATE] Using binary at: ${binaryPath}`);
       console.error(`[DELEGATE] Command args: ${args.join(" ")}`);
     }
-    return new Promise((resolve, reject) => {
+    return new Promise((resolve2, reject) => {
       const delegationSpan = tracer ? tracer.createDelegationSpan(sessionId, task) : null;
       const process2 = (0, import_child_process5.spawn)(binaryPath, args, {
         stdio: ["pipe", "pipe", "pipe"],
@@ -1608,7 +1608,7 @@ async function delegate({ task, timeout = 300, debug = false, currentIteration =
               delegationSpan.end();
             }
           }
-          resolve(response);
+          resolve2(response);
         } else {
           const errorMessage = stderr.trim() || `Delegate process failed with exit code ${code}`;
           if (debug) {
@@ -2868,7 +2868,7 @@ function createMockProvider() {
       provider: "mock",
       // Mock the doGenerate method used by Vercel AI SDK
       doGenerate: async ({ messages, tools: tools2 }) => {
-        await new Promise((resolve) => setTimeout(resolve, 10));
+        await new Promise((resolve2) => setTimeout(resolve2, 10));
         return {
           text: "This is a mock response for testing",
           toolCalls: [],
@@ -4693,7 +4693,7 @@ var ProbeAgent_exports = {};
 __export(ProbeAgent_exports, {
   ProbeAgent: () => ProbeAgent
 });
-var import_anthropic, import_openai, import_google, import_ai2, import_crypto5, import_events2, MAX_TOOL_ITERATIONS, MAX_HISTORY_MESSAGES, ProbeAgent;
+var import_anthropic, import_openai, import_google, import_ai2, import_crypto5, import_events2, import_fs5, import_promises, import_path7, MAX_TOOL_ITERATIONS, MAX_HISTORY_MESSAGES, SUPPORTED_IMAGE_EXTENSIONS, MAX_IMAGE_FILE_SIZE, ProbeAgent;
 var init_ProbeAgent = __esm({
   "src/agent/ProbeAgent.js"() {
     "use strict";
@@ -4703,6 +4703,9 @@ var init_ProbeAgent = __esm({
     import_ai2 = require("ai");
     import_crypto5 = require("crypto");
     import_events2 = require("events");
+    import_fs5 = require("fs");
+    import_promises = require("fs/promises");
+    import_path7 = require("path");
     init_tokenCounter();
     init_tools2();
     init_common();
@@ -4713,6 +4716,8 @@ var init_ProbeAgent = __esm({
     init_mcp();
     MAX_TOOL_ITERATIONS = parseInt(process.env.MAX_TOOL_ITERATIONS || "30", 10);
     MAX_HISTORY_MESSAGES = 100;
+    SUPPORTED_IMAGE_EXTENSIONS = ["png", "jpg", "jpeg", "webp", "gif", "bmp", "svg"];
+    MAX_IMAGE_FILE_SIZE = 20 * 1024 * 1024;
     ProbeAgent = class {
       /**
        * Create a new ProbeAgent instance
@@ -4756,6 +4761,8 @@ var init_ProbeAgent = __esm({
         }
         this.initializeTools();
         this.history = [];
+        this.pendingImages = /* @__PURE__ */ new Map();
+        this.currentImages = [];
         this.events = new import_events2.EventEmitter();
         this.enableMcp = !!options.enableMcp || process.env.ENABLE_MCP === "1";
         this.mcpConfigPath = options.mcpConfigPath || null;
@@ -4881,6 +4888,175 @@ var init_ProbeAgent = __esm({
           console.log(`Using Google API with model: ${this.model}${apiUrl ? ` (URL: ${apiUrl})` : ""}`);
         }
       }
+      /**
+       * Process assistant response content and detect/load image references
+       * @param {string} content - The assistant's response content
+       * @returns {Promise<void>}
+       */
+      async processImageReferences(content) {
+        if (!content) return;
+        const extensionsPattern = `(?:${SUPPORTED_IMAGE_EXTENSIONS.join("|")})`;
+        const imagePatterns = [
+          // Direct file path mentions: "./screenshot.png", "/path/to/image.jpg", etc.
+          new RegExp(`(?:\\.?\\.\\/)?[^\\s"'<>\\[\\]]+\\.${extensionsPattern}(?!\\w)`, "gi"),
+          // Contextual mentions: "look at image.png", "the file screenshot.jpg shows"
+          new RegExp(`(?:image|file|screenshot|diagram|photo|picture|graphic)\\s*:?\\s*([^\\s"'<>\\[\\]]+\\.${extensionsPattern})(?!\\w)`, "gi"),
+          // Tool result mentions: often contain file paths
+          new RegExp(`(?:found|saved|created|generated).*?([^\\s"'<>\\[\\]]+\\.${extensionsPattern})(?!\\w)`, "gi")
+        ];
+        const foundPaths = /* @__PURE__ */ new Set();
+        for (const pattern of imagePatterns) {
+          let match;
+          while ((match = pattern.exec(content)) !== null) {
+            const imagePath = match[1] || match[0];
+            if (imagePath && imagePath.length > 0) {
+              foundPaths.add(imagePath.trim());
+            }
+          }
+        }
+        if (foundPaths.size === 0) return;
+        if (this.debug) {
+          console.log(`[DEBUG] Found ${foundPaths.size} potential image references:`, Array.from(foundPaths));
+        }
+        for (const imagePath of foundPaths) {
+          await this.loadImageIfValid(imagePath);
+        }
+      }
+      /**
+       * Load and cache an image if it's valid and accessible
+       * @param {string} imagePath - Path to the image file
+       * @returns {Promise<boolean>} - True if image was loaded successfully
+       */
+      async loadImageIfValid(imagePath) {
+        try {
+          if (this.pendingImages.has(imagePath)) {
+            if (this.debug) {
+              console.log(`[DEBUG] Image already loaded: ${imagePath}`);
+            }
+            return true;
+          }
+          const allowedDirs = this.allowedFolders && this.allowedFolders.length > 0 ? this.allowedFolders : [process.cwd()];
+          let absolutePath;
+          let isPathAllowed = false;
+          if ((0, import_path7.isAbsolute)(imagePath)) {
+            absolutePath = imagePath;
+            isPathAllowed = allowedDirs.some((dir) => absolutePath.startsWith((0, import_path7.resolve)(dir)));
+          } else {
+            for (const dir of allowedDirs) {
+              const resolvedPath = (0, import_path7.resolve)(dir, imagePath);
+              if (resolvedPath.startsWith((0, import_path7.resolve)(dir))) {
+                absolutePath = resolvedPath;
+                isPathAllowed = true;
+                break;
+              }
+            }
+          }
+          if (!isPathAllowed) {
+            if (this.debug) {
+              console.log(`[DEBUG] Image path outside allowed directories: ${imagePath}`);
+            }
+            return false;
+          }
+          let fileStats;
+          try {
+            fileStats = await (0, import_promises.stat)(absolutePath);
+          } catch (error) {
+            if (this.debug) {
+              console.log(`[DEBUG] Image file not found: ${absolutePath}`);
+            }
+            return false;
+          }
+          if (fileStats.size > MAX_IMAGE_FILE_SIZE) {
+            if (this.debug) {
+              console.log(`[DEBUG] Image file too large: ${absolutePath} (${fileStats.size} bytes, max: ${MAX_IMAGE_FILE_SIZE})`);
+            }
+            return false;
+          }
+          const extension = absolutePath.toLowerCase().split(".").pop();
+          if (!SUPPORTED_IMAGE_EXTENSIONS.includes(extension)) {
+            if (this.debug) {
+              console.log(`[DEBUG] Unsupported image format: ${extension}`);
+            }
+            return false;
+          }
+          const mimeTypes = {
+            "png": "image/png",
+            "jpg": "image/jpeg",
+            "jpeg": "image/jpeg",
+            "webp": "image/webp",
+            "gif": "image/gif",
+            "bmp": "image/bmp",
+            "svg": "image/svg+xml"
+          };
+          const mimeType = mimeTypes[extension];
+          const fileBuffer = await (0, import_promises.readFile)(absolutePath);
+          const base64Data = fileBuffer.toString("base64");
+          const dataUrl = `data:${mimeType};base64,${base64Data}`;
+          this.pendingImages.set(imagePath, dataUrl);
+          if (this.debug) {
+            console.log(`[DEBUG] Successfully loaded image: ${imagePath} (${fileBuffer.length} bytes)`);
+          }
+          return true;
+        } catch (error) {
+          if (this.debug) {
+            console.log(`[DEBUG] Failed to load image ${imagePath}: ${error.message}`);
+          }
+          return false;
+        }
+      }
+      /**
+       * Get all currently loaded images as an array for AI model consumption
+       * @returns {Array<string>} - Array of base64 data URLs
+       */
+      getCurrentImages() {
+        return Array.from(this.pendingImages.values());
+      }
+      /**
+       * Clear loaded images (useful for new conversations)
+       */
+      clearLoadedImages() {
+        this.pendingImages.clear();
+        this.currentImages = [];
+        if (this.debug) {
+          console.log("[DEBUG] Cleared all loaded images");
+        }
+      }
+      /**
+       * Prepare messages for AI consumption, adding images to the latest user message if available
+       * @param {Array} messages - Current conversation messages
+       * @returns {Array} - Messages formatted for AI SDK with potential image content
+       */
+      prepareMessagesWithImages(messages) {
+        const loadedImages = this.getCurrentImages();
+        if (loadedImages.length === 0) {
+          return messages;
+        }
+        const messagesWithImages = [...messages];
+        const lastUserMessageIndex = messagesWithImages.map((m) => m.role).lastIndexOf("user");
+        if (lastUserMessageIndex === -1) {
+          if (this.debug) {
+            console.log("[DEBUG] No user messages found to attach images to");
+          }
+          return messages;
+        }
+        const lastUserMessage = messagesWithImages[lastUserMessageIndex];
+        if (typeof lastUserMessage.content === "string") {
+          messagesWithImages[lastUserMessageIndex] = {
+            ...lastUserMessage,
+            content: [
+              { type: "text", text: lastUserMessage.content },
+              ...loadedImages.map((imageData) => ({
+                type: "image",
+                image: imageData
+              }))
+            ]
+          };
+          if (this.debug) {
+            console.log(`[DEBUG] Added ${loadedImages.length} images to the latest user message`);
+          }
+        }
+        return messagesWithImages;
+      }
       /**
        * Initialize mock model for testing
        */
@@ -5262,9 +5438,10 @@ You are working with a repository located at: ${searchDirectory}
             let assistantResponseContent = "";
             try {
               const executeAIRequest = async () => {
+                const messagesForAI = this.prepareMessagesWithImages(currentMessages);
                 const result = await (0, import_ai2.streamText)({
                   model: this.provider(this.model),
-                  messages: currentMessages,
+                  messages: messagesForAI,
                   maxTokens: maxResponseTokens,
                   temperature: 0.3
                 });
@@ -5298,6 +5475,9 @@ You are working with a repository located at: ${searchDirectory}
               const assistantPreview = createMessagePreview(assistantResponseContent);
               console.log(`[DEBUG] Assistant response (${assistantResponseContent.length} chars): ${assistantPreview}`);
             }
+            if (assistantResponseContent) {
+              await this.processImageReferences(assistantResponseContent);
+            }
             const validTools = [
               "search",
               "query",
@@ -5422,12 +5602,17 @@ ${toolResultContent}
                       throw toolError;
                     }
                     currentMessages.push({ role: "assistant", content: assistantResponseContent });
+                    const toolResultContent = typeof toolResult === "string" ? toolResult : JSON.stringify(toolResult, null, 2);
+                    const toolResultMessage = `<tool_result>
+${toolResultContent}
+</tool_result>`;
                     currentMessages.push({
                       role: "user",
-                      content: `<tool_result>
-${typeof toolResult === "string" ? toolResult : JSON.stringify(toolResult, null, 2)}
-</tool_result>`
+                      content: toolResultMessage
                     });
+                    if (toolResultContent) {
+                      await this.processImageReferences(toolResultContent);
+                    }
                     if (this.debug) {
                       console.log(`[DEBUG] Tool ${toolName} executed successfully. Result length: ${typeof toolResult === "string" ? toolResult.length : JSON.stringify(toolResult).length}`);
                     }
@@ -5925,12 +6110,12 @@ function initializeSimpleTelemetryFromOptions(options) {
   });
   return telemetry;
 }
-var import_fs5, import_path7, SimpleTelemetry, SimpleAppTracer;
+var import_fs6, import_path8, SimpleTelemetry, SimpleAppTracer;
 var init_simpleTelemetry = __esm({
   "src/agent/simpleTelemetry.js"() {
     "use strict";
-    import_fs5 = require("fs");
-    import_path7 = require("path");
+    import_fs6 = require("fs");
+    import_path8 = require("path");
     SimpleTelemetry = class {
       constructor(options = {}) {
         this.serviceName = options.serviceName || "probe-agent";
@@ -5944,11 +6129,11 @@ var init_simpleTelemetry = __esm({
       }
       initializeFileExporter() {
         try {
-          const dir = (0, import_path7.dirname)(this.filePath);
-          if (!(0, import_fs5.existsSync)(dir)) {
-            (0, import_fs5.mkdirSync)(dir, { recursive: true });
+          const dir = (0, import_path8.dirname)(this.filePath);
+          if (!(0, import_fs6.existsSync)(dir)) {
+            (0, import_fs6.mkdirSync)(dir, { recursive: true });
           }
-          this.stream = (0, import_fs5.createWriteStream)(this.filePath, { flags: "a" });
+          this.stream = (0, import_fs6.createWriteStream)(this.filePath, { flags: "a" });
           this.stream.on("error", (error) => {
             console.error(`[SimpleTelemetry] Stream error: ${error.message}`);
           });
@@ -6009,20 +6194,20 @@ var init_simpleTelemetry = __esm({
       }
       async flush() {
         if (this.stream) {
-          return new Promise((resolve) => {
-            this.stream.once("drain", resolve);
+          return new Promise((resolve2) => {
+            this.stream.once("drain", resolve2);
             if (!this.stream.writableNeedDrain) {
-              resolve();
+              resolve2();
             }
           });
         }
       }
       async shutdown() {
         if (this.stream) {
-          return new Promise((resolve) => {
+          return new Promise((resolve2) => {
             this.stream.end(() => {
               console.log(`[SimpleTelemetry] File stream closed: ${this.filePath}`);
-              resolve();
+              resolve2();
             });
           });
         }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@probelabs/probe",
-  "version": "0.6.0-rc102",
+  "version": "0.6.0-rc103",
   "description": "Node.js wrapper for the probe code search tool",
   "main": "src/index.js",
   "module": "src/index.js",

package/src/agent/ProbeAgent.js

@@ -5,6 +5,9 @@ import { createGoogleGenerativeAI } from '@ai-sdk/google';
 import { streamText } from 'ai';
 import { randomUUID } from 'crypto';
 import { EventEmitter } from 'events';
+import { existsSync } from 'fs';
+import { readFile, stat } from 'fs/promises';
+import { resolve, isAbsolute } from 'path';
 import { TokenCounter } from './tokenCounter.js';
 import { 
   createTools,
@@ -46,6 +49,12 @@ import {
 const MAX_TOOL_ITERATIONS = parseInt(process.env.MAX_TOOL_ITERATIONS || '30', 10);
 const MAX_HISTORY_MESSAGES = 100;
 
+// Supported image file extensions
+const SUPPORTED_IMAGE_EXTENSIONS = ['png', 'jpg', 'jpeg', 'webp', 'gif', 'bmp', 'svg'];
+
+// Maximum image file size (20MB) to prevent OOM attacks
+const MAX_IMAGE_FILE_SIZE = 20 * 1024 * 1024;
+
 /**
  * ProbeAgent class to handle AI interactions with code search capabilities
  */
@@ -105,6 +114,10 @@ export class ProbeAgent {
     // Initialize chat history
     this.history = [];
 
+    // Initialize image tracking for agentic loop
+    this.pendingImages = new Map(); // Map<imagePath, base64Data> to avoid reloading
+    this.currentImages = []; // Currently active images for AI calls
+
     // Initialize event emitter for tool execution updates
     this.events = new EventEmitter();
 
@@ -268,6 +281,226 @@ export class ProbeAgent {
     }
   }
 
+  /**
+   * Process assistant response content and detect/load image references
+   * @param {string} content - The assistant's response content
+   * @returns {Promise<void>}
+   */
+  async processImageReferences(content) {
+    if (!content) return;
+
+    // Enhanced pattern to detect image file mentions in various contexts
+    // Looks for: "image", "file", "screenshot", etc. followed by path-like strings with image extensions
+    const extensionsPattern = `(?:${SUPPORTED_IMAGE_EXTENSIONS.join('|')})`;
+    const imagePatterns = [
+      // Direct file path mentions: "./screenshot.png", "/path/to/image.jpg", etc.
+      new RegExp(`(?:\\.?\\.\\/)?[^\\s"'<>\\[\\]]+\\\.${extensionsPattern}(?!\\w)`, 'gi'),
+      // Contextual mentions: "look at image.png", "the file screenshot.jpg shows"
+      new RegExp(`(?:image|file|screenshot|diagram|photo|picture|graphic)\\s*:?\\s*([^\\s"'<>\\[\\]]+\\.${extensionsPattern})(?!\\w)`, 'gi'),
+      // Tool result mentions: often contain file paths
+      new RegExp(`(?:found|saved|created|generated).*?([^\\s"'<>\\[\\]]+\\.${extensionsPattern})(?!\\w)`, 'gi')
+    ];
+
+    const foundPaths = new Set();
+
+    // Extract potential image paths using all patterns
+    for (const pattern of imagePatterns) {
+      let match;
+      while ((match = pattern.exec(content)) !== null) {
+        // For patterns with capture groups, use the captured path; otherwise use the full match
+        const imagePath = match[1] || match[0];
+        if (imagePath && imagePath.length > 0) {
+          foundPaths.add(imagePath.trim());
+        }
+      }
+    }
+
+    if (foundPaths.size === 0) return;
+
+    if (this.debug) {
+      console.log(`[DEBUG] Found ${foundPaths.size} potential image references:`, Array.from(foundPaths));
+    }
+
+    // Process each found path
+    for (const imagePath of foundPaths) {
+      await this.loadImageIfValid(imagePath);
+    }
+  }
+
+  /**
+   * Load and cache an image if it's valid and accessible
+   * @param {string} imagePath - Path to the image file
+   * @returns {Promise<boolean>} - True if image was loaded successfully
+   */
+  async loadImageIfValid(imagePath) {
+    try {
+      // Skip if already loaded
+      if (this.pendingImages.has(imagePath)) {
+        if (this.debug) {
+          console.log(`[DEBUG] Image already loaded: ${imagePath}`);
+        }
+        return true;
+      }
+
+      // Security validation: check if path is within any allowed directory
+      const allowedDirs = this.allowedFolders && this.allowedFolders.length > 0 ? this.allowedFolders : [process.cwd()];
+      
+      let absolutePath;
+      let isPathAllowed = false;
+      
+      // If absolute path, check if it's within any allowed directory
+      if (isAbsolute(imagePath)) {
+        absolutePath = imagePath;
+        isPathAllowed = allowedDirs.some(dir => absolutePath.startsWith(resolve(dir)));
+      } else {
+        // For relative paths, try resolving against each allowed directory
+        for (const dir of allowedDirs) {
+          const resolvedPath = resolve(dir, imagePath);
+          if (resolvedPath.startsWith(resolve(dir))) {
+            absolutePath = resolvedPath;
+            isPathAllowed = true;
+            break;
+          }
+        }
+      }
+      
+      // Security check: ensure path is within at least one allowed directory
+      if (!isPathAllowed) {
+        if (this.debug) {
+          console.log(`[DEBUG] Image path outside allowed directories: ${imagePath}`);
+        }
+        return false;
+      }
+
+      // Check if file exists and get file stats
+      let fileStats;
+      try {
+        fileStats = await stat(absolutePath);
+      } catch (error) {
+        if (this.debug) {
+          console.log(`[DEBUG] Image file not found: ${absolutePath}`);
+        }
+        return false;
+      }
+
+      // Validate file size to prevent OOM attacks
+      if (fileStats.size > MAX_IMAGE_FILE_SIZE) {
+        if (this.debug) {
+          console.log(`[DEBUG] Image file too large: ${absolutePath} (${fileStats.size} bytes, max: ${MAX_IMAGE_FILE_SIZE})`);
+        }
+        return false;
+      }
+
+      // Validate file extension
+      const extension = absolutePath.toLowerCase().split('.').pop();
+      if (!SUPPORTED_IMAGE_EXTENSIONS.includes(extension)) {
+        if (this.debug) {
+          console.log(`[DEBUG] Unsupported image format: ${extension}`);
+        }
+        return false;
+      }
+
+      // Determine MIME type
+      const mimeTypes = {
+        'png': 'image/png',
+        'jpg': 'image/jpeg',
+        'jpeg': 'image/jpeg',
+        'webp': 'image/webp',
+        'gif': 'image/gif',
+        'bmp': 'image/bmp',
+        'svg': 'image/svg+xml'
+      };
+      const mimeType = mimeTypes[extension];
+
+      // Read and encode file asynchronously
+      const fileBuffer = await readFile(absolutePath);
+      const base64Data = fileBuffer.toString('base64');
+      const dataUrl = `data:${mimeType};base64,${base64Data}`;
+
+      // Cache the loaded image
+      this.pendingImages.set(imagePath, dataUrl);
+
+      if (this.debug) {
+        console.log(`[DEBUG] Successfully loaded image: ${imagePath} (${fileBuffer.length} bytes)`);
+      }
+
+      return true;
+    } catch (error) {
+      if (this.debug) {
+        console.log(`[DEBUG] Failed to load image ${imagePath}: ${error.message}`);
+      }
+      return false;
+    }
+  }
+
+  /**
+   * Get all currently loaded images as an array for AI model consumption
+   * @returns {Array<string>} - Array of base64 data URLs
+   */
+  getCurrentImages() {
+    return Array.from(this.pendingImages.values());
+  }
+
+  /**
+   * Clear loaded images (useful for new conversations)
+   */
+  clearLoadedImages() {
+    this.pendingImages.clear();
+    this.currentImages = [];
+    if (this.debug) {
+      console.log('[DEBUG] Cleared all loaded images');
+    }
+  }
+
+  /**
+   * Prepare messages for AI consumption, adding images to the latest user message if available
+   * @param {Array} messages - Current conversation messages
+   * @returns {Array} - Messages formatted for AI SDK with potential image content
+   */
+  prepareMessagesWithImages(messages) {
+    const loadedImages = this.getCurrentImages();
+    
+    // If no images loaded, return messages as-is
+    if (loadedImages.length === 0) {
+      return messages;
+    }
+
+    // Clone messages to avoid mutating the original
+    const messagesWithImages = [...messages];
+    
+    // Find the last user message to attach images to
+    const lastUserMessageIndex = messagesWithImages.map(m => m.role).lastIndexOf('user');
+    
+    if (lastUserMessageIndex === -1) {
+      if (this.debug) {
+        console.log('[DEBUG] No user messages found to attach images to');
+      }
+      return messages;
+    }
+
+    const lastUserMessage = messagesWithImages[lastUserMessageIndex];
+    
+    // Convert to multimodal format if we have images
+    if (typeof lastUserMessage.content === 'string') {
+      messagesWithImages[lastUserMessageIndex] = {
+        ...lastUserMessage,
+        content: [
+          { type: 'text', text: lastUserMessage.content },
+          ...loadedImages.map(imageData => ({
+            type: 'image',
+            image: imageData
+          }))
+        ]
+      };
+
+      if (this.debug) {
+        console.log(`[DEBUG] Added ${loadedImages.length} images to the latest user message`);
+      }
+    }
+
+    return messagesWithImages;
+  }
+
   /**
    * Initialize mock model for testing
    */
@@ -695,9 +928,12 @@ When troubleshooting:
         try {
           // Wrap AI request with tracing if available
           const executeAIRequest = async () => {
+            // Prepare messages with potential image content
+            const messagesForAI = this.prepareMessagesWithImages(currentMessages);
+            
             const result = await streamText({
               model: this.provider(this.model),
-              messages: currentMessages,
+              messages: messagesForAI,
               maxTokens: maxResponseTokens,
               temperature: 0.3,
             });
@@ -741,6 +977,11 @@ When troubleshooting:
           console.log(`[DEBUG] Assistant response (${assistantResponseContent.length} chars): ${assistantPreview}`);
         }
 
+        // Process image references in assistant response for next iteration
+        if (assistantResponseContent) {
+          await this.processImageReferences(assistantResponseContent);
+        }
+
         // Parse tool call from response with valid tools list
         const validTools = [
           'search', 'query', 'extract', 'listFiles', 'searchFiles', 'attempt_completion'
@@ -898,11 +1139,20 @@ When troubleshooting:
                 
                 // Add assistant response and tool result to conversation
                 currentMessages.push({ role: 'assistant', content: assistantResponseContent });
+                
+                const toolResultContent = typeof toolResult === 'string' ? toolResult : JSON.stringify(toolResult, null, 2);
+                const toolResultMessage = `<tool_result>\n${toolResultContent}\n</tool_result>`;
+                
                 currentMessages.push({
                   role: 'user',
-                  content: `<tool_result>\n${typeof toolResult === 'string' ? toolResult : JSON.stringify(toolResult, null, 2)}\n</tool_result>`
+                  content: toolResultMessage
                 });
 
+                // Process tool result for image references
+                if (toolResultContent) {
+                  await this.processImageReferences(toolResultContent);
+                }
+
                 if (this.debug) {
                   console.log(`[DEBUG] Tool ${toolName} executed successfully. Result length: ${typeof toolResult === 'string' ? toolResult.length : JSON.stringify(toolResult).length}`);
                 }