@probelabs/probe 0.6.0-rc165 → 0.6.0-rc167

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@probelabs/probe",
-  "version": "0.6.0-rc165",
+  "version": "0.6.0-rc167",
   "description": "Node.js wrapper for the probe code search tool",
   "main": "src/index.js",
   "module": "src/index.js",
@@ -77,6 +77,7 @@
     "@ai-sdk/anthropic": "^2.0.8",
     "@ai-sdk/google": "^2.0.14",
     "@ai-sdk/openai": "^2.0.10",
+    "@anthropic-ai/claude-agent-sdk": "^0.1.46",
     "@modelcontextprotocol/sdk": "^1.0.0",
     "@probelabs/maid": "^0.0.21",
     "adm-zip": "^0.5.16",

package/src/agent/ProbeAgent.js

@@ -27,6 +27,7 @@ import {
   bashToolDefinition,
   listFilesToolDefinition,
   searchFilesToolDefinition,
+  readImageToolDefinition,
   attemptCompletionToolDefinition,
   implementToolDefinition,
   editToolDefinition,
@@ -220,6 +221,9 @@ export class ProbeAgent {
     this.fallbackConfig = options.fallback || null;
     this.fallbackManager = null; // Will be initialized in initializeModel
 
+    // Engine support - minimal interface for multi-engine compatibility
+    this.engine = null; // Will be set in initializeModel or getEngine
+
     // Initialize the AI model
     this.initializeModel();
 
@@ -300,6 +304,31 @@ export class ProbeAgent {
    * This method initializes MCP and merges MCP tools into the tool list, and loads history from storage
    */
   async initialize() {
+    // Check if we need to auto-detect claude-code provider
+    // This happens when no API keys are set and no provider is specified
+    if (!this.provider && !this.clientApiProvider && this.apiType !== 'claude-code') {
+      // Check if initializeModel marked as uninitialized (no API keys)
+      if (this.apiType === 'uninitialized') {
+        const claudeAvailable = await this.isClaudeCommandAvailable();
+        if (claudeAvailable) {
+          if (this.debug) {
+            console.log('[DEBUG] No API keys found, but claude command detected');
+            console.log('[DEBUG] Auto-switching to claude-code provider');
+          }
+          // Set provider to claude-code
+          this.clientApiProvider = 'claude-code';
+          this.provider = null;
+          this.model = this.clientApiModel || 'claude-3-5-sonnet-20241022';
+          this.apiType = 'claude-code';
+        } else {
+          // Neither API keys nor claude command available
+          throw new Error('No API key provided and claude command not found. Please either:\n' +
+            '1. Set an API key: ANTHROPIC_API_KEY, OPENAI_API_KEY, GOOGLE_GENERATIVE_AI_API_KEY, or AWS credentials\n' +
+            '2. Install claude command from https://docs.claude.com/en/docs/claude-code');
+        }
+      }
+    }
+
     // Load history from storage adapter
     try {
       const history = await this.storageAdapter.loadHistory(this.sessionId);
@@ -399,6 +428,23 @@ export class ProbeAgent {
       delegate: wrappedTools.delegateToolInstance,
       listFiles: listFilesToolInstance,
       searchFiles: searchFilesToolInstance,
+      readImage: {
+        execute: async (params) => {
+          const imagePath = params.path;
+          if (!imagePath) {
+            throw new Error('Image path is required');
+          }
+
+          // Load the image using the existing loadImageIfValid method
+          const loaded = await this.loadImageIfValid(imagePath);
+
+          if (!loaded) {
+            throw new Error(`Failed to load image: ${imagePath}. The file may not exist, be too large, have an unsupported format, or be outside allowed directories.`);
+          }
+
+          return `Image loaded successfully: ${imagePath}. The image is now available for analysis in the conversation.`;
+        }
+      }
     };
 
     // Add bash tool if enabled
@@ -434,6 +480,21 @@ export class ProbeAgent {
     }
   }
 
+  /**
+   * Check if claude command is available on the system
+   * @returns {Promise<boolean>} True if claude command is available
+   * @private
+   */
+  async isClaudeCommandAvailable() {
+    try {
+      const { execSync } = await import('child_process');
+      execSync('claude --version', { stdio: 'ignore' });
+      return true;
+    } catch (error) {
+      return false;
+    }
+  }
+
   /**
    * Initialize the AI model based on available API keys and forced provider setting
    */
@@ -447,6 +508,19 @@ export class ProbeAgent {
       return;
     }
 
+    // Skip API key requirement for Claude Code (uses built-in access in Claude Code)
+    if (this.clientApiProvider === 'claude-code' || process.env.USE_CLAUDE_CODE === 'true') {
+      // Claude Code engine will be initialized lazily in getEngine()
+      // Set minimal defaults for compatibility
+      this.provider = null;
+      this.model = modelName || 'claude-3-5-sonnet-20241022';
+      this.apiType = 'claude-code';
+      if (this.debug) {
+        console.log('[DEBUG] Claude Code engine selected - will use built-in access if available');
+      }
+      return;
+    }
+
     // Get API keys from environment variables
     // Support both ANTHROPIC_API_KEY and ANTHROPIC_AUTH_TOKEN (used by Z.AI)
     const anthropicApiKey = process.env.ANTHROPIC_API_KEY || process.env.ANTHROPIC_AUTH_TOKEN;
@@ -516,7 +590,12 @@ export class ProbeAgent {
       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.');
+      // No API keys found - mark for potential claude-code auto-detection in initialize()
+      this.apiType = 'uninitialized';
+      if (this.debug) {
+        console.log('[DEBUG] No API keys found - will check for claude command in initialize()');
+      }
+      // Don't throw error yet - will be checked in initialize() method
     }
   }
 
@@ -591,6 +670,60 @@ export class ProbeAgent {
    * @private
    */
   async streamTextWithRetryAndFallback(options) {
+    // Check if we should use Claude Code engine
+    if (this.clientApiProvider === 'claude-code' || process.env.USE_CLAUDE_CODE === 'true') {
+      try {
+        const engine = await this.getEngine();
+        if (engine && engine.query) {
+          // Convert Vercel AI SDK format to engine format
+          // Extract the ORIGINAL user message as the main prompt (skip any warning messages)
+          // Look for user messages that are NOT the warning message
+          const userMessages = options.messages.filter(m =>
+            m.role === 'user' &&
+            !m.content.includes('WARNING: You have reached the maximum tool iterations limit')
+          );
+          const lastUserMessage = userMessages[userMessages.length - 1];
+          const prompt = lastUserMessage ? lastUserMessage.content : '';
+
+          // Pass system message and other options
+          const engineOptions = {
+            maxTokens: options.maxTokens,
+            temperature: options.temperature,
+            messages: options.messages,
+            systemPrompt: options.messages.find(m => m.role === 'system')?.content
+          };
+
+          // Get the engine's query result (async generator)
+          const engineStream = engine.query(prompt, engineOptions);
+
+          // Create a text stream that extracts text from engine messages
+          async function* createTextStream() {
+            for await (const message of engineStream) {
+              if (message.type === 'text' && message.content) {
+                yield message.content;
+              } else if (typeof message === 'string') {
+                // If engine returns plain strings, pass them through
+                yield message;
+              }
+              // Ignore other message types for the text stream
+            }
+          }
+
+          // Wrap the engine result to match streamText interface
+          return {
+            textStream: createTextStream(),
+            usage: Promise.resolve({}), // Engine should handle its own usage tracking
+            // Add other streamText-compatible properties as needed
+          };
+        }
+      } catch (error) {
+        if (this.debug) {
+          console.log(`[DEBUG] Failed to use Claude Code engine, falling back to Vercel:`, error.message);
+        }
+        // Fall through to use Vercel engine as fallback
+      }
+    }
+
     // Initialize retry manager if not already created
     if (!this.retryManager) {
       this.retryManager = new RetryManager({
@@ -734,6 +867,56 @@ export class ProbeAgent {
     }
   }
 
+  /**
+   * Get or create the AI engine based on configuration
+   * @returns {Promise<Object>} Engine interface
+   * @private
+   */
+  async getEngine() {
+    // If engine already created, return it
+    if (this.engine) {
+      return this.engine;
+    }
+
+    // Try Claude Code engine if requested
+    if (this.clientApiProvider === 'claude-code' || process.env.USE_CLAUDE_CODE === 'true') {
+      try {
+        const { createEnhancedClaudeCLIEngine } = await import('./engines/enhanced-claude-code.js');
+
+        // For Claude Code, use a cleaner system prompt without XML formatting
+        // since it has native MCP support for tools
+        const systemPrompt = this.customPrompt || this.getClaudeNativeSystemPrompt();
+
+        this.engine = await createEnhancedClaudeCLIEngine({
+          agent: this, // Pass reference to ProbeAgent for tool access
+          systemPrompt: systemPrompt,
+          customPrompt: this.customPrompt,
+          sessionId: this.options?.sessionId,
+          debug: this.debug,
+          allowedTools: this.allowedTools  // Pass tool filtering configuration
+        });
+        if (this.debug) {
+          console.log('[DEBUG] Using Claude Code engine with Probe tools');
+          if (this.customPrompt) {
+            console.log('[DEBUG] Using custom prompt/persona');
+          }
+        }
+        return this.engine;
+      } catch (error) {
+        console.warn('[WARNING] Failed to load Claude Code engine:', error.message);
+        console.warn('[WARNING] Falling back to Vercel AI SDK');
+        this.clientApiProvider = null;
+      }
+    }
+    // Default to enhanced Vercel AI SDK (wraps existing logic)
+    const { createEnhancedVercelEngine } = await import('./engines/enhanced-vercel.js');
+    this.engine = createEnhancedVercelEngine(this);
+    if (this.debug) {
+      console.log('[DEBUG] Using Vercel AI SDK engine');
+    }
+    return this.engine;
+  }
+
   /**
    * Process assistant response content and detect/load image references
    * @param {string} content - The assistant's response content
@@ -1120,6 +1303,56 @@ export class ProbeAgent {
     }
   }
 
+  /**
+   * Get system prompt for Claude native engines (CLI/SDK) without XML formatting
+   * These engines have native MCP support and don't need XML instructions
+   */
+  getClaudeNativeSystemPrompt() {
+    let systemPrompt = '';
+
+    // Add persona/role if configured
+    if (this.predefinedPrompt) {
+      const personaPrompt = getPromptByType(this.predefinedPrompt);
+      if (personaPrompt?.system) {
+        systemPrompt += personaPrompt.system + '\n\n';
+      }
+    }
+
+    // Add high-level instructions about when to use tools
+    systemPrompt += `You have access to powerful code search and analysis tools through MCP:
+- search: Find code patterns using semantic search
+- extract: Extract specific code sections with context
+- query: Use AST patterns for structural code matching
+- listFiles: Browse directory contents
+- searchFiles: Find files by name patterns`;
+
+    if (this.enableBash) {
+      systemPrompt += `\n- bash: Execute bash commands for system operations`;
+    }
+
+    systemPrompt += `\n
+When exploring code:
+1. Start with search to find relevant code patterns
+2. Use extract to get detailed context when needed
+3. Prefer focused, specific searches over broad queries
+4. Combine multiple tools to build complete understanding`;
+
+    // Add workspace context
+    if (this.allowedFolders && this.allowedFolders.length > 0) {
+      systemPrompt += `\n\nWorkspace: ${this.allowedFolders.join(', ')}`;
+    }
+
+    // Add repository structure if available
+    if (this.fileList) {
+      systemPrompt += `\n\n# Repository Structure\n`;
+      systemPrompt += `You are working with a repository located at: ${this.allowedFolders[0]}\n\n`;
+      systemPrompt += `Here's an overview of the repository structure (showing up to 100 most relevant files):\n\n`;
+      systemPrompt += '```\n' + this.fileList + '\n```\n';
+    }
+
+    return systemPrompt;
+  }
+
   /**
    * Get the system message with instructions for the AI (XML Tool Format)
    */
@@ -1172,6 +1405,9 @@ export class ProbeAgent {
     if (isToolAllowed('searchFiles')) {
       toolDefinitions += `${searchFilesToolDefinition}\n`;
     }
+    if (isToolAllowed('readImage')) {
+      toolDefinitions += `${readImageToolDefinition}\n`;
+    }
 
     // Edit tools (require both allowEdit flag AND allowedTools permission)
     if (this.allowEdit && isToolAllowed('implement')) {
@@ -1262,6 +1498,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.
+- readImage: Read and load an image file for AI analysis.
 ${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).
@@ -1510,6 +1747,79 @@ When troubleshooting:
       const baseMaxIterations = this.maxIterations || MAX_TOOL_ITERATIONS;
       const maxIterations = options.schema ? baseMaxIterations + 4 : baseMaxIterations;
 
+      // Check if we're using Claude Code engine which handles its own agentic loop
+      const isClaudeCode = this.clientApiProvider === 'claude-code' || process.env.USE_CLAUDE_CODE === 'true';
+
+      if (isClaudeCode) {
+        // For Claude Code, bypass the tool loop entirely - it handles its own internal dialogue
+        if (this.debug) {
+          console.log(`[DEBUG] Using Claude Code engine - bypassing tool loop (black box mode)`);
+          console.log(`[DEBUG] Sending question directly to Claude Code: ${message.substring(0, 100)}...`);
+        }
+
+        // Send the message directly to Claude Code and collect the response
+        try {
+          const engine = await this.getEngine();
+          if (engine && engine.query) {
+            let assistantResponseContent = '';
+            let toolBatch = null;
+
+            // Query Claude Code directly with the message and schema
+            for await (const chunk of engine.query(message, options)) {
+              if (chunk.type === 'text' && chunk.content) {
+                assistantResponseContent += chunk.content;
+                if (options.onStream) {
+                  options.onStream(chunk.content);
+                }
+              } else if (chunk.type === 'toolBatch' && chunk.tools) {
+                // Store tool batch for processing after response
+                toolBatch = chunk.tools;
+                if (this.debug) {
+                  console.log(`[DEBUG] Received batch of ${chunk.tools.length} tool events from Claude Code`);
+                }
+              } else if (chunk.type === 'error') {
+                throw chunk.error;
+              }
+            }
+
+            // Emit tool events after response is complete (batch mode)
+            if (toolBatch && toolBatch.length > 0 && this.events) {
+              if (this.debug) {
+                console.log(`[DEBUG] Emitting ${toolBatch.length} tool events from Claude Code batch`);
+              }
+              for (const toolEvent of toolBatch) {
+                this.events.emit('toolCall', toolEvent);
+              }
+            }
+
+            // Update history with the exchange
+            this.history.push(userMessage);
+            this.history.push({
+              role: 'assistant',
+              content: assistantResponseContent
+            });
+
+            // Store conversation history
+            // TODO: storeConversationHistory is not yet implemented for Claude Code
+            // await this.storeConversationHistory(this.history, oldHistoryLength);
+
+            // Emit completion hook
+            await this.hooks.emit(HOOK_TYPES.COMPLETION, {
+              sessionId: this.sessionId,
+              prompt: message,
+              response: assistantResponseContent
+            });
+
+            return assistantResponseContent;
+          }
+        } catch (error) {
+          if (this.debug) {
+            console.error('[DEBUG] Claude Code error:', error);
+          }
+          throw error;
+        }
+      }
+
       if (this.debug) {
         console.log(`[DEBUG] Starting agentic flow for question: ${message.substring(0, 100)}...`);
         if (options.schema) {
@@ -1517,7 +1827,7 @@ When troubleshooting:
         }
       }
 
-      // Tool iteration loop
+      // Tool iteration loop (only for non-Claude Code engines)
       while (currentIteration < maxIterations && !completionAttempted) {
         currentIteration++;
         if (this.cancelled) throw new Error('Request was cancelled by the user');
@@ -1590,7 +1900,7 @@ When troubleshooting:
               const messagesForAI = this.prepareMessagesWithImages(currentMessages);
 
               const result = await this.streamTextWithRetryAndFallback({
-                model: this.provider(this.model),
+                model: this.provider ? this.provider(this.model) : this.model,
                 messages: messagesForAI,
                 maxTokens: maxResponseTokens,
                 temperature: 0.3,
@@ -1700,10 +2010,8 @@ When troubleshooting:
           console.log(`[DEBUG] Assistant response (${assistantResponseContent.length} chars): ${assistantPreview}`);
         }
 
-        // Process image references in assistant response for next iteration
-        if (assistantResponseContent) {
-          await this.processImageReferences(assistantResponseContent);
-        }
+        // Images in assistant responses are not automatically processed
+        // AI can use the readImage tool to explicitly request reading an image
 
         // Parse tool call from response with valid tools list
         // Build validTools based on allowedTools configuration (same pattern as getSystemMessage)
@@ -1713,6 +2021,7 @@ When troubleshooting:
         if (this.allowedTools.isEnabled('extract')) validTools.push('extract');
         if (this.allowedTools.isEnabled('listFiles')) validTools.push('listFiles');
         if (this.allowedTools.isEnabled('searchFiles')) validTools.push('searchFiles');
+        if (this.allowedTools.isEnabled('readImage')) validTools.push('readImage');
         if (this.allowedTools.isEnabled('attempt_completion')) validTools.push('attempt_completion');
 
         // Edit tools (require both allowEdit flag AND allowedTools permission)