@probelabs/probe 0.6.0-rc278 → 0.6.0-rc279

package/src/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;
@@ -792,6 +814,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,
@@ -1362,6 +1385,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 +1765,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) {
@@ -3369,6 +3406,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 +3460,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 +3672,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 +4589,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 +4623,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/src/delegate.js

@@ -386,12 +386,18 @@ export async function delegate({
 	mcpConfig = null,
 	mcpConfigPath = null,
 	delegationManager = null,  // Optional per-instance manager, falls back to default singleton
-	concurrencyLimiter = null  // Optional global AI concurrency limiter
+	concurrencyLimiter = null,  // Optional global AI concurrency limiter
+	parentAbortSignal = null   // Optional AbortSignal from parent to cancel this delegation
 }) {
 	if (!task || typeof task !== 'string') {
 		throw new Error('Task parameter is required and must be a string');
 	}
 
+	// Check if parent has already been cancelled
+	if (parentAbortSignal?.aborted) {
+		throw new Error('Delegation cancelled: parent operation was aborted');
+	}
+
 	// Support runtime timeout override via environment variables when timeout not explicitly passed
 	// This allows operators to configure delegation timeouts without code changes
 	// Priority: DELEGATION_TIMEOUT_MS (milliseconds) > DELEGATION_TIMEOUT_SECONDS > DELEGATION_TIMEOUT (seconds)
@@ -481,24 +487,47 @@ export async function delegate({
 			console.error(`[DELEGATE] Subagent config: promptType=${promptType}, enableDelegate=false, maxIterations=${remainingIterations}`);
 		}
 
-		// Set up timeout with proper cleanup
-		// TODO: Implement AbortController support in ProbeAgent.answer() for proper cancellation
-		// Current limitation: When timeout occurs, subagent.answer() continues running in background
-		// This is acceptable since:
-		// 1. The promise will eventually resolve/reject and be garbage collected
-		// 2. The delegation slot is properly released on timeout
-		// 3. The parent receives timeout error and can handle it
-		// Future improvement: Add signal parameter to ProbeAgent.answer(task, [], { signal })
+		// Set up timeout and parent abort handling.
+		// When timeout fires or parent aborts, we cancel the subagent so it
+		// stops making API calls and releases resources promptly.
 		const timeoutPromise = new Promise((_, reject) => {
 			timeoutId = setTimeout(() => {
+				subagent.cancel();
 				reject(new Error(`Delegation timed out after ${timeout} seconds`));
 			}, timeout * 1000);
 		});
 
-		// Execute the task with timeout
+		// Listen for parent abort signal
+		let parentAbortHandler;
+		const parentAbortPromise = new Promise((_, reject) => {
+			if (parentAbortSignal) {
+				if (parentAbortSignal.aborted) {
+					subagent.cancel();
+					reject(new Error('Delegation cancelled: parent operation was aborted'));
+					return;
+				}
+				parentAbortHandler = () => {
+					subagent.cancel();
+					reject(new Error('Delegation cancelled: parent operation was aborted'));
+				};
+				parentAbortSignal.addEventListener('abort', parentAbortHandler, { once: true });
+			}
+		});
+
+		// Execute the task with timeout and parent abort
 		const answerOptions = schema ? { schema } : undefined;
 		const answerPromise = answerOptions ? subagent.answer(task, [], answerOptions) : subagent.answer(task);
-		const response = await Promise.race([answerPromise, timeoutPromise]);
+		const racers = [answerPromise, timeoutPromise];
+		if (parentAbortSignal) racers.push(parentAbortPromise);
+		let response;
+		try {
+			response = await Promise.race(racers);
+		} finally {
+			// Clean up parent abort listener to prevent memory leaks
+			if (parentAbortHandler && parentAbortSignal) {
+				parentAbortSignal.removeEventListener('abort', parentAbortHandler);
+			}
+		}
 
 		// Clear timeout immediately after race completes to prevent memory leak
 		// Note: timeoutId is always set by this point (synchronous in Promise constructor)

package/src/tools/analyzeAll.js

@@ -192,7 +192,8 @@ Instructions:
 			promptType: 'code-researcher',
 			allowedTools: ['extract'],
 			maxIterations: 5,
-			delegationManager: options.delegationManager  // Per-instance delegation limits
+			delegationManager: options.delegationManager,  // Per-instance delegation limits
+			parentAbortSignal: options.parentAbortSignal || null
 			// timeout removed - inherit default from delegate (300s)
 		});
 
@@ -328,7 +329,8 @@ Organize all findings into clear categories with items listed under each.${compl
 			promptType: 'code-researcher',
 			allowedTools: [],
 			maxIterations: 5,
-			delegationManager: options.delegationManager  // Per-instance delegation limits
+			delegationManager: options.delegationManager,  // Per-instance delegation limits
+			parentAbortSignal: options.parentAbortSignal || null
 			// timeout removed - inherit default from delegate (300s)
 		});
 
@@ -404,7 +406,8 @@ CRITICAL: Do NOT guess keywords. Actually run searches and see what returns resu
 			promptType: 'code-researcher',
 			// Full tool access for exploration and experimentation
 			maxIterations: 15,
-			delegationManager: options.delegationManager  // Per-instance delegation limits
+			delegationManager: options.delegationManager,  // Per-instance delegation limits
+			parentAbortSignal: options.parentAbortSignal || null
 			// timeout removed - inherit default from delegate (300s)
 		});
 
@@ -475,7 +478,8 @@ When done, use the attempt_completion tool with your answer as the result.`;
 			promptType: 'code-researcher',
 			allowedTools: [],
 			maxIterations: 5,
-			delegationManager: options.delegationManager  // Per-instance delegation limits
+			delegationManager: options.delegationManager,  // Per-instance delegation limits
+			parentAbortSignal: options.parentAbortSignal || null
 			// timeout removed - inherit default from delegate (300s)
 		});
 

package/src/tools/vercel.js

@@ -277,7 +277,8 @@ export const searchTool = (options = {}) => {
 					promptType: 'code-searcher',
 					allowedTools: ['search', 'extract', 'listFiles', 'attempt_completion'],
 					searchDelegate: false,
-					schema: CODE_SEARCH_SCHEMA
+					schema: CODE_SEARCH_SCHEMA,
+					parentAbortSignal: options.parentAbortSignal || null
 				});
 
 				const delegateResult = options.tracer?.withSpan
@@ -581,7 +582,7 @@ export const delegateTool = (options = {}) => {
 		name: 'delegate',
 		description: delegateDescription,
 		inputSchema: delegateSchema,
-		execute: async ({ task, currentIteration, maxIterations, parentSessionId, path, provider, model, tracer, searchDelegate }) => {
+		execute: async ({ task, currentIteration, maxIterations, parentSessionId, path, provider, model, tracer, searchDelegate, parentAbortSignal }) => {
 			// Validate required parameters - throw errors for consistency
 			if (!task || typeof task !== 'string') {
 				throw new Error('Task parameter is required and must be a non-empty string');
@@ -673,7 +674,8 @@ export const delegateTool = (options = {}) => {
 				enableMcp,
 				mcpConfig,
 				mcpConfigPath,
-				delegationManager  // Per-instance delegation limits
+				delegationManager,  // Per-instance delegation limits
+				parentAbortSignal
 			});
 
 			return result;
@@ -733,7 +735,8 @@ export const analyzeAllTool = (options = {}) => {
 					provider: options.provider,
 					model: options.model,
 					tracer: options.tracer,
-					delegationManager  // Per-instance delegation limits
+					delegationManager,  // Per-instance delegation limits
+					parentAbortSignal: options.parentAbortSignal || null
 				});
 
 				return result;