@probelabs/probe 0.6.0-rc167 → 0.6.0-rc169

package/build/agent/ProbeAgent.d.ts

@@ -11,6 +11,8 @@ export interface ProbeAgentOptions {
   sessionId?: string;
   /** Custom system prompt to replace the default system message */
   customPrompt?: string;
+  /** Alias for customPrompt. More intuitive naming for system prompts. */
+  systemPrompt?: string;
   /** Predefined prompt type (persona) */
   promptType?: 'code-explorer' | 'engineer' | 'code-review' | 'support' | 'architect';
   /** Allow the use of the 'implement' tool for code editing */
@@ -35,6 +37,10 @@ export interface ProbeAgentOptions {
   mcpConfig?: any;
   /** @deprecated Use mcpConfig instead */
   mcpServers?: any[];
+  /** 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']). */
+  allowedTools?: string[] | null;
+  /** Convenience flag to disable all tools (equivalent to allowedTools: []). Takes precedence over allowedTools if set. */
+  disableTools?: boolean;
   /** Retry configuration for handling transient API failures */
   retry?: RetryOptions;
   /** Fallback configuration for multi-provider support */
@@ -227,4 +233,4 @@ export interface ProbeAgentEvents {
 }
 
 // Default export
-export { ProbeAgent as default };
+export { ProbeAgent as default };

package/build/agent/ProbeAgent.js

@@ -55,6 +55,7 @@ import {
   validateAndFixMermaidResponse
 } from './schemaUtils.js';
 import { removeThinkingTags } from './xmlParsingUtils.js';
+import { predefinedPrompts } from './shared/prompts.js';
 import {
   MCPXmlBridge,
   parseHybridXmlToolCall,
@@ -89,6 +90,7 @@ export class ProbeAgent {
    * @param {Object} options - Configuration options
    * @param {string} [options.sessionId] - Optional session ID
    * @param {string} [options.customPrompt] - Custom prompt to replace the default system message
+   * @param {string} [options.systemPrompt] - Alias for customPrompt; takes precedence when both are provided
    * @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
@@ -125,7 +127,8 @@ export class ProbeAgent {
   constructor(options = {}) {
     // Basic configuration
     this.sessionId = options.sessionId || randomUUID();
-    this.customPrompt = options.customPrompt || null;
+    // Support systemPrompt alias (overrides customPrompt when both are provided)
+    this.customPrompt = options.systemPrompt || options.customPrompt || null;
     this.promptType = options.promptType || 'code-explorer';
     this.allowEdit = !!options.allowEdit;
     this.enableDelegate = !!options.enableDelegate;
@@ -304,12 +307,14 @@ 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
+    // Check if we need to auto-detect claude-code or codex provider
     // This happens when no API keys are set and no provider is specified
-    if (!this.provider && !this.clientApiProvider && this.apiType !== 'claude-code') {
+    if (!this.provider && !this.clientApiProvider && this.apiType !== 'claude-code' && this.apiType !== 'codex') {
       // Check if initializeModel marked as uninitialized (no API keys)
       if (this.apiType === 'uninitialized') {
         const claudeAvailable = await this.isClaudeCommandAvailable();
+        const codexAvailable = await this.isCodexCommandAvailable();
+
         if (claudeAvailable) {
           if (this.debug) {
             console.log('[DEBUG] No API keys found, but claude command detected');
@@ -320,11 +325,22 @@ export class ProbeAgent {
           this.provider = null;
           this.model = this.clientApiModel || 'claude-3-5-sonnet-20241022';
           this.apiType = 'claude-code';
+        } else if (codexAvailable) {
+          if (this.debug) {
+            console.log('[DEBUG] No API keys found, but codex command detected');
+            console.log('[DEBUG] Auto-switching to codex provider');
+          }
+          // Set provider to codex
+          this.clientApiProvider = 'codex';
+          this.provider = null;
+          this.model = this.clientApiModel || 'gpt-4o';
+          this.apiType = 'codex';
         } else {
-          // Neither API keys nor claude command available
-          throw new Error('No API key provided and claude command not found. Please either:\n' +
+          // Neither API keys nor CLI commands available
+          throw new Error('No API key provided and neither claude nor codex command 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');
+            '2. Install claude command from https://docs.claude.com/en/docs/claude-code\n' +
+            '3. Install codex command from https://openai.com/codex');
         }
       }
     }
@@ -404,14 +420,20 @@ export class ProbeAgent {
    * Initialize tools with configuration
    */
   initializeTools() {
+    const isToolAllowed = (toolName) => this.allowedTools.isEnabled(toolName);
+
     const configOptions = {
       sessionId: this.sessionId,
       debug: this.debug,
       defaultPath: this.allowedFolders.length > 0 ? this.allowedFolders[0] : process.cwd(),
       allowedFolders: this.allowedFolders,
       outline: this.outline,
+      allowEdit: this.allowEdit,
+      enableDelegate: this.enableDelegate,
       enableBash: this.enableBash,
-      bashConfig: this.bashConfig
+      bashConfig: this.bashConfig,
+      allowedTools: this.allowedTools,
+      isToolAllowed
     };
 
     // Create base tools
@@ -420,15 +442,33 @@ export class ProbeAgent {
     // Create wrapped tools with event emission
     const wrappedTools = createWrappedTools(baseTools);
 
-    // Store tool instances for execution
-    this.toolImplementations = {
-      search: wrappedTools.searchToolInstance,
-      query: wrappedTools.queryToolInstance,
-      extract: wrappedTools.extractToolInstance,
-      delegate: wrappedTools.delegateToolInstance,
-      listFiles: listFilesToolInstance,
-      searchFiles: searchFilesToolInstance,
-      readImage: {
+    // Store tool instances for execution (respect allowedTools + feature flags)
+    this.toolImplementations = {};
+
+    if (wrappedTools.searchToolInstance && isToolAllowed('search')) {
+      this.toolImplementations.search = wrappedTools.searchToolInstance;
+    }
+    if (wrappedTools.queryToolInstance && isToolAllowed('query')) {
+      this.toolImplementations.query = wrappedTools.queryToolInstance;
+    }
+    if (wrappedTools.extractToolInstance && isToolAllowed('extract')) {
+      this.toolImplementations.extract = wrappedTools.extractToolInstance;
+    }
+    if (this.enableDelegate && wrappedTools.delegateToolInstance && isToolAllowed('delegate')) {
+      this.toolImplementations.delegate = wrappedTools.delegateToolInstance;
+    }
+
+    // File browsing tools
+    if (isToolAllowed('listFiles')) {
+      this.toolImplementations.listFiles = listFilesToolInstance;
+    }
+    if (isToolAllowed('searchFiles')) {
+      this.toolImplementations.searchFiles = searchFilesToolInstance;
+    }
+
+    // Image loading tool
+    if (isToolAllowed('readImage')) {
+      this.toolImplementations.readImage = {
         execute: async (params) => {
           const imagePath = params.path;
           if (!imagePath) {
@@ -444,20 +484,20 @@ export class ProbeAgent {
 
           return `Image loaded successfully: ${imagePath}. The image is now available for analysis in the conversation.`;
         }
-      }
-    };
+      };
+    }
 
-    // Add bash tool if enabled
-    if (this.enableBash && wrappedTools.bashToolInstance) {
+    // Add bash tool if enabled and allowed
+    if (this.enableBash && wrappedTools.bashToolInstance && isToolAllowed('bash')) {
       this.toolImplementations.bash = wrappedTools.bashToolInstance;
     }
 
-    // Add edit and create tools if enabled
+    // Add edit and create tools if enabled and allowed
     if (this.allowEdit) {
-      if (wrappedTools.editToolInstance) {
+      if (wrappedTools.editToolInstance && isToolAllowed('edit')) {
         this.toolImplementations.edit = wrappedTools.editToolInstance;
       }
-      if (wrappedTools.createToolInstance) {
+      if (wrappedTools.createToolInstance && isToolAllowed('create')) {
         this.toolImplementations.create = wrappedTools.createToolInstance;
       }
     }
@@ -482,13 +522,34 @@ export class ProbeAgent {
 
   /**
    * Check if claude command is available on the system
+   * Uses execFile instead of exec to avoid shell injection risks
    * @returns {Promise<boolean>} True if claude command is available
    * @private
    */
   async isClaudeCommandAvailable() {
     try {
-      const { execSync } = await import('child_process');
-      execSync('claude --version', { stdio: 'ignore' });
+      const { execFile } = await import('child_process');
+      const { promisify } = await import('util');
+      const execFileAsync = promisify(execFile);
+      await execFileAsync('claude', ['--version'], { timeout: 5000 });
+      return true;
+    } catch (error) {
+      return false;
+    }
+  }
+
+  /**
+   * Check if codex command is available on the system
+   * Uses execFile instead of exec to avoid shell injection risks
+   * @returns {Promise<boolean>} True if codex command is available
+   * @private
+   */
+  async isCodexCommandAvailable() {
+    try {
+      const { execFile } = await import('child_process');
+      const { promisify } = await import('util');
+      const execFileAsync = promisify(execFile);
+      await execFileAsync('codex', ['--version'], { timeout: 5000 });
       return true;
     } catch (error) {
       return false;
@@ -521,6 +582,25 @@ export class ProbeAgent {
       return;
     }
 
+    // Skip API key requirement for Codex CLI (uses built-in access in Codex CLI)
+    if (this.clientApiProvider === 'codex' || process.env.USE_CODEX === 'true') {
+      // Codex CLI engine will be initialized lazily in getEngine()
+      // Set minimal defaults for compatibility
+      this.provider = null;
+      // Only set model if explicitly provided, otherwise let Codex use account default
+      this.model = modelName || null;
+      this.apiType = 'codex';
+      if (this.debug) {
+        console.log('[DEBUG] Codex CLI engine selected - will use built-in access if available');
+        if (this.model) {
+          console.log(`[DEBUG] Using model: ${this.model}`);
+        } else {
+          console.log('[DEBUG] Using Codex account default model');
+        }
+      }
+      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;
@@ -724,6 +804,60 @@ export class ProbeAgent {
       }
     }
 
+    // Check if we should use Codex engine
+    if (this.clientApiProvider === 'codex' || process.env.USE_CODEX === '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 Codex 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({
@@ -908,6 +1042,39 @@ export class ProbeAgent {
         this.clientApiProvider = null;
       }
     }
+
+    // Try Codex CLI engine if requested
+    if (this.clientApiProvider === 'codex' || process.env.USE_CODEX === 'true') {
+      try {
+        const { createCodexEngine } = await import('./engines/codex.js');
+
+        // For Codex CLI, use a cleaner system prompt without XML formatting
+        // since it has native MCP support for tools
+        const systemPrompt = this.customPrompt || this.getCodexNativeSystemPrompt();
+
+        this.engine = await createCodexEngine({
+          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
+          model: this.model  // Pass model name (e.g., gpt-4o, o3, etc.)
+        });
+        if (this.debug) {
+          console.log('[DEBUG] Using Codex CLI 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 Codex CLI 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);
@@ -917,6 +1084,35 @@ export class ProbeAgent {
     return this.engine;
   }
 
+  /**
+   * Get session information including thread ID for resumability
+   * @returns {Object} Session info with sessionId, threadId, messageCount
+   */
+  getSessionInfo() {
+    if (this.engine && this.engine.getSession) {
+      return this.engine.getSession();
+    }
+    return {
+      id: this.sessionId,
+      threadId: null,
+      messageCount: 0
+    };
+  }
+
+  /**
+   * Close the agent and clean up resources (e.g., MCP servers)
+   * @returns {Promise<void>}
+   */
+  async close() {
+    if (this.engine && this.engine.close) {
+      await this.engine.close();
+    }
+    if (this.mcpBridge) {
+      // Clean up MCP bridge if needed
+      this.mcpBridge = null;
+    }
+  }
+
   /**
    * Process assistant response content and detect/load image references
    * @param {string} content - The assistant's response content
@@ -1311,11 +1507,64 @@ export class ProbeAgent {
     let systemPrompt = '';
 
     // Add persona/role if configured
-    if (this.predefinedPrompt) {
-      const personaPrompt = getPromptByType(this.predefinedPrompt);
-      if (personaPrompt?.system) {
-        systemPrompt += personaPrompt.system + '\n\n';
-      }
+    if (this.customPrompt) {
+      systemPrompt += this.customPrompt + '\n\n';
+    } else if (this.promptType && predefinedPrompts[this.promptType]) {
+      systemPrompt += predefinedPrompts[this.promptType] + '\n\n';
+    } else {
+      // Use default code-explorer prompt
+      systemPrompt += predefinedPrompts['code-explorer'] + '\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 system prompt for Codex CLI (similar to Claude but optimized for Codex)
+   */
+  getCodexNativeSystemPrompt() {
+    let systemPrompt = '';
+
+    // Add persona/role if configured
+    if (this.customPrompt) {
+      systemPrompt += this.customPrompt + '\n\n';
+    } else if (this.promptType && predefinedPrompts[this.promptType]) {
+      systemPrompt += predefinedPrompts[this.promptType] + '\n\n';
+    } else {
+      // Use default code-explorer prompt
+      systemPrompt += predefinedPrompts['code-explorer'] + '\n\n';
     }
 
     // Add high-level instructions about when to use tools
@@ -1747,8 +1996,9 @@ 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
+      // Check if we're using CLI-based engines which handle their own agentic loop
       const isClaudeCode = this.clientApiProvider === 'claude-code' || process.env.USE_CLAUDE_CODE === 'true';
+      const isCodex = this.clientApiProvider === 'codex' || process.env.USE_CODEX === 'true';
 
       if (isClaudeCode) {
         // For Claude Code, bypass the tool loop entirely - it handles its own internal dialogue
@@ -1820,6 +2070,77 @@ When troubleshooting:
         }
       }
 
+      // Handle Codex engine (same pattern as Claude Code)
+      if (isCodex) {
+        // For Codex, bypass the tool loop entirely - it handles its own internal dialogue
+        if (this.debug) {
+          console.log(`[DEBUG] Using Codex engine - bypassing tool loop (black box mode)`);
+          console.log(`[DEBUG] Sending question directly to Codex: ${message.substring(0, 100)}...`);
+        }
+
+        // Send the message directly to Codex and collect the response
+        try {
+          const engine = await this.getEngine();
+          if (engine && engine.query) {
+            let assistantResponseContent = '';
+            let toolBatch = null;
+
+            // Query Codex 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 Codex`);
+                }
+              } 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 Codex 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 Codex
+            // 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] Codex error:', error);
+          }
+          throw error;
+        }
+      }
+
       if (this.debug) {
         console.log(`[DEBUG] Starting agentic flow for question: ${message.substring(0, 100)}...`);
         if (options.schema) {
@@ -1827,7 +2148,7 @@ When troubleshooting:
         }
       }
 
-      // Tool iteration loop (only for non-Claude Code engines)
+      // Tool iteration loop (only for non-CLI engines like Vercel/Anthropic/OpenAI)
       while (currentIteration < maxIterations && !completionAttempted) {
         currentIteration++;
         if (this.cancelled) throw new Error('Request was cancelled by the user');