@probelabs/probe 0.6.0-rc88 → 0.6.0-rc89

package/build/agent/ProbeAgent.d.ts

@@ -23,6 +23,14 @@ export interface ProbeAgentOptions {
   debug?: boolean;
   /** Optional telemetry tracer instance */
   tracer?: any;
+  /** Enable MCP (Model Context Protocol) tool integration */
+  enableMcp?: boolean;
+  /** Path to MCP configuration file */
+  mcpConfigPath?: string;
+  /** MCP configuration object (overrides mcpConfigPath) */
+  mcpConfig?: any;
+  /** @deprecated Use mcpConfig instead */
+  mcpServers?: any[];
 }
 
 /**

package/build/agent/ProbeAgent.js

@@ -19,14 +19,15 @@ import {
   parseXmlToolCallWithThinking
 } from './tools.js';
 import { createMessagePreview } from '../tools/common.js';
-import { 
-  createWrappedTools, 
-  listFilesToolInstance, 
+import {
+  createWrappedTools,
+  listFilesToolInstance,
   searchFilesToolInstance,
-  clearToolExecutionData 
+  clearToolExecutionData
 } from './probeTool.js';
+import { createMockProvider } from './mockProvider.js';
 import { listFilesByLevel } from '../index.js';
-import { 
+import {
   cleanSchemaResponse,
   isJsonSchema,
   validateJsonResponse,
@@ -35,6 +36,11 @@ import {
   createSchemaDefinitionCorrectionPrompt,
   validateAndFixMermaidResponse
 } from './schemaUtils.js';
+import {
+  MCPXmlBridge,
+  parseHybridXmlToolCall,
+  loadMCPConfigurationFromPath
+} from './mcp/index.js';
 
 // Maximum tool iterations to prevent infinite loops - configurable via MAX_TOOL_ITERATIONS env var
 const MAX_TOOL_ITERATIONS = parseInt(process.env.MAX_TOOL_ITERATIONS || '30', 10);
@@ -58,6 +64,10 @@ export class ProbeAgent {
    * @param {boolean} [options.outline] - Enable outline-xml format for search results
    * @param {number} [options.maxResponseTokens] - Maximum tokens for AI responses
    * @param {boolean} [options.disableMermaidValidation=false] - Disable automatic mermaid diagram validation and fixing
+   * @param {boolean} [options.enableMcp=false] - Enable MCP tool integration
+   * @param {string} [options.mcpConfigPath] - Path to MCP configuration file
+   * @param {Object} [options.mcpConfig] - MCP configuration object (overrides mcpConfigPath)
+   * @param {Array} [options.mcpServers] - Deprecated, use mcpConfig instead
    */
   constructor(options = {}) {
     // Basic configuration
@@ -92,14 +102,29 @@ export class ProbeAgent {
     // Initialize tools
     this.initializeTools();
 
-    // Initialize the AI model
-    this.initializeModel();
-
     // Initialize chat history
     this.history = [];
-    
+
     // Initialize event emitter for tool execution updates
     this.events = new EventEmitter();
+
+    // MCP configuration
+    this.enableMcp = !!options.enableMcp || process.env.ENABLE_MCP === '1';
+    this.mcpConfigPath = options.mcpConfigPath || null;
+    this.mcpConfig = options.mcpConfig || null;
+    this.mcpServers = options.mcpServers || null; // Deprecated, keep for backward compatibility
+    this.mcpBridge = null;
+
+    // Initialize the AI model
+    this.initializeModel();
+
+    // Initialize MCP if enabled
+    if (this.enableMcp) {
+      this.initializeMCP().catch(error => {
+        console.error('[MCP] Failed to initialize MCP:', error);
+        this.mcpBridge = null;
+      });
+    }
   }
 
   /**
@@ -138,6 +163,12 @@ export class ProbeAgent {
    * Initialize the AI model based on available API keys and forced provider setting
    */
   initializeModel() {
+    // Check if we're in test mode and should use mock provider
+    if (process.env.NODE_ENV === 'test' || process.env.USE_MOCK_AI === 'true') {
+      this.initializeMockModel();
+      return;
+    }
+
     // Get API keys from environment variables
     const anthropicApiKey = process.env.ANTHROPIC_API_KEY;
     const openaiApiKey = process.env.OPENAI_API_KEY;
@@ -231,12 +262,83 @@ export class ProbeAgent {
     });
     this.model = modelName || 'gemini-2.5-pro';
     this.apiType = 'google';
-    
+
     if (this.debug) {
       console.log(`Using Google API with model: ${this.model}${apiUrl ? ` (URL: ${apiUrl})` : ''}`);
     }
   }
 
+  /**
+   * Initialize mock model for testing
+   */
+  initializeMockModel(modelName) {
+    this.provider = createMockProvider();
+    this.model = modelName || 'mock-model';
+    this.apiType = 'mock';
+
+    if (this.debug) {
+      console.log(`Using Mock API with model: ${this.model}`);
+    }
+  }
+
+  /**
+   * Initialize MCP bridge and load tools
+   */
+  async initializeMCP() {
+    if (!this.enableMcp) return;
+
+    try {
+      let mcpConfig = null;
+
+      // Priority order: mcpConfig > mcpConfigPath > mcpServers (deprecated) > auto-discovery
+      if (this.mcpConfig) {
+        // Direct config object provided (SDK usage)
+        mcpConfig = this.mcpConfig;
+        if (this.debug) {
+          console.log('[DEBUG] Using provided MCP config object');
+        }
+      } else if (this.mcpConfigPath) {
+        // Explicit config path provided
+        try {
+          mcpConfig = loadMCPConfigurationFromPath(this.mcpConfigPath);
+          if (this.debug) {
+            console.log(`[DEBUG] Loaded MCP config from: ${this.mcpConfigPath}`);
+          }
+        } catch (error) {
+          throw new Error(`Failed to load MCP config from ${this.mcpConfigPath}: ${error.message}`);
+        }
+      } else if (this.mcpServers) {
+        // Backward compatibility: convert old mcpServers format
+        mcpConfig = { mcpServers: this.mcpServers };
+        if (this.debug) {
+          console.warn('[DEBUG] Using deprecated mcpServers option. Consider using mcpConfig instead.');
+        }
+      }
+      // Note: auto-discovery fallback is removed - user must explicitly provide config
+
+      // Initialize the MCP XML bridge
+      this.mcpBridge = new MCPXmlBridge({ debug: this.debug });
+      await this.mcpBridge.initialize(mcpConfig);
+
+      const mcpToolCount = this.mcpBridge.getToolNames().length;
+      if (mcpToolCount > 0) {
+        if (this.debug) {
+          console.log(`[DEBUG] Loaded ${mcpToolCount} MCP tools`);
+        }
+      } else {
+        // For backward compatibility: if no tools were loaded, set bridge to null
+        // This maintains the behavior expected by existing tests
+        if (this.debug) {
+          console.log('[DEBUG] No MCP tools loaded, setting bridge to null');
+        }
+        this.mcpBridge = null;
+      }
+    } catch (error) {
+      console.error('[MCP] Error initializing MCP:', error);
+      this.mcpBridge = null;
+    }
+  }
+
   /**
    * Get the system message with instructions for the AI (XML Tool Format)
    */
@@ -428,6 +530,13 @@ When troubleshooting:
     // Add Tool Definitions
     systemMessage += `\n# Tools Available\n${toolDefinitions}\n`;
 
+    // Add MCP tools if available
+    if (this.mcpBridge && this.mcpBridge.getToolNames().length > 0) {
+      systemMessage += `\n## MCP Tools (JSON parameters in <params> tag)\n`;
+      systemMessage += this.mcpBridge.getXmlToolDefinitions();
+      systemMessage += `\n\nFor MCP tools, use JSON format within the params tag, e.g.:\n<mcp_tool>\n<params>\n{"key": "value"}\n</params>\n</mcp_tool>\n`;
+    }
+
     // Add folder information
     const searchDirectory = this.allowedFolders.length > 0 ? this.allowedFolders[0] : process.cwd();
     if (this.debug) {
@@ -640,7 +749,11 @@ When troubleshooting:
           validTools.push('implement');
         }
         
-        const parsedTool = parseXmlToolCallWithThinking(assistantResponseContent, validTools);
+        // Try parsing with hybrid parser that supports both native and MCP tools
+        const nativeTools = validTools;
+        const parsedTool = this.mcpBridge
+          ? parseHybridXmlToolCall(assistantResponseContent, nativeTools, this.mcpBridge)
+          : parseXmlToolCallWithThinking(assistantResponseContent, validTools);
         if (parsedTool) {
           const { toolName, params } = parsedTool;
           if (this.debug) console.log(`[DEBUG] Parsed tool call: ${toolName} with params:`, params);
@@ -654,7 +767,9 @@ When troubleshooting:
               const lastAssistantMessage = [...currentMessages].reverse().find(msg =>
                 msg.role === 'assistant' &&
                 msg.content &&
-                !parseXmlToolCallWithThinking(msg.content, validTools)
+                !(this.mcpBridge
+                  ? parseHybridXmlToolCall(msg.content, validTools, this.mcpBridge)
+                  : parseXmlToolCallWithThinking(msg.content, validTools))
               );
 
               if (lastAssistantMessage) {
@@ -677,8 +792,32 @@ When troubleshooting:
             }
             break;
           } else {
-            // Execute the tool
-            if (this.toolImplementations[toolName]) {
+            // Check tool type and execute accordingly
+            const { type } = parsedTool;
+
+            if (type === 'mcp' && this.mcpBridge && this.mcpBridge.isMcpTool(toolName)) {
+              // Execute MCP tool
+              try {
+                if (this.debug) console.log(`[DEBUG] Executing MCP tool '${toolName}' with params:`, params);
+
+                // Execute MCP tool through the bridge
+                const executionResult = await this.mcpBridge.mcpTools[toolName].execute(params);
+
+                const toolResultContent = typeof executionResult === 'string' ? executionResult : JSON.stringify(executionResult, null, 2);
+                const preview = createMessagePreview(toolResultContent);
+                if (this.debug) {
+                  console.log(`[DEBUG] MCP tool '${toolName}' executed successfully. Result preview: ${preview}`);
+                }
+
+                currentMessages.push({ role: 'user', content: `<tool_result>\n${toolResultContent}\n</tool_result>` });
+              } catch (error) {
+                console.error(`Error executing MCP tool ${toolName}:`, error);
+                const toolResultContent = `Error executing MCP tool ${toolName}: ${error.message}`;
+                if (this.debug) console.log(`[DEBUG] MCP tool '${toolName}' execution FAILED.`);
+                currentMessages.push({ role: 'user', content: `<tool_result>\n${toolResultContent}\n</tool_result>` });
+              }
+            } else if (this.toolImplementations[toolName]) {
+              // Execute native tool
               try {
                 // Add sessionId to params for tool execution
                 const toolParams = { ...params, sessionId: this.sessionId };
@@ -778,9 +917,15 @@ When troubleshooting:
             } else {
               console.error(`[ERROR] Unknown tool: ${toolName}`);
               currentMessages.push({ role: 'assistant', content: assistantResponseContent });
+
+              // Build list of available tools including MCP tools
+              const nativeTools = Object.keys(this.toolImplementations);
+              const mcpTools = this.mcpBridge ? this.mcpBridge.getToolNames() : [];
+              const allAvailableTools = [...nativeTools, ...mcpTools];
+
               currentMessages.push({
                 role: 'user',
-                content: `<tool_result>\nError: Unknown tool '${toolName}'. Available tools: ${Object.keys(this.toolImplementations).join(', ')}\n</tool_result>`
+                content: `<tool_result>\nError: Unknown tool '${toolName}'. Available tools: ${allAvailableTools.join(', ')}\n</tool_result>`
               });
             }
           }
@@ -1309,6 +1454,26 @@ Convert your previous response content into actual JSON data that follows this s
     }
   }
 
+  /**
+   * Clean up resources (including MCP connections)
+   */
+  async cleanup() {
+    // Clean up MCP bridge
+    if (this.mcpBridge) {
+      try {
+        await this.mcpBridge.cleanup();
+        if (this.debug) {
+          console.log('[DEBUG] MCP bridge cleaned up');
+        }
+      } catch (error) {
+        console.error('Error cleaning up MCP bridge:', error);
+      }
+    }
+
+    // Clear history and other resources
+    this.clearHistory();
+  }
+
   /**
    * Cancel the current request
    */

package/build/agent/acp/tools.js

@@ -269,7 +269,7 @@ export class ACPToolManager {
         name: 'search',
         description: 'Search for code patterns and content using flexible text search with stemming and stopword removal. Supports regex patterns and elastic search query syntax.',
         kind: ToolCallKind.search,
-        parameters: {
+        inputSchema: {
           type: 'object',
           properties: {
             query: {
@@ -296,7 +296,7 @@ export class ACPToolManager {
         name: 'query',
         description: 'Perform structural queries using AST patterns to find specific code structures like functions, classes, or methods.',
         kind: ToolCallKind.query,
-        parameters: {
+        inputSchema: {
           type: 'object',
           properties: {
             pattern: {
@@ -323,7 +323,7 @@ export class ACPToolManager {
         name: 'extract',
         description: 'Extract specific code blocks from files based on file paths and optional line numbers.',
         kind: ToolCallKind.extract,
-        parameters: {
+        inputSchema: {
           type: 'object',
           properties: {
             files: {
@@ -352,7 +352,7 @@ export class ACPToolManager {
         name: 'delegate',
         description: 'Automatically delegate big distinct tasks to specialized probe subagents within the agentic loop. Use when complex requests can be broken into focused, parallel tasks.',
         kind: ToolCallKind.execute,
-        parameters: {
+        inputSchema: {
           type: 'object',
           properties: {
             task: {

package/build/agent/acp/tools.test.js

@@ -288,26 +288,26 @@ describe('ACPToolManager', () => {
       const searchTool = definitions.find(d => d.name === 'search');
       expect(searchTool).toBeDefined();
       expect(searchTool.kind).toBe(ToolCallKind.search);
-      expect(searchTool.parameters.properties.query).toBeDefined();
-      expect(searchTool.parameters.required).toContain('query');
-      
+      expect(searchTool.inputSchema.properties.query).toBeDefined();
+      expect(searchTool.inputSchema.required).toContain('query');
+
       const queryTool = definitions.find(d => d.name === 'query');
       expect(queryTool).toBeDefined();
       expect(queryTool.kind).toBe(ToolCallKind.query);
-      expect(queryTool.parameters.properties.pattern).toBeDefined();
-      expect(queryTool.parameters.required).toContain('pattern');
-      
+      expect(queryTool.inputSchema.properties.pattern).toBeDefined();
+      expect(queryTool.inputSchema.required).toContain('pattern');
+
       const extractTool = definitions.find(d => d.name === 'extract');
       expect(extractTool).toBeDefined();
       expect(extractTool.kind).toBe(ToolCallKind.extract);
-      expect(extractTool.parameters.properties.files).toBeDefined();
-      expect(extractTool.parameters.required).toContain('files');
-      
+      expect(extractTool.inputSchema.properties.files).toBeDefined();
+      expect(extractTool.inputSchema.required).toContain('files');
+
       const delegateTool = definitions.find(d => d.name === 'delegate');
       expect(delegateTool).toBeDefined();
       expect(delegateTool.kind).toBe(ToolCallKind.execute);
-      expect(delegateTool.parameters.properties.task).toBeDefined();
-      expect(delegateTool.parameters.required).toContain('task');
+      expect(delegateTool.inputSchema.properties.task).toBeDefined();
+      expect(delegateTool.inputSchema.required).toContain('task');
     });
   });