@probelabs/probe 0.6.0-rc154 → 0.6.0-rc161

package/build/agent/ProbeAgent.js

@@ -18,15 +18,19 @@ import { TokenCounter } from './tokenCounter.js';
 import { InMemoryStorageAdapter } from './storage/InMemoryStorageAdapter.js';
 import { HookManager, HOOK_TYPES } from './hooks/HookManager.js';
 import { SUPPORTED_IMAGE_EXTENSIONS, IMAGE_MIME_TYPES } from './imageConfig.js';
-import { 
+import {
   createTools,
   searchToolDefinition,
   queryToolDefinition,
   extractToolDefinition,
+  delegateToolDefinition,
+  bashToolDefinition,
   listFilesToolDefinition,
   searchFilesToolDefinition,
   attemptCompletionToolDefinition,
   implementToolDefinition,
+  editToolDefinition,
+  createToolDefinition,
   attemptCompletionSchema,
   parseXmlToolCallWithThinking
 } from './tools.js';
@@ -54,9 +58,18 @@ import {
   parseHybridXmlToolCall,
   loadMCPConfigurationFromPath
 } from './mcp/index.js';
+import { RetryManager, createRetryManagerFromEnv } from './RetryManager.js';
+import { FallbackManager, createFallbackManagerFromEnv, buildFallbackProvidersFromEnv } from './FallbackManager.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)
@@ -75,6 +88,7 @@ export class ProbeAgent {
    * @param {string} [options.customPrompt] - Custom prompt to replace the default system message
    * @param {string} [options.promptType] - Predefined prompt type (architect, code-review, support)
    * @param {boolean} [options.allowEdit=false] - Allow the use of the 'implement' tool
+   * @param {boolean} [options.enableDelegate=false] - Enable the delegate tool for task distribution to subagents
    * @param {string} [options.path] - Search directory path
    * @param {string} [options.provider] - Force specific AI provider
    * @param {string} [options.model] - Override model name
@@ -90,6 +104,18 @@ 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 {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
@@ -97,11 +123,18 @@ export class ProbeAgent {
     this.customPrompt = options.customPrompt || null;
     this.promptType = options.promptType || 'code-explorer';
     this.allowEdit = !!options.allowEdit;
+    this.enableDelegate = !!options.enableDelegate;
     this.debug = options.debug || process.env.DEBUG === '1';
     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;
@@ -168,6 +201,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();
 
@@ -281,6 +322,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;
 
@@ -349,15 +400,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.`);
@@ -366,17 +421,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
    */
@@ -888,6 +1072,14 @@ ${attemptCompletionToolDefinition}
 `;
     if (this.allowEdit) {
       toolDefinitions += `${implementToolDefinition}\n`;
+      toolDefinitions += `${editToolDefinition}\n`;
+      toolDefinitions += `${createToolDefinition}\n`;
+    }
+    if (this.enableBash) {
+      toolDefinitions += `${bashToolDefinition}\n`;
+    }
+    if (this.enableDelegate) {
+      toolDefinitions += `${delegateToolDefinition}\n`;
     }
 
     // Build XML tool guidelines
@@ -952,7 +1144,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.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).
 `;
@@ -968,7 +1160,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>
 `;
 
@@ -1257,8 +1452,8 @@ When troubleshooting:
           const executeAIRequest = async () => {
             // Prepare messages with potential image content
             const messagesForAI = this.prepareMessagesWithImages(currentMessages);
-            
-            const result = await streamText({
+
+            const result = await this.streamTextWithRetryAndFallback({
               model: this.provider(this.model),
               messages: messagesForAI,
               maxTokens: maxResponseTokens,
@@ -1321,7 +1516,13 @@ 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');
         }
         
         // Try parsing with hybrid parser that supports both native and MCP tools
@@ -1462,18 +1663,24 @@ When troubleshooting:
                 
                 // Execute tool with tracing if available
                 const executeToolCall = async () => {
-                  // For delegate tool, pass current iteration and max iterations
+                  // For delegate tool, pass current iteration, max iterations, session ID, and config
                   if (toolName === 'delegate') {
                     const enhancedParams = {
                       ...toolParams,
                       currentIteration,
                       maxIterations,
+                      parentSessionId: this.sessionId,  // Pass parent session ID for tracking
+                      path: this.searchPath,            // Inherit search path
+                      provider: this.provider,          // Inherit AI provider
+                      model: this.model,                // Inherit model
                       debug: this.debug,
                       tracer: this.tracer
                     };
-                    
+
                     if (this.debug) {
                       console.log(`[DEBUG] Executing delegate tool at iteration ${currentIteration}/${maxIterations}`);
+                      console.log(`[DEBUG] Parent session: ${this.sessionId}`);
+                      console.log(`[DEBUG] Inherited config: path=${this.searchPath}, provider=${this.provider}, model=${this.model}`);
                       console.log(`[DEBUG] Delegate task: ${toolParams.task?.substring(0, 100)}...`);
                     }
                     
@@ -2205,6 +2412,7 @@ Convert your previous response content into actual JSON data that follows this s
       customPrompt: this.customPrompt,
       promptType: this.promptType,
       allowEdit: this.allowEdit,
+      enableDelegate: this.enableDelegate,
       path: this.allowedFolders[0], // Use first allowed folder as primary path
       allowedFolders: [...this.allowedFolders],
       provider: this.clientApiProvider,

package/build/agent/RetryManager.d.ts

@@ -0,0 +1,157 @@
+/**
+ * TypeScript type definitions for RetryManager
+ */
+
+/**
+ * Retry configuration options
+ */
+export interface RetryOptions {
+  /** Maximum number of retry attempts (0-100) */
+  maxRetries?: number;
+  /** Initial delay in milliseconds (0-60000) */
+  initialDelay?: number;
+  /** Maximum delay in milliseconds (0-300000) */
+  maxDelay?: number;
+  /** Exponential backoff multiplier (1-10) */
+  backoffFactor?: number;
+  /** List of retryable error patterns */
+  retryableErrors?: string[];
+  /** Enable debug logging */
+  debug?: boolean;
+  /** Add random jitter to delays (default: true) */
+  jitter?: boolean;
+}
+
+/**
+ * Context information for retry operations
+ */
+export interface RetryContext {
+  /** Provider name for logging */
+  provider?: string;
+  /** Model name for logging */
+  model?: string;
+  /** AbortSignal for cancellation */
+  signal?: AbortSignal;
+  /** Additional context data */
+  [key: string]: any;
+}
+
+/**
+ * Retry statistics
+ */
+export interface RetryStats {
+  /** Total number of attempts made */
+  totalAttempts: number;
+  /** Total number of retries (excluding initial attempts) */
+  totalRetries: number;
+  /** Number of successful retries */
+  successfulRetries: number;
+  /** Number of failed retries (exhausted max retries) */
+  failedRetries: number;
+}
+
+/**
+ * Error information extracted from an error object
+ */
+export interface ErrorInfo {
+  /** Error message */
+  message: string;
+  /** Error type or constructor name */
+  type: string;
+  /** Error code if available */
+  code?: string;
+  /** HTTP status code if available */
+  statusCode?: number;
+  /** Provider name if available */
+  provider?: string;
+  /** Whether the error is retryable */
+  isRetryable: boolean;
+}
+
+/**
+ * RetryManager class for handling retry logic with exponential backoff
+ */
+export class RetryManager {
+  /** Maximum retry attempts */
+  maxRetries: number;
+  /** Initial delay in milliseconds */
+  initialDelay: number;
+  /** Maximum delay in milliseconds */
+  maxDelay: number;
+  /** Exponential backoff multiplier */
+  backoffFactor: number;
+  /** List of retryable error patterns */
+  retryableErrors: string[];
+  /** Debug logging enabled */
+  debug: boolean;
+  /** Jitter enabled */
+  jitter: boolean;
+  /** Retry statistics */
+  stats: RetryStats;
+
+  /**
+   * Create a new RetryManager
+   * @param options - Retry configuration options
+   */
+  constructor(options?: RetryOptions);
+
+  /**
+   * Execute a function with retry logic
+   * @param fn - Async function to execute
+   * @param context - Context information for logging
+   * @returns Result from the function
+   * @throws Error if all retries are exhausted or operation is aborted
+   */
+  executeWithRetry<T>(
+    fn: () => Promise<T>,
+    context?: RetryContext
+  ): Promise<T>;
+
+  /**
+   * Check if an error is retryable
+   * @param error - The error to check
+   * @returns True if error should be retried
+   */
+  isRetryable(error: Error): boolean;
+
+  /**
+   * Get retry statistics
+   * @returns Statistics object (copy)
+   */
+  getStats(): RetryStats;
+
+  /**
+   * Reset statistics
+   */
+  resetStats(): void;
+}
+
+/**
+ * Check if an error is retryable based on error patterns
+ * @param error - The error to check
+ * @param retryableErrors - List of retryable error patterns
+ * @returns True if error should be retried
+ */
+export function isRetryableError(
+  error: Error,
+  retryableErrors?: string[]
+): boolean;
+
+/**
+ * Extract meaningful error information for logging
+ * @param error - The error to extract info from
+ * @returns Error information object
+ */
+export function extractErrorInfo(error: Error): ErrorInfo;
+
+/**
+ * Create a RetryManager from environment variables
+ * @param debug - Enable debug logging
+ * @returns Configured RetryManager instance
+ */
+export function createRetryManagerFromEnv(debug?: boolean): RetryManager;
+
+/**
+ * Default retryable error patterns
+ */
+export const DEFAULT_RETRYABLE_ERRORS: string[];