@probelabs/probe 0.6.0-rc207 → 0.6.0-rc209

package/build/agent/mcp/client.js

@@ -9,6 +9,47 @@ import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js';
 import { WebSocketClientTransport } from '@modelcontextprotocol/sdk/client/websocket.js';
 import { loadMCPConfiguration, parseEnabledServers, DEFAULT_TIMEOUT } from './config.js';
 
+/**
+ * Check if a method is allowed based on server's method filter configuration
+ * Supports wildcard patterns (e.g., "*_read", "search_*", "prefix_*_suffix")
+ * @param {string} methodName - The method name to check
+ * @param {string[]|null} allowedMethods - Array of allowed method patterns (null = all allowed)
+ * @param {string[]|null} blockedMethods - Array of blocked method patterns (null = none blocked)
+ * @returns {boolean} Whether the method is allowed
+ */
+export function isMethodAllowed(methodName, allowedMethods, blockedMethods) {
+  /**
+   * Check if a method name matches a pattern
+   * Supports * wildcard which matches any characters
+   * @param {string} name - Method name to check
+   * @param {string} pattern - Pattern to match against (may contain *)
+   * @returns {boolean} Whether the name matches the pattern
+   */
+  const matchesPattern = (name, pattern) => {
+    if (!pattern.includes('*')) {
+      return name === pattern;
+    }
+    // Convert pattern to regex: escape special chars, replace * with .*
+    const regexPattern = pattern
+      .replace(/[.+?^${}()|[\]\\]/g, '\\$&')  // Escape special regex chars
+      .replace(/\*/g, '.*');                    // Replace * with .*
+    return new RegExp(`^${regexPattern}$`).test(name);
+  };
+
+  // If allowedMethods is specified (whitelist mode), only those methods are allowed
+  if (allowedMethods && allowedMethods.length > 0) {
+    return allowedMethods.some(pattern => matchesPattern(methodName, pattern));
+  }
+
+  // If blockedMethods is specified (blacklist mode), all methods except those are allowed
+  if (blockedMethods && blockedMethods.length > 0) {
+    return !blockedMethods.some(pattern => matchesPattern(methodName, pattern));
+  }
+
+  // No filter specified - all methods are allowed
+  return true;
+}
+
 /**
  * Create transport based on configuration
  * @param {Object} serverConfig - Server configuration
@@ -117,6 +158,18 @@ export class MCPClientManager {
     this.tools = new Map();
     this.debug = options.debug || process.env.DEBUG_MCP === '1';
     this.config = null;
+    this.tracer = options.tracer || null;
+  }
+
+  /**
+   * Record an MCP telemetry event if tracer is available
+   * @param {string} eventType - Event type (e.g., 'server.connect', 'tool.discovered')
+   * @param {Object} data - Event data
+   */
+  recordMcpEvent(eventType, data = {}) {
+    if (this.tracer && typeof this.tracer.recordMcpEvent === 'function') {
+      this.tracer.recordMcpEvent(eventType, data);
+    }
   }
 
   /**
@@ -128,12 +181,24 @@ export class MCPClientManager {
     this.config = config || loadMCPConfiguration();
     const servers = parseEnabledServers(this.config);
 
+    // Record initialization start
+    this.recordMcpEvent('initialization.started', {
+      serverCount: servers.length,
+      serverNames: servers.map(s => s.name)
+    });
+
     // Always log the number of servers found
     console.error(`[MCP INFO] Found ${servers.length} enabled MCP server${servers.length !== 1 ? 's' : ''}`);
 
     if (servers.length === 0) {
       console.error('[MCP INFO] No MCP servers configured or enabled');
       console.error('[MCP INFO] 0 MCP tools available');
+      this.recordMcpEvent('initialization.completed', {
+        connected: 0,
+        total: 0,
+        toolCount: 0,
+        tools: []
+      });
       return {
         connected: 0,
         total: 0,
@@ -178,10 +243,19 @@ export class MCPClientManager {
       });
     }
 
+    // Record initialization completion
+    const toolNames = Array.from(this.tools.keys());
+    this.recordMcpEvent('initialization.completed', {
+      connected: connectedCount,
+      total: servers.length,
+      toolCount: this.tools.size,
+      tools: toolNames
+    });
+
     return {
       connected: connectedCount,
       total: servers.length,
-      tools: Array.from(this.tools.keys())
+      tools: toolNames
     };
   }
 
@@ -192,6 +266,14 @@ export class MCPClientManager {
   async connectToServer(serverConfig) {
     const { name } = serverConfig;
 
+    // Record connection attempt
+    this.recordMcpEvent('server.connecting', {
+      serverName: name,
+      transport: serverConfig.transport,
+      hasAllowedMethods: !!(serverConfig.allowedMethods && serverConfig.allowedMethods.length > 0),
+      hasBlockedMethods: !!(serverConfig.blockedMethods && serverConfig.blockedMethods.length > 0)
+    });
+
     try {
       if (this.debug) {
         console.error(`[MCP DEBUG] Connecting to ${name} via ${serverConfig.transport}...`);
@@ -223,10 +305,34 @@ export class MCPClientManager {
 
       // Fetch and register tools
       const toolsResponse = await client.listTools();
-      const toolCount = toolsResponse?.tools?.length || 0;
+      const totalToolCount = toolsResponse?.tools?.length || 0;
+      let registeredCount = 0;
+      let filteredCount = 0;
+      const registeredTools = [];
+      const filteredTools = [];
 
       if (toolsResponse && toolsResponse.tools) {
+        const { allowedMethods, blockedMethods } = serverConfig;
+        const allToolNames = toolsResponse.tools.map(t => t.name);
+
+        // Record tools discovered from server
+        this.recordMcpEvent('tools.discovered', {
+          serverName: name,
+          toolCount: totalToolCount,
+          tools: allToolNames
+        });
+
         for (const tool of toolsResponse.tools) {
+          // Apply method filtering based on server config
+          if (!isMethodAllowed(tool.name, allowedMethods, blockedMethods)) {
+            filteredCount++;
+            filteredTools.push(tool.name);
+            if (this.debug) {
+              console.error(`[MCP DEBUG]     Filtered out tool: ${tool.name} (not allowed by method filter)`);
+            }
+            continue;
+          }
+
           // Add server prefix to avoid conflicts
           const qualifiedName = `${name}_${tool.name}`;
           this.tools.set(qualifiedName, {
@@ -234,14 +340,68 @@ export class MCPClientManager {
             serverName: name,
             originalName: tool.name
           });
+          registeredCount++;
+          registeredTools.push(qualifiedName);
 
           if (this.debug) {
             console.error(`[MCP DEBUG]     Registered tool: ${qualifiedName}`);
           }
         }
+
+        // Record method filtering results if any filtering was applied
+        if (filteredCount > 0) {
+          this.recordMcpEvent('tools.filtered', {
+            serverName: name,
+            filteredCount,
+            filteredTools,
+            allowedMethods: allowedMethods || [],
+            blockedMethods: blockedMethods || []
+          });
+        }
+
+        // Check for unmatched patterns in allowedMethods and warn users
+        if (allowedMethods && allowedMethods.length > 0) {
+          const unmatchedPatterns = allowedMethods.filter(pattern => {
+            // Check if this pattern matches at least one tool
+            return !allToolNames.some(toolName => isMethodAllowed(toolName, [pattern], null));
+          });
+
+          if (unmatchedPatterns.length > 0) {
+            console.error(`[MCP WARN] Server '${name}': The following allowedMethods patterns did not match any tools: ${unmatchedPatterns.join(', ')}`);
+            console.error(`[MCP WARN] Available methods from '${name}': ${allToolNames.join(', ')}`);
+          }
+        }
+
+        // Check for unmatched patterns in blockedMethods and warn users
+        if (blockedMethods && blockedMethods.length > 0) {
+          const unmatchedPatterns = blockedMethods.filter(pattern => {
+            // Check if this pattern matches at least one tool
+            return !allToolNames.some(toolName => !isMethodAllowed(toolName, null, [pattern]));
+          });
+
+          if (unmatchedPatterns.length > 0) {
+            console.error(`[MCP WARN] Server '${name}': The following blockedMethods patterns did not match any tools: ${unmatchedPatterns.join(', ')}`);
+            console.error(`[MCP WARN] Available methods from '${name}': ${allToolNames.join(', ')}`);
+          }
+        }
+      }
+
+      // Log connection result with filtering info
+      if (filteredCount > 0) {
+        console.error(`[MCP INFO] Connected to ${name}: ${registeredCount} tool${registeredCount !== 1 ? 's' : ''} loaded (${filteredCount} filtered out)`);
+      } else {
+        console.error(`[MCP INFO] Connected to ${name}: ${registeredCount} tool${registeredCount !== 1 ? 's' : ''} loaded`);
       }
 
-      console.error(`[MCP INFO] Connected to ${name}: ${toolCount} tool${toolCount !== 1 ? 's' : ''} loaded`);
+      // Record successful connection
+      this.recordMcpEvent('server.connected', {
+        serverName: name,
+        transport: serverConfig.transport,
+        totalToolCount,
+        registeredCount,
+        filteredCount,
+        registeredTools
+      });
 
       return true;
     } catch (error) {
@@ -249,6 +409,14 @@ export class MCPClientManager {
       if (this.debug) {
         console.error(`[MCP DEBUG] Full error details:`, error);
       }
+
+      // Record connection failure
+      this.recordMcpEvent('server.connection_failed', {
+        serverName: name,
+        transport: serverConfig.transport,
+        error: error.message
+      });
+
       return false;
     }
   }
@@ -261,14 +429,32 @@ export class MCPClientManager {
   async callTool(toolName, args) {
     const tool = this.tools.get(toolName);
     if (!tool) {
+      this.recordMcpEvent('tool.call_failed', {
+        toolName,
+        error: 'Unknown tool'
+      });
       throw new Error(`Unknown tool: ${toolName}`);
     }
 
     const clientInfo = this.clients.get(tool.serverName);
     if (!clientInfo) {
+      this.recordMcpEvent('tool.call_failed', {
+        toolName,
+        serverName: tool.serverName,
+        error: 'Server not connected'
+      });
       throw new Error(`Server ${tool.serverName} not connected`);
     }
 
+    const startTime = Date.now();
+
+    // Record tool call start
+    this.recordMcpEvent('tool.call_started', {
+      toolName,
+      serverName: tool.serverName,
+      originalToolName: tool.originalName
+    });
+
     try {
       if (this.debug) {
         console.error(`[MCP DEBUG] Calling ${toolName} with args:`, JSON.stringify(args, null, 2));
@@ -288,24 +474,48 @@ export class MCPClientManager {
       });
 
       // Race between the actual call and timeout
+      // Pass timeout to SDK's callTool to override its default 60s timeout
       const result = await Promise.race([
         clientInfo.client.callTool({
           name: tool.originalName,
           arguments: args
-        }),
+        }, undefined, { timeout }),
         timeoutPromise
       ]);
 
+      const durationMs = Date.now() - startTime;
+
       if (this.debug) {
         console.error(`[MCP DEBUG] Tool ${toolName} executed successfully`);
       }
 
+      // Record successful tool call
+      this.recordMcpEvent('tool.call_completed', {
+        toolName,
+        serverName: tool.serverName,
+        originalToolName: tool.originalName,
+        durationMs
+      });
+
       return result;
     } catch (error) {
+      const durationMs = Date.now() - startTime;
+
       console.error(`[MCP ERROR] Error calling tool ${toolName}:`, error.message);
       if (this.debug) {
         console.error(`[MCP DEBUG] Full error details:`, error);
       }
+
+      // Record failed tool call
+      this.recordMcpEvent('tool.call_failed', {
+        toolName,
+        serverName: tool.serverName,
+        originalToolName: tool.originalName,
+        error: error.message,
+        durationMs,
+        isTimeout: error.message.includes('timeout')
+      });
+
       throw error;
     }
   }
@@ -357,6 +567,7 @@ export class MCPClientManager {
    */
   async disconnect() {
     const disconnectPromises = [];
+    const serverNames = Array.from(this.clients.keys());
 
     if (this.clients.size === 0) {
       if (this.debug) {
@@ -365,6 +576,12 @@ export class MCPClientManager {
       return;
     }
 
+    // Record disconnection start
+    this.recordMcpEvent('disconnection.started', {
+      serverCount: this.clients.size,
+      serverNames
+    });
+
     if (this.debug) {
       console.error(`[MCP DEBUG] Disconnecting from ${this.clients.size} MCP server${this.clients.size !== 1 ? 's' : ''}...`);
     }
@@ -376,9 +593,16 @@ export class MCPClientManager {
             if (this.debug) {
               console.error(`[MCP DEBUG] Disconnected from ${name}`);
             }
+            this.recordMcpEvent('server.disconnected', {
+              serverName: name
+            });
           })
           .catch(error => {
             console.error(`[MCP ERROR] Error disconnecting from ${name}:`, error.message);
+            this.recordMcpEvent('server.disconnect_failed', {
+              serverName: name,
+              error: error.message
+            });
           })
       );
     }
@@ -387,6 +611,12 @@ export class MCPClientManager {
     this.clients.clear();
     this.tools.clear();
 
+    // Record disconnection completion
+    this.recordMcpEvent('disconnection.completed', {
+      serverCount: serverNames.length,
+      serverNames
+    });
+
     if (this.debug) {
       console.error('[MCP DEBUG] All MCP connections closed');
     }
@@ -405,5 +635,6 @@ export async function createMCPManager(options = {}) {
 export default {
   MCPClientManager,
   createMCPManager,
-  createTransport
+  createTransport,
+  isMethodAllowed
 };

package/build/agent/mcp/config.js

@@ -29,6 +29,61 @@ export function validateTimeout(value) {
   return Math.min(num, MAX_TIMEOUT); // Cap at max timeout
 }
 
+/**
+ * Validate and normalize method filter configuration
+ * @param {Object} serverConfig - Server configuration
+ * @param {string} serverName - Server name for logging
+ * @returns {Object} Object with allowedMethods and blockedMethods (null if not configured)
+ */
+export function validateMethodFilter(serverConfig, serverName = 'unknown') {
+  const result = { allowedMethods: null, blockedMethods: null };
+  const debug = process.env.DEBUG === '1' || process.env.DEBUG_MCP === '1';
+
+  // Check if both are specified - allowedMethods takes precedence
+  if (serverConfig.allowedMethods && serverConfig.blockedMethods) {
+    console.error(`[MCP WARN] Server '${serverName}' has both allowedMethods and blockedMethods - using allowedMethods only`);
+  }
+
+  // Process allowedMethods
+  if (serverConfig.allowedMethods) {
+    if (!Array.isArray(serverConfig.allowedMethods)) {
+      console.error(`[MCP WARN] Server '${serverName}' allowedMethods must be an array, ignoring`);
+    } else {
+      const validMethods = serverConfig.allowedMethods.filter(m => typeof m === 'string' && m.length > 0);
+      if (validMethods.length !== serverConfig.allowedMethods.length) {
+        console.error(`[MCP WARN] Server '${serverName}' allowedMethods contains non-string values, skipping those`);
+      }
+      if (validMethods.length > 0) {
+        result.allowedMethods = validMethods;
+        if (debug) {
+          console.error(`[MCP DEBUG] Server '${serverName}' allowedMethods: ${validMethods.join(', ')}`);
+        }
+      }
+    }
+    return result; // If allowedMethods is specified (even if invalid), don't process blockedMethods
+  }
+
+  // Process blockedMethods (only if allowedMethods not specified)
+  if (serverConfig.blockedMethods) {
+    if (!Array.isArray(serverConfig.blockedMethods)) {
+      console.error(`[MCP WARN] Server '${serverName}' blockedMethods must be an array, ignoring`);
+    } else {
+      const validMethods = serverConfig.blockedMethods.filter(m => typeof m === 'string' && m.length > 0);
+      if (validMethods.length !== serverConfig.blockedMethods.length) {
+        console.error(`[MCP WARN] Server '${serverName}' blockedMethods contains non-string values, skipping those`);
+      }
+      if (validMethods.length > 0) {
+        result.blockedMethods = validMethods;
+        if (debug) {
+          console.error(`[MCP DEBUG] Server '${serverName}' blockedMethods: ${validMethods.join(', ')}`);
+        }
+      }
+    }
+  }
+
+  return result;
+}
+
 /**
  * Default MCP configuration structure
  */
@@ -187,6 +242,16 @@ function mergeWithEnvironment(config) {
             console.error(`[MCP WARN] Invalid timeout value for ${normalizedName}: ${value}`);
           }
           break;
+        case 'ALLOWLIST':
+          // Comma-separated list of allowed method names (supports wildcards)
+          // e.g., MCP_SERVERS_MYSERVER_ALLOWLIST=method1,method2,prefix_*
+          config.mcpServers[normalizedName].allowedMethods = value.split(',').map(m => m.trim()).filter(Boolean);
+          break;
+        case 'BLOCKLIST':
+          // Comma-separated list of blocked method names (supports wildcards)
+          // e.g., MCP_SERVERS_MYSERVER_BLOCKLIST=dangerous_*,risky_method
+          config.mcpServers[normalizedName].blockedMethods = value.split(',').map(m => m.trim()).filter(Boolean);
+          break;
       }
     }
   }
@@ -256,6 +321,11 @@ export function parseEnabledServers(config) {
       server.timeout = validatedTimeout;
     }
 
+    // Validate and normalize method filter configuration
+    const methodFilter = validateMethodFilter(serverConfig, name);
+    server.allowedMethods = methodFilter.allowedMethods;
+    server.blockedMethods = methodFilter.blockedMethods;
+
     servers.push(server);
   }
 
@@ -321,6 +391,22 @@ export function createSampleConfig() {
         enabled: false,
         timeout: 120000,
         description: 'Example server with custom 2-minute timeout (overrides global setting)'
+      },
+      'filtered-server-example': {
+        command: 'npx',
+        args: ['-y', '@example/mcp-server'],
+        transport: 'stdio',
+        enabled: false,
+        allowedMethods: ['safe_read', 'safe_query'],
+        description: 'Example server with method allowlist - only safe_read and safe_query are available'
+      },
+      'blocklist-server-example': {
+        command: 'npx',
+        args: ['-y', '@example/mcp-server'],
+        transport: 'stdio',
+        enabled: false,
+        blockedMethods: ['dangerous_delete', 'dangerous_*'],
+        description: 'Example server with method blocklist - all methods except dangerous ones (supports wildcards)'
       }
     },
     // Global settings (apply to all servers unless overridden per-server)
@@ -356,6 +442,7 @@ export default {
   createSampleConfig,
   saveConfig,
   validateTimeout,
+  validateMethodFilter,
   DEFAULT_TIMEOUT,
   MAX_TIMEOUT
 };

package/build/agent/mcp/xmlBridge.js

@@ -127,6 +127,7 @@ export function parseXmlMcpToolCall(xmlString, mcpToolNames = []) {
 export class MCPXmlBridge {
   constructor(options = {}) {
     this.debug = options.debug || false;
+    this.tracer = options.tracer || null;
     this.mcpTools = {};
     this.mcpManager = null;
     this.xmlDefinitions = {};
@@ -179,8 +180,8 @@ export class MCPXmlBridge {
         console.error('[MCP DEBUG] Initializing MCP client manager...');
       }
 
-      // Initialize the MCP client manager
-      this.mcpManager = new MCPClientManager({ debug: this.debug });
+      // Initialize the MCP client manager with tracer support
+      this.mcpManager = new MCPClientManager({ debug: this.debug, tracer: this.tracer });
       const result = await this.mcpManager.initialize(mcpConfigs);
 
       // Get tools from the manager
@@ -211,11 +212,20 @@ export class MCPXmlBridge {
   }
 
   /**
-   * Get all XML tool definitions for inclusion in system prompt
+   * Get XML tool definitions for inclusion in system prompt
+   * @param {Array<string>|null} filterToolNames - Optional list of tool names to include (if null, include all)
    * @returns {string} Combined XML tool definitions
    */
-  getXmlToolDefinitions() {
-    return Object.values(this.xmlDefinitions).join('\n\n');
+  getXmlToolDefinitions(filterToolNames = null) {
+    if (filterToolNames === null) {
+      return Object.values(this.xmlDefinitions).join('\n\n');
+    }
+
+    // Filter definitions based on provided tool names
+    return Object.entries(this.xmlDefinitions)
+      .filter(([name]) => filterToolNames.includes(name))
+      .map(([, def]) => def)
+      .join('\n\n');
   }
 
   /**

package/build/agent/simpleTelemetry.js

@@ -231,6 +231,32 @@ export class SimpleAppTracer {
     });
   }
 
+  /**
+   * Record MCP (Model Context Protocol) events
+   * Tracks server connections, tool discovery, method filtering, and tool execution
+   */
+  recordMcpEvent(eventType, data = {}) {
+    if (!this.isEnabled()) return;
+
+    this.addEvent(`mcp.${eventType}`, {
+      'session.id': this.sessionId,
+      ...data
+    });
+  }
+
+  /**
+   * Record bash tool events
+   * Tracks command permission checks, allowed/denied commands, and execution
+   */
+  recordBashEvent(eventType, data = {}) {
+    if (!this.isEnabled()) return;
+
+    this.addEvent(`bash.${eventType}`, {
+      'session.id': this.sessionId,
+      ...data
+    });
+  }
+
   setAttributes(attributes) {
     // For simplicity, just log attributes when no active span
     if (this.telemetry && this.telemetry.enableConsole) {

package/build/tools/bash.js

@@ -31,16 +31,18 @@ export const bashTool = (options = {}) => {
     bashConfig = {},
     debug = false,
     cwd,
-    allowedFolders = []
+    allowedFolders = [],
+    tracer = null
   } = options;
 
-  // Create permission checker
+  // Create permission checker with tracer for telemetry
   const permissionChecker = new BashPermissionChecker({
     allow: bashConfig.allow,
     deny: bashConfig.deny,
     disableDefaultAllow: bashConfig.disableDefaultAllow,
     disableDefaultDeny: bashConfig.disableDefaultDeny,
-    debug
+    debug,
+    tracer
   });
 
   // Determine default working directory