@probelabs/probe 0.6.0-rc295 → 0.6.0-rc297

package/build/agent/mcp/client.js

@@ -58,6 +58,11 @@ export function isMethodAllowed(methodName, allowedMethods, blockedMethods) {
 export function createTransport(serverConfig) {
   const { transport, command, args, url, env } = serverConfig;
 
+  // Allow pre-created transport instances (e.g., InMemoryTransport for testing)
+  if (serverConfig.transportInstance) {
+    return serverConfig.transportInstance;
+  }
+
   switch (transport) {
     case 'stdio':
       return new StdioClientTransport({
@@ -159,6 +164,8 @@ export class MCPClientManager {
     this.debug = options.debug || process.env.DEBUG_MCP === '1';
     this.config = null;
     this.tracer = options.tracer || null;
+    // Optional event emitter for broadcasting tool call lifecycle to the agent (#522)
+    this.agentEvents = options.agentEvents || null;
   }
 
   /**
@@ -447,6 +454,7 @@ export class MCPClientManager {
     }
 
     const startTime = Date.now();
+    const toolCallId = `mcp-${toolName}-${startTime}`;
 
     // Record tool call start
     this.recordMcpEvent('tool.call_started', {
@@ -455,6 +463,17 @@ export class MCPClientManager {
       originalToolName: tool.originalName
     });
 
+    // Emit toolCall event so the agent's activeTools map tracks MCP tool calls (#522)
+    if (this.agentEvents) {
+      this.agentEvents.emit('toolCall', {
+        toolCallId,
+        name: toolName,
+        args,
+        status: 'started',
+        timestamp: new Date().toISOString(),
+      });
+    }
+
     try {
       if (this.debug) {
         console.error(`[MCP DEBUG] Calling ${toolName} with args:`, JSON.stringify(args, null, 2));
@@ -497,6 +516,16 @@ export class MCPClientManager {
         durationMs
       });
 
+      // Emit toolCall completion so agent's activeTools removes this entry (#522)
+      if (this.agentEvents) {
+        this.agentEvents.emit('toolCall', {
+          toolCallId,
+          name: toolName,
+          status: 'completed',
+          timestamp: new Date().toISOString(),
+        });
+      }
+
       return result;
     } catch (error) {
       const durationMs = Date.now() - startTime;
@@ -516,10 +545,71 @@ export class MCPClientManager {
         isTimeout: error.message.includes('timeout')
       });
 
+      // Emit toolCall error so agent's activeTools removes this entry (#522)
+      if (this.agentEvents) {
+        this.agentEvents.emit('toolCall', {
+          toolCallId,
+          name: toolName,
+          status: 'error',
+          timestamp: new Date().toISOString(),
+        });
+      }
+
       throw error;
     }
   }
 
+  /**
+   * Call graceful_stop on all MCP servers that expose it.
+   * This signals agent-type MCP servers to wrap up their work.
+   * @returns {Promise<Array<{server: string, success: boolean, error?: string}>>}
+   */
+  async callGracefulStopAll() {
+    const results = [];
+    for (const [serverName, clientInfo] of this.clients) {
+      // Look for a graceful_stop tool on this server (qualified name: serverName_graceful_stop)
+      const qualifiedName = `${serverName}_graceful_stop`;
+      if (this.tools.has(qualifiedName)) {
+        if (this.debug) {
+          console.log(`[DEBUG] MCP callGracefulStopAll: calling graceful_stop on server "${serverName}"`);
+        }
+        try {
+          // Short timeout — this is a signal, not a long operation
+          const timeoutMs = 5000;
+          const timeoutPromise = new Promise((_, reject) =>
+            setTimeout(() => reject(new Error('graceful_stop timeout')), timeoutMs)
+          );
+          await Promise.race([
+            clientInfo.client.callTool({ name: 'graceful_stop', arguments: {} }, undefined, { timeout: timeoutMs }),
+            timeoutPromise
+          ]);
+          results.push({ server: serverName, success: true });
+          if (this.debug) {
+            console.log(`[DEBUG] MCP callGracefulStopAll: server "${serverName}" acknowledged graceful_stop`);
+          }
+        } catch (e) {
+          results.push({ server: serverName, success: false, error: e.message });
+          if (this.debug) {
+            console.log(`[DEBUG] MCP callGracefulStopAll: server "${serverName}" graceful_stop failed: ${e.message}`);
+          }
+        }
+      }
+    }
+    if (this.debug) {
+      const withStop = results.length;
+      const total = this.clients.size;
+      console.log(`[DEBUG] MCP callGracefulStopAll: ${withStop}/${total} servers had graceful_stop tool`);
+    }
+    // Record telemetry event for the graceful_stop sweep
+    this.recordMcpEvent('graceful_stop.sweep_completed', {
+      servers_total: this.clients.size,
+      servers_with_graceful_stop: results.length,
+      servers_acknowledged: results.filter(r => r.success).length,
+      servers_failed: results.filter(r => !r.success).length,
+    });
+    return results;
+  }
+
   /**
    * Get all available tools with their schemas
    * @returns {Object} Map of tool name to tool definition
@@ -550,11 +640,32 @@ export class MCPClientManager {
         inputSchema: tool.inputSchema,
         execute: async (args) => {
           const result = await this.callTool(name, args);
-          // Extract text content from MCP response
-          if (result.content && result.content[0]) {
-            return result.content[0].text;
+          if (!result.content || !result.content[0]) {
+            return JSON.stringify(result);
+          }
+          // Check if response contains image content blocks
+          const hasImage = result.content.some(block => block.type === 'image');
+          if (hasImage) {
+            // Return the full content array so toModelOutput can convert it
+            return { _mcpContent: result.content };
+          }
+          // Text-only: return just the text
+          return result.content[0].text;
+        },
+        // Convert MCP content blocks (including images) to Vercel AI SDK format
+        toModelOutput: ({ output }) => {
+          if (output && typeof output === 'object' && output._mcpContent) {
+            const parts = [];
+            for (const block of output._mcpContent) {
+              if (block.type === 'text') {
+                parts.push({ type: 'text', text: block.text });
+              } else if (block.type === 'image') {
+                parts.push({ type: 'image-data', data: block.data, mediaType: block.mimeType });
+              }
+            }
+            return { type: 'content', value: parts };
           }
-          return JSON.stringify(result);
+          return { type: 'text', value: typeof output === 'string' ? output : JSON.stringify(output) };
         }
       };
     }

package/build/agent/mcp/xmlBridge.js

@@ -36,6 +36,7 @@ export class MCPXmlBridge {
   constructor(options = {}) {
     this.debug = options.debug || false;
     this.tracer = options.tracer || null;
+    this.agentEvents = options.agentEvents || null;
     this.mcpTools = {};
     this.mcpManager = null;
     this.toolDescriptions = {};
@@ -84,7 +85,7 @@ export class MCPXmlBridge {
         console.error('[MCP DEBUG] Initializing MCP client manager...');
       }
 
-      this.mcpManager = new MCPClientManager({ debug: this.debug, tracer: this.tracer });
+      this.mcpManager = new MCPClientManager({ debug: this.debug, tracer: this.tracer, agentEvents: this.agentEvents });
       const result = await this.mcpManager.initialize(mcpConfigs);
 
       // Get tools from the manager (already in Vercel format)
@@ -145,6 +146,17 @@ export class MCPXmlBridge {
     return toolName in this.mcpTools;
   }
 
+  /**
+   * Call graceful_stop on all MCP servers that expose it.
+   * @returns {Promise<Array>}
+   */
+  async callGracefulStopAll() {
+    if (this.mcpManager) {
+      return this.mcpManager.callGracefulStopAll();
+    }
+    return [];
+  }
+
   /**
    * Clean up MCP connections
    */

package/build/agent/otelLogBridge.js

@@ -0,0 +1,184 @@
+/**
+ * OTEL Log Bridge — patches console.log/info/warn/error to:
+ * 1. Append trace context [trace_id=... span_id=...] to output (like visor2)
+ * 2. Emit each log as an OTEL Log Record via @opentelemetry/api-logs
+ *
+ * Lazy-loads @opentelemetry/api and @opentelemetry/api-logs.
+ * If packages are not installed, patching is a no-op.
+ *
+ * Usage:
+ *   import { patchConsole, unpatchConsole } from './otelLogBridge.js';
+ *   patchConsole();  // Call once at startup
+ */
+
+import { createRequire } from 'module';
+
+// createRequire for loading optional @opentelemetry packages in ESM context
+const _require = createRequire(import.meta.url);
+
+// OTel severity mapping (OpenTelemetry SeverityNumber values)
+const OTEL_SEVERITY = {
+  log: 9,     // INFO
+  info: 9,    // INFO
+  warn: 13,   // WARN
+  error: 17,  // ERROR
+  debug: 5,   // DEBUG
+};
+
+// Track patch state
+let patched = false;
+const originals = {};
+
+// Lazy-loaded OTel references
+let otelApi = null;
+let otelApiAttempted = false;
+let otelLogger = null;
+let otelLoggerAttempted = false;
+
+/**
+ * Try to load @opentelemetry/api lazily.
+ * Returns { trace, context } or null if not available.
+ */
+function getOtelApi() {
+  if (otelApiAttempted) return otelApi;
+  otelApiAttempted = true;
+  try {
+    // Dynamic require wrapped in IIFE to prevent bundler from resolving
+    otelApi = (function(name) { return _require(name); })('@opentelemetry/api');
+  } catch {
+    // @opentelemetry/api not installed
+  }
+  return otelApi;
+}
+
+/**
+ * Try to get an OTEL Logger from @opentelemetry/api-logs lazily.
+ * Returns a logger instance or null if not available.
+ */
+function getOtelLogger() {
+  if (otelLoggerAttempted) return otelLogger;
+  otelLoggerAttempted = true;
+  try {
+    const { logs } = (function(name) { return _require(name); })('@opentelemetry/api-logs');
+    otelLogger = logs.getLogger('probe-agent');
+  } catch {
+    // @opentelemetry/api-logs not installed
+  }
+  return otelLogger;
+}
+
+/**
+ * Extract trace context suffix from the active OTel span.
+ * Returns '' if no active span or OTel is not available.
+ */
+function getTraceSuffix() {
+  try {
+    const api = getOtelApi();
+    if (!api) return '';
+    const span = api.trace.getSpan(api.context.active());
+    const ctx = span?.spanContext?.();
+    if (!ctx?.traceId) return '';
+    return ` [trace_id=${ctx.traceId} span_id=${ctx.spanId}]`;
+  } catch {
+    return '';
+  }
+}
+
+/**
+ * Emit a log record to the OTEL Logs pipeline.
+ * Non-blocking, best-effort — errors are silently ignored.
+ */
+function emitOtelLog(msg, level) {
+  try {
+    const logger = getOtelLogger();
+    if (!logger) return;
+
+    const api = getOtelApi();
+    let traceId, spanId;
+    if (api) {
+      const span = api.trace.getSpan(api.context.active());
+      const ctx = span?.spanContext?.();
+      if (ctx?.traceId) {
+        traceId = ctx.traceId;
+        spanId = ctx.spanId;
+      }
+    }
+
+    logger.emit({
+      severityNumber: OTEL_SEVERITY[level] || 9,
+      severityText: level.toUpperCase(),
+      body: msg,
+      attributes: {
+        'probe.logger': true,
+        ...(traceId ? { trace_id: traceId, span_id: spanId } : {}),
+      },
+    });
+  } catch {
+    // OTel logs not available; ignore
+  }
+}
+
+/**
+ * Patch console.log/info/warn/error to:
+ * - Append trace context to output
+ * - Emit OTEL log records
+ *
+ * Safe to call multiple times — only patches once.
+ */
+export function patchConsole() {
+  if (patched) return;
+
+  const methods = ['log', 'info', 'warn', 'error'];
+  const c = globalThis.console;
+
+  for (const m of methods) {
+    const orig = c[m].bind(c);
+    originals[m] = orig;
+
+    c[m] = (...args) => {
+      // Build the message string for OTEL log emission
+      const msgParts = args.map(a =>
+        typeof a === 'string' ? a : (a instanceof Error ? a.message : JSON.stringify(a))
+      );
+      const msg = msgParts.join(' ');
+
+      // Emit to OTEL Logs pipeline (non-blocking, best-effort)
+      emitOtelLog(msg, m === 'log' ? 'log' : m);
+
+      // Append trace context suffix to console output
+      const suffix = getTraceSuffix();
+      if (suffix) {
+        if (typeof args[0] === 'string') {
+          args[0] = args[0] + suffix;
+        } else {
+          args.push(suffix);
+        }
+      }
+
+      return orig(...args);
+    };
+  }
+
+  patched = true;
+}
+
+/**
+ * Restore original console methods.
+ * Useful for testing or cleanup.
+ */
+export function unpatchConsole() {
+  if (!patched) return;
+
+  const c = globalThis.console;
+  for (const [m, orig] of Object.entries(originals)) {
+    c[m] = orig;
+  }
+  patched = false;
+}
+
+/**
+ * Check if console is currently patched.
+ */
+export function isConsolePatched() {
+  return patched;
+}

package/build/agent/simpleTelemetry.js

@@ -1,5 +1,6 @@
 import { existsSync, mkdirSync, createWriteStream } from 'fs';
 import { dirname } from 'path';
+import { patchConsole } from './otelLogBridge.js';
 
 /**
  * Simple telemetry implementation for probe-agent
@@ -512,5 +513,12 @@ export function initializeSimpleTelemetryFromOptions(options) {
     filePath: options.traceFile || './traces.jsonl'
   });
 
+  // Patch console methods to append trace context and emit OTEL log records.
+  // This bridges ALL console.log/info/warn/error calls to the OTEL Logs pipeline
+  // automatically — no code changes needed at existing call sites.
+  // Safe to call multiple times (idempotent). No-op if @opentelemetry/api-logs
+  // is not installed.
+  patchConsole();
+
   return telemetry;
 }

package/build/delegate.js

@@ -397,7 +397,14 @@ export async function delegate({
 	mcpConfigPath = null,
 	delegationManager = null,  // Optional per-instance manager, falls back to default singleton
 	concurrencyLimiter = null,  // Optional global AI concurrency limiter
-	parentAbortSignal = null   // Optional AbortSignal from parent to cancel this delegation
+	parentAbortSignal = null,  // Optional AbortSignal from parent to cancel this delegation
+	// Timeout settings inherited from parent agent
+	timeoutBehavior = undefined,
+	requestTimeout = undefined,
+	gracefulTimeoutBonusSteps = undefined,
+	// Subagent lifecycle callbacks for graceful stop coordination
+	onSubagentCreated = null,
+	onSubagentCompleted = null,
 }) {
 	if (!task || typeof task !== 'string') {
 		throw new Error('Task parameter is required and must be a string');
@@ -489,12 +496,38 @@ export async function delegate({
 			enableMcp,   // Inherit from parent (subagent creates own MCPXmlBridge)
 			mcpConfig,   // Inherit from parent
 			mcpConfigPath, // Inherit from parent
-			concurrencyLimiter // Inherit global AI concurrency limiter
+			concurrencyLimiter, // Inherit global AI concurrency limiter
+			// Inherit timeout behavior from parent — subagent gets its own graceful wind-down
+			// so it can produce partial results instead of being hard-killed by the external timer.
+			// The external delegate timeout (capped to parent's remaining budget) is the hard limit;
+			// maxOperationTimeout on the subagent is set slightly shorter so its own wind-down
+			// fires before the external kill.
+			maxOperationTimeout: Math.max(10000, (timeout * 1000) - 15000), // 15s before external kill
+			timeoutBehavior: timeoutBehavior || 'graceful',
+			requestTimeout,
+			gracefulTimeoutBonusSteps: gracefulTimeoutBonusSteps ?? 2, // fewer steps for subagents
 		});
 
+		// Register subagent with parent for graceful stop coordination
+		if (onSubagentCreated) {
+			onSubagentCreated(sessionId, subagent);
+		}
+
 		if (debug) {
 			console.error(`[DELEGATE] Created subagent with session ${sessionId}`);
 			console.error(`[DELEGATE] Subagent config: promptType=${promptType}, enableDelegate=false, maxIterations=${remainingIterations}`);
+			console.error(`[DELEGATE] Timeout inheritance: externalTimeout=${timeout}s, maxOperationTimeout=${Math.max(10000, (timeout * 1000) - 15000)}ms, behavior=${timeoutBehavior || 'graceful'}, bonusSteps=${gracefulTimeoutBonusSteps ?? 2}`);
+		}
+		if (tracer) {
+			tracer.addEvent('delegation.subagent_created', {
+				'delegation.session_id': sessionId,
+				'delegation.parent_session_id': parentSessionId,
+				'delegation.external_timeout_s': timeout,
+				'delegation.internal_timeout_ms': Math.max(10000, (timeout * 1000) - 15000),
+				'delegation.timeout_behavior': timeoutBehavior || 'graceful',
+				'delegation.bonus_steps': gracefulTimeoutBonusSteps ?? 2,
+				'delegation.max_iterations': remainingIterations,
+			});
 		}
 
 		// Set up timeout and parent abort handling.
@@ -507,8 +540,11 @@ export async function delegate({
 			}, timeout * 1000);
 		});
 
-		// Listen for parent abort signal
+		// Listen for parent abort signal — use two-phase shutdown:
+		// Phase 1: Trigger graceful wind-down so subagent can summarize its work
+		// Phase 2: Hard cancel after deadline if subagent hasn't finished
 		let parentAbortHandler;
+		let parentAbortHardCancelId = null;
 		const parentAbortPromise = new Promise((_, reject) => {
 			if (parentAbortSignal) {
 				if (parentAbortSignal.aborted) {
@@ -517,8 +553,33 @@ export async function delegate({
 					return;
 				}
 				parentAbortHandler = () => {
-					subagent.cancel();
-					reject(new Error('Delegation cancelled: parent operation was aborted'));
+					// Phase 1: graceful wind-down — let subagent finish its current step
+					subagent.triggerGracefulWindDown();
+					if (debug) {
+						console.error(`[DELEGATE] Parent abort signal received — triggered graceful wind-down on subagent ${sessionId}`);
+					}
+					if (tracer) {
+						tracer.addEvent('delegation.parent_abort_phase1', {
+							'delegation.session_id': sessionId,
+							'delegation.parent_session_id': parentSessionId,
+							'delegation.action': 'graceful_wind_down',
+						});
+					}
+					// Phase 2: hard cancel after 30s if subagent hasn't finished
+					parentAbortHardCancelId = setTimeout(() => {
+						if (debug) {
+							console.error(`[DELEGATE] Graceful wind-down deadline expired — hard cancelling subagent ${sessionId}`);
+						}
+						if (tracer) {
+							tracer.addEvent('delegation.parent_abort_phase2', {
+								'delegation.session_id': sessionId,
+								'delegation.parent_session_id': parentSessionId,
+								'delegation.action': 'hard_cancel',
+							});
+						}
+						subagent.cancel();
+						reject(new Error('Delegation cancelled: parent operation was aborted (graceful wind-down deadline expired)'));
+					}, 30000);
 				};
 				parentAbortSignal.addEventListener('abort', parentAbortHandler, { once: true });
 			}
@@ -533,10 +594,18 @@ export async function delegate({
 		try {
 			response = await Promise.race(racers);
 		} finally {
-			// Clean up parent abort listener to prevent memory leaks
+			// Clean up parent abort listener and hard cancel timer to prevent memory leaks
 			if (parentAbortHandler && parentAbortSignal) {
 				parentAbortSignal.removeEventListener('abort', parentAbortHandler);
 			}
+			if (parentAbortHardCancelId) {
+				clearTimeout(parentAbortHardCancelId);
+				parentAbortHardCancelId = null;
+			}
+			// Unregister subagent from parent
+			if (onSubagentCompleted) {
+				onSubagentCompleted(sessionId);
+			}
 		}
 
 		// Clear timeout immediately after race completes to prevent memory leak

package/build/index.js

@@ -45,7 +45,7 @@ import { createExecutePlanTool, createCleanupExecutePlanTool } from './tools/exe
 import { bashTool } from './tools/bash.js';
 import { editTool, createTool, multiEditTool } from './tools/edit.js';
 import { FileTracker } from './tools/fileTracker.js';
-import { ProbeAgent } from './agent/ProbeAgent.js';
+import { ProbeAgent, ENGINE_ACTIVITY_TIMEOUT_DEFAULT, ENGINE_ACTIVITY_TIMEOUT_MIN, ENGINE_ACTIVITY_TIMEOUT_MAX } from './agent/ProbeAgent.js';
 import { SimpleTelemetry, SimpleAppTracer, initializeSimpleTelemetryFromOptions } from './agent/simpleTelemetry.js';
 import { listFilesToolInstance, searchFilesToolInstance } from './agent/probeTool.js';
 import { StorageAdapter, InMemoryStorageAdapter } from './agent/storage/index.js';
@@ -68,8 +68,12 @@ export {
 	listFilesByLevel,
 	tools,
 	DEFAULT_SYSTEM_MESSAGE,
-	// Export AI Agent (NEW!)
+	// Export AI Agent
 	ProbeAgent,
+	// Export timeout constants
+	ENGINE_ACTIVITY_TIMEOUT_DEFAULT,
+	ENGINE_ACTIVITY_TIMEOUT_MIN,
+	ENGINE_ACTIVITY_TIMEOUT_MAX,
 	// Export storage adapters
 	StorageAdapter,
 	InMemoryStorageAdapter,