@probelabs/probe 0.6.0-rc103 → 0.6.0-rc104

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@probelabs/probe",
-  "version": "0.6.0-rc103",
+  "version": "0.6.0-rc104",
   "description": "Node.js wrapper for the probe code search tool",
   "main": "src/index.js",
   "module": "src/index.js",
@@ -68,6 +68,7 @@
   "author": "probelabs",
   "license": "ISC",
   "dependencies": {
+    "@ai-sdk/amazon-bedrock": "^1.0.8",
     "@ai-sdk/anthropic": "^2.0.8",
     "@ai-sdk/google": "^2.0.14",
     "@ai-sdk/openai": "^2.0.10",

package/src/agent/ProbeAgent.js

@@ -2,6 +2,7 @@
 import { createAnthropic } from '@ai-sdk/anthropic';
 import { createOpenAI } from '@ai-sdk/openai';
 import { createGoogleGenerativeAI } from '@ai-sdk/google';
+import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock';
 import { streamText } from 'ai';
 import { randomUUID } from 'crypto';
 import { EventEmitter } from 'events';
@@ -91,8 +92,14 @@ export class ProbeAgent {
     this.maxResponseTokens = options.maxResponseTokens || parseInt(process.env.MAX_RESPONSE_TOKENS || '0', 10) || null;
     this.disableMermaidValidation = !!options.disableMermaidValidation;
 
-    // Search configuration
-    this.allowedFolders = options.path ? [options.path] : [process.cwd()];
+    // Search configuration - support both path (single) and allowedFolders (array)
+    if (options.allowedFolders && Array.isArray(options.allowedFolders)) {
+      this.allowedFolders = options.allowedFolders;
+    } else if (options.path) {
+      this.allowedFolders = [options.path];
+    } else {
+      this.allowedFolders = [process.cwd()];
+    }
 
     // API configuration
     this.clientApiProvider = options.provider || null;
@@ -186,12 +193,18 @@ export class ProbeAgent {
     const anthropicApiKey = process.env.ANTHROPIC_API_KEY;
     const openaiApiKey = process.env.OPENAI_API_KEY;
     const googleApiKey = process.env.GOOGLE_API_KEY;
+    const awsAccessKeyId = process.env.AWS_ACCESS_KEY_ID;
+    const awsSecretAccessKey = process.env.AWS_SECRET_ACCESS_KEY;
+    const awsRegion = process.env.AWS_REGION;
+    const awsSessionToken = process.env.AWS_SESSION_TOKEN;
+    const awsApiKey = process.env.AWS_BEDROCK_API_KEY;
 
     // Get custom API URLs if provided
     const llmBaseUrl = process.env.LLM_BASE_URL;
     const anthropicApiUrl = process.env.ANTHROPIC_API_URL || llmBaseUrl;
     const openaiApiUrl = process.env.OPENAI_API_URL || llmBaseUrl;
     const googleApiUrl = process.env.GOOGLE_API_URL || llmBaseUrl;
+    const awsBedrockBaseUrl = process.env.AWS_BEDROCK_BASE_URL || llmBaseUrl;
 
     // Get model override if provided
     const modelName = process.env.MODEL_NAME;
@@ -200,7 +213,12 @@ export class ProbeAgent {
     const forceProvider = this.clientApiProvider || (process.env.FORCE_PROVIDER ? process.env.FORCE_PROVIDER.toLowerCase() : null);
 
     if (this.debug) {
-      console.log(`[DEBUG] Available API keys: Anthropic=${!!anthropicApiKey}, OpenAI=${!!openaiApiKey}, Google=${!!googleApiKey}`);
+      const hasAwsCredentials = !!(awsAccessKeyId && awsSecretAccessKey && awsRegion);
+      const hasAwsApiKey = !!awsApiKey;
+      console.log(`[DEBUG] Available API keys: Anthropic=${!!anthropicApiKey}, OpenAI=${!!openaiApiKey}, Google=${!!googleApiKey}, AWS Bedrock=${hasAwsCredentials || hasAwsApiKey}`);
+      if (hasAwsCredentials) console.log(`[DEBUG] AWS credentials: AccessKey=${!!awsAccessKeyId}, SecretKey=${!!awsSecretAccessKey}, Region=${awsRegion}, SessionToken=${!!awsSessionToken}`);
+      if (hasAwsApiKey) console.log(`[DEBUG] AWS API Key provided`);
+      if (awsBedrockBaseUrl) console.log(`[DEBUG] AWS Bedrock base URL: ${awsBedrockBaseUrl}`);
       console.log(`[DEBUG] Force provider: ${forceProvider || '(not set)'}`);
       if (modelName) console.log(`[DEBUG] Model override: ${modelName}`);
     }
@@ -216,6 +234,9 @@ export class ProbeAgent {
       } else if (forceProvider === 'google' && googleApiKey) {
         this.initializeGoogleModel(googleApiKey, googleApiUrl, modelName);
         return;
+      } else if (forceProvider === 'bedrock' && ((awsAccessKeyId && awsSecretAccessKey && awsRegion) || awsApiKey)) {
+        this.initializeBedrockModel(awsAccessKeyId, awsSecretAccessKey, awsRegion, awsSessionToken, awsApiKey, awsBedrockBaseUrl, modelName);
+        return;
       }
       console.warn(`WARNING: Forced provider "${forceProvider}" selected but required API key is missing or invalid! Falling back to auto-detection.`);
     }
@@ -227,8 +248,10 @@ export class ProbeAgent {
       this.initializeOpenAIModel(openaiApiKey, openaiApiUrl, modelName);
     } else if (googleApiKey) {
       this.initializeGoogleModel(googleApiKey, googleApiUrl, modelName);
+    } else if ((awsAccessKeyId && awsSecretAccessKey && awsRegion) || awsApiKey) {
+      this.initializeBedrockModel(awsAccessKeyId, awsSecretAccessKey, awsRegion, awsSessionToken, awsApiKey, awsBedrockBaseUrl, modelName);
     } else {
-      throw new Error('No API key provided. Please set ANTHROPIC_API_KEY, OPENAI_API_KEY, or GOOGLE_API_KEY environment variable.');
+      throw new Error('No API key provided. Please set ANTHROPIC_API_KEY, OPENAI_API_KEY, GOOGLE_API_KEY, AWS credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION), or AWS_BEDROCK_API_KEY environment variables.');
     }
   }
 
@@ -281,6 +304,46 @@ export class ProbeAgent {
     }
   }
 
+  /**
+   * Initialize AWS Bedrock model
+   */
+  initializeBedrockModel(accessKeyId, secretAccessKey, region, sessionToken, apiKey, baseURL, modelName) {
+    // Build configuration object, only including defined values
+    const config = {};
+    
+    // Authentication - prefer API key if provided, otherwise use AWS credentials
+    if (apiKey) {
+      config.apiKey = apiKey;
+    } else if (accessKeyId && secretAccessKey) {
+      config.accessKeyId = accessKeyId;
+      config.secretAccessKey = secretAccessKey;
+      if (sessionToken) {
+        config.sessionToken = sessionToken;
+      }
+    }
+    
+    // Region is required for AWS credentials but optional for API key
+    if (region) {
+      config.region = region;
+    }
+    
+    // Optional base URL
+    if (baseURL) {
+      config.baseURL = baseURL;
+    }
+    
+    this.provider = createAmazonBedrock(config);
+    this.model = modelName || 'anthropic.claude-sonnet-4-20250514-v1:0';
+    this.apiType = 'bedrock';
+
+    if (this.debug) {
+      const authMethod = apiKey ? 'API Key' : 'AWS Credentials';
+      const regionInfo = region ? ` (Region: ${region})` : '';
+      const baseUrlInfo = baseURL ? ` (Base URL: ${baseURL})` : '';
+      console.log(`Using AWS Bedrock API with model: ${this.model}${regionInfo} [Auth: ${authMethod}]${baseUrlInfo}`);
+    }
+  }
+
   /**
    * Process assistant response content and detect/load image references
    * @param {string} content - The assistant's response content
@@ -1060,8 +1123,12 @@ When troubleshooting:
             } else if (this.toolImplementations[toolName]) {
               // Execute native tool
               try {
-                // Add sessionId to params for tool execution
-                const toolParams = { ...params, sessionId: this.sessionId };
+                // Add sessionId and workingDirectory to params for tool execution
+                const toolParams = { 
+                  ...params, 
+                  sessionId: this.sessionId,
+                  workingDirectory: (this.allowedFolders && this.allowedFolders[0]) || process.cwd()
+                };
                 
                 // Emit tool start event
                 this.events.emit('toolCall', {
@@ -1214,8 +1281,10 @@ Remember: Use proper XML format with BOTH opening and closing tags:
 <parameter>value</parameter>
 </tool_name>
 
-Or for quick completion if your previous response was already correct:
-<attempt_complete>`;
+Or for quick completion if your previous response was already correct and complete:
+<attempt_complete>
+
+IMPORTANT: When using <attempt_complete>, this must be the ONLY content in your response. No additional text, explanations, or other content should be included. This tag signals to reuse your previous response as the final answer.`;
           }
 
           currentMessages.push({
@@ -1733,4 +1802,4 @@ Convert your previous response content into actual JSON data that follows this s
       console.log(`[DEBUG] Agent cancelled for session ${this.sessionId}`);
     }
   }
-}
+}

package/src/agent/index.js

@@ -112,6 +112,7 @@ function parseArgs() {
     acp: false,
     question: null,
     path: null,
+    allowedFolders: null,
     prompt: null,
     systemPrompt: null,
     schema: null,
@@ -145,6 +146,8 @@ function parseArgs() {
       config.allowEdit = true;
     } else if (arg === '--path' && i + 1 < args.length) {
       config.path = args[++i];
+    } else if (arg === '--allowed-folders' && i + 1 < args.length) {
+      config.allowedFolders = args[++i].split(',').map(dir => dir.trim());
     } else if (arg === '--prompt' && i + 1 < args.length) {
       config.prompt = args[++i];
     } else if (arg === '--system-prompt' && i + 1 < args.length) {
@@ -201,6 +204,7 @@ Usage:
 
 Options:
   --path <dir>                     Search directory (default: current)
+  --allowed-folders <dirs>         Comma-separated list of allowed directories for file operations
   --prompt <type>                  Persona: code-explorer, engineer, code-review, support, architect
   --system-prompt <text|file>      Custom system prompt (text or file path)
   --schema <schema|file>           Output schema (JSON, XML, any format - text or file path)
@@ -296,17 +300,22 @@ class ProbeAgentMcpServer {
       tools: [
         {
           name: 'search_code',
-          description: "Search code and answer questions about the codebase using an AI agent. This tool provides intelligent responses based on code analysis.",
+          description: "AI agent that answers free-form questions about codebases. Ask detailed questions in natural language.",
           inputSchema: {
             type: 'object',
             properties: {
               query: {
                 type: 'string',
-                description: 'The question or request about the codebase.',
+                description: 'A detailed, free-form question about the codebase in natural language. Be specific and descriptive. Example: "How does the authentication system work and where is user session management implemented?"',
               },
               path: {
                 type: 'string',
-                description: 'Optional path to the directory to search in. Defaults to current directory.',
+                description: 'Absolute path to the directory to search in (e.g., "/Users/username/projects/myproject").',
+              },
+              allowed_folders: {
+                type: 'array',
+                items: { type: 'string' },
+                description: 'Optional list of allowed directories for file operations. Defaults to current directory if not specified.',
               },
               prompt: {
                 type: 'string',
@@ -315,34 +324,6 @@ class ProbeAgentMcpServer {
               system_prompt: {
                 type: 'string',
                 description: 'Optional custom system prompt (text or file path).',
-              },
-              provider: {
-                type: 'string',
-                description: 'Optional AI provider to force: anthropic, openai, google.',
-              },
-              model: {
-                type: 'string',
-                description: 'Optional model name override.',
-              },
-              allow_edit: {
-                type: 'boolean',
-                description: 'Enable code modification capabilities.',
-              },
-              max_iterations: {
-                type: 'number',
-                description: 'Maximum number of tool iterations (default: 30).',
-              },
-              max_response_tokens: {
-                type: 'number',
-                description: 'Maximum tokens for AI response (overrides model defaults).',
-              },
-              schema: {
-                type: 'string',
-                description: 'Optional output schema (JSON, XML, or any format - text or file path).',
-              },
-              no_mermaid_validation: {
-                type: 'boolean',
-                description: 'Disable automatic mermaid diagram validation and fixing.',
               }
             },
             required: ['query']
@@ -404,7 +385,7 @@ class ProbeAgentMcpServer {
           
           // Create agent with configuration
           const agentConfig = {
-            path: args.path || process.cwd(),
+            path: args.path || (args.allowed_folders && args.allowed_folders[0]) || process.cwd(),
             promptType: args.prompt || 'code-explorer',
             customPrompt: systemPrompt,
             provider: args.provider,
@@ -643,6 +624,7 @@ async function main() {
     // Create and configure agent
     const agentConfig = {
       path: config.path,
+      allowedFolders: config.allowedFolders,
       promptType: config.prompt,
       customPrompt: systemPrompt,
       allowEdit: config.allowEdit,

package/src/agent/probeTool.js

@@ -207,14 +207,24 @@ export function createWrappedTools(baseTools) {
 // Simple file listing tool
 export const listFilesTool = {
   execute: async (params) => {
-    const { directory = '.' } = params;
+    const { directory = '.', workingDirectory } = params;
+    
+    // Use the provided working directory, or fall back to process.cwd()
+    const baseCwd = workingDirectory || process.cwd();
+    
+    // Security: Validate path to prevent traversal attacks
+    const secureBaseDir = path.resolve(baseCwd);
+    const targetDir = path.resolve(secureBaseDir, directory);
+    if (!targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
+      throw new Error('Path traversal attempt detected. Access denied.');
+    }
     
     try {
       const files = await listFilesByLevel({
-        directory,
+        directory: targetDir,
         maxFiles: 100,
         respectGitignore: !process.env.PROBE_NO_GITIGNORE || process.env.PROBE_NO_GITIGNORE === '',
-        cwd: process.cwd()
+        cwd: secureBaseDir
       });
       
       return files;
@@ -227,15 +237,23 @@ export const listFilesTool = {
 // Simple file search tool
 export const searchFilesTool = {
   execute: async (params) => {
-    const { pattern, directory = '.', recursive = true } = params;
+    const { pattern, directory = '.', recursive = true, workingDirectory } = params;
     
     if (!pattern) {
       throw new Error('Pattern is required for file search');
     }
     
+    // Security: Validate path to prevent traversal attacks
+    const baseCwd = workingDirectory || process.cwd();
+    const secureBaseDir = path.resolve(baseCwd);
+    const targetDir = path.resolve(secureBaseDir, directory);
+    if (!targetDir.startsWith(secureBaseDir + path.sep) && targetDir !== secureBaseDir) {
+      throw new Error('Path traversal attempt detected. Access denied.');
+    }
+    
     try {
       const options = {
-        cwd: directory,
+        cwd: targetDir,
         ignore: ['node_modules/**', '.git/**'],
         absolute: false
       };

package/src/agent/tools.js

@@ -141,12 +141,45 @@ export function parseXmlToolCallWithThinking(xmlString, validTools) {
   // Remove thinking tags and their content from the XML string
   let cleanedXmlString = xmlString.replace(/<thinking>[\s\S]*?<\/thinking>/g, '').trim();
 
-  // Check for attempt_complete shorthand (single tag with no closing tag and no parameters)
-  const attemptCompleteMatch = cleanedXmlString.match(/^<attempt_complete>\s*$/);
-  if (attemptCompleteMatch) {
-    // Convert shorthand to full attempt_completion format with special marker
+  // Enhanced recovery logic for attempt_complete shorthand
+  // Check for various forms of attempt_complete tags:
+  // 1. Perfect shorthand: <attempt_complete>
+  // 2. Self-closing: <attempt_complete/>  
+  // 3. Empty with closing tag: <attempt_complete></attempt_complete>
+  // 4. Incomplete: <attempt_complete (missing closing bracket)
+  // 5. With trailing content that should be ignored
+  const attemptCompletePatterns = [
+    // Standard shorthand with optional whitespace
+    /^<attempt_complete>\s*$/,
+    // Empty with proper closing tag (common case from the logs)
+    /^<attempt_complete>\s*<\/attempt_complete>\s*$/,
+    // Self-closing variant
+    /^<attempt_complete\s*\/>\s*$/,
+    // Incomplete opening tag (missing closing bracket)
+    /^<attempt_complete\s*$/,
+    // With trailing content (extract just the tag part) - must come after empty tag pattern
+    /^<attempt_complete>(.*)$/s,
+    // Self-closing with trailing content
+    /^<attempt_complete\s*\/>(.*)$/s
+  ];
+
+  for (const pattern of attemptCompletePatterns) {
+    const match = cleanedXmlString.match(pattern);
+    if (match) {
+      // Convert any form of attempt_complete to the standard format
+      return {
+        toolName: 'attempt_completion',
+        params: { result: '__PREVIOUS_RESPONSE__' }
+      };
+    }
+  }
+
+  // Additional recovery: check if the string contains attempt_complete anywhere
+  // and treat the entire response as a completion signal if no other tool tags are found
+  if (cleanedXmlString.includes('<attempt_complete') && !hasOtherToolTags(cleanedXmlString, validTools)) {
+    // This handles malformed cases where attempt_complete appears but is broken
     return {
-      toolName: 'attempt_completion',
+      toolName: 'attempt_completion', 
       params: { result: '__PREVIOUS_RESPONSE__' }
     };
   }
@@ -160,4 +193,23 @@ export function parseXmlToolCallWithThinking(xmlString, validTools) {
   }
 
   return parsedTool;
+}
+
+/**
+ * Helper function to check if the XML string contains other tool tags
+ * @param {string} xmlString - The XML string to check
+ * @param {string[]} validTools - List of valid tool names
+ * @returns {boolean} - True if other tool tags are found
+ */
+function hasOtherToolTags(xmlString, validTools = []) {
+  const defaultTools = ['search', 'query', 'extract', 'listFiles', 'searchFiles', 'implement', 'attempt_completion'];
+  const toolsToCheck = validTools.length > 0 ? validTools : defaultTools;
+  
+  // Check for any tool tags other than attempt_complete variants
+  for (const tool of toolsToCheck) {
+    if (tool !== 'attempt_completion' && xmlString.includes(`<${tool}`)) {
+      return true;
+    }
+  }
+  return false;
 }

package/src/utils.js

@@ -43,7 +43,7 @@ export async function getBinaryPath(options = {}) {
 	
 	// Check bin directory
 	const isWindows = process.platform === 'win32';
-	const binaryName = isWindows ? 'probe.exe' : 'probe';
+	const binaryName = isWindows ? 'probe.exe' : 'probe-binary';
 	const binaryPath = path.join(binDir, binaryName);
 
 	if (fs.existsSync(binaryPath) && !forceDownload) {