@probelabs/probe 0.6.0-rc302 → 0.6.0-rc304

package/build/agent/FallbackManager.js

@@ -8,10 +8,7 @@
  * - Custom fallback chains with full configuration
  */
 
-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 { createProviderInstance, DEFAULT_MODELS as SHARED_DEFAULT_MODELS } from '../utils/provider.js';
 
 /**
  * Fallback strategies
@@ -40,12 +37,7 @@ export const FALLBACK_STRATEGIES = {
 /**
  * Default model mappings for each provider
  */
-const DEFAULT_MODELS = {
-  anthropic: 'claude-sonnet-4-6',
-  openai: 'gpt-5.2',
-  google: 'gemini-2.5-flash',
-  bedrock: 'anthropic.claude-sonnet-4-6'
-};
+const DEFAULT_MODELS = SHARED_DEFAULT_MODELS;
 
 /**
  * FallbackManager class for handling provider and model fallback
@@ -138,53 +130,7 @@ export class FallbackManager {
    */
   _createProviderInstance(config) {
     try {
-      switch (config.provider) {
-        case 'anthropic':
-          return createAnthropic({
-            apiKey: config.apiKey,
-            ...(config.baseURL && { baseURL: config.baseURL })
-          });
-
-        case 'openai':
-          return createOpenAI({
-            compatibility: 'strict',
-            apiKey: config.apiKey,
-            ...(config.baseURL && { baseURL: config.baseURL })
-          });
-
-        case 'google':
-          return createGoogleGenerativeAI({
-            apiKey: config.apiKey,
-            ...(config.baseURL && { baseURL: config.baseURL })
-          });
-
-        case 'bedrock': {
-          const bedrockConfig = {};
-
-          if (config.apiKey) {
-            bedrockConfig.apiKey = config.apiKey;
-          } else if (config.accessKeyId && config.secretAccessKey) {
-            bedrockConfig.accessKeyId = config.accessKeyId;
-            bedrockConfig.secretAccessKey = config.secretAccessKey;
-            if (config.sessionToken) {
-              bedrockConfig.sessionToken = config.sessionToken;
-            }
-          }
-
-          if (config.region) {
-            bedrockConfig.region = config.region;
-          }
-
-          if (config.baseURL) {
-            bedrockConfig.baseURL = config.baseURL;
-          }
-
-          return createAmazonBedrock(bedrockConfig);
-        }
-
-        default:
-          throw new Error(`FallbackManager: Unknown provider "${config.provider}"`);
-      }
+      return createProviderInstance(config);
     } catch (error) {
       // Re-throw with more context
       const providerName = this._getProviderDisplayName(config);

package/build/agent/ProbeAgent.js

@@ -27,10 +27,7 @@ export const ENGINE_ACTIVITY_TIMEOUT_MIN = 5000;
  */
 export const ENGINE_ACTIVITY_TIMEOUT_MAX = 600000;
 
-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 { createProviderInstance, DEFAULT_MODELS } from '../utils/provider.js';
 import { streamText, generateText, tool, stepCountIs, jsonSchema, Output } from 'ai';
 import { randomUUID } from 'crypto';
 import { EventEmitter } from 'events';
@@ -1673,13 +1670,10 @@ export class ProbeAgent {
    * Initialize Anthropic model
    */
   initializeAnthropicModel(apiKey, apiUrl, modelName) {
-    this.provider = createAnthropic({
-      apiKey: apiKey,
-      ...(apiUrl && { baseURL: apiUrl }),
-    });
-    this.model = modelName || 'claude-sonnet-4-6';
+    this.provider = createProviderInstance({ provider: 'anthropic', apiKey, ...(apiUrl && { baseURL: apiUrl }) });
+    this.model = modelName || DEFAULT_MODELS.anthropic;
     this.apiType = 'anthropic';
-    
+
     if (this.debug) {
       console.log(`Using Anthropic API with model: ${this.model}${apiUrl ? ` (URL: ${apiUrl})` : ''}`);
     }
@@ -1689,14 +1683,10 @@ export class ProbeAgent {
    * Initialize OpenAI model
    */
   initializeOpenAIModel(apiKey, apiUrl, modelName) {
-    this.provider = createOpenAI({
-      compatibility: 'strict',
-      apiKey: apiKey,
-      ...(apiUrl && { baseURL: apiUrl }),
-    });
-    this.model = modelName || 'gpt-5.2';
+    this.provider = createProviderInstance({ provider: 'openai', apiKey, ...(apiUrl && { baseURL: apiUrl }) });
+    this.model = modelName || DEFAULT_MODELS.openai;
     this.apiType = 'openai';
-    
+
     if (this.debug) {
       console.log(`Using OpenAI API with model: ${this.model}${apiUrl ? ` (URL: ${apiUrl})` : ''}`);
     }
@@ -1706,10 +1696,7 @@ export class ProbeAgent {
    * Initialize Google model
    */
   initializeGoogleModel(apiKey, apiUrl, modelName) {
-    this.provider = createGoogleGenerativeAI({
-      apiKey: apiKey,
-      ...(apiUrl && { baseURL: apiUrl }),
-    });
+    this.provider = createProviderInstance({ provider: 'google', apiKey, ...(apiUrl && { baseURL: apiUrl }) });
     this.model = modelName || 'gemini-2.5-pro';
     this.apiType = 'google';
 
@@ -2245,32 +2232,10 @@ 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-6';
+    this.provider = createProviderInstance({
+      provider: 'bedrock', apiKey, accessKeyId, secretAccessKey, sessionToken, region, baseURL
+    });
+    this.model = modelName || DEFAULT_MODELS.bedrock;
     this.apiType = 'bedrock';
 
     if (this.debug) {
@@ -3012,7 +2977,7 @@ export class ProbeAgent {
 
     // Add high-level instructions about when to use tools
     const searchToolDesc1 = this.searchDelegate
-      ? '- search: Ask natural language questions to find code (e.g., "How does authentication work?"). A subagent handles keyword searches and returns extracted code blocks. Do NOT formulate keyword queries — just ask questions.'
+      ? '- search: Ask natural language questions to find code locations (e.g., "How does authentication work?"). Returns structured JSON with file locations grouped by relevance. Use extract() on the returned files to read the actual code. Do NOT formulate keyword queries — just ask questions.'
       : '- search: Find code patterns using keyword queries with Elasticsearch syntax. Handles stemming and case variations automatically — do NOT try manual keyword variations.';
     systemPrompt += `You have access to powerful code search and analysis tools through MCP:
 ${searchToolDesc1}
@@ -3025,10 +2990,10 @@ ${searchToolDesc1}
     }
 
     const searchGuidance1 = this.searchDelegate
-      ? '1. Start with search — ask a question about what you want to understand. It returns extracted code blocks directly.'
+      ? '1. Start with search — ask a question about what you want to understand. It returns file locations grouped by relevance (JSON with confidence and groups).'
       : '1. Start with search to find relevant code patterns. One search per concept is usually enough — probe handles stemming and case variations.';
     const extractGuidance1 = this.searchDelegate
-      ? '2. Use extract only if you need more context or a full file'
+      ? '2. Use extract on the file locations returned by search to read the actual code. Each group has a "reason" explaining why those files matter.'
       : '2. Use extract to get detailed context when needed';
 
     systemPrompt += `\n
@@ -3078,7 +3043,7 @@ ${extractGuidance1}
 
     // Add high-level instructions about when to use tools
     const searchToolDesc2 = this.searchDelegate
-      ? '- search: Ask natural language questions to find code (e.g., "How does authentication work?"). A subagent handles keyword searches and returns extracted code blocks. Do NOT formulate keyword queries — just ask questions.'
+      ? '- search: Ask natural language questions to find code locations (e.g., "How does authentication work?"). Returns structured JSON with file locations grouped by relevance. Use extract() on the returned files to read the actual code. Do NOT formulate keyword queries — just ask questions.'
       : '- search: Find code patterns using keyword queries with Elasticsearch syntax. Handles stemming and case variations automatically — do NOT try manual keyword variations.';
     systemPrompt += `You have access to powerful code search and analysis tools through MCP:
 ${searchToolDesc2}
@@ -3091,10 +3056,10 @@ ${searchToolDesc2}
     }
 
     const searchGuidance2 = this.searchDelegate
-      ? '1. Start with search — ask a question about what you want to understand. It returns extracted code blocks directly.'
+      ? '1. Start with search — ask a question about what you want to understand. It returns file locations grouped by relevance (JSON with confidence and groups).'
       : '1. Start with search to find relevant code patterns. One search per concept is usually enough — probe handles stemming and case variations.';
     const extractGuidance2 = this.searchDelegate
-      ? '2. Use extract only if you need more context or a full file'
+      ? '2. Use extract on the file locations returned by search to read the actual code. Each group has a "reason" explaining why those files matter.'
       : '2. Use extract to get detailed context when needed';
 
     systemPrompt += `\n
@@ -3160,10 +3125,10 @@ ${extractGuidance2}
 Follow these instructions carefully:
 1. Analyze the user's request.
 2. Use the available tools step-by-step to fulfill the request.
-3. You MUST use the search tool before answering ANY code-related question. NEVER answer from memory or general knowledge — your answers must be grounded in actual code found via search/extract.${this.searchDelegate ? ' Ask natural language questions — the search subagent handles keyword formulation and returns extracted code blocks. Use extract only to expand context or read full files.' : ' Search handles stemming and case variations automatically — do NOT try keyword variations manually. Read full files only if really necessary.'}
+3. You MUST use the search tool before answering ANY code-related question. NEVER answer from memory or general knowledge — your answers must be grounded in actual code found via search/extract.${this.searchDelegate ? ' Ask natural language questions — the search subagent handles keyword formulation and returns file locations grouped by relevance. Then use extract() on those locations to read the actual code.' : ' Search handles stemming and case variations automatically — do NOT try keyword variations manually. Read full files only if really necessary.'}
 4. Ensure to get really deep and understand the full picture before answering. Follow call chains — if function A calls B, search for B too. Look for related subsystems (e.g., if asked about rate limiting, also check for quota, throttling, smoothing).
 5. Once the task is fully completed, provide your final answer directly as text. Always cite specific files and line numbers as evidence. Do NOT output planning or thinking text — go straight to the answer.
-6. ${this.searchDelegate ? 'Ask clear, specific questions when searching. Each search should target a distinct concept or question.' : 'Prefer concise and focused search queries. Use specific keywords and phrases to narrow down results.'}
+6. ${this.searchDelegate ? 'Ask clear, specific questions when searching. Each search should target a distinct concept or question. NEVER re-search the same concept with different phrasing — if you already searched for "wrapToolWithEmitter", do NOT search again for "definition of wrapToolWithEmitter" or "how wrapToolWithEmitter works". Use extract() on the files already found instead. Limit yourself to one search per distinct concept. When formulating queries, describe WHAT you are looking for, not WHERE — the search agent will search the full codebase. Do NOT include file names or class names in the query unless that IS the concept (e.g., say "search dedup logic" not "search dedup ProbeAgent").' : 'Prefer concise and focused search queries. Use specific keywords and phrases to narrow down results.'}
 7. NEVER use bash for code exploration (no grep, cat, find, head, tail, awk, sed) — always use search and extract tools instead. Bash is only for system operations like building, running tests, or git commands.${this.allowEdit ? `
 7. When modifying files, choose the appropriate tool:
     - Use 'edit' for all code modifications:
@@ -4088,9 +4053,16 @@ or
                 const searchSummary = searchesTried.length > 0
                   ? `\nSearches attempted: ${searchesTried.join(', ')}`
                   : '';
+
+                // For code-searcher subagents: instruct to output structured JSON even on partial results
+                const isCodeSearcher = this.promptType === 'code-searcher';
+                const lastIterMessage = isCodeSearcher
+                  ? `⚠️ LAST ITERATION — you are out of tool calls. Output your JSON response NOW with whatever files you have verified so far. Set confidence to "low" if your search was incomplete. Include the "searches" array listing all search queries you made with their paths and outcomes.${searchSummary}`
+                  : `⚠️ LAST ITERATION — you are out of tool calls. Provide your BEST answer NOW with the information gathered so far. If you could not find what was requested, explain exactly what you searched for and why it did not work, so the caller can try a different approach.${searchSummary}`;
+
                 return {
                   toolChoice: 'none',
-                  userMessage: `⚠️ LAST ITERATION — you are out of tool calls. Provide your BEST answer NOW with the information gathered so far. If you could not find what was requested, explain exactly what you searched for and why it did not work, so the caller can try a different approach.${searchSummary}`
+                  userMessage: lastIterMessage
                 };
               }
 
@@ -4766,13 +4738,16 @@ Double-check your response based on the criteria above. If everything looks good
         if (!finalResult || finalResult === DEFAULT_MAX_ITER_MSG) {
           try {
             const searchQueries = [];
+            const searchDetails = [];
             const toolCounts = {};
             for (const tc of _toolCallLog) {
               toolCounts[tc.name] = (toolCounts[tc.name] || 0) + 1;
               if (tc.name === 'search') {
                 const q = tc.args.query || '';
+                const p = tc.args.path || '.';
                 const exact = tc.args.exact ? ' (exact)' : '';
                 searchQueries.push(`"${q}"${exact}`);
+                searchDetails.push({ query: q, path: p, had_results: false });
               }
             }
             const toolBreakdown = Object.entries(toolCounts)
@@ -4780,13 +4755,24 @@ Double-check your response based on the criteria above. If everything looks good
               .join(', ');
             const uniqueSearches = [...new Set(searchQueries)];
 
-            let summary = `I was unable to complete your request after ${currentIteration} tool iterations.\n\n`;
-            summary += `Tool calls made: ${toolBreakdown || 'none'}\n`;
-            if (uniqueSearches.length > 0) {
-              summary += `Search queries tried: ${uniqueSearches.join(', ')}\n`;
+            // For code-searcher subagents: produce structured JSON so the parent
+            // can still use partial results instead of getting a plain error string.
+            if (this.promptType === 'code-searcher') {
+              finalResult = JSON.stringify({
+                confidence: 'low',
+                reason: 'Search incomplete — iteration limit reached',
+                groups: [],
+                searches: searchDetails
+              });
+            } else {
+              let summary = `I was unable to complete your request after ${currentIteration} tool iterations.\n\n`;
+              summary += `Tool calls made: ${toolBreakdown || 'none'}\n`;
+              if (uniqueSearches.length > 0) {
+                summary += `Search queries tried: ${uniqueSearches.join(', ')}\n`;
+              }
+              summary += `\nThe search approach may be fundamentally wrong for this query. Consider: using exact=true for literal string matching, using bash/grep for pattern-based file searches, or trying a completely different strategy instead of repeating similar searches.`;
+              finalResult = summary;
             }
-            summary += `\nThe search approach may be fundamentally wrong for this query. Consider: using exact=true for literal string matching, using bash/grep for pattern-based file searches, or trying a completely different strategy instead of repeating similar searches.`;
-            finalResult = summary;
           } catch {
             finalResult = DEFAULT_MAX_ITER_MSG;
           }

package/build/delegate.js

@@ -20,13 +20,14 @@ import { ProbeAgent } from './agent/ProbeAgent.js';
  */
 class DelegationManager {
 	constructor(options = {}) {
+		const parseSafe = (val, fallback) => { const n = parseInt(val, 10); return Number.isNaN(n) ? fallback : n; };
 		this.maxConcurrent = options.maxConcurrent
-			?? parseInt(process.env.MAX_CONCURRENT_DELEGATIONS || '3', 10);
+			?? parseSafe(process.env.MAX_CONCURRENT_DELEGATIONS, 3);
 		this.maxPerSession = options.maxPerSession
-			?? parseInt(process.env.MAX_DELEGATIONS_PER_SESSION || '10', 10);
+			?? parseSafe(process.env.MAX_DELEGATIONS_PER_SESSION, 10);
 		// Default queue timeout: 60 seconds. Set DELEGATION_QUEUE_TIMEOUT=0 to disable.
 		this.defaultQueueTimeout = options.queueTimeout
-			?? parseInt(process.env.DELEGATION_QUEUE_TIMEOUT || '60000', 10);
+			?? parseSafe(process.env.DELEGATION_QUEUE_TIMEOUT, 60000);
 
 		// Track delegations per session with timestamp for potential TTL cleanup
 		// Map<string, { count: number, lastUpdated: number }>
@@ -545,6 +546,9 @@ export async function delegate({
 		// Phase 2: Hard cancel after deadline if subagent hasn't finished
 		let parentAbortHandler;
 		let parentAbortHardCancelId = null;
+		// Track whether the race has settled so late abort signals don't create
+		// unhandled promise rejections.
+		let raceSettled = false;
 		const parentAbortPromise = new Promise((_, reject) => {
 			if (parentAbortSignal) {
 				if (parentAbortSignal.aborted) {
@@ -553,6 +557,8 @@ export async function delegate({
 					return;
 				}
 				parentAbortHandler = () => {
+					// If the race already settled, the answer won — don't reject.
+					if (raceSettled) return;
 					// Phase 1: graceful wind-down — let subagent finish its current step
 					subagent.triggerGracefulWindDown();
 					if (debug) {
@@ -567,6 +573,7 @@ export async function delegate({
 					}
 					// Phase 2: hard cancel after 30s if subagent hasn't finished
 					parentAbortHardCancelId = setTimeout(() => {
+						if (raceSettled) return;
 						if (debug) {
 							console.error(`[DELEGATE] Graceful wind-down deadline expired — hard cancelling subagent ${sessionId}`);
 						}
@@ -594,6 +601,8 @@ export async function delegate({
 		try {
 			response = await Promise.race(racers);
 		} finally {
+			// Mark race as settled so late abort signals are ignored
+			raceSettled = true;
 			// Clean up parent abort listener and hard cancel timer to prevent memory leaks
 			if (parentAbortHandler && parentAbortSignal) {
 				parentAbortSignal.removeEventListener('abort', parentAbortHandler);
@@ -650,10 +659,12 @@ export async function delegate({
 			});
 
 			if (delegationSpan) {
+				const { truncateForSpan } = await import('./agent/simpleTelemetry.js');
 				delegationSpan.setAttributes({
 					'delegation.result.success': true,
 					'delegation.result.response_length': response.length,
-					'delegation.result.duration_ms': duration
+					'delegation.result.duration_ms': duration,
+					'delegation.result': truncateForSpan(response, 4096)
 				});
 				delegationSpan.setStatus({ code: 1 }); // OK
 				delegationSpan.end();

package/build/tools/common.js

@@ -7,6 +7,11 @@ import { z } from 'zod';
 import { resolve, isAbsolute } from 'path';
 
 // Common schemas for tool parameters (used for internal execution after XML parsing)
+export const searchDelegateSchema = z.object({
+	query: z.string().describe('Natural language question about the code (e.g., "How does authentication work?", "Where is the rate limiting middleware?"). Do NOT use keyword syntax — just describe what you are looking for in plain English. A subagent will handle keyword searches for you.'),
+	path: z.string().optional().default('.').describe('Path to search in.'),
+});
+
 export const searchSchema = z.object({
 	query: z.string().describe('Search query — natural language questions or Elasticsearch-style keywords both work. For keywords: use quotes for exact phrases, AND/OR for boolean logic, - for negation. Probe handles stemming and camelCase/snake_case splitting automatically, so do NOT try case or style variations of the same keyword.'),
 	path: z.string().optional().default('.').describe('Path to search in. For dependencies use "go:github.com/owner/repo", "js:package_name", or "rust:cargo_name" etc.'),
@@ -100,7 +105,17 @@ export const attemptCompletionSchema = {
 // Tool descriptions (used by Vercel tool() definitions)
 
 export const searchDescription = 'Search code in the repository. Free-form questions are accepted, but Elasticsearch-style keyword queries work best. Use this tool first for any code-related questions. NOTE: By default, search handles stemming, case-insensitive matching, and camelCase/snake_case splitting automatically — do NOT manually try keyword variations like "getAllUsers" then "get_all_users" then "GetAllUsers". One search covers all variations.';
-export const searchDelegateDescription = 'Search code in the repository by asking a question. Accepts natural language questions (e.g., "How does authentication work?", "Where is the user validation logic?"). A specialized subagent breaks down your question into targeted keyword searches and returns extracted code blocks. Do NOT formulate keyword queries yourself — just ask the question naturally.';
+export const searchDelegateDescription = `Find where relevant code is located by asking a natural language question. A subagent searches the codebase and returns file locations grouped by relevance, with reasons explaining why each group matters. Use extract() to read the actual code from the returned locations.
+
+Returns JSON: { "confidence": "high|medium|low", "groups": [{ "reason": "why these files matter", "files": ["path#Symbol", ...] }] }
+
+IMPORTANT — each call spawns a subagent (expensive, takes minutes). Be deliberate:
+- Ask plain English questions about WHERE code is, NOT keyword queries. Good: "How are user sessions extracted from cookies?" Bad: "ctxGetSession OR GetSession"
+- Each call should explore a DIFFERENT ANGLE of the problem. Don't rephrase — reframe:
+  Good: 1) "How are sessions extracted from HTTP requests?" 2) "What middleware runs before route handlers?" 3) "How is the session cookie parsed and validated?"
+  Bad: 1) "How does session extraction work?" 2) "Where is the session extracted?" 3) "Find session extraction code" ← same question reworded
+- If a search returned no useful results, ask about a DIFFERENT part of the system. Think: what upstream/downstream component touches this?
+- After getting results, use extract() to read the files you need — search only locates, extract reads.`;
 export const queryDescription = 'Search code using ast-grep structural pattern matching. Use this tool to find specific code structures like functions, classes, or methods.';
 export const extractDescription = 'Extract code blocks from files based on file paths and optional line numbers. Use this tool to see complete context after finding relevant files. Line numbers from output can be used with edit start_line/end_line for precise editing.';
 export const delegateDescription = 'Automatically delegate big distinct tasks to specialized probe subagents within the agentic loop. Used by AI agents to break down complex requests into focused, parallel tasks.';