@probelabs/probe 0.6.0-rc278 → 0.6.0-rc280

package/build/agent/ProbeAgent.js

@@ -125,6 +125,27 @@ const MAX_HISTORY_MESSAGES = 100;
 // Maximum image file size (20MB) to prevent OOM attacks
 const MAX_IMAGE_FILE_SIZE = 20 * 1024 * 1024;
 
+/**
+ * Truncate a string for debug logging, showing first and last portion.
+ */
+export function debugTruncate(s, limit = 200) {
+  if (s.length <= limit) return s;
+  const half = Math.floor(limit / 2);
+  return s.substring(0, half) + ` ... [${s.length} chars] ... ` + s.substring(s.length - half);
+}
+
+/**
+ * Log tool results details for debug output.
+ */
+export function debugLogToolResults(toolResults) {
+  if (!toolResults || toolResults.length === 0) return;
+  for (const tr of toolResults) {
+    const argsStr = JSON.stringify(tr.args || {});
+    const resultStr = typeof tr.result === 'string' ? tr.result : JSON.stringify(tr.result || '');
+    console.log(`[DEBUG]   tool: ${tr.toolName} | args: ${debugTruncate(argsStr)} | result: ${debugTruncate(resultStr)}`);
+  }
+}
+
 /**
  * ProbeAgent class to handle AI interactions with code search capabilities
  */
@@ -193,6 +214,7 @@ export class ProbeAgent {
     this.enableExecutePlan = !!options.enableExecutePlan;
     this.debug = options.debug || process.env.DEBUG === '1';
     this.cancelled = false;
+    this._abortController = new AbortController();
     this.tracer = options.tracer || null;
     this.outline = !!options.outline;
     this.searchDelegate = options.searchDelegate !== undefined ? !!options.searchDelegate : true;
@@ -232,6 +254,7 @@ export class ProbeAgent {
     // Supports exclusion with '!' prefix: ['*', '!bash'] = all tools except bash
     // disableTools is a convenience flag that overrides allowedTools to []
     const effectiveAllowedTools = options.disableTools ? [] : options.allowedTools;
+    this._rawAllowedTools = options.allowedTools; // Keep raw value for explicit tool checks
     this.allowedTools = this._parseAllowedTools(effectiveAllowedTools);
 
     // Storage adapter (defaults to in-memory)
@@ -459,6 +482,17 @@ export class ProbeAgent {
     return mcpToolNames.filter(toolName => this._isMcpToolAllowed(toolName));
   }
 
+  /**
+   * Check if query tool was explicitly listed in allowedTools (not via wildcard).
+   * Query (ast-grep) is excluded by default because models struggle with AST pattern syntax.
+   * @returns {boolean}
+   * @private
+   */
+  _isQueryExplicitlyAllowed() {
+    if (!this._rawAllowedTools) return false;
+    return Array.isArray(this._rawAllowedTools) && this._rawAllowedTools.includes('query');
+  }
+
   /**
    * Check if tracer is AppTracer (expects sessionId as first param) vs SimpleAppTracer
    * @returns {boolean} - True if tracer is AppTracer style (requires sessionId)
@@ -792,6 +826,7 @@ export class ProbeAgent {
       searchDelegateProvider: this.searchDelegateProvider,
       searchDelegateModel: this.searchDelegateModel,
       delegationManager: this.delegationManager,  // Per-instance delegation limits
+      parentAbortSignal: this._abortController.signal,  // Propagate cancellation to delegations
       outputBuffer: this._outputBuffer,
       concurrencyLimiter: this.concurrencyLimiter,  // Global AI concurrency limiter
       isToolAllowed,
@@ -814,7 +849,9 @@ export class ProbeAgent {
     if (wrappedTools.searchToolInstance && isToolAllowed('search')) {
       this.toolImplementations.search = wrappedTools.searchToolInstance;
     }
-    if (wrappedTools.queryToolInstance && isToolAllowed('query')) {
+    // query tool (ast-grep) is not exposed to AI by default — models struggle with AST pattern syntax.
+    // Only register it when explicitly listed in allowedTools (not via wildcard '*').
+    if (wrappedTools.queryToolInstance && isToolAllowed('query') && this._isQueryExplicitlyAllowed()) {
       this.toolImplementations.query = wrappedTools.queryToolInstance;
     }
     if (wrappedTools.extractToolInstance && isToolAllowed('extract')) {
@@ -1362,6 +1399,19 @@ export class ProbeAgent {
     const controller = new AbortController();
     const timeoutState = { timeoutId: null };
 
+    // Link agent-level abort to this operation's controller
+    // so that cancel() / cleanup() stops the current streamText call
+    if (this._abortController.signal.aborted) {
+      controller.abort();
+    } else {
+      const onAgentAbort = () => controller.abort();
+      this._abortController.signal.addEventListener('abort', onAgentAbort, { once: true });
+      // Clean up listener when this controller aborts (from any source)
+      controller.signal.addEventListener('abort', () => {
+        this._abortController.signal.removeEventListener('abort', onAgentAbort);
+      }, { once: true });
+    }
+
     // Set up overall operation timeout (default 5 minutes)
     if (this.maxOperationTimeout && this.maxOperationTimeout > 0) {
       timeoutState.timeoutId = setTimeout(() => {
@@ -1729,7 +1779,8 @@ export class ProbeAgent {
                 allowEdit: this.allowEdit,
                 allowedTools: allowedToolsForDelegate,
                 debug: this.debug,
-                tracer: this.tracer
+                tracer: this.tracer,
+                parentAbortSignal: this._abortController.signal
               };
 
               if (this.debug) {
@@ -1971,12 +2022,15 @@ export class ProbeAgent {
     const toolMap = {
       search: {
         schema: searchSchema,
-        description: 'Search code in the repository using keyword queries with Elasticsearch syntax.'
-      },
-      query: {
-        schema: querySchema,
-        description: 'Search code using ast-grep structural pattern matching.'
+        description: this.searchDelegate
+          ? 'Search code in the repository by asking a question. Accepts natural language questions — a subagent breaks them into targeted keyword searches and returns extracted code blocks. Do NOT formulate keyword queries yourself.'
+          : 'Search code in the repository using keyword queries with Elasticsearch syntax. Handles stemming, case-insensitive matching, and camelCase/snake_case splitting automatically — do NOT try keyword variations manually.'
       },
+      // query tool (ast-grep) removed from AI-facing tools — models struggle with pattern syntax
+      // query: {
+      //   schema: querySchema,
+      //   description: 'Search code using ast-grep structural pattern matching.'
+      // },
       extract: {
         schema: extractSchema,
         description: 'Extract code blocks from files based on file paths and optional line numbers.'
@@ -2812,10 +2866,12 @@ 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: 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:
-- search: Find code patterns using semantic search
+${searchToolDesc1}
 - 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`;
 
@@ -2823,19 +2879,21 @@ export class ProbeAgent {
       systemPrompt += `\n- bash: Execute bash commands for system operations`;
     }
 
-    const searchGuidance = this.searchDelegate
-      ? '1. Start with search to retrieve extracted code blocks'
-      : '1. Start with search to find relevant code patterns';
-    const extractGuidance = this.searchDelegate
+    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 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 to get detailed context when needed';
 
     systemPrompt += `\n
 When exploring code:
-${searchGuidance}
-${extractGuidance}
+${searchGuidance1}
+${extractGuidance1}
 3. Prefer focused, specific searches over broad queries
-4. Combine multiple tools to build complete understanding`;
+4. Do NOT repeat the same search or try trivial keyword variations — probe handles stemming and case variations automatically
+5. If 2-3 consecutive searches return no results for a concept, stop searching for it — the term likely does not exist in that codebase
+6. Combine multiple tools to build complete understanding`;
 
     // Add workspace context
     if (this.allowedFolders && this.allowedFolders.length > 0) {
@@ -2874,10 +2932,12 @@ ${extractGuidance}
     }
 
     // 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: 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:
-- search: Find code patterns using semantic search
+${searchToolDesc2}
 - 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`;
 
@@ -2885,19 +2945,21 @@ ${extractGuidance}
       systemPrompt += `\n- bash: Execute bash commands for system operations`;
     }
 
-    const searchGuidance = this.searchDelegate
-      ? '1. Start with search to retrieve extracted code blocks'
-      : '1. Start with search to find relevant code patterns';
-    const extractGuidance = this.searchDelegate
+    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 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 to get detailed context when needed';
 
     systemPrompt += `\n
 When exploring code:
-${searchGuidance}
-${extractGuidance}
+${searchGuidance2}
+${extractGuidance2}
 3. Prefer focused, specific searches over broad queries
-4. Combine multiple tools to build complete understanding`;
+4. Do NOT repeat the same search or try trivial keyword variations — probe handles stemming and case variations automatically
+5. If 2-3 consecutive searches return no results for a concept, stop searching for it — the term likely does not exist in that codebase
+6. Combine multiple tools to build complete understanding`;
 
     // Add workspace context
     if (this.allowedFolders && this.allowedFolders.length > 0) {
@@ -2953,10 +3015,10 @@ ${extractGuidance}
 Follow these instructions carefully:
 1. Analyze the user's request.
 2. Use the available tools step-by-step to fulfill the request.
-3. You should always prefer the search tool for code-related questions.${this.searchDelegate ? ' It already returns extracted code blocks; use extract only to expand context or read full files.' : ' Read full files only if really necessary.'}
+3. You should always prefer the search tool for code-related questions.${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.'}
 4. Ensure to get really deep and understand the full picture before answering.
 5. Once the task is fully completed, use the attempt_completion tool to provide the final result.
-6. Prefer concise and focused search queries. Use specific keywords and phrases to narrow down results.${this.allowEdit ? `
+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.'}${this.allowEdit ? `
 7. When modifying files, choose the appropriate tool:
     - Use 'edit' for all code modifications:
       * PREFERRED: Use start_line (and optionally end_line) for line-targeted editing — this is the safest and most precise approach.${this.hashLines ? ' Use the line:hash references from extract/search output (e.g. "42:ab") for integrity verification.' : ''} Always use extract first to see line numbers${this.hashLines ? ' and hashes' : ''}, then edit by line reference.
@@ -3369,6 +3431,11 @@ Follow these instructions carefully:
         completionAttempted = true;
       }, toolContext);
 
+      if (this.debug) {
+        const toolNames = Object.keys(tools);
+        console.log(`[DEBUG] Agent tools registered (${toolNames.length}): ${toolNames.join(', ')}`);
+      }
+
       let maxResponseTokens = this.maxResponseTokens;
       if (!maxResponseTokens) {
         maxResponseTokens = 4000;
@@ -3418,6 +3485,7 @@ Follow these instructions carefully:
 
               if (this.debug) {
                 console.log(`[DEBUG] Step ${currentIteration}/${maxIterations} finished (reason: ${finishReason}, tools: ${toolResults?.length || 0})`);
+                debugLogToolResults(toolResults);
               }
             }
           };
@@ -3629,6 +3697,7 @@ Double-check your response based on the criteria above. If everything looks good
               }
               if (this.debug) {
                 console.log(`[DEBUG] Completion prompt step finished (reason: ${finishReason}, tools: ${toolResults?.length || 0})`);
+                debugLogToolResults(toolResults);
               }
             }
           };
@@ -4545,6 +4614,11 @@ Convert your previous response content into actual JSON data that follows this s
    * Clean up resources (including MCP connections)
    */
   async cleanup() {
+    // Abort any in-flight operations (delegations, streaming, etc.)
+    if (!this._abortController.signal.aborted) {
+      this._abortController.abort();
+    }
+
     // Clean up MCP bridge
     if (this.mcpBridge) {
       try {
@@ -4574,12 +4648,24 @@ Convert your previous response content into actual JSON data that follows this s
   }
 
   /**
-   * Cancel the current request
+   * Cancel the current request and all in-flight delegations.
+   * Aborts the internal AbortController so streamText, subagents,
+   * and any code checking the signal will stop.
    */
   cancel() {
     this.cancelled = true;
+    this._abortController.abort();
     if (this.debug) {
       console.log(`[DEBUG] Agent cancelled for session ${this.sessionId}`);
     }
   }
+
+  /**
+   * Get the abort signal for this agent.
+   * Delegations and subagents should check this signal.
+   * @returns {AbortSignal}
+   */
+  get abortSignal() {
+    return this._abortController.signal;
+  }
 }

package/build/agent/dsl/environment.js

@@ -284,6 +284,7 @@ export function generateSandboxGlobals(options) {
       results.push(p);
 
       if (executing.size >= mapConcurrency) {
+        console.error(`[map] Concurrency limit reached (${executing.size}/${mapConcurrency}), waiting for a slot...`);
         await Promise.race(executing);
       }
     }