@probelabs/probe 0.6.0-rc101 → 0.6.0-rc103

package/build/agent/index.js

@@ -1821,7 +1821,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((resolve2, reject) => {
+    return new Promise((resolve3, reject) => {
       const delegationSpan = tracer ? tracer.createDelegationSpan(sessionId, task) : null;
       const process2 = spawn(binaryPath, args, {
         stdio: ["pipe", "pipe", "pipe"],
@@ -1886,7 +1886,7 @@ async function delegate({ task, timeout = 300, debug = false, currentIteration =
               delegationSpan.end();
             }
           }
-          resolve2(response);
+          resolve3(response);
         } else {
           const errorMessage = stderr.trim() || `Delegate process failed with exit code ${code}`;
           if (debug) {
@@ -2246,13 +2246,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:**
 
@@ -2510,20 +2516,20 @@ var init_simpleTelemetry = __esm({
       }
       async flush() {
         if (this.stream) {
-          return new Promise((resolve2) => {
-            this.stream.once("drain", resolve2);
+          return new Promise((resolve3) => {
+            this.stream.once("drain", resolve3);
             if (!this.stream.writableNeedDrain) {
-              resolve2();
+              resolve3();
             }
           });
         }
       }
       async shutdown() {
         if (this.stream) {
-          return new Promise((resolve2) => {
+          return new Promise((resolve3) => {
             this.stream.end(() => {
               console.log(`[SimpleTelemetry] File stream closed: ${this.filePath}`);
-              resolve2();
+              resolve3();
             });
           });
         }
@@ -2969,7 +2975,7 @@ function createMockProvider() {
       provider: "mock",
       // Mock the doGenerate method used by Vercel AI SDK
       doGenerate: async ({ messages, tools: tools2 }) => {
-        await new Promise((resolve2) => setTimeout(resolve2, 10));
+        await new Promise((resolve3) => setTimeout(resolve3, 10));
         return {
           text: "This is a mock response for testing",
           toolCalls: [],
@@ -4859,7 +4865,10 @@ import { createGoogleGenerativeAI } from "@ai-sdk/google";
 import { streamText } from "ai";
 import { randomUUID as randomUUID4 } from "crypto";
 import { EventEmitter as EventEmitter2 } from "events";
-var MAX_TOOL_ITERATIONS, MAX_HISTORY_MESSAGES, ProbeAgent;
+import { existsSync as existsSync3 } from "fs";
+import { readFile, stat } from "fs/promises";
+import { resolve, isAbsolute } from "path";
+var MAX_TOOL_ITERATIONS, MAX_HISTORY_MESSAGES, SUPPORTED_IMAGE_EXTENSIONS, MAX_IMAGE_FILE_SIZE, ProbeAgent;
 var init_ProbeAgent = __esm({
   "src/agent/ProbeAgent.js"() {
     "use strict";
@@ -4873,6 +4882,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
@@ -4916,6 +4927,8 @@ var init_ProbeAgent = __esm({
         }
         this.initializeTools();
         this.history = [];
+        this.pendingImages = /* @__PURE__ */ new Map();
+        this.currentImages = [];
         this.events = new EventEmitter2();
         this.enableMcp = !!options.enableMcp || process.env.ENABLE_MCP === "1";
         this.mcpConfigPath = options.mcpConfigPath || null;
@@ -5041,6 +5054,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 (isAbsolute(imagePath)) {
+            absolutePath = imagePath;
+            isPathAllowed = allowedDirs.some((dir) => absolutePath.startsWith(resolve(dir)));
+          } else {
+            for (const dir of allowedDirs) {
+              const resolvedPath = resolve(dir, imagePath);
+              if (resolvedPath.startsWith(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 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 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
        */
@@ -5140,7 +5322,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
@@ -5167,7 +5349,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:
@@ -5422,9 +5604,10 @@ You are working with a repository located at: ${searchDirectory}
             let assistantResponseContent = "";
             try {
               const executeAIRequest = async () => {
+                const messagesForAI = this.prepareMessagesWithImages(currentMessages);
                 const result = await streamText({
                   model: this.provider(this.model),
-                  messages: currentMessages,
+                  messages: messagesForAI,
                   maxTokens: maxResponseTokens,
                   temperature: 0.3
                 });
@@ -5458,6 +5641,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",
@@ -5582,12 +5768,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}`);
                     }
@@ -6084,11 +6275,12 @@ import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"
 import {
   CallToolRequestSchema,
   ErrorCode as ErrorCode2,
+  InitializeRequestSchema,
   ListToolsRequestSchema,
   McpError
 } from "@modelcontextprotocol/sdk/types.js";
-import { readFileSync as readFileSync2, existsSync as existsSync3 } from "fs";
-import { resolve } from "path";
+import { readFileSync as readFileSync2, existsSync as existsSync4 } from "fs";
+import { resolve as resolve2 } from "path";
 
 // src/agent/acp/server.js
 import { randomUUID as randomUUID5 } from "crypto";
@@ -6347,8 +6539,8 @@ var ACPConnection = class extends EventEmitter3 {
     if (params !== null) {
       message.params = params;
     }
-    return new Promise((resolve2, reject) => {
-      this.pendingRequests.set(id, { resolve: resolve2, reject });
+    return new Promise((resolve3, reject) => {
+      this.pendingRequests.set(id, { resolve: resolve3, reject });
       this.sendMessage(message);
       setTimeout(() => {
         if (this.pendingRequests.has(id)) {
@@ -6756,8 +6948,8 @@ import { randomUUID as randomUUID6 } from "crypto";
 function readInputContent(input) {
   if (!input) return null;
   try {
-    const resolvedPath = resolve(input);
-    if (existsSync3(resolvedPath)) {
+    const resolvedPath = resolve2(input);
+    if (existsSync4(resolvedPath)) {
       return readFileSync2(resolvedPath, "utf-8").trim();
     }
   } catch (error) {
@@ -6765,7 +6957,7 @@ function readInputContent(input) {
   return input;
 }
 function readFromStdin() {
-  return new Promise((resolve2, reject) => {
+  return new Promise((resolve3, reject) => {
     let data = "";
     let hasReceivedData = false;
     let dataChunks = [];
@@ -6790,7 +6982,7 @@ function readFromStdin() {
       if (!trimmed && dataChunks.length === 0) {
         reject(new Error("No input received from stdin"));
       } else {
-        resolve2(trimmed);
+        resolve3(trimmed);
       }
     });
     process.stdin.on("error", (error) => {
@@ -6948,7 +7140,7 @@ var ProbeAgentMcpServer = class {
   constructor() {
     this.server = new Server(
       {
-        name: "@buger/probe-agent",
+        name: "@probelabs/probe agent",
         version: "1.0.0"
       },
       {
@@ -6957,6 +7149,7 @@ var ProbeAgentMcpServer = class {
         }
       }
     );
+    this.agent = null;
     this.setupToolHandlers();
     this.server.onerror = (error) => console.error("[MCP Error]", error);
     process.on("SIGINT", async () => {
@@ -6965,6 +7158,18 @@ var ProbeAgentMcpServer = class {
     });
   }
   setupToolHandlers() {
+    this.server.setRequestHandler(InitializeRequestSchema, async (request) => {
+      return {
+        protocolVersion: "2024-11-05",
+        capabilities: {
+          tools: {}
+        },
+        serverInfo: {
+          name: "@probelabs/probe agent",
+          version: "1.0.0"
+        }
+      };
+    });
     this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
       tools: [
         {
@@ -7056,18 +7261,24 @@ var ProbeAgentMcpServer = class {
             throw new Error("Schema could not be read");
           }
         }
-        const agentConfig = {
-          path: args.path || process.cwd(),
-          promptType: args.prompt || "code-explorer",
-          customPrompt: systemPrompt,
-          provider: args.provider,
-          model: args.model,
-          allowEdit: !!args.allow_edit,
-          debug: process.env.DEBUG === "1",
-          maxResponseTokens: args.max_response_tokens,
-          disableMermaidValidation: !!args.no_mermaid_validation
-        };
-        const agent = new ProbeAgent(agentConfig);
+        if (!this.agent) {
+          if (process.env.DEBUG === "1") {
+            console.error("[DEBUG] Initializing AI agent on first MCP tool call");
+          }
+          const agentConfig2 = {
+            path: args.path || process.cwd(),
+            promptType: args.prompt || "code-explorer",
+            customPrompt: systemPrompt,
+            provider: args.provider,
+            model: args.model,
+            allowEdit: !!args.allow_edit,
+            debug: process.env.DEBUG === "1",
+            maxResponseTokens: args.max_response_tokens,
+            disableMermaidValidation: !!args.no_mermaid_validation
+          };
+          this.agent = new ProbeAgent(agentConfig2);
+        }
+        const agent = this.agent;
         let result = await agent.answer(query2, [], { schema });
         if (schema) {
           const schemaPrompt = `Now you need to respond according to this schema:
@@ -7128,7 +7339,7 @@ Please reformat your previous response to match this schema exactly. Only return
           } catch (error) {
           }
         }
-        const tokenUsage = agent.getTokenUsage();
+        const tokenUsage = this.agent.getTokenUsage();
         console.error(`Token usage: ${JSON.stringify(tokenUsage)}`);
         return {
           content: [
@@ -7253,7 +7464,7 @@ async function main() {
         process.exit(1);
       }
     }
-    const agentConfig = {
+    const agentConfig2 = {
       path: config.path,
       promptType: config.prompt,
       customPrompt: systemPrompt,
@@ -7264,7 +7475,7 @@ async function main() {
       maxResponseTokens: config.maxResponseTokens,
       disableMermaidValidation: config.noMermaidValidation
     };
-    const agent = new ProbeAgent(agentConfig);
+    const agent = new ProbeAgent(agentConfig2);
     let result;
     if (appTracer) {
       const sessionSpan = appTracer.createSessionSpan({
@@ -7386,7 +7597,12 @@ Please reformat your previous response to match this schema exactly. Only return
         }
       }
     }
-    console.log(result);
+    const resultMatch = result.match(/<result>([\s\S]*?)<\/result>/);
+    if (resultMatch) {
+      console.log(resultMatch[1].trim());
+    } else {
+      console.log(result);
+    }
     if (config.verbose) {
       const tokenUsage = agent.getTokenUsage();
       console.error(`