@probelabs/probe 0.6.0-rc102 → 0.6.0-rc103

package/build/mcp/index.ts

@@ -112,37 +112,19 @@ console.log(`Bin directory: ${binDir}`);
 interface SearchCodeArgs {
   path: string;
   query: string | string[];
-  filesOnly?: boolean;
-  ignore?: string[];
-  excludeFilenames?: boolean;
   exact?: boolean;
   maxResults?: number;
   maxTokens?: number;
   allowTests?: boolean;
   session?: string;
-  timeout?: number;
   noGitignore?: boolean;
 }
 
-interface QueryCodeArgs {
-  path: string;
-  pattern: string;
-  language?: string;
-  ignore?: string[];
-  allowTests?: boolean;
-  maxResults?: number;
-  format?: 'markdown' | 'plain' | 'json' | 'color';
-  timeout?: number;
-  noGitignore?: boolean;
-}
 
 interface ExtractCodeArgs {
   path: string;
   files: string[];
   allowTests?: boolean;
-  contextLines?: number;
-  format?: 'markdown' | 'plain' | 'json';
-  timeout?: number;
   noGitignore?: boolean;
 }
 
@@ -193,38 +175,22 @@ class ProbeServer {
               },
               query: {
                 type: 'string',
-                description: 'Elastic search query. Supports logical operators (AND, OR, NOT), and grouping with parentheses. Examples: "config", "(term1 OR term2) AND term3". Use quotes for exact matches, like function or type names.',
-              },
-              filesOnly: {
-                type: 'boolean',
-                description: 'Skip AST parsing and just output unique files',
-              },
-              ignore: {
-                type: 'array',
-                items: { type: 'string' },
-                description: 'Custom patterns to ignore (in addition to .gitignore and common patterns)'
-              },
-              excludeFilenames: {
-                type: 'boolean',
-                description: 'Exclude filenames from being used for matching'
+                description: 'Elastic search query. Supports logical operators (AND, OR, NOT), and grouping with parentheses. For exact matches of specific identifiers, use quotes: "MyFunction", "SpecificStruct", "exact_variable_name". Examples: "config", "(term1 OR term2) AND term3", "getUserData", "struct Config".',
               },
               exact: {
                 type: 'boolean',
-                description: 'Perform exact search without tokenization (case-insensitive)'
+                description: 'When you exactly know what you are looking for, like known function name, struct name, or variable name, set this flag for precise matching'
               },
               allowTests: {
                 type: 'boolean',
-                description: 'Allow test files and test code blocks in results (disabled by default)'
+                description: 'Allow test files and test code blocks in results (enabled by default)',
+                default: true
               },
               session: {
                 type: 'string',
                 description: 'Session identifier for caching. Set to "new" if unknown, or want to reset cache. Re-use session ID returned from previous searches',
                 default: "new",
               },
-              timeout: {
-                type: 'number',
-                description: 'Timeout for the search operation in seconds (default: 30)',
-              },
               noGitignore: {
                 type: 'boolean',
                 description: 'Skip .gitignore files (will use PROBE_NO_GITIGNORE environment variable if not set)',
@@ -233,50 +199,6 @@ class ProbeServer {
             required: ['path', 'query']
           },
         },
-        {
-          name: 'query_code',
-          description: "Search code using ast-grep structural pattern matching. Use this tool to find specific code structures like functions, classes, or methods.",
-          inputSchema: {
-            type: 'object',
-            properties: {
-              path: {
-                type: 'string',
-                description: 'Absolute path to the directory to search in (e.g., "/Users/username/projects/myproject").',
-              },
-              pattern: {
-                type: 'string',
-                description: 'The ast-grep pattern to search for. Examples: "fn $NAME($$$PARAMS) $$$BODY" for Rust functions, "def $NAME($$$PARAMS): $$$BODY" for Python functions.',
-              },
-              language: {
-                type: 'string',
-                description: 'The programming language to search in. If not specified, the tool will try to infer the language from file extensions. Supported languages: rust, javascript, typescript, python, go, c, cpp, java, ruby, php, swift, csharp, yaml.',
-              },
-              ignore: {
-                type: 'array',
-                items: { type: 'string' },
-                description: 'Custom patterns to ignore (in addition to common patterns)',
-              },
-              maxResults: {
-                type: 'number',
-                description: 'Maximum number of results to return'
-              },
-              format: {
-                type: 'string',
-                enum: ['markdown', 'plain', 'json', 'color'],
-                description: 'Output format for the query results'
-              },
-              timeout: {
-                type: 'number',
-                description: 'Timeout for the query operation in seconds (default: 30)',
-              },
-              noGitignore: {
-                type: 'boolean',
-                description: 'Skip .gitignore files (will use PROBE_NO_GITIGNORE environment variable if not set)',
-              }
-            },
-            required: ['path', 'pattern']
-          },
-        },
         {
           name: 'extract_code',
           description: "Extract code blocks from files based on line number, or symbol name. Fetch full file when line number is not provided.",
@@ -294,22 +216,8 @@ class ProbeServer {
               },
               allowTests: {
                 type: 'boolean',
-                description: 'Allow test files and test code blocks in results (disabled by default)',
-              },
-              contextLines: {
-                type: 'number',
-                description: 'Number of context lines to include before and after the extracted block when AST parsing fails to find a suitable node',
-                default: 0
-              },
-              format: {
-                type: 'string',
-                enum: ['markdown', 'plain', 'json'],
-                description: 'Output format for the extracted code',
-                default: 'markdown'
-              },
-              timeout: {
-                type: 'number',
-                description: 'Timeout for the extract operation in seconds (default: 30)',
+                description: 'Allow test files and test code blocks in results (enabled by default)',
+                default: true
               },
               noGitignore: {
                 type: 'boolean',
@@ -323,8 +231,8 @@ class ProbeServer {
     }));
 
     this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
-      if (request.params.name !== 'search_code' && request.params.name !== 'query_code' && request.params.name !== 'extract_code' &&
-          request.params.name !== 'probe' && request.params.name !== 'query' && request.params.name !== 'extract') {
+      if (request.params.name !== 'search_code' && request.params.name !== 'extract_code' &&
+          request.params.name !== 'probe' && request.params.name !== 'extract') {
         throw new McpError(
           ErrorCode.MethodNotFound,
           `Unknown tool: ${request.params.name}`
@@ -356,9 +264,6 @@ class ProbeServer {
           }
           
           result = await this.executeCodeSearch(args);
-        } else if (request.params.name === 'query_code' || request.params.name === 'query') {
-          const args = request.params.arguments as unknown as QueryCodeArgs;
-          result = await this.executeCodeQuery(args);
         } else { // extract_code or extract
           const args = request.params.arguments as unknown as ExtractCodeArgs;
           result = await this.executeCodeExtract(args);
@@ -409,13 +314,15 @@ class ProbeServer {
       };
       
       // Add optional parameters only if they exist
-      if (args.filesOnly !== undefined) options.filesOnly = args.filesOnly;
-      if (args.ignore !== undefined) options.ignore = args.ignore;
-      if (args.excludeFilenames !== undefined) options.excludeFilenames = args.excludeFilenames;
       if (args.exact !== undefined) options.exact = args.exact;
       if (args.maxResults !== undefined) options.maxResults = args.maxResults;
       if (args.maxTokens !== undefined) options.maxTokens = args.maxTokens;
-      if (args.allowTests !== undefined) options.allowTests = args.allowTests;
+      // Set allowTests to true by default if not specified
+      if (args.allowTests !== undefined) {
+        options.allowTests = args.allowTests;
+      } else {
+        options.allowTests = true;
+      }
       // Use noGitignore from args, or fall back to PROBE_NO_GITIGNORE environment variable
       if (args.noGitignore !== undefined) {
         options.noGitignore = args.noGitignore;
@@ -427,12 +334,6 @@ class ProbeServer {
       } else {
         options.session = "new";
       }
-      // Use timeout from args, or fall back to instance default
-      if (args.timeout !== undefined) {
-        options.timeout = args.timeout;
-      } else if (this.defaultTimeout !== undefined) {
-        options.timeout = this.defaultTimeout;
-      }
 
       // Handle format options
       if (this.defaultFormat === 'outline-xml') {
@@ -467,50 +368,6 @@ class ProbeServer {
     }
   }
 
-  private async executeCodeQuery(args: QueryCodeArgs): Promise<string> {
-    try {
-      // Validate required parameters
-      if (!args.path) {
-        throw new Error("Path is required");
-      }
-      if (!args.pattern) {
-        throw new Error("Pattern is required");
-      }
-
-      // Create a single options object with both pattern and path
-      const options: any = {
-        path: args.path,
-        pattern: args.pattern,
-        language: args.language,
-        ignore: args.ignore,
-        allowTests: args.allowTests,
-        maxResults: args.maxResults,
-        format: args.format,
-        timeout: args.timeout || this.defaultTimeout
-      };
-      
-      // Use noGitignore from args, or fall back to PROBE_NO_GITIGNORE environment variable
-      if (args.noGitignore !== undefined) {
-        options.noGitignore = args.noGitignore;
-      } else if (process.env.PROBE_NO_GITIGNORE) {
-        options.noGitignore = process.env.PROBE_NO_GITIGNORE === 'true';
-      }
-      
-      console.log("Executing query with options:", JSON.stringify({
-        path: options.path,
-        pattern: options.pattern
-      }));
-      
-      const result = await query(options);
-      return result;
-    } catch (error: any) {
-      console.error('Error executing code query:', error);
-      throw new McpError(
-        'MethodNotFound' as unknown as ErrorCode,
-        `Error executing code query: ${error.message || String(error)}`
-      );
-    }
-  }
 
   private async executeCodeExtract(args: ExtractCodeArgs): Promise<string> {
     try {
@@ -526,12 +383,16 @@ class ProbeServer {
       const options: any = {
         files: args.files,
         path: args.path,
-        allowTests: args.allowTests,
-        contextLines: args.contextLines,
-        format: args.format,
-        timeout: args.timeout || this.defaultTimeout
+        format: 'xml'
       };
       
+      // Set allowTests to true by default if not specified
+      if (args.allowTests !== undefined) {
+        options.allowTests = args.allowTests;
+      } else {
+        options.allowTests = true;
+      }
+      
       // Use noGitignore from args, or fall back to PROBE_NO_GITIGNORE environment variable
       if (args.noGitignore !== undefined) {
         options.noGitignore = args.noGitignore;

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) {
@@ -2782,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: [],
@@ -4608,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");
@@ -4617,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();
@@ -4627,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
@@ -4670,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;
@@ -4795,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
        */
@@ -5176,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
                 });
@@ -5212,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",
@@ -5336,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}`);
                     }