@probelabs/probe 0.6.0-rc101 → 0.6.0-rc103

package/cjs/agent/ProbeAgent.cjs

@@ -1840,7 +1840,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"],
@@ -1905,7 +1905,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) {
@@ -2268,13 +2268,19 @@ For GitHub-compatible mermaid diagrams, avoid single quotes and parentheses in n
 
 **Rules:**
 - NO single quotes in any node labels: 'text' \u2192 "text" or text
-- NO parentheses in square brackets: [Text (detail)] \u2192 [Text - detail]
+- NO parentheses in square brackets: [Text (detail)] \u2192 [Text - detail]  
 - NO complex expressions in diamonds: {a && b} \u2192 {condition}
+- NO HTML tags in node labels: [<pre>code</pre>] \u2192 ["code block"] or [Code Block]
 - USE single quotes for styles and classes: classDef highlight fill:'#ff9999'
+- CRITICAL: When using quotes in node labels, place them INSIDE the brackets: ["quoted text"], NOT [quoted text"]
 
 **Examples:**
 - \u2705 [Load Config] ["Run command"] {Valid?}
+- \u2705 ["depends_on: [generate-items]"] (correct quote placement)
+- \u2705 ["Code Block"] (clean text instead of HTML)
 - \u274C [Load (config)] [Run 'command'] {isValid('x')}
+- \u274C [depends_on: [generate-items"] (incorrect quote placement - quote ends inside bracket)
+- \u274C [<pre>depends_on: [generate-items]</pre>] (HTML tags in node labels)
 
 **Diagram Type Selection:**
 
@@ -2776,7 +2782,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: [],
@@ -4602,7 +4608,7 @@ __export(ProbeAgent_exports, {
   ProbeAgent: () => ProbeAgent
 });
 module.exports = __toCommonJS(ProbeAgent_exports);
-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_fs6, import_promises, import_path8, MAX_TOOL_ITERATIONS, MAX_HISTORY_MESSAGES, SUPPORTED_IMAGE_EXTENSIONS, MAX_IMAGE_FILE_SIZE, ProbeAgent;
 var init_ProbeAgent = __esm({
   "src/agent/ProbeAgent.js"() {
     import_anthropic = require("@ai-sdk/anthropic");
@@ -4611,6 +4617,9 @@ var init_ProbeAgent = __esm({
     import_ai2 = require("ai");
     import_crypto5 = require("crypto");
     import_events2 = require("events");
+    import_fs6 = require("fs");
+    import_promises = require("fs/promises");
+    import_path8 = require("path");
     init_tokenCounter();
     init_tools2();
     init_common();
@@ -4621,6 +4630,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
@@ -4664,6 +4675,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;
@@ -4789,6 +4802,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_path8.isAbsolute)(imagePath)) {
+            absolutePath = imagePath;
+            isPathAllowed = allowedDirs.some((dir) => absolutePath.startsWith((0, import_path8.resolve)(dir)));
+          } else {
+            for (const dir of allowedDirs) {
+              const resolvedPath = (0, import_path8.resolve)(dir, imagePath);
+              if (resolvedPath.startsWith((0, import_path8.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
        */
@@ -4888,7 +5070,7 @@ Examples:
 </extract>
 
 <attempt_completion>
-<result>The configuration is loaded from src/config.js lines 15-25 which contains the database settings.</result>
+The configuration is loaded from src/config.js lines 15-25 which contains the database settings.
 </attempt_completion>
 
 # Special Case: Quick Completion
@@ -4915,7 +5097,7 @@ I need to find code related to error handling in the search module. The most app
 6. Wait for the tool execution result, which will be provided in the next message (within a <tool_result> block).
 7. Analyze the tool result and decide the next step. If more tool calls are needed, repeat steps 2-6.
 8. If the task is fully complete and all previous steps were successful, use the \`<attempt_completion>\` tool to provide the final answer. This is the ONLY way to finish the task.
-9. If you cannot proceed (e.g., missing information, invalid request), explain the issue clearly before using \`<attempt_completion>\` with an appropriate message in the \`<result>\` tag.
+9. If you cannot proceed (e.g., missing information, invalid request), use \`<attempt_completion>\` to explain the issue clearly with an appropriate message directly inside the tags.
 10. If your previous response was already correct and complete, you may use \`<attempt_complete>\` as a shorthand.
 
 Available Tools:
@@ -5170,9 +5352,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
                 });
@@ -5206,6 +5389,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",
@@ -5330,12 +5516,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}`);
                     }