@probelabs/probe 0.6.0-rc159 → 0.6.0-rc162

package/build/agent/ProbeAgent.js

@@ -24,10 +24,13 @@ import {
   queryToolDefinition,
   extractToolDefinition,
   delegateToolDefinition,
+  bashToolDefinition,
   listFilesToolDefinition,
   searchFilesToolDefinition,
   attemptCompletionToolDefinition,
   implementToolDefinition,
+  editToolDefinition,
+  createToolDefinition,
   attemptCompletionSchema,
   parseXmlToolCallWithThinking
 } from './tools.js';
@@ -55,9 +58,19 @@ import {
   parseHybridXmlToolCall,
   loadMCPConfigurationFromPath
 } 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 = parseInt(process.env.MAX_TOOL_ITERATIONS || '30', 10);
+const MAX_TOOL_ITERATIONS = (() => {
+  const val = parseInt(process.env.MAX_TOOL_ITERATIONS || '30', 10);
+  if (isNaN(val) || val < 1 || val > 200) {
+    console.warn('[ProbeAgent] MAX_TOOL_ITERATIONS must be between 1 and 200, using default: 30');
+    return 30;
+  }
+  return val;
+})();
 const MAX_HISTORY_MESSAGES = 100;
 
 // Supported image file extensions (imported from shared config)
@@ -92,6 +105,20 @@ 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
+   * @param {number} [options.retry.maxDelay=30000] - Maximum delay in ms
+   * @param {number} [options.retry.backoffFactor=2] - Exponential backoff multiplier
+   * @param {Array<string>} [options.retry.retryableErrors] - List of retryable error patterns
+   * @param {Object} [options.fallback] - Fallback configuration
+   * @param {string} [options.fallback.strategy] - Fallback strategy: 'same-model', 'same-provider', 'any', 'custom'
+   * @param {Array<string>} [options.fallback.models] - List of models for same-provider fallback
+   * @param {Array<Object>} [options.fallback.providers] - List of provider configurations for custom fallback
+   * @param {boolean} [options.fallback.stopOnSuccess=true] - Stop on first success
+   * @param {number} [options.fallback.maxTotalAttempts=10] - Maximum total attempts across all providers
    */
   constructor(options = {}) {
     // Basic configuration
@@ -104,11 +131,24 @@ export class ProbeAgent {
     this.cancelled = false;
     this.tracer = options.tracer || null;
     this.outline = !!options.outline;
-    this.maxResponseTokens = options.maxResponseTokens || parseInt(process.env.MAX_RESPONSE_TOKENS || '0', 10) || null;
+    this.maxResponseTokens = options.maxResponseTokens || (() => {
+      const val = parseInt(process.env.MAX_RESPONSE_TOKENS || '0', 10);
+      if (isNaN(val) || val < 0 || val > 200000) {
+        return null;
+      }
+      return val || null;
+    })();
     this.maxIterations = options.maxIterations || null;
     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();
 
@@ -171,6 +211,14 @@ export class ProbeAgent {
     this.mcpBridge = null;
     this._mcpInitialized = false; // Track if MCP initialization has been attempted
 
+    // Retry configuration
+    this.retryConfig = options.retry || {};
+    this.retryManager = null; // Will be initialized lazily when needed
+
+    // Fallback configuration
+    this.fallbackConfig = options.fallback || null;
+    this.fallbackManager = null; // Will be initialized in initializeModel
+
     // Initialize the AI model
     this.initializeModel();
 
@@ -178,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
@@ -210,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`);
+            }
           }
         }
 
@@ -284,6 +405,16 @@ export class ProbeAgent {
       this.toolImplementations.bash = wrappedTools.bashToolInstance;
     }
 
+    // Add edit and create tools if enabled
+    if (this.allowEdit) {
+      if (wrappedTools.editToolInstance) {
+        this.toolImplementations.edit = wrappedTools.editToolInstance;
+      }
+      if (wrappedTools.createToolInstance) {
+        this.toolImplementations.create = wrappedTools.createToolInstance;
+      }
+    }
+
     // Store wrapped tools for ACP system
     this.wrappedTools = wrappedTools;
 
@@ -352,15 +483,19 @@ export class ProbeAgent {
     if (forceProvider) {
       if (forceProvider === 'anthropic' && anthropicApiKey) {
         this.initializeAnthropicModel(anthropicApiKey, anthropicApiUrl, modelName);
+        this.initializeFallbackManager(forceProvider, modelName);
         return;
       } else if (forceProvider === 'openai' && openaiApiKey) {
         this.initializeOpenAIModel(openaiApiKey, openaiApiUrl, modelName);
+        this.initializeFallbackManager(forceProvider, modelName);
         return;
       } else if (forceProvider === 'google' && googleApiKey) {
         this.initializeGoogleModel(googleApiKey, googleApiUrl, modelName);
+        this.initializeFallbackManager(forceProvider, modelName);
         return;
       } else if (forceProvider === 'bedrock' && ((awsAccessKeyId && awsSecretAccessKey && awsRegion) || awsApiKey)) {
         this.initializeBedrockModel(awsAccessKeyId, awsSecretAccessKey, awsRegion, awsSessionToken, awsApiKey, awsBedrockBaseUrl, modelName);
+        this.initializeFallbackManager(forceProvider, modelName);
         return;
       }
       console.warn(`WARNING: Forced provider "${forceProvider}" selected but required API key is missing or invalid! Falling back to auto-detection.`);
@@ -369,17 +504,146 @@ export class ProbeAgent {
     // If no provider is forced or forced provider failed, use the first available API key
     if (anthropicApiKey) {
       this.initializeAnthropicModel(anthropicApiKey, anthropicApiUrl, modelName);
+      this.initializeFallbackManager('anthropic', modelName);
     } else if (openaiApiKey) {
       this.initializeOpenAIModel(openaiApiKey, openaiApiUrl, modelName);
+      this.initializeFallbackManager('openai', modelName);
     } else if (googleApiKey) {
       this.initializeGoogleModel(googleApiKey, googleApiUrl, modelName);
+      this.initializeFallbackManager('google', modelName);
     } else if ((awsAccessKeyId && awsSecretAccessKey && awsRegion) || awsApiKey) {
       this.initializeBedrockModel(awsAccessKeyId, awsSecretAccessKey, awsRegion, awsSessionToken, awsApiKey, awsBedrockBaseUrl, modelName);
+      this.initializeFallbackManager('bedrock', modelName);
     } else {
       throw new Error('No API key provided. Please set ANTHROPIC_API_KEY (or ANTHROPIC_AUTH_TOKEN), OPENAI_API_KEY, GOOGLE_GENERATIVE_AI_API_KEY (or GOOGLE_API_KEY), AWS credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION), or AWS_BEDROCK_API_KEY environment variables.');
     }
   }
 
+  /**
+   * Initialize fallback manager based on configuration
+   * @param {string} primaryProvider - The primary provider being used
+   * @param {string} primaryModel - The primary model being used
+   * @private
+   */
+  initializeFallbackManager(primaryProvider, primaryModel) {
+    // Skip fallback initialization if explicitly disabled or in test mode
+    if (this.fallbackConfig === false || process.env.DISABLE_FALLBACK === '1') {
+      return;
+    }
+
+    // If fallback config is provided explicitly, use it
+    if (this.fallbackConfig && this.fallbackConfig.providers) {
+      try {
+        this.fallbackManager = new FallbackManager({
+          ...this.fallbackConfig,
+          debug: this.debug
+        });
+
+        if (this.debug) {
+          console.log(`[DEBUG] Fallback manager initialized with ${this.fallbackManager.providers.length} providers`);
+        }
+      } catch (error) {
+        console.error('[WARNING] Failed to initialize fallback manager:', error.message);
+      }
+      return;
+    }
+
+    // Try to load from environment variables
+    const envFallbackManager = createFallbackManagerFromEnv(this.debug);
+    if (envFallbackManager) {
+      this.fallbackManager = envFallbackManager;
+      if (this.debug) {
+        console.log(`[DEBUG] Fallback manager initialized from environment variables`);
+      }
+      return;
+    }
+
+    // Auto-build fallback from available providers if enabled
+    if (process.env.AUTO_FALLBACK === '1' || this.fallbackConfig?.auto) {
+      const providers = buildFallbackProvidersFromEnv({
+        primaryProvider,
+        primaryModel
+      });
+
+      if (providers.length > 1) {
+        try {
+          this.fallbackManager = new FallbackManager({
+            strategy: 'custom',
+            providers,
+            debug: this.debug
+          });
+
+          if (this.debug) {
+            console.log(`[DEBUG] Auto-fallback enabled with ${providers.length} providers`);
+          }
+        } catch (error) {
+          console.error('[WARNING] Failed to initialize auto-fallback:', error.message);
+        }
+      }
+    }
+  }
+
+  /**
+   * Execute streamText with retry and fallback support
+   * @param {Object} options - streamText options
+   * @returns {Promise<Object>} - streamText result
+   * @private
+   */
+  async streamTextWithRetryAndFallback(options) {
+    // Initialize retry manager if not already created
+    if (!this.retryManager) {
+      this.retryManager = new RetryManager({
+        maxRetries: this.retryConfig.maxRetries ?? 3,
+        initialDelay: this.retryConfig.initialDelay ?? 1000,
+        maxDelay: this.retryConfig.maxDelay ?? 30000,
+        backoffFactor: this.retryConfig.backoffFactor ?? 2,
+        retryableErrors: this.retryConfig.retryableErrors,
+        debug: this.debug
+      });
+    }
+
+    // If no fallback manager, just use retry with current provider
+    if (!this.fallbackManager) {
+      return await this.retryManager.executeWithRetry(
+        () => streamText(options),
+        {
+          provider: this.apiType,
+          model: this.model
+        }
+      );
+    }
+
+    // Use fallback manager with retry for each provider
+    return await this.fallbackManager.executeWithFallback(
+      async (provider, model, config) => {
+        // Create options with the fallback provider
+        const fallbackOptions = {
+          ...options,
+          model: provider(model)
+        };
+
+        // Create a retry manager for this specific provider
+        const providerRetryManager = new RetryManager({
+          maxRetries: config.maxRetries ?? this.retryConfig.maxRetries ?? 3,
+          initialDelay: this.retryConfig.initialDelay ?? 1000,
+          maxDelay: this.retryConfig.maxDelay ?? 30000,
+          backoffFactor: this.retryConfig.backoffFactor ?? 2,
+          retryableErrors: this.retryConfig.retryableErrors,
+          debug: this.debug
+        });
+
+        // Execute with retry for this provider
+        return await providerRetryManager.executeWithRetry(
+          () => streamText(fallbackOptions),
+          {
+            provider: config.provider,
+            model: model
+          }
+        );
+      }
+    );
+  }
+
   /**
    * Initialize Anthropic model
    */
@@ -866,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) {
@@ -880,19 +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.enableDelegate) {
+    if (this.allowEdit && isToolAllowed('edit')) {
+      toolDefinitions += `${editToolDefinition}\n`;
+    }
+    if (this.allowEdit && isToolAllowed('create')) {
+      toolDefinitions += `${createToolDefinition}\n`;
+    }
+
+    // Bash tool (require both enableBash flag AND allowedTools permission)
+    if (this.enableBash && isToolAllowed('bash')) {
+      toolDefinitions += `${bashToolDefinition}\n`;
+    }
+
+    // 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`;
     }
 
@@ -958,7 +1261,7 @@ Available Tools:
 - extract: Extract specific code blocks or lines from files.
 - listFiles: List files and directories in a specified location.
 - searchFiles: Find files matching a glob pattern with recursive search capability.
-${this.allowEdit ? '- implement: Implement a feature or fix a bug using aider.\n' : ''}${this.enableDelegate ? '- delegate: Delegate big distinct tasks to specialized probe subagents.\n' : ''}
+${this.allowEdit ? '- implement: Implement a feature or fix a bug using aider.\n- edit: Edit files using exact string replacement.\n- create: Create new files with specified content.\n' : ''}${this.enableDelegate ? '- delegate: Delegate big distinct tasks to specialized probe subagents.\n' : ''}${this.enableBash ? '- bash: Execute bash commands for system operations.\n' : ''}
 - attempt_completion: Finalize the task and provide the result to the user.
 - attempt_complete: Quick completion using previous response (shorthand).
 `;
@@ -974,7 +1277,10 @@ Follow these instructions carefully:
 6. You MUST respond with exactly ONE tool call per message, using the specified XML format, until the task is complete.
 7. Wait for the tool execution result (provided in the next user message in a <tool_result> block) before proceeding to the next step.
 8. Once the task is fully completed, use the '<attempt_completion>' tool to provide the final result. This is the ONLY way to signal completion.
-9. Prefer concise and focused search queries. Use specific keywords and phrases to narrow down results.
+9. Prefer concise and focused search queries. Use specific keywords and phrases to narrow down results.${this.allowEdit ? `
+10. When modifying files, choose the appropriate tool:
+    - Use 'edit' for precise changes to existing files (requires exact string match)
+    - Use 'create' for new files or complete file rewrites` : ''}
 </instructions>
 `;
 
@@ -1068,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
@@ -1258,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 streamText({
-              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
@@ -1327,7 +1702,10 @@ When troubleshooting:
           'search', 'query', 'extract', 'listFiles', 'searchFiles', 'attempt_completion'
         ];
         if (this.allowEdit) {
-          validTools.push('implement');
+          validTools.push('implement', 'edit', 'create');
+        }
+        if (this.enableBash) {
+          validTools.push('bash');
         }
         if (this.enableDelegate) {
           validTools.push('delegate');
@@ -1745,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) {
@@ -1815,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
@@ -1939,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],
@@ -1957,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) {
@@ -1972,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)) {
@@ -2185,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
@@ -2215,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,
@@ -2231,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,