@probelabs/probe 0.6.0-rc161 → 0.6.0-rc162

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@probelabs/probe",
-  "version": "0.6.0-rc161",
+  "version": "0.6.0-rc162",
   "description": "Node.js wrapper for the probe code search tool",
   "main": "src/index.js",
   "module": "src/index.js",

package/src/agent/ProbeAgent.js

@@ -60,6 +60,7 @@ import {
 } from './mcp/index.js';
 import { RetryManager, createRetryManagerFromEnv } from './RetryManager.js';
 import { FallbackManager, createFallbackManagerFromEnv, buildFallbackProvidersFromEnv } from './FallbackManager.js';
+import { handleContextLimitError } from './contextCompactor.js';
 
 // Maximum tool iterations to prevent infinite loops - configurable via MAX_TOOL_ITERATIONS env var
 const MAX_TOOL_ITERATIONS = (() => {
@@ -104,6 +105,8 @@ export class ProbeAgent {
    * @param {Array} [options.mcpServers] - Deprecated, use mcpConfig instead
    * @param {Object} [options.storageAdapter] - Custom storage adapter for history management
    * @param {Object} [options.hooks] - Hook callbacks for events (e.g., {'tool:start': callback})
+   * @param {Array<string>|null} [options.allowedTools] - List of allowed tool names. Use ['*'] for all tools (default), [] or null for no tools (raw AI mode), or specific tool names like ['search', 'query', 'extract']. Supports exclusion with '!' prefix (e.g., ['*', '!bash'])
+   * @param {boolean} [options.disableTools=false] - Convenience flag to disable all tools (equivalent to allowedTools: []). Takes precedence over allowedTools if set.
    * @param {Object} [options.retry] - Retry configuration
    * @param {number} [options.retry.maxRetries=3] - Maximum retry attempts per provider
    * @param {number} [options.retry.initialDelay=1000] - Initial delay in ms
@@ -139,6 +142,13 @@ export class ProbeAgent {
     this.disableMermaidValidation = !!options.disableMermaidValidation;
     this.disableJsonValidation = !!options.disableJsonValidation;
 
+    // Tool filtering configuration
+    // Parse allowedTools option: ['*'] = all tools, [] or null = no tools, ['tool1', 'tool2'] = specific tools
+    // Supports exclusion with '!' prefix: ['*', '!bash'] = all tools except bash
+    // disableTools is a convenience flag that overrides allowedTools to []
+    const effectiveAllowedTools = options.disableTools ? [] : options.allowedTools;
+    this.allowedTools = this._parseAllowedTools(effectiveAllowedTools);
+
     // Storage adapter (defaults to in-memory)
     this.storageAdapter = options.storageAdapter || new InMemoryStorageAdapter();
 
@@ -216,6 +226,74 @@ export class ProbeAgent {
     // Constructor must remain synchronous for backward compatibility
   }
 
+  /**
+   * Parse allowedTools configuration
+   * @param {Array<string>|null|undefined} allowedTools - Tool filtering configuration
+   * @returns {Object} Parsed configuration with isEnabled method
+   * @private
+   */
+  _parseAllowedTools(allowedTools) {
+    // Helper to check if tool matches a pattern (supports * wildcard)
+    const matchesPattern = (toolName, pattern) => {
+      if (!pattern.includes('*')) {
+        return toolName === pattern;
+      }
+      const regexPattern = pattern.replace(/\*/g, '.*');
+      return new RegExp(`^${regexPattern}$`).test(toolName);
+    };
+
+    // Default: all tools allowed
+    if (!allowedTools || (Array.isArray(allowedTools) && allowedTools.includes('*'))) {
+      const exclusions = Array.isArray(allowedTools)
+        ? allowedTools.filter(t => t.startsWith('!')).map(t => t.slice(1))
+        : [];
+
+      return {
+        mode: 'all',
+        exclusions,
+        isEnabled: (toolName) => !exclusions.some(pattern => matchesPattern(toolName, pattern))
+      };
+    }
+
+    // Empty array or null: no tools (raw AI mode)
+    if (Array.isArray(allowedTools) && allowedTools.length === 0) {
+      return {
+        mode: 'none',
+        isEnabled: () => false
+      };
+    }
+
+    // Specific tools allowed (with wildcard support)
+    const allowedPatterns = allowedTools.filter(t => !t.startsWith('!'));
+    return {
+      mode: 'whitelist',
+      allowed: allowedPatterns,
+      isEnabled: (toolName) => allowedPatterns.some(pattern => matchesPattern(toolName, pattern))
+    };
+  }
+
+  /**
+   * Check if an MCP tool is allowed based on allowedTools configuration
+   * Uses mcp__ prefix convention (like Claude Code)
+   * @param {string} toolName - The MCP tool name (without mcp__ prefix)
+   * @returns {boolean} - Whether the tool is allowed
+   * @private
+   */
+  _isMcpToolAllowed(toolName) {
+    const mcpToolName = `mcp__${toolName}`;
+    return this.allowedTools.isEnabled(mcpToolName) || this.allowedTools.isEnabled(toolName);
+  }
+
+  /**
+   * Filter MCP tools based on allowedTools configuration
+   * @param {string[]} mcpToolNames - Array of MCP tool names
+   * @returns {string[]} - Filtered array of allowed MCP tool names
+   * @private
+   */
+  _filterMcpTools(mcpToolNames) {
+    return mcpToolNames.filter(toolName => this._isMcpToolAllowed(toolName));
+  }
+
   /**
    * Initialize the agent asynchronously (must be called after constructor)
    * This method initializes MCP and merges MCP tools into the tool list, and loads history from storage
@@ -248,10 +326,15 @@ export class ProbeAgent {
         await this.initializeMCP();
 
         // Merge MCP tools into toolImplementations for unified access
+        // Apply allowedTools filtering using mcp__ prefix (like Claude Code)
         if (this.mcpBridge) {
           const mcpTools = this.mcpBridge.mcpTools || {};
           for (const [toolName, toolImpl] of Object.entries(mcpTools)) {
-            this.toolImplementations[toolName] = toolImpl;
+            if (this._isMcpToolAllowed(toolName)) {
+              this.toolImplementations[toolName] = toolImpl;
+            } else if (this.debug) {
+              console.error(`[DEBUG] MCP tool '${toolName}' filtered out by allowedTools`);
+            }
           }
         }
 
@@ -1047,10 +1130,15 @@ export class ProbeAgent {
         await this.initializeMCP();
 
         // Merge MCP tools into toolImplementations for unified access
+        // Apply allowedTools filtering using mcp__ prefix (like Claude Code)
         if (this.mcpBridge) {
           const mcpTools = this.mcpBridge.mcpTools || {};
           for (const [toolName, toolImpl] of Object.entries(mcpTools)) {
-            this.toolImplementations[toolName] = toolImpl;
+            if (this._isMcpToolAllowed(toolName)) {
+              this.toolImplementations[toolName] = toolImpl;
+            } else if (this.debug) {
+              console.error(`[DEBUG] MCP tool '${toolName}' filtered out by allowedTools`);
+            }
           }
         }
       } catch (error) {
@@ -1061,24 +1149,53 @@ export class ProbeAgent {
       }
     }
 
-    // Build tool definitions
-    let toolDefinitions = `
-${searchToolDefinition}
-${queryToolDefinition}
-${extractToolDefinition}
-${listFilesToolDefinition}
-${searchFilesToolDefinition}
-${attemptCompletionToolDefinition}
-`;
-    if (this.allowEdit) {
+    // Build tool definitions based on allowedTools configuration
+    let toolDefinitions = '';
+
+    // Helper to check if a tool is allowed
+    const isToolAllowed = (toolName) => this.allowedTools.isEnabled(toolName);
+
+    // Core tools (filtered by allowedTools)
+    if (isToolAllowed('search')) {
+      toolDefinitions += `${searchToolDefinition}\n`;
+    }
+    if (isToolAllowed('query')) {
+      toolDefinitions += `${queryToolDefinition}\n`;
+    }
+    if (isToolAllowed('extract')) {
+      toolDefinitions += `${extractToolDefinition}\n`;
+    }
+    if (isToolAllowed('listFiles')) {
+      toolDefinitions += `${listFilesToolDefinition}\n`;
+    }
+    if (isToolAllowed('searchFiles')) {
+      toolDefinitions += `${searchFilesToolDefinition}\n`;
+    }
+
+    // Edit tools (require both allowEdit flag AND allowedTools permission)
+    if (this.allowEdit && isToolAllowed('implement')) {
       toolDefinitions += `${implementToolDefinition}\n`;
+    }
+    if (this.allowEdit && isToolAllowed('edit')) {
       toolDefinitions += `${editToolDefinition}\n`;
+    }
+    if (this.allowEdit && isToolAllowed('create')) {
       toolDefinitions += `${createToolDefinition}\n`;
     }
-    if (this.enableBash) {
+
+    // Bash tool (require both enableBash flag AND allowedTools permission)
+    if (this.enableBash && isToolAllowed('bash')) {
       toolDefinitions += `${bashToolDefinition}\n`;
     }
-    if (this.enableDelegate) {
+
+    // Always include attempt_completion (unless explicitly disabled in raw AI mode)
+    if (isToolAllowed('attempt_completion')) {
+      toolDefinitions += `${attemptCompletionToolDefinition}\n`;
+    }
+
+    // Delegate tool (require both enableDelegate flag AND allowedTools permission)
+    // Place after attempt_completion as it's an optional tool
+    if (this.enableDelegate && isToolAllowed('delegate')) {
       toolDefinitions += `${delegateToolDefinition}\n`;
     }
 
@@ -1257,11 +1374,17 @@ When troubleshooting:
     // Add Tool Definitions
     systemMessage += `\n# Tools Available\n${toolDefinitions}\n`;
 
-    // Add MCP tools if available
+    // Add MCP tools if available (filtered by allowedTools)
     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`;
+      const allMcpTools = this.mcpBridge.getToolNames();
+      const allowedMcpTools = this._filterMcpTools(allMcpTools);
+
+      if (allowedMcpTools.length > 0) {
+        systemMessage += `\n## MCP Tools (JSON parameters in <params> tag)\n`;
+        // Get only allowed MCP tool definitions
+        systemMessage += this.mcpBridge.getXmlToolDefinitions(allowedMcpTools);
+        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
@@ -1447,57 +1570,120 @@ When troubleshooting:
 
         // Make AI request
         let assistantResponseContent = '';
-        try {
-          // Wrap AI request with tracing if available
-          const executeAIRequest = async () => {
-            // Prepare messages with potential image content
-            const messagesForAI = this.prepareMessagesWithImages(currentMessages);
-
-            const result = await this.streamTextWithRetryAndFallback({
-              model: this.provider(this.model),
-              messages: messagesForAI,
-              maxTokens: maxResponseTokens,
-              temperature: 0.3,
-            });
+        let compactionAttempted = false;
+
+        // Retry loop for context compaction - separate from streamTextWithRetryAndFallback
+        // which handles transient errors (rate limits, network issues, etc.)
+        while (true) {
+          try {
+            // Wrap AI request with tracing if available
+            const executeAIRequest = async () => {
+              // Prepare messages with potential image content
+              const messagesForAI = this.prepareMessagesWithImages(currentMessages);
+
+              const result = await this.streamTextWithRetryAndFallback({
+                model: this.provider(this.model),
+                messages: messagesForAI,
+                maxTokens: maxResponseTokens,
+                temperature: 0.3,
+              });
+
+              // Get the promise reference BEFORE consuming stream (doesn't lock it)
+              const usagePromise = result.usage;
 
-            // Get the promise reference BEFORE consuming stream (doesn't lock it)
-            const usagePromise = result.usage;
+              // Collect the streamed response - stream all content for now
+              for await (const delta of result.textStream) {
+                assistantResponseContent += delta;
+                // For now, stream everything - we'll handle segmentation after tools execute
+                if (options.onStream) {
+                  options.onStream(delta);
+                }
+              }
 
-            // Collect the streamed response - stream all content for now
-            for await (const delta of result.textStream) {
-              assistantResponseContent += delta;
-              // For now, stream everything - we'll handle segmentation after tools execute
-              if (options.onStream) {
-                options.onStream(delta);
+              // Record token usage - await the promise AFTER stream is consumed
+              const usage = await usagePromise;
+              if (usage) {
+                this.tokenCounter.recordUsage(usage, result.experimental_providerMetadata);
               }
+
+              return result;
+            };
+
+            if (this.tracer) {
+              await this.tracer.withSpan('ai.request', executeAIRequest, {
+                'ai.model': this.model,
+                'ai.provider': this.clientApiProvider || 'auto',
+                'iteration': currentIteration,
+                'max_tokens': maxResponseTokens,
+                'temperature': 0.3,
+                'message_count': currentMessages.length
+              });
+            } else {
+              await executeAIRequest();
             }
 
-            // Record token usage - await the promise AFTER stream is consumed
-            const usage = await usagePromise;
-            if (usage) {
-              this.tokenCounter.recordUsage(usage, result.experimental_providerMetadata);
+            // Success - break out of compaction retry loop
+            break;
+
+          } catch (error) {
+            // Check if this is a context limit error (only try compaction once per iteration)
+            if (!compactionAttempted && handleContextLimitError) {
+              const compactionResult = handleContextLimitError(error, currentMessages, {
+                keepLastSegment: true,
+                minSegmentsToKeep: 1
+              });
+
+              if (compactionResult) {
+                // Context limit error detected - compact and retry once
+                const { messages: compactedMessages, stats } = compactionResult;
+
+                // Check if compaction actually reduced message count
+                if (stats.removed === 0) {
+                  // No messages removed - compaction won't help, fail immediately
+                  console.error(`[ERROR] Context window exceeded but no messages can be compacted.`);
+                  console.error(`[ERROR] The conversation history is already minimal (${stats.originalCount} messages).`);
+                  finalResult = `Error: Context window limit exceeded and conversation cannot be compacted further. Consider starting a new session or reducing system message size.`;
+                  throw new Error(finalResult);
+                }
+
+                compactionAttempted = true;
+
+                console.log(`[INFO] Context window limit exceeded. Compacting conversation...`);
+                console.log(`[INFO] Removed ${stats.removed} messages (${stats.reductionPercent}% reduction)`);
+                console.log(`[INFO] Estimated token savings: ${stats.tokensSaved} tokens`);
+
+                if (this.debug) {
+                  console.log(`[DEBUG] Compaction stats:`, stats);
+                  console.log(`[DEBUG] Original message count: ${stats.originalCount}`);
+                  console.log(`[DEBUG] Compacted message count: ${stats.compactedCount}`);
+                }
+
+                // Replace currentMessages with compacted version (creates new array reference)
+                // This ensures we don't mutate the original history array
+                currentMessages = [...compactedMessages];
+
+                // Log compaction event if tracer is available
+                if (this.tracer) {
+                  this.tracer.addEvent('context.compacted', {
+                    'iteration': currentIteration,
+                    'original_count': stats.originalCount,
+                    'compacted_count': stats.compactedCount,
+                    'reduction_percent': stats.reductionPercent,
+                    'tokens_saved': stats.tokensSaved
+                  });
+                }
+
+                // Continue to retry with compacted messages
+                continue;
+              }
             }
 
-            return result;
-          };
-
-          if (this.tracer) {
-            await this.tracer.withSpan('ai.request', executeAIRequest, {
-              'ai.model': this.model,
-              'ai.provider': this.clientApiProvider || 'auto',
-              'iteration': currentIteration,
-              'max_tokens': maxResponseTokens,
-              'temperature': 0.3,
-              'message_count': currentMessages.length
-            });
-          } else {
-            await executeAIRequest();
+            // Not a context limit error, compaction already attempted, or compaction not available
+            // IMPORTANT: This break prevents infinite loop if compacted messages still exceed limit
+            console.error(`Error during streamText (Iter ${currentIteration}):`, error);
+            finalResult = `Error: Failed to get response from AI model during iteration ${currentIteration}. ${error.message}`;
+            throw new Error(finalResult);
           }
-
-        } catch (error) {
-          console.error(`Error during streamText (Iter ${currentIteration}):`, error);
-          finalResult = `Error: Failed to get response from AI model during iteration ${currentIteration}. ${error.message}`;
-          throw new Error(finalResult);
         }
 
         // Log preview of assistant response for debugging loops
@@ -1937,15 +2123,13 @@ NOT: {"type": "object", "properties": {"name": {"type": "string"}}}
 Convert your previous response content into actual JSON data that follows this schema structure.`;
           
           // Call answer recursively with _schemaFormatted flag to prevent infinite loop
-          finalResult = await this.answer(schemaPrompt, [], { 
-            ...options, 
-            _schemaFormatted: true 
+          finalResult = await this.answer(schemaPrompt, [], {
+            ...options,
+            _schemaFormatted: true
           });
-          
-          // Step 2: Clean the response (remove code blocks)
-          finalResult = cleanSchemaResponse(finalResult);
-          
-          // Step 3: Validate and fix Mermaid diagrams if present
+
+          // Step 2: Validate and fix Mermaid diagrams if present (BEFORE cleaning schema)
+          // This ensures mermaid validation sees the full response before JSON extraction strips content
           if (!this.disableMermaidValidation) {
             try {
               if (this.debug) {
@@ -2007,19 +2191,16 @@ Convert your previous response content into actual JSON data that follows this s
           } else if (this.debug) {
             console.log(`[DEBUG] Mermaid validation: Skipped due to disableMermaidValidation option`);
           }
-          
+
+          // Step 3: Clean the response (remove code blocks, extract JSON)
+          // This happens AFTER mermaid validation to preserve full content for validation
+          finalResult = cleanSchemaResponse(finalResult);
+
           // Step 4: Validate and potentially correct JSON responses
           if (isJsonSchema(options.schema)) {
             if (this.debug) {
               console.log(`[DEBUG] JSON validation: Starting validation process for schema response`);
-              console.log(`[DEBUG] JSON validation: Response length: ${finalResult.length} chars`);
-            }
-
-            // Clean the response first to extract JSON from markdown/code blocks
-            finalResult = cleanSchemaResponse(finalResult);
-
-            if (this.debug) {
-              console.log(`[DEBUG] JSON validation: After cleaning, length: ${finalResult.length} chars`);
+              console.log(`[DEBUG] JSON validation: Cleaned response length: ${finalResult.length} chars`);
             }
 
             // Record JSON validation start in telemetry
@@ -2131,17 +2312,16 @@ Convert your previous response content into actual JSON data that follows this s
       } else if (reachedMaxIterations && options.schema && this.debug) {
         console.log('[DEBUG] Skipping schema formatting due to max iterations reached without completion');
       } else if (completionAttempted && options.schema && !options._schemaFormatted && !options._skipValidation) {
-        // For attempt_completion results with schema, still clean markdown if needed
+        // For attempt_completion results with schema, validate mermaid diagrams BEFORE cleaning schema
+        // This ensures mermaid validation sees the full response before JSON extraction strips content
         // Skip this validation if we're in a recursive correction call (_skipValidation flag)
         try {
-          finalResult = cleanSchemaResponse(finalResult);
-          
-          // Validate and fix Mermaid diagrams if present
+          // Validate and fix Mermaid diagrams if present (BEFORE schema cleaning)
           if (!this.disableMermaidValidation) {
             if (this.debug) {
-              console.log(`[DEBUG] Mermaid validation: Validating attempt_completion result...`);
+              console.log(`[DEBUG] Mermaid validation: Validating attempt_completion result BEFORE schema cleaning...`);
             }
-            
+
             const mermaidValidation = await validateAndFixMermaidResponse(finalResult, {
               debug: this.debug,
               path: this.allowedFolders[0],
@@ -2149,7 +2329,7 @@ Convert your previous response content into actual JSON data that follows this s
               model: this.model,
               tracer: this.tracer
             });
-            
+
             if (mermaidValidation.wasFixed) {
               finalResult = mermaidValidation.fixedResponse;
               if (this.debug) {
@@ -2164,6 +2344,9 @@ Convert your previous response content into actual JSON data that follows this s
           } else if (this.debug) {
             console.log(`[DEBUG] Mermaid validation: Skipped for attempt_completion result due to disableMermaidValidation option`);
           }
+
+          // Now clean the schema response (may extract JSON and discard other content)
+          finalResult = cleanSchemaResponse(finalResult);
           
           // Validate and potentially correct JSON for attempt_completion results
           if (isJsonSchema(options.schema)) {
@@ -2377,6 +2560,75 @@ Convert your previous response content into actual JSON data that follows this s
     }
   }
 
+  /**
+   * Manually compact conversation history
+   * Removes intermediate monologues from older segments while preserving
+   * user messages, final answers, and the most recent segment
+   *
+   * @param {Object} options - Compaction options
+   * @param {boolean} [options.keepLastSegment=true] - Keep the most recent segment intact
+   * @param {number} [options.minSegmentsToKeep=1] - Number of recent segments to preserve fully
+   * @returns {Object} Compaction statistics
+   */
+  async compactHistory(options = {}) {
+    const { compactMessages, calculateCompactionStats } = await import('./contextCompactor.js');
+
+    if (this.history.length === 0) {
+      if (this.debug) {
+        console.log(`[DEBUG] No history to compact for session ${this.sessionId}`);
+      }
+      return {
+        originalCount: 0,
+        compactedCount: 0,
+        removed: 0,
+        reductionPercent: 0,
+        originalTokens: 0,
+        compactedTokens: 0,
+        tokensSaved: 0
+      };
+    }
+
+    // Perform compaction
+    const compactedMessages = compactMessages(this.history, options);
+    const stats = calculateCompactionStats(this.history, compactedMessages);
+
+    // Update history
+    this.history = compactedMessages;
+
+    // Save to storage (clear old history first, then save compacted messages)
+    try {
+      // Clear existing history to avoid duplicates
+      await this.storageAdapter.clearHistory(this.sessionId);
+
+      // Save compacted messages
+      // Note: Using sequential saves as storage adapter interface doesn't support batch operations
+      // For large histories, consider implementing a batch save method in your custom adapter
+      for (const message of compactedMessages) {
+        await this.storageAdapter.saveMessage(this.sessionId, message);
+      }
+    } catch (error) {
+      console.error(`[ERROR] Failed to save compacted messages to storage:`, error);
+    }
+
+    // Log results
+    console.log(`[INFO] Manually compacted conversation history`);
+    console.log(`[INFO] Removed ${stats.removed} messages (${stats.reductionPercent}% reduction)`);
+    console.log(`[INFO] Estimated token savings: ${stats.tokensSaved} tokens`);
+
+    if (this.debug) {
+      console.log(`[DEBUG] Compaction stats:`, stats);
+    }
+
+    // Emit hook
+    await this.hooks.emit(HOOK_TYPES.STORAGE_SAVE, {
+      sessionId: this.sessionId,
+      compacted: true,
+      stats
+    });
+
+    return stats;
+  }
+
   /**
    * Clone this agent's session to create a new agent with shared conversation history
    * @param {Object} options - Clone options
@@ -2407,6 +2659,17 @@ Convert your previous response content into actual JSON data that follows this s
     }
 
     // Create new agent with same configuration
+    // Reconstruct the original allowedTools array from the parsed configuration
+    let allowedToolsArray = null;
+    if (this.allowedTools.mode === 'whitelist') {
+      allowedToolsArray = [...this.allowedTools.allowed];
+    } else if (this.allowedTools.mode === 'none') {
+      allowedToolsArray = [];
+    } else if (this.allowedTools.mode === 'all' && this.allowedTools.exclusions.length > 0) {
+      allowedToolsArray = ['*', ...this.allowedTools.exclusions.map(t => '!' + t)];
+    }
+    // If mode is 'all' with no exclusions, leave as null (default)
+
     const clonedAgent = new ProbeAgent({
       // Copy current agent's config
       customPrompt: this.customPrompt,
@@ -2423,6 +2686,7 @@ Convert your previous response content into actual JSON data that follows this s
       maxIterations: this.maxIterations,
       disableMermaidValidation: this.disableMermaidValidation,
       disableJsonValidation: this.disableJsonValidation,
+      allowedTools: allowedToolsArray,
       enableMcp: !!this.mcpBridge,
       mcpConfig: this.mcpConfig,
       enableBash: this.enableBash,